Understanding In-context Learning of Addition via Activation Subspaces

TL;DR

This paper investigates how transformer language models perform in in-context learning of addition, revealing that a few attention heads with low-dimensional subspaces encode the addition process, and introduces methods to analyze these mechanisms.

Contribution

The paper introduces a novel optimization and analysis framework that localizes and interprets the in-context learning mechanism in transformer models, especially focusing on attention head subspaces.

Findings

Few attention heads encode addition in low-dimensional subspaces.

Identified a self-correction mechanism in the model's in-context learning.

Reduced model complexity to three heads with interpretable subspaces.

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.