Revisiting In-context Learning Inference Circuit in Large Language Models

Abstract: In-context Learning (ICL) is an emerging few-shot learning paradigm on LLMs (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 LLMs. 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, LLM 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.

• The paper presents a novel three-step inference circuit that integrates summarization, semantic merging, and feature copying for enhanced in-context learning.
• Experiments and ablation studies demonstrate that disabling any circuit component severely degrades few-shot learning performance.
• The research also uncovers bypass mechanisms, offering deeper insights into the operational nuances of large language models.

The paper "Revisiting In-context Learning Inference Circuit in LLMs" explores the intricacies of in-context learning (ICL), a prominent few-shot learning paradigm in LLMs. While previous works have attempted to unravel the inner workings of ICL, they have often fallen short in capturing the full range of inference phenomena these models exhibit. This study proposes a novel and comprehensive inference circuit to address these limitations and better explain the observed behaviors in ICL.

The proposed inference circuit is divided into three main operations:

1. Summarize: LLMs encode each input text, including both demonstrations and queries, into a linear representation within their hidden states. This encoding contains sufficient information to tackle ICL tasks effectively.
2. Semantics Merge: This step involves merging the encoded representations of demonstrations with their respective label tokens. As a result, it creates joint representations that unify the labels and demonstrations, which help in understanding and addressing the tasks at hand.
3. Feature Retrieval and Copy: Here, the LLMs search through the joint representations to find similarities with the query representation within a task-specific subspace. The models then copy these relevant representations into the query context. Finally, the LLM's heads leverage the copied label representations and decode them into the predicted labels.

The authors demonstrate that their proposed circuit captures many phenomena observed during the ICL process, offering a comprehensive and practical understanding of how LLMs perform in-context learning. Furthermore, through ablation studies, they show that if any of these steps are disabled, the ICL performance is significantly impaired, indicating that the inference circuit proposed is a crucial mechanism. Additionally, the study identifies some bypass mechanisms that can solve ICL tasks alongside the proposed circuit, adding another layer of understanding to how these models operate.

This work enriches the understanding of ICL in LLMs by presenting a detailed model of its inner workings, which could potentially guide future research and optimization in few-shot learning approaches.