Researchers map how transformers learn to add by tracking six hidden dimensions

When a transformer model learns from a handful of examples to perform a new task, it must somehow extract patterns from limited input and generalize them immediately. Researchers have now pinpointed exactly where and how this "few-shot learning" happens—at least for a fundamental operation like addition.

In a new arXiv paper, researchers investigated how language models implement in-context learning by examining what happens inside transformer networks during the forward pass. Rather than treating the model as a black box, they developed techniques to locate specific attention heads responsible for learning tasks and then analyzed the geometric structure of the information flowing through them.

Localizing the learning to a few key heads#

The team started by studying a clean test case: models trained to add an integer to an input number. Through their localization method, they discovered that this capability wasn't distributed across many parts of the model. Instead, just a handful of attention heads—as few as three in Llama-3-8B-Instruct—carried out the entire computation.

This finding generalizes beyond simple arithmetic. The same principle held for four other task families spanning both mathematical and semantic domains, suggesting that transformer models use a modular strategy for few-shot adaptation.

Decoding the hidden geometry#

The deeper insight came from analyzing the internal structure of these crucial heads. Using dimensionality reduction and decomposition techniques, the researchers discovered that the three heads in Llama-3 operated within just six-dimensional subspaces—far simpler than the model's full representational capacity.

Within this compressed space, the team identified an elegant decomposition: four dimensions track the units digit of the answer using trigonometric functions with specific periods (2, 5, and 10), while the remaining two dimensions encode magnitude using low-frequency components. This mathematical structure mirrors how humans might internally represent multi-digit addition.

Tracking information flow#

To connect individual examples to the final learned rule, the researchers derived a mathematical identity linking what they call "aggregator" and "extractor" subspaces across attention heads. This relationship allowed them to trace how information from separate input-label pairs flows through the model and gets consolidated into a unified prediction mechanism—providing a more complete picture of the computational pipeline.