📊 statistics

Assign and Add: A Mechanistic Study of Compositional Arithmetic

This paper investigates how small transformers achieve compositional generalization in a controlled setting of variable assignment and modular addition, revealing that the models reuse internal mechanisms for both direct and indirect inputs and progress through distinct learning phases to naturally generalize to unseen combinations.

Original authors: Brady Exoo, Alberto Bietti, John Sous

Published 2026-06-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Brady Exoo, Alberto Bietti, John Sous

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a very smart, but very literal, robot to do math. You want it to solve a puzzle like this: "If 'A' is 5 and 'B' is 10, what is A + B?"

The tricky part is that the robot has never seen the specific letters "A" and "B" used in that exact order before. It has only seen numbers added together (like 5 + 10) or other letters used in different ways. The big question the researchers asked is: How does the robot learn to combine these two different skills—understanding that a letter stands for a number, and then actually doing the math—to solve a brand-new puzzle it hasn't seen yet?

Here is what the paper "Assign and Add" discovered, explained simply:

1. The Setup: A Two-Step Dance

The researchers built a small, simple "brain" (a neural network) and gave it a specific job. They fed it sentences like:

c=17, b=42, b+19=?

The robot had to figure out that b is actually 42, so the math is really 42 + 19, and then give the answer.

To see how the robot learned, they split the training into different "flavors" of puzzles:

The Easy Way: Just numbers (e.g., 17 + 19 = ?).
The Medium Way: One letter and one number (e.g., b+19 = ?).
The Hard Way: Two letters (e.g., b+c = ?).

2. The Three Phases of Learning

The robot didn't learn everything at once. It went through three distinct phases, like a student mastering a subject step-by-step:

Phase 1: The Math Master. First, the robot learned how to do the actual addition. It got really good at adding numbers together. At this stage, it was great at 17 + 19 but completely confused by letters. It was like a calculator that didn't know what a variable was.
Phase 2: The Translator. Next, the robot learned the "assignment" skill. It figured out how to look at c=17 and remember that c means 17. It learned to swap the letter for the number.
Phase 3: The Conductor. Finally, the robot learned to put the two skills together. It realized, "Oh, I can use my 'Translator' skill to find the numbers, and then use my 'Math Master' skill to add them."

The Surprise: The robot learned the math before it fully learned the translation. It could do the addition perfectly even while it was still struggling to understand the letters!

3. How the Robot's "Brain" Works (The Mechanism)

The researchers didn't just look at the final score; they looked inside the robot's "brain" to see how it was thinking. They found two specific tools working together:

Tool A: The "Previous Token" Head (The Translator).
Imagine a robot arm that reaches back to grab the item immediately before the current one. In the robot's first layer, it learned to use this arm to grab the number assigned to a letter. If it sees b, it reaches back to find the number b was assigned to earlier in the sentence.
Tool B: The "Fourier" Math Module (The Calculator).
In the second layer, the robot uses a special mathematical trick (involving waves and patterns) to do the addition. This is a known way for robots to do modular math (math that wraps around, like a clock).

The Magic Moment: The robot learned to connect these two tools. The "Translator" arm grabs the correct numbers and passes them to the "Calculator" module. Once connected, the robot can solve b + c even if it has never seen b and c added together before. It's like having a translator who speaks "Letter" and a calculator who speaks "Number," and finally building a bridge between them.

4. Why This Matters

The paper shows that generalization (solving new problems) isn't magic. It happens naturally when a model learns small, atomic skills (like adding numbers and finding variables) and then figures out how to wire them together.

The "Grokking" Effect: The researchers noticed a phenomenon called "grokking." Sometimes the robot memorizes the training data perfectly but fails on new tests. Then, suddenly, it "clicks," and it starts understanding the rules instead of just memorizing answers.
The Result: Because the robot learned the internal mechanisms separately and then combined them, it could handle puzzles it had never seen, like using a letter in a position it wasn't allowed to be in during training.

Summary Analogy

Think of the robot as a chef.

First, it learns how to chop vegetables (the math).
Then, it learns how to read a recipe card that says "Use the onion from the blue bowl" (the variable assignment).
Finally, it learns to combine these: Read the card, find the onion, and chop it.

The paper proves that if you teach a chef these two skills separately, they will eventually figure out how to combine them to cook a dish you've never asked them to make before. The "composition" of the skills is what makes the robot smart.

Technical Summary: Assign and Add: A Mechanistic Study of Compositional Arithmetic

Problem Statement

Large language models (LLMs) demonstrate an ability to compose skills to perform complex tasks unseen during training, yet the specific mechanisms enabling this "compositional generalization" remain elusive. While generalization is often attributed to the composition of atomic mechanisms, the interplay between distinct mechanisms—specifically variable binding (associative recall via induction heads) and modular arithmetic (structured representations via Fourier features)—has not been systematically studied in a controlled setting.

This paper addresses the question: How do transformers learn to compose variable binding and modular addition mechanisms to generalize beyond the training distribution, and how do these capabilities emerge during training?

Methodology

The authors employ a controlled experimental setup involving a small, mechanistically interpretable transformer trained on a synthetic task combining variable assignment and modular addition.

Model Architecture: A 2-layer, single-head transformer with $d_{model} = d_{head} = 128$ and $d_{mlp} = 512$ . The architecture deliberately omits biases, normalization layers, and the Layer-1 MLP to simplify circuit discovery. It uses GPT-2 style absolute positional embeddings and causal masking.
Task Definition: The model is trained to perform modular addition ( $N=59$ ) where operands can be constants or variables dynamically assigned in the context (e.g., c=1, b=17, a=42, b+19=?).
Data Distribution: Sequences are tokenized with random padding to prevent reliance on absolute positions. The training set is partitioned to evaluate generalization:
- 0-variable (0-var): Pure constant addition.
- 1-variable (1-var) & 2-variable (2-var): Addition involving assigned variables.
- Held-out Sets: Approximately 30% of addition pairs are removed from training (e.g., if 1+3 is unseen, all sequences containing it are excluded). Additionally, specific variables are restricted from certain positions in the training set to test positional generalization.
Analysis Techniques: The study utilizes mechanistic interpretability tools (via the TransformerLens library) to analyze attention patterns, residual stream vectors, and weight matrices. It also employs linear probes and theoretical gradient analysis to track training dynamics.

Key Contributions

Tractable Model for Compositional Arithmetic: The authors present a simplified model that allows for the mechanistic interpretation of how variable assignment composes with modular arithmetic.
Decoupled Training Dynamics: The analysis reveals that modular arithmetic is learned prior to and independently of variable assignment. The model achieves near-perfect accuracy on 0-var tasks while still learning to perform variable binding.
Mechanistic Explanation of Generalization: The paper identifies that generalization arises from the composition of two specific circuits:
- A modified induction head in Layer 1 that retrieves assigned variables and routes their values.
- A modular addition circuit in Layer 2 (specifically the Output-Value matrix and the Layer-2 MLP) that performs the arithmetic on the retrieved values.
Theoretical Framework: The authors provide a theoretical analysis demonstrating how gradient dynamics facilitate the composition of variable binding and modular addition circuits, showing that the negative gradient on the Layer-2 Query-Key matrix naturally leads to the formation of an associative memory block.

Results

Training Dynamics and Phases

The training process exhibits three distinct phases:

Modular Addition Learning: The model first learns to perform addition on sequences with no variables (0-var), achieving high accuracy before variable binding is fully established.
Variable Assignment Learning: The model subsequently learns the structure required for variable assignment. A sharp "spike" in accuracy (around step 5,300–5,400) coincides with the emergence of the variable-assignment routing circuit, where the Layer-1 attention correctly routes constants to their assigned variables.
Refinement Phase: A second phase (steps 14,500–16,000) involves refining the routing for "hard" sequences (e.g., specific variables in held-out positions), correcting initial failures to attend to certain variable positions.

Generalization Capabilities

Positional Generalization: The model generalizes to novel variable orderings (e.g., predicting b+14 after seeing 14+b).
Out-of-Distribution (OOD) Generalization: The model successfully predicts held-out addition pairs (e.g., 1+3) that were never seen in any form during training, achieving 99.4% accuracy on curated test sets.
Data Requirements: Generalization to unseen constant pairs requires training on at least ~25% of the total distribution ( $f \approx 0.25$ ). Generalization to 2-variable sequences requires a relative frequency of ~20% ( $r \approx 0.2$ ) of 2-var sequences in the training data.

Mechanistic Findings

Shared Representation: The residual stream vectors for variable-assigned addition and pure constant addition are nearly identical (cosine similarity $\approx 1$ ) before the Layer-2 MLP, indicating a shared computational mechanism.
Fourier Features: The Layer-2 MLP neurons exhibit preactivations matching the theoretical form of Fourier features ( $\cos(\frac{2\pi k}{N}n + \alpha_k) + \cos(\frac{2\pi k}{N}m + \beta_k)$ ), confirming the use of established periodic mechanisms for modular arithmetic.
Induction Head Modification: The Layer-1 attention acts as a "previous-token head" for constants, attending to the variable immediately preceding them. The Layer-2 Query-Key matrix then forms a strong identity pattern on variable embeddings, allowing the = token to attend directly to the correct constants associated with the variables.

Significance

The paper claims that compositional generalization in transformers can be a natural consequence of the compositionality of internal mechanisms. By demonstrating that distinct mechanisms (induction heads and Fourier-based arithmetic) can be learned independently and then composed via attention routing, the work provides a mechanistic basis for understanding how transformers generalize to complex reasoning tasks.

The authors maintain modesty regarding their scope, noting that the study is limited to a small, controlled setting to facilitate analysis. They do not claim these findings directly explain all LLM behaviors but suggest they offer a framework for understanding how generalizing capabilities arise in transformers, potentially motivating future investigations into more complex tasks and scaling laws.