Improving Few-shot Prompting using Cluster-based Sample Retrieval for Medical NER in Clinical Text (2024)

Abstract: Named Entity Recognition (NER) in the medical domain is crucial for extracting essential information from clinical text. Large Language Models (LLMs) have demonstrated remarkable capabilities in this task, but their performance is highly dependent on the quality of the prompts. Few-shot prompting or prompt-by-example, where the input query to LLM is augmented with one or more sample outputs, is a well-known technique in guiding the LLMs to the expected result. The quality of the sample in the prompt plays an important role in this task. This paper proposes to improve the performance of few-shot prompting for medical NER on clinical text using a cluster-based strategy for sample selection. We employ the concepts from Retrieval Augmented Generation (RAG) and K-means clustering to identify the most similar annotated examples for any given input text. Using these contextually relevant yet divergent training samples as examples, we guide the LLM toward extracting more accurate medical entities. Our experiments using the llama-2 model show that this approach significantly outperforms zero-shot prompting and random sampled few-shot prompting in two data sets chosen for this study, demonstrating the efficacy of cluster-based retrieval in improving few-shot prompting for medical NER tasks.

## introduction large language models (llm) are deep neural networks trained heavily on massive datasets and are very effective for a wide variety of generative tasks including machine translation @xcite , question-answering @xcite , and human-computer conversations @xcite . the potential of llms in various extraction tasks like named entity recognition (ner) has also been widely explored in the literature @xcite . llms are initially pre-trained on language modeling tasks and later fine-tuned to adapt to ner tasks. their success depends on the quality and quantity of the annotated datasets used for this supervised fine-tuning. instructiontuned llms used for ner tasks are trained to take the input query as a component of their prompt and respond with the ner annotated text using techniques like slot-filling. engineering the right prompt for such models is still an active research area. llm prompts need accurate instructions and appropriate context information to guide them to the correct output for any specific tasks. few-shot prompting technique @xcite ) that provides the llm with sample outputs in their context has been shown to be effective in improving llm responses. medical named entity recognition (ner) is the process of identifying and classifying medical entities from unstructured text. recognizing medical entities is essential for applications like electronic health record (ehr) analysis, clinical decision support, and biomedical research @xcite . medical text contains a wide range of medical terms, including diseases, medications, symptoms, and treatments, which can be expressed in diverse ways. many studies show that the performance of domain-agnostic llms for medical ner is significantly lower than that of models trained specifically for the medical domain. even domain-specific models require fine-tuning and retraining to adapt to new datasets, and their performance gain is only marginal if not none to traditional extraction techniques. few-shot prompting in llms for medical ner offers opportunities to adapt these models to unseen data with mini-mal sample data and training effort. considering the context length limit, identifying the most relevant few-shot samples for the given input text is a crucial step. this challenge has not been much explored in the medical domain. the main objective of this work is to improve few-shot prompting in pre-trained llms for medical ner by proposing a novel approach in sample selection. our approach makes use of vectorbased similarity to select the most appropriate kshot demonstrations required for model prompting. to ensure relevance and divergence in the retrieved samples, we select samples belonging to k-different clusters whose centroids match the most with the given input query. we experimented with the opensource llama2 model for different numbers of clusters and with multiple values for k, on two clinical datasets. our findings are reported in the result section which are promising and call for further research in this area. ## related work in recent years, large language models (llms) have demonstrated immense potential in fewshot and zero-shot clinical information extraction. @xcite explored the use of instruct-gpt, showing that llms can effectively reduce dependency on large annotated datasets and perform well on re-annotated datasets like casi. this is particularly relevant in healthcare, where annotated clinical datasets are limited. similarly, techniques using pre-trained language models combined with domain-specific dictionaries and bilstm-crf architectures have enhanced ner in unlabelled medical texts, significantly improving entity recognition rates @xcite . the integration of llms such as bert in methods like nessma, which uses surrounding sequence matching, has further improved contextual understanding in clinical text mining @xcite . however, token-level ner, particularly in rare diseases, remains a challenge. @xcite identified limitations in token-level ner for llms like meditron and universalner but found that llama2-medtuned-7b shows promising results, highlighting the potential of open-source models in this field. moreover, advanced systems like bern2 combine traditional rule-based methods with llms, enhancing both recognition and normalization tasks @xcite . the use of retrieval-augmented generation (rag) models has also significantly improved knowledge-intensive nlp tasks, offering better factual accuracy and interpretability @xcite . biomistral, an opensource model for medical domains, further showcases performance improvements through finetuning and model merging techniques @xcite . the retrieving and chain-of-thought (rt) framework @xcite enhances few-shot biomedical named entity recognition (ner) by combining retrieval and reasoning. the retrieval module selects relevant examples, providing context for better sentence understanding, while the chain-of-thought module applies systematic reasoning to improve entity prediction. together, they significantly boost model accuracy. a comparative study on bc5cdr and ncbi datasets demonstrated superior performance, while error analysis revealed challenges with long and out-ofvocabulary entities. the rt framework offers a promising approach for advancing ner in biomedical applications. a prompting method for few-shot named entity recognition (ner) using large language models (llms) features a standardized prompt structure with task definition, few-shot demonstrations, and a constrained output format to ensure accurate and structured responses @xcite . the method also includes error mitigation prompts to correct recognition mistakes and optimizes demonstration selection for better performance. ablation studies assess the impact of each prompt component. overall, this approach improves ner accuracy by leveraging llms while addressing fewshot learning challenges across datasets. the clustered retrieved augmented generation (crag) approach @xcite enhances question-answering systems using large language models (llms) by optimizing data processing. it begins with collecting product reviews and generating embeddings using a sentence transformer, followed by clustering these embeddings via k-means. summarization of each cluster is performed with the mistral 7b model to reduce redundancy. the cluster summaries are then aggregated into a comprehensive knowledge base. this method reduces token usage while maintaining response quality, improving the system's efficiency. two clinical ner datasets -mtsamples and vaers used in this paper are sourced from @xcite . mtsamples dataset consists of discharge summaries from mtsamples, which were annotated based on the guidelines from the 2010 i2b2 challenge, focusing on extracting medical problems, treatments, and tests @xcite . vaers dataset is a collection of publicly accessible safety reports from vaccine adverse event reporting system (vaers), focused on extracting events related to nervous system disorders @xcite . the mtsamples dataset is completely synthetic, meaning it was generated artificially and does not include any real patient information. the vaers dataset comes from publicly available post-market safety reports that are anonymized and do not contain any personal data. as a result, no sensitive information was shared, ensuring that this study poses no privacy concerns. the two datasets were split into training, validation, and test subsets. the training subset is used as the external knowledge base to provide the few-shot samples, whereas the test subset is for evaluating the model's performance and for comparative analysis. ## methodology the architecture for the proposed cluster-based retrieval augmented extraction (crae) approach is shown in figure 1 . for our experiments, we use the training data split of the respective dataset as the external knowledge store, although it can be any named entity annotated data. these documents are first split into manageable chunks and are embedded using a pre-trained hugging face model. the embeddings are stored in faiss, a vector database designed for efficient retrieval. the nearest k-shots retrieved from this vector store are not only very similar to the input query but also each other. our experiment showed that augmenting more than one such sample in the prompt did not improve the performance of the model but rather had an inverse effect. to address this issue, we clustered these embeddings into n clusters out of which k clusters were used in prompting. for each input query in the test dataset, we retrieve the most similar example from each cluster by performing a similarity search between the cluster centroids and the input query embedding. these selected examples were then used as context for the query. both the query and the context are passed to a large language model (llm) through a pre-defined prompt template to generate informed responses. the task specific prompt template for mtsamples is shown in ## result and discussion we have experimented the models for different values of n (number of clusters created) and k (number of clusters selected) and those results are discussed. figure 2 shows the impact of the number of fewshot demonstrations (k) on the performance for mt-samples with 126 clusters, with k clusters chosen based on the nearest centroids to the input embedding. the evaluation metrics vary as the number of selected clusters changes. the best performance can be seen at @xmath0. figure 3 shows the evaluation results on mtsamples using a 2-shot prompt. the results show the performance for different numbers of clusters created, with two clusters being selected based on the centroids closest to the input query embedding. the best performance can be seen at a cluster count of 126. the the ## conclusion and futureworks this paper presents a novel approach to enhance few-shot prompting for named entity recognition (ner) in the medical domain. by leveraging kmeans clustering and embedding similarity, we propose a cluster-based sample selection strategy to identify contextually relevant yet diverse training examples in few-shot prompting. our experiments demonstrate that this approach significantly outperforms zero-shot prompting and random sampled few-shot prompting on two medical datasets. however, further research is needed to explore this approach using more open-source domain-agnostic models. a thorough qualitative analysis of the results could give more insights to improve the results. our study was limited by the availability of computing power to run the model. the usage of more advanced embedding models or fine-tuning the current models can make the similarity search more accurate. this would help retrieve better and more relevant context. exploring the use of hybrid search methods that combine semantic and keyword-based approaches might further enhance retrieval performance. this work could potentially lead to even more robust and generalizable solutions for ner few-shot prompting across various domains.