Rapid Selection and Ordering of In-Context Demonstrations via Prompt Embedding Clustering
Authors: Kha Pham
Hung Le
Man Ngo
Truyen Tran
Venue: ICLR 2025 Poster (2025)
Keywords: in-context learning
order sensitivity
LLMs
clustering
cluster-based search
positional encoding
attention mask
serial-position effect
cluster-based search
Primary area: generative models
TL;DR: We accelerate selection and ordering of in-context demonstrations in self-adaptive ICL settings by leveraging our newfound clustering property in prompt embedding spaces.

Abstract: While Large Language Models (LLMs) excel at in-context learning (ICL) using just a few demonstrations, their performances are sensitive to demonstration orders. The reasons behind this sensitivity remain poorly understood. In this paper, we investigate the prompt embedding space to bridge the gap between the order sensitivity of ICL with inner workings of decoder-only LLMs, uncovering the clustering property: prompts sharing the first and last demonstrations have closer embeddings, with first-demonstration clustering usually being stronger in practice. We explain this property through extensive theoretical analyses and empirical evidences. Our finding suggests that the positional encoding and the causal attention mask are key contributors to the clustering phenomenon. Leveraging this clustering insight, we introduce Cluster-based Search, a novel method that accelerates the selection and ordering of demonstrations in self-adaptive ICL settings. Our approach substantially decreases the time complexity from factorial to quadratic, saving 92% to nearly 100% execution time while maintaining comparable performance to exhaustive search.

Published as a conference paper at ICLR 2025
RAPID SELECTION AND ORDERING OF IN-CONTEXT DEMONSTRATIONS VIA PROMPT EMBEDDING CLUSTERING
Kha Pham- Hung Le
Man Ngo² Truyen Tran
Applied Artificial Intelligence Institute, Deakin University
2 Faculty of Data Science in Business, Ho Chi Minh University of Banking, Vietnam
1 (phti, thai .e, truyen. tran eereakin.edd.... 2 mannm@hub. edu.vn vn
ABSTRACT
While Large Language Models (LLMs) excel at in-context learning (ICL) using just few demonstrations, their performances are sensitive to demonstration or- ders. The reasons behind this sensitivity remain poorly understood. In this paper, we investigate the prompt embedding space to bridge the gap between the order sensitivity of ICL with inner workings of decoder-only LLMs, uncovering the clustering property: prompts sharing the first and last demonstrations have closer embeddings, with fist-demonstrattin clustering usually being stronger in practice. We explain this property through extensive theoretical analyses and empirical evi- dences. Our finding suggests that the positional encoding and the causal attention mask are key contributors to the clustering phenomenon. Leveraging this clustering insight, we introduce Cluster-based Search, a novel method that accelerates the se- lection and ordering of demonstrations in adlf-tdeptive ICL settings. Our approach substantially decreases the time complexity from factorial to quadratic, saving 92% to nearly 100% execution time while maintaining comparable performance cxhaustive search.
INTRODUCTION
In-context Learning (ICL) is a remarkable emergent ability of Large Language Models (LLMs) to perform few-shot learning. This means they answer queries based just handful of related demonstrations, even without explicit training for the specific task (Brown et al., 2020). While highly effective, research has shown ICL performances can be extremely sensitive to demonstration orders (Zhao et al., 2021; Liu et al., 2021; Lu et al., 2022; Liu et al., 2023). To the best of our knowledge, the reasons behind this sensitivity remain unclear.
In this paper, we investigate the prompt embedding space to explain the order sensitivity of decoder- only LLMs in ICL. Specifically, we focus on the last token embedding in the final LLM layer. which is the input for the next-token prediction head. Our investigation reveals the clustering property: prompts with the same first and last demonstrations form clusters the embedding space (Figure 1), with drrsssmmmottitioon clustering usually being stronger in practice. Our finding aligns with recent research on the special roles of the first and last demonstrations ICL (Liu et al., 2023; Janik, 2023). But unlike prior works, our clustering property links the order sensitivity of ICL and inner workings of LLMs. allowing more systematic insights into LLMs' behaviors.
We provide extensive analyses confirm the clustering property. In particular, visualize prompt embeddings in 2D spaces using UMAP, run K-Means clustering on high-dimensional embedding spaces, and quantify the importance of input tokens by their partial derivative norms. Experimental results consistently support the existence of clusters.
Further, we provide theoretical and empirical explanations for the clustering behaviors. Theoretically, we prove that demonssmmsttiiin clustering is associated with causal attention mask under ideal conditions. To validate this relationship in practice, we train decoder-only Transformers from scratch

---

Published as a conference paper at ICLR 2025
Tolk: Reverse 1LM GPT12
Trele Repeat LLM: Limio 50B
Taak: Sentimeat- Symbuilic LLM Linua 12-13B
Figure 1: Prompts are clustered by first and last demonstrations. We consider prompts with the same query and demonstrations but different orders. For each prompt, we show its UMAP 2D projection of the last embedding in the final LLM layer. Points with the same colors have the same first demonstrations, and points with the same shapes have the same last demonstrations. Interestingly, points with similar colors shapes form cluster. This property holds for various LLMs and ICI. tasks.
on WikiText2 dataset (Merity et al., 2017) with and without causal attention mask, achieving similar outcomes: causal attention mask contributes the importance of beginning tokens. We also perform similar experiments with and without positional encoding, and find that the interplay of the causal structure and the positional encoding contributes to the importance of ending tokens.
Finally, we leverage the clustering property to accelerate self adaptive approaches (Wu et al., 2022) when dealing with the costly problem of demonstration selection and ordering. Our proposed method, dubbed Cluster-based Search, substantially decreases the time complexity from the factorial order of exhaustive search to quadratie order. We apply Cluster-based Search two selection scenarios: one with the ideal selection criterion and another with the practical entropy-basee criterion (Lu et al., 2022). In both cases, our proposed method achieves competitive accuracies compared to exhaustive search while being significantly faster saving 92% nearly 100% execution time.
In short, our contributions include: (1) Establishing the clustering property by first & last demonstra- tions and verifying it from various angles, (2) analyzing its underlying mechanisms both theoretically and empirically, and (3) leveraging it to propose Cluster-based Search, which significantly accelerates the selection and ordering of in-context demonstrations in self-adaptive ICL settings.
PRELIMINARIES
2.1 TERMINOLOGIES AND NOTATIONS
Transformer We consider decoder-only LLMs built upon the Transformer architecture. The Transformer takes the input as a sequence of tokens (o,(0)) E (RT)) This input sequence is transformed through T Transformer layers, of which the (-th layer (1 1, T) receives the sequence {a; 1)}2
(Rd)) " as inputs and outputs another sequence (a,(t)}[ E (Rd)". We primarily focus on the last embeddings of Transformer layers, i.e. In (1), which we will denote as x- (1) for < T and _- when ( T. We denote (Ii, w5) as the inner product of vectors æ; and zz--
In-context learning (ICL) An ICL prompt consists of ordered sequence of k demonstrations E F= (ei,.... ex) and a query Each demonstration e; includes an input e.th (e.g. a sentence) and its associated output eşet (e.g. the sentence's sentiment). The query q only consists of an input qdh and the LLM will predict its associated output, i.e. q LLM((E F, q) The prediction qour is compared with the ground-truth output qutl to determine whether the LLM answers the query correctly. In Section and Section 4, we investigate LLMs' behaviors on ICL prompts built upon the same query and the same set of demonstrations but different orders.
2.2 TASKS
We consider two types of tasks: classification and reasoning.
For text classification, consider tasks of sentiment classification and language identification. We use data from SST-2 (Socher et al., 2013) dataset for sentiment classification, where each demon- stration is pair of natural language sentence and its associated sentiment. Prompts in language identification task are similar, except that each demonstration a pair a sentence and its associated language. Dataset for language identification is taken from HuggingFace (HuggingFace, 2021). We
2

---

Published as a conference paper at ICLR 2025
also consider tasks of symbolic sentiment classification and symbolic language identification. The labels of demonstrations in symbolic sentiment classification are not "positive" or "negative", but replaced by random strings. For example, "positive" becomes "a b c" and "negative" becomes "x y z", so the model should predict "a c" or "x y z instead of "positive" or "negative". Similarly for symbolie language identification.
For reasoning tasks, we consider symbolic reasoning, common-sense reasoning. and mathematical arithmetic reasoning. The symbolic reasoning tasks include reverse and repeat. For reverse, the demonstration input is a string of three random letters, e.g. "a b c", and the output is its reverse "c b a". For repeat, the imput contains possibly duplicated letters, e.g. "a a b c", and the output is the string without duplication, i.c. "a b c". This kind of letter-string logical tasks remains a challenge for current Al systems (Mitchell, 2021). For the common-senss reasoning task, we leverage question-answar pairs from the CommonsenseeQA dataset (Talmor et al., 2019), where each pair forms an ICL. demonstration. For the mathematical arithmetic task, we use questions and answers from the AddSub dataset (Hosseini et al., 2014).
In all tasks, the input and the output in a demonstration demonstration in the reverse task should be "a b c - c b a". Demonstrations are in turn separated by "An". Finally, the query consists of the input string and the arrow, e.g. y y
separated by an arrow "-"", e.g.
2.3 EXPERIMENTAL LLMs
We conduct experiments various decoder-only open-source L.Mss, including GPT-2 (Radford et al., 2019), GPT-Neo (Gao et al., 2020), LJama-v1 (Touvron et al., 2023a) & Ilama-v2 (Touvron al., 2023b), MPT (Team, 2023), Phi-2 (Javaheripi al., 2023), and Qwen-2.5 (Hui et al., 2024).
3 PROMPT CLUSTERING
In this section, we introduce and verify the clustering property: x-1's 's of prompts with the same first and last demonstrations tend to stay closer in the embedding space, with fist-seeoosttattion clustering showing notably stronger effects. This means the embedding space contains clusters, each of which is associated with prompts with the same first and last demonstrations, though the clustering by first demonstration is more pronounced. We will show that the clustering property consistently occurs across different LLMs and ICL tasks.
Formally, consider set of & demonstrations fei....,ey, exk } and query q. We consider A in this section. Let P be the permutation set of (ei,-..,ea, ex,,, hence [P] kt. Each permuta- tion p E P and the query q compose a prompt. Let 2 ddd be the last-layer embedding of the last input token of the prompt associated with permutation p. We verify the clustering property of the set {-==_ E P) with different approaches: oow-dimensional visualization (Sec- tion 3.1), K-Means clustering and quantifying token importance (Section 3.2). Experimental results consistently confirm the occurrence of clusters. We also show that clustering appears from early layers of LLMs (Section 3.3).
Senticest Classdtttion
1LM: GPT22
LLM: Llatta 12-13B
Figure 2: Histogram of pairwise distances of points in
{x" : p E P}. Both histograms contains more than one peak, indicating the existence of clusters.
3.1 LOW-DIMENSIONAL VISUALIZATIONS OF CLUSTERING
Figure shows the 2D UMAP (McInnes et al., 2018) projections of x=, 's with various LLMs and ICL tasks (similar results for t-SNE projection shown Appendix E). Projections with the same colors are of permutations sharing the first demonstrations. Similarly, projections with the same shapes are of permutations sharing the last demonstrations.
We observe consistent property across different cases: projections are clustered by colors and shapes, i.e. if p and p' share the first or last demonstrations, then the projections of x", and wh tend to stay closer. This also means when p and p' have different first and last demonstrations, x." and
3

---

Published as a conference paper at ICLR 2025
Task: Reverse
(a)
Tusk Repeat LLM: Liana 30B
Task:
(b)
Tiok Repeat
LLM: GPT-2
GPT-2
30B
(M0)
20
Language Identifcation LLM: CPT--No-2..7B
Clasification Llamarr2-13B
TTask: Language Ideuniecation LLM: GPP-Neo-22788
Tass:: Sentiment Classitication
L.I.M. Llama-v2-13B
50
25
(a) Percentage frequencies demonstration positions. The z-axis is the index of demonstration positions; y y-axis the pereentage frequency. We observe U-shapes of the percentage frequency curves across different tasks and LLMs. This indicates first and last demonstrations of elements in cluster tend to be the same. (b) Partial derivative norms of various LLMs on different tasks. The r-axis is the chunk index; g-axis is the chunk-averaged Frohenius norms of partial derivatives. We can observe consistent U-shapes across different LLMs and tasks, meaning that beginming tokens and ending tokens play important roles in prompt.
Figure 3:
r" are likely to stay in different clusters and thus far away from each other. Since x-1 is the input for the prediction head, this indicates the next-token prediction of p is likely differ from one of p'. This explains why LLMs' responses vary with demonstration orders. Interestingly, our analysis reveals that while both types of clustering are present, the clustering effect tends to be stronger for shared first demonstrations compared to shared last demonstrations.
3.2 CLUSTERING ON HIGH DIMENSIONAL SPACES
We verify the clustering property of x=1's on Rd. To do so, we employ two different methods: K-Means clustering, and quantifying token importance. Experimental results consistently support the clustering property, while also confirming the relatively stronger clustering effect for shared first demonstrations.
K-Means clustering First, we examine the existence of clusters in {- : p E PP based on the idea in (Steinbach et al., 2004): we plot the histograms of pairwise Euclidean distances between points in {-- p P). If the histogram has more than one peak, {x" : p P} is likely to contain clusters. Figure 2a shows the histograms of the reverse task on GPT-2 and the sentiment
classification task on Llama-v2-13B. Both histograms have more than one peak, meaning clusters exist in {a-= p E P).
Next, we investigate how fz" :p P} is clustered. We run K-Means clustering algorithm (Lloyd, 1982) with K = kik- 1) centers. each cluster, we compute the percentage frequency of the most frequent demonstration at each demonstration position. For example, with k = 5, each cluster is associated with a 5-dimensional array of percentage frequencies, e.g. (80%, 40% 30%, 30%, 55%). This means at most 80% of the permutations in the cluster share the same first demonstrations, at most 40% them share the same second demonstrations, and so on. We then average the percentage frequency arrays over all K clusters to come up with a k-dimensional array representing the current prompt. Finally, we take the average of those k-dimensional arrays over 100 prompts with different LLMs and tasks. Figure 3a shows the results. Overall, we observe U-shape percentage frequency curves across different tasks and LLMs, with notably higher percentages for the first position. This indicates both first and last demonstrations of permutations in cluster tend to be the same, with stronger effect for first demonstrations, which supports our clustering property.
Token importance The property of clustering by first and last demonstrations suggests beginning and ending tokens play important roles. We thus verify the clustering property by quantifying token importance. Specifically, we quantify how much x will change when each (0) changes: the more . is affected, the more important 2;; (0) is. A popular quantity measure the input-output
8x_1 (Novak et al., 2018), where ||-|| is the Frobenius sensitivity is the partial derivative norm (a)
norm. If the clustering property holds, the partial derivative norms of beginning and ending tokens should be the highest among all tokens.
We randomly select 50 prompts and compute the partial derivative norms with respect to their tokens. Here omit the last input token (0) from consideration, since perturbing it obviously leads to
4

---

Published as a conference paper at ICLR 2025
Layes
Loyer
Layer
Lager 3
Loyer
Layer
Figure 4: 2D UMAP projections of x-1 (1)'s with 2 on reverse task and GPT-2. Layer is input for GPT-2, while layer 48-th is the last layer of GPT-2. We observe interesting property that (t)'ss are clustered from early layers. A thcoretical explanation will given in Section 4.1.
the most dramatic changes in 1 As prompt lengths are different, we divide each prompt into 10 chunks: the first chunk includes the first 10% tokens, and on, to the last chunk including the last 10% tokens. For each chunk, we compute the partial derivative norms with respect to all of its tokens, then averaging out to obtain representative norm for that chunk. Doing this way, we always have 10 chunk partial derivative norms for prompt, regardless of how long the prompt is.
Figure 3b shows the results. We observe consistent U-shapes across LLMs and tasks, indicating that both beginning and ending tokens play important roles in prompts, with beginning tokens showing relatively stronger effects. These results again support the existence of clustering Rd Note that these results do not imply that middle tokens are not important (see Table 8 in Appendix F).
3.3 CLUSTERING AT INTERMEDIATE LAYERS
In previous sections, we have merely studied the behaviors of x=-'s and discovered their clustering property. We wonder if the clustering property still occurs at intermediate Transformer layers. In this section, we conduct similar experiments with the 2D UMAP projections, this time for the last input token embeddings at intermediate layers (i.e. (t)'s 's with < T). We report the results of sentiment classification task on GPT-2 in Figure 4. Interestingly, the clustering property appears from early Transformer layers, and it is maintained in the afterwards layers.
The emergence of clustering in early LLM layers suggests Transformer operations themselves may significantly contribute clustering. In the next section, we provide theoretical and empirical evidences to support this hypothesis.
4 MECHANISMS BEHIND CLUSTERING PROPERTY
We are motivated by findings Section 3 to further explore the mechanisms behind the clustering property. In this section, we provide theoretical analysis (Section 4.1) and empirical experiments (Section 4.2) to explain the observed behaviors clustering. Our analysis suggests that while prompts with shared first or last demonstrations form clusters in the embedding space, the underlying mechanisms may differ. The first-demonstration clustering appears to be partially explained by the causal attention mask mechanism. In contrast, the clustering by last demonstrations seems to involve a more complex interaction between the causal structure and positional encoding. This difference in mechanisms aligns with empirical observation that fist-demonstration clustering tends to be more pronounced, though further investigation would be needed for a complete understanding.
4.1 THFORETICAL ANALYSIS
We provide theoretical evidence for the clustering property. More importantly, we would like to point that causal attention mask is one important factor leading clustering by beginning tokens. This, in turn, results in the property of clustering by the first demonstration.
Suppose (1) is on the unit sphere gd-1 for every i I, and 2 The dynamics of Transformer can be written as
dx:
(000000000000000000000000
= P Pz((0)
(I)
dt
(t, ==1
5

---

Published as a conference paper at ICLR 2025
Thsl:: Matbortth
Taslo Reverse
Thel: Symbolioeeettmme
02
04-
90-
26
Figure 5: Attention weights from the last token the first token (blue) and last token (orange) with Qwen- 2.5-72B across layers and different tasks. The x-axis shows layer indices ranging 0-80; y-axis shows atlention weights averaged over 100 prompts. Notably, attention first token dominates early layers (0.8-0.9) aligning with Proposition 1, while attention to last token stcadily increases and peaks final layers (around 0.2). (Left) Mathematical arithmetic task. (Middle) Reverse task. (Right) Symbolic sentiment classification task.
for i = 1, 72 and f 2 0. Here Q(), K(), and V(.) are the query, key, and value matrices, P is the projection of E gd-1 onto the tangent space T, and (t, = Σi= (8))())) the partition function. We further assume (t) is identity for all t 2 O.
Consider an initial states (0)) (IR) " that satisfies the following hypothesis (H): there exists E Sd-1 such that (we, 1r;;;;00 (0) for all I, n. Under the hypothesis (H). it has been proved that (.)(-) E Co (R2o: (s=-1)") is the unique solution of the corresponding Cauchy problem (1), then there exists xx E gd-1 and constants C, > such that [tx,,t) x"1 S Ce- (Geshkovski et al., 2023). This means all items in the sequence (tr become identical exponentially fast w.r.t. the depth & under appropriate conditions. We call x the convergence point of the Cauchy problem corresponding
the initial state {0) The following proposition shows that the convergence points of two different initial states with the same first token coincide when casusal attention mask is applied.
Proposition 4.1. Let {zx(0)};_ and {z((0) be different initial states satisfying (H) and I (0) = 2 (0). Note that n may be different from n'. Let x and x'* are the convergence points of the Cauchy problems Eq. (1) corresponding to (0)) and (0)) respectively. causal attention mask is applied, then x" x² This means all states of the two Cauchy equations corre- sponding (0)( and (z)(0)) eventually coincide, even though they only share the same first input token.
Proof. The key notice here is if causal attention mask is applied, then 2, (t) for all 2 0. This leads to
|| - 2111 || S ||= - 21 (t)l + ||z - z''ll
Ce 2 X't 11 0,
which means x² = x²".
Proposition 4. is closely related to the prompts share the same first demonstrations, they also share the same first token (i.e. xn (0) = x (0). Under appropriate circumstances as in Proposition 4.1, the convergence points these prompts coincide, which means their associated w_1 's also coincide if there are infinitely many Transformer layers. Moreover, with xx = x'",, we further have
property. Concretely, if two
| | - -000='10011 < [|=- - 211l + ||=L((() x'1ll
< Ce + C'e--Xe Vt 2 0.
which means the last embeddings become identical exponentially fast. This explains why clustering occurs from early layers of Transformer as shown in Section 3.3.
Proposition 4.1 draws a close theoretical relationship between causal attention mask and first- demonstration clustering. We further empirically investigate this connection in the next section, together with the one between positional encoding and last-demonsttation clustering.

---

Published as a conference paper at ICLR 2025
(8)
(b)
(c)
P441
No pasteaaleecdde
141
No ceardii cocrding
241
No periisal
No www.kaicccooo twok
Napa esc, als calk
Nockadsencakk
Nopa oea and wad.
No
No pii nx addes soil
Figure 6: Partial derivative norms w.E.. chunks in froied-fron-scrateee Transformers with different types of positional encodings. The æ-axis is chunk indices ranging 0-9; y-axis is the partial derivative norms. (a) Sinusoidal positional encoding. (b) Rotary positional encoding. (c) Trainable positional encoding.
Does Proposition 4.1 violate last-demonstration clustering? While Proposition 4.1 suggests that first-token clustering theoretically possible under ideal conditions, it does not address contradict our empirical observation of last-demonstration clustering. Our analysis of attention weights across layers reveals more nuanced picture of how positional information is processed in practice.
Figure 5 shows the attention weights from the first token (blue) and last token (orange) to the last token across different layers, measured on the 2w5--2.5-72B model (Hui al., 2024). The results are averaged over 100 different prompts of varying lengths across multiple tasks including mathematical arithmetic, reverse and symbolic sentiment classification. This comprehensive sampling ensures our findings robust across different task types and input lengths. The attention pattern exhibits three distinct phases:
Initial Phase (layers 1-40): The first token dominates attention with consistently high weights (0.8-0.9), showing remarkable stability in early layers. This strong initial convergence precisely aligns with our theoretical prediction in Proposition 4. and reflects the attention sink phenomenon (Xiao et al., 2024).
Transition Phase (layers 40-60): We observe a sharp decline in first-token attention, dropping from 0.8-0.9 to 0.2-0.4. During this phase, attention begins redistribute across tokens, marking shift from the theoretical behavior to practical requirements.
Final Phase (layers 60-80); The first-token attention exhibits oscillating patterns (between 0.4-0.7), while attention to the last token steadily increases and reaches its peak (0.1-0.2) in the final layers. This consistent pattern of increasing last-token attention across tasks, despite variations in first-token attention, suggests fundamental architectural behavior aligned with the next-token prediction objective.
This layerwise progression offers insights into how our theoretical and empirical findings coexist: while the theoretical tendency toward first-token clustering manifests in early layers (with attention weights of 0.8-0.9), the practical requirements of causal language modeling lead to attention redis- tribution in later layers, where we observe increased attention to the last token (reaching 0. 1-1-0. This pattern suggests a dynamic balance between the model's architectural bias toward first-token clustering and its need to capture sequence- final information for next-token prediction. This observa- tion aligns with prior findings that causal attention mechanisms can infer positional information even without explicit positional encoding (Kazemnejad et al., 2024) a property that emerges naturally from the causal structure. When combined with positional encodings, this evolving attention pattern appears to contribute to the observed dual clustering behavior both first-demonstration and last- demonstration clustering though our experiments in Section 4.2 suggest this interaction is complex and dependent on the specific type of positional encoding used.
4.2 EMPIRICAL ANALYSIS
While our theoretical analysis in Section 4.1 suggests how causal attention mask could enable firt-demoosstrtion clustering under ideal conditions, we need to examine how these mechanisms manifest in practice without theoretical restrictions. We investigate both first-demonstration and last-demonstration clustering through empirical experiments, following a similar methodology to Section 3.2. Our experiments suggest that different architectural components contribute distinctly to clustering behavior: the causal attention mask appears particularly important for firt-demmonttration clustering, while both causal attention mask and positional encoding seem to play roles in last- demonstration clustering.
7

---

Published as a conference paper at ICLR 2025
Exhaustive
Cluster
Relative decrease
GPI--Ne-2.7B
97.4 (1.8)
94.0
3.5
Llama-v2-7B
98.5 (4.5)
95.6
2.9
MPT-7B
98.7 (11.9)
97.8
0.9
Table 1: Accuracies (%) of Exhaustive and Cluster-bused Search with the ideal selecting criterion with different LLMs on reverse task. The first two columns show accuracies, while the third one shows the relative decreased percentage of Cluster-based Scarch compared Exhaustive Search. Exhaustive Search achieves the best accuracies on all LLMs, but the performances of Cluster-based Search are comparable. Numbers in parentheses are accuracies with dumb selection criterion.
To investigate the role of causal attention mask and positional encoding, train different Trans- formers with or without those components. Specifically, we train from scratch four Transformers: without positional encoding,
one with full components, one without both. Other Transformer components are always included. Each Transformer has 12 self-attention layers, each with 12 attention heads; the token embedding size is 768, and the hidden size in MLP layers is 2048. We train the Transformers with different types of positional encodings, namely the sinusoidal, rotary, and trainable positional encoding, on WikiTex12 dataset with SGD optimizer with learning rate 5 10 in 100 epochs. For fair comparisons, the training task for all Transformers next-token prediction.
without causal attention mask, and
We prepare 100 randomized prompts and compute the partial derivative norms similarly to Section 3.2. To ensure the prompts are differently distributed to training ones, we build each prompt as sequence of 50 to 100 random words, resulting in meaningless sentences. For each prompt, we compute its chunk partial derivative norms, then average over 100 prompts. Figure 6 shows interesting results. reveals a robust correlation between first-demonstration clustering and the utilization of the causal attention mask. Specifically, the importance of beginning tokens markedly elevated when, and only when, the causal attention mask is applied, which aligns with the findings presented in Proposition 4.1. On the other hand, the case for last-demonstration is more complex. While the importance of ending tokens remains distinctively high when sinusoidal positional encoding is employed in the absence of a causal attention mask, this phenomenon is not observed for rotary and trainable positional encoding. This suggests that the importance ending tokens is influenced by the interplay between the causal structure and the choice of positional encoding method.
ACCELERATING SELECTION AND ORDERING OF IN- IN-CONTEXT DEMONSTRATIONS IN SELF-ADAPTIV ICL SETTINGS
Based on the clustering property. propose an efficient approach to improve self-adaptive ICL methods (Wu et al., 2022). In particular, self-adaptive methods aim to optimize the selection and ordering of demonstrations based on model's own predictions, without relying external knowledge or supervision. However, this process typically suffers from factorial complexity. For example, when exhaustively selecting and ordering out of krotal demonstrations (referred to as Exhaustive Search), there are Al = [Kj possibilities. In contrast, the firs-demmonsraation clustering property suggests that prompts sharing the first demonstration are likely to have the same next-token prediction. Consequently, our approach, called Cluster-based Search, only requires selecting the first demonstration, while the rest can randomly selected, resulting in merely kiotal possibilities.
Auni!
Formally, consider an L.M,, pool of kiotal demonstrations E and a query q. Any algorithm for demonstration selection and ordering will construct set of possible prompts, P(F, q). and uses pre-defined criterion C to select among this set potentially correct prompt, i.c. with which the LLM can answer the query correctly. In this section, we consider two types of criterion C: the ideal criterion and the practical entropy-basee one. In both cases, Cluster-based Search achieves comparable performances with Exhaustive Search while being significantly faster.
We consider Asxaa 6 and k 4 for as it is feasible to Exhaustive Search. We also report results when Aipcal = 16 and = in Appendix to show the scalability of Cluster-based Search. Furthermore, in this section, focus non-instructed prompt to better align with previous
8

---

Published as a conference paper at ICLR 2025
Classification
Reasoning
Avg
SymSen
SymLan
Rev.
Rep.
ComSen
Math
Random
51.34.7 51.5
65.2a.0
56.6s.
49.5
17.3a.1
1.0p.s
40.2
GPT-Neo-2.7B
Exhaustive
77.0
66.5
49.0
18.7
0.5
43.9
Cluster
51.55.6
74.63.3
76.3s.1
55.3s.6
16.63.5
1.0p.s
45.9
Random
23.13.7
59.34.8
77.94.9
53.0s.4
61.24.0
7..
58.3
Phi-2 (2.7B)
Exhaustive 31.5
76.0
80.5
71.0
63.5
77.0
66.6
Cluster
30.8
72.3s.1
77.33.o
65.04.6
64.23.9
77.94.2
64.6
Random 71.73.0
76.32.8
99.1o.s
93.32.4
84.33.4
85.13.2
85.0
Qwen-2.5-14B
Exhaustive
78.0
86.0
100.0
96.0
84.2
85.5
88.3
Cluster
82.32.9
87.13.3
99.30.9
97.71.6
84.32.5
87.13.0
87.
89.6
Table Accuracies (%) of Random Selection, Exhaustive Search, and Cluster-based Search with entropy-based selecting criterion different I.IMs and ICI. tasks. The subscript numbers indicate the standard deviation over 10 runs. Due to computational constraints, we do not report standard deviations for Exhaustive Search. Bold: best; Underline: second best. Scarching is obviously
effective than random selection. Morcover, the performances of Cluster-based Search are comparable to Exhaustive Search.
discussions on the positions of demonstrations. Additional results regarding the instructed prommp case given in Table 7 in Appendix F.
5.1 IDEAL SELECTION CRITERION
In this section, we consider C be the ideal criterion, which always selects the correct prompt in P(E, 4) if exists. This criterion is alternatively called the oracle (Lu et al., 2022). The accuracy achieved by this ideal criterion sets an upper bound for any other selecting criterion. We also consider the dumb criterion, which always select the wrong prompt if exists. In the following, show the accuracy upper bounds of cluseer-based Search are comparable to Exhaustive Search while significantly reducing the search time.
Accuracies of different LLMs on reverse task are reported in Table 1. Here take the average accu- racy over 1, 000 tuples of (E, q) See Appendix for results of other tasks. While Exhaustive Search achieves the best accuracies on all LLMs, it worth noticing that the performances of Cluster-based Search are comparable. Concretely, performances of Cluster-based Search only decrease relatively by 2.4% on average compared to Exhaustive Search, decreases while search time of Cluster-based Search by 91.7% (Figure 7). Note that with larger kiotal, the time saving is almost 100%. Additionally, the perfor- mance gaps between the ideal and the dumb criterion (number in parentheses in Table 1) are huge for all LLMs, further highlighting the efficiency of Cluster- based Search.
Nemalieed PeamateeePefffeeeee Peffemmnnc
Rduartie
Chise
81-
Accemayy
Figure Normalized accuracies and search times of Exhaustive Search and Cluster-based Search with the ideal selecting criterion.
5.2 ENTROPY-BASER SELECTION CRITERION
We consider 2022), to show the effectiveness of Cluster-based Search practice. To be concrete, denote ε- to be the logits of the first prediction step. The size of {_ is equal to the length of the token dictionary, and it is the output of the prediction head whose input is -- We define the confidence score of f_, to be c(__1)== == entropy (softmax (E) A permutation p of demonstrations is more confident than another permutation p' if its associated logits {_1 more confident than f' i.c. e (f-1) (f",). Among all prompt candidates, we will select the most confident one. See Algorithm in Appendix for the pseudocode of this criterion.
practical selecting criterion C. i.e. the popular entropy-based criterion (Lu et al.,
We run the Exhaustive Search and Cluster-based Search with Algorithm on various ICL tasks introduced in Section 2.2. Each task includes 1, 000 tuples (E, q). Average accuracies of different Random Selection are the worst,
LLMs are reported Table 2. On average, the performances with large gaps compared Exhaustive Search, showing the effectiveness of selecting and ordering.

---

Published as a conference paper at ICLR 2025
Classification SymSen
Reasoning
Avg
SymLan
Rev.
Rep.
ComSen
Math
First Cluster
51.5
74.6
76.3
55.3
16.6
1.0
45.9
GPT-Neo-2.7B
First-Last Cluster
55.5
75.7
63.7
53.2
17.8
1.0
44.5
First Cluster
30.8
72.3
77.3
65.0
64.2
77.9
64.6
Phi-2 (2.7B)
First-Last Cluster
28.1
72.2
78.9
69.3
62.7
78.5
64.9
First Cluster
82.3
87.1
99.3
97.7
84.3
87.1
89.6
Qwen-2.5-14B
First-Last Cluster
79.3
83.7
99.2
97.0
84.0
86.2
88.2
Table 3:
Comparison between First-only and First-Last clustering search strategies across different models and lasks, showing accuracy (%) ) on classification and reasoning tasks.
However, it is interesting to note that performances of Cluster-based Search are comparable to Exhaustive Search, with slight absolute performance gap of 0.3% average. This result once again highlights the efficiency of Cluster-based Search.
5.2.1 ABLATION STUDIES
We further investigate whether selecting only the first demonstration, rather than both first and last demonstrations, could maintain performance while improving efficiency. Table presents the results of this investigation.
The results demonstrate that selecting only the first demonstration for clustering-baset search achieves comparable or sometimes better accuracy compared to selecting both first and last demonstrations. This is particularly evident in GPT-Neo-2.78 where first-only clustering achieves an average accuracy of 45.9% compared to 44.5% for first-last clustering, and in Qwen-2.5-14P where it achieves 89.6% versus 88.2%. Most notably, this simpler approach reduces the computational complexity significantly from (Khea (konal 1)) when selecting both first and last demonstrations to just ((kuaa) when selecting only the first demonstration. This order of magnitude improvement in efficiency, combined with the maintained improved accuracy, makes first-only clustering clearly superior choice for demonstration selection.
RELATED WORK
Recent research has unveiled the sensitivity to the order of in-context demonstrations in large language models (L.LMs). has demonstrated that simply reshuffling the examples in sentiment analysis prompt can dramatically impact accuracy, ranging from near-random guessing to state-of-the-ar performance (Zhao et al., 2021). Subsequent studies have confirmed this "order sensitivity" is widespread, affecting models of all sizes and independent of specific examples (Liu et al., 2021). Further work noted a negative correlation between sensitivity and accuracy, suggesting highly sensitive outputs less reliable (Chen al., 2022). Another study finds that language models' performance can significantly decline when relevant information appears in the middle of the context, suggesting they may struggle effectively utilize long contexts (Liu et al., 2023). To the best of our knowledge, our paper is the first to investigate the prompt embeddings to explain the order sensitivity of LLMs ICL, Furthermore, we also provide thcoretical and empirical evidences support our findings.
Oversmoothing is problem where information is lost and representations become identical. As explored in Section 3.2, this issue is closely related to the clustering property of LLMs. Previous work shows Transformer prone oversmoothing during training. especially with deep architectures. and this ultimately lowers model performance (Chen et al., 2022). Further research finds the underlying cause of oversmoothing in Transformer, and hypothesizes that layer normalization plays a crucial role in oversmoothing (Shi et al., 2022) Recently, phenomenon called attention sink has been discovered, which states that a large portion of the attention weights stays at the first token, starting from early layers of the Transformer (Xiao et al., 2024). This may partly explain the oversmoothing problem in Transformer.
The feature importance in neural networks can be assessed by the Jacobian matrix, which maps minor input changes to output shifts. This sensitivity measure has motivated research on model robustness and knowledge transfer. Recent work reveal that even the simple optimization algorithm SGD subtly
10

---

Published as a conference paper at ICLR 2025
regularizes the model's sensitivity (Lee al., 2023), while other studies manipulate the Jacobian to boost noise and attack resistance (Hoffman et al., 2019). Beyond robustness, it has been highlighted how Jacobian regularization supports more realistic dynamics in the system (Finlay et al., 2020).
The scrial-positiie effect, a fascinating psychological phenomenon where first and last items in a sequence recalled best, has been fruitful area of study. Pioneering work highlights how position influences recall, suggesting more rehearsal for beginning and ending items (Murdock and Bennet, 1962). Later research identifies two driving forces: the primacy effect, boosting recall for initial items, and the recency effect, favoring recent ones (Glanzer and Cunitz, 1966). We give a brief discussion on the connection between the perial-posiiinn effect and our clustering property in Appendix A, and we believe that this research direction is interesting to explore further.
CONCLUSION AND LIMITATIONS
We study the prompt embedding space to understand the order sensitivity of ICL in decoder-only LLMs. Our analysis reveals that prompts sharing first or last demonstrations tend to form clusters in the embedding space, with first demonsttmtttttinn clustering showing notably stronger effects. Our theoretical analysis suggests this asymmetry may be partially explained by the causal structure, as first-token information is reused O(n²) times in self-attention layer. The last-demonstration clustering appears to emerge through a more complex interaction between the causal structure and positional encoding, though its precise mechanisms warrant further investigation. Based on these insights, we introduce Cluster- Search for demonstration selection and ordering in aclf-adaptive ICL methods, achieving comparable performances to exhaustive search while reducing computational complexity from factorial to quadratic.
Our work contributes to understanding the operational mechanisms of LLMs, area that would benefit from increased research attention. The clustering property we identify may represent one of several important characteristics that could enhance our comprehension of these models. Such insights can lead to practical improvements in L.M applications, as demonstrated by our Cluster-hased Search method.
Several limitations of this study should be acknowledged. First, while Cluster-based Search shows promise in self-adaptiv settings, its effectiveness in other contexts requires further evaluation. Second, our theoretical framework primarily addresses disttteemontttttion clustering, while the mechanisms behind dast-demonstration clustering remain incompletely understood. Third, verifying the clustering property poses challenges in black-bor box models where internal representations are not accessible. Future research should focus on developing techniques to better understand these behaviors and extend our findings to broader contexts.
REFERENCES
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language Models are Few-Shot Learners. Advances neural information processing systems, 33:1877-199 1901, 2020.
Lihu Chen, Gaël Varoquaux, and Fabian M. Suchanek. The Locality and Symmetry of Positional Encodings. Conference on Empirical Methods in Natural Language Processing, 2023.
Yanda Chen, Chen Zhao, Zhou Yu, Kathleen McKeown, and He He. On the Relation Between Sensitivity and Accuracy in In-Context Learning. Findings of the Association for Computational Linguistics: EMNLP 2023, abs/2209.0766 07661,,,,.. 2022.
Fergus I. M. Craik, J. Matthew Gardiner, and Michael J. Watkins. Further Evidence for a Negative Recency Effect in Free Recall. Journal of Verbal Learning and Verbal Behavior, 9:554-560, 1970.
Chris Finlay, Jörn-Henrik Jacobsen, Levon Nurbekyan, and Adam Oberman. How to Train Your Neural ODE: The World of Jacobian and Kinetic Regularization. In International conference machine learning, pages 3154-3166 PMLR, 2020.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The Pile: An 800GB Dataset of Diverse Text for Language Modeling. arXiv preprint arXiv:2 100.00027, 2020.

---

Published as a conference paper at ICLR 2025
Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A Mathematical Perspec- tive on Transformers. arXiv preprint arXiv:2112.10794, 2023.
Murray Glanzer and Anita R. Cunitz. Two Storage Mechanisms in Free Recall. Journal of Verbal Learning and Verbal Behavior, 5:351-360, 1966.
Judy Hoffman, Daniel A. Roberts, and Sho Yaida. Robust Learning With Jacobian Regularization. ArXiv, htss/998.002299, 2019.
Mohammad Javad Hosseini, Hannanch Hajishirzi, Oren Etzioni, and Nate Kushman. Learning to Solve Arithmetic Word Problems With Verb Categorization. In Proceedings of the 2014 Conference Empirical Methods in Natural Language Processing (EMNLP), pages 523-533, 2014.
HuggingFace. Language Identification Dataset, 2021. URL https: http:ss//hhggiigggfeeeee co/
batasetssppapaacaLLagguue Languugggaaaaaeegeeeaaaaaa itteeeeeeeeetiiiiiiiiiii
Binyuan Hui, Jian Yang, Zeyu Cui, Jiaxi Yang, Dayiheng Liu, Lei Zhang, Tianyu Liu, Jiajun Zhang,
Bowen Yu, Keming Lu, et al. Qwen2. 5-coder technical report. arXiv preprint wXXv:2409.121886 2024.
Romuald A Janik. Aspects of Human Memory and Large Language Models. arXiv preprint arXiv:2311.0388/ 2023.
Mojan Javaheripi, Sébastien Bubeck, Marah Abdin, Jyoti Aneja, Sebastien Bubeck, Caio César Teodoro Mendes, Weizhu Chen, Allie Del Giorno, Ronen Eldan, Sivakanth Gopi, et al. Phi-2: The surprising power of small language models. Microsoft Research Blog, 13):3, 2023.
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 36, 2024.
Sungyoon Lee, Jinseong Park, and Jaewook Lee. Implicit Jacobian Regularization Weighted With Impurity of Probability Output. In International Conference on Machine Learning, 2023.
Jiachang Liu, Dinghan Shen, Yizhe Zhang, Bill Dolan, Lawrence Carin, and Weizhu Chen. What Makes Good In-Context Examples for GPT-3? In Workshop on Knowledge Extraction and Integration for Deep Learning Architectures; Deep Learning Inside Out, 2021.
Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the Middle: How Language Models Use Long Contexts. Transactions of the Association for Computational Linguistics (TACL). 2023.
Stuart P. Lloyd. Least Squares Quantization in PCM. IEEE Trans. Inf. Theory, 28:129-136, 1982.
Yao Lu, Max Bartolo, Alastair Moore, Sebastian Riedel, and Pontus Stenetorp. Fantastically Ordered Prompts and Where to Find Them: Overcoming Few-Shot Prompt Order Sensitivity. Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), 2022.
Leland McInnes, John Healy, and James Melville. UMAP: Uniform Manifold Approximation and 2018.
Projection for Dimension Reduction. arXiv preprint
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer Sentinel Mixture Models. International Conference on Learning Representations, 2017.
Melanie Mitchell, Abstraction and Analogy-Making in Artificial Intelligence. Annals of the New York Academy of Sciences, 1505(1):7--001, 2021.
Murdock and B Bennet. The Serial Position Effect of Free Recall. Journal of Experimental Psychology, 64:482-488, 1962.
Roman Novak, Yasaman Bahri, Daniel A. Abolafia, Jeffrey Pennington, and Jascha Narain Sohl- Dickstein. Sensitivity and Generalization in Neural Networks: An Empirical Study. International Conference Learning Representations, 2018
12

---

Published as a conference paper at ICLR 2025
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language Models Are Unsupervised Multitask Learners. OpenAl blog, 1(8)99 2019.
Dewey Rundus. Analysis of Rehearsal Processes in Free Recall. Journal of Experimental Psychology, 89:63-77, 1971.
Han Shi, Jiahui Gao, Hang Xu, Xiaodan Liang. Zhenguo Li, Lingpeng Kong, Stephen M. S. Lee, and James Tin Kwok. Revisiting Dver-Smoothing in BERT From the Perspective of Graph. International Conference on Learning Representations, 2022.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng. and Christopher Potts. Recursive Deep Models for Semantic Compositionality Over Sentiment Treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, pages 1631-1642, Seattle, Washington, USA, October 2013. Association for Computa- tional Linguistics.
Michael S. Steinbach, Levent Ertoz, and Vipin Kumar. The Challenges of Clustering High Dimen- sional Data. 2004.
Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. Co 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 (Long and Short Papers), pages 4149 4158, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
MosaicML NLP Team. Introducing MPT-7B: A New Standard for Open-Source, Commercially Usable LLMs, 2023. URL www www www.wmosaiicoocmm cm1
h/ppt//ppt-///// 7b. Accessed: 2023-05-05.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. LLLam: Open and Efficient Foundation Language Models. arXiv preprint arXiv:2302.1397.1 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Bashlykov, Nikolay, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, al. I.Lama 2: Open Foundation and Fine-Tuuned Chat Models. arXiv preprint arXiv:2307.09288, 2023b.
Zhiyong Wu, Yaoxiang Wang, Jiacheng Ye, and Lingpeng Kong. Self-Adaptive In-Context Learning. arXiv preprint a1X11:2212.10375, 2022.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient Streaming Language Models With Attention Sinks. International Conference on Learning Representations, 2024.
Tony Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. Calibrate Before Use: Improving Few-Shot Performance of Language Models. In International Conference Machine Learning, 2021.
13

---

Published as a conference paper at ICLR 2025
APPENDIX
A RELATION OF CLUSTERING PROPERTY AND SERIAL-POSITION EFFECT
Analyses in Section 4.1 and 4.2 allow to link the clustering property with the well-known serial- position effect (Murdock and Bennet, 1962), which has two components: primacy effect and recency effect. The primacy effect means that the first item in sequence is remembered better than middle ones. Previous research attributes this to the long-term memory: the first item gets more rehearsal when the person tries to recall the sequence, and thus it is more likely to be stored in the long-term memory (Rundus, 1971). This is similar to causal attention mask, where all tokens attend to the first one for updating. Moreover, it has been shown that that the primacy effect weakens with longer sequences (Murdock and Bennet, 1962). We see the same pattern in LLMs: the ratio between the importance of beginning tokens and ending tokens shrinks as the prompts get longer (see Figure 8). On the other hand, the recency effect means that the last item in sequence is remembered better than middle ones. This is due to the short-term memory, which keeps a few recent items (Craik et al., 1970). Recent work suggests that positional encoding adds this locality language models (Chen et al., 2023). Following Secion 4.2, this implies that the dast-demonstration clustering property is very similar to the recency effect.
Tink: Sentiment Clavsiication
Task: Reverse
Repeat 1.M GPT-Noo-2.7
L.M GPT-3
11M: Elama-22-78
10
Raie
1.0
26
255
08
67
Figure 8: Ratio of the first and last token chunk importance. The r-axis is the number of demonstra- tions. The ratios consistently decrease with different LLMs and ICL tasks.
B ENTROPY-BASED SELECTION CRITERION PSEUDO-CODE
Here we give the pseudo-code for the entropy-based selection criterion presented in Section 5.2. The algorithm computes the confidence score for each permutation of demonstrations based on the entropy of the softmax output of the first prediction step's 's logits. It then selects the permutation with the highest confidence score, i.c., the lowest entropy, as the most confident prompt candidate. This approach allows for the identification of the optimal demonstration order that maximizes the model's confidence in its predictions.
Algorithm 1 Entropy-Basee Selecting Criterion
Input: set of prompt candidates P
set cBest = -inf
set pBest None for p in P do
compute logits é_1
compute confidence score e (4-1)
if (f-1)))) > cBest then
cBest = c (f-1)
pBest = p
end if end for Output: pBest
14

---

Published as a conference paper at ICLR 2025
Classification
Reasoning
Avg
Sent.
Lang.
Sym.Sen. Jym.Sen.Sym.Lan Rev.
Rep.
Exhaustive
99.9
100.0
52.0
100.0
97.4
94.6
90.7
GPT-Nen-2.7B
Cluster
99.6
99.8
51,9
100.0
94.0
90.3
89.3
Exhaustive
99.7
100.0
100.0
100.0
98.5
92.6
98.5
Llama-v2-7B
Cluster
99.4
99.9
100.0
99.8
95.6
89.0
97.3
Exhaustive
99.7
99.9
52.0
100.0
98.7
97.8
91.4
MPT-7B
Cluster
99.2
99.9
52.0
99.9
97.8
96.0
90.8
Table 4: Accuracies (%) of Exhaustive and Cluster Search with the ideal selecting criterion with different LLMs on various task. The case know 6 and k 4.
Classification
Reasoning
Avg
SymSen SymLan Rev.
Rep.
ComSen
Math
Random
51.0
54.0
58.0
47.0
13.0
2.0
37.5
GPT-Nco-2.7B
Cluster
55.5
80.0
82.0
56.5
17.0
48.8
Random
22.5
53.0
77.0
53.5
57.5
80.0
57.3
Phi-2 (2.7B)
Cluster
36.0
82.0
78.5
69.5
62.5
76.5
67.5
Random
65.0
74.0
98.0
98.5
86.5
88.0
85.0
Qwen-2,5-14B
Cluster
88.0
96.0
100.0
99.5
85.5
86.0
92.5
Table Accuracies (%) of Random Selection, Exhaustive Search, and Cluster Search with entropy- based selecting criterion on different LLMs and ICI. tasks. The case knmaa = 16 and k = 4.
C SEARCH WITH IDEAL SELECTION CRITERION
We report in Table 4 results of similar experiments as in Section 5.1 on different LLMs and ICL tasks. By comparing the performance of Cluster-based Search with Exhaustive Search, we can assess the effectiveness of our proposed method in finding optimal demonstration orders. Overall, the performances of cluster-based Search are comparable to Exhaustive Search, while significantly reducing the search time complexity.
D CLUSTER SEARCH WITH KTOTAL = 16 AND
report in Table results of similar experiments as Section 5.2 with kiotal = 16 and A = 4. The table presents the average accuracies of different LLMs various ICL tasks using the entropy-basee selection criterion. By comparing the performances of Cluster-based Search and Random Selection, assess the effectiveness and scalability of proposed method when the total number of demonstrations is increased to 16. Moreover, comparing to Table 2, we observe that the performance of Cluster-based Search is better when the pool is larger. These results provide further evidence for the efficiency of Cluster-based Search in selecting and ordering demonstrations for in-context learning, even when dealing with larger pool of available demonstrations.
We further conduct similar experiments with kiotal = 16 and A: = 10 to show the scalability of our proposed Cluster Search. this case, the Exhaustive Search clearly infeasible. Experimental results are reported in Table 6. On average, the Cluster Search method is better than Random Search by 1.5% 5.9%.
E LOW-DIMENSIONAL PROJECTIONS WITH T-SNE
We re-build Figure with UMAP projection replaced by t-SNE projection. The t-SNE projections, shown in Figure 9, exhibit similar clustering patterns as observed in the UMAP projections. Permuta- tions sharing the first demonstrations (indicated by colors) the last demonstrations (indicated by shapes) tend to form distinct clusters in the low-dimensional space. This consistency across different dimensionality reduction techniques further supports clustering property.
15

---

theoretical Published as a conference paper at ICLR 2025
Classification SymSen SymLan
Reasoning
Avg
Rev.
Rep.
ComSen
Math
Random Cluster
54.1
75.1
78.9
62.1
20.6
2.7
48.9
GPT-Neo-2.7B
52.6
82.1
85.8
62.0
19.1
2.9
50.8
Random
28.3
71.2
77.4
62.7
67.0
79.4
64.3
Phi-2 (2.7B)
Cluster
38.3
80.3
68.1
67.6
82.0
Random
84.6
90.4
99.8
99.3
85.8
87.8
91.3
Qwen-2.5-14B
Cluster
88.5
95.9
99.6
99.2
85.5
88.0
92.8
Table 6: Accuracies ( (%) of Random Selection, Exhaustive Search, and Cluster Search with entropy- based selecting criterion on different LLMs and ICL tasks. The case Kroua 16 and 10.
Taok: Revene
Took: Repeat
Tack: Sentimen--Smmhan
L1M
Figure 9: r-SNE 2D projections of x-1 's with different tasks and LLMs. Points with the same color and shape tend to form clusters.
F OTHER PROBLEMS REGARDING PROMPT DESIGN
Prompt with instruction In this paper, only focus prompts without instructions to highlight the clustering property with respect to the orders of in-context examples. However, when an instruc- tion is included, i.., when the first token of the prompt is part of the instruction, result in Proposition 4.1 still holds. Furthermore, Figure 3 also gives us evidences that the first and last demonstrations are still more important than the other instruction rather than a demonstration. We provide additional empirical evidence Table 7 for the case of reverse task with instructed prompt. Overall, the performances of our cluster-bate based method are comparable to the costly exhaustive method.
even when the prompt starts with
GPT-2
GPT-Nco-2.7B
Llama-v2- -2BBBBBBBBBBBBBBBB
Random
30.3
67.2
94.3
Exhaustive
37.0
83.8
94.8
Cluster 35.3
82.7
95.1
Table 7: Performances of diferent methods with different LLMs in the reverse task with instructed prompt. Here we consider kirota = 6 and k: = 4. The cluster-based method still comparable to the exhautive method.
About the importance of middle in-context demonstrations The clustering property does imply that the middle in-context demonstrations are not important. Indeed, it only shows that first and last demonstrations are more important than the others. However, the middle demonstrations still play role in shaping the model's understanding of the task and the context. They provide additional examples and variety, which can help the model generalize better. Empirical evidence is given in Table 8, in which we run different LLMs on the reverse task with 4 and 8 demonstrations, respectively. The performances of the case of 8 demonstrations are consistently better, showing the importance of middle demonstrations.
16

---

Published as a conference paper at ICLR 2025
GPT-2
Llama-vl-7B
Llama-v2-13B
Llama-v1-300
demonstrations
32.0
19.7
95.1
80.4
demonstrations
44.4
27.5
98.0
86.7
Table 8: Performances of diferent LLMs on reverse task with different number of demonstrations.
17


Reviews:
[{"review_id": "bFPLY3zbmU", "reviewer": "Reviewer_Q71R", "summary": "The authors investigate the few-shot setting and the order of given demonstrations/examples in the prompt. Analyzing the last hidden state of the last layer of decoder-only transformers, they study the clustering property, which prompts sharing the first and last demonstration. Experiments are conducted in two domains: classification and reasoning. Each is divided into two tasks, and the classification sub-tasks are further modified into symbolic ones.\nThe explanation proposed is that this property depends highly on the causal attention mask and the positional encoding. The first demonstration clustering depends on the causal attention mask. However, the last demonstration clustering depends on a more complex interplay of the causal attention mask and the positional encoding.\nFollowing their findings, the authors propose a selection and ordering method based on the uncovered clusters. Experiments are conducted using their methods with an already-used entropy-based search. They compare their methods with an oracle and unmodified entropy methods. Their findings show that the clustering-based method while suffering a slight drop in performance, their method is more than 90% faster.", "questions": "* How does a variant of Figure 3b with demonstrations instead of chunks of text compare?\n* Do the prompts that share a close representation get similar scores? (what is the standard deviation ?)\n* How does the performance change with the number of intermediate demonstrations? ( Some insights are given, but more results would significantly improve the demonstration)", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "* Running large language models is costly, and few-shot in-context learning is a common approach to alleviate the cost. The proposed method is simple and greatly reduces search time, making a practical contribution. \n* Even though the theoretical assumptions are strong, their partial derivative analysis is original and clearly advocates for the clustering property.\n* The cluster-based search proposed by the authors is well explained.", "weaknesses": "* Too little evidence of clustering is given on the classification tasks, and clustering is unclear on The 2D projection (Figure 1, Figure 4).\n* Few experiments have been done varying the number of demonstrations and the pool size; it would be really beneficial to give some insight on the scaling possibility of the method. \n * A more thorough analysis of the results would be appreciated to confirm the findings, for example: Do the prompts sharing a close representation share similar scores? (what is the standard deviation ?) How does the performance change with the number of intermediate demonstrations? ( Some insights are given, but more results would greatly improve the demonstration). \n* Not enough selection methods are considered for comparison in terms of time and scores.\n* A table showing time performance and or gap with other methods is needed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6BUMyzJosM", "reviewer": "Reviewer_nbs8", "summary": "The paper explores the issue of demonstration order sensitivity in large language models (LLMs) during in-context learning (ICL) and uncovers a clustering phenomenon in the embedding space, where prompts with the same first and last demonstrations tend to cluster together. Through theoretical analysis and empirical evidence, the paper identifies that this clustering effect stems from the interaction of causal attention masks and positional encoding. Moreover, they propose a \"Cluster-based Search\" method that significantly reduces the computational complexity of selecting and ordering demonstrations while maintaining high model performance.", "questions": "My main concerns have been listed above. I look forward to the authors' response and am willing to reopen and adjust the score upward.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. Clear Argumentation: The paper is well-structured, with clear explanations that make the objectives and contributions easy to follow.\n2. Robust Proofs: The theoretical analysis is thorough, supporting the proposed mechanisms in in-context learning.\n3. Comprehensive Experiments: The experiments are detailed and varied, effectively demonstrating the method’s efficacy across multiple tasks.", "weaknesses": "1. The models used in this study seem somewhat outdated. Models with the equivalent size should include newer architectures, such as LLaMA 3, Phi, or similar. Why were these not used?\n2. The datasets and tasks included in the study are limited. For instance, why is there no mathematical task such as GSM8k included in the paper\n3. While the authors highlight the importance of the first and last demonstrations in ICL, the figures in the paper suggest that the first demonstration may be particularly or even most significant. However, in the cluster-based method, the authors did not conduct an ablation study that uses only the first or only the last demonstration in clustering to analyze the contributions of the first and last demonstrations independently.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bIgDrDmc29", "reviewer": "Reviewer_NLJn", "summary": "This paper studied the ordering effects of demonstrations in in-context learning (ICL) and claimed that the first and last demonstrations are the most crucial ones for effective demonstrations using both empirical and theoretical analyses. Based on this observation, this paper proposed a cluster-based search method to find out effective demonstration orders (considering only the first and last demonstrations instead of all demonstrations), which will not suffer from the efficiency issue in Exhaustive Search. The experiments showed that the proposed method achieve small drop in accuracy but significant improvement in efficiency.", "questions": "Please refer to my concerns in the weakness part.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* The proposed idea of cluster-based search is simple yet effective for ICL.\n* The performance of the proposed method, especially efficiency improvement, is very promising.", "weaknesses": "* Some claims are not well supported by the empirical analyses. The cluster structure of GPT-2 model in Figure 1 seems unclear, compared to the other two LLMs. Figure 3 (a) shows that the clusters also share the same second demonstrations with high percentage, and for the two bottom figures, the percentage of sharing the same second demonstrations is even higher than the percentage of sharing the same last demonstrations. These observations may be conflict with the main claim of this work. Also, the analyses about the last demonstration seem to be less convincing, e.g., lines 340-346.\n* The theoretical analyses are counter intuitive. According to Prop. 4.1, the embedding of the transformer layers will eventially the same if two promopts share the same first input token. I cannot understand this claim in the proof also, in which the authors mentioned that \"if causal attention mask is applied, then x_1(t) = x'_1(t) for all t >= 0.\" I am not sure why this assumption holds. Intuitively, if this proposition holds, I may infer that only the first demonstration will affect the performance and the last demonstration will not matter too much, which is different from the authors' claim.\n* More comprehensive experiments are required. In Table 1, the case of Random demostrations is not included. It would be useful to also compare with Random ordering as in Table 2. Also, they authors used k=4 in the experiments, it might be also important to evaluate larger k values, e.g., 10 or 20. The main claim of this paper is that the demonstrations in the middle are not very important to the performance of ICL, but using only a few demonstrations in the middle (as in the experiments) may not be as convincing as using many demonstrations in the middle.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)