Revisiting In-context Learning Inference Circuit in Large Language Models

Authors: Hakaze Cho, Mariko Kato, Yoshihiro Sakai, Naoya Inoue
Primary area: interpretability and explainable AI
Keywords: In-context Learning; Induction Circuit; Mechanistic Interpretability

Abstract:
In-context Learning (ICL) is an emerging few-shot learning paradigm on Language Models (LMs) with inner mechanisms un-explored. There are already existing works describing the inner processing of ICL, while they struggle to capture all the inference phenomena in large language models. Therefore, this paper proposes a comprehensive circuit to model the inference dynamics and try to explain the observed phenomena of ICL. In detail, we divide ICL inference into 3 major operations: (1) Input Text Encode: LMs encode every input text (in the demonstrations and queries) into linear representation in the hidden states with sufficient information to solve ICL tasks. (2) Semantics Merge: LMs merge the encoded representations of demonstrations with their corresponding label tokens to produce joint representations of labels and demonstrations. (3) Feature Retrieval and Copy: LMs search the joint representations of demonstrations similar to the query representation on a task subspace, and copy the searched representations into the query. Then, language model heads capture these copied label representations to a certain extent and decode them into predicted labels. Through careful measurements, the proposed inference circuit successfully captures and unifies many fragmented phenomena observed during the ICL process, making it a comprehensive and practical explanation of the ICL inference process. Moreover, ablation analysis by disabling the proposed steps seriously damages the ICL performance, suggesting the proposed inference circuit is a dominating mechanism. Additionally, we confirm and list some bypass mechanisms that solve ICL tasks in parallel with the proposed circuit.

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

Reviews and discussion:
<thread>
creation date:
1730127457233

thread_type:
official_review

<thread_object>
creator: reviewer_yfjp

<content>
summary:
This paper proposes a three-stage ICL circuit hypothesis and provides thorough empirical examinations of the existence and significance of these stages. Within this circuit framework, many phenomena are explained, such as how Forerunner Tokens encode input text representations and the bias in input text encoding towards position. These findings present intriguing insights.

strengths:
- **Originality:** To my knowledge, the three-stage circuit proposed by the authors is a novel contribution.
- **Quality:** The hypothesis put forward is reasonable, and the experiments are thorough with a well-crafted methodology.
- **Clarity:** The arguments and evidence presented in the paper are clear, and the experimental descriptions are appropriately detailed.
- **Significance:** Currently, ICL is one of the most important applications in the LLM field, and understanding the mechanisms behind ICL will greatly aid in enhancing its performance.

weaknesses:
The three-stage ICL framework appears to have implicit applicability conditions, which I believe should be clarified.

For example, in Fig. 1 on page 2, a few-shot scenario with $ k=2 $ is presented, which indeed fits the three-stage ICL circuit framework. However, in a zero-shot scenario ($ k=0 $), step 1 may still exist, but steps 2 and 3 would not be applicable. In a few-shot scenario with $ k=1 $, steps 1 and 2 might still apply, but step 3 cannot exist.

Therefore, the framework proposed in this paper should be limited to discussions of scenarios where $ k \geq 2 $. A related question arises: if the focus is restricted to this scenario, what potential issues might emerge?

Furthermore, if we condition on $ k \geq C $ (where $ C $ is a fixed value), could this value vary depending on the problem type? For instance, in tasks like SST-2 and SST-5, which have different label set sizes, might the value of $ C $ differ across these scenarios?

questions:
The three-stage framework proposed in the paper is quite interesting. The questions I have raised are mentioned in the weaknesses section. Here, I would like to know what inspired you to propose this framework. Each part of the framework consists of very specific ideas—were they derived from repeated trial and error, or were they inspired by something else? Alternatively, are they improvements on a significant prior work?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730264535347

thread_type:
official_review

<thread_object>
creator: reviewer_gt4r

<content>
summary:
This paper investigates the mechanisms within large language models (LLMs) that enable in-context learning (ICL) tasks, breaking down the process into three distinct stages: summarization, semantic merging, and feature retrieval/copying. The study employs a variety of experiments across multiple LLMs to validate its findings. Overall, this paper presents valuable insights that can contribute to the field of LLM research, particularly within the ICL community.

strengths:
1. The findings in this paper are clearly explained. Experimental results and visualizations enhance readability and help the audience follow the study's progression easily.

2. This work is well-connected to existing ICL research, with discussions that compare its findings to prior studies on ICL explainability and demonstration selection.

3. The study uses multiple LLMs, strengthening the generalizability of the findings across different model architectures.

4. The insights provided are thought-provoking and have potential practical implications for ICL applications.

weaknesses:
1. [Section 3.1] The authors used mutual nearest-neighbor kernel alignment to evaluate LLMs' summarization abilities. However, the term “summarize” lacks clarity. Does it refer to encoding capabilities similar to those in BGE?

2. [Section 3.1] Additionally, the kernel alignment metric may not be sufficiently robust, as alignment scores in Figure 2 range only from 0.25 to 0.35, which is not significant enough. Consequently, the finding on “summarization” may hold only to a limited extent.

3. [Section 4.1 – Copying from Text Feature to Label Token] It is unclear whether the copying mechanism is applied solely to label tokens or if it extends to other tokens within the input. Using results from other tokens as a baseline could provide a more nuanced understanding of the copying process.

4. [Figure 5, Right] After layer 40, the classification accuracy drops significantly. The authors did not investigate potential reasons for this decline. Could it be due to the gradual degradation of copied information?

5. The experimental setup in Section 5.1 is insufficiently detailed. For instance, how many attention heads are disconnected at each layer? Additionally, the experiments lack certain baselines, such as randomly disconnecting some attention heads to observe the impact on model performance.

Despite these questions and weaknesses, I believe this paper still offers meaningful insights.

questions:
Please see weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730696256496

thread_type:
official_review

<thread_object>
creator: reviewer_z26b

<content>
summary:
The authors propose a three-step inference circuit to capture the in-context learning (ICL) process in large language models (LLMs):

1. Summarize: Each input (both demonstration and query) is encoded into linear representations in the model's hidden states.

2. Semantics Merge: The encoded representation of each demonstration is merged with its label, creating a joint representation for the label and demonstration.

3. Feature Retrieval and Copy: The model retrieves and copies the label representation most similar to the query's representation, using this merged representation to predict the query's label.

This circuit explains various ICL phenomena, such as position bias, robustness to noisy labels, and demonstration saturation. Ablation studies show that removing steps in this process significantly reduces performance, supporting the dominance of this circuit in ICL. The paper also identifies some bypass mechanisms that operate in parallel, assisting the inference circuit through residual connections.

strengths:
1. The authors proposed to use the mutual nearest-neighbor kernel alignment of the intermediate representations of LLMs and sentence embeddings produced by another pre-training model to assess the quality of these representations. This method is novel.

2. Extensive analysis has been performed on all three steps of the proposed framework. Possible explanations have also been provided for many phenomena.

3. The experiments are performed with real-world LLMs and datasets, which makes the insights more likely to be useful in practice.

4. The paper is well-written and easy to follow.

weaknesses:
1. The majority of the analysis is based on associations without verifying their strength and whether those effects are causal. For example, Figure 2 right does not look significant enough for me. The peaks highlighted in Figure 5 also look pretty noisy to me.

2. The causal evidence that the authors provided in the ablation study only shows the effect of deleting the hypothesized important components in ICL. What if unimportant components are deleted? Would they have a similar effect? Only if the unimportant components have a significantly weaker effect on ICL performance, can we draw a causal conclusion that the proposed three-step process dominates ICL.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731247840487

thread_type:
official_review

<thread_object>
creator: reviewer_kes9

<content>
summary:
This paper aims to explain the mechanisms behind in-context learning (ICL) using the inference circuit framework.

According to the authors, the ICL process consists of three internal steps:
1. Summarize: Large language models (LLMs) encode each demonstration within its corresponding forerunner token,  $s_i$ .
2. Semantics Merge: The semantics of each demonstration and its label are combined into the representation of the label  $y_i$ .
3. Feature Retrieval and Copy: LLMs rely on the accumulated labels  $y_{1:k}$  to respond accurately to the query  $s_q$ , yielding the most appropriate answer.

Each step is empirically validated using methods such as kernel alignment and embedding comparisons. The authors also seek to align their findings with those of prior research, reinforcing the credibility of the arguments presented in this work.

strengths:
- Attempts to explain the inner workings of ICL, based on reasonable assumptions and investigative tools.
- The findings align with previous work, encouraging readers to accept the claims presented in the paper.
- Visualized results help readers quickly grasp the core concepts and findings of the paper.

weaknesses:
- While the proposed framework is logical and reasonable, it remains challenging to argue definitively that the core mechanism of ICL follows the assumptions presented in the paper. As noted in Section 5.2, there are exceptions that do not align well with the proposed framework, raising concerns that the explanations may be superficial and fail to capture the essence of ICL. This is understandable, as fully explaining the inner workings of neural networks is inherently difficult, if not nearly impossible.
- I am somewhat unclear about the core novelty of this paper. As I understand it, the primary contribution seems to be the attempt to apply the existing inference circuit framework to the ICL of specific LLMs, including LLaMA 3. In Section 2.1, I did not find explanations that clarify why the procedure conducted in this paper is particularly innovative, compelling, or novel. More comprehensive comparisons with prior work employing induction or inference circuits to illustrate the inner workings of ICL would be helpful to underscore the merits and uniqueness of this study.
- In Section 3.1, sentence embeddings generated by an external encoder (BGE M3) are compared to hidden representations computed by an LLM. Since these two representations come from different models, without any modification or fine-tuning, there is a risk that their vector spaces are not aligned or compatible. This raises concerns about whether this experiment is sufficiently reasonable.

questions:
Please refer to the Weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732192505071

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Revision Notes

comment:
Dear AC and Reviewers,

Thanks for your valuable time in improving our paper. We have tried to address the review comments and submitted a revised version.

Note: Due to the inclusion of new content, we have extended the main body to a 10th page and reformatted the manuscript for better organization. **While this may appear as an extensive change in some PDF comparison tools, we assure you that no content beyond what is explicitly mentioned in these notes has been altered.**


**Revisions According to Review**

(Revised on November 21, 2024)

1. **[Reviewer Z26B Q2, Reviewer gt4r Q5] The baseline value of ablation experiment.**


   *Revision: Reformatting Table 1, 2, 6*


   The captions and small numbers in Tables 1, 2, and 6 have been highlighted with color to clarify that they represent baseline results. Additionally, we have included the standard deviation data for these baseline results, ensuring they no longer appear as the standard deviation of a main result.

2. **[Reviewer gt4r Q5] The detailed settings of ablation experiments.**


   *Revision: Experiment setting added in Appendix A.7*
   
   Appendix A.7 has been added to provide detailed information on the ablation experiment settings.


3. **[Reviewer gt4r Q3] The token selectivity of the copy processing.**


   *Revision: Additional context in Sec 4.1 and Appendix D*


   Our experiments demonstrate that the copying mechanism is inherent to the model and extends to other tokens than the label tokens. These findings are now referenced in Section 4.1 (“Text Representations are Copied without Selectivity”) and detailed in Appendix D.


4. **[Reviewer Z26B Q1] Limitation about the connectivity of attention head.**


   *Revision: Limitation Section*


   Limitation 1 has been expanded to discuss future work focusing on measuring the connection magnitude of model components.


5. **[Reviewer yFjP] Limitation about the application condition.**


   *Revision: Limitation Section, Appendix G*
   
   The original Limitation 2 has been merged into Appendix B.1, and a new Limitation 2 has been added to address Reviewer yFjP's concerns about the application conditions of this work. Further details are provided in the newly added Appendix G.


6. **[Reviewer keS9 Q2] Relationship with Previous Works.**


   *Revision: Additional context “Comparison with Previous Works” in Discussion*
   
   We have added comparisons with previous interpretability research to highlight the novelty and contributions of our work.


7. **[Reviewer gt4r Q4] The reason for accuracy dropping in late layers in Fig. 5 Right.**

   *Revision: Additional context at the end of Sec. 4.1*

   *(November 27) Replaced at the end of the paragraph: "Hidden States of Label Tokens are Joint Representations of Text Encodings and Label Semantics."* 

   We have added a brief description of the accuracy decline, which is aligned with our findings of Fig. 5 Middle.

----

(Revised on November 22, 2024)

8. **[Reviewer gt4r Q1] Term "Summarize" lacks clarity.**

   *Revision: Rename the "Summarize" to "Input Text Encode"*

   We have renamed the "Summarize" to "Input Text Encode" throughout the whole paper, and some expressions have been revised to fit this.

**Other Minor Revisions**


- The Reproducibility Statement / Author Contributions (reserved) / Acknowledgments (reserved) have been moved to the end of the main body. (Revised on November 21)
- Some figures and tables have been moved / resized to make better use of the available space. (Revised on November 21, 27)
- Minor word-level fixes (e.g., typos and grammatical errors) have been made without affecting the semantical content. (Revised on November 21, 22, 27)

We believe we have carefully revised our paper to address the key concerns raised in the reviews, and issues not reflected in the revisions have been discussed in our response in detail. Should there be any further questions or points requiring clarification, we are more than willing to engage in additional discussions.

If you have any other concerns, please kindly let us know.
</content>
</thread_object>

</thread>


Metareview:
metareview:
**Summary:**
The paper investigates the mechanisms behind In-Context Learning in large language models. The authors introduce a three-step inference circuit comprising (1) input text encoding, (2) semantic merge, and (3) feature retrieval and copying. Through empirical analysis, they demonstrate that this circuit effectively captures observed phenomena like positional bias, noise robustness, and demonstration saturation in ICL. Ablation studies reveal the circuit's dominance in driving ICL while identifying auxiliary bypass mechanisms enabled by residual connections. 

**Strength:**
- Explains the inner workings of In-Context Learning (ICL) based on reasonable assumptions and robust investigative tools.
- Extensively analyzes all three steps of the proposed three-stage inference circuit.
- Conducts experiments on real-world LLMs and datasets, ensuring the insights are applicable in practice.

**Weakness:**
- While the proposed framework is logical, it is challenging to argue definitively that the core mechanism of ICL aligns with the assumptions presented.
- Limited clarification why the procedure conducted in this paper is particularly innovative, compelling, or novel.