KITE: Kernelized and Information Theoretic Exemplars for In-Context Learning

## Abstract

Abstract In-context learning (ICL) has emerged as a powerful paradigm for adapting large language models (LLMs) to new and data-scarce tasks using only a few carefully selected task-specific examples presented in the prompt. However, given the limited context size of LLMs, a fundamental question arises: Which examples should be selected to maximize performance on a given user query? While nearest-neighbor-based methods like KATE have been widely adopted for this purpose, they suffer from well-known drawbacks in high-dimensional embedding spaces, including poor generalization and a lack of diversity. In this work, we study this problem of example selection in ICL from a principled, information theory-driven perspective. We first model an LLM as a linear function over input embeddings and frame the example selection task as a query-specific optimization problem: selecting a subset of exemplars from a larger example bank that minimizes the prediction error on a specific query. This formulation departs from traditional generalization-focused learning theoretic approaches by targeting accurate prediction for a specific query instance. We derive a principled surrogate objective that is approximately submodular, enabling the use of a greedy algorithm with an approximation guarantee. We further enhance our method by (i) incorporating the kernel trick to operate in high-dimensional feature spaces without explicit mappings, and (ii) introducing an optimal design-based regularizer to encourage diversity in the selected examples. Empirically, we demonstrate significant improvements over standard retrieval methods across a suite of classification tasks, highlighting the benefits of structure-aware, diverse example selection for ICL in real-world, label-scarce scenarios. We conduct extensive experiments over multiple classification datasets using several LLM models and demonstrate significantly improved performance achieved by our exemplar retrieval algorithm.

## Introduction

Figure: Figure 1: This figure is divided into two halves, illustrating the Kite’s selection pipeline on the left and the kernelization trick on the right. Left: The selection module (Alg. 1) maintains an inverse design matrix $V^{-1}$ and, for each candidate $x$, computes relevance (Eq. (8)) and diversity (Eq. (10)), combines them into a total score (Eq. (11)), selects the highest‐scoring example, and updates $V^{-1}$ via Sherman–Morrison (Eq. (12)). The chosen exemplars are then fed into the LLM to produce the final prediction $y$. Right: The kernel trick panel shows how every inner product in feature space is replaced by a kernel evaluation, $\langle\phi(x),\phi(x^{\prime})\rangle\;\to\;k(x,x^{\prime})$ (Lemma 1, Eq. (13)), enabling use of Gaussian RBF and polynomial kernels to work implicitly in the reproducing kernel hilbert space (RKHS) without ever computing high‐dimensional feature vectors.
Refer to caption: x1.png

With the advent of highly capable large language models (LLMs) (Adiwardana et al. 2020; Wang et al. 2019; Zhang, Rudra, and Anand 2021; Wang, Dou, and Che 2022), in-context learning (ICL) (Rubin, Herzig, and Berant 2022; Liu et al. 2022; Wu et al. 2022) via prompt optimization has emerged as a powerful technique for generating responses to complex user queries in data-scarce settings. In this popular and practical paradigm, we assume access to a small bank of high-quality task-specific examples. For a given user query, a few unique and relevant demonstrations are selected to form additional context for the language model. Since LLMs are pre-trained on large corpora, even a small number of carefully chosen exemplars can often suffice to guide the model toward producing accurate and task-consistent responses (Luo et al. 2024). The key is to select examples that are both representative and query-relevant, enabling the model to implicitly infer the task. Compared to fine-tuning (see Wang et al. (2025) and references therein), ICL offers a lightweight and efficient alternative, especially when labeled data is limited.

Given the limited context window of large language models, a natural and important question arises for ICL:
“Given a specific user query, how can we optimally select and order a subset of task-specific examples from an associated example bank to include in the prompt to maximize performance?”
While several existing methods address this question empirically (Dong et al. 2022; Luo et al. 2024), our work takes a principled, information-theoretic approach to tackle this problem.

Motivated by data-scarce settings and the need for task generalization, we focus on common and practical scenarios where the exemplar retriever is frozen and non-trainable. Among unsupervised retrieval methods, the most widely used is KATE (Liu et al. 2021), which pioneered a $k$-nearest neighbor (kNN)-based strategy for selecting in-context examples. KATE identifies the $k$ examples most similar to the user query in a pre-trained embedding space (e.g., BERT), drawing on the intuition that nearby points in this space are likely to be the most informative.

However, as with classical kNN, this approach suffers from the curse of dimensionality—a well-documented issue in high-dimensional spaces where distances become less informative (Köppen 2000). In traditional machine learning, this problem is often addressed by assuming a structured hypothesis class (e.g., linear models), which effectively reduces the model complexity and improves generalization guarantees with fewer samples. This motivates our central theoretical question in the ICL context: *“Can we design an example selection algorithm—i.e., a retriever—that, given a user query, operates under a structured modeling assumption to improve selection quality?”* Such an information-theoretic modeling-driven approach can not only mitigate the limitations of high-dimensional retrieval but also implicitly encourage diversity among selected examples - another desirable property for generalization in ICL.

From a theoretical standpoint, we model the LLM as a function that behaves linearly in its input,
and the goal of ICL in the context of LLMs translates to algorithmically selecting a small number of datapoints to train the linear model (conditioned on an input test query).

Suppose we are given a high-quality example bank of $n$ labeled datapoints. For a specific $d$-dimensional user query $\mathbf{z}$, we consider the following core problem:
*“Given a test point $\mathbf{z}$ and an example bank ${(\mathbf{x}_{i},y_{i})}_{i=1}^{n}\in\mathbb{R}^{d}\times\mathbb{R}$, which $k$ examples should be selected to train a linear model such that the prediction error on $z\in\mathbb{R}^{d}$ is minimized?”* Note that a selection algorithm for the aforementioned problem can be directly used for selecting ICL examples for LLMs. (^1^11In fact, if we were using kNN to predict the response for $\mathbf{z}$, then the selected datapoints would have been the $k-$nearest datapoints to $\mathbf{z}$. This is also the intuition for choosing examples in KATE (Liu et al. 2021).)

Theoretically speaking, this departs from standard learning-theoretic goals, which typically aim for generalization over a distribution of test points. In contrast, we focus on query-specific generalization, i.e., training a model that performs well on a single (and potentially arbitrary) test query using a carefully selected subset of training examples. For the linear model, we derive a surrogate loss objective that quantifies this query-specific prediction error and show that this objective is approximately submodular. This key structural insight allows us to apply a greedy selection algorithm with a provable $1-e^{-\gamma}$ approximation guarantee relative to the optimal subset - here $\gamma\in(0,1)$ is the approximate sub-modularity ratio of the objective given the example bank.

Building on this formulation, we design a computationally efficient and fully unsupervised example selection algorithm. We evaluate it across multiple tasks and find that it consistently outperforms strong baselines such as kNN-based top-k retrieval (Liu et al. 2021), DPP-based retrieval (Ye et al. 2023a), and BM25 (TF–IDF) keyword retrieval, demonstrating the effectiveness of our theory-guided retrieval in in-context learning. For example, our algorithm surpasses the strongest baseline, DPP (Ye et al. 2023a) by $2.44\%$ on HellaSwag benchmark (Zellers et al. 2019) with Qwen-2.5-1.5B model (Hui et al. 2024). In this study, our focus is primarily on classification tasks, where the output is a single label, aligning with the assumptions of our theoretical framework. Note that classification tasks encompass a large fraction of use-cases where LLMs are used in a widespread fashion; for instance, LLM-as-a-judge (Zheng et al. 2023), toxicity detection (Gehman et al. 2020), and intent detection (Larson et al. 2019).

In addition to our base algorithm, we incorporate two complementary techniques to further enhance example selection. First, in the theoretical setting, we extend our greedy algorithm under linear models to the general nonlinear models using the well-known kernel trick. This allows our method to operate in high (possibly infinite) dimensional feature spaces without explicitly computing the feature mappings. By replacing inner products in the objective with kernel evaluations, we enable richer representations that capture nonlinear relationships between datapoints. Empirically, we observe that using a well-calibrated kernel can lead to noticeable improvements in LLM performance, suggesting that the choice of feature space plays a critical role in guiding in-context learning. Second, we introduce a mechanism for enforcing diversity in the selected examples. Inspired by maximum information gain theory—a well-established framework in experimental design, bandits, and Bayesian optimization (Lattimore and Szepesvári 2020) — we consider the scenario where the goal is to select a subset of examples such that the resulting trained model performs well across all queries in the example bank. This naturally leads to the classical optimal experimental design problem, whose objective function is well-understood. We incorporate this design objective as a regularizer into our selection criterion to encourage diversity among selected examples. Intuitively, and as supported by our experiments, promoting diversity improves the generalizability of the model and boosts the quality of LLM responses.

## Related Work

#### In-Context Learning (ICL).

In-context learning, introduced by Brown et al. (2020b), is a powerful paradigm where large language models (LLMs) learn new tasks by conditioning on a few input-output examples provided in the prompt, without any parameter updates. The underlying mechanisms of ICL have been extensively studied; some works suggest that ICL enables the prediction task to become linearly separable (Saunshi, Malladi, and Arora 2020) or that it allows the model to infer a shared latent concept from demonstrations (Xie et al. 2021). Further research indicates that ICL is a nuanced process, as models do not always rely strictly on the provided input-output mappings (Min et al. 2022). The ICL capabilities of LLMs can be enhanced through dedicated self-supervised or supervised training procedures (Chen et al. 2022; Min et al. 2021; Wei et al. 2023a). A broad range of studies has also investigated the factors that influence ICL performance, such as demonstration calibration and corpora effects (Zhao et al. 2021; Shin et al. 2022; Wei et al. 2022; Yoo et al. 2022; Wei et al. 2023b), as well as the working mechanisms themselves, including how models perform implicit gradient descent (Olsson et al. 2022; Li et al. 2023b; Pan 2023; Dai et al. 2022).

Example retrieval for few-shot learning.
The effectiveness of ICL is highly sensitive to the choice of in-context examples, a variance quantified in early works (Brown et al. 2020b; Min et al. 2022), which has motivated a significant line of research on demonstration selection. Initial approaches focused on unsupervised heuristics, such as retrieving semantically similar examples using embedding-based nearest neighbors (Liu et al. 2022) or lexical overlap, which showed sizeable gains over random sampling. Subsequent research has advanced beyond fixed heuristics to instead learn a retriever model. These methods often fine-tune a dense retriever to select effective demonstrations, using labels distilled from a language model (Rubin, Herzig, and Berant 2022) or employing reinforcement and contrastive objectives to balance relevance and coverage (Li et al. 2023a; Luo et al. 2023; Wang et al. 2023; Ye et al. 2023b; Ghosal et al. 2025b, a). Orthogonal to retrieval, information‑theoretic criteria such as mutual information (Sorensen, Ding, and Chang 2022) or LM perplexity (Gonen, Hosking, and Augenstein 2023) serve as lightweight proxies for example quality. Influence‑function analysis by Li and Qiu (2023) selects training points that exert the greatest effect on the language model’s prediction. The current state-of-the-art involves subset-level methods that explicitly model interactions between examples, often using determinantal point processes (DPPs) to promote diversity and avoid redundancy (Yang et al. 2023; Ye et al. 2023a). Based on an information-theoretic standpoint, in this paper, we propose a computationally efficient and fully unsupervised example selection framework.

## Theoretical Model

Consider access to a dataset $\mathcal{X}$ containing $n$ known input-output pairs $(\mathbf{x}_{1},y_{1}),(\mathbf{x}_{2},y_{2}),\ldots,(\mathbf{x}_{n},y_{n})\in\mathbb{R}^{d}\times\mathbb{R}$, where features are $d$-dimensional vectors and the output are real scalars.

Further, consider an input test query $\mathbf{z}\in\mathbb{R}^{d}$.
For any $k\ll d$, the ICL problem is defined as selecting the optimal subset of examples from $\mathcal{X}$ to provide as context (to the LLM) or train a model for predicting the response to the input test query $\mathbf{z}$. The goal is to ensure that the prediction error is minimized via an optimal choice of examples. Note that in ICL, the selected subset contains unique elements and cannot be a multi-set. This is because, in an LLM, duplicating the examples does not provide additional context.

To make the problem mathematically tractable, conditioned on the input test query $\mathbf{z}\in\mathbb{R}^{d}$, assume that the model response (in this instance, the LLM) is linear in its input with a distinct underlying parameter vector $\theta\in\mathbb{R}^{d}$. The selected subset of $k$ examples is used to train the model to generate a prediction for the test query $z\in\mathbb{R}^{d}$.(^2^22We emphasize that the underlying parameter vector $\theta\in\mathbb{R}^{d}$ can vary conditioned on the test query $\mathbf{z}$.)
Specifically, conditioned on $\mathbf{z}$, there exists an unknown parameter vector $\boldsymbol{\theta}\in\mathbb{R}^{d}$ such that the response $y$ for any feature vector $\mathbf{x}\in\mathbb{R}^{d}$ is generated as

$$ $y=\langle\mathbf{x},\boldsymbol{\theta}\rangle+\eta,\,\,\text{where}\,\,\eta\sim\mathcal{N}(0,1)~.$ (1) $$

In the above framework, the ICL problem reduces to the following question: For a given test query $\mathbf{z}$, which subset of examples $\mathcal{S}$ of size at most $k$ from $\mathcal{X}$ should we select so that a least squares estimator fitted on $\mathcal{S}$ minimizes the expected error for the prediction on $\mathbf{z}$? Mathematically, given a test query $\mathbf{z}$, we aim to solve the optimization problem:

$$ $\min_{\mathcal{S}\subseteq\mathcal{X}:|\mathcal{S}|\leq k}|\langle\mathbf{z},\boldsymbol{\theta}-\widehat{\boldsymbol{\theta}}_{\mathcal{S}}\rangle|~,$ (2) $$

where $\widehat{\boldsymbol{\theta}}_{\mathcal{S}}=\Big(\beta\mathbf{I}_{d}+\sum_{i\in\mathcal{S}}\mathbf{x}_{i}\mathbf{x}_{i}^{\top}\Big)^{-1}\Big(\sum_{i\in\mathcal{S}}\mathbf{x}_{i}y_{i}\Big)$ is the $\beta$-regularized least squares estimator.

This discrete optimization problem is computationally challenging, with complexity that is exponential in $k$ and polynomial in $n$, stemming from the combinatorial nature of selecting subsets—specifically, the ${n\choose k}$ possible combinations. Note that this poses a severe computational challenge since this runtime will be required per user query. On the other hand, fast algorithms during inference without hurting the accuracy are of paramount importance. Our goal is to develop an efficient algorithm that can find high-quality solutions to the optimization problem in equation ([2](https://arxiv.org/html/2509.15676v1#Sx3.E2)).

## Algorithm Design

We develop a principled approach to solve the in-context learning subset selection problem by leveraging submodular optimization theory (Das and Kempe 2011). Our methodology consists of three main components: problem reformulation using concentration inequalities, theoretical analysis of submodularity, and a greedy algorithm with provable approximation guarantees.

#### Problem Reformulation

To make the optimization problem in equation ([2](https://arxiv.org/html/2509.15676v1#Sx3.E2)) tractable, we use concentration inequalities to bound the prediction error. Applying the Chernoff bound for sub-Gaussian random variables, we can bound, for any $\delta\in(0,1)$, the prediction error as

$$ $|\langle\mathbf{z},\boldsymbol{\theta}-\widehat{\boldsymbol{\theta}}_{\mathcal{S}}\rangle|\lessapprox\sqrt{\|\mathbf{z}\|_{\mathbf{V}_{\mathcal{S}}^{-1}}^{2}\log(1/\delta)}~,$ (3) $$

with probability greater than $1-\delta$, where $\mathbf{V}_{\mathcal{S}}=\beta\mathbf{I}_{d}+\sum_{i\in\mathcal{S}}\mathbf{x}_{i}\mathbf{x}_{i}^{\top}$ is the design covariance matrix, and a constant $\beta$ dependent factor is hidden under $\lessapprox$. This is a well-known result on the prediction error in linear models - for instance, see (Lattimore and Szepesvári 2020, Chapter 20).

The above theoretical bound shows that the upper bound on prediction error is related to the chosen subset $\mathcal{S}$ of training datapoints via the term $\|\mathbf{z}\|_{\mathbf{V}_{\mathcal{S}}^{-1}}^{2}$. Hence, we need to minimize $\|\mathbf{z}\|_{\mathbf{V}_{\mathcal{S}}^{-1}}^{2}$ for a fixed $\mathbf{z}\in\mathbb{R}^{d}$, or, equivalently solve the following optimization problem

$$ $\max_{\mathcal{S}\subseteq\mathcal{X}:|\mathcal{S}|\leq k}f_{\mathbf{z}}(\mathcal{S})~,\,\,\text{where}\,\,f_{\mathbf{z}}(\mathcal{S})=-\mathbf{z}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{z}~.$ (4) $$

The set function $f_{\mathbf{z}}$ is monotonically increasing, which follows from the fact that for any two sets $\mathcal{S}\subseteq\mathcal{L}$, the corresponding design matrices satisfy $V_{\mathcal{S}}\preceq V_{\mathcal{L}}$ in the Lowener order. Along with monotonicity, a desirable property for maximizing set functions is submodularity, defined as:

###### Definition 1 (Submodular Function) .

A set function $f:2^{\mathcal{X}}\rightarrow\mathbb{R}$ is called *submodular* if for all $\mathcal{S}\subseteq\mathcal{L}\subseteq\mathcal{X}$ and $\mathbf{x}\in\mathcal{X}\setminus\mathcal{L}$, it satisfies the diminishing returns property:

$$ $f(\mathcal{S}\cup\{\mathbf{x}\})-f(\mathcal{S})\geq f(\mathcal{L}\cup\{\mathbf{x}\})-f(\mathcal{L})~.$ $$

For monotone and submodular functions, a greedy algorithm admits an $(1-1/e)$-factor approximation.

The set function $f_{\mathbf{z}}$ in equation ([4](https://arxiv.org/html/2509.15676v1#Sx4.E4)) doesn’t satisfy sub-modularity. However, it exhibits approximate submodularity and retains the near-optimality guarantee of greedy algorithms. The notion of approximate submodularity can be captured by the *submodularity ratio* (Das and Kempe 2011), formally defined below.

###### Definition 2 (Submodularity ratio) .

For a ground set $\mathcal{X}$ and a parameter $k\in\mathbb{N}$, the submodularity ratio of a set function $f:2^{\mathcal{X}}\rightarrow\mathbb{R}$ is defined as

$$ $\gamma_{k}(f)=\min_{\begin{subarray}{c}\mathcal{S}\subseteq\mathcal{X},\\ \mathcal{L}\subseteq\mathcal{X}:|\mathcal{L}|\leq k,\\ \mathcal{L}\cap\mathcal{S}=\emptyset\end{subarray}}\frac{\sum_{\mathbf{x}\in\mathcal{L}}\left(f(\mathcal{S}\cup\{\mathbf{x}\})-f(\mathcal{S})\right)}{f(\mathcal{S}\cup\mathcal{L})-f(\mathcal{S})}.$ $$

The submodularity ratio captures how much more the value of the function $f$ can increase by adding any subset $\mathcal{L}$ of size at most $k$ to $\mathcal{S}$, compared to the
combined benefits of adding its elements to $\mathcal{S}$. Hence, it quantifies the degree to which $f$ satisfies the diminishing returns property, and hence how well greedy algorithms are expected to perform in maximizing $f$ under cardinality constraints. If the set function $f$ is submodular, then $\gamma=1$. When $0<\gamma<1$, the function is said to be approximately submodular. The next result lower bounds the submodularity ratio of $f_{\mathbf{z}}$.

###### Lemma 3 (Submodularity ratio of $f_{\mathbf{z}}$ ) .

For any $\mathbf{z}\in\mathbb{R}^{d}$, the submodularity ratio of $f_{\mathbf{z}}(\mathcal{S})=-\mathbf{z}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{z}$ satisfies

$$ $\displaystyle\gamma_{k}(f_{\mathbf{z}})\geq\frac{1}{1+(k-1)\mu}~,$ $$

where $\mu$ is the maximum coherence between any pair of elements in $\mathcal{X}\setminus\mathcal{S}$, given by $\mu=\max_{\mathbf{x}_{i},\mathbf{x}_{j}\notin\mathcal{S}}|\mu_{i,j}|$, with

$$ $\displaystyle\mu_{i,j}=\frac{\mathbf{x}_{i}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{j}}{\sqrt{1+\mathbf{x}_{i}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{i}}\sqrt{1+\mathbf{x}_{j}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{j}}}~.$ $$

Lemma [3](https://arxiv.org/html/2509.15676v1#Thmtheorem3) gives a query-independent lower bound on the submodularity ratio of $f_{\mathbf{z}}$. In practice, for a given query $\mathbf{z}$, sets $\mathcal{S}$ and $\mathcal{L}=\{\mathbf{x}_{1},\ldots,\mathbf{x}_{k}\}\subset\mathcal{X}\setminus\mathcal{S}$, one can compute the submodularity ratio (with abuse of notation) as

$$ $\displaystyle\gamma_{k}(f,\mathbf{z},\mathcal{L},\mathcal{S})=\frac{\sum_{i=1}^{k}\Delta_{i}}{\sum_{i=1}^{k}\Delta_{i}-\sum_{i\neq j}\sqrt{\Delta_{i}\Delta_{j}}\mu_{i,j}}~,$ $$

where $\Delta_{i}=\frac{(\mathbf{z}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{i})^{2}}{1+\mathbf{x}_{i}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{i}}$
denotes the marginal gain in $f_{\mathbf{z}}$ of adding $\mathbf{x}_{i}$ to $\mathcal{S}$. The lower bound in Lemma [3](https://arxiv.org/html/2509.15676v1#Thmtheorem3) is derived by bounding this ratio using the Cauchy-Schwartz inequality and the maximum coherence $\mu$. Detailed proof of Lemma [3](https://arxiv.org/html/2509.15676v1#Thmtheorem3) is deferred to the appendix. The lower bound is empirically verified for all the datasets used in our experiments.

###### Definition 1 (Submodular Function) .

###### Definition 2 (Submodularity ratio) .

###### Lemma 3 (Submodularity ratio of f 𝐳 f_{\mathbf{z}} ) .

### Greedy Algorithm

We are now ready to present our Greedy algorithm for in-context example selection.
Using the Sherman-Morrison formula, for any $\mathcal{S^{\prime}}\subset\mathcal{S}$ such that $\mathcal{S}\setminus\mathcal{S^{\prime}}=\{\mathbf{x}\}$, we get

$$ $f_{\mathbf{z}}(\mathcal{S})=f_{\mathbf{z}}(\mathcal{S^{\prime}})+\frac{(\mathbf{z}^{\top}\mathbf{V}_{\mathcal{S^{\prime}}}^{-1}\mathbf{x})^{2}}{1+\mathbf{x}^{\top}\mathbf{V}_{\mathcal{S^{\prime}}}^{-1}\mathbf{x}}~.$ (5) $$

This naturally yields a greedy algorithm. We start with an emty set $\mathcal{S}_{0}=\emptyset$ and at each step $i\in\{1,2,\ldots,k\}$, we select the
example that provides the maximum marginal gain over already selected examples $\mathcal{S}_{i-1}=\{\mathbf{x}_{1},\ldots,\mathbf{x}_{i-1}\}$:

$$ $\mathbf{x}_{i}=\operatorname*{argmax}_{\mathbf{x}\in\mathcal{X}\setminus\mathcal{S}_{i-1}}\frac{(\mathbf{z}^{\top}\mathbf{V}_{\mathcal{S}_{i-1}}^{-1}\mathbf{x})^{2}}{1+\mathbf{x}^{\top}\mathbf{V}_{\mathcal{S}_{i-1}}^{-1}\mathbf{x}}~.$ (6) $$

Although equation ([6](https://arxiv.org/html/2509.15676v1#Sx4.E6)) involves inverting matrices $\mathbf{V}_{i}$ of dimension $d$, the rank-one nature of their sequential updates $\mathbf{V}_{\mathcal{S}_{i}}=\mathbf{V}_{\mathcal{S}_{i-1}}+\mathbf{x}_{i}\mathbf{x}_{i}^{\top}$ helps compute their inverse efficiently, again using the Sherman-Morrison formula:

$$ $\mathbf{V}_{\mathcal{S}_{i}}^{-1}=\mathbf{V}_{\mathcal{S}_{i-1}}^{-1}-\frac{\mathbf{V}_{\mathcal{S}_{i-1}}^{-1}\mathbf{x}_{i}\mathbf{x}_{i}^{\top}\mathbf{V}_{\mathcal{S}_{i-1}}^{-1}}{1+\mathbf{x}_{i}^{\top}\mathbf{V}_{\mathcal{S}_{i-1}}^{-1}\mathbf{x}_{i}}~.$ (7) $$

This allows us to maintain the inverse matrix in $O(d^{2})$ time per iteration, rather than recomputing it from scratch.

The following adaptation of a result from Das and Kempe (2011) using our lower bound on the submodularity ratio (Lemma [3](https://arxiv.org/html/2509.15676v1#Thmtheorem3)) states the performance of the greedy algorithm under approximate submodularity.

###### Theorem 4 (Greedy Algorithm Performance with Approximate Submodularity) .

For any $\mathbf{z}\in\mathbb{R}^{d}$, let the greedy algorithm return a set $\mathcal{S}_{\text{greedy}}$ for the optimization problem in equation ([4](https://arxiv.org/html/2509.15676v1#Sx4.E4)). Then,
$\mathcal{S}_{\text{greedy}}$ satisfies

$$ $f_{\mathbf{z}}(\mathcal{S}_{\text{greedy}})\geq\left(1-e^{-\frac{1}{1+(k-1)\mu}}\right)\cdot f(\mathcal{S}^{\ast}),$ $$

where $S^{\ast}$ is an optimal solution of size at most $k$.

Thus, the approximation guarantee degrades gracefully with the submodularity ratio $\gamma_{k}(f_{\mathbf{z}})$. When the submodularity ratio $\gamma_{k}(f_{\mathbf{z}})$ is close to 1 (i.e., when $k=1$ or $\mu\approx 0$), greedy algorithms perform nearly as well as optimal algorithms (upto an $(1-1/e)$ factor).
In practice, the value of $\gamma$ can be estimated, providing useful certificates of near-optimality for solutions obtained via greedy selection.

###### Remark 1 .

For functions of the form $f_{\mathbf{z}}(\mathcal{S})=-\mathbf{z}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{z}$, where $\mathbf{V}_{\mathcal{S}}=\sum_{i\in\mathcal{S}}\mathbf{x}_{i}\mathbf{x}_{i}^{\top}+\beta\mathbf{I}_{d}$, increasing the regularization strength $\beta$ tends to make $f_{\mathbf{z}}(\mathcal{S})$ nearly submodular—that is, $\gamma_{k}(f)\to 1$. This shift further enhances the performance of the greedy algorithm in practice.

###### Theorem 4 (Greedy Algorithm Performance with Approximate Submodularity) .

###### Remark 1 .

#### Selecting Diverse Examples

Equation ([6](https://arxiv.org/html/2509.15676v1#Sx4.E6)) seeks the example $\mathbf{x}_{i}$ which is most relevant to the input query $\mathbf{z}$ in the geometry induced by the matrix $\mathbf{V}_{\mathcal{S}_{i-1}}^{-1}$. However, for in-context learning, both relevance (i.e., choosing examples similar to the input) and diversity
(i.e., choosing similar examples) are essential, and need to be carefully balanced.

To select diverse examples, we seek the set $\mathcal{S}$ that fetches the maximum information about the unknown parameter vector $\theta\in\mathbb{R}^{d}$ from noisy linear responses of the form ([1](https://arxiv.org/html/2509.15676v1#Sx3.E1)). To this end, let $\mathcal{S}\subset\mathcal{X}$ be a subset of size $k$. For a parameter vector $\mathbf{\theta}\sim\mathcal{N}(\mathbf{0}_{d},\beta\mathbf{I}_{d})$, a design matrix $\mathbf{X}_{\mathcal{S}}\in\mathbb{R}^{k\times d}$ formed from the elements of $\mathcal{S}$, and response vector $\mathbf{y}_{\mathcal{S}}\in\mathbb{R}^{k}$ formed from responses $y=\mathbf{x}^{\top}\mathbf{\theta}+\eta$, $\mathbf{x}\in\mathcal{S}$, with $\eta\sim\mathcal{N}(0,\beta)$, the information gain about $\mathbf{\theta}$ from $\mathbf{y}_{\mathcal{S}}$ is given by the mutual information between $\mathbf{\theta}$ and $\mathbf{y}_{\mathcal{S}}$:

$$ $\displaystyle I(\mathbf{\theta};\mathbf{y}_{\mathcal{S}}):=H(\mathbf{y}_{\mathcal{S}})-H(\mathbf{y}_{\mathcal{S}}|\mathbf{\theta})~,$ $$

where $H(\cdot)$ denotes the Shannon entropy (Cover 1999).
Since $\mathbf{y}_{\mathcal{S}}|\mathbf{\theta}\sim\mathcal{N}(\mathbf{0}_{k},\beta\mathbf{I}_{k})$ and $\mathbf{y}_{\mathcal{S}}\sim\mathcal{N}(\mathbf{0}_{k},\mathbf{X}_{\mathcal{S}}^{\top}\mathbf{X}_{\mathcal{S}}+\beta\mathbf{I}_{k})$, the information gain simplifies to

$$ $I(\mathbf{\theta};\mathbf{y}_{\mathcal{S}})=\frac{1}{2}\log\det\Big(\mathbf{I}_{d}+\frac{1}{\beta}\mathbf{X}_{\mathcal{S}}^{\top}\mathbf{X}_{\mathcal{S}}\Big)~,$ $$

which follows from the fact that the entropy of a Gaussian distribution with covariance matrix $\Sigma$ is $\frac{1}{2}\log\det(2\pi e\Sigma)$. Therefore, to find the most informative (or, equivalently, diverse) set, we find the set with the maximum information gain. This, along with the fact that $\mathbf{V}_{\mathcal{S}}=\mathbf{X}_{\mathcal{S}}^{\top}\mathbf{X}_{\mathcal{S}}+\beta\mathbf{I}_{d}$, leads to the following optimization problem:

$$ $\displaystyle\max_{\mathcal{S}\subset\mathcal{X},\,|\mathcal{S}|\leq k}g(\mathcal{S})~,\,\,\text{where}\,\,g(\mathcal{S})=\log\det(\mathbf{V}_{\mathcal{S}})~.$ (8) $$

The function $g$ is known as the D-optimal design in the experimental design literature (Pukelsheim 2006).

###### Lemma 5 (Submodularity of $g$ ) .

For any $\beta>0$, the function $g(\mathcal{S})=\log\det(\mathbf{V}_{\mathcal{S}})$ is monotone and submodular.

The result holds
from the matrix-determinant lemma, which, for any $\mathbf{x}\notin\mathcal{S}$, yields

$$ $\log\det(\mathbf{V}_{\mathcal{S}}+\mathbf{x}\mathbf{x}^{\top})-\log\det(\mathbf{V}_{\mathcal{S}})=\log(1+\mathbf{x}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x})~.$ $$

Monotonicity follows since this increment $\log(1+\mathbf{x}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x})$ is positive. Submodularity holds since this increment becomes smaller as $\mathbf{V}_{\mathcal{S}}$ grows (i.e., as the set $S$ grows), due to redundancy in the directions already spanned. This also admits a natural greedy selection rule:

$$ $\mathbf{x}_{i}=\operatorname*{argmax}_{\mathbf{x}\in\mathcal{X}\setminus\mathcal{S}_{i-1}}\log(1+\mathbf{x}^{\top}\mathbf{V}_{\mathcal{S}_{i-1}}^{-1}\mathbf{x})~,$ (9) $$

which can also be computed efficiently using equation ([7](https://arxiv.org/html/2509.15676v1#Sx4.E7)).

###### Lemma 5 (Submodularity of g g ) .

#### Combined Relevance and Diversity

To select relevant as well as diverse examples, we maximize the combined objective
$f_{\mathbf{z}}(\mathcal{S})+g(\mathcal{S})$, which maintains both monotonicity and approximate submodularity. Therefore, we combine the corresponding greedy selection rules ([6](https://arxiv.org/html/2509.15676v1#Sx4.E6)) and ([9](https://arxiv.org/html/2509.15676v1#Sx4.E9)) to obtain

$$ $\displaystyle\mathbf{x}_{i}=\operatorname*{argmax}_{\mathbf{x}\in\mathcal{X}\setminus\mathcal{S}_{i-1}}\frac{(\mathbf{z}^{\top}\mathbf{V}_{\mathcal{S}_{i-1}}^{-1}\mathbf{x})^{2}}{1+\mathbf{x}^{\top}\mathbf{V}_{\mathcal{S}_{i-1}}^{-1}\mathbf{x}}+\lambda\log(1+\mathbf{x}^{\top}\mathbf{V}_{\mathcal{S}_{i-1}}^{-1}\mathbf{x})~,$ $$

where $\lambda\geq 0$ controls the trade-off between relevance and diversity. The set $\mathcal{S}$ returned by this greedy algorithm maintains the approximation guarantee of Theorem [4](https://arxiv.org/html/2509.15676v1#Thmtheorem4). We call this algorithm Linear Information Theoretic Exemplars (Lite) for ICL and present the pseudo-code in Algorithm [1](https://arxiv.org/html/2509.15676v1#alg1).

Figure: Algorithm 1 Lite

### Moving Beyond Linearity: Kernel Trick

In this section, we generalize our approach to the practical setting of non-linear models by assuming that conditioned on a test query $\mathbf{z}\in\mathbb{R}^{d}$, the response $y$ for any feature vector $\mathbf{x}\in\mathbb{R}^{d}$ is generated as $y=h(\mathbf{x})+\eta$, where $\eta\sim\mathcal{N}(0,1)$ and $h$ is an unknown element of a reproducing kernel Hilbert space (RKHS, see Schölkopf and Smola (2002)). An RKHS, denoted by $\mathcal{H}_{k}$, is completely characterized by a symmetric and positive semi-definite kernel $k:\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}$ and vice-versa. Two commonly used kernels are the polynomial kernel $k_{\text{Poly}}(\mathbf{x},\mathbf{x}^{\prime})=(\mathbf{x}^{\top}\mathbf{x}^{\prime}+c)^{m}$ and the Gaussian kernel
$k_{\text{Gauss}}(\mathbf{x},\mathbf{x}^{\prime})=\exp\left(-\frac{\|\mathbf{x}-\mathbf{x}^{\prime}\|^{2}}{2\sigma^{2}}\right)$, where $c,m,\sigma$ are hyperparameters of the kernels. The kernel trick says that there exists a feature map $\phi:\mathbb{R}^{d}\to\mathcal{H}_{k}$ such that $k(\mathbf{x},\mathbf{x}^{\prime})=\phi(\mathbf{x})^{\top}\phi(\mathbf{x}^{\prime})$, where the inner product is associated with $\mathcal{H}_{k}$. This trick helps us to perform all computations in the high (possibly infinite) dimensional Hilbert space $\mathcal{H}_{k}$ analogously to those in the Euclidean space, but without the need for computing the inner product explicitly.

To see this, the greedy selection rule in the Euclidean space $\mathbb{R}^{d}$ can be lifted to the RKHS as

$$ $\displaystyle\mathbf{x}_{i}=\operatorname*{argmax}_{\mathbf{x}\in\mathcal{X}\setminus\mathcal{S}_{i-1}}$ $\displaystyle\frac{\Big(\phi(\mathbf{z})^{\top}\mathbf{V}_{\phi,\mathcal{S}_{i-1}}^{-1}\phi(\mathbf{x})\Big)^{2}}{1+\phi(\mathbf{x})^{\top}\mathbf{V}_{\phi,\mathcal{S}_{i-1}}^{-1}\phi(\mathbf{x})}$ $\displaystyle+\lambda\log\big(1+\phi(\mathbf{x})^{\top}\mathbf{V}_{\phi,\mathcal{S}_{i-1}}^{-1}\phi(\mathbf{x})\big)~,$ (10) $$

where $\mathbf{V}_{\mathcal{\phi,S}}=\beta\mathbf{I}+\sum_{i\in\mathcal{S}}\phi(\mathbf{x}_{i})\phi(\mathbf{x}_{i})^{\top}=\beta\mathbf{I}+\Phi_{\mathcal{S}}^{\top}\Phi_{\mathcal{S}}$ is the design covariance matrix in $\mathcal{H}_{k}$. Forming a design matrix $\Phi_{\mathcal{S}}$ from $\phi(\mathbf{x})$, $\mathbf{x}\in\mathcal{S}$ and using kernel trick, we obtain

$$ $\displaystyle\phi(\mathbf{x})^{\top}\mathbf{V}_{\phi,\mathcal{S}}^{-1}\phi(\mathbf{x})=\phi(\mathbf{x})^{\top}(\Phi_{\mathcal{S}}^{\top}\Phi_{\mathcal{S}}+\beta\mathbf{I})^{-1}\phi(\mathbf{x})$ $\displaystyle=$ $\displaystyle\frac{1}{\beta}\Big(\phi(\mathbf{x})^{\top}\phi(\mathbf{x})-\phi(\mathbf{x})^{\top}\Phi_{\mathcal{S}}^{\top}(\Phi_{\mathcal{S}}\Phi_{\mathcal{S}}^{\top}+\beta\mathbf{I})^{-1}\Phi_{\mathcal{S}}\phi(\mathbf{x})\Big)$ $\displaystyle=$ $\displaystyle\frac{1}{\beta}\Big(k(\mathbf{x},\mathbf{x})-\mathbf{k}_{\mathcal{S}}(\mathbf{x})^{\top}(\mathbf{K}_{\mathcal{S}}+\beta\mathbf{I})^{-1}\mathbf{k}_{\mathcal{S}}(\mathbf{x})\Big)~,$ (11) $$

where $\mathbf{K}_{\mathcal{S}}\!=\![k(\mathbf{x}_{i},\mathbf{x}_{j})]_{\mathbf{x}_{i},\mathbf{x}_{j}\in\mathcal{S}}$ is the gram matrix of size $|\mathcal{S}|\times|\mathcal{S}|$ and $\mathbf{k}_{\mathcal{S}}(x)\!=\!\left[k(\mathbf{x},\mathbf{x}_{i})\right]_{\mathbf{x}_{i}\in\mathcal{S}}$ is a vector of size $|\mathcal{S}|$.

Equation ([Moving Beyond Linearity: Kernel Trick](https://arxiv.org/html/2509.15676v1#Sx4.Ex12)) helps the selection rule ([Moving Beyond Linearity: Kernel Trick](https://arxiv.org/html/2509.15676v1#Sx4.Ex11)) to be implemented with only the kernel computations, avoiding matrix inverses and matrix-vector multiplications in (possibly) infinite-dimensional RKHS. This leads to our framework for in-context example selection, summarized below:

Note that Kite is a strict generalization of Lite (Algorithm [1](https://arxiv.org/html/2509.15676v1#alg1)) and captures it as a special case when the kernel is linear, i.e., when $k_{\text{lin}}(\mathbf{x},\mathbf{x}^{\prime})=\mathbf{x}^{\top}\mathbf{x}^{\prime}$. In our experiments, we keep the kernel as a tunable hyperparameter of Kite.

## Experiments

We present a comprehensive empirical analysis of our proposed approach, evaluating its performance across a variety of open-source datasets and models.

**Table 1: Evaluation Results. We compare classification accuracy (%) of KITE against retrieval baselines on five few‐shot benchmarks. Model abbreviations: GN (GPT-Neo 2.7B), QW (Qwen 2.5–1.5B), LL (Llama 3.2–3B).**
|  | SST-5 | CMSQA | MRPC | QNLI | HellaSwag |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Method | GN | QW | LL | GN | QW | LL | GN | QW | LL | GN | QW | LL | GN | QW | LL |
| Random | 32.31 | 40.53 | 42.41 | 42.25 | 65.90 | 67.22 | 66.38 | 70.45 | 71.12 | 57.56 | 70.21 | 68.43 | 41.75 | 66.30 | 65.21 |
| BM25 | 36.05 | 47.86 | 47.85 | 42.91 | 70.02 | 70.92 | 66.96 | 73.77 | 70.58 | 62.12 | 71.11 | 69.38 | 41.34 | 67.85 | 66.32 |
| Dense | 35.96 | 46.64 | 44.14 | 43.98 | 70.18 | 72.23 | 65.83 | 73.92 | 71.81 | 64.42 | 70.73 | 70.04 | 40.44 | 68.80 | 67.88 |
| DPP | 36.84 | 47.88 | 46.45 | 37.53 | 70.05 | 71.65 | 68.04 | 74.25 | 70.63 | 63.62 | 70.53 | 70.08 | 40.56 | 69.32 | 68.55 |
| Kite | 40.60 | 49.59 | 47.59 | 46.60 | 71.25 | 72.89 | 68.38 | 75.27 | 71.98 | 65.32 | 71.05 | 71.68 | 41.68 | 71.02 | 70.12 |

**Table 2: Analysis on Submodularity ratio $\gamma_{\min}$. Left: We vary $k$ while fixing $\beta=1$; Right: We vary $\beta$ while fixing $k=20$.**
|  | Vary $k$ (fixed $\beta=1$) | Vary $\beta$ (fixed $k=20$) |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Dataset | 5 | 10 | 15 | 20 | 25 | 30 | 35 | 1 | 3 | 5 | 7 | 9 | 11 |
| SST–5 | 0.748 | 0.685 | 0.705 | 0.689 | 0.703 | 0.716 | 0.696 | 0.689 | 0.744 | 0.719 | 0.758 | 0.792 | 0.827 |
| CMSQA | 0.852 | 0.811 | 0.816 | 0.841 | 0.830 | 0.797 | 0.836 | 0.841 | 0.804 | 0.884 | 0.843 | 0.907 | 0.886 |
| MRPC | 0.876 | 0.840 | 0.793 | 0.816 | 0.850 | 0.887 | 0.848 | 0.816 | 0.765 | 0.803 | 0.824 | 0.824 | 0.898 |
| QNLI | 0.844 | 0.801 | 0.900 | 0.893 | 0.875 | 0.837 | 0.877 | 0.893 | 0.829 | 0.823 | 0.919 | 0.912 | 0.882 |
| HellaSwag | 0.395 | 0.601 | 0.525 | 0.596 | 0.839 | 0.908 | 0.806 | 0.596 | 0.803 | 0.767 | 0.806 | 0.935 | 0.852 |

Implementation Details. We conduct our evaluations using three open-source, state-of-the-art language models: GPT-Neo-2.7B (Black et al. 2021), Qwen 2.5-1.5B (Hui et al. 2024), and Llama-3.2-3B (Grattafiori et al. 2024). For all baselines, we fix the number of in-context examples to $50$ during inference, truncating as necessary based on the model’s maximum context length. We represent textual inputs using embeddings obtained from bert-base-uncased (Devlin et al. 2019). We experiment with three kernels: (i) Linear kernel $k_{\text{lin}}(\mathbf{x},\mathbf{x}^{\prime})$ (which reduces to Lite; see Algorithm [1](https://arxiv.org/html/2509.15676v1#alg1)), (ii) Polynomial kernel $k_{\text{poly}}(\mathbf{x},\mathbf{x}^{\prime})$ with varying degree $m$ and (iii) Gaussian kernel $k_{\text{Gauss}}(\mathbf{x},\mathbf{x}^{\prime})$ (length-scale $\sigma=1.0$). We use a regularization parameter of $\beta=0.02$ and a diversity parameter of $\lambda=0.5$. For Kite in Table [1](https://arxiv.org/html/2509.15676v1#Sx5.T1), we report the best result across the three kernel choices. A detailed ablation study on the impact of the kernel choice is presented in Table [3](https://arxiv.org/html/2509.15676v1#Sx5.T3). We defer the mathematical formulation for each kernel to Appendix.

Datasets. To empirically demonstrate the efficacy of Kite, we evaluate on five few-shot classification datasets: SST-5 (Socher et al. 2013) (single-sentence sentiment), CMSQA (Talmor et al. 2019) (question answering), MRPC (Dolan, Quirk, and Brockett 2004) (sentence-pair paraphrase), QNLI (Wang et al. 2018) (binary entailment classification), and HellaSwag (Zellers et al. 2019) (common sense NLI).

For each task, we use the corresponding validation split for evaluation and employ the deduplicated training split as the candidate exemplar set.

Evaluation Metrics.
Following prior work (Brown et al. 2020a), we formulate all classification tasks as multiple-choice problems. we construct input prompts by concatenating the given context with each candidate label (i.e., “context + label”). We then compute the conditional log-likelihood of generating the label tokens given this combined input prompt, and select the candidate with the highest log-likelihood as the model’s prediction. Model performance is evaluated using accuracy, defined as the fraction of correctly predicted examples on the validation split.

Baselines. For a fair and comprehensive evaluation, we compare Kite against several baselines, including Random, BM25 (Robertson, Zaragoza et al. 2009), top-$k$ with Dense embddings (Dense) (Liu et al. 2021), and DPP-based retrieval strategies (Ye et al. 2023a).

Results. We report our evaluation results in Table [1](https://arxiv.org/html/2509.15676v1#Sx5.T1). We find that Kite consistently outperforms all baselines across most datasets and model architectures. On GPT-Neo-2.7B, for example, Kite delivers an accuracy improvement of +$4.55\%$ on SST-5 (Socher et al. 2013) and +$3.69\%$ on CMSQA (Talmor et al. 2019) over BM25 (Robertson, Zaragoza et al. 2009). Further, this trend is consistently observed across all evaluated models. With Qwen-2.5-1.5B, Kite surpasses the strongest baseline, DPP (Ye et al. 2023a), on four out of five datasets, achieving notable gains of +$1.71\%$ on SST-5 and +$1.70\%$ on HellaSwag. Similarly, on Llama-3B, Kite achieves the highest accuracy on four datasets, including a substantial +$2.24\%$ improvement on HellaSwag over the state-of-the-art (Ye et al. 2023a). Across all $15$ evaluation settings (five datasets and three models), Kite attains the highest accuracy in $13$ cases, underscoring its robustness and efficacy as an exemplar retrieval framework.

Ablation study on kernel function.
The choice of kernel is a critical hyperparameter in our Kite framework, as it defines the geometry of the feature space in which relevance and diversity are measured. To clearly understand its impact, we conducted an ablation study comparing the performance of three distinct kernels: Linear (which reduces Kite to the Lite algorithm), Polynomial (with degree $m=3$), and Gaussian RBF ($\sigma=1.0$). The results, detailed in Table [3](https://arxiv.org/html/2509.15676v1#Sx5.T3), reveal that no single kernel is universally optimal; the best choice is contingent on the specific dataset and its underlying data distribution.

The strong performance of the non-linear kernels across most datasets validates our core motivation for moving beyond the linear model of Lite. It underscores the importance of the kernel trick in empowering the model to capture richer, non-linear relationships between exemplars, which is essential for effective in-context learning. The main results reported in Table 1 represent the best performance achieved across these kernels for each task, highlighting the consistent advantage conferred by a well-tuned kernelized approach.

**Table 3: Ablation study on the kernel function. We report accuracy for Kite using different kernels on Qwen-1.5B.**
| Dataset | Linear | Polynomial | Gaussian RBF |
| --- | --- | --- | --- |
| SST-5 | 47.95 | 48.38 | 49.59 |
| CMSQA | 69.86 | 71.06 | 71.25 |
| MRPC | 75.27 | 71.22 | 67.15 |
| QNLI | 70.46 | 70.38 | 71.05 |
| HellaSwag | 67.18 | 71.02 | 70.67 |

Empirical validation of submodularity.
We empirically validate the submodularity of our set function, $f_{z}(\mathcal{S})=-\,\mathbf{z}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{z}$ by estimating its submodularity ratio (Def. [2](https://arxiv.org/html/2509.15676v1#Thmtheorem2)) on real-world text embeddings. For this analysis, we used the first 500 examples from each dataset to generate embeddings. We created two distinct types: demonstration embeddings, by concatenating an input with its gold label, and query embeddings, from the input alone.

To estimate the ratio, we ran a Monte Carlo simulation across a grid of hyperparameters for subset size $k$ and regularization $\beta$. In each trial, we sampled a disjoint triplet $(\mathcal{S},\mathcal{L},\mathbf{z})$, where $\mathcal{S}$ is a random set of demonstration embeddings, $\mathcal{L}$ is a diverse set of up to $k$ demonstration embeddings selected via farthest-point sampling, and $\mathbf{z}$ is a query embedding. We then computed the ratio and recorded the minimum value observed ($\gamma_{\min}$) for each configuration.

As shown in Table [2](https://arxiv.org/html/2509.15676v1#Sx5.T2), our experiments reveal a consistently high $\gamma_{\min}$ across all datasets and settings. These values, typically at or above 0.8, provide strong empirical evidence that our objective function is approximately submodular in practice. This finding justifies our use of a greedy algorithm, as it is expected to yield near-optimal results.

## Conclusion

In this paper, we introduce Kite, a principled, information-theoretic framework that treats in-context exemplar selection as a query-specific optimization problem. By leveraging an approximately submodular objective and the kernel trick, our algorithm efficiently selects exemplars that balance relevance and diversity, an approach supported by theoretical performance guarantees. Extensive experiments demonstrate that Kite consistently outperforms strong retrieval baselines across multiple classification datasets and language models, validating its efficacy and effectiveness for in-context example selection. Extending Kite to generative tasks, which involve producing multiple dependent output tokens, is an important direction for future work.

## References

- Adiwardana et al. (2020)
Adiwardana, D.; Luong, M.-T.; So, D. R.; Hall, J.; Fiedel, N.; Thoppilan, R.; Yang, Z.; Kulshreshtha, A.; Nemade, G.; Lu, Y.; et al. 2020.
Towards a human-like open-domain chatbot.
*arXiv preprint arXiv:2001.09977*.
- Black et al. (2021)
Black, S.; Gao, L.; Wang, P.; Leahy, C.; and Biderman, S. 2021.
GPT-Neo: Large Scale Autoregressive Language Modeling with Mesh-Tensorflow.
- Brown et al. (2020a)
Brown, T. B.; Mann, B.; Ryder, N.; Subbiah, M.; Kaplan, J.; Dhariwal, P.; Neelakantan, A.; Shyam, P.; Sastry, G.; and Amanda Askell, e. a. 2020a.
Language Models are Few-Shot Learners.
arXiv:2005.14165.
- Brown et al. (2020b)
Brown, T. B.; Mann, B.; Ryder, N.; Subbiah, M.; Kaplan, J.; Dhariwal, P.; Neelakantan, A.; Shyam, P.; Sastry, G.; Askell, A.; and et al. 2020b.
Language Models are Few-Shot Learners.
In *Advances in Neural Information Processing Systems 33*, 1877–1901.
- Chen et al. (2022)
Chen, M.; Du, J.; Pasunuru, R.; Mihaylov, T.; Iyer, S.; Stoyanov, V.; and Kozareva, Z. 2022.
Improving in-context few-shot learning via self-supervised training.
*arXiv preprint arXiv:2205.01703*.
- Cover (1999)
Cover, T. M. 1999.
*Elements of information theory*.
John Wiley & Sons.
- Dai et al. (2022)
Dai, D.; Sun, Y.; Dong, L.; Hao, Y.; Ma, S.; Sui, Z.; and Wei, F. 2022.
Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers.
*arXiv preprint arXiv:2212.10559*.
- Das and Kempe (2011)
Das, A.; and Kempe, D. 2011.
Submodular meets spectral: Greedy algorithms for subset selection, sparse approximation and dictionary selection.
*arXiv preprint arXiv:1102.3975*.
- Devlin et al. (2019)
Devlin, J.; Chang, M.-W.; Lee, K.; and Toutanova, K. 2019.
BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.
arXiv:1810.04805.
- Dolan, Quirk, and Brockett (2004)
Dolan, B.; Quirk, C.; and Brockett, C. 2004.
Unsupervised Construction of Large Paraphrase Corpora: Exploiting Massively Parallel News Sources.
In *COLING 2004: Proceedings of the 20th International Conference on Computational Linguistics*, 350–356. Geneva, Switzerland: COLING.
- Dong et al. (2022)
Dong, Q.; Li, L.; Dai, D.; Zheng, C.; Ma, J.; Li, R.; Xia, H.; Xu, J.; Wu, Z.; Liu, T.; et al. 2022.
A survey on in-context learning.
*arXiv preprint arXiv:2301.00234*.
- Gehman et al. (2020)
Gehman, S.; Gururangan, S.; Sap, M.; Choi, Y.; and Smith, N. A. 2020.
Realtoxicityprompts: Evaluating neural toxic degeneration in language models.
*arXiv preprint arXiv:2009.11462*.
- Ghosal et al. (2025a)
Ghosal, S. S.; Pal, S.; Mukherjee, K.; and Manocha, D. 2025a.
PromptRefine: Enhancing Few-Shot Performance on Low-Resource Indic Languages with Example Selection from related Example Banks.
In *Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers)*, 351–365.
- Ghosal et al. (2025b)
Ghosal, S. S.; Singh, V.; Ghosh, A.; Pal, S.; Baidya, S.; Saha, S.; and Manocha, D. 2025b.
Relic: Enhancing Reward Model Generalization for Low-Resource Indic Languages with Few-Shot Examples.
*arXiv preprint arXiv:2506.16502*.
- Gonen, Hosking, and Augenstein (2023)
Gonen, H.; Hosking, J.; and Augenstein, I. 2023.
Perplexity-Based Prompt Selection for Large Language Models.
In *Proceedings of the 2023 Conference of the Association for Computational Linguistics*.
- Grattafiori et al. (2024)
Grattafiori, A.; Dubey, A.; Jauhri, A.; Pandey, A.; Kadian, A.; Al-Dahle, A.; Letman, A.; Mathur, A.; Schelten, A.; and Alex Vaughan, e. a. 2024.
The Llama 3 Herd of Models.
arXiv:2407.21783.
- Hui et al. (2024)
Hui, B.; Yang, J.; Cui, Z.; Yang, J.; Liu, D.; Zhang, L.; Liu, T.; Zhang, J.; Yu, B.; Lu, K.; Dang, K.; Fan, Y.; Zhang, Y.; Yang, A.; Men, R.; Huang, F.; Zheng, B.; Miao, Y.; Quan, S.; Feng, Y.; Ren, X.; Ren, X.; Zhou, J.; and Lin, J. 2024.
Qwen2.5-Coder Technical Report.
arXiv:2409.12186.
- Köppen (2000)
Köppen, M. 2000.
The curse of dimensionality.
In *5th online world conference on soft computing in industrial applications (WSC5)*, volume 1, 4–8.
- Larson et al. (2019)
Larson, S.; Mahendran, A.; Peper, J. J.; Clarke, C.; Lee, A.; Hill, P.; Kummerfeld, J. K.; Leach, K.; Laurenzano, M. A.; Tang, L.; et al. 2019.
An evaluation dataset for intent classification and out-of-scope prediction.
*arXiv preprint arXiv:1909.02027*.
- Lattimore and Szepesvári (2020)
Lattimore, T.; and Szepesvári, C. 2020.
*Bandit algorithms*.
Cambridge University Press.
- Li and Qiu (2023)
Li, L.; and Qiu, M. 2023.
In-Context Learning Demonstration Selection via Influence Analysis.
arXiv:2402.11750.
- Li et al. (2023a)
Li, X.; Lv, K.; Yan, H.; Lin, T.; Zhu, W.; Ni, Y.; Xie, G.; Wang, X.; and Qiu, X. 2023a.
Unified demonstration retriever for in-context learning.
*arXiv preprint arXiv:2305.04320*.
- Li et al. (2023b)
Li, Y.; Ildiz, M. E.; Papailiopoulos, D.; and Oymak, S. 2023b.
Transformers as algorithms: Generalization and stability in in-context learning.
In *International Conference on Machine Learning*, 19565–19594. PMLR.
- Liu et al. (2021)
Liu, J.; Shen, D.; Zhang, Y.; Dolan, B.; Carin, L.; and Chen, W. 2021.
What Makes Good In-Context Examples for GPT-$3$?
*arXiv preprint arXiv:2101.06804*.
- Liu et al. (2022)
Liu, J.; Shen, D.; Zhang, Y.; Dolan, B.; Carin, L.; and Chen, W. 2022.
What Makes Good In-Context Examples for GPT-3?
In *Proceedings of the 3rd Workshop on Deep Learning for Low-Resource Natural Language Processing (DeepLo) at ACL*, 100–114.
- Luo et al. (2023)
Luo, M.; Xu, X.; Dai, Z.; Pasupat, P.; Kazemi, M.; Baral, C.; Imbrasaite, V.; and Zhao, V. Y. 2023.
Dr. icl: Demonstration-retrieved in-context learning.
*arXiv preprint arXiv:2305.14128*.
- Luo et al. (2024)
Luo, M.; Xu, X.; Liu, Y.; Pasupat, P.; and Kazemi, M. 2024.
In-context learning with retrieved demonstrations for language models: A survey.
*arXiv preprint arXiv:2401.11624*.
- Min et al. (2021)
Min, S.; Lewis, M.; Zettlemoyer, L.; and Hajishirzi, H. 2021.
Metaicl: Learning to learn in context.
*arXiv preprint arXiv:2110.15943*.
- Min et al. (2022)
Min, S.; Lyu, X.; Holtzman, A.; Artetxe, M.; Lewis, M.; Hajishirzi, H.; and Zettlemoyer, L. 2022.
Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?
In *Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing*, 11048–11064.
- Olsson et al. (2022)
Olsson, C.; Elhage, N.; Nanda, N.; Joseph, N.; DasSarma, N.; Henighan, T.; Mann, B.; Askell, A.; Bai, Y.; Chen, A.; et al. 2022.
In-context learning and induction heads.
*arXiv preprint arXiv:2209.11895*.
- Pan (2023)
Pan, J. 2023.
*What in-context learning “learns” in-context: Disentangling task recognition and task learning*.
Master’s thesis, Princeton University.
- Pukelsheim (2006)
Pukelsheim, F. 2006.
*Optimal design of experiments*.
SIAM.
- Robertson, Zaragoza et al. (2009)
Robertson, S.; Zaragoza, H.; et al. 2009.
The probabilistic relevance framework: BM25 and beyond.
*Foundations and Trends® in Information Retrieval*, 3(4): 333–389.
- Rubin, Herzig, and Berant (2022)
Rubin, O.; Herzig, J.; and Berant, J. 2022.
Learning to Retrieve Prompts for In-Context Learning.
In *Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics*, 2655–2671.
- Saunshi, Malladi, and Arora (2020)
Saunshi, N.; Malladi, S.; and Arora, S. 2020.
A mathematical exploration of why language models help solve downstream tasks.
*arXiv preprint arXiv:2010.03648*.
- Schölkopf and Smola (2002)
Schölkopf, B.; and Smola, A. J. 2002.
*Learning with kernels: support vector machines, regularization, optimization, and beyond*.
MIT press.
- Shin et al. (2022)
Shin, S.; Lee, S.-W.; Ahn, H.; Kim, S.; Kim, H.; Kim, B.; Cho, K.; Lee, G.; Park, W.; Ha, J.-W.; et al. 2022.
On the effect of pretraining corpora on in-context learning by a large-scale language model.
*arXiv preprint arXiv:2204.13509*.
- Socher et al. (2013)
Socher, R.; Perelygin, A.; Wu, J.; Chuang, J.; Manning, C. D.; Ng, A.; and Potts, C. 2013.
Recursive Deep Models for Semantic Compositionality Over a Sentiment Treebank.
In Yarowsky, D.; Baldwin, T.; Korhonen, A.; Livescu, K.; and Bethard, S., eds., *Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing*, 1631–1642. Seattle, Washington, USA: Association for Computational Linguistics.
- Sorensen, Ding, and Chang (2022)
Sorensen, M.; Ding, N.; and Chang, M.-W. 2022.
Information-Theoretic Demonstration Selection for Few-Shot Prompting.
In *Proceedings of the Findings of EMNLP*.
- Talmor et al. (2019)
Talmor, A.; Herzig, J.; Lourie, N.; and Berant, J. 2019.
CommonsenseQA: A Question Answering Challenge Targeting Commonsense Knowledge.
In Burstein, J.; Doran, C.; and Solorio, T., eds., *Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)*, 4149–4158. Minneapolis, Minnesota: Association for Computational Linguistics.
- Wang et al. (2019)
Wang, A.; Pruksachatkun, Y.; Nangia, N.; Singh, A.; Michael, J.; Hill, F.; Levy, O.; and Bowman, S. 2019.
Superglue: A stickier benchmark for general-purpose language understanding systems.
*Advances in neural information processing systems*, 32.
- Wang et al. (2018)
Wang, A.; Singh, A.; Michael, J.; Hill, F.; Levy, O.; and Bowman, S. 2018.
GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding.
In Linzen, T.; Chrupała, G.; and Alishahi, A., eds., *Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP*, 353–355. Brussels, Belgium: Association for Computational Linguistics.
- Wang, Dou, and Che (2022)
Wang, D.; Dou, L.; and Che, W. 2022.
A survey on table-and-text hybridqa: Concepts, methods, challenges and future directions.
*arXiv preprint arXiv:2212.13465*.
- Wang et al. (2025)
Wang, L.; Chen, S.; Jiang, L.; Pan, S.; Cai, R.; Yang, S.; and Yang, F. 2025.
Parameter-efficient fine-tuning in large language models: a survey of methodologies.
*Artificial Intelligence Review*, 58(8): 227.
- Wang et al. (2023)
Wang, X.; Zhu, W.; Saxon, M.; Steyvers, M.; and Wang, W. Y. 2023.
Large Language Models are Latent Variable Models: Explaining and Finding Good Demonstrations for In-Context Learning.
In *Advances in Neural Information Processing Systems 36*.
- Wei et al. (2023a)
Wei, J.; Hou, L.; Lampinen, A.; Chen, X.; Huang, D.; Tay, Y.; Chen, X.; Lu, Y.; Zhou, D.; Ma, T.; et al. 2023a.
Symbol tuning improves in-context learning in language models.
*arXiv preprint arXiv:2305.08298*.
- Wei et al. (2022)
Wei, J.; Tay, Y.; Bommasani, R.; Raffel, C.; Zoph, B.; Borgeaud, S.; Yogatama, D.; Bosma, M.; Zhou, D.; Metzler, D.; et al. 2022.
Emergent abilities of large language models.
*arXiv preprint arXiv:2206.07682*.
- Wei et al. (2023b)
Wei, J.; Wei, J.; Tay, Y.; Tran, D.; Webson, A.; Lu, Y.; Chen, X.; Liu, H.; Huang, D.; Zhou, D.; et al. 2023b.
Larger language models do in-context learning differently.
*arXiv preprint arXiv:2303.03846*.
- Williams, Nangia, and Bowman (2017)
Williams, A.; Nangia, N.; and Bowman, S. R. 2017.
A broad-coverage challenge corpus for sentence understanding through inference.
*arXiv preprint arXiv:1704.05426*.
- Wu et al. (2022)
Wu, Z.; Wang, Y.; Ye, J.; and Kong, L. 2022.
Self-adaptive in-context learning: An information compression perspective for in-context example selection and ordering.
*arXiv preprint arXiv:2212.10375*.
- Xie et al. (2021)
Xie, S. M.; Raghunathan, A.; Liang, P.; and Ma, T. 2021.
An explanation of in-context learning as implicit bayesian inference.
*arXiv preprint arXiv:2111.02080*.
- Yang et al. (2023)
Yang, Z.; Zhang, Y.; Sui, D.; Liu, C.; Zhao, J.; and Liu, K. 2023.
Representative Demonstration Selection for In-Context Learning with Two-Stage Determinantal Point Process.
In *Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing*, 5443–5456.
- Ye et al. (2023a)
Ye, J.; Wu, Z.; Feng, J.; Yu, T.; and Kong, L. 2023a.
Compositional Exemplars for In-Context Learning.
In *Proceedings of the 40th International Conference on Machine Learning*.
ICML.
- Ye et al. (2023b)
Ye, J.; Wu, Z.; Feng, J.; Yu, T.; and Kong, L. 2023b.
Compositional exemplars for in-context learning.
In *International Conference on Machine Learning*, 39818–39833. PMLR.
- Yoo et al. (2022)
Yoo, K. M.; Kim, J.; Kim, H. J.; Cho, H.; Jo, H.; Lee, S.-W.; Lee, S.-g.; and Kim, T. 2022.
Ground-truth labels matter: A deeper look into input-label demonstrations.
*arXiv preprint arXiv:2205.12685*.
- Zellers et al. (2019)
Zellers, R.; Holtzman, A.; Bisk, Y.; Farhadi, A.; and Choi, Y. 2019.
HellaSwag: Can a Machine Really Finish Your Sentence?
In Korhonen, A.; Traum, D.; and Màrquez, L., eds., *Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics*, 4791–4800. Florence, Italy: Association for Computational Linguistics.
- Zhang, Rudra, and Anand (2021)
Zhang, Z.; Rudra, K.; and Anand, A. 2021.
Explain and predict, and then predict again.
In *Proceedings of the 14th ACM international conference on web search and data mining*, 418–426.
- Zhao et al. (2021)
Zhao, Z.; Wallace, E.; Feng, S.; Klein, D.; and Singh, S. 2021.
Calibrate before use: Improving few-shot performance of language models.
In *International conference on machine learning*, 12697–12706. PMLR.
- Zheng et al. (2023)
Zheng, L.; Chiang, W.-L.; Sheng, Y.; Zhuang, S.; Wu, Z.; Zhuang, Y.; Lin, Z.; Li, Z.; Li, D.; Xing, E.; et al. 2023.
Judging llm-as-a-judge with mt-bench and chatbot arena.
*Advances in neural information processing systems*, 36: 46595–46623.

## Appendix

## Appendix A Extended Ablation Study

Figure: Figure 2: Ablation study on the MNLI dataset illustrating the effect of varying $\lambda$ values. We report accuracy for KITE using different kernel-based selection strategies on the Qwen-1.5B model. Since MNLI has a large example bank (over 400k examples), a purely relevance-based selection ($\lambda=0$) underperforms. Incorporating diversity through higher $\lambda$ values (e.g., $\lambda=1$) leads to notable improvements in accuracy by promoting a more diverse selection of in-context examples.
Refer to caption: images/mnli_lambda_accuracy_plot.png

#### Ablation study on λ \lambda .

In this section, we present an ablation study over $\lambda$ on the MNLI dataset (Williams, Nangia, and Bowman 2017). As shown in Figure 2, we evaluate the impact of the trade-off parameter $\lambda$, which balances relevance and diversity in our selection process. For the MNLI dataset, which has a very large and varied example bank (over 400k examples), relying solely on relevance-based selection ($\lambda$=0) results in suboptimal performance. This is because many examples can be semantically similar, leading to a redundant in-context set. By increasing $\lambda$, we introduce diversity into the selection, which proves crucial for this task. The results indicate that a balanced approach, with $\lambda$ set to approximately 1, achieves the highest accuracy. This demonstrates that for complex, large-scale datasets, a combination of both relevance and diversity is essential for selecting the most effective in-context examples.

#### Ablation study on number of ICE examples.

We also experimented with the number of in-context examples to understand how it affects model performance. In Figure 3, we report the accuracy of the Qwen-1.5B model on the SST-5 dataset as a function of the number of provided examples. As expected, performance for all methods generally improves as more examples are added to the context. However, our proposed KITE algorithm consistently and significantly outperforms all baseline methods, including BM25, Dense retrieval, and DPP. The advantage of KITE is particularly pronounced in low-resource settings (e.g., with 2 or 4 ICE), where it establishes a substantial performance gap. This highlights KITE’s superior ability to select high-quality, informative examples, making it highly effective even when the number of in-context examples is limited.

Figure: Figure 3: Accuracy comparison of the Qwen-1.5B model on the SST-5 dataset across varying numbers of in-context examples (ICE). The KITE algorithm, using kernel-based selection strategies, significantly outperforms baseline methods (BM25, Dense, DPP) in low-ICE settings.
Refer to caption: images/sst5_ice_accuracy_plot.png

## Appendix B Synthetic Experiments for Linear Model

### Implementation Details

#### Evaluation on synthetic setup.

We empirically verify that our proposed method KITE replicates its theoretical behavior under the stylized linear generative model. The goal of this setup is to ensure that the synthetic
controlled setting removes confounding factors such as representation drift or non‑Gaussian noise and isolates the effect of the subset–selection rule itself.

#### Data generation.

For each run, we sample a ground‑truth parameter $\boldsymbol{\theta}^{\star}\!\sim\!\mathcal{N}(\mathbf{0},\sigma^{2}I)$ and draw training features
$\mathbf{x}^{(i)}\!\sim\!\mathcal{N}(\mu_{\text{train}}\mathbf{1},\sigma^{2}I)$, where $\sigma\!=\!5.0$.
The responses are generated by the linear model $y^{(i)}=\langle\mathbf{x}^{(i)},\boldsymbol{\theta}^{\star}\rangle+\varepsilon^{(i)}$ with $\varepsilon^{(i)}\sim\mathcal{N}(0,1)$.
Test queries are drawn from the distribution $\mathbf{z}^{(i)}\!\sim\!\mathcal{N}((\mu_{\text{train}}+\mu_{\text{test}})\mathbf{1},\sigma^{2}I)$ but with a configurable mean shift $\mu_{\text{test}}\in[0,5]$ (default $=0$) to probe distribution shift. Recall that our theoretical guarantees hold for any arbitrary test query $\mathbf{z}$.

#### Subset selectors compared.

We evaluate three greedy scoring rules: (i) Dense retriever, (ii) DPP retriever, and (iii) Kite with linear kernel, e.g., Lite. We sweep $\lambda\!\in\!\{1,\dots,10\}$ (that controls the tradeoff between relevance and diversity), while the ridge regulariser in the downstream estimator remains fixed at $\beta=0.02$ throughout all trials.

#### Evaluation protocol.

Given a single test query $\mathbf{z}$, each selector chooses a subset $S$ of size $k$.
We fit the ridge estimator $\hat{\boldsymbol{\theta}}_{S}$ on $(\mathbf{X}_{S},\mathbf{y}_{S})$ and report the mean absolute prediction error on the test set (comprising of $N_{\text{test}}$ datapoints)
$\mathcal{L}=\tfrac{1}{N_{\text{test}}}\sum_{j\in[N_{\text{test}}]}\!\bigl|\langle\mathbf{z}^{(j)},\boldsymbol{\theta}^{\star}-\hat{\boldsymbol{\theta}}_{S}\rangle\bigr|$.
Every configuration is repeated $M$ times with new random seeds.
Unless stated otherwise we use dimension $d=5$, subset size $k=5$, and sweep (i) the training/test size $N\!\in\!\{500,\dots,5000\}$ and (ii) the mean shift $\mu_{\text{test}}$. Note that due to the Sherman–Morrison update (Alg. [1](https://arxiv.org/html/2509.15676v1#alg1), each greedy iteration costs $O(d^{2})$.

**Table 4: Synthetic linear model: the average of mean absolute test error across the $M$ distinct runs. Left block varies $N$ - the number of train/test datapoints; right block varies $\mu_{\text{test}}$ capturing the distribution shift. In both of these experiments, our method LITE of selecting datapoints for training the linear model (conditioned on the test query) performs significantly better than existing baselines for ICL (Dense via top-k selection; DPP based retrieval).**
| Vary $N$ | Vary $\mu_{\text{test}}$ |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| $\mathbf{N}$ | Dense | DPP | Lite | $\boldsymbol{\mu_{\text{test}}}$ | Dense | DPP | Lite |
| 1000 | 3.915 | 2.718 | 1.086 | 0.0 | 67.716 | 67.726 | 68.200 |
| 2000 | 3.753 | 2.875 | 0.748 | 1.0 | 68.692 | 68.714 | 67.834 |
| 3000 | 4.204 | 2.862 | 0.842 | 2.0 | 72.868 | 72.810 | 71.598 |
| 3500 | 4.210 | 2.744 | 0.765 | 3.0 | 76.636 | 76.677 | 76.094 |
| 4000 | 4.134 | 2.818 | 0.789 | 4.0 | 81.515 | 81.553 | 80.787 |
| 4500 | 4.085 | 2.917 | 0.833 | 4.5 | 83.929 | 84.005 | 82.284 |
| 5000 | 4.253 | 2.981 | 0.844 | 5.0 | 86.908 | 86.849 | 83.958 |

### Results

#### On the effect of training size N N .

As shown in Table [4](https://arxiv.org/html/2509.15676v1#A2.T4), enlarging the train/test pool from $N=500$ to $N=5000$ monotonically reduces error for every selector; nevertheless, the margin between methods remains almost constant. The error of Lite is the lowest throughout (mean $0.90$), outperforming DPP and Dense benchmarks.

#### On the effect of test‑distribution mean μ test \mu_{\text{test}} .

Table [4](https://arxiv.org/html/2509.15676v1#A2.T4) illustrates covariate‑shift robustness. As we move the test mean from 0 to $5$, errors rise across the board. Dense and DPP track each other almost exactly—their scores differ only by a small diversity addend—whereas Lite degrades more gently.

The results above empirically validate the good theoretical properties of our framework (in the linear setting) for selecting datapoints for training the linear model as compared to existing retriever methods.

## Appendix C Theoretical Proofs and Results

In this section, we present the missing proofs and results from the main paper.

### Proof of Lemma 3

For a set $\mathcal{S}\subseteq\mathcal{X}$ and an element $\mathbf{x}_{i}\notin\mathcal{S}$, define the marginal gain of adding $\mathbf{x}_{i}$ to $\mathcal{S}$ in the set function $f_{\mathbf{z}}(\mathcal{S})$ as

$$ $\displaystyle\Delta_{i}:=f_{\mathbf{z}}(\mathcal{S}\cup\{\mathbf{x}_{i}\})-f(\mathcal{S})=\frac{(\mathbf{z}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{i})^{2}}{1+\mathbf{x}_{i}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{i}},$ $$

where $f_{\mathbf{z}}(\mathcal{S})=-\mathbf{z}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{z}$ and $\mathbf{V}_{\mathcal{S}}=\sum_{i\in S}\mathbf{x}_{i}\mathbf{x}_{i}^{\top}+\lambda I$. Now consider a set $L=\{\mathbf{x}_{1},\dots,\mathbf{x}_{k}\}\subseteq\mathcal{X}\setminus S$. We compare

$$ $\sum_{\mathbf{x}_{i}\in\mathcal{L}}\Delta_{i},\qquad\text{vs.}\qquad f_{\mathbf{z}}(\mathcal{S}\cup\mathcal{L})-f_{\mathbf{z}}(\mathcal{S}).$ $$

The ratio of these two quantities defines the *submodularity ratio* $\gamma$, which quantifies how closely the function $f$ adheres to submodularity.

Let $\mathbf{X}_{\mathcal{L}}=[\mathbf{x}_{1},\mathbf{x}_{2},\dots,\mathbf{x}_{k}]\in\mathbb{R}^{d\times k}$
be the matrix formed by stacking the vectors $\mathbf{x}_{j}$ corresponding to elements in $L=\{\mathbf{x}_{1},\dots,\mathbf{x}_{k}\}\subseteq\mathcal{X}\setminus S$. Then, the updated covariance matrix becomes
$\mathbf{V}_{\mathcal{S}\cup\mathcal{L}}=\mathbf{V}_{\mathcal{S}}+\mathbf{X}_{\mathcal{L}}\mathbf{X}_{\mathcal{L}}^{\top}$.
To compute $\mathbf{V}_{\mathcal{S}\cup\mathcal{L}}^{-1}$ efficiently, we apply the Woodbury matrix identity:

$$ $\mathbf{V}_{\mathcal{S}\cup\mathcal{L}}^{-1}=\mathbf{V}_{\mathcal{S}}^{-1}-\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{X}_{\mathcal{L}}\left(\mathbf{I}+\mathbf{X}_{\mathcal{L}}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{X}_{\mathcal{L}}\right)^{-1}\mathbf{X}_{\mathcal{L}}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}.$ $$

This yields

$$ $\displaystyle f_{\mathbf{z}}(\mathcal{S}\cup\mathcal{L})=-\mathbf{z}^{\top}\mathbf{V}_{\mathcal{S}\cup\mathcal{L}}^{-1}\mathbf{z}$ $\displaystyle=-\mathbf{z}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{z}+\mathbf{z}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{X}_{\mathcal{L}}\left(\mathbf{I}+\mathbf{X}_{\mathcal{L}}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{X}_{\mathcal{L}}\right)^{-1}\mathbf{X}_{\mathcal{L}}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{z}.$ $$

Hence, the Total gain from adding $\mathcal{L}$ is

$$ $\displaystyle f_{\mathbf{z}}(\mathcal{S}\cup\mathcal{L})-f_{\mathbf{z}}(\mathcal{S})$ $\displaystyle=\mathbf{z}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{X}_{\mathcal{L}}\left(\mathbf{I}+\mathbf{X}_{\mathcal{L}}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{X}_{\mathcal{L}}\right)^{-1}\mathbf{X}_{\mathcal{L}}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{z}.$ $$

Letting $\mathbf{w}=\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{z}\in\mathbb{R}^{d}$, we get

$$ $\displaystyle f_{\mathbf{z}}(\mathcal{S}\cup\mathcal{L})-f_{\mathbf{z}}(\mathcal{S})$ $\displaystyle=\mathbf{w}^{\top}\mathbf{X}_{\mathcal{L}}\left(\mathbf{I}+\mathbf{X}_{\mathcal{L}}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{X}_{\mathcal{L}}\right)^{-1}\mathbf{X}_{\mathcal{L}}^{\top}\mathbf{w}~,$ $$

which can be interpreted as a bilinear form involving the projection of the vector $\mathbf{w}$ onto the subspace spanned by the columns of $\mathbf{X}_{\mathcal{L}}$, i.e., the set $\{\mathbf{x}_{j}\}_{j\in\mathcal{L}}$. The matrix $\left(\mathbf{I}+\mathbf{X}_{\mathcal{L}}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{X}_{\mathcal{L}}\right)^{-1}$ serves as a correction factor that adjusts for the interaction between the new directions $\mathbf{x}_{j}$ and the current inverse covariance structure $\mathbf{V}_{\mathcal{S}}^{-1}$. This reflects how the total gain from adding a group of elements depends not just on their alignment with $\mathbf{w}$, but also on how orthogonal or redundant they are relative to the current set $\mathcal{S}$.

Let us define the matrix

$$ $\mathbf{G}=\mathbf{X}_{\mathcal{L}}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{X}_{\mathcal{L}}=\mathbf{D}+\mathbf{N}~,$ $$

where $\mathbf{D}=\operatorname{diag}\left(\mathbf{x}_{j}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{j}\right)$ is a diagonmal matrix and $\mathbf{N}$ is a matrix with off-diagonal terms $\left(\mathbf{x}_{j}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{k}\right)_{j\neq k}$. We ideally want these terms to be low, i.e., low mutual coherence among the columns of $\mathbf{X}_{\mathcal{L}}$ in the $\mathbf{V}_{\mathcal{S}}^{-1}$-inner product. Now see that

$$ $\displaystyle(\mathbf{D}+\mathbf{N})^{-1}=\left(\mathbf{D}\left(\mathbf{I}+\mathbf{D}^{-1}\mathbf{N}\right)\right)^{-1}=\left(\mathbf{I}+\mathbf{D}^{-1}\mathbf{N}\right)^{-1}\mathbf{D}^{-1}.$ $$

If $\|\mathbf{D}^{-1}\mathbf{N}\|\ll 1$, then we can expand the inverse using a Neumann series:(^3^33This condition holds when vectors $\mathbf{x}_{i}$ and $\mathbf{x}_{j}$ are nearly orthogonal under the $\mathbf{V}_{\mathcal{S}}^{-1}$-induced geometry, which ensures that the off-diagonal elements of $\mathbf{G}$ are small relative to the diagonal ones.)

$$ $\left(\mathbf{I}+\mathbf{D}^{-1}\mathbf{N}\right)^{-1}=\sum_{k=0}^{\infty}(-1)^{k}(\mathbf{D}^{-1}\mathbf{N})^{k}.$ $$

Then we obtain

$$ $(\mathbf{D}+\mathbf{N})^{-1}=\left(\sum_{k=0}^{\infty}(-1)^{k}(\mathbf{D}^{-1}\mathbf{N})^{k}\right)\mathbf{D}^{-1}.$ $$

Keeping only the first two terms in the series, we get

$$ $(\mathbf{D}+\mathbf{N})^{-1}\approx\mathbf{D}^{-1}-\mathbf{D}^{-1}\mathbf{N}\mathbf{D}^{-1}.$ $$

This approximation is accurate when the norm of $\mathbf{N}$ is small relative to $\mathbf{D}$, i.e., $\|\mathbf{D}^{-1}\mathbf{N}\|=\|\mathbf{D}^{-1/2}\mathbf{N}\mathbf{D}^{-1/2}\|\ll 1$ as $\mathbf{N}$ is a symmetric matrix.
Therefore, we get

$$ $(\mathbf{I}+\mathbf{G})^{-1}\approx\operatorname{diag}\left(\frac{1}{1+\mathbf{x}_{j}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{j}}\right)+\mathbf{D}^{-1}\mathbf{N}\mathbf{D}^{-1}.$ $$

Now, note that

$$ $\displaystyle\mathbf{w}^{\top}\mathbf{X}_{\mathcal{L}}\operatorname{diag}\left(\frac{1}{1+\mathbf{x}_{j}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{j}}\right)\mathbf{X}_{\mathcal{L}}^{\top}\mathbf{w}$ $\displaystyle=\sum_{i=1}^{k}\frac{(\mathbf{z}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{i})^{2}}{1+\mathbf{x}_{i}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{i}}=\sum_{i=1}^{k}\Delta_{i}$ $$

Similarly, for the diagonal matrix
$\mathbf{D}=\operatorname{diag}(d_{1},d_{2},\dots,d_{n})$ and the off-diagonal matrix
$\mathbf{N}=[n_{ij}]\in\mathbb{R}^{n\times n}$, we have

$$ $\mathbf{w}^{\top}\mathbf{D}^{-1}\mathbf{N}\mathbf{D}^{-1}\mathbf{w}=\sum_{i\neq j}\frac{w_{i}w_{j}}{d_{i}d_{j}}n_{ij}.$ $$

This gives us

$$ $\displaystyle\mathbf{w}^{\top}\mathbf{X}_{\mathcal{L}}\mathbf{D}^{-1}\mathbf{N}\mathbf{D}^{-1}\mathbf{X}_{\mathcal{L}}^{\top}\mathbf{w}$ $\displaystyle=\sum_{i\neq j}\frac{(\mathbf{z}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{i})(\mathbf{z}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{j})}{(1+\mathbf{x}_{i}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{i})(1+\mathbf{x}_{j}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{j})}(\mathbf{x}_{i}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{j})$ $\displaystyle=\sum_{i\neq j}\sqrt{\Delta_{i}\Delta_{j}}\mu_{i,j}~,$ $$

where $\mu_{i,j}=\frac{\mathbf{x}_{i}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{j}}{\sqrt{1+\mathbf{x}_{i}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{i}}\sqrt{1+\mathbf{x}_{j}^{\top}\mathbf{V}_{\mathcal{S}}^{-1}\mathbf{x}_{j}}}$.
Therefore

$$ $f_{\mathbf{z}}(\mathcal{S}\cup\mathcal{L})-f_{\mathbf{z}}(\mathcal{S})=\sum_{i=1}^{k}\Delta_{i}-\sum_{i\neq j}\sqrt{\Delta_{i}\Delta_{j}}\mu_{i,j}~.$ $$

This implies

$$ $\gamma_{k}(f,\mathbf{z},\mathcal{L},\mathcal{S})=\frac{\sum_{i=1}^{k}\Delta_{i}}{\sum_{i=1}^{k}\Delta_{i}-\sum_{i\neq j}\sqrt{\Delta_{i}\Delta_{j}}\mu_{i,j}}~.$ $$

Let us denote $a_{i}:=\sqrt{\Delta_{i}}$. Then we have

$$ $\gamma_{k}(f,\mathbf{z},\mathcal{L},\mathcal{S})=\frac{\sum\limits_{i=1}^{k}a_{i}^{2}}{\sum\limits_{i=1}^{k}a_{i}^{2}-\sum\limits_{i\neq j}a_{i}a_{j}\mu_{i,j}}.$ $$

Using the bound $|\mu_{i,j}|\leq\mu$, we get

$$ $\displaystyle|\sum_{i\neq j}a_{i}a_{j}\mu_{i,j}|\leq\mu\sum_{i\neq j}|a_{i}a_{j}|$ $\displaystyle\implies-\sum_{i\neq j}a_{i}a_{j}\mu_{i,j}\leq\mu\sum_{i\neq j}|a_{i}a_{j}|~.$ $$

To bound $\sum_{i\neq j}|a_{i}a_{j}|$, note that

$$ $\displaystyle\left(\sum_{i=1}^{k}|a_{i}|\right)^{2}=\sum_{i=1}^{k}a_{i}^{2}+\sum_{i\neq j}|a_{i}a_{j}|$ $\displaystyle\implies\sum_{i\neq j}|a_{i}a_{j}|=\left(\sum_{i=1}^{k}|a_{i}|\right)^{2}-\sum_{i=1}^{k}a_{i}^{2}$ $\displaystyle\implies\sum_{i\neq j}|a_{i}a_{j}|\leq(k-1)\sum_{i=1}^{k}a_{i}^{2}~,$ $$

since by the Cauchy-Schwarz inequality
$\left(\sum_{i=1}^{k}|a_{i}|\right)^{2}\leq k\sum_{i=1}^{k}a_{i}^{2}$.
Therefore, we can bound the denominator as

$$ $\sum_{i=1}^{k}a_{i}^{2}-\sum_{i\neq j}a_{i}a_{j}\mu_{i,j}\leq\left(1+(k-1)\mu\right)\sum_{i=1}^{k}a_{i}^{2}~,$ $$

which yields

$$ $\displaystyle\gamma_{k}(f,\mathbf{z},\mathcal{L},\mathcal{S})\geq\frac{1}{1+(k-1)\mu}~.$ $$

Since the bound holds for any $\mathbf{z}$, any $\mathcal{S}\subseteq\mathcal{X}$ and any $\mathcal{L}\subseteq\mathcal{X}$ such that $\cal{L}\cap\mathcal{S}=\emptyset$, we can conclude the proof.

### Proof of Theorem 4

From Das and Kempe (2011), we know that for any $\mathbf{z}\in\mathbb{R}^{d}$, if the greedy algorithm return a set $\mathcal{S}_{\text{greedy}}$ for the optimization problem in equation ([4](https://arxiv.org/html/2509.15676v1#Sx4.E4)), then
$\mathcal{S}_{\text{greedy}}$ satisfies

$$ $f_{\mathbf{z}}(\mathcal{S}_{\text{greedy}})\geq\left(1-e^{-\gamma}\right)\cdot f(\mathcal{S}^{\ast}),$ $$

where $S^{\ast}$ is an optimal solution of size at most $k$. Substituting the lower bound on submodularity ratio from Lemma [3](https://arxiv.org/html/2509.15676v1#Thmtheorem3), i.e., $\gamma\geq\frac{1}{1+(k-1)\mu}$, we get the desired bound.

### Operator Identity

The following well-known result on linear operators helps us apply the kernel trick to compute the greedy selection rule of Kite, see equation ([Moving Beyond Linearity: Kernel Trick](https://arxiv.org/html/2509.15676v1#Sx4.Ex12)).

###### Lemma 6 .

Let $\mathbf{A}$ be a linear operator. Then, for any $\beta>0$, the following holds:

$$ $\displaystyle(\mathbf{A}^{\top}\mathbf{A}+\beta\mathbf{I})^{-1}\mathbf{A}^{\top}$ $\displaystyle=\mathbf{A}^{\top}(\mathbf{A}\mathbf{A}^{\top}+\beta\mathbf{I})^{-1},$ $\displaystyle\mathbf{I}-\mathbf{A}^{\top}(\mathbf{A}\mathbf{A}^{\top}+\beta\mathbf{I})^{-1}\mathbf{A}$ $\displaystyle=\beta(\mathbf{A}^{\top}\mathbf{A}+\beta\mathbf{I})^{-1}.$ $$

**Table 5: Datasets with corresponding prompts and examples used in the experiments.**
| Dataset | Prompt | Example | #Train | #Validation |
| --- | --- | --- | --- | --- |
| SST-5 | {input} It is {output} | Input: this is a stunning film, a one-of-a-kind tour de force. | 8,534 | 1,101 |
|  |  | Output: very positive |  |  |
| MRPC | {input1} Can we say “{input2}”? {output} | Input1: The company didn’t detail the costs of the replacement and repairs. | 3,668 | 408 |
|  |  | Input2: But company officials expect the costs of the replacement work to run into the millions of dollars. |  |  |
|  |  | Output: No |  |  |
| MNLI | {input1} Can we say “{input2}”? {output} | Input1: yeah i know and i did that all through college and it worked too | 392,568 | 19,647 |
|  |  | Input2: I did that all through college but it never worked |  |  |
|  |  | Output: No |  |  |
| QNLI | {input1} Can we know “{input2}”? {output} | Input1: As of that day, the new constitution heralding the Second Republic came into force. | 104,707 | 5,463 |
|  |  | Input2: What came into force after the new constitution was herald? |  |  |
|  |  | Output: Yes |  |  |
| CMSQA | {input}{output} | Input: Sammy wanted to go to where the people were. Where might he go? | 9,740 | 1,221 |
|  |  | Output: populated areas |  |  |
| HellaSwag | {input}{output} | Input: Members of the procession walk down the street holding small horn brass instruments. A drum line | 52,611 | 20,006 |
|  |  | Output: passes by walking down the street playing their instruments |  |  |

###### Lemma 6 .

### Dataset Descriptions

#### SST-5

is a sentiment classification benchmark containing five fine-grained classes: *very positive*, *positive*, *neutral*, *negative*, and *very negative*.

#### MRPC

is a corpus of sentence pairs automatically extracted from online news sources, with human annotations indicating whether the sentences in each pair are semantically equivalent.

#### MNLI

is a crowdsourced collection of premise–hypothesis sentence pairs annotated for textual entailment. The task is to predict whether the premise *entails*, *contradicts*, or is *neutral* with respect to the hypothesis.

#### QNLI

is a question–sentence dataset derived from QA pairs, where the task is to determine whether the given context sentence contains the answer to the question.

#### CMSQA

is a multiple-choice question-answering dataset requiring diverse forms of commonsense knowledge. Given a question and five candidate answers, the model must select the correct one.

#### HellaSwag

is a large-scale benchmark for grounded commonsense reasoning. Each example pairs a context with four candidate endings: one true video caption (from ActivityNet Captions and the Large Scale Movie Description Challenge) and three adversarially generated distractors designed to fool machines.

## Peer Reviews

confidence: 4
contribution: 1
initial_rating: 2
presentation: 3
questions: 1. Computational scalability. The proposed approach requires computing the in-context examples for each test query, which is computationally expensive. What’s the computational overhead introduced by this method over existing baselines? 
2. How does the model perform if the embedding is changed from BERT to other embeddings, such as E5 or RoBERTa?
3. It is known that small LLMs exhibit weaker in-context learning ability. Will the method work for larger models?
4. It is unclear how the approach works for multi-class classification. Does the author have one kernel for all classes, or each class has its own kernel?
soundness: 3
strengths: 1. The paper is well-written and of mathematical rigor. 
2. Under the linear model assumption, the methodology and analysis are reasonable and well justified. 
3. Empirically, the proposed algorithm. KITE outperforms multiple traditional baselines.
summary: The paper proposes KITE, a theoretically grounded framework for exemplar selection in in-context learning (ICL). Specifically, the paper first proposes LITE, which models the ICL exemplar selection problem as a query-specific optimization problem under a linear model assumption, and derives a surrogate objective that is approximately submodular. This enables a greedy selection algorithm. To incorporate more complex non-linear scenarios, the paper further introduced KITE, which is a kernelized version of the algorithm. Empirically, KITE outperforms several existing retrieval baselines on six classification datasets across multiple LLMs.
weaknesses: 1. The paper lacks novelty and excitement. Theoretically, the analysis is straightforward in a standard linear model context. Methodologically, using an approximate submodular function for in-context learning has already been explored in prior works, such as [1]. 
2. Missing baselines. The paper only compares with very classic retrieval baselines, but most of the recent advances for exemplar selection in LLMs are missing, such as [1] [2]. In addition, the paper should compare with the most basic baseline where no example but only a task description is provided. 
3. The focus is primarily on classification tasks developed for classic NLP, and many LLMs can already perform very well on them. Nowadays, there are many different evaluation approaches, such as instruction following and reasoning. 
4. The proposed approach is model-agnostic, meaning that the exemplars are selected only using BERT without knowing the target LLM. 

[1] Wu, Z., Lin, X., Dai, Z., Hu, W., Shu, Y., Ng, S. K., ... & Low, B. K. H. (2024). Prompt optimization with EASE? efficient ordering-aware automated selection of exemplars. Advances in Neural Information Processing Systems, 37, 122706-122740.

[2] Nguyen, T., & Wong, E. (2023). In-context example selection with influences. arXiv preprint arXiv:2302.11042.

---

confidence: 2
contribution: 2
initial_rating: 6
presentation: 3
questions: See above.
soundness: 2
strengths: 1. The motivation and writing are clear, easy to follow, and grounded in theory. 
2. This paper provides a clear mathematical formulation with proofs of approximate submodularity and corresponding approximation guarantees.
 3. The core idea is fresh and well-motivated. It goes beyond the usual heuristic retrieval methods and gives a principled way to think about the problem.
summary: This paper presents KITE, a theoretically grounded framework for exemplar selection in in-context learning (ICL). The method formulates exemplar retrieval as a query-specific optimization problem under an information-theoretic view. By modeling the LLM as approximately linear in its input embeddings, the authors derive a submodular surrogate objective and design a greedy selection algorithm with provable approximation guarantees. KITE extends this linear model via the kernel trick, allowing flexible, nonlinear feature mappings. Experiments on six few-shot classification benchmarks (SST-2/5, CMSQA, MRPC, QNLI, HellaSwag) and three relatively small-scale LLMs (GPT-Neo-2.7B, Qwen-1.5B, Llama-3B) show consistent improvements over common retrieval baselines such as BM25, dense embeddings, and DPP retrieval.
weaknesses: 1. Empirical scope is limited. Evaluation is restricted to classification tasks, which aligns with the linear assumption but leaves uncertainty about generalization to generation or structured prediction tasks (e.g., summarization, reasoning). Extending experiments beyond classification would strengthen the practical relevance.
2. Dense baseline clarity. Table 1 refers to a “Dense” retriever, but the implementation details (e.g., whether it uses bi-encoder fine-tuning or off-the-shelf embeddings) are under-specified. This is critical since modern dense retrievers (e5, bge) are highly competitive for ICL retrieval.
3. Model scale. The largest tested model is 3B parameters. Given recent trends, evaluating at least one 7B-scale model (e.g., Llama-3-8B or Qwen-2.5-7B) would better demonstrate scalability.

---

confidence: 3
contribution: 3
initial_rating: 4
presentation: 2
questions: * The KITE framework's entire theoretical foundation is built on modeling the LLM's behavior as a (kernelized) linear regressor. This model inherently assumes a smooth performance landscape, where embedding-space proximity correlates with performance similarity. However, what if the true LLM performance landscape is highly discontinuous, where minor variations in exemplar embeddings can lead to drastic and non-linear shifts in task accuracy?

* You motivate your work by highlighting the "curse of dimensionality" as a key drawback of kNN-based methods, a flaw typically associated with Euclidean distance. However, your "Dense" baseline (KATE) presumably uses **Cosine Similarity**, as is standard in NLP to mitigate this very issue. Could you clarify which metric your "Dense" baseline actually uses? Furthermore, could you provide an ablation study comparing KITE against both a Euclidean-kNN baseline and a Cosine-kNN baseline, to demonstrate that the choice of retrieval algorithm (KITE vs. kNN), and not just the distance metric, is the true source of performance gain?
soundness: 3
strengths: 1. Originality：
KITE’s most original contribution is the query-specific optimization framing—unlike prior work (e.g., KATE, DPPs) that optimizes for distribution-wide generalization, KITE targets error reduction for individual user queries. This aligns with real-world ICL, where performance on a single query (not a dataset) determines utility. 
2. Quality：
The theoretical analysis is rigorous: the derivation of the surrogate objective from prediction error bounds (via Chernoff inequalities) is logically sound. Empirically, the paper uses diverse datasets (sentiment, QA, entailment) and LLMs, demonstrating KITE’s robustness. The kernel ablation provides actionable insights, adding depth to the results.
3. Clarity：
The paper avoids excessive jargon in key sections (e.g., the Introduction clearly explains ICL and exemplar selection). Algorithm 1 is well-documented, with explicit steps for initializing and computing relevance/diversity scores. 
4. Significance:
KITE addresses a high-impact problem: ICL performance is highly sensitive to exemplar choice, and existing methods struggle in high-dimensional spaces. By providing a computationally efficient, unsupervised framework with theoretical guarantees, KITE enables practical deployment of ICL in data-scarce domains.
summary: This paper addresses the critical problem of exemplar selection for in-context learning (ICL), a core challenge given the limited context windows of large language models (LLMs). The key limitation of existing methods (e.g., kNN-based KATE, DPPs) is their poor generalization in high-dimensional embedding spaces and insufficient diversity in selected examples. The authors model the LLM as a linear function in the embedding space and aim to select a subset of examples that minimizes prediction error on a given query. The key contributions are:
1. A theoretical formulation of ICL exemplar selection as a subset selection problem with a surrogate objective that is approximately submodular.
2. A greedy algorithm with an approximation guarantee, leveraging the Sherman–Morrison formula for efficient updates.
3. Kernelization of the approach to handle nonlinear relationships via RKHS, without explicit high-dimensional mappings.
4. A diversity regularizer inspired by optimal experimental design to improve generalization.
5. Extensive experiments across multiple classification tasks and models, showing consistent improvements over strong baselines like DPP, BM25, and dense retrieval.
weaknesses: * The KITE framework is fundamentally order-agnostic, as it is designed as a subset selection algorithm based on optimizing set functions. This theoretical formulation completely overlooks the well-documented sensitivity of large language models to exemplar ordering, a crucial factor in ICL performance that this method cannot account for.

* The algorithm's generalizability is constrained by its reliance on a fixed, pre-compiled example bank. As a frozen and non-trainable retriever, it cannot update its parameters for a new task (like fine-tuning) or adapt by collecting new data (like online learning). Its applicability seems restricted to scenarios where a high-quality example bank for the target task is already available.