A Foundation Model for Entity Recognition

Entity recognition is a widely used information extraction task, yet publicly available foundation models are not well suited for it. We leverage modern LLMs to create a small-yet-powerful foundation model for this task. This BERT-size model can be used to create custom entity recognizers with typically 5x less annotated data than before. This model is powering NuMind and we open-source it with an MIT license for everyone to use.

Entity Recognition

Entity recognition (a.k.a. NER) is the task of detecting entity types and other concepts mentioned in text — cities, people, companies, or any instance of human concepts. It powers automatic news analysis, medical coding, legal document analysis, and many other applications.

You could try GPT-4 with a good prompt — fine for some cases. But for high volumes, better performance, or confidentiality, the traditional deep learning approach (fine-tuning a BERT-like foundation model on hand-annotated data) is preferable. The catch: you typically need hundreds of annotated documents to reach (and surpass) GPT-4 performance.

Two ways to reduce that effort: annotate automatically with an LLM, or use a better BERT-size foundation model. This post is about the second approach.

RoBERTa Out-of-the-Box

Transformers like BERT/RoBERTa are pre-trained self-supervised on missing-word prediction. The last-layer embeddings carry contextual semantic information — visible when "Amazon-the-river" and "Amazon-the-company" cluster differently.

The classic approach is to attach a linear classifier on top of these embeddings to predict, per token, whether it belongs to a given concept.

We evaluate on MIT Movie, MIT Restaurant, OntoNotes 5, and BioNLP 2004, training only the linear classifier (other layers frozen). This is our baseline.

Leveraging Human-Annotated Datasets

Last-layer embeddings carry contextual word meaning, but there's no reason for them to encode human concepts in a linearly accessible way. We want a foundation model that explicitly knows about human concepts, with this information surfaced in the last layer.

The simple idea: train an entity recognizer on a dataset annotated with a large, diverse set of concepts. Some researchers did this using the NER Corpus (16M examples, 475M tokens, 315 unique concepts from Wikipedia). Fine-tuning the last 6 layers of RoBERTa on it gives a small improvement over the baseline — better in the few-shot regime, equivalent for larger training sets.

A 315-concept set is small. Wikipedia is one domain. We want tens of thousands of concepts on highly diverse data. Human annotation would be too costly. Modern LLMs solve this.

Using LLMs to Annotate Human Concepts

Annotating from a pre-defined ontology gives mediocre results, even with GPT-4. Our key insight: let the LLM figure out the ontology as it annotates, introducing concepts on the fly. This gives both ontology diversity and high annotation quality — even GPT-3.5 works in this setting.

The prompt asks the model to "label as many entities, concepts, and ideas as possible", invent new entity types where needed, and output entity from the text -|- entity concept -|- description.

We applied this to 160k English sentences from the C4 dataset, producing about 800k annotations and 80k unique concepts — long-tailed distribution, with the 100 most common concepts accounting for 43% of annotations and ~50k rare concepts appearing once.

The 100 most common concepts include classics ("person", "location") and richer ones ("medical procedure", "operating system"). 28k unique words appear in concept names. This dataset has much higher concept diversity than any available human-labeled dataset.

Annotations contain mistakes but few false positives — the main issue is false negatives (missed concepts). As we'll see, this isn't a problem for the foundation model.

Learning From LLM-Annotated Concepts

Naive approach (linear classifier on top of RoBERTa) doesn't work — too many concepts, many similar, most rare.

Our trick: instead of independent weight vectors per concept, compute the weight vector as f(concept_name + description) where f is a sentence encoder. This lets the network leverage similarities between concepts and scales to many more concepts.

In practice, we ignore concepts not present in the current batch — we don't learn from "negatives" of other concepts. Probabilities become uncalibrated, but we only care about token embeddings. False negatives in the data also become a non-issue. The setup is essentially contrastive learning.

We use RoBERTa for both networks and fine-tune the last six layers. After training, we keep only the token-encoder network — that's our foundation model.

Results

Compared to RoBERTa-base and RoBERTa-fine-tuned-on-NER-Corpus, our foundation model is much better both in few-shot and as data scales. F1 differences are large, but the better metric is data efficiency:

• ~30 examples per concept needed for previous models to reach F1=0.65
• ~5 examples per concept needed for ours — 6× more data efficient in this regime

Data efficiency keeps increasing with training size.

Per-dataset breakdown shows our model is substantially superior across all datasets and data regimes. Best case: >10× data efficiency on BioNLP2004 / MIT Movie. Even on the favorable MIT Restaurant: 2-3× improvement.

The large gain (compared to the moderate gain from training on NER Corpus) likely comes from the combination of a large diverse concept set, domain-diverse data, and the specific training procedure.

Let's Put It to Work

This foundation model is an important step forward, allowing accurate entity recognizers to be trained with substantially less annotated data. We open-source it under MIT license:

• English model
• Multilingual model

Of course, the best way to use it is through NuMind 🙂.