Data-efficient biomedical in-context learning: a diversity-enhanced submodular perspective

Abstract

Recent progress in large language models (LLMs) has leveraged their in-context learning (ICL) abilities to enable quick adaptation to unseen biomedical NLP tasks. By incorporating only a few input-output examples into prompts, LLMs can rapidly perform these new tasks. While the impact of these demonstrations on LLM performance has been extensively studied, most existing approaches prioritize representativeness over diversity when selecting examples from large corpora. To address this gap, we propose Dual-Div, a diversity-enhanced data-efficient framework for demonstration selection in biomedical ICL. Dual-Div employs a two-stage retrieval and ranking process: First, it identifies a limited set of candidate examples from a corpus by optimizing both representativeness and diversity (with optional annotation for unlabeled data). Second, it ranks these candidates against test queries to select the most relevant and non-redundant demonstrations. Evaluated on three biomedical NLP tasks (named entity recognition (NER), relation extraction (RE), and text classification (TC)) using LLaMA 3.1 and Qwen 2.5 for inference, along with three retrievers (BGE-Large, BMRetriever, MedCPT), Dual-Div consistently outperforms baselines-achieving up to 5% higher macro-F1 scores—while demonstrating robustness to prompt permutations and class imbalance. Our findings establish that diversity in initial retrieval is more critical than ranking-stage optimization, and limiting demonstrations to 3-5 examples maximizes performance efficiency.

Introduction

Recent years have witnessed the rapid advancement of large language models (LLMs), exemplified by prominent instances like ChatGPT¹, Llama², and Qwen³ series. These models have significantly enhanced few-shot capabilities⁴ across numerous natural language processing (NLP) tasks. Within the biomedical domain, however, a critical challenge persists: the scarcity of high-quality training data. This scarcity arises from two primary factors. Firstly, stringent privacy regulations, such as HIPAA, and patient consent requirements strictly limit access to sensitive patient information. Secondly, rare diseases often suffer from a paucity of structured clinical records and well-defined features⁵, demanding robust and generalizable algorithmic solutions⁶. Consequently, the few-shot learning paradigm, particularly its ability to perform tasks without task-specific training data, becomes critical⁷.

In-Context Learning (ICL)⁸, a prominent few-shot technique, offers significant promise for biomedical NLP. ICL leverages task-specific prompts containing a few annotated examples, enabling pre-trained LLMs to perform unseen tasks without parameter updates. This approach typically involves retrieving a small set of relevant examples from a large unlabeled corpus and annotating them according to task requirements. ICL confers several key advantages for biomedical NLP. Firstly, it simplifies the deployment of pretrained LLMs by replacing computationally intensive fine-tuning with prompt design. More importantly, since a task-specific fine-tuning dataset is no longer required, ICL also reduces the need for labeled data for downstream tasks. By avoiding parameter updates, ICL can potentially improve model stability across diverse tasks.

Crucially, the performance of ICL is highly dependent on the quality of the selected demonstrations. Effective demonstrations provide essential context, supplementing label information and relationships beyond what was learned during pre-training⁹. This is especially vital in biomedicine, where texts are dense with specialized terminology (e.g., genes, cells, pathways). Since domain-specific knowledge cannot be easily integrated post-hoc into general LLMs, carefully curated demonstrations become the primary mechanism for activating or supplementing the relevant internal knowledge encoded during pre-training.

A significant challenge for existing ICL frameworks lies in selecting optimal demonstrations. Traditional methods often rely on a single metric, typically prioritizing examples most representative of or semantically similar to the test query. More recent approaches¹⁰ have adopted active learning principles to select groups that maximally cover the semantic space of the corpus. However, a key limitation persists: the neglect of diversity among the selected examples. Providing diverse demonstrations enhances model robustness by exposing it to varied scenarios. In biomedicine, this diversity is critical not only for performance but also for improving fairness in disease representation and supporting more comprehensive clinical decision-making.

To address this gap, we propose Dual-Div, a novel diversity-enhanced ICL framework for biomedical NLP that explicitly optimizes demonstration selection across multiple dimensions – representativeness and diversity. Dual-Div operates in two stages (Fig. 1), i.e., demonstration retrieval and ranking. In the first stage, we recall a limited set of candidate examples from a large (often unlabeled) corpus, followed by annotation if necessary. In the second stage, these candidates are ranked in conjunction with test queries, selecting the top demonstrations for prompt construction. Dual-Div provides a systematic and effective solution for balancing broad semantic coverage with intrinsic diversity within the selected demonstrations. Our comprehensive evaluations demonstrate that this approach yields significant improvements in ICL performance over prior methods. Our key contributions are summarized as follows.

• To the best of our knowledge, this is the first study to systematically integrate diversity metrics into the demonstration selection process for in-context learning.

• We propose a novel two-stage submodular optimization framework that jointly maximizes semantic coverage and diversity of selected demonstrations.

• We conduct extensive experiments, evaluating the framework across 2 LLMs and 3 retrievers on 3 datasets, demonstrating state-of-the-art results and providing comprehensive insights.

Fig. 1. The workflow of Dual-Div, a two-stage in-context learning framework.

Source: Created by the authors using draw.io (https://www.draw.io/).

Results

Experimental Setup

We evaluate the ICL performance on three core biomedical NLP tasks, including named entity recognition (NER), relation extraction (RE), and text classification (TC). These tasks can comprehensively demonstrate the ability of LLMs in biomedical natural language understanding and clinical decision support. For each task, we select one typical dataset – ChemProt for NER, DDI for RE, and HealthAdvice for TC.

• ChemProt¹¹ This dataset consists of 1,820 PubMed abstracts with chemical-protein interactions annotated by domain experts. Therefore, it can also be used to evaluate the systems that are able to detect the biomedical terminology related to chemical compounds/drugs and genes/proteins.

• DDI¹². This dataset is identified as a typical biomedical relation extraction benchmark to evaluate the ability of extracting drug-drug interactions. The potential types includes advice, effect, mechanism, and int. All drug entities and interactions in sentences are annotated from biomedical literature and drug product information sources.

• HealthAdvice¹³. This dataset is a diverse collection of health-related advice, annotated for text classification tasks related to health information and advisory content. Structured to facilitate automatic classification of health advice into relevant categories, it can support applications such as misinformation detection, personalized health recommendations, and automated triaging of medical inquiries.

We note that the label distributions in the selected datasets are highly imbalanced. For instance, in the DDI dataset, only 979 out of 5,761 test queries contain one of the four target relations (effect, advice, mechanism, or int), whereas this ratio is 23.6% for the HealthAdvice dataset. Given this imbalance, we employ macro-F1 score – in addition to accuracy – to better reflect in-context learning performance across tasks. Unlike micro-F1, macro-F1 equally weights all classes, making it more rational for biomedical applications where rare classes are often as critical as frequent ones¹⁴.

For baseline comparisons, we employ the heuristic Random-Similar algorithm in addition to the most recent Div-S3¹⁰ method. Similar to our Dual-Div approach, both of these methods follow a two-stage procedure. In the first stage, Random refers to randomly selecting examples from the entire set V, whereas Div defines f(S) solely as $\mathcal{C}\left(S\right)$, omitting the diversity term $\mathcal{D}\left(S\right)$. In the second stage, Similar chooses the most similar examples conditioned with test set Q (with the highest average cosine similarity scores), while the S3 simplifies the utility function by reducing f(T ∣ Q) to $\mathcal{C}(T\mid Q)$. To clearly evaluate the impact of the diversity metrics $\mathcal{D}\left(S\right)$ and $\mathcal{D}\left(T\mid Q\right)$, we also introduce two variants of Div-S3: Div*-S3 and Div-S3*. Div*-S3 adopts the same objective as our Dual-Div in the first stage while keeping the second stage unchanged, whereas Div-S3* applies the original first stage of Div-S3 but uses our proposed second-stage objective.

The involved language models cover two categories. The smaller retriever models are utilized to quickly obtain the semantic representations of biomedical resources. With these semantic vectors, we can determine the demonstration examples in natural language and derive the final prompts for LLMs. In contrast, the LLMs, with many more parameters, are more powerful in solving the concrete downstream tasks.

For retriever models, we select three different ones, including BMRetriever¹⁵, MedCPT¹⁶, and BGE-Large¹⁷. Among them, BMRetriever and MedCPT are trained on biomedical resources, whereas BGE-Large is a family of embedding models trained on general data. For the inference LLMs, we select the representative Qwen2.5-7B³ and Llmama 3.1-8B² to balance the performance and efficiency during evaluation.

For the remaining implementation details, we supplement them in Supplementary Section 1. In particular, we list the prompt template we utilize in Supplementary Section 2.

ICL Performance

Table 1 presents the in-context learning performance (accuracy and macro-F1) across different tasks and methods. Utilizing the BGE-large retriever model, Dual-Div achieves the highest macro-F1 score regardless of whether Qwen2.5 or Llama serves as the inference LLM. This result indicates that the demonstration examples selected by Dual-Div effectively enhance the inference capability of LLMs on biomedical resources. Notably, Dual-Div also maintains a strong balance between accuracy and macro-F1 performance. This balance likely stems from our method’s trade-off between coverage and diversity when selecting demonstrations. Given the significant class imbalance inherent in biomedical queries (where instances with meaningful relations are vastly outnumbered by non-biomedical ones), we argue that macro-F1 here serves as a more reliable and persuasive evaluation metric than accuracy. Results for the other two retriever models are provided in Supplementary Table 2 and Supplementary Table 3; overall, these findings align with our observations using the BGE-Large retriever. We additionally compare the three retriever models and observe that no single model consistently outperforms the others across all datasets. For instance, BMRetriever achieves the best results on ChemProt, MedCPT performs strongest on DDI, and BGE-Large excels on HealthAdvice—highlighting dataset-dependent performance variations.

Table 1.

ICL Performance on different NLP datasets using Llama 3.1 (8B) and Qwen 2.5 (7B) as inference LLMs and BGE-Large as the retriever (Diff rows report Dual-Div − Div-S3 as mean difference with 95% CI and paired t-test p-value from 10-fold cross-validation)

LLM	Method	ChemProt (NER)		DDI (RE)		HealthAdvice (TC)
		acc.	macro-F1	acc.	macro-F1	acc.	macro-F1
Qwen 2.5	Random-Simliar	0.6488	0.6764	0.7837	0.2515	0.7592	0.3306
	Div-S3	0.7585	0.7067	0.7004	0.2843	0.7661	0.3330
	Div*-S3	0.7691	0.7029	0.8047	0.2794	0.7615	0.3604
	Div-S3*	0.7212	0.7004	0.7086	0.2849	0.7598	0.3623
	Dual-Div	0.7598	0.7117	0.7477	0.2949	0.7642	0.3686
	Diff of acc.	[0.0015, 0.0139]**		[0.0309, 0.0473]***		[− 0.0125, 0.0018]
	Diff of macro-F1	[0.0006, 0.0118]*		[− 0.0007, 0.0225]		[0.0101, 0.0449]**
Llama 3.1	Random-Simliar	0.6046	0.5819	0.6888	0.2646	0.6832	0.2780
	Div-S3	0.6921	0.6026	0.7301	0.2653	0.7160	0.2892
	Div*-S3	0.6912	0.6345	0.7351	0.3086	0.7586	0.2886
	Div-S3*	0.6558	0.6183	0.7256	0.2664	0.6872	0.2875
	Dual-Div	0.6910	0.6387	0.7283	0.3129	0.7241	0.2948
	Diff of acc.	[− 0.0133, 0.0147]		[−0.0019, 0.0208]*		[0.0075, 0.0393]***
	Diff of macro-F1	[0.0211, 0.0432]***		[0.0272, 0.0640]***		[− 0.0015, 0.0108]

Note: ^*p < 0.05, ^**p < 0.01, ^***p < 0.001.

To further validate the statistical significance of the improvements achieved by our Dual–Div method over Div–S3, we performed a 10–fold paired t–test based on a 10–fold partitioning of the entire dataset. Specifically, for given method pair, we partitioned the full dataset into 10 folds using a fixed random seed, and in each fold we evaluated the algorithms on the held–out fold. We then computed per–fold performance metrics for both methods on the same test split, derived fold–wise performance differences, and conducted a paired t–test, reporting the mean difference, the standard deviation of differences, the p–value, and the 95% confidence interval of the mean difference (based on the t–distribution). Among all 36 comparisons across different settings, Dual-Div outperforms Div-S3 in 31 cases, with 23 of these advantages being statistically significant. Focusing on the more critical macro-F1 metric, our method shows superior performance in all 18 results, 11 of which are statistically significant.

For the comparison of the two LLMs, Qwen2.5 significantly outperforms Llama 3.1 on NER and TC tasks, regardless of the retrievers used. This suggests Qwen2.5 may have a stronger potential for transferring biomedical knowledge from external inputs. In practice, we also observed more formatting errors in Llama 3.1’s outputs than in Qwen2.5’s. For example, Llama 3.1 responses for relation types sometimes include (partial) parentheses. This issue diminishes as the number of fine-tuning steps over the corpus V increases.

Another critical finding emerges from our detailed comparison of introducing the diversity term $\mathcal{D}\left(S\right)$ at different stages of the pipeline, as reflected by the performance of Div*-S3 and Div-S3*. We find that incorporating diversity regularization in Stage 1 consistently leads to substantially improved ICL performance, often yielding the highest accuracy across settings. This indicates that retrieving a highly diverse candidate set from the original corpus V is more influential than applying a sophisticated ranking strategy in Stage 2. The effect is particularly pronounced when the reduction in candidate size during Stage 1 (N − k₁) is larger than that in Stage 2 (k₁ − k), highlighting the decisive role of early-stage diversity in shaping final outcomes.

More broadly, the superior accuracy achieved by our Div variants (even over Dual-Div) across most experimental configurations robustly validates our central hypothesis: explicitly incorporating diversity considerations at different stages of ICL yields tangible performance benefits. However, a closer examination reveals an important trade-off. Despite their strong accuracy, these variants often exhibit lower macro-F1 scores, indicating a tendency to favor majority-class predictions while being less effective on queries associated with infrequent, minority labels.

This observation highlights both the strengths and limitations of the Div variants in comparison to our proposed Dual-Div: although they improve overall correctness by emphasizing representative and frequent patterns, their performance can diminish on rarer yet semantically important cases. This is precisely why we retain the diversity regularization in Stage 2. The finding further justifies our use of macro-F1 as the primary evaluation metric, since it more accurately reflects performance on minority classes—which often carry greater practical significance.

Visualization of diversity

To better illustrate the effect of the diversity-enhanced technique, we also visualize the semantic representations of the 100 candidate queries from the first retrieval stage. Specifically, we use UMAP¹⁸ to project the high-dimensional semantic embedding vectors from the retriever models onto a 2D plane. We then apply MinMaxScaler for normalization to facilitate visualization. The results, summarized in Fig. 2, show that queries retrieved using the diversity metric exhibit greater dispersion in their vector representations compared to those retrieved without it. This difference is particularly pronounced on the DDI dataset for the relation extraction task. This outcome aligns with our intuition of $\mathcal{D}\left(S\right)$ regarding the determinant of the Gram matrix of these vectors. By maximizing the “volume” of the space spanned by the semantic vectors, we encourage the model to avoid selecting queries clustered in a small neighborhood, which are likely to represent the same biomedical terminology or phenomenon.

Fig. 2. Visualization of retrieved queries in S^* after Stage 1.

Semantic embedding vectors from various retriever models on a ChemProt w/o $\mathcal{D}\left(S\right)$, b DDI w/o $\mathcal{D}\left(S\right)$, c HealthAdvice w/o $\mathcal{D}\left(S\right)$, d ChemProt w/ $\mathcal{D}\left(S\right)$, e DDI w/ $\mathcal{D}\left(S\right)$, and f HealthAdvice w/ $\mathcal{D}\left(S\right)$ during optimization are reduced and projected by UMAP.

In the comparison of the three retriever models, we observe that without diversity-enhancing techniques, the embeddings generated by BGE-Large and MedCPT exhibit a more scattered and uniform distribution compared to those from BMRetriever. This pattern may indicate a stronger capability of these two models in shaping the embedding space and facilitating instruction-agnostic retrieval. However, when diversity enhancement is applied, the distributions of demonstrations retrieved by all three models become considerably more competitive.

Sensitivity Analysis

In order to investigate the sensitivity of our Dual-Div algorithm, we mainly study the effect of internal ordering of final demonstration sets, as well as the value of the budget constraint k, which limits the number of examples in the prompts.

Previous studies^19–21 have shown that the ICL performance of LLMs is highly sensitive to the ordering of demonstrations within input prompts. To investigate this effect, we evaluate the performance distributions across all possible permutations of the examples in the final set T (default cardinality 3). Fig. 3 presents the results for our Dual-Div method and all baselines, using BGE-Large as the retriever. Results for other retriever models are provided in Supplementary Section 4.2.

Fig. 3. Sensitivity analysis regarding the permutation ordering in the LLM’s prompts of the final demonstration set.

Given BGE-Large as the retriever model, for each set of 3 demonstrations from different methods, we plot the violin graphics of macro-F1 scores of all 6 possible permutations on the a ChemProt, b DDI, and c HealthAdvice datasets.

Overall, Dual-Div exhibits the strongest robustness against prompt permutations. On ChemProt and HealthAdvice, Dual-Div not only achieves the highest median macro-F1 scores but also shows the tightest performance distribution (narrowest violin shape), indicating minimal sensitivity to permutation order. Conversely, Random-Similar displays the broadest performance spread across all datasets, likely due to the inherent randomness introduced in its first stage.

On the DDI dataset, Dual-Div shows a much wider distribution than Div-S3 and Div-S3*, but achieves superior average performance. This suggests that prioritizing high diversity in candidate queries during the first stage may reduce stability when using BGE-Large. However, we posit that this effect is both task-specific and retriever-dependent. Supplemental experiments confirm that incorporating diversity significantly improves stability intervals: for ChemProt using BMRetriever (Supplementary Fig. 1a) and for DDI using MedCPT (Supplementary Fig. 2b).

Figure 4 and Figure 5 illustrate the trend of macro-F1 scores as the budget limit k for the second stage varies. Overall, Qwen 2.5 consistently outperforms Llama 3.1 across all datasets and retriever models, indicating superior inference capabilities from in-context information for biomedical tasks. This finding aligns with our previous observations. Regarding specific tasks, the performance improvement is most pronounced on ChemProt (approximately 9.1% at the optimal k) and HealthAdvice (approximately 47.5%). Concerning the choice of k, increasing its value generally enhances performance across all LLMs and tasks. The optimal macro-F1 performance is typically achieved at k = 3 or 5. However, performance gains diminish beyond k = 3, becoming marginal. Notably, further increasing k from 5 to 10 often results in negligible or even negative gains across nearly all experiments. This pattern suggests that the quality of demonstrations is more critical than quantity. An excessive number of examples may potentially impair LLM performance by introducing redundant or irrelevant knowledge.

Fig. 4. Sensitivity analysis of the budget constraint in the second stage using Qwen 2.5 as the inference LLM.

(a) Macro-F1 performance on the ChemProt dataset under different values of K. (b) Macro-F1 performance on the DDI dataset under different values of K. (c) Macro-F1 performance on the HealthAdvice dataset under different values of K. Symbols: red circles = BGE-Large; blue squares = BMRetriever; green triangles = MedCPT.

Fig. 5. Sensitivity analysis of the budget constraint in the second stage using Llama 3.1 as the inference LLM.

(a) Macro-F1 performance on the ChemProt dataset under different values of K. (b) Macro-F1 performance on the DDI dataset under different values of K. (c) Macro-F1 performance on the HealthAdvice dataset under different values of K. Symbols: red circles = BGE-Large; blue squares = BMRetriever; green triangles = MedCPT.

Furthermore, regarding the fluctuation of macro-F1 scores with varying k across different retriever models, BGE-Large generally exhibits greater robustness compared to BMRetriever and MedCPT. This robustness likely stems from BGE-Large’s training on general-domain data, rather than being specialized solely for biomedical contexts.

Case Study

We further provide case studies to demonstrate the enhanced ability in extracting diverse examples from training corpora. The comparison in Figure 6 showcases that the previous Div-S3 method is more likely to select top examples with the same NER tags, while our proposed Dual-Div can greatly alleviate this issue by reducing the similarity within the ranked examples. See more comparisons regarding the other two tasks in Supplementary Fig. 3 and Supplementary Fig. 4.

Fig. 6.

Case study of ranked demonstrations on the ChemProt dataset for NER task from (diversity-enhanced) Dual-Div and Div-S3, with BGE-Large as the retriever and Qwen 2.5 as the inference LLM.

Discussion

In this work, we introduce Dual-Div, a two-stage framework that optimizes demonstration selection for biomedical ICL by jointly maximizing semantic coverage and diversity through submodular optimization. Leveraging a determinantal point process (DPP) for diversity and a facility location function for representativeness, Dual-Div first retrieves a diverse candidate set from large corpora, then ranks examples conditioned on test queries to minimize redundancy. Evaluated across three biomedical NLP tasks (NER, RE, TC) using multiple LLMs (Llama 3.1, Qwen 2.5) and retrievers (BGE-Large, BMRetriever, MedCPT), Dual-Div consistently outperforms baselines—achieving higher macro-F1 scores (up to +5%) — while proving robust to prompt permutations and class imbalance. Key insights reveal that diversity in initial retrieval (Stage 1) is more critical than ranking (Stage 2), and limiting demonstrations to 3-5 examples optimizes performance, establishing Dual-Div as a data-efficient solution for biomedical ICL.

For future research, we plan to explore the dynamic adjustment of the trade-off parameter λ in Eq. (10). Another important direction is improving scalability to extremely large corpora, which could be addressed by developing distributed variants of the lazy greedy algorithm. Furthermore, as indicated in ^22,23, a more principled approach to balancing similarity metrics and word frequency in queries involves using discounted cosine similarity, where the discounting parameters can be optimized via Monte Carlo Bayesian optimization. Applying this framework may lead to further performance improvements.

Methods

As model and data scales increase, ICL emerges as a distinctive capability of LLMs^4,8. This paradigm enables LLMs to perform unseen biomedical tasks by learning directly from demonstration examples provided within their input context. Mirroring human analogical reasoning, ICL offers an interpretable interaction mechanism: it explicitly activates domain-specific knowledge encoded within LLM parameters through natural language demonstrations, bypassing traditional parameter updates.

Unlike retrieval-augmented generation (RAG)²⁴, which dynamically retrieves external, up-to-date knowledge relevant to a query, ICL primarily queries the LLM’s internal knowledge. Its efficacy hinges on the model’s ability to infer task requirements solely from the provided context demonstrations. Consequently, ICL performance is highly sensitive to the quality, relevance, and composition of these selected demonstrations.

Crucially, while ICL was originally conceived as a training-free method to reduce computational cost, many current approaches incur significant auxiliary costs during demonstration selection. For instance, some methods require data preparation and supervised training to learn specialized retrievers^25,26, while others leverage the LLM’s inference capabilities during selection (e.g., using confidence scores)²⁷, adding latency and compute overhead. This paper specifically focuses on learning-free, end-to-end ICL frameworks. We prioritize methods where the LLM itself is not utilized within the demonstration selection algorithm; its role is confined to the final inference step after demonstrations are fully determined.

Retrieving effective demonstrations from large corpora (labeled or unlabeled) often starts with simple heuristics like cosine similarity, which remains surprisingly effective and widely used^28–30. Moving beyond basic similarity, research has explored optimization objectives like the facility location function³¹, which emphasizes the coverage or representativeness of the selected set. While Div-S3¹⁰ offers a theoretically grounded framework using this function, it primarily addresses coverage and neglects diversity. The more recent diversity-guided search³² addresses this gap by incorporating heuristic beam search to promote diverse examples. Finally, it is essential to note that once demonstrations are selected, factors like prompt formatting and the permutation order of examples within the context significantly influence the accuracy of the LLM’s response^19–21.

Mathematical Formulation

ICL is a kind of few-shot learning method enabling LLMs to adapt to new tasks with only a small number of task-specific demonstration examples⁴. Let V denote the corpus, with or without task-specific labels, composed of N = ∣V∣ datapoints ${\left(x_i,y_i,{\mathbf{e}}_i\right)}_{i=1}^N$, where x_i is the input sentence, y_i is the label, and e_i is the representation vector from a retriever model. Here, we denote this retriever model by g( ⋅ ), which means e_i = g(x_i).

Given a test query set Q that includes samples (x_test, y_test), the target is to select a subset T^* ⊆ V such that the likelihood of the target answer y_test from the scoring distribution outputted by LLMs, denoted by ${\mathrm{\Pi }}_{\left(x_{\mathrm{test}},y_{\mathrm{test}}\right)\in Q}P\left(y_{\mathrm{test}}\mid x_{\mathrm{test}}\right)$, is maximized. During this process, a major destination is to limit the cardinality of T^*, i.e., ensuring that ∣T^*∣≤k with k < < N. Equivalently, this constraint can also be considered as the maximal sequence length of the prompt instruction³³.

In general, this process can be further split into two procedures, as illustrated in Figure 1. For the first one, a limited number (namely k₁, k < k₁ < < N) of candidate examples, denoted by S^*, are retrieved from the corpus V, in terms of the coverage and diversity. Intuitively, what we expect is to select a subset S that contains the most relevant and non-redundant examples from the corpus set V. As a result, the two metrics for coverage and diversity are critical. Let $\mathcal{C}\left(S\right)$ and $\mathcal{D}\left(S\right)$ be the corresponding notations. The final utility function that is composed of these two perspectives can be written as

$$f\left(S\right)=\mathcal{C}\left(S\right)+\lambda \mathcal{D}\left(S\right),$$ 1

where λ is a trade-off parameter. With f(S) as our optimization objective, we can conclude the first procedure as

$$S^{*}=\underset{S\subseteq V,\mid S\mid \le k_1}{\arg \max }f\left(S\right).$$ 2

If the samples in S^* are unlabeled, an extra annotation step is required. However, this approach has already reduced costs by avoiding annotation of the entire corpus V.

In the second stage, given the test query set Q, the informative examples in the subset S^* are ranked using the same metrics — now applied conditionally as f(T ∣ Q) with T ⊆ S^*. The optimization problem associated with this step can be written as

$$T^{*}=\underset{T\subseteq S^{*},\mid T\mid \le k}{\arg \max }f(T\mid Q).$$ 3

Finally, the selected demonstrations in T^* and the test query x_test in Q will be used to construct the prompt for LLM inference. The evaluation is based on the likelihood

$$\prod _{(x_{\mathrm{test}},y_{\mathrm{test}})\in Q}P(y_{\mathrm{test}}\mid x_{\mathrm{test}})=\prod _{(x_{\mathrm{test}},y_{\mathrm{test}})\in Q}f(y_{\mathrm{test}}\mid \mathcal{T}(T^{*},x_{\mathrm{test}})),$$ 4

where $\mathcal{T}$ denotes the instruction prompt template regarding the examples, test query, and task descriptions.

Again, our proposed framework aims to address two important issues associated with efficient in-context learning:

1. Given a large corpus V, how can we identify a subset with cardinality constraint to keep the most information while ensuring the internal diversity at the same time?

2. Given a query Q, how can we decide the most relevant and non-redundant demonstrations from S^*?

Metric Selection

Ideally, we desire to find suitable monotone submodular functions to describe diversity and coverage such that the weighted form f(S) still belongs to monotone submodular functions. With this assumption, we can obtain a theoretically guaranteed optimal solution.

To seek an elegant probabilistic model to describe the probabilities of subsets S, we leverage the algebraic properties of determinantal point processes (DPPs)³⁴. Specifically, we select the L-ensemble representation of DPPs with a semidefinite matrix L. Given a subset S ⊆ V, the probability that it would be selected is

$$\mathcal{P}\left(S\right)=\frac{\det \left({\mathbf{L}}_S\right)}{\det \left(\mathbf{L}+\mathbf{I}\right)}$$ 5

where ${\mathbf{L}}_S={\left[{\mathbf{L}}_{ij}\right]}_{i,j\in S}$ denotes the restriction of L to the entries indexed by elements in S, det( ⋅ ) denotes the determinant of a matrix, and I represents the identity matrix. Note that $\mathcal{P}\left(S\right)$ is normalized because of the equation ∑_S⊆Vdet(L_S) = det(L + I) according to Theorem 2.1 in ³⁴.

Similar to other scenarios like recommendation³⁵, we can construct the kernel matrix in the context of ICL by the Gram matrix L = E^TE, where the i-th columns of E corresponds the vector e_i representing the query i in V. Intuitively, the determinant of the gram matrix L describes the “volume” of the space spanned by all query representations in E. Therefore, if the queries in S are more various, the spanned space will be greater, resulting in a greater determinant of L_S. Furthermore, if we rewrite e_i as the product of the L₂ norm r_i≥0 and a normalized vector ${\overline{\mathbf{e}}}_i$ with $\mid \mid {\overline{\mathbf{e}}}_i\mid {\mid }_2=1$. The elements of kernel L can be written as

$$\begin{array}{ccc}\hfill {\mathbf{L}}_{ij}& \hfill =\hfill & ⟨{\mathbf{E}}_i,{\mathbf{E}}_j⟩\hfill \\ \hfill & \hfill =\hfill & ⟨r_i{\overline{\mathbf{e}}}_i,r_j{\overline{\mathbf{e}}}_j⟩\hfill \\ \hfill & \hfill =\hfill & r_i{r}_j⟨{\overline{\mathbf{e}}}_i,{\overline{\mathbf{e}}}_j⟩\hfill \\ \hfill & \hfill =\hfill & r_i{r}_j{w}_{i,j}.\hfill \end{array}$$ 6

Here, the term w_i,j is exactly the cosine similarity of query i and j. Due to the equation $\mathbf{L}=\mathrm{Diag}\left(\mathbf{r}\right)\cdot \mathbf{W}\cdot \mathrm{Diag}\left(\mathbf{r}\right)$, the problem of maximizing $\mathcal{P}\left(S\right)$ can be thought of as

$$\begin{array}{cc}\underset{S\subseteq V}{\max }\log \mathcal{P}\left(S\right)\hfill & =\underset{S\subseteq V}{\max }\log det\left({\mathbf{L}}_S\right)\hfill \\ \hfill & =\underset{S\subseteq V}{\max }\left[\sum _{i\in S}\log \left(r_i^2\right)+\log det\left({\mathbf{W}}_S\right)\right].\hfill \end{array}$$ 7

This decomposition reveals that the probabilities of DPP models can be further split into a diversity-related factor $\log \det \left({\mathbf{W}}_S\right)$ and a norm-related factor ${\sum }_{i\in S}\log \left(r_i^2\right)$. Meanwhile, prior studies^36,37 have established that the L₂ norm of word embeddings exhibits a linear correlation with the word’s log-frequency in the training corpus. To mitigate the influence of word frequency inherent in the queries, we therefore propose to define the diversity of a subset S independent of the L₂ norms of the corresponding query embeddings ${\left\{{\mathbf{e}}_i\right\}}_{i\in S}$, i.e.,

$$\mathcal{D}\left(S\right)=\log det\left({\mathbf{W}}_S\right),$$ 8

which also aligns with the definition of $\mathcal{C}\left(S\right)$ below.

To describe the coverage, we follow the previous studies^10,31 and define the coverage of a subset S regarding the entire set V as the facility location function

$$\mathcal{C}\left(S\right)=\sum _{i\in V}\underset{j\in S}{\max }{w}_{i,j}=\sum _{i\in V}\underset{j\in S}{\max }⟨{\overline{\mathbf{e}}}_i,{\overline{\mathbf{e}}}_j⟩.$$ 9

As such, the ultimate objective is a balance of diversity and coverage metric as

$$\begin{array}{cc}f\left(S\right)\hfill & =\mathcal{C}\left(S\right)+\lambda \mathcal{D}\left(S\right)\hfill \\ \hfill & =\sum _{i\in V}\underset{j\in S}{\max }{w}_{i,j}+\lambda \log \det \left({\mathbf{W}}_S\right),\hfill \end{array}$$ 10

where ${\left[{\mathbf{W}}_S\right]}_{ij}=w_{i,j}$ for i, j ∈ S. We can prove that f(S) is monotone and submodular for a small trade-off constant λ (See Below).

Submodularity and Monotonicity of f(S)

Proof

First of all, $\mathcal{C}\left(S\right)$ and $\mathcal{D}\left(S\right)$ are both submodular functions. $\mathcal{C}\left(S\right)$ is submodular because the linear combinations of submodular functions are still submodular. $\mathcal{D}\left(S\right)$ is submodular because W_S is positive semi-definite and the determinant of positive semi-definite matrices is log-submodular³⁸. It can also be observed that $\mathcal{C}\left(S\right)$ is monotonically increasing and $\mathcal{D}\left(S\right)$ is monotonically decreasing. The former holds because if we add a new element $x^{\prime }\notin S$ into S, we have ∀ i ∈ V, $\underset{j\in S\cup \left\{x^{\prime }\right\}}{\max }{w}_{i,j}\ge \underset{j\in S}{\max }{w}_{i,j}$. Therefore $\mathcal{C}(S\cup \left\{x^{\prime }\right\})\ge \mathcal{C}\left(S\right)$. Next, we consider the latter. Again, det(W_S) is non-negative since W_S is positive semi-definite. If we add a new element $x^{\prime }\notin S$ into S, we denote $q_S\left({\overline{\mathbf{e}}}_{x^{\prime }}\right)$ by the orthogonal projection residual of the normalized embedding vector ${\overline{\mathbf{e}}}_{x^{\prime }}$ from the retriever model $g\left(x^{\prime }\right)$ in terms of the space spanned by vectors ${\left\{{\overline{\mathbf{e}}}_i\right\}}_{i\in S}$, namely ${\overline{\mathbf{e}}}_{x^{\prime }}-{\mathrm{proj}}_{\mathrm{span}\left(S\right)}\left({\overline{\mathbf{e}}}_{x^{\prime }}\right)$. It holds that $\mid \mid q_S\left({\overline{\mathbf{e}}}_{x^{\prime }}\right)\mid {\mid }^2\le 1$ because $\mid \mid {\overline{\mathbf{e}}}_{x^{\prime }}\mid \mid =1$. Hence, we have $\det \left({\mathbf{W}}_{S\cup \left\{x^{\prime }\right\}}\right)=\det \left({\mathbf{W}}_S\right)\cdot \mid \mid q_S\left({\overline{\mathbf{e}}}_{x^{\prime }}\right)\mid {\mid }^2\le \det \left({\mathbf{W}}_S\right)$, leading to $\mathcal{D}(S\cup \left\{x^{\prime }\right\})\le \mathcal{D}\left(S\right)$.

Intuitively, during the balance of the coverage and diversity of S, if we ensure that λ is small enough that $0<\lambda \le \frac{\mathcal{C}(S\cup \left\{x^{\prime }\right\})-\mathcal{C}\left(S\right)}{\mathcal{D}\left(S\right)-\mathcal{D}(S\cup \left\{x^{\prime }\right\})}$ for all subset S and $x^{\prime }\notin S$, we can obtain a monotonically increasing submodular objective function $f\left(S\right)=\mathcal{C}\left(S\right)+\lambda \mathcal{D}\left(S\right)$. In our context, if we assume the space is high-dimensional and all vectors in E are linearly independent, there exists a constant $\epsilon =1-\underset{i,j\in V}{\min }{w}_{i,j}>0$ such that $\mathcal{C}(S\cup \left\{x^{\prime }\right\})-\mathcal{C}\left(S\right)\le \epsilon$ for all subset S and $x^{\prime }\notin S$. Since the upper bound of $\mathcal{D}\left(S\right)-\mathcal{D}(S\cup \left\{x^{\prime }\right\})$ is also limited, we can guarantee the existence of the constant λ. □

Diversity-Enhanced Retrieval

With the utility function f(S) defined, we first obtain the embedding vector for each sentence in V using retriever models, and compute the pairwise cosine similarity scores to instantiate the value of f(S) for every singleton set $S={\left\{x_i\right\}}_{x_i\in V}$, as specified in Eq. (10). Subsequently, we iteratively select the sample that yields the highest marginal gain f(S ∪ {x}) − f(S), adding it to the set S until the cardinality constraint k₁ is met. During this process, the lazy greedy³⁹ strategy is applied to improve computational efficiency.

Since we only consider the special cardinality constraint so far, where the costs of all elements are identical, it has been shown in ⁴⁰ that the greedy optimization provides a (1 − 1/e)-approximation guarantee. In other words, we can make sure the f(S^*)≥(1 − 1/e)f(S_opt) where S_opt denotes the ideal optimal subset. For a more general case where the constraint ∣S∣≤k₁ in Eq. (2) is replaced with ∑_x∈Sc(s)≤k₁, a modified greedy algorithm proposed in ⁴¹ alternatively offer a $\left(1-1/\sqrt{e}\right)$ performance guarantee factor.

Diversity-Enhanced Ranking

Once we have retrieved a subset S that covers as many semantic meanings of V as possible, we now consider the conditional ranking process given the test query set Q. At a high level, we expect to generate a subset with the largest marginal advantage given the query set Q, where the marginal advantage of a subset T ⊆ S^* is denoted by

$$f(T\mid Q):=f(T\cup Q)-f\left(Q\right).$$ 11

A higher value of f(T ∣ Q) indicates a higher utility gain if we combine the samples in T with those in Q, where the utility depends on our objective, such as the relevance or non-redundancy. Here, we use the same definition of f(T) as used in the previous stage – a balance of coverage and diversity.

In practice, we utilize a modular approximation ${\sum }_{x_i\in T}f\left(\left\{x_i\right\}\mid Q\right)$ as the upper bound of f(T ∣ Q), resulting the final objective of this stage to be

$$\begin{array}{cc}{T}^{*}\hfill & =\underset{\mathrm{T}\subseteq {\mathrm{S}}^{*},\mid \mathrm{T}\mid \le \mathrm{k}}{\mathrm{argmax}}\sum _{x_i\in T}f(\left\{x_i\right\}\mid Q)\hfill \\ \hfill & =\underset{T\subseteq S^{*},\mid T\mid \le k}{\mathrm{argmax}}\sum _{x_i\in T}f(\left\{x_i\right\}\cup Q)-f\left(Q\right).\hfill \end{array}\begin{array}{}\\ \end{array}$$ 12

We also conclude the detailed algorithms for the entire process, including the retrieval and ranking in Supplementary Algorithm 1.

Author contributions

J.W. and Q.Z. conceptualized and designed the study. J.W. and Z.Z. curated the data and executed the experiments. J.W. drafted the initial manuscript. Z.Z., Q.Z., M.L., M.S. and R.Z. reviewed and finalized the manuscript. R.Z. supervised the whole project.

Data availability

The datasets involved in this study are publicly available from the following links: ChemProt (https://huggingface.co/datasets/bigbio/chemprot), DDI (https://github.com/isegura/DDICorpus), and HealthAdvice (https://huggingface.co/datasets/medalpaca/medical_meadow_health_advice).

Code availability

The code is publicly available on GitHub at https://github.com/huoxusg/submodular_icl.

Competing interests

Author Rui Zhang is Editorial Board Members of Journal of Biomedical Informatics and npj Health Systems. Author Rui Zhang is also Associate Editor of npj Digital Medicine. Rui Zhang was not involved in the journal’s review of, or decisions related to, this manuscript.

Supplementary information

The online version contains supplementary material available at https://doi.org/10.1038/s44387-026-00123-0.