Understanding In-context Learning of Addition via Activation Subspaces

TL;DR

The paper tackles how in-context learning emerges in forward passes of large transformers by studying an arithmetic add-$k$ task. It introduces activation patching and a sparse-head optimization to localize ICL to a few heads, revealing a six-dimensional aggregator subspace per head that encodes unit-digit and magnitude information via periodic patterns, with a parity direction, and a self-correction mechanism across demonstrations. The key contributions include a novel optimization method identifying 33 significant heads (reduced to three critical heads in practice), a PCA-based subspace localization, a decomposition into interpretable periodic feature directions, and a causal link between extractor and aggregator subspaces that explains how signals flow across tokens. The findings illuminate how fine-grained, low-dimensional activation subspaces support ICL in large models, offering a principled lens for analyzing and potentially engineering ICL mechanisms in future systems.

Abstract

To perform few-shot learning, language models extract signals from a few input-label pairs, aggregate these into a learned prediction rule, and apply this rule to new inputs. How is this implemented in the forward pass of modern transformer models? To explore this question, we study a structured family of few-shot learning tasks for which the true prediction rule is to add an integer $k$ to the input. We introduce a novel optimization method that localizes the model's few-shot ability to only a few attention heads. We then perform an in-depth analysis of individual heads, via dimensionality reduction and decomposition. As an example, on Llama-3-8B-instruct, we reduce its mechanism on our tasks to just three attention heads with six-dimensional subspaces, where four dimensions track the unit digit with trigonometric functions at periods $2$, $5$, and $10$, and two dimensions track magnitude with low-frequency components. To deepen our understanding of the mechanism, we also derive a mathematical identity relating ``aggregation'' and ``extraction'' subspaces for attention heads, allowing us to track the flow of information from individual examples to a final aggregated concept. Using this, we identify a self-correction mechanism where mistakes learned from earlier demonstrations are suppressed by later demonstrations. Our results demonstrate how tracking low-dimensional subspaces of localized heads across a forward pass can provide insight into fine-grained computational structures in language models.

Figures (16)

• Figure 1: Key findings of our methods in the specific case of Llama-3-8B-instruct (illustrated using an example add-k prompt): (1) out of 1024 attention heads, only three are important for add-k ICL (§\ref{['sec:agg']}); (2) each head encodes the task information $k$ in a six-dimensional subspace (§\ref{['sec:6dim']}); (3) the six-dimensional subspace further factors into four dimensions capturing the unit digit of $k$ (encoding periodic functions at periods 2,5, and 10) and two dimensions capturing its tens digit (encoding higher frequency functions) (§\ref{['sec:4+2dim']}); and (4) task information flows from output tokens to final token via a self-correction mechanism (§\ref{['sec:extract']}).
• Figure 2: The chain of localization in §\ref{['sec:id']} and §\ref{['sec:agg']}. We first identify $33$ significant attention heads (out of 1024) via a global optimization method (§\ref{['sec:33head']}), then narrow down to 3 main heads while mean-ablating the remaining 30 (§\ref{['sec:3head']}). We next study the structure of the representation of each main head by localizing it to a six-dimensional subspace (§\ref{['sec:6dim']}), and decompose it into a four-dimensional subspace encoding the unit digit and a two-dimensional subspace encoding the tens digit (§\ref{['sec:4+2dim']}).
• Figure 3: Comparison of significant heads identified by our optimized coefficients (left) and by average indirect effects (AIE) from the previous method fv (right). Colors indicate the magnitude of each head's importance (coefficients or AIE) on Llama-3-8B-instruct. The top $33$ heads identified by both methods are highlighted with frames ($13$ heads common across both methods in red and other $20$ heads in blue). Our identified heads yield an intervention accuracy of $0.85$, compared to the previous method's accuracy of $0.31$. Both methods select heads from similar layers, but our optimization approach is significantly more effective.
• Figure 4: Explained variance ratio vs. number of PCs for each head across models. The first six PCs make up most of the explained variance (97%) for Llama-3 models, and the first eight PCs do so for Qwen-2.5-7B.
• Figure 5: Coordinates of three heads’ vectors (inner products with PCs) for the first six PCs across different add-$k$ tasks on Llama-3-8B-instruct. Periodic patterns are visible in the first few PCs of each head.