Unified retrieval method improves entity recognition for small AI models

URA-NER: A Unified Retrieval-Augmented Framework with Retrieval Alignment and Uncertainty Reduction for Low-Resource NER

Computation and LanguageArtificial Intelligence

Summary

Named Entity Recognition (NER) systems identify important names or terms in texts but usually need lots of labeled examples to work well. The authors developed URA-NER, a new method that helps smaller language AI models do this better with only a few examples by carefully choosing helpful examples and reducing confusion during the process. Their approach involves two stages of finding examples, improving how AI understands these examples, and double-checking the AI's reasoning with external knowledge. Tests showed this method significantly boosts recognition accuracy on tasks where training data is scarce, especially benefiting smaller AI models.

What this means in practice

  • For natural language processing engineers: Improve NER model accuracy in domains with limited labeled data by integrating URA-NER’s retrieval and verification techniques.
  • For medical data analysts: Enhance entity recognition in specialized medical texts with few annotations by applying the retrieval alignment and uncertainty reduction in URA-NER.

Authors

Jingyu Wang, Shijie Wu, Fusheng Jin

Abstract

In-context learning (ICL) based on large language models (LLMs) has shown promising potential in alleviating performance bottlenecks caused by the limited availability of annotated data in Named Entity Recognition (NER). However, existing methods still face issues of retrieval misalignment and generation uncertainty, making their performance heavily dependent on the LLM's capabilities. As the parameter scale of LLMs decreases, their performance in few-shot settings deteriorates significantly. In this paper, we propose a novel unified retrieval-augmented framework, URA-NER, including three key components: Progressive Granularity Retrieval (PGR), Model-aware Representation Enhancement (MaRE), and Reason-aware Knowledge Verification. PGR is a two-stage retrieval mechanism that achieves stage alignment. It first retrieves demonstrations for span detection based on the query's global semantics, and then for type classification based on the specific entity context, providing fine-grained local information. Moreover, MaRE employs entity pre-recognition to guide the construction of representations, ensuring the query and demonstrations are aligned within the LLM's semantic space and attention pattern. In addition, to mitigate generation uncertainty, we propose RaKV, a closed-loop "generation-retrieval-verification" process. It explicates the LLM's reasoning paths, leverages them for the retrieval of external knowledge, and reorganizes the knowledge into verification evidence aligned with the original reasoning paths. We conduct extensive experiments on multiple low-resource NER datasets. Results demonstrate that URA-NER significantly enhances the performance of LLMs under low-resource settings, with particularly pronounced gains for smaller LLMs, achieving new state-of-the-art results on several benchmarks.