Assign and Add: A Mechanistic Study of Compositional Arithmetic

Unlocking Compositionality: How Transformers Learn to Bind Variables and Perform Arithmetic

AI models learn complex tasks in stages. First, they learn how to do math. Then, they learn how to remember which numbers belong to which labels. Finally, they combine these two skills to solve new problems they have not seen during training.

In the field of mechanistic interpretability (the study of reverse-engineering the internal weights of neural networks), a major open question is how "compositionality" emerges. We know that large language models (LLMs) can tackle complex tasks by combining smaller, atomic skills. However, the precise mechanism by which these disparate circuits integrate remains elusive. Specifically, we do not fully understand how a model moves from simple arithmetic to performing that same arithmetic on variables dynamically assigned in the context.

This paper provides a mechanistic answer. The authors demonstrate that compositional generalization is not a monolithic emergence. Instead, it is a result of integrating two independent, pre-learned modules: a modular addition circuit and a variable assignment mechanism.

The Problem

Current understanding of LLM generalization often treats it as a black box. While we know models can perform out-of-distribution (OOD) tasks (tasks involving data distributions different from the training set), we lack a clear picture of how functional circuits actually integrate. Specifically, we need to know how associative recall (retrieving information) and algorithmic reasoning (performing math) shake hands in the residual stream (the primary communication channel between transformer layers).

When a model encounters a sequence like c=1, b=17, a=42, b+19=?, it must perform two distinct operations. It must bind the variable b to the value 17. It must then feed that value into an addition circuit. If these circuits fail to compose, the model fails the task. This happens even if the model is a "math expert" or a "memory expert" in isolation. Previous research has looked at these phenomena separately. There has been little clarity on the training dynamics that force these two different types of logic to merge.

How It Works

The researchers use a 2-layer transformer trained on a controlled task of modular addition (modulo $N=59$) combined with variable assignment. The architecture is intentionally stripped down. It uses a single attention head and omits the first-layer MLP (multi-layer perceptron, a type of feed-forward neural network) to make circuit discovery tractable.

The composition happens through a specific interplay of two circuits:

1. The Variable Assignment Module: This is implemented via a modified induction head (a circuit that helps models copy information). Traditionally, induction heads implement a copying mechanism. Here, the authors find that the first layer develops a "previous-token head" behavior. Constant tokens look back at the variable assigned to them. This effectively remaps the variable's identity into the residual stream .

2. The Modular Addition Module: This is located in the second-layer MLP. The authors show that the MLP learns to represent modular arithmetic using Fourier features (periodic, sinusoidal representations). These allow the network to compute sums through constructive interference .

The "magic" of composition lies in the routing. The second-layer attention (the QK circuit, which determines where information is sent) acts as the glue. It identifies the correct operands—whether they are raw constants or retrieved variable values. It then routes them to the second-layer MLP. The authors prove that the same MLP module is used regardless of whether the input is a direct constant or an indirect variable. This confirms the model has truly composed the two skills.

Numbers

The authors report that the model achieves over 98% test accuracy on complex 2-variable sequences. This means it succeeds even on unseen combinations of variables and numbers. However, this capability is gated by specific data requirements.

The paper finds that generalization is highly sensitive to the density of the training distribution: * Arithmetic Density: To generalize to the entire 0-variable (pure constant) set, the model requires training on at least $f \approx 0.25$ of all possible addition pairs [Figure 3b]. This means the model needs to see roughly a quarter of all possible math pairs to master basic arithmetic. * Task Complexity: To successfully generalize to 2-variable sequences, the model requires a relative frequency of $r \approx 0.2$ in the training set [Figure 3a]. This indicates the model needs a minimum threshold of complex examples to learn the composition.

Crucially, the training dynamics reveal a clear hierarchy of learning. The model first masters 0-variable modular addition. Then, it builds the structure for variable assignment. Finally, it enters a "refinement" phase to resolve difficult routing cases .

This suggests that arithmetic competence is a prerequisite for variable binding in this architecture.

What's Missing

While the mechanistic insights are deep, the scope is narrow.

First, the study relies on a highly controlled, low-complexity synthetic task. This is necessary for the mathematical proofs. However, it is a far cry from the messy, high-dimensional reasoning required in natural language. We do not know if these clean "circuit handshakes" survive the introduction of linguistic ambiguity or massive scale.

Second, the theoretical analysis assumes a "disentangled" transformer setup with orthogonal embeddings (vectors that are perpendicular and do not overlap). This is a significant idealization. In a real-world production model, embeddings are dense and highly correlated. The assumption that the gradient is dominated by a single associative-memory term might not hold when the noise from other features is higher.

Finally, the paper does not address scaling laws. It shows what happens in a 2-layer, 128-dimension model. It leaves open whether these circuits become more or less modular as the model depth and width increase.

Should You Prototype This

Not yet.

If you are looking for a way to improve the reasoning capabilities of a production LLM, this paper does not provide a new loss function or a better architecture to ship. It is a foundational piece of mechanistic interpretability. It explains why current models behave the way they do.

However, if you are building specialized symbolic reasoners, the methodology is invaluable. The researchers released their code under the MIT License (available via MechCompose). If you want to test if your own small-scale models are developing these modular circuits, the codebase is a solid starting point. For now, treat this as a map of the terrain, not a blueprint for a new engine.

Figures from the paper