# Unraveling the Mechanics of Learning-Based Demonstration Selection for In-
Context Learning

Hui Liu
City University of Hong Kong
Hong Kong
<EMAIL_ADDRESS>
&Wenya Wang
Nanyang Technological University
Singapore
<EMAIL_ADDRESS>
&Hao Sun
Peking University
Beijing, China
<EMAIL_ADDRESS>
&Chris XING TIAN
City University of Hong Kong
Hong Kong
<EMAIL_ADDRESS>
&Chenqi Kong
Nanyang Technological University
Singapore
<EMAIL_ADDRESS>
&Xin Dong
NVIDIA
America
<EMAIL_ADDRESS>
&Haoliang Li
City University of Hong Kong
Hong Kong
<EMAIL_ADDRESS>

###### Abstract

Large Language Models (LLMs) have demonstrated impressive in-context learning
(ICL) capabilities from few-shot demonstration exemplars. While recent
learning-based demonstration selection methods have proven beneficial to ICL
by choosing more useful exemplars, their underlying mechanisms are opaque,
hindering efforts to address limitations such as high training costs and poor
generalization across tasks. These methods generally assume the selection
process captures similarities between the exemplar and the target instance,
however, it remains unknown what kinds of similarities are captured and vital
to performing ICL. To dive into this question, we analyze the working
mechanisms of the learning-based demonstration selection methods and
empirically identify two important factors related to similarity measurement:
1) The ability to integrate different levels of task-agnostic text
similarities between the input of exemplars and test cases enhances
generalization power across different tasks. 2) Incorporating task-specific
labels when measuring the similarities significantly improves the performance
on each specific task. We validate these two findings through extensive
quantitative and qualitative analyses across ten datasets and various LLMs.
Based on our findings, we introduce two effective yet simplified exemplar
selection methods catering to task-agnostic and task-specific demands,
eliminating the costly LLM inference overhead.

##  1 Introduction

While advanced large language models have demonstrated impressive zero-shot
capabilities across various tasks, in-context learning (ICL) remains a
promising paradigm that employs a sequence of demonstration exemplars as
prompts to assist LMs in effectively performing unseen tasks [6, 35]. However,
the performance of ICL can be sensitive to the choice, format, and order of
the in-context examples [51, 52, 38, 22]. Our work delves into the choice
perspective and unravels important factors contributing to selecting good in-
context exemplars to enhance LLMs’ performances [33, 44, 1, 37, 18, 24].

In-context exemplar selection assumes the availability of a demonstration set
$\mathcal{D}$ containing input-output examples $(x,y)$ for a specific task.
Given a test case $x^{t}$, ICL relies on selecting effective examples from
$\mathcal{D}$ to inform the prediction of $x^{t}$. To achieve that, most
studies select demonstration exemplars based on a similarity measurement
between $x$ and $x^{t}$. Some work utilizes surface similarity computed by
BM25 and semantic similarity111Surface, syntactic and semantic similarity
focus on comparing texts based on word or phrase overlap, structural and
grammatical similarity and the underlying meaning of texts, respectively [13,
23] computed by off-the-shelf text encoders [21, 1]. Recent studies (which we
refer to as learning-based methods) [33, 44, 19], however, separately train a
retriever for each task by composing positive exemplars $x^{+}$ and negative
exemplars $x^{-}$ via interacting with LLMs. This data creation process often
requires hundreds of thousands of queries to LLMs for a single task to collect
sufficient positive/negative data.

The use of learning-based methods, as discussed in references [33, 44, 19],
has demonstrated notable performance improvements over traditional models such
as BM25 or standard text encoders in specific tasks. However, these approaches
are not without limitations, including the costly process of data collection
and poor generalization across different tasks. These challenges underscore
the need for a deeper and more thorough understanding of the underlying
similarities these methods capture and how these similarities correlate with
their performance. By exploring these aspects, it will be beneficial to
develop more efficient and generalizable methods for exemplar selection in the
future.

In this work, we identify two dimensions of similarity: task-agnostic and
task-specific. Task-agnostic similarity encompasses multi-level aspects
including surface, syntactic, and semantic features. These are general
linguistic features that can be shared across various tasks. In contrast,
task-specific similarity measures the distance from an arbitrary exemplar to a
test case sample from a specific task. We find that existing learning-based
methods learn a subset of these similarities, which explains their varied
performance across different tasks. We summarize our findings as follows:

$\mathcal{F}_{1}$: After training, the retriever is analogous to an ensemble
model which adaptively integrates multi-level similarities from different
layers of the encoder between the exemplar input ($x$) and test cases
($x^{t}$) for different tasks. Compared to off-the-shelf encoders which only
adopt the last layer or BM25 which focuses on surface-level similarities, the
learning-based retriever encompasses diverse task-agnostic similarities.
$\mathcal{F}_{2}$: Beyond input similarities ($x$ and $x^{t}$), the training
process of learning-based methods further encourages selecting exemplars with
similar output ($y$) to the hidden output of the test case ($y^{t}$),
implicitly predicted during retrieval. This enhances the model’s
discriminative power for a specific task.

Extensive quantitative experiments are conducted to support our findings.
Drawing insights from these findings, we further propose two cost-effective
variants for learning-based methods: 1) Multi-level Similarity Maximization
(MLSM) retriever that aims to maximize agreement across different similarities
during the inference time of LLMs. 2) Test Task Fine-tuning (TTF) retriever,
which uses labeled data from the demonstration set to fine-tune the retriever
to learn task-specific discriminative information. Both retrievers eliminate
the need for expensive interactions with LLMs, catering to cross-task and
task-specific demands respectively. To validate the effectiveness of these
methods, we conducted experiments across three distinct LLMs and a range of
tasks, which not only confirms the versatility of our assumptions but also
provides valuable insights for future research in ICL.

##  2 Preliminary

###  2.1 Learning-based Demonstration Selection

Demonstration selection aims to identify a sequence of high-quality in-context
examples from the demonstration set as prompt to enhance test case accuracy on
LLMs. Prior studies [21, 9] find that good exemplar exhibit similarities with
the test case. They employ the pre-trained text encoder like BERT [7] as the
retriever to encode inputs and take the average embedding of all tokens from
the final layer of this model to represent test cases and exemplars.
Subsequently, cosine similarity measures are computed between test cases and
exemplars to retrieve the top-K most similar exemplars as prompts.

While the pipeline of learning-based demonstration selection methods[33, 44,
19] is similar to the above strategy, they further exploit LLMs to label
positive and negative exemplars to construct a proxy task to fine-tune BERT,
aiming to learn a better similarity metric. Specifically, let $\mathcal{D}$
denote the demonstration set. Given an exemplar $(x_{i},y_{i})$ in
$\mathcal{D}$, Rubin et al. [33] propose EPR to sample a sequence of candidate
examples from $\mathcal{D}$, denoted as
$S=\\{(\overline{x}_{1},\overline{y}_{1}),\ldots,(\overline{x}_{m},\overline{y}_{m})\\}$
and score them by
$s(\overline{x},\overline{y})=P_{\mathrm{LLM}}(Y=y_{i}|(\overline{x},\overline{y}),x_{i})$,
representing the probability of producing correct output $y_{i}$ for $x_{i}$
conditioned on $(\overline{x},\overline{y})$ using an LLM. Subsequently, the
$\overline{x}$ with the highest score is selected as the positive sample,
denoted as $x^{+}$ and the lowest as the hard negative sample, denoted as
$x^{-}$ for $x_{i}$. These samples are then used to train the retriever (BERT)
by increasing the similarity between $x$ and $x^{+}$ and decreasing the
similarity between $x$ and $x^{-}$ via contrastive learning. In subsequent
sections, without special note, we conduct analysis on EPR to unravel the
mechanics of learning-based demonstration selection. We adopt
BERT222https://huggingface.co/google-bert/bert-base-uncased consisting of
twelve transformer layers as the retriever.

###  2.2 Layers of BERT as Anchors of Multi-level Linguistic Similarity

Previous studies [13, 23] have empirically showed that the intermediate layers
of BERT encode a rich hierarchy of linguistic information with surface
features at the bottom, syntactic features in the middle, and semantic
features at the top through probing tasks. For each layer $l$, given two texts
$s_{1}$ and $s_{2}$, we can extract the text embeddings of all tokens from
this layer and compute the average pooling of these tokens to represent each
text, denoted as $x^{l}_{1}$ and $x^{l}_{2}$. The similarity corresponding to
different layers can then be calculated by computing the cosine similarity
between these embeddings: $\cos(x^{l}_{1},x^{l}_{2})$. While surface,
syntactic, and semantic similarities aim to compare texts based on word or
phrase overlap, structural and grammatical patterns, and underlying meanings,
we propose that the different layers of the original BERT (i.e., BERT without
task-specific fine-tuning) can serve as anchors to represent such task-
agnostic multi-level linguistic similarities between different text inputs
because BERT has been pre-trained on a huge corpus capturing general
linguistic features that can be utilized for a wide range of tasks.

##  3 Rethinking Learning-based Demonstration Selection

This section proposes two key findings regarding the underlying similarity
mechanisms of learning-based exemplar section methods: ($\mathcal{F}_{1}$) The
learning-based retriever is analogous to an ensemble model which adaptively
aggregate multi-level similarities from different layers between the input of
exemplar and test cases ($x$ and $x^{t}$) for different tasks.
($\mathcal{F}_{2}$) The learning-based retriever favors selecting exemplars
with similar output ($y$) to the output of the test case ($y^{t}$). Moreover,
the quantitative validation is conducted for both findings.

###  3.1 Multi-level Similarity ($\mathcal{F}_{1}$)

Figure 1: Left: Top-10 retrieval accuracy using each of the twelve layers of
the original BERT to retrieve positive exemplars to solve the proxy task of
EPR across ten tasks. Different colors represents different layers. Top-10
accuracy refers to the percentage of times the positive exemplar is retrieved
within the top 10 predictions. Middle: CKA scores between twelve layers of
original BERT and the final layer of BERT of EPR trained on ten tasks and the
training-free BERT. Right: CKA scores between each layer of the original BERT.
These CKA scores are min-max normalized for better visualization. We use GPT-
Neo [5] as the LLM.

When considering similarity-based methods for demonstration selection, it’s
important to note there exists multiple levels of linguistic similarities,
encompassing surface, syntactic, and semantic levels. Although semantic
similarity generally excels in text retrieval, our observations show that
surface-level similarity (e.g., BM25) can sometimes outperform semantic
counterpart in demonstration retrieval for ICL, especially on Nl2Bash [20] and
SWAG [46]. Thus, we speculate that a critical aspect which makes the learning-
based exemplar retriever effective lies in its ability to potentially learn
different levels of similarities within the text encodings which is task-
agnostic and integrate these similarities ($\mathcal{F}_{1}$). We empirically
find that the learning-based method EPR acts as a dynamic similarity selector
that ensembles different layers from an off-the-shelf BERT encoder, detailed
in the following sections.

#### Quantitative Validation.

As a first step, we validate our assumption that different layers of the text
encoder (representing different levels of linguistic characteristics), when
used as a retriever, exhibit different behaviors for different tasks. To do
that, following EPR which builds a positive set
$\\{(x_{i},x_{i}^{+})\\}_{i}^{N}$ based on $x_{i}$ using an LLM (Section 2.1),
we treat each $x_{i}^{+}$ as a gold exemplar to be retrieved for $x_{i}$. Then
we utilize different layers of the original pre-trained BERT as anchors to
retrieve/rank exemplars (as discussed in Section 2.2) given each $x_{i}$ and
evaluate the Top-10 retrieval accuracy compared with the gold exemplar. The
results on four tasks are illustrated in Fig. 1 (Left), which reveal that
different tasks exhibit distinct preferences towards specific layers,
emphasizing different linguistic properties. More information on these tasks
can be found in Section 5. While it is a prevalent practice to employ features
from the final layer of original BERT for exemplar retrieval [21, 48], it is
not consistently optimal, likely due to potential inclusion of irrelevant
information relevant to BERT’s pre-training tasks.

In the next step, we investigate what is encapsulated in the retriever learned
using EPR. As the retriever utilises the last layer to compute similarities,
we extract the representations from the last layer and compare those with the
representations given by each layer of the original BERT to study the
correlations between the learned retriever and each prior linguistic knowledge
encoded in the original BERT layers. For this purpose, We introduce CKA [14],
which effectively identifies correspondences between representations in
different networks. Let $X^{a}\in\mathbb{R}^{n\times p_{1}}$ denote a matrix
of activations of $p_{1}$ neurons for $n$ examples and
$X^{b}\in\mathbb{R}^{n\times p_{2}}$ denote a matrix of activations of $p_{2}$
neurons for the same $n$ examples. The core insight of CKA lies in measuring
the similarity between two matrices $X^{a}$ and $X^{b}$ by considering the
inter-sample similarities. Specifically, CKA computes $K^{a}$ and $K^{b}$ to
derive the inter-example similarity structures for $X^{a}$ and $X^{b}$, where
$K^{a}_{ij}=k^{a}(x^{a}_{i},x^{a}_{j})$,
$K^{b}_{ij}=k^{b}(x^{b}_{i},x^{b}_{j})$, and $k^{a}$, $k^{b}$ represents two
kernels. Then CKA metric can be formulated as follows:

$\mathrm{CKA}(K^{a},K^{b})=\frac{\mathrm{HSIC(K^{a},K^{b})}}{\sqrt{\mathrm{HSIC}(K^{a},K^{a})\mathrm{HSIC}(K^{b},K^{b})}},\quad\mathrm{HSIC(K^{a},K^{b})}=\frac{1}{(n-1)^{2}}tr(K^{a}HK^{b}H).$

Here $H$ is the centering matrix ${H_{n}=I_{n}-{\tfrac{1}{n}}J_{n}}$ where
$I_{n}$ is the identity matrix of size $n$ and $J_{n}$ is a $n$-by-$n$ all-
ones matrix.

We use CKA to measure the similarity between the last layer of the trained
retriever and each layer of the original BERT. To do that, we randomly sample
$n=2000$ instances from the demonstration set $\mathcal{D}$ (If
$|\mathcal{D}|<2000$, $n=|\mathcal{D}|$) for each task. Denote $X^{EPR}$ as
the matrix composing $n$ rows of last-layer representations from the trained
retriever on each of the ten tasks, and $X^{l}$ as the matrix composing $n$
rows of tje $l$th-layer representations from the orginal BERT model. Then we
can calculate the CKA scores $CKA(K^{EPR},K^{l})$ for each task.

The results are illustrated in Fig. 1 (Middle). Each row reflects the CKA
similarity between the trained retriever for a specific task and each layer of
the original BERT (regardless of tasks). Obviously, the CKA distribution
across various tasks exhibits significant diversity among different pretrained
BERT layers. This finding supports $\mathcal{F}_{1}$ that learning-based
methods can learn to adaptively aggregate multi-level(layer) linguistic
similarities catering to different tasks. For instance, the results suggest
that the exemplar retriever trained on Nl2Bash and SWAG tasks may prioritize
low-level features, corroborating our experimental results where the
BM25-based method outperforms higher-level semantic-based ones on these
datasets.

###  3.2 Output Simialirty ($\mathcal{F}_{2}$)

Figure 2: Left: Comparison of similarity between the input/output of positive
and negative demonstration examples and the input/output of the test case
across ten tasks for EPR. Right: Difference in average similarity between the
output of test case and retrieved example for EPR and each of the three
learning-free prompt retrieval methods. We use GPT-Neo [5] as the LLM.

When employing the learning-based paradigm to acquire better similarity
mechanics between exemplars and test cases for ICL, the expectation is that
such mechanics can perform well on unseen tasks, thus justifying the high cost
of data collection. However, the sub-optimal generalization performance
revealed by Ye et al. [44] suggests that the exemplar retriever, trained on
the proxy task, primarily learns task-specific information. As data and
training objective can serve as a lens to analyze the behaviors of neural
network models, we first investigate the data generated based on LLM feedback
involving positive and negative pairs, as shown in Fig. 2 (Left), which
reveals the similarities between the input of positive/negative exemplars and
the test case as well as the output of positive/negative exemplars and the
test case. Specifically, for the input similarity, we compute text similarity
using sentence-transformers333https://huggingface.co/sentence-
transformers/paraphrase-MiniLM-L6-v2 for all tasks, while we compute the exact
match for the first three classification tasks and the text similarity for
other QA and generation tasks for the output similarity. Let $(x,y)$,
$(x^{+},y^{+})$, $(x^{-},y^{-})$ denote the test case and corresponding
positive and negative exemplars. The results indicate that the similarity
between $x^{+}$ and $x$ is significantly higher than that between $x$ and
$x^{-}$, affirming the efficacy of similarity-based heuristic methods.
Moreover, it is noteworthy that the similarity between $y^{+}$ and $y$ is also
markedly higher than that between $y$ and $y^{-}$. Additionally, in the SST-5
dataset, even when the output of positive and negative exemplars are extremely
similar to $y$ to provide useful patterns, $x^{+}$ remains more similar to $x$
than $x^{-}$, demonstrating the importance of input similarities.

Acknowledge that the training objective of the proxy task is to push the
embedding of $x$ and $x^{+}$ closer and push $x$ and $x^{-}$ away through
contrastive learning in the embedding space. As a result, during the training
phase, demonstration examples with similar outputs will resemble each other in
this space due to the strong correlation between $y$ and $y^{+}$. This leads
to a higher probability of selecting exemplars with outputs similar to the
test case as prompts when the test case’s output is unknown. Therefore, we
suggest that the success of learning-based approaches partly stems from the
implicit prediction of the output of test cases (i.e., $P(Y|X)$) during
exemplar retrieval ($\mathcal{F}_{2}$), which can be viewed as computing
similarity of the joint input-output distribution between the test case and
exemplar.

#### Quantitative Validation.

After training the examplar retriever on the proxy task, we utilize it to
assess the similarity between the input of test case and exemplars from the
demonstration set and select top-K exemplars as prompts. To validate
$\mathcal{F}_{2}$, we evaluate the retriever’s ability to learn the output
similarity by computing the average similarity between the output of test
cases and that of these retrieved exemplars. We compare EPR trained using GPT-
Neo against learning-free methods (i.e., Random, Top-K BM25 and Top-K BERT),
as illustrated in Fig. 2 (Right). We compute the output similarity for all
tasks in the same way as experiments in 2 (Left). The results demonstrate that
the exemplar chosen by EPR have outputs more akin to the test case compared to
the learning-free competitors, particularly in classification tasks where the
output similarity can be well-captured by exact match. Similar validation is
conducted for EPR with GP2-XL, and the results in Fig. 4 align with those
obtained for GPT-Neo, thus providing consistent support for $\mathcal{F}_{2}$.

##  4 Methodology

Given the above two findings and analyses, we propose two simple yet effective
methods inspired from the two findings, respectively. Specifically, we propose
1) Multi-level Linguistic Similarity Maximization (MLSM) accounting for a
better ensemble of task-agnostic layer-wise anchors from BERT to achieve
better task generalization given $\mathcal{F}_{1}$, and 2) Test Task Fine-
tuning (TTF) to explicitly infuse task-specific information to the retriever
according to $\mathcal{F}_{2}$which significantly enhances performances on
classification tasks. These methods and results serve as further validation of
the effectiveness of both findings. In addition, both methods do not require
interacting with LLMs to construct positive and negative exemplar sets,
significantly reducing implementation cost.

###  4.1 Multi-level Linguistic Similarity Maximization (MLSM)

$\mathcal{F}_{1}$ emphasizes that learning-based methods are capable of
adaptively integrating diverse linguistic similarities which can be captured
through different layers of a pretrained text encoder (e.g., BERT) from bottom
to top. Inspired by ensemble learning [28, 3, 49], each layer can work as an
expert for exemplar selection. Our goal is to integrate the insights from all
these experts by maximizing their agreement during the inference of LLMs.

However, as depicted in Fig. 1 (Right), each layer of BERT (without task-
specific fine-tuning on the proxy task of EPR) shows a higher similarity to
adjacent layers due to the residual design of transformers. Hence, our initial
step involves filtering out redundant layers to avoid overfitting to specific
levels of linguistic similarity and reduce computational overhead.
Specifically, given a task and its corresponding demonstration set
$\mathcal{D}$, we sample a subset of unlabeled exemplars from $\mathcal{D}$
and compute layer-wise CKA scores between every pair of BERT layers, forming a
similarity matrix $\mathbf{S}\in\mathbb{R}^{12\times 12}$ where $S_{i,j}$
signifies the similarity between the $i$-th and $j$-th layers of BERT. We then
employ unsupervised K-means clustering algorithm to derive $n_{l}$ clusters,
maximizing the intra-cluster CKA score while minimizing the inter-cluster CKA
score, and designate the central node in each cluster as the representative
layer. Finally, we attain a set of refined layers, denoted as
$L=\\{l_{i}\\}^{n_{l}}_{i=1}$, to represent the linguistic similarity at
varying levels.

For a given test case $x^{t}$, we first sample a mini training set
$\mathcal{D}_{p}=\\{x_{j}\\}^{n_{p}}_{j=1}$ and validation set
$\mathcal{D}_{v}=\\{x_{j}\\}^{n_{v}}_{j=1}$ from $\mathcal{D}$. Then, for each
$l_{i}\in L$, we compute the average of all token embeddings extracted from
the corresponding layer ($l_{i}$) of the original BERT model as the
representation of $x^{t}$ (denoted as $\mathbf{h}^{t}$) and demonstration
examples $x_{j}\in\mathcal{D}_{p}$ (denoted as $\mathbf{h}_{j}$). Following
this, we compute their cosine similarity as
$\mathbf{r}_{i}=[\cos(\mathbf{h}^{t},\mathbf{h}_{1}),...,\cos(\mathbf{h}^{t},\mathbf{h}_{n_{p}})]$,
and normalize it to obtain the probability distribution of these exemplars via
$\mathbf{y}_{i}=\mathrm{softxmax}(\frac{\mathbf{r}_{i}}{\tau})$ for layer
$l_{i}$, where $\tau$ is the temperature parameter. Intuitively, such
distribution represents the ranking distribution of the demonstration examples
when using the linguistic knowledge captured at $l_{i}$ to perform the
similarity metric. After collecting the output distribution of all experts in
$L$, we aggregate experts with learnable aggregation weights, denoted as
$\mathbf{w}\in\mathbb{R}^{n_{l}}$ and get the ensembled prediction as
$\hat{\mathbf{y}}=\mathrm{sofxmax}(\frac{\sum_{i=1}^{n_{l}}{w_{i}\mathbf{r}_{i}}}{\tau})$,
where $\mathbf{w}$ is normalized before aggregation, i.e.,
$\sum_{i}^{n_{l}}w_{i}=1$. Subsequently, to encourage agreement among all
levels of experts, we minimize the loss
$\mathcal{L}=-\sum_{i=1}^{n_{l}}\hat{\mathbf{y}}\cdot\mathbf{y}_{i}$. The
optimal $\mathbf{w}$ can be determined based on the loss on the validation set
$\mathcal{D}_{v}$ by early stopping strategy to prevent overfitting.

In our work, we focus on the scenario involving online streaming test data,
where only one test point is observed during inference to align with the
demand of real-world applications [47, 39, 49]. However, MLSM can enable batch
inference by updating $\mathbf{w}$ using a batch of test cases during the
training phase, thereby enhancing computational efficiency. Notably, MLSM
focuses on general layer-wise linguistic information regardless of specific
tasks, and thus does not rely on any label information from the demonstration
set.

###  4.2 Test Task Fine-tuning (TTF)

$\mathcal{F}_{2}$ posits that the learning-based demonstration retriever
inherently acquires the output similarity between exemplars and test cases for
one specific test task when training on the proxy task. However, such proxy
task relies on costly interactions with LLMs for each distinct task. To
alleviate this issue, we propose to infuse the output information to the
retriever by fine-tuning it with additional modules customized for distinct
tasks using the labeled data from the demonstration set $\mathcal{D}$. We
consider two different fine-tuning frameworks tailoring classification tasks
and generation tasks, respectively.

For convenience, let $f_{\bm{\theta}}$ denote the retriever and
$q_{\bm{\phi}}$ denote the extra module, containing $\bm{\theta}$ and
$\bm{\phi}$ as learnable paramters. For classification tasks, $q_{\bm{\phi}}$
will be instantiated through various classification heads. Given a test input
$x$, assuming a linear classifier, the predication is derived by taking the
argmax over the approximated probability distribution:

$\operatorname*{arg\,max}\limits_{y_{i}}q_{\bm{\phi}}(Y=y_{i}|f_{\bm{\theta}}(x))=\frac{\exp(\mathbf{z}\cdot\bm{\phi}_{i})}{\sum_{j}\exp(\mathbf{z\cdot\bm{\phi}_{j}})},$
(1)

where $\mathbf{z}=f_{\bm{\theta}}(x)$ and $\bm{\phi}_{i}$ is the $i$-th
component of the weights $\bm{\phi}$ corresponding to label $y_{i}$.
Essentially, the prediction is determined by evaluating the distance between
$\bm{\phi}_{i}$ and $\mathbf{z}$. Consequently, test cases with a similar
output are more likely to exhibit a smaller distance in the semantic space, as
they are situated closer to their corresponding $\bm{\phi}$. Furthermore,
previous research [50, 12] has leveraged $\mathbf{z}$ as a pseudo-prototype
for each label to construct non-parameter classifiers, providing evidence that
TTF can effectively encapsulate the input-output relationship (i.e., $P(X|Y)$)
in the wild classification application.

For generation tasks, while decoder-only frameworks are unsuitable to derive
sentence embeddings without prompting or fine-tuning [25], we adopt the
encoder-decoder architecture, where $q_{\bm{\phi}}$ is instantiated by the
decoder and the retriever $f_{\bm{\theta}}$ works as the encoder. However,
determining which part of the model contains effective input-output
relationship remains challenging. Considering that the decoder generates new
tokens based on the encoder’s output, allowing the encoder’s output to
naturally capture pertinent input-output information, we follow Ni et al. [26]
to use the average pooling of all token embeddings extracted from the last
layer of the encoder to represent test cases and exemplars. However, such
method dose not perform well for demonstration selection for generation task,
which will be discussed in Sec. 5.

##  5 Experiments

Table 1: The statistics of ten datasets. We report the number of training
instances after deduplicating.

Type | Dataset | Task | Train | Validation
---|---|---|---|---
Classification | SST-5 [34] | Sentiment Analysis | 8,534 | 1,101
MRPC [8] | Paraphrase Detection | 3,668 | 408
QNLI [40] | Natural Language Inference | 104,707 | 5,463
CMSQA [36] | Commonsense Reasoning | 9,740 | 1,221
HellaSwag [46] | Commonsense Reasoning | 52,611 | 20,006
Generation | WebQs [4] | Open-Domain QA | 3,778 | 2,032
GeoQuery [45] | Code Generation | 404 | 280
Nl2Bash [20] | Code Generation | 7,441 | 609
MTOP [17] | Semantic Parsing | 15,564 | 2,235
SMCalFlow [2] | Semantic Parsing | 102,491 | 14,751

#### Datasets.

We conduct experiments on ten datasets spanning seven distinct categories of
NLP tasks: sentiment analysis, paraphrase detection, natural language
inference, commonsense reasoning, open-domain question answering, code
generation and semantic parsing. As certain datasets lack a test set, we take
the training split as the demonstration set and the validation split for
evaluation across all datasets. The statistics of all dataset are listed in
Table 1. A detailed description of these datasets and prompts to reproduce our
experimental results are shown in Appendix A.1.

#### Baselines.

In line with previous studies [33, 44, 19], we consider the following prompt
retrievers:

  * •

Random: The method randomly selects in-context examples from the demonstration
set without repetition.

  * •

Top-K BM25: This sparse retriever employs BM25 [31], an extension of TF-IDF,
to retrieve the Top-K most similar exemplars based on surface text similarity.

  * •

Top-K BERT: Utilizing BERT [7] as a text encoder, this dense retriever
generates text representations by averaging token embeddings from the final
layer and retrieve the Top-K most similar exemplars based on semantic
similarity.

  * •

EPR [33]: This learning-based retriever uses Top-K BM25 to generate
demonstration candidates and scores them using LLMs to construct a proxy task.
It further fintunes BERT in Top-K BERT using this task.

  * •

CEIL [44]: This method uses EPR to generate demonstration sequence candidates,
scores them using LLMs to construct a proxy task, and also further finetunes
BERT using this task. It balances diversity and relevance using a trade-off
parameter and search optimal exemplar combination using Determinantal Point
Processes [16].

Table 2: Main results on the classification task.

Method SST-5 MRPC QNLI CMSQA SWAG Avg. Unsupervised Random 28.61 65.93 55.08
42.34 41.39 46.67 Top-K BM25 32.06 65.93 60.11 35.79 43.35 47.45 Top-K BERT
32.70 69.12 60.94 35.87 41.09 47.94 2-7[2pt/3pt] MLSM 33.15 69.87 65.02 37.26
41.49 49.36 Supervised EPR 36.88 81.37 77.87 38.74 43.39 55.65 CEIL 37.69
77.94 80.58 38.90 43.84 55.79 2-7[2pt/3pt] TTF 42.14 74.51 85.08 47.83 55.72
61.06

We categorize the baselines into two groups, namely unsupervised and
supervised methods, based on their utilization of labeled data from the
demonstration set. While we mainly utilize BERT as the retriever model of MLSM
and TTF, we also try $T5$ for TTF, denoted as TTF (T5) on the generation
tasks. The implementation detail of our methods and all baselines can be found
in Appendix. A.2.

#### Experiment settings.

We employ GPT-Neo [5] with 2.7B parameter as the main LLM for inference in our
study. Additionally, we conduct experiments on a smaller GPT-2 XL [29] (1.5B)
and text-davinci-002 to verify the transferability of our findings and
corresponding methods. Due to computational constraints and different maximum
context sizes among LMs, we restrict the number of in-context examples to 20.
These exemplars are sorted based on their similarities to the input of test
cases in ascending order following prior practices [33, 1, 21]. Additionally,
in accordance with Ye et al. [44], we access the performance of all methods by
comparing the predicted output with ground truth and report Accuracy (Acc.)
and Exact Match (EM) for classification tasks and generation tasks,
respectively.

#### Main Results.

Table 3: Main results on the generation tasks.

Method WebQs GeoQ. NL2B. MTOP SMCal. Avg. Unsupervised Random 3.79 25.36 31.27
3.98 3.70 13.62 Top-K BM25 14.17 65.71 58.81 49.66 44.02 46.48 Top-K BERT
14.17 64.64 52.45 51.36 44.76 45.48 Top-K T5 11.52 59.64 46.37 38.12 39.05
38.94 2-7[2pt/3pt] MLSM 16.14 68.93 56.11 54.05 47.72 48.59 Supervised EPR
17.62 73.21 77.87 60.82 60.49 53.43 CEIL 17.08 70.71 53.66 63.40 56.30 52.23
2-7[2pt/3pt] TTF 10.97 47.86 28.69 31.59 13.36 26.49 TTF (T5) 12.11 63.57
48.73 43.27 40.52 41.64

We compare MLSM and TTF with existing unsupervised (using off-the-shelf models
directly) and supervised learning-based baselines on classification tasks
(Table 2) and generation tasks (Table 3). The results demonstrate that MLSM
consistently outperforms all unsupervised baselines in most cases, achieving
an average improvement of 1.42% over the best baseline, Top-K BERT (semantic
similarity), on classification tasks, and an average improvement of 2.11% over
the best baseline, Top-K BM25 (surface similarity), on generation tasks. This
indicates that while different similarities excel at different tasks, MLSM can
adaptively integrate multi-level similarities for various tasks by updating
the aggregation weight of each expert for each test case, thus providing
evidence for $\mathcal{F}_{1}$. Additionally, learning-based methods like EPR
and CEIL show a clear advantage over MLSM, especially on classification tasks.
However, our TTF surpasses both learning-based methods, achieving over 5%
absolute improvements without costly integration with LLMs. This suggests that
fine-tuning on the target task is a more effective way than proxy tasks to
learn input-output mapping, further validating $\mathcal{F}_{2}$.
Nevertheless, TTF underperforms on generation tasks. We speculate this is
primarily because fine-tuning in generation tasks typically requires
substantial data, especially when the decoder initialized by BERT lacks
sufficient training due to the random initialization of external cross-
attention models. When transitioning from BERT to T5, TTF still obtains
further improvements over Top-K T5, thereby providing additional validation
for our speculation regarding the failure of TTF in generation tasks.

#### Transfer across Tasks for MLSM.

We compare EPR and MLSM on cross-task experiments, where EPR is trained on a
source task and then transferred to a target task, as shown in Fig. 3 (Left).
The results show that EPR generally underperforms compared to MLSM,
particularly when transferring between classification and generation tasks.
This suggests that learning-based methods like EPR may overfit to task-
specific characteristics during proxy task training, making it challenging to
justify the high costs of data collection and the need for extensive labeled
data in the demonstration set. In contrast, MLSM proves to be a practical
solution for cross-task demands, as it only leverages information from the
test case to adapt to different tasks during LLM inference.

#### Ablation of Batchsize for MLSM.

Although we assume only one test case is available to learn the aggregation
weight $\mathbf{w}$ for different similarity levels, we conduct an ablation
study to investigate the effect of increasing the batch size in Fig. 3
(Right). The results indicate that MLSM generally benefits from a larger batch
size, especially on classification tasks, showing over 4% average improvements
when the batchsize equals to 8. This improvement can be attributed to the fact
that different test cases in the same batch share a similar pattern of multi-
level analogs (i.e., similar $\mathbf{w}$). These findings further demonstrate
that multi-level analogs are versatile characteristics for selecting good
demonstration exemplars.

Figure 3: Left: Comparison of transferability between EPR and MLSM. We show
the absolute improvement of MLSM over EPR. Right: Comparisons of different
batchsize for MLSM.

#### Transfer across LLMs.

Table 4: Results of cross-LLM Transferability Validation. We show the absolute
improvement of TTF and MLSM over Top-K BERT.

TTF MLSM LLM SST-5 MRPC QNLI CMSQA Avg. SST-5 MRPC GeoQ. NL2Bash Avg GPT-2 XL
(1.5B) 3.54 0.00 5.35 6.38 3.82 1.54 0.00 1.07 4.57 1.74 GPT NEO (2.7B) 9.45
5.39 24.14 11.96 12.73 0.05 0.75 4.29 3.66 2.29 text-davinci-002 3.27 1.51
18.52 1.15 6.11 1.82 1.47 3.21 3.02 2.38

We validate the versatility of our proposed TTF and MLSM methods on GPT-2 XL
(1.5B), GPT-NEO (2.7B), and text-davinci-002, comparing the results with Top-K
BERT in Table 4. The results indicate that both TTF and MLSM can enhance ICL
performance across different large language models and datasets. TTF
consistently outperforms MLSM, demonstrating the effectiveness of acquiring
task-specific output similarity between exemplars and test cases. However, TTF
exhibits high variance in performance across different LLMs, whereas MLSM
provides more stable enhancements. This suggests that different LLMs have
varying abilities to exploit exemplars with similar outputs to the test case.
Additionally, the better performance of TTF on GPT-NEO compared to text-
davinci-002 suggests that the latter’s stronger inference ability makes it
more resilient to prompt choices.

#### Limitations.

Due to page limitations, we discuss two main limitations of our work in
Appendix B: 1) Why not combine MLSM and TTF? and 2) Are there better methods
to achieve $\mathcal{F}_{2}$? We recommend reading this section for detailed
information.

##  6 Conclusion

In this work, we delve into the inner-working mechanism of learning-based
demonstration exemplar selection methods. Through a series of experiments and
analyses, we speculate the advantages of these methods possibly stem from
their ability to integrate diverse linguistic similarities for exemplar
selection ($\mathcal{F}_{1}$) and their capacity to choose exemplars with
similar outputs to the test case ($\mathcal{F}_{2}$). Drawing insight from
both findings, we introduce two effective but simple exemplar section methods,
MLSM and TTF, tailored to address task-agnostic and task-specific demands
without requiring costly interactions with LLMs. The effectiveness of both
methods, as well as quantitative validations, provide substantial evidence in
favor of these assumptions. Our work not only provides insights into the
efficient and effective utilization of LLMs in practical applications but also
contributes empirical evidence to existing explanation of ICL, which we
discuss in Appendix C. While we acknowledge limitations of our work, we leave
it for further exploration.

## References

  * An et al. [2023] S. An, B. Zhou, Z. Lin, Q. Fu, B. Chen, N. Zheng, W. Chen, and J. Lou.  Skill-based few-shot selection for in-context learning.  In _EMNLP_ , pages 13472–13492. Association for Computational Linguistics, 2023.
  * Andreas et al. [2020] J. Andreas, J. Bufe, D. Burkett, C. Chen Jr, J. Clausman, J. Crawford, K. Crim, J. DeLoach, L. Dorner, J. Eisner, et al.  Task-oriented dialogue as dataflow synthesis.  _Transactions of the Association for Computational Linguistics_ , 8:556–571, 2020.
  * Barber and Bishop [1997] D. Barber and C. M. Bishop.  Ensemble learning for multi-layer networks.  In _NIPS_ , pages 395–401. The MIT Press, 1997.
  * Berant et al. [2013] J. Berant, A. Chou, R. Frostig, and P. Liang.  Semantic parsing on Freebase from question-answer pairs.  In _Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing_ , pages 1533–1544, Seattle, Washington, USA, Oct. 2013. Association for Computational Linguistics.  URL https://www.aclweb.org/anthology/D13-1160.
  * Black et al. [2021] S. Black, L. Gao, P. Wang, C. Leahy, and S. Biderman.  GPT-Neo: Large Scale Autoregressive Language Modeling with Mesh-Tensorflow, Mar. 2021.  URL https://doi.org/10.5281/zenodo.5297715.
  * Brown et al. [2020] T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei.  Language models are few-shot learners.  In _NeurIPS_ , 2020.
  * Devlin et al. [2019] J. Devlin, M. Chang, K. Lee, and K. Toutanova.  BERT: pre-training of deep bidirectional transformers for language understanding.  In _NAACL-HLT (1)_ , pages 4171–4186. Association for Computational Linguistics, 2019.
  * Dolan et al. [2004] W. B. Dolan, C. Quirk, and C. Brockett.  Unsupervised construction of large paraphrase corpora: Exploiting massively parallel news sources.  In _COLING 2004: Proceedings of the 20th International Conference on Computational Linguistics_ , pages 350–356, 2004.
  * Gao et al. [2021] T. Gao, A. Fisch, and D. Chen.  Making pre-trained language models better few-shot learners.  In _ACL/IJCNLP (1)_ , pages 3816–3830. Association for Computational Linguistics, 2021.
  * Halawi et al. [2023] D. Halawi, J.-S. Denain, and J. Steinhardt.  Overthinking the truth: Understanding how language models process false demonstrations.  _arXiv preprint arXiv:2307.09476_ , 2023.
  * Heilbron et al. [2015] F. C. Heilbron, V. Escorcia, B. Ghanem, and J. C. Niebles.  Activitynet: A large-scale video benchmark for human activity understanding.  In _2015 IEEE conference on computer vision and pattern recognition (CVPR)_ , pages 961–970. IEEE, 2015.
  * Iwasawa and Matsuo [2021] Y. Iwasawa and Y. Matsuo.  Test-time classifier adjustment module for model-agnostic domain generalization.  In _NeurIPS_ , pages 2427–2440, 2021.
  * Jawahar et al. [2019] G. Jawahar, B. Sagot, and D. Seddah.  What does BERT learn about the structure of language?  In _ACL (1)_ , pages 3651–3657. Association for Computational Linguistics, 2019.
  * Kornblith et al. [2019] S. Kornblith, M. Norouzi, H. Lee, and G. E. Hinton.  Similarity of neural network representations revisited.  In _ICML_ , volume 97 of _Proceedings of Machine Learning Research_ , pages 3519–3529. PMLR, 2019.
  * Kossen et al. [2023] J. Kossen, T. Rainforth, and Y. Gal.  In-context learning in large language models learns label relationships but is not conventional learning.  _CoRR_ , abs/2307.12375, 2023.
  * Kulesza and Taskar [2011] A. Kulesza and B. Taskar.  k-dpps: Fixed-size determinantal point processes.  In _ICML_ , 2011.
  * Li et al. [2021] H. Li, A. Arora, S. Chen, A. Gupta, S. Gupta, and Y. Mehdad.  Mtop: A comprehensive multilingual task-oriented semantic parsing benchmark.  In _Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume_ , pages 2950–2962, 2021.
  * Li and Qiu [2023] X. Li and X. Qiu.  Finding support examples for in-context learning.  In _EMNLP (Findings)_ , pages 6219–6235. Association for Computational Linguistics, 2023.
  * Li et al. [2023] X. Li, K. Lv, H. Yan, T. Lin, W. Zhu, Y. Ni, G. Xie, X. Wang, and X. Qiu.  Unified demonstration retriever for in-context learning.  In _Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 4644–4668, Toronto, Canada, July 2023. Association for Computational Linguistics.  doi: 10.18653/v1/2023.acl-long.256.
  * Lin et al. [2018] X. V. Lin, C. Wang, L. Zettlemoyer, and M. D. Ernst.  Nl2bash: A corpus and semantic parser for natural language interface to the linux operating system.  In _LREC_. European Language Resources Association (ELRA), 2018.
  * Liu et al. [2022] J. Liu, D. Shen, Y. Zhang, B. Dolan, L. Carin, and W. Chen.  What makes good in-context examples for gpt-3?  In _DeeLIO@ACL_ , pages 100–114. Association for Computational Linguistics, 2022.
  * Lu et al. [2022] Y. Lu, M. Bartolo, A. Moore, S. Riedel, and P. Stenetorp.  Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity.  In _ACL (1)_ , pages 8086–8098. Association for Computational Linguistics, 2022.
  * Ma et al. [2019] X. Ma, Z. Wang, P. Ng, R. Nallapati, and B. Xiang.  Universal text representation from BERT: an empirical study.  _CoRR_ , abs/1910.07973, 2019.
  * Milios et al. [2023] A. Milios, S. Reddy, and D. Bahdanau.  In-context learning for text classification with many labels.  _CoRR_ , abs/2309.10954, 2023.
  * Muennighoff [2022] N. Muennighoff.  SGPT: GPT sentence embeddings for semantic search.  _CoRR_ , abs/2202.08904, 2022.
  * Ni et al. [2022] J. Ni, G. H. Ábrego, N. Constant, J. Ma, K. B. Hall, D. Cer, and Y. Yang.  Sentence-t5: Scalable sentence encoders from pre-trained text-to-text models.  In _ACL (Findings)_ , pages 1864–1874. Association for Computational Linguistics, 2022.
  * Olsson et al. [2022] C. Olsson, N. Elhage, N. Nanda, N. Joseph, N. DasSarma, T. Henighan, B. Mann, A. Askell, Y. Bai, A. Chen, T. Conerly, D. Drain, D. Ganguli, Z. Hatfield-Dodds, D. Hernandez, S. Johnston, A. Jones, J. Kernion, L. Lovitt, K. Ndousse, D. Amodei, T. Brown, J. Clark, J. Kaplan, S. McCandlish, and C. Olah.  In-context learning and induction heads.  _CoRR_ , abs/2209.11895, 2022.
  * Polikar [2009] R. Polikar.  Ensemble learning.  _Scholarpedia_ , 4(1):2776, 2009.
  * Radford et al. [2019] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever.  Language models are unsupervised multitask learners.  2019\.
  * Reddy [2023] G. Reddy.  The mechanistic basis of data dependence and abrupt learning in an in-context classification task.  _CoRR_ , abs/2312.03002, 2023.
  * Robertson and Zaragoza [2009] S. Robertson and H. Zaragoza.  The probabilistic relevance framework: Bm25 and beyond.  _Foundations and Trends in Information Retrieval_ , 3:333–389, 01 2009.  doi: 10.1561/1500000019.
  * Rohrbach et al. [2017] A. Rohrbach, A. Torabi, M. Rohrbach, N. Tandon, C. Pal, H. Larochelle, A. Courville, and B. Schiele.  Movie description.  _International Journal of Computer Vision_ , 123:94–120, 2017.
  * Rubin et al. [2022] O. Rubin, J. Herzig, and J. Berant.  Learning to retrieve prompts for in-context learning.  In _NAACL-HLT_ , pages 2655–2671. Association for Computational Linguistics, 2022.
  * Socher et al. [2013] R. Socher, A. Perelygin, J. Wu, J. Chuang, C. D. Manning, A. Y. Ng, and C. Potts.  Recursive deep models for semantic compositionality over a sentiment treebank.  In _EMNLP_ , pages 1631–1642. ACL, 2013.
  * Su et al. [2023] H. Su, J. Kasai, C. H. Wu, W. Shi, T. Wang, J. Xin, R. Zhang, M. Ostendorf, L. Zettlemoyer, N. A. Smith, and T. Yu.  Selective annotation makes language models better few-shot learners.  In _ICLR_. OpenReview.net, 2023.
  * Talmor et al. [2019] A. Talmor, J. Herzig, N. Lourie, and J. Berant.  CommonsenseQA: A question answering challenge targeting commonsense knowledge.  In _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)_ , pages 4149–4158, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.  doi: 10.18653/v1/N19-1421.  URL https://aclanthology.org/N19-1421.
  * Tonglet et al. [2023] J. Tonglet, M. Reusens, P. Borchert, and B. Baesens.  SEER : A knapsack approach to exemplar selection for in-context hybridqa.  In _EMNLP_ , pages 13569–13583. Association for Computational Linguistics, 2023.
  * Voronov et al. [2024] A. Voronov, L. Wolf, and M. Ryabinin.  Mind your format: Towards consistent evaluation of in-context learning improvements.  _CoRR_ , abs/2401.06766, 2024.
  * VS et al. [2023] V. VS, P. Oza, and V. M. Patel.  Towards online domain adaptive object detection.  In _Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision_ , pages 478–488, 2023.
  * Wang et al. [2018] A. Wang, A. Singh, J. Michael, F. Hill, O. Levy, and S. Bowman.  Glue: A multi-task benchmark and analysis platform for natural language understanding.  In _Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_ , pages 353–355, 2018.
  * Williams et al. [2017] A. Williams, N. Nangia, and S. R. Bowman.  A broad-coverage challenge corpus for sentence understanding through inference.  _arXiv preprint arXiv:1704.05426_ , 2017.
  * Wolfson et al. [2020] T. Wolfson, M. Geva, A. Gupta, M. Gardner, Y. Goldberg, D. Deutch, and J. Berant.  Break it down: A question understanding benchmark.  _Transactions of the Association for Computational Linguistics_ , 8:183–198, 2020.
  * Yan et al. [2023] J. Yan, J. Xu, C. Song, C. Wu, Y. Li, and Y. Zhang.  Understanding in-context learning from repetitions.  _CoRR_ , abs/2310.00297, 2023.
  * Ye et al. [2023] J. Ye, Z. Wu, J. Feng, T. Yu, and L. Kong.  Compositional exemplars for in-context learning.  In _ICML_ , volume 202 of _Proceedings of Machine Learning Research_ , pages 39818–39833. PMLR, 2023.
  * Zelle and Mooney [1996] J. M. Zelle and R. J. Mooney.  Learning to parse database queries using inductive logic programming.  In _AAAI/IAAI_ , pages 1050–1055, Portland, OR, August 1996. AAAI Press/MIT Press.  URL http://www.cs.utexas.edu/users/ai-lab?zelle:aaai96.
  * Zellers et al. [2019] R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi.  Hellaswag: Can a machine really finish your sentence?  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , 2019.
  * Zhang et al. [2022a] M. Zhang, S. Levine, and C. Finn.  Memo: Test time robustness via adaptation and augmentation.  _Advances in Neural Information Processing Systems_ , 35:38629–38642, 2022a.
  * Zhang et al. [2023a] S. Zhang, X. Xia, Z. Wang, L. Chen, J. Liu, Q. Wu, and T. Liu.  IDEAL: influence-driven selective annotations empower in-context learners in large language models.  _CoRR_ , abs/2310.10873, 2023a.
  * Zhang et al. [2022b] Y. Zhang, B. Hooi, L. Hong, and J. Feng.  Self-supervised aggregation of diverse experts for test-agnostic long-tailed recognition.  In _NeurIPS_ , 2022b.
  * Zhang et al. [2023b] Y. Zhang, X. Wang, K. Jin, K. Yuan, Z. Zhang, L. Wang, R. Jin, and T. Tan.  Adanpc: Exploring non-parametric classifier for test-time adaptation.  In _ICML_ , volume 202 of _Proceedings of Machine Learning Research_ , pages 41647–41676. PMLR, 2023b.
  * Zhao et al. [2021] Z. Zhao, E. Wallace, S. Feng, D. Klein, and S. Singh.  Calibrate before use: Improving few-shot performance of language models.  In _ICML_ , volume 139 of _Proceedings of Machine Learning Research_ , pages 12697–12706. PMLR, 2021.
  * Zhou et al. [2023] H. Zhou, X. Wan, L. Proleev, D. Mincu, J. Chen, K. A. Heller, and S. Roy.  Batch calibration: Rethinking calibration for in-context learning and prompt engineering.  _CoRR_ , abs/2309.17249, 2023.

##  Appendix A Experimental Setup

Table 5: Datasets with corresponding prompts and examples used in the
experiments.

###  A.1 Datasets

Following existing work [44], we conduct experiments on five classification
and and five generation tasks444We exclude MNLI [41] to reduce computation
cost and Break [42] because of failure to reproduce its evaluation method.
While we advise readers to refer to the detail of each dataset in the original
work [44], we provide the prompts and examples for each dataset in Table 5 and
offer a detailed description of each dataset below for completeness.

#### SST-5 [34]

is a sentiment classification benchmark containing five fine-grained classes
including ‘very positive’, ‘positive’ ‘neutral’, ‘negative’, and ‘very
negative’.

#### MRPC [8]

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

#### MNLI [41]

is a crowdsourced collection of sentence pairs with textual entailment
annotations. Given a premise sentence and a hypothesis sentence, the task is
to predict whether the premise entails the hypothesis (entailment),
contradicts the hypothesis (contradiction), or neither (neutral).

#### QNLI [40]

is a question-answering dataset consisting of question-paragraph pairs, and
the task is to determine whether the context sentence contains the answer to
the question.

#### CMSQA [36]

(short for CommonsenseQA) is a multiple-choice question-answering dataset that
requires different types of commonsense knowledge. The task is to predict the
correct answer out of five provided candidate answers.

#### HellaSwag [46]

is a large-scale dataset of grounded commonsense reasoning. There are four
candidate answers for each question: a video caption from ActivityNet Captions
[11] and the Large Scale Movie Description Challenge [32]. The three incorrect
answers are adversarially generated and human validated to deceive machines.
The correct answer is the actual video caption for the subsequent occurrence
in the video.

#### WebQs [4]

is question-answer pairs obtained from the web. The questions are selected
using Google Suggest API, and the answers are entities in Freebase.

#### Nl2Bash [20]

is a dataset for the problem of mapping English sentences to Bash commands.
The corpus consists of text–command pairs, where each pair consists of a Bash
command scraped from the web and an expert-generated natural language
description.

#### GeoQuery [45]

contains a parallel corpus of 880 English questions about US geography paired
with Prolog queries.

#### Break [42]

is a dataset that maps complex natural language questions into a language-
based meaning representation. The question is decomposed into an ordered list
of atomic steps, which is used as the target sequence. We use the low-level
Break subset following [33].

#### MTOP [17]

is a multilingual task-oriented semantic parsing dataset covering 6 languages
and 11 domains. The target commands are complex queries featuring nested
intent-slot prediction. Similar to past work [33], we use the English subset
of MTOP.

#### SMCalFlow [2]

is a large dialogue dataset, featuring natural conversations about tasks
involving calendars, weather, places, and people. The meaning representation
is an executable dataflow program featuring API calls, function composition,
and complex constraints.

###  A.2 Implementation Details

We employ the implementation555https://github.com/HKUNLP/icl-ceil from Ye et
al. [44] for all baselines. Specifically, for EPR and CEIL, we limit the
maximum instances in the proxy task to 4,000 ($|\mathcal{D}^{s}|=4,000$) and
sample 50 candidates for each instance to create positive and negative pairs.
It is worth noting it is pretty expensive and time-consuming to collect these
data for both methods (i.e., 200,000 queries to LLMs), especially for CEIL
where each candidate sequence involves 16 exemplars.

For our proposed MLSM, we randomly sample 1,000 examples ($n_{c}=1,000$) from
the demonstration set $\mathcal{D}$ to compute layer-wise CKA scores and
obtain three representative layers through clustering ($n_{l}=3$). Then for
each test case, we randomly sample 256 and 64 examples ($n_{t}=256$ and
$n_{v}=64$) form $\mathcal{D}$ as mini training and validation set,
respectively. The temperature of softmax function is set to 0.01
($\tau=0.01$). We utilize Adam optimizer with batch size 32 and learning rate
0.1 to learn the aggregation weight $\mathbf{w}$ in fewer epochs.

For our proposed TTF, we instantiate $f_{\bm{\theta}}$ with BERT and
$q_{\bm{\phi}}$ with different task heads for classification tasks.
Concretely, for SST-5, MRPC and QNLI, we utilize the sequential classification
head666https://huggingface.co/docs/transformers/model_doc/bert#transformers.BertForSequenceClassification
and train the model using Adam optimizer with batchsize 32, learning rate 5e-4
and weight decay 1e-4. For SWAG and CMSQA, we adopt the multi-choice
head777https://huggingface.co/docs/transformers/model_doc/bert#transformers.BertForMultipleChoice
and also train the model using Adam optimizer with batchsize 8, learning rate
5e-4 and weight decay 1e-4. Additionally, we instantiate $f_{\bm{\theta}}$
with BERT and $q_{\bm{\phi}}$ with decoder for generation tasks. For all
tasks, we utilize Adam optimizer with batchsize 16, learning rate 5e-5.
Especially, for TTF (T5), we instantiate $f_{\bm{\theta}}$ and $q_{\bm{\phi}}$
using the encoder and decoder of T5888https://huggingface.co/google-t5/t5-base
and utilize to Adam optimizer with batchsize 8, learning rate 4e-5 and weight
decay 0.01.

Furthermore, we conducted all experiments for EPR and CEIL on two NVIDIA A100
GPUs (40GB), while the remaining experiments were performed on two NVIDIA V100
GPUs (30GB). Each experiment is repeated three times using different random
seeds to mitigate the effects of randomness.

##  Appendix B Limitations

The section discusses two main limitations of our work.

#### Why not combine MLSM and TTF?

MSLM and TTF are designed for distinct scenarios: MSLM addresses cross-task
situations by maximizing agreement across different similarities, while TTF
targets task-specific situations by fine-tuning the retriever with labeled
data to learn task-specific information. However, it is worth exploring
whether combining these methods could enhance performance in task-specific
situations. To investigate this, we utilize the BERT model trained in TTF to
replace the original BERT in MSLM and conduct experiments on five
classification tasks using the same implementation detailed in Appendix A.2.
As shown in Table 6, although the combined approach of TTF and MSLM
significantly outperforms MSLM with an average improvement of around 4%, it
falls short of TTF by over 6%. This performance drop suggests that input-
output-based similarity is superior to similarities from other layers. While
TTF’s final layer effectively captures task-specific input-output-based
similarity, integrating it with other sub-optimal similarities introduces
noise, negatively impacting exemplar selection for ICL.

#### Any better method to achieve $\mathcal{F}_{2}$?

We empirically find that learning-based methods acquire output similarity
between exemplars and test cases, significantly enhancing task-specific
performance in demonstration selection. However, when we simulate this
similarity mechanism using TTF, it proves ineffective for generation tasks due
to the difficulty in pinpointing which model components contain effective
input-output relationships and the extensive data required for fine-tuning
generation task heads. In a further attempt, akin to EPR, we tried selecting
the exemplar with the most similar output to the test case as a positive pair
and the most dissimilar one as a negative pair. We then fine-tuned BERT on
this task, similar to EPR. However, this approach led to a performance
collapse. We suggest this failure is primarily due to the complexity of
accurately modeling the nuanced input-output relationships necessary for
generation tasks and the inherent limitations of fine-tuning process (e.g.,
the lack of comprehensive hyperparameter and framework searching ).

##  Appendix C Connection with Explanatory Work of ICL

#### Qualitative Validation of $\mathcal{F}_{1}$.

While learning-based prompt retrieval methods are capable of retrieving
exemplars with multi-level analogs to the test case, we demonstrate that such
exemplar will help LLMs to make correct predictions at a higher possibility
than dissimilar ones when all of them include correct patterns pertinent to
the test case. Previous investigation [27, 30] proposed a possible inner
working of ICL that LLMs can learn from surface patterns in the demonstration
sequence such as copying tokens from contextual prompts. Furthermore, recent
research [43] empirically demonstrated that the connection between two tokens
will be strengthened by increased contextual co-occurrences during generation
caused by the maximizing likelihood objective of LLMs. These insights suggest
that effective demonstration examples may manifest more token or phrase-level
correspondence with the test case corresponding to anchors at lower and middle
layers, exerting a more significant influence on LLMs’ output, further
supporting $\mathcal{F}_{1}$.

#### Qualitative Validation of $\mathcal{F}_{2}$.

In line with the qualitative validation of $\mathcal{F}_{1}$, we illustrate
the exemplar with similar output to test cases contributes to the performance
of ICL. Prior work have demonstrated ICL typically learn input-output relation
from exemplars even for a genuinely novel task the LLM cannot know from pre-
training [15, 10]. Moreover, Kossen et al. [15] further proposed LLMs exhibit
a preference for utilizing information closer to the query rather than
treating all available information equally. Hence, if the exemplar selection
method successfully learns the output similarity via the proxy task, it not
only selects demonstration examples exhibiting useful input-output
correlations for the test case due to their shared relevant input-output
correlations but also positions it closely to the test query in the prompt.
These advantages align with the previously mentioned underlying working
mechanisms of LLMs, thereby validating $\mathcal{F}_{2}$.

Table 6: Experimental results for the combination of MLSM and TTF

Method SST-5 MRPC QNLI CMSQA SWAG Avg. MLSM 33.15 69.87 65.02 37.26 41.49
49.36 TTF 42.14 74.51 85.08 47.83 55.72 61.06 TTF and MLSM 36.14 71.07 65.31
45.61 50.27 53.69

Figure 4: Left: Comparison of similarity between the input/output of positive
and negative demonstration examples and the input/output of the test case
across ten tasks for EPR. Right: Difference in average similarity between the
output of test case and retrieved example for EPR and three learning-free
prompt retrieval methods. We use GPT-2 XL [5] as the LLM.

## NeurIPS Paper Checklist

  1. 1.

Claims

  2. Question: Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

  3. Answer: [Yes]

  4. Justification: In this work, we delve into the inner-working mechanism of learning-based

  5. 2.

Limitations

  6. Question: Does the paper discuss the limitations of the work performed by the authors?

  7. Answer: [Yes]

  8. Justification: We discuss two main limitations of our work in Appendix B.

  9. 3.

Theory Assumptions and Proofs

  10. Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

  11. Answer: [N/A]

  12. Justification: Our paper does not include theoretical results.

  13. 4.

Experimental Result Reproducibility

  14. Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

  15. Answer: [Yes]

  16. Justification: The implementation detail of our work is discussed in Appendix A.2.

  17. 5.

Open access to data and code

  18. Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

  19. Answer: [Yes]

  20. Justification: We submit our codes in supplemental material with instructions. We will clean up the code after paper accepted.

  21. 6.

Experimental Setting/Details

  22. Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

  23. Answer: [Yes]

  24. Justification: The implementation detail of our work is discussed in Appendix A.2.

  25. 7.

Experiment Statistical Significance

  26. Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

  27. Answer: [No]

  28. Justification: We reported the averaging results of three runs with different random seeds for all experiments.

  29. 8.

Experiments Compute Resources

  30. Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

  31. Answer: [Yes]

  32. Justification: We discuses this point in Appendix A.2.

  33. 9.

Code Of Ethics

  34. Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?

  35. Answer: [Yes]

  36. Justification: This work conform with the NeurIPS Code of Ethics.

  37. 10.

Broader Impacts

  38. Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

  39. Answer: [N/A]

  40. Justification: Our work aims to unravel the mechanics of learning-Based demonstration Selection to improve the efficiency of ICL. As such, our work will not produce potential malicious impacts to our society.

  41. 11.

Safeguards

  42. Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

  43. Answer: [N/A]

  44. Justification: Our paper poses no such risks.

  45. 12.

Licenses for existing assets

  46. Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

  47. Answer: [Yes]

  48. Justification: We have cited the original paper that produced the code package or dataset.

  49. 13.

New Assets

  50. Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

  51. Answer: [Yes]

  52. Justification: We follow the requirements of NIPS.

  53. 14.

Crowdsourcing and Research with Human Subjects

  54. Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

  55. Answer: [N/A]

  56. Justification: Our paper does not involve crowdsourcing nor research with human subjects.

  57. 15.

Institutional Review Board (IRB) Approvals or Equivalent for Research with
Human Subjects

  58. Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

  59. Answer: [N/A]

  60. Justification: Our paper does not involve crowdsourcing or research with human subjects.