MMRAG: multi-mode retrieval-augmented generation with large language models for biomedical in-context learning

Abstract

Objectives

To optimize in-context learning in biomedical natural language processing by improving example selection.

Materials and Methods

We introduce a novel multi-mode retrieval-augmented generation (MMRAG) framework, which integrates 4 retrieval strategies: (1) Random Mode, selecting examples arbitrarily; (2) Top Mode, retrieving the most relevant examples based on similarity; (3) Diversity Mode, ensuring variation in selected examples; and (4) Class Mode, selecting category-representative examples. This study evaluates MMRAG on 3 core biomedical NLP tasks: Named Entity Recognition (NER), Relation Extraction (RE), and Text Classification (TC). The datasets used include BC2GM for gene and protein mention recognition (NER), DDI for drug-drug interaction extraction (RE), GIT for general biomedical information extraction (RE), and HealthAdvice for health-related text classification (TC). The framework is tested with 2 large language models (Llama-2-7B and Llama-3-8B) and 3 retrievers (Contriever, MedCPT, and BGE-Large) to assess performance across different retrieval strategies.

Results

The results from the Random Mode indicate that providing more examples in the prompt improves the model’s generation performance. Meanwhile, Top Mode and Diversity Mode significantly outperform Random Mode on the RE (DDI) task, achieving an F1 score of 0.9669—a 26.4% improvement. Among the 3 retrievers tested, Contriever outperformed the other 2 in a greater number of experiments. Additionally, Llama 2 and Llama 3 demonstrated varying capabilities across different tasks, with Llama 3 showing a clear advantage in handling NER tasks.

Conclusion

MMRAG effectively enhances biomedical in-context learning by refining example selection, mitigating data scarcity issues, and demonstrating superior adaptability for NLP-driven healthcare applications.

Introduction

Large language models (LLMs) such as the GPT series^1,2 and Llama models³ are now playing an increasingly pivotal role in healthcare,⁴ demonstrating remarkable capabilities in processing biomedical text and electronic health records (EHRs).^5–7 Their applications include extracting clinical insights,^8,9 supporting patient monitoring,¹⁰ performing clinical text classification,¹¹ addressing natural language inference,¹² and facilitating medical information extraction.^13–15

However, due to the need to protect patient privacy,¹⁶ research in medical NLP faces significant challenges related to insufficient training data.^17,18 Although hospitals and healthcare institutions often possess large amounts of raw EHR data, much of these data are unstructured and unlabeled.⁴ The process of curating, annotating, and structuring this data into usable formats for training machine learning models is time-consuming, money-consuming, and resource-intensive. Furthermore, for rare diseases^19,20—such as genetic disorders, orphan diseases, and conditions with limited geographic or demographic prevalence—the lack of documented cases exacerbates the data scarcity issue. These conditions inherently have fewer reported instances, making it difficult to collect even raw EHR data, let alone labeled datasets. Finally, privacy regulations such as HIPAA or GDPR further restrict access to sensitive patient information,^21,22 making it even harder to obtain and share data for research purposes. As a result, novel approaches are urgently needed to address these challenges.

Given these considerations, in-context learning (ICL) offers a promising approach. Along with the development of LLMs, the size of these models has grown significantly, the amount of training data has increased, and the length of prompts has expanded.²³ As a result, LLMs have begun to demonstrate the ability to learn from context, a capability known as ICL.²⁴ This approach allows the model to learn patterns from a few examples provided within the prompt, enabling it to deliver high-quality responses. ICL has several advantages. First, it eliminates the need for training or fine-tuning, which requires substantial computational resources and large amounts of data; instead, ICL only requires a few examples.²⁵ Second, designing prompts is much simpler than training a model,²⁶ making this approach accessible to individuals without specialized knowledge of LLMs, such as doctors in hospitals. Third, providing examples aligns with the human brain’s natural way of thinking—learning through analogy. This paradigm allows for the seamless incorporation of human knowledge into LLMs by simply adjusting demonstrations and templates.^27,28 Finally, a single model can address a wide range of tasks by designing different prompts and providing corresponding examples, showcasing its remarkable versatility.²⁹ Therefore, ICL is indispensable in the field of healthcare.

By leveraging the ability of large language models to learn patterns from a few carefully designed examples embedded in the prompt, ICL could mitigate the insufficient data challenges in the biomedical domain. However, despite its potential, the process of selecting relevant examples from existing verified or labeled data for ICL remains underexplored.³⁰ The careful curation of examples is crucial in biomedical contexts, as the examples not only encapsulate rich medical knowledge—such as clinical guidelines, rare disease manifestations, or treatment pathways—but also significantly influence the quality and accuracy of the model’s responses based on their relevance to the input.^25,27,30 Therefore, developing effective methodologies for selecting high-quality examples becomes essential to unlocking the full potential of ICL in medical NLP.

Several studies have explored the selection of examples for ICL. For instance, Liu et al. proposed kNN-augmented in-context example selection to improve ChatGPT-3’s performance. Kumari et al.³¹ developed an end-to-end submodular framework to enhance the example selection process. Zhang et al.²⁵ formulated example selection for ICL as a sequential decision problem and introduced a reinforcement learning approach for identifying generalizable policies to select demonstration examples. Li et al.³² presented LENS, which finds informative in-context examples for better performance. Similarly, Mavromatis et al.³⁰ proposed AdaICL to improve example diversity. However, to the best of our knowledge, no one has yet explored this area in the biomedical domain.

Fortunately, retrieval-augmented generation (RAG),^33,34 a technique designed to extract information directly relevant to an input query, offers a promising approach to address this challenge. RAG can enhance the example selection process by efficiently retrieving contextually similar data points,^35,36 making it a natural complement to ICL. To further advance the application of ICL in medical NLP, we propose an innovative multi-mode retrieval-augmented generation (MMRAG) framework that leverages RAG techniques to extract and select similar examples using different strategies tailored to diverse use cases. This framework introduces 4 distinct and independent modes for example selection: (1) Random Mode, where examples are selected without any specific prioritization; (2) Top Mode, which retrieves the most relevant examples based on similarity metrics; (3) Diverse Mode, which ensures a broad range of examples to increase diversity; and (4) Class Mode, which focuses on selecting examples representative of specific categories or conditions. Each mode corresponds to a unique rule for selecting examples, providing flexibility and adaptability to a variety of medical NLP tasks. By integrating these strategies, the proposed framework offers a systematic and effective approach to maximize the potential of ICL in addressing critical challenges in medical NLP, particularly in domains constrained by data scarcity and privacy considerations. Our key contributions are summarized as follows:

• To the best of our knowledge, this is the first study to investigate the selection of examples in the biomedical domain for ICL.

• We propose a novel RAG-based framework with 4 distinct modes (Random, Top, Diverse, and Class) to optimize example selection and explore the potential of ICL in biomedical contexts.

• We conduct extensive experiments, evaluating the framework across 2 LLMs on 4 datasets with 3 retrievers, providing comprehensive insights into its performance.

Methods

Overview of methods

This study introduces MMRAG (multi-mode retrieval-augmented generation) framework, designed to enhance ICL in biomedical NLP by optimizing example selection. MMRAG employs 4 retrieval strategies: Random Mode (arbitrary selection), Top Mode (most similar examples), Diversity Mode (balanced similarity and variation), and Class Mode (representative examples from different categories). The framework integrates RAG with LLMs (Llama-2-7B³⁷ and Llama3-8B³⁸) and 3 retrievers (Contriever,³⁹ MedCPT,⁴⁰ and BGE-Large⁴¹) to systematically evaluate retrieval effectiveness. Our experiments cover Named Entity Recognition (NER), Relation Extraction (RE), and Text Classification (TC) using 4 biomedical datasets (BC2GM,⁴² DDI,⁴³ GIT,⁴⁴ and HealthAdvice⁴⁵). This framework demonstrates how tailored retrieval strategies can significantly improve ICL performance in biomedical NLP, addressing challenges such as data scarcity and privacy constraints.

Tasks and datasets

In this work, we focus on 3 core information extraction tasks in the biomedical domain: NER, RE, and TC. The primary objective of these tasks is to extract meaningful information from raw, unstructured text, thereby transforming it into structured data suitable for downstream applications. Detailed statistics and representative examples of each dataset used can be found in Supplemental Material 1.

Named Entity Recognition (NER)—BC2GM ⁴² : The BC2GM dataset is a widely used benchmark for biomedical named entity recognition (NER), specifically targeting gene and protein mentions in scientific literature. Derived from the BioCreative II Gene Mention task, it comprises manually annotated gene mentions sourced from PubMed abstracts. This dataset plays a crucial role in biomedical text mining by supporting the development and evaluation of NER models capable of accurately identifying gene-related entities.

Relation Extraction (RE)—DDI ⁴³ : The DDI (Drug-Drug Interaction) dataset is a key resource for relation extraction in the biomedical domain, particularly designed to identify and classify drug-drug interactions from text. It includes sentences from biomedical literature and drug product information sources, annotated with drug entities and their interactions. This dataset is crucial for pharmacovigilance and drug safety, as accurate extraction of drug interactions supports clinical decision-making, prevents adverse drug reactions, and enhances patient safety.

Relation Extraction (RE)—GIT ⁴⁴ : The GIT (General BioMedical and Complementary and Integrative Health Triples) dataset is a high-quality biomedical triple extraction dataset specifically focused on non-drug therapies. It is characterized by its high-quality annotations and comprehensive coverage of relation types, making it a valuable resource for biomedical relation extraction. The dataset includes 22 relation types derived from SemMedDB, providing structured representations of relationships within biomedical texts. By supporting the extraction of meaningful entity-relation triples, GIT facilitates advancements in knowledge graph construction, automated reasoning, and biomedical text mining applications related to non-drug interventions.

Text Classification (TC)—HealthAdvice ⁴⁵ : The HealthAdvice dataset is designed for text classification tasks related to health information and advisory content. It consists of a diverse collection of health-related advice. The dataset is structured to facilitate automatic classification of health advice into relevant categories, supporting applications such as misinformation detection, personalized health recommendations, and automated triaging of medical inquiries. By leveraging this dataset, researchers can enhance natural language processing (NLP) models tailored for health communication and decision support.

The combination of these datasets—BC2GM for entity recognition, DDI and GIT for relation extraction, and HealthAdvice for text classification—provides a comprehensive foundation for biomedical NLP research. These datasets enable the development of advanced NLP models capable of extracting critical biomedical knowledge from vast textual sources. By improving the accuracy of entity recognition, relationship extraction, and health-related text classification, they contribute significantly to biomedical informatics, clinical decision support, and evidence-based medicine.

MMRAG framework

The motivation for the proposed framework is to release the potential of ICL using the RAG technique for biomedical information extraction with LLMs. As depicted in Figure 1, the retrieval process begins by embedding both the training set sentences and the input sentence. Differed from traditional RAG, our MMRAG framework uses training set as the external knowledge database. A retriever then ranks the training sentences based on similarity scores. The selection mode determines how the final examples are chosen from this ranked list. Finally, the retrieved examples are integrated with the input sentence in a prompt which then was fed into LLM to get response. In this paper, we introduced 4 different modes for selecting examples: Random Mode, Top Mode, Diversity Mode, and Class Mode. These modes determine how retrieved sentences are selected based on their similarity rankings, striking different balances between similarity and diversity.

Figure 1.

An overview of the multi-mode retrieval-augmented generation framework. Compared with the traditional retrieval-augmented generation framework, our new MMRAG framework introduces 4 independent modes for better in-context learning.

Random Mode (Few-shot): Random Mode selects examples arbitrarily from the training set without considering similarity ranking. This introduces randomness, which can improve model generalization but may result in irrelevant examples. While this mode maximizes diversity, it does not ensure the selection of highly relevant sentences.

Top Mode: Top Mode selects the most similar examples by directly choosing the top-ranked sentences. This ensures high contextual similarity between the input sentence and the retrieved examples. However, it may lead to redundancy, as the examples are often very similar to each other, limiting the diversity of information. In this paper, we used cosine similarity to measure the similarity between the input sentence and examples.

Diversity Mode: Diversity Mode introduces a skip-step selection mechanism, where examples are picked at intervals from the ranked list instead of consecutively selecting the top sentences. By skipping certain highly ranked examples, this method ensures that the retrieved examples maintain relevance while increasing the diversity of information.

Class Mode: Class Mode selects the most relevant example from each predefined class rather than purely ranking by similarity. This ensures that different types of sentences are represented in the final selection, improving coverage across multiple categories. While it maintains relevance, some selected examples may have slightly lower similarity scores compared to those chosen in Top Mode.

Each mode provides a different trade-off between similarity and diversity. Random Mode cannot guarantee diversity and relevance, but it represents the most common method. Top Mode guarantees high similarity but lacks diversity. Diversity Mode balances both by selecting non-consecutive similar examples. Class Mode further enhances diversity by ensuring category coverage. Choosing the appropriate mode depends on the specific requirements of the retrieval task.

Prompts

In retrieval-augmented ICL, prompt design plays a crucial role in guiding LLMs to understand the task and generate appropriate outputs. While our MMRAG framework maintains consistent procedures for retrieval, generation, and evaluation across tasks, the task-specific prompt templates vary slightly depending on the nature of the NLP task. For each task—NER, RE, and TC—we constructed different instruction in prompts as follows:

• Please do a named entity recognition task. You need to accurately recognize and label these entities for the input sentence. The input sentence has been divided into a list form. The response should be a list and each element is corresponding to each word and sign in the input sentence. Response should follow this rule: {\”no disease\”: 0 , \”first token of a disease\”: 1 , \”the subsequent disease tokens\”: 2}

• Please do a relation extraction task. You need to extract the relation type from the input sentence. The response should be in predefine set: {'NA', 'advise’, 'effect’, 'int’, 'mechanism’}.

• Please do a classification task. You need to answer the question based on the input sentence. The response should be in predefine set: {'yes’, 'no’}.

Experimental settings

Our experiments were conducted using 2 large language models, Llama-2-7B[38] and Llama-3-8B,³⁸ both of which have demonstrated strong performance in natural language processing tasks. To retrieve relevant examples, we employed 3 different retrievers: Contriever,³⁹ a dense retriever trained with contrastive learning for general-purpose retrieval; MedCPT,⁴⁰ a retriever optimized for medical and clinical text retrieval; and BGE-Large,⁴¹ an embedding-based retriever known for its effectiveness in similarity search. These retrievers were chosen to test the adaptability of our framework across different retrieval strategies.

All experiments were performed on NVIDIA A100 GPUs with 40GB memory, ensuring sufficient computational resources for training and inference. The batch size per device was set to 4 for both training and evaluation, and the inference was carried out sentence by sentence to optimize efficiency. To make sure the output of LLM following the same format, we applied Low-Rank Adaptation,⁴⁶ a parameter-efficient fine-tuning method, with a rank of 64, alpha set to 32, and a dropout rate of 0.1, to fine-tune the model. The models were optimized using the AdamW optimizer with a learning rate of 1e-5. Fine-tuning was conducted for 5000 steps, with evaluations performed every 1000 steps, and the best-performing model was selected for inference.

For evaluation, we adopted Micro Precision, Micro Recall, and Micro F1-score, in line with established studies.^36,47,48 A prediction was considered correct only if the entire output exactly matched the ground truth, ensuring a strict and reliable assessment of model performance. This evaluation approach allowed us to measure the effectiveness of different retrieval strategies and selection modes in a controlled and reproducible manner.

Results

Performance in Random Mode

Table 1 shows the mean and standard deviation of F1 scores for different tasks and datasets since the Random Mode introduces huge randomness. For the NER task on BC2GM, Llama-2-7B achieved an F1 score of 0.8766 with one example, improving to 0.9172 with 10 examples. Llama-3-8B showed superior performance, reaching 0.9660 with one example and further increasing to 0.9782 with 10 examples. In the RE task on the DDI dataset, Llama-2-7B fluctuated between 0.7049 and 0.7050 across different example counts, while Llama-3-8B exhibited marginal improvement, achieving a maximum F1 score of 0.7114 with 5 examples. In text classification on HealthAdvice, Llama-2-7B increased from 0.8954 (1 example) to 0.9171 (10 examples), while Llama-3-8B ranged from 0.8848 to 0.9084. These results indicate that while Random Mode can yield reasonable performance improvements, it lacks the reliability and control of retrieval-based selection.

Table 1.

Comparison of 2 LLMs across 3 datasets in Random Mode.

Tasks	Datasets	Models	Examples	Precision	Recall	F1
NER	BC2GM	Non-ICL-based SOTA49		–	–	0.980
		ICL-based SOTA50		–	–	0.856
		Llama-2-7B	1	0.9232 ± 0.0001	0.8345 ± 0.0007	0.8766 ± 0.0004
			5	0.9267 ± 0.0008	0.8878 ± 0.0032	0.9069 ± 0.0018
			10	0.9318 ± 0.0007	0.9032 ± 0.0006	0.9172 ± 0.0006
		Llama-3-8B	1	0.9733 ± 0.0006	0.9589 ± 0.0008	0.9660 ± 0.0001
			5	0.9806 ± 0.0001	0.9712 ± 0.0004	0.9759 ± 0.0003
			10	0.9814 ± 0.0007	0.9749 ± 0.0010	0.9782 ± 0.0004
RE	DDI	Non-ICL-based SOTA51		0.864	0.788	0.824
		ICL-based SOTA52		–	–	0.8392 ± 0.0002
		Llama-2-7B	1	0.7049 ± 0.0031	0.7049 ± 0.0031	0.7049 ± 0.0031
			5	0.7081 ± 0.0002	0.7081 ± 0.0002	0.7081 ± 0.0002
			10	0.7036 ± 0.0128	0.7064 ± 0.0083	0.7050 ± 0.0105
		Llama-3-8B	1	0.7004 ± 0.0018	0.7004 ± 0.0018	0.7004 ± 0.0018
			5	0.7114 ± 0.0045	0.7114 ± 0.0045	0.7114 ± 0.0045
			10	0.7005 ± 0.0066	0.7005 ± 0.0066	0.7005 ± 0.0066
TC	HealthAdvice	Non-ICL-based SOTA6		0.91	0.91	0.91
		Llama-2-7B	1	0.8954 ± 0.0027	0.8954 ± 0.0027	0.8954 ± 0.0027
			5	0.9154 ± 0.0017	0.9159 ± 0.0020	0.9156 ± 0.0018
			10	0.9171 ± 0.0021	0.9171 ± 0.0021	0.9171 ± 0.0021
		Llama-3-8B	1	0.8848 ± 0.0025	0.8848 ± 0.0025	0.8848 ± 0.0025
			5	0.9040 ± 0.0019	0.9040 ± 0.0019	0.9040 ± 0.0019
			10	0.9084 ± 0.0006	0.9084 ± 0.0006	0.9084 ± 0.0006

The reported values represent the mean and standard deviation over 3 repeated experiments. The bolded numbers denote the highest scores in each setting.

Performance in Top Mode

Table 2 presents the results for this mode, demonstrating notable performance improvements over Random Mode. In NER, Llama-2-7B reached an F1 score of 0.9098 with 5 examples when using the Contriever retriever, while Llama-3-8B achieved a peak score of 0.9726. The MedCPT retriever also contributed to strong results, with Llama-3-8B scoring 0.9706. For RE on the DDI dataset, Contriever retrieval led to the highest F1 score of 0.9669 for Llama-2-7B and 0.9573 for Llama-3-8B. Similarly, for text classification, Top Mode improved performance, with Llama-2-7B reaching 0.8969 and Llama-3-8B achieving 0.8923 when using Contriever. These results confirm that retrieval-driven selection significantly enhances model effectiveness, particularly when using domain-optimized retrievers.

Table 2.

Comparison of 2 LLMs across 3 datasets using 3 retrievers in Top Mode.

Tasks	Datasets	Retrievers	Num	Llama 2			Llama 3
				Precision	Recall	F1	Precision	Recall	F1
NER	BC2GM	BGE-Large	1	0.9187	0.7310	0.8142	0.9573	0.9365	0.9468
			5	0.9154	0.7318	0.8134	0.9080	0.8820	0.8948
			10	0.9130	0.7197	0.8049	0.9052	0.8763	0.8905
		Contriever	1	0.9030	0.7755	0.8344	0.9468	0.9180	0.9322
			5	0.9242	0.8958	0.9098	0.9747	0.9686	0.9716
			10	0.9167	0.8633	0.8892	0.9780	0.9672	0.9726
		MedCPT	1	0.8608	0.7302	0.7902	0.9500	0.9204	0.9350
			5	0.9324	0.8527	0.8908	0.9784	0.9629	0.9706
			10	0.9228	0.8828	0.9024	0.9803	0.9493	0.9646
RE	DDI	BGE-Large	1	0.7358	0.7358	0.7358	0.7305	0.7305	0.7305
			5	0.7745	0.7745	0.7745	0.7589	0.7589	0.7589
			10	0.8026	0.8026	0.8026	0.7795	0.7795	0.7795
		Contriever	1	0.9662	0.9662	0.9662	0.7532	0.7550	0.7541
			5	0.9504	0.9649	0.9576	0.9646	0.9646	0.9646
			10	0.9669	0.9669	0.9669	0.9573	0.9573	0.9573
		MedCPT	1	0.8146	0.8149	0.8148	0.7050	0.7106	0.7078
			5	0.8152	0.8152	0.8152	0.8205	0.8205	0.8205
			10	0.8264	0.8264	0.8264	0.8175	0.8175	0.8175
TC	HealthAdvice	BGE-Large	1	0.8836	0.8836	0.8836	0.8180	0.8180	0.8180
			5	0.8906	0.8906	0.8906	0.8612	0.8612	0.8612
			10	0.8998	0.8998	0.8998	0.8646	0.8646	0.8646
		Contriever	1	0.8900	0.8900	0.8900	0.8116	0.8116	0.8116
			5	0.8892	0.8969	0.8930	0.8923	0.8923	0.8923
			10	0.8969	0.8969	0.8969	0.8571	0.8606	0.8589
		MedCPT	1	0.7454	0.7454	0.7454	0.7028	0.7028	0.7028
			5	0.7999	0.8128	0.8063	0.8243	0.8243	0.8243
			10	0.7797	0.7869	0.7833	0.7954	0.8105	0.8029

The bolded numbers denote the highest scores in each setting.

Performance in Diversity Mode

Table 3 presents the results, demonstrating that this approach improved model robustness across datasets. In NER, the highest F1 scores were achieved using the Contriever retriever, with Llama-2-7B reaching 0.9134 and Llama-3-8B achieving 0.9731 with 10 examples. For the DDI dataset, Diversity Mode yielded stable results, with Llama-2-7B reaching 0.9666 and Llama-3-8B scoring 0.9623 using Contriever retrieval. Text classification results remained competitive, with Llama-2-7B achieving an F1 score of 0.8980 and Llama-3-8B reaching 0.8765. The results indicate that Diversity Mode helps mitigate overfitting while maintaining high task relevance.

Table 3.

Comparison of 2 LLMs across 3 datasets using 3 retrievers in Diversity Mode.

Tasks and datasets	Retrievers	Gap	Llama2 (5 examples)			Llama2 (10 examples)			Llama3 (5 examples)			Llama3 (10 examples)
			Precision	Recall	F1	Precision	Recall	F1	Precision	Recall	F1	Precision	Recall	F1
NER, BC2GM	BGE-Large	1	0.9154	0.7318	0.8134	0.9130	0.7197	0.8049	0.9080	0.8820	0.8948	0.9052	0.8763	0.8905
		2	0.9145	0.7360	0.8156	0.9147	0.7235	0.8079	0.9109	0.8815	0.8960	0.9059	0.8762	0.8908
		3	0.9125	0.7415	0.8182	0.9133	0.7192	0.8047	0.9081	0.8828	0.8953	0.9052	0.8797	0.8923
	Contriever	1	0.9242	0.8958	0.9098	0.9167	0.8633	0.8892	0.9747	0.9686	0.9716	0.9780	0.9672	0.9726
		2	0.9292	0.8980	0.9134	0.9192	0.8710	0.8945	0.9712	0.9710	0.9711	0.9788	0.9675	0.9731
		3	0.9288	0.8798	0.9036	0.9168	0.8921	0.9043	0.9747	0.9636	0.9691	0.9782	0.9688	0.9735
	MedCPT	1	0.9324	0.8527	0.8908	0.9228	0.8828	0.9024	0.9784	0.9629	0.9706	0.9803	0.9493	0.9646
		2	0.9411	0.8386	0.8869	0.9206	0.8615	0.8900	0.9780	0.9601	0.9690	0.9811	0.9520	0.9663
		3	0.9408	0.8272	0.8804	0.9159	0.8481	0.8807	0.9751	0.9440	0.9593	0.9762	0.9634	0.9698
RE, DDI	BGE-Large	1	0.7745	0.7745	0.7745	0.8026	0.8026	0.8026	0.7589	0.7589	0.7589	0.7795	0.7795	0.7795
		2	0.7811	0.7811	0.7811	0.7768	0.7768	0.7768	0.7450	0.7450	0.7450	0.7385	0.7385	0.7385
		3	0.7666	0.7666	0.7666	0.7611	0.7611	0.7611	0.7440	0.7440	0.7440	0.7487	0.7487	0.7487
	Contriever	1	0.9504	0.9649	0.9576	0.9669	0.9669	0.9669	0.9646	0.9646	0.9646	0.9573	0.9573	0.9573
		2	0.9666	0.9666	0.9666	0.9669	0.9669	0.9669	0.9623	0.9623	0.9623	0.9334	0.9334	0.9334
		3	0.9659	0.9659	0.9659	0.9649	0.9662	0.9656	0.9613	0.9613	0.9613	0.8574	0.8583	0.8579
	MedCPT	1	0.8152	0.8152	0.8152	0.8264	0.8264	0.8264	0.8205	0.8205	0.8205	0.8175	0.8175	0.8175
		2	0.8199	0.8199	0.8199	0.8170	0.8170	0.8170	0.8136	0.8136	0.8136	0.8228	0.8228	0.8228
		3	0.8252	0.8252	0.8252	0.8211	0.8211	0.8211	0.8162	0.8162	0.8162	0.8319	0.8319	0.8319
TC, HealthAdvice	BGE-Large	1	0.8906	0.8906	0.8906	0.8998	0.8998	0.8998	0.8612	0.8612	0.8612	0.8646	0.8646	0.8646
		2	0.8796	0.8796	0.8796	0.8940	0.8940	0.8940	0.8606	0.8606	0.8606	0.8445	0.8445	0.8445
		3	0.8894	0.8894	0.8894	0.8980	0.8980	0.8980	0.8583	0.8583	0.8583	0.8623	0.8623	0.8623
	Contriever	3	0.8888	0.8934	0.8911	0.8899	0.8899	0.8899	0.9062	0.9067	0.9064	0.8762	0.8767	0.8765
		1	0.8892	0.8969	0.8930	0.8969	0.8969	0.8969	0.8923	0.8923	0.8923	0.8571	0.8606	0.8589
		2	0.8779	0.8825	0.8802	0.8928	0.8928	0.8928	0.9026	0.9026	0.9026	0.8562	0.8675	0.8618
	MedCPT	1	0.7999	0.8128	0.8063	0.7797	0.7869	0.7833	0.8243	0.8243	0.8243	0.7954	0.8105	0.8029
		2	0.8179	0.8203	0.8191	0.8154	0.8173	0.8164	0.8324	0.8324	0.8324	0.8397	0.8416	0.8406
		3	0.8084	0.8116	0.8100	0.7522	0.7752	0.7636	0.8101	0.8111	0.8106	0.8393	0.8393	0.8393

The bolded numbers denote the highest scores in each setting.

Performance in Class Mode

Table 4 compares Llama2 and Llama3 across 3 datasets using 3 retrievers in Class Mode. Llama3 generally outperforms Llama2, especially in RE (GIT) where its F1-score with MedCPT reaches 0.6129, significantly higher than Llama2’s 0.3786. In RE (DDI), Contriever achieves the best performance for both models, with Llama3 scoring 0.7060. For TC (HealthAdvice) both models perform well, with Llama2 slightly ahead using random retrieval, but Llama3 excels with Contriever (0.8917 F1-score). These results highlight Llama3’s superior retrieval utilization and the impact of effective retrievers like MedCPT and Contriever on model performance.

Table 4.

Comparison of 2 LLMs across 3 datasets using 3 retrievers in Class Mode.

Tasks, datasets, examples	Retrievers	Llama 2			Llama 3
		Precision	Recall	F1	Precision	Recall	F1
RE, DDI, examples: 4	Random	0.6725	0.6725	0.6725	0.6738	0.6738	0.6738
	MedCPT	0.6825	0.6825	0.6825	0.6841	0.6841	0.6841
	Contriever	0.6915	0.6917	0.6916	0.7060	0.7060	0.7060
	BGE-Large	0.6739	0.6755	0.6747	0.6881	0.6881	0.6881
RE, GIT, examples: 22	Random	0.3004	0.3004	0.3004	0.5075	0.5075	0.5075
	MedCPT	0.3786	0.3786	0.3786	0.6129	0.6129	0.6129
	Contriever	0.2354	0.2354	0.2354	0.5978	0.5978	0.5978
	BGE-Large	0.3143	0.3143	0.3143	0.6108	0.6108	0.6108
TC, HealthAdvice, examples: 3	Random	0.8733	0.8733	0.8733	0.8658	0.8658	0.8658
	MedCPT	0.8986	0.8986	0.8986	0.8842	0.8842	0.8842
	Contriever	0.8940	0.8940	0.8940	0.8917	0.8917	0.8917
	BGE-Large	0.8865	0.8865	0.8865	0.8785	0.8785	0.8785

The bolded numbers denote the highest scores in each setting.

Discussion

The proposed MMRAG framework is particularly valuable for the biomedical domain, where labeled data are scarce due to privacy constraints and annotation costs. Biomedical NLP tasks require high precision, as errors in entity recognition or relation extraction can lead to misinformation in clinical settings. By integrating retrieval-augmented generation, this framework improves ICL by dynamically selecting relevant and diverse examples. The results demonstrate that Top Mode provided the highest F1 scores in precision-sensitive tasks such as NER and RE, demonstrating that high-relevance retrieval improves performance. However, the Diversity Mode achieved more stable generalization across datasets while preserving task relevance. Class Mode, while slightly reducing F1 scores in some cases, ensured broader category representation, which is crucial for classification and multi-relation extraction tasks. Unlike traditional fine-tuning, this approach allows models to adapt to evolving biomedical knowledge. Its flexibility makes it suitable for applications in clinical decision support, drug discovery, and genomic research.

Figure 2 illustrates the impact of different numbers of examples on model performance in Random Mode. For the NER (BC2GM) and TC (HealthAdvice) tasks, an increase in the number of examples generally improves performance. However, the performance gain from 5 to 10 examples is noticeably smaller compared to the significant improvement observed from 1 to 5 examples. This suggests that while the jump from 1 to 5 examples provides substantial additional information, the marginal benefit of increasing from 5 to 10 examples is limited. Interestingly, the trend for RE (DDI) is different. The model’s performance declines when increasing from 5 to 10 examples. This can be attributed to the random selection of examples, which may include irrelevant or misleading instances. As a result, these less informative or even contradictory examples could lead the model to make incorrect predictions, ultimately degrading performance.

Figure 2.

Performance comparison across 3 datasets evaluated under Random Mode (Few-shot). The mean and standard deviation values are computed from 3 experimental repetitions for each configuration.

Figure 3 presents the F1-score comparison across different numbers of examples and retrievers in Top Mode. The results reveal that Llama3 significantly outperforms Llama2 on the NER (BC2GM) task. However, in contrast, Llama2 performs better than Llama3 on the TC (HealthAdvice) and RE (DDI) tasks. In terms of retriever performance, all 3 retrievers are well-suited for NER (BC2GM). For TC (HealthAdvice), Contriever and BGE-Large outperform MedCPT, indicating their stronger retrieval effectiveness in this task. In the RE (DDI) task, Contriever demonstrates superior performance compared to the other retrievers. Overall, these findings highlight that the effectiveness of different models and retrievers varies across tasks, emphasizing the importance of choosing the appropriate combination for optimal results.

Figure 3.

Performance comparison of 2 models under the Top Mode across 3 datasets and 3 few-shot settings.

Figures 4 present the results for Diversity Mode with different gap values using 5 examples. Across most experiments, we observe that the gap parameter causes slight fluctuations in performance, but the variations are generally small. This is because Diversity Mode still retrieves examples that are highly similar to the input sentence. However, the nature of these similar examples can influence performance in different ways. If the retrieved examples contain useful and relevant information, they contribute positively, leading to performance improvements. Conversely, if they contain irrelevant or misleading information, they may slightly degrade model performance. Regardless of the gap value, the most similar example is always included as one of the retrieved examples, ensuring that the performance variation remains anchored to the Top Mode baseline with one example setting.

Figure 4.

Trend of the gap parameter increase in Diversity Mode when 5 examples are provided in the prompt.

Class Mode results, as shown in Figure 5, demonstrate its effectiveness in classification and relation extraction tasks. While the absolute F1 scores are slightly lower than those in Top Mode, this approach ensures that the retrieved examples cover a broader range of biomedical categories, providing a more diverse reference set for model inference. The retrieval strategy in Class Mode closely resembles human thinking. In real-world classification tasks, people tend to perform better when they have observed a sufficient number of examples across all categories. However, current LLMs and most NLP algorithms are primarily trained for next-word prediction, which limits their ability to generalize in classification scenarios. Enhancing LLM inference by aligning it more closely with human reasoning processes could significantly improve performance. Therefore, we can anticipate that if future LLMs evolve to better mimic human cognitive patterns, strategies like Class Mode could further enhance their effectiveness in complex classification tasks.

Figure 5.

Class Mode performance comparison across 3 datasets on 2 models.

Statistical analysis

Based on the results from Tables 1-4, we further use boxplots to present a statistical comparison across different datasets, models, and retrievers, as illustrated in Figure 6. From the DDI dataset (Figure 6, A-D), we can see that the Top Mode and Diversity Mode perform best. On the health advice dataset, all 4 modes deliver competitive results. Notably, on the health advice dataset under Llama2, the Random Mode surpasses the others, whereas under Llama3, the Top Mode performs worse than the other modes. Meanwhile, Figure 6, E-H compares the 2 models under the 4 modes, revealing that overall, both models achieve a relatively similar level of performance. Lastly, Figure 6, I-K shows the statistical results for different retrievers. Here, Contriever achieves comparatively strong performance across all 3 tasks, standing out noticeably on the RE task. In contrast, BGE-large and MedCPT both underperform the other 2 retrievers on NER and TC, respectively.

Figure 6.

Statistical comparison by datasets, models, and retrievers.

Error analysis

Diving deeper into these promising results, our observations of the generated outputs reveal that LLMs mainly suffer from the following 2 errors:

• Generation Errors: These include misclassifying entities, incorrectly assigning relationships between head and tail entities, and extracting non-entity words as entities. Example 1: LLMs wrongly identify the word “Genetic” as gene type but it is not. Example 2: for sentence “[‘Bort’, ‘-‘, ‘dex’, ‘was’, ‘an’, ‘effective’, ‘salvage’, ‘treatment’, ‘for’, ‘MM’]”, the correct response should be [1, 0, 1, 0, 0, 0, 0, 0, 0, 2]. However, the generated response was [1, 0, 1, 0, 0, 0, 0, 2], which missed some tokens.

• Over-Extraction: Due to our strict evaluation criteria, extracting additional entities beyond the ground truth or producing boundary errors in entity spans can lead to a decline in performance scores. Example 1: The term “CRE-binding-protein” might be tokenized as [‘CRE’, ‘-’, ‘binding’, ‘-’, ‘protein’], but the LLM may not recognize the hyphens as part of the entity. As a result, it predicts “CRE binding protein” as the disease entity, which does not exactly match the annotated “CRE-binding-protein.” Example 2: The correct disease entity is “depressive disorder,” but the LLM predicts “major depressive disorder,” resulting in an over-extraction beyond the ground truth.

Limitations and future directions

While this study focuses on example retrieval strategies, several limitations remain:

• Limited retriever diversity: Although the MMRAG framework is compatible with all retrievers, this study only evaluates 3 retrievers. Future work should explore the effectiveness of retrieval strategies across a wider range of retrievers, including dense and sparse retrievers.

• Fixed prompt format: the integration of retrieved examples into prompts lacks systematic exploration of formatting strategies, which can influence LLM performance. Future work should investigate how to improve the performance by sophisticated designed prompt format.

• Four modes in our framework are evaluated separately. In the future, weighing or prioritization mechanisms could be considered to integrate these 4 modes.

By addressing these limitations, retrieval-augmented ICL can be further enhanced for broader applications in AI-driven biomedical and other large-scale NLP tasks.

Conclusion

In this study, we proposed MMRAG, a novel multi-mode retrieval-augmented framework that enhances ICL in biomedical NLP by optimizing example selection. By systematically evaluating 4 retrieval strategies—Random, Top, Diversity, and Class Mode—our approach demonstrates superior adaptability across multiple biomedical tasks. The results show that retrieval-enhanced selection significantly improves information extraction, particularly in precision-critical tasks. As AI-driven healthcare evolves, MMRAG provides a flexible and efficient method to address data scarcity and privacy challenges, paving the way for future advancements in clinical decision support, drug discovery, and biomedical text processing.

Author contributions

Zaifu Zhan (Conceptualization, Methodology, Software, Data curation, Formal analysis, Writing – original draft, Visualization), Jun Wang (Writing – review & editing), Shuang Zhou (Writing – review & editing), Jiawen Deng (Writing – review & editing), Rui Zhang (Conceptualization, Supervision, Project administration, Funding acquisition, Writing – review & editing)

Supplementary material

Supplementary material is available at Journal of the American Medical Informatics Association online.

Funding

This work was supported by the National Institutes of Health’s National Center for Complementary and Integrative Health (grant numbers R01AT009457 and U01AT012871), the National Institute on Aging (grant number R01AG078154), the National Cancer Institute (grant number R01CA287413), the National Institute of Diabetes and Digestive and Kidney Diseases (grant number R01DK115629), and the National Institute on Minority Health and Health Disparities (grant number 1R21MD019134-01).

Conflicts of interest

The authors have no competing interests to declare.

Data availability

All the datasets in this paper are public. Readers could find them in the following links:

BC2GM: https://biocreative.bioinformatics.udel.edu/tasks/

Health Advice: https://huggingface.co/datasets/medalpaca/medical_meadow_health_advice

DDI: https://github.com/isegura/DDICorpus

GIT: https://github.com/ToneLi/BIoMedRAG/tree/main/dataset/0_GM-CIHT