Rapid Selection and Ordering of In-Context Demonstrations via Prompt Embedding Clustering

Authors: Kha Pham, Hung Le, Man Ngo, Truyen Tran
Primary area: generative models
Keywords: in-context learning, order sensitivity, LLMs, clustering, cluster-based search, positional encoding, attention mask, serial-position effect, cluster-based search

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.

Decision: accept (poster)
Reviewer scores: 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1729838255925

thread_type:
official_review

<thread_object>
creator: reviewer_nljn

<content>
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.

strengths:
* The proposed idea of cluster-based search is simple yet effective for ICL.
* 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.
* 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.
* 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.

questions:
Please refer to my concerns in the weakness part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730356723662

thread_type:
official_review

<thread_object>
creator: reviewer_nbs8

<content>
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.

strengths:
1. Clear Argumentation: The paper is well-structured, with clear explanations that make the objectives and contributions easy to follow.
2. Robust Proofs: The theoretical analysis is thorough, supporting the proposed mechanisms in in-context learning.
3. 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?
2. 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
3. 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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730659312434

thread_type:
official_review

<thread_object>
creator: reviewer_q71r

<content>
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.
The 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.
Following 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.

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. 
* Even though the theoretical assumptions are strong, their partial derivative analysis is original and clearly advocates for the clustering property.
* 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).
* 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. 
 * 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). 
* Not enough selection methods are considered for comparison in terms of time and scores.
* A table showing time performance and or gap with other methods is needed.

questions:
* How does a variant of Figure 3b with demonstrations instead of chunks of text compare?
* Do the prompts that share a close representation get 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 significantly improve the demonstration)
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper investigates the prompt embedding space to address the order sensitivity of in-context learning (ICL) in decoder-only LLMs, revealing that prompts with shared first and last demonstrations exhibit closer embeddings, particularly with stronger clustering around the first demonstration. By analyzing the role of positional encoding and causal attention masks in this clustering phenomenon, the authors propose cluster-based search, a novel method that enhances the selection and ordering of demonstrations in self-adaptive ICL settings. The paper is well-written, and in the discussion phase, the authors addressed almost all of the reviewers' concerns.