Counting Hypothesis: Potential Mechanism of In-Context Learning

TL;DR

This work proposes the Counting Hypothesis to explain In-Context Learning in pretrained LLMs, suggesting that fixed transformer computations encode context-dependent subspaces in residual streams that allow competing potential answers to be inferred from prompt examples. It argues FFNs function as associative memories storing keys and values, enabling parallel encoding of multiple candidate answers, and that last-separator predictions leverage common subspace components shared with answer tokens. The authors support this with dictionary learning and ICA analyses across six models, showing shared components between separators and answers and increasing alignment as depth grows; they also demonstrate a diagnostic embedding signal (ratio R) that distinguishes correct from incorrect predictions. If validated, this framework offers a mechanistic lens on ICL, enabling reliability diagnostics and guiding future architectural and training strategies to enhance in-context adaptability without fine-tuning.

Abstract

In-Context Learning (ICL) indicates that large language models (LLMs) pretrained on a massive amount of data can learn specific tasks from input prompts' examples. ICL is notable for two reasons. First, it does not need modification of LLMs' internal structure. Second, it enables LLMs to perform a wide range of tasks/functions with a few examples demonstrating a desirable task. ICL opens up new ways to utilize LLMs in more domains, but its underlying mechanisms still remain poorly understood, making error correction and diagnosis extremely challenging. Thus, it is imperative that we better understand the limitations of ICL and how exactly LLMs support ICL. Inspired by ICL properties and LLMs' functional modules, we propose 1the counting hypothesis' of ICL, which suggests that LLMs' encoding strategy may underlie ICL, and provide supporting evidence.

Figures (13)

• Figure 1: Cosine distance ($D$) between $300$ atoms obtained from dictionary learning. $x$-axis and $y$-axis denote $v^S_i$ and $v^A_j$.
• Figure 2: Cosine distance $D$ between $20$ independent components obtained from ICA. $x$-axis and $y$-axis denote $v^S_i$ and $v^A_j$, respectively.
• Figure 3: The minimum distance between $v^S_i$ and $v^A_j$. (A), $D'$ between independent components from separators and answer tokens. (B), $D$ estimated using dictionary learning.
• Figure 4: Alignment between the last separator and answer tokens, which is analyzed by dictionary learning. The residual streams are recorded from LLMs performing the 'country-capital' task. We choose the top $20$ atoms out of $300$ for 6 models and show their coding coefficients across layers shown in $x$-axis. $y$-axis denotes the top $20$ atoms. The names of the models are displayed above the plots.
• Figure 5: Alignment between the last separator and answer tokens, which is analyzed by ICA. The residual streams are recorded from LLMs performing the 'country-capital' task. Each plot shows the coding coefficients of $20$ independent components for 6 models. $x$-axis and $y$-axis denote independent components and layers, respectively. As ICA is not sensitive to the component and consequently coding coefficient, we display the absolute value of coding coefficients. We present the alignment observed during different tasks in Supplemental Figs. \ref{['fig:sub2']}, \ref{['fig:sub3']}, \ref{['fig:sub4']}, \ref{['fig:sub5']} and \ref{['fig:sub6']}.