[Paper Review] NuNER: Entity Recognition Encoder Pre-training via LLM-Annotated Data

Sergei Bogdanov, Alexandre Constantin+3 more|arXiv (Cornell University)|Feb 23, 2024

Data Quality and ManagementDecision Sciences3 citations

TL;DR

This paper proposes NuNER, a task-specific foundation model for Named Entity Recognition (NER) trained by fine-tuning RoBERTa on a 24.4M-word, LLM-annotated dataset derived from C4. NuNER achieves state-of-the-art few-shot performance among similarly sized models, outperforming both its base model and larger models like NER-BERT, while competing with much larger LLMs such as GPT-4 despite being 56× smaller.

ABSTRACT

Large Language Models (LLMs) have shown impressive abilities in data annotation, opening the way for new approaches to solve classic NLP problems. In this paper, we show how to use LLMs to create NuNER, a compact language representation model specialized in the Named Entity Recognition (NER) task. NuNER can be fine-tuned to solve downstream NER problems in a data-efficient way, outperforming similar-sized foundation models in the few-shot regime and competing with much larger LLMs. We find that the size and entity-type diversity of the pre-training dataset are key to achieving good performance. We view NuNER as a member of the broader family of task-specific foundation models, recently unlocked by LLMs.

Motivation & Objective

• To develop a data-efficient, task-specific foundation model for Named Entity Recognition (NER) that reduces reliance on costly human-annotated data.
• To investigate whether LLM-annotated data can effectively pre-train a compact encoder model for downstream NER tasks.
• To evaluate the performance of the resulting model, NuNER, in few-shot and zero-shot regimes compared to large language models and existing NER-specific models.
• To identify key factors influencing the performance of such task-specific foundation models, particularly dataset size and entity-type diversity.

Proposed method

• Annotate a subset of the C4 dataset (24.4M words) using GPT-3.5 to generate entity annotations across 200k unique concepts.
• Fine-tune a RoBERTa-base model on the LLM-annotated dataset using a contrastive learning objective to improve representation learning.
• Train a downstream classifier head via fine-tuning on various NER benchmarks using few-shot and zero-shot protocols.
• Compare NuNER’s performance to GPT-3.5, GPT-4, and UniversalNER via in-context learning and fine-tuning across multiple datasets.
• Systematically ablate the impact of dataset size, entity-type diversity, and text diversity on model performance.
• Use a two-layer feedforward network with dropout as the classification head, trained for 30 epochs during fine-tuning.

Figure 1: NuNER creation procedure. RoBERTa is further pre-trained on a subset of C4 automatically annotated by GPT-3.5. The resulting model can be fine-tuned on various downstream NER problems.

Experimental results

Research questions

• RQ1 Can LLM-annotated data effectively pre-train a compact, task-specific foundation model for NER?
• RQ2 How does the performance of NuNER compare to large language models like GPT-3.5 and GPT-4 in few-shot NER settings?
• RQ3 What factors—dataset size, entity-type diversity, or text diversity—most significantly influence the performance of a task-specific foundation model?
• RQ4 Can a smaller, encoder-based model like NuNER match or exceed the performance of larger, generative models such as UniversalNER?
• RQ5 Does the pre-training procedure of NuNER enhance the accessibility of entity representations in the final network layers, improving few-shot fine-tuning?

Key findings

• NuNER outperforms its base RoBERTa model and the similarly sized NER-BERT pre-trained model in the few-shot regime, demonstrating the effectiveness of LLM-annotated data.
• NuNER surpasses GPT-3.5 in performance when more than 8 words per entity type are used in fine-tuning, and competes with GPT-4, suggesting limitations in in-context learning for large models.
• Despite being 56× smaller, NuNER achieves comparable performance to UniversalNER (70.30±0.35 F1-score on $64\sim128$ words) in the $k\sim2k$ few-shot setting.
• Entity-type diversity and dataset size are the most influential factors for NuNER’s performance, while text diversity shows minimal impact.
• The contrastive pre-training procedure enhances the emergence of human concepts in the final network layers, potentially improving few-shot fine-tuning efficiency.
• NuNER serves as a strong drop-in replacement for RoBERTa in NER tasks, offering better few-shot generalization with minimal computational cost.

Figure 4: Frequency of each concept assigned by GPT-3.5, sorted from most to least common. We observe a heavy-tailed distribution.