Source-linked AI summary
NE-R1: Enhancing Named Entity Recognition Model via Reinforcement Learning
Meixuan Chen, Hehan Li, Ruizhi Zhao, Xin Lu, peizhi xu, Liwei Qian, LI Meifang, shuanglong li, Hanmeng Liu, Xin Pei, Yanbiao Ma
TL;DR
LLM-based NER remains challenged by long-tail and domain-specific entities because of deficient parametric knowledge, while always-on retrieval can add noise and unnecessary cost. NE-R1 introduces retrieval-on-demand NER with two-stage training and end-to-end RL using Chain-of-Thought and rewards for accuracy and retrieval benefit. It improves average in-domain F1 by 2.52 points and zero-shot cross-domain F1 by 1.18 points over the strongest baselines, while reducing MIT-Movie latency by approximately 55% versus always-retrieve RAG.
Problem
Parametric-knowledge deficiencies make long-tail, domain-specific, or ambiguous entities difficult for LLM-based NER, while full retrieval can introduce noise and unnecessary latency.
Method
NE-R1 uses retrieval-on-demand with multi-task instruction tuning followed by end-to-end RL optimization using Chain-of-Thought and a multi-dimensional reward.
Results
NE-R1 improves average in-domain F1 by 2.52 points and zero-shot cross-domain F1 by 1.18 points over the strongest baselines, while reducing MIT-Movie latency by approximately 55% versus always-retrieve RAG.
Takeaways & Limitations
NE-R1 combines parametric knowledge for familiar inputs with external evidence for knowledge-intensive or ambiguous entities, with retrieval behavior varying across domains and sample difficulties.
Takeaways & Limitations
Evidence is limited to evaluated English NER benchmarks and the current multi-source Web-search setting, without establishing generalization to other languages, domains, or retrieval backends.
Abstract
from arXiv · showhide
Named Entity Recognition (NER) has achieved substantial progress since the advent of large language models (LLMs). Nevertheless, the recognition of long-tail and domain-specific entities remains challenging due to the deficiency in parametric knowledge. Retrieval-augmented generation (RAG) offers a promising remedy by injecting external knowledge, but it also introduces noise and unnecessary cost when dealing with familiar cases. In this paper, we propose NE-R1, a novel framework for adaptive retrieval-augmented NER. We design a "retrieval-on-demand" mechanism for NER. Then we integrate it into models by a two-stage training method: (1) multi-task instruction tuning initialization; (2) end-to-end RL optimization with CoT. To achieve reasonable selection between parameterized and external knowledge, we design a multi-dimensional reward considering both accuracy and retrieval benefit. NE-R1 achieves state-of-the-art performance on various benchmarks, with an average F1 score gain of 2.52% in in-domain evaluation and 1.18% in zero-shot cross-domain evaluation.
1 Introduction
NE-R1 addresses parametric-knowledge gaps in NER with retrieval only when needed, avoiding the noise and latency of indiscriminate retrieval. Its two-stage training combines capability initialization with end-to-end reinforcement learning, and experiments report gains across in-domain and zero-shot evaluations.
- Long-tail and domain-specific entities expose parametric-knowledge deficiencies in LLM-based NER, causing hallucinations or omissions.
- Full retrieval can inject irrelevant information for familiar cases and increase inference latency without meaningful performance gains.On MIT-Movie, full retrieval increases latency 4.85× versus no retrieval while yielding negligible performance gains.
- NE-R1 triggers retrieval based on input complexity, allowing retrieval only when the model’s own knowledge falls short.
- NE-R1 uses multi-task instruction tuning for capability initialization, followed by end-to-end RL optimization with Chain-of-Thought and a multi-dimensional reward.The initialized capabilities are parametric inference, retrieval triggering, and evidence fusion; the reward considers prediction correctness and retrieval benefit.
- 2.52 points is the average in-domain F1 improvement over the strongest overall baseline, while zero-shot cross-domain performance exceeds the strongest baseline by 1.18 F1 points.The cross-domain result improves performance in four of five CrossNER domains and achieves the best five-domain average.
2 Related Work
Related work spans prompting and supervised fine-tuning for LLM-based NER, retrieval augmentation, and reinforcement learning. NE-R1 is situated at their intersection while addressing the limited adaptability of static retrieval strategies.
- LLM-based NER uses in-context learning through prompt optimization, domain knowledge, reflection mechanisms, and example selection.
- Supervised LLM-based NER methods use unified instruction templates, Chain-of-Thought, task decomposition, and domain-adaptation techniques.
- RAG supplements model knowledge through retrieval optimization, document refinement, and multi-round mechanisms, but most systems use static retrieval strategies.
- Reinforcement learning has been applied to LLMs and extended to machine translation, text-to-SQL, and question answering.
3 Methodology
NE-R1 combines multi-task capability initialization with end-to-end reinforcement learning to support adaptive retrieval-augmented NER. The framework trains parametric inference, retrieval triggering, and evidence fusion, then optimizes their selection using CoT-guided reasoning and multi-dimensional rewards.
- NE-R1 uses two stages: multi-task instruction tuning initializes capabilities, while GRPO-based end-to-end RL optimizes adaptive retrieval and entity extraction.The RL stage uses CoT-guided reasoning and a multi-dimensional reward mechanism.
- Multi-Task Capability Initialization: The initialization stage trains parametric inference, retrieval triggering, and fusion inference as complementary supervised tasks.Parametric inference uses internal knowledge, retrieval triggering decides whether to search and generates a rewritten query, and fusion inference integrates retrieved evidence.
- Multi-Task Capability Initialization: Pass-rate-based selection uses teacher-generated candidate responses to retain task-specific training samples and partition examples into easy and hard groups.The resulting groups define retrieval-benefit rewards during reinforcement learning.
- End-to-End RL Optimization: CoT-guided reasoning branches from introspective planning into either direct answering or search followed by evidence fusion.The model analyzes ambiguity, domain-knowledge needs, and context sufficiency before generating a query or producing entity predictions directly.
- End-to-End RL Optimization: The multi-dimensional reward jointly accounts for format validity, entity-level Micro-F1 accuracy, and retrieval benefit.Retrieval is encouraged on hard instances and discouraged when unnecessary or harmful on easy instances; GRPO optimizes the resulting trajectory-level objective.
4 Experiments
NE-R1 is evaluated across in-domain and zero-shot cross-domain NER benchmarks, with ablations and behavioral analyses examining its performance, efficiency, and adaptive retrieval decisions.
- 4.1 Setup: Evaluation uses four in-domain benchmarks and five zero-shot CrossNER domains, with strict entity-level Micro-F1 as the metric.In-domain datasets cover general, short-text, and specialized domains; cross-domain models are trained on CoNLL-03.
- 4.2 Main Results (In-Domain Evaluation): NE-R1 outperforms comparable and larger baselines across all four in-domain benchmarks, including 81.70 F1 on terminology-dense GENIA.It achieves 91.05 and 85.44 F1 on MIT-Movie and MIT-Restaurant, and 91.89 F1 on OntoNotes 5.0.
- 4.3 Cross-Domain Evaluation: NE-R1 achieves the best five-domain CrossNER average of 78.15 after improving four of five domains, while scoring 1.92 points lower on Music than IF-WRANER.Its gains over IF-WRANER are 1.55, 0.52, 0.78, and 4.97 F1 points on Politics, Natural Science, Literature, and AI.
- 4.4 Ablation Study: Removing either multi-task capability initialization or end-to-end RL causes notable performance degradation, indicating that both training stages contribute to NE-R1.Removing retrieval triggering and fusion initialization produces a 6.8-point F1 drop from 87.52 to 80.73.
- 4.4 Ablation Study: The retrieval-benefit reward is more important than CoT removal in the ablation, supporting its central role in balancing retrieval costs.The ablation compares removing Chain-of-Thought reasoning with removing the Retrieval Benefit Reward.
- 4.5 Analysis of Adaptive Search Behavior: Retrieval frequency rises from 14.25% on OntoNotes 5.0 to 52.18% on GENIA, while NE-R1 reaches 91.05 Micro-F1 at 2.19× relative latency on MIT-Movie.On MIT-Movie, this is approximately 55% lower latency than Standard RAG, which has 4.85× latency and 85.56 Micro-F1.
5 Conclusion
NE-R1 implements retrieval on demand for NER by combining parametric inference with selective external evidence. Across in-domain and zero-shot evaluations, it improves performance while reducing latency relative to always-retrieve RAG.
- 5 Conclusion: NE-R1 combines multi-task initialization with end-to-end reinforcement learning to jointly optimize NER predictions and retrieval decisions.The design initializes parametric inference, retrieval triggering, and evidence fusion before RL optimization.
- 5 Conclusion: NE-R1 relies on parametric knowledge for familiar inputs and retrieves external evidence for knowledge-intensive or ambiguous entities.This is the framework’s retrieval-on-demand behavior.
- 5 Conclusion: NE-R1 improves average F1 by 2.52 points in-domain and 1.18 F1 points in zero-shot cross-domain evaluation while reducing MIT-Movie latency by approximately 55% versus always-retrieve RAG.It improves performance in four of five CrossNER domains and achieves the best five-domain average.
Limitations
NE-R1’s evaluation is constrained by live web-search variability, English benchmark coverage, the current retrieval setting, and limited statistical analysis.
- Live, multi-source Web search can change with API updates, indexing changes, and regional variation, preventing guaranteed exact reproduction.
- Empirical conclusions are limited to evaluated English NER benchmarks and the current multi-source Web-search setting.
- The experiments do not establish generalization to other languages, domains, or retrieval backends.
- Heterogeneous search engines may return multilingual, conflicting, or uneven-quality evidence, whose effects are not separately quantified.
- The framework supports only single-turn retrieval, while uniform run-level variance and paired significance analyses were unavailable across many baselines.
A.1 Dataset Statistics
The experiments use diverse in-domain and cross-domain NER benchmarks spanning standard, biomedical, short-text, and five specialized domains.
- The dataset statistics include train, development, and test sentence counts, entity-type counts, average token length, and entities per sentence.
- In-domain evaluation covers CoNLL-03 and Ontonotes 5.0 as standard or general-domain corpora.
- Ontonotes 5.0 contains over 59k training samples and 18 fine-grained entity types.
- GENIA represents the biomedical domain and has an average of 3.5 entities per sentence.
- MIT-Movie and MIT-Restaurant contain informal user queries with average sentence lengths of approximately 9–11 tokens.
- CrossNER evaluates zero-shot cross-domain performance across AI, Literature, Music, Politics, and Science, using CoNLL-03 for training.
A.2 Implementation Details
NE-R1 uses Qwen-2.5-7B with separate supervised initialization and reinforcement-learning phases, plus multi-source retrieval through a unified interface.
- All experiments use Qwen-2.5-7B as the backbone model on 8 NVIDIA A800 GPUs.
- MTCI/SFT phase: The supervised phase uses AdamW, a 1e-5 learning rate, cosine scheduling, 5 epochs, 0.1 warmup, batch size 16, and bfloat16 precision.
- RL phase: The reinforcement-learning phase uses critic-free GRPO with a 1e-6 learning rate, 500 steps, batch sizes of 512 and 256, clipping ratio 0.2, and KL coefficient 0.001.
- Retrieval module: The retrieval module uses multi-source web search engines and dynamically incorporates results through a unified search interface.
A.3 Supplementary Experiments and Analyses
Supplementary analyses compare NE-R1 with random retrieval policies, document dataset and training configurations, describe the search interface, and discuss retrieval leakage risk.
- Matched-rate random trigger baseline: Random retrieval policies with matched search rates improve over naive decisions in some settings but remain consistently inferior to NE-R1’s adaptive policy.
- Dataset statistics: Table 4 presents dataset statistics, including split sizes, entity-type counts, average token length, and average entities per sentence.
- Hyperparameters: Table 5 details hyperparameter configurations for the MTCI and RL phases.
- Search engines: Table 6 lists the search engines used for retrieval, all accessed through a unified retrieval interface.
- Retrieval leakage risk: NE-R1 mitigates retrieval leakage by using entity-centric rewritten queries, natural-language evidence without BIO labels, and adaptive rather than always-on retrieval.
- Retrieval leakage risk: Zero-shot cross-domain evaluation supports reliance on transferable retrieval and reasoning behavior rather than memorized benchmark-specific labels.
A.4 MIT-Movie Prompt Templates
The MIT-Movie pipeline initializes direct extraction, retrieval triggering, and evidence fusion before unifying them in an adaptive evaluation prompt. Case studies show retrieval used for ambiguity while avoiding it for clear cases, though false-positive and false-negative retrieval remain limitations.
- A.4 MIT-Movie Prompt Templates: The MIT-Movie prompting scheme covers parametric inference, retrieval triggering, fusion inference, and adaptive evaluation, with task descriptions and entity labels adapted across datasets.The four prompts are shown in Figures 6–9.
- A.4 MIT-Movie Prompt Templates: Parametric inference trains direct NER extraction without retrieval, preserving original spans in compact JSON annotations.
- A.4 MIT-Movie Prompt Templates: Retrieval triggering separates clear cases from ambiguous or knowledge-intensive cases and produces a search query only when external evidence is likely to help.
- A.4 MIT-Movie Prompt Templates: Fusion inference uses retrieved snippets for disambiguation while keeping final entity spans grounded in the original query.
- A.4 MIT-Movie Prompt Templates: The evaluation prompt permits either search-and-fusion or direct prediction, while requiring the final answer to be valid JSON inside the answer tag.
- A.5 Case Study: For the ambiguous movie title “Crash,” NE-R1 recognizes insufficient context, retrieves movie information, and correctly outputs the entity as a Title.
- A.5 Case Study: For the clear entity “Steven Spielberg,” NE-R1 directly predicts without retrieval, indicating uncertainty-driven rather than heuristic retrieval decisions.
- A.5 Case Study: The case study includes redundant retrieval on an easy instance and missed retrieval for an ambiguous or knowledge-intensive mention, identifying remaining policy limitations.The false-positive example leaves the final prediction unchanged despite potentially introducing noise.