Understanding In-context Learning of Addition via Activation Subspaces
Xinyan Hu · Kayo Yin · Michael I. Jordan · Jacob Steinhardt · Lijie Chen
We mechanistically interpret few-shot in-context learning of addition in large language models by developing a general pipeline to localize its computation to three attention heads with six-dimensional subspaces.
摘要
To perform few-shot learning, language models extract signals from a few input-label pairs, aggregate them 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 method that localizes the model's few-shot learning ability to only a few attention heads. We find that this method and the localization finding generalize to four additional task families spanning arithmetic and semantic tasks. We then perform an in-depth analysis of individual heads via dimensionality reduction and decomposition. For example, in Llama-3-8B-Instruct, we reduce the mechanism underlying these tasks to just three attention heads with six-dimensional subspaces, in which four dimensions track the unit digit using trigonometric functions with periods $2$, $5$, and $10$, while two dimensions track magnitude using low-frequency components. To deepen our understanding of this mechanism, we also derive a mathematical identity relating the “aggregation” and “extraction” subspaces of attention heads, allowing us to track the flow of information from individual examples to a final aggregated concept. Our results demonstrate how tracking low-dimensional subspaces of localized heads throughout a forward pass can provide insight into fine-grained computational structures in language models. Our code is available at [https://github.com/xyVickyHu/addition-subspaces](https://github.com/xyVickyHu/addition-subspaces).