Improving Few-Shot Cross-Domain Named Entity Recognition by Instruction Tuning a Word-Embedding based Retrieval Augmented Large Language Model (2024)

Abstract: Few-Shot Cross-Domain NER is the process of leveraging knowledge from data-rich source domains to perform entity recognition on data-scarce target domains. Most previous state-of-the-art (SOTA) approaches use pre-trained language models (PLMs) for cross-domain NER. However, these models are often domain specific. To successfully use these models for new target domains, we need to modify either the model architecture or perform model fine-tuning using data from the new domains. Both of these result in the creation of entirely new NER models for each target domain which is infeasible for practical scenarios. Recently, several works have attempted to use LLMs to solve Few-Shot Cross-Domain NER. However, most of these are either too expensive for practical purposes or struggle to follow LLM prompt instructions. In this paper, we propose IF-WRANER (Instruction Finetuned Word-embedding based Retrieval Augmented large language model for Named Entity Recognition), a retrieval augmented LLM, finetuned for the NER task. By virtue of the regularization techniques used during LLM finetuning and the adoption of word-level embedding over sentence-level embedding during the retrieval of in-prompt examples, IF-WRANER is able to outperform previous SOTA Few-Shot Cross-Domain NER approaches. We have demonstrated the effectiveness of our model by benchmarking its performance on the open source CrossNER dataset, on which it shows more than 2% F1 score improvement over the previous SOTA model. We have deployed the model for multiple customer care domains of an enterprise. Accurate entity prediction through IF-WRANER helps direct customers to automated workflows for the domains, thereby reducing escalations to human agents by almost 15% and leading to millions of dollars in yearly savings for the company.

## introduction named entity recognition (ner) @xcite is a key process in information extraction, designed to identify and categorize entities in natural language into predefined entity types. due to the large variations in entities and the way they are used across domains, ner has been a challenging task in nlp. most traditional ner models require large volumes of labelled data for training @xcite @xcite . however, collecting large volumes of labeled data is both costly and time consuming. therefore, we need a model that can perform ner on multiple domains with minimal labeled examples from that domain. to tackle this problem, several solutions have been proposed, which attempt to transfer knowledge from data rich source domain to perform ner on data-scarce target domain. this is referred to as few-shot cross-domain ner. the traditional way of solving this involves training plms with entity-tagged source domain data, followed by fine-tuning them on target domain data, thereby transferring knowledge from source to target domain. this approach fails to address the semantic gap that may exist between source and target domains. to address this, some previous studies utilize adding auxiliary objects @xcite or designing new model architectures @xcite @xcite to train with both source and target domain data. @xcite with target domain data for a better understanding of domain-specific data. another line of research @xcite focuses on modeling the label relationship across domains to improve label information transfer. specifically, laner @xcite utilizes an architecture to better leverage semantic relationships among domain labels to increase cross-domain performance. most few-shot cross-domain ner models however, have one or both of the following weaknesses: our approach involves a single model architecture, finetuned only using entity tagged source domain data. for adaptation to target domain, our model does not require further fine-tuning. recently, with the advent of generative ai, many researchers have tried to use llms to solve the few shot cross-domain ner problem. gpt-ner @xcite and promptner @xcite have experimented with different llm prompting strategies for the task with varying degrees of success. gpt-ner further demonstrates that using the retrieval augmented generation (rag) @xcite framework to select the in-prompt examples further boosts ner performance. one common theme that has emerged from these works is that most of these approaches demonstrate good performance with gpt4 as the backbone llm. open source alternatives typically fall well short of sota performance as they do not seem to closely follow the prompt instructions. this is a serious problem for real world scenarios as the use of proprietary software like gpt4 can be cost-prohibitive, especially for applications operating at scale. in our approach, we finetune open source llms @xcite , so that they can follow domain specific prompt instructions for the ner task. like gpt-ner @xcite , we too utilise the rag framework for selecting in-prompt examples. we store labelled domain data and their vector embeddings in a vector store and extract relevant domain examples from the store during inference based on similarity between the inference query embedding and the embeddings stored. most applications using the rag framework use sentence level embeddings for similarity score calculations. in our work, we show that for the ner task, retrieving examples based on word-level em-bedding similarity performs much better than that based on sentence-level embedding similarity. ## background and related work traditional approaches to solve the ner problem typically fall into one of the two categories: cross-domain named entity recognition (ner) algorithms @xcite help to address the issue of data scarcity in the target domain by leveraging data from source domains. one commonly used strategy to tackle few-shot cross-domain ner is multitask learning, which involves the use of auxiliary objects @xcite or the development of fresh model structures @xcite @xcite . these methods aim to enhance the ner performance in the target domain by training on data from both the source and target domains. another facet of cross-domain ner research concentrates on the transfer of label information across domains. as noted by @xcite , the relationship of labels can be represented as a probability distribution to facilitate the transfer of cross-domain knowledge in ner more effectively. @xcite suggests a method to capitalise on the semantic relationships between domains more efficiently by utilising previous labels (from the source) and the corresponding token. chen et al., 2023 introduces collaborative prefix tuning as a solution to the cross-domain ner issue. although prefix-tuning is considerably quicker than complete fine-tuning, it still requires the addition and modification of new model parameters. unlike the above approaches, if-wraner does not necessitate architectural alterations or model fine-tuning/prefix-tuning to adapt to new target domains, making it more suitable for practical applications. large language models (llms) @xcite have demonstrated remarkable proficiency in in-context learning, where they can generate results for a new test input using only a handful of task-specific examples. operating under the in-context learning framework, llms have yielded encouraging outcomes across a range of nlp tasks, including machine translation (mt) @xcite , question answering (qa) @xcite etc. for few-shot cross domain ner, promptner @xcite and gpt-ner @xcite have utilized llms, achieving a performance level comparable to the industry benchmark through extensive prompt engineering. gpt-ner has redefined the ner problem from a sequence labeling task to a generation task that llms can easily adapt to. on the other hand, promptner uses the chain-of-thought prompting technique, which offers a precise, adaptable, and user-friendly way to carry out few-shot ner and requires prompting the llm only once. in our approach, we instruct the llm to provide responses in a structured format and both the extraction of entities and their categorization into entity types happens in a single llm call. another recent innovation that has gained tremendous popularity and adoption post the advent of llms is retrieval-augmented generation (rag) @xcite with llms. for few-shot cross-domain ner, we show that using proprietary llm like gpt4 with the rag framework we are able to obtain results comparable to sota models. however, given the cost implications of using gpt4, we have developed a strategy to finetune open-source llms. with our approach of finetuning open-source llm, coupled with regularization techniques and replacing the similarity between sentence level embeddings with similarity between word level embedding as criteria for selecting relevant examples using the retriever, we are able to achieve better performance on most domains compared to previous sota models. ## results in ## model deployment with if-wraner, we have built a model that new domains can use off the shelf, simply by adding entity type definitions and a few labelled examples from the respective domains. we use the tensorrt framework on triton inference server @xcite for serving our model. depending on the traffic and latency requirements for each domain, we create separate instances of our model and serve them with triton. using if-wraner we are able to achieve reasonable latency and throughput numbers on a100 gpus. some domains however, have very low latency requirements and as per our experiments, a 7b-parameter if-wraner cannot meet these latency requirements. for such domains, we create a new model with tinyllama @xcite as the base llm. tinyllama is a 1.1b model with the same architecture and tokenizer as 7b meta ## conclusion in this work, we have introduced if-wraner, a retrieval augmented instruction following llm, that outperforms sota models for few-shot cross-domain ner. unlike many of the models developed for cross-domain ner, we do not need to finetune or make structural modifications to our model to adapt to new domains. also, if-wraner manages to attain sota performance using non-proprietary llm, making it much more cost effective compared to proprietary llm based cross-domain ner models. our model is flexible and can be easily used by end users with no technical expertise. all they have to do is provide definitions for their domain's entity types and a few labelled examples. for serving domains with very low latency requirements, we have proposed tiny-if-wraner which uses tinyllama instead of 7b meta llm as its base llm.