Counting Hypothesis: Potential Mechanism of In-Context Learning

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.

💡 Research Summary

The paper tackles the still‑mysterious mechanism behind In‑Context Learning (ICL) in large language models (LLMs). While prior work has suggested Bayesian inference or dual‑gradient descent as possible explanations, none have provided a concrete, model‑level account of how a frozen transformer can extract task information from a few prompt examples and apply it to a new query.
The authors propose the “counting hypothesis.” Their central claim is that the feed‑forward networks (FFNs) inside each transformer layer act as associative memories that store a set of possible next‑token answers as key‑value pairs. Because every transformer layer adds its output to the residual stream, the representations of all possible answers accumulate across layers, forming a superposition of candidate answers.
Three core assumptions underlie the hypothesis:

1. FFN decomposition – the output of an FFN can be expressed as a linear combination of possible answers for the current input; these candidates then compete in later layers.
2. Constant attention weights – all examples in an ICL prompt are equally important, so the self‑attention scores can be approximated by constants that depend only on token type (question, separator, answer).
3. Context‑dependent subspaces – different semantic contexts (e.g., geographic vs. fashion) are encoded in distinct subspaces of the model’s d‑dimensional hidden space, allowing the model to keep multiple answer candidates separate.
   With these assumptions, the authors analytically derive two empirical properties of ICL. First, the order of examples has little effect on the final answer because the attention contribution of each example collapses to a constant term. Second, adding more examples improves accuracy because the shared “common component” of the residual streams is summed more often, amplifying the subspace that corresponds to the correct context. In other words, the model “counts” how many times each context‑subspace appears in the prompt and selects the answer from the most frequently reinforced subspace.
   To test the hypothesis, the authors collect residual streams from six publicly available LLMs (GPT‑j‑6B, LLaMA‑3.1‑8B, OLMo‑2‑0325‑32B, Phythia‑12B/6.9B, GPT‑NEOX‑20B) across 200 ICL prompts covering six tasks (antonyms, synonyms, country‑capital, English‑French, product‑company, person‑sport). They apply two unsupervised component‑extraction techniques: dictionary learning (which yields “atoms” without imposing orthogonality) and Independent Component Analysis (ICA). Both methods decompose the normalized residual vectors into a set of basis vectors.
   The analysis shows that separator tokens (the colon after “A:”) and answer tokens share a substantial number of identical atoms and independent components, indicating that consecutive tokens indeed occupy overlapping subspaces. Moreover, the attention scores derived from the data are approximately constant across token types, supporting assumption 2. These findings corroborate the counting hypothesis: the model stores multiple potential answers in its FFNs, aggregates them through residual connections, and selects the answer that receives the highest “count” from the context‑specific subspace.
   The paper acknowledges limitations. If attention weights deviate significantly from constancy, the counting process becomes noisy, which may explain higher error rates in some models. The analysis focuses on linear combinations of residual streams and does not fully capture nonlinear interactions introduced by activation functions. Finally, only pre‑trained (not fine‑tuned) models are examined, leaving open how adaptation or architectural changes would affect the hypothesis.
   Overall, the work offers a concrete, mathematically grounded framework for understanding how frozen transformers can perform ICL. By linking FFN associative memory, residual superposition, and context‑dependent subspace counting, it provides a fresh perspective that bridges empirical observations (order invariance, example‑number scaling) with the internal dynamics of modern LLMs. This contributes valuable insight for researchers aiming to diagnose, improve, or extend ICL capabilities.