Research map / Architectures and efficiency / Attention and transformers

Assign and Add: A Mechanistic Study of Compositional Arithmetic

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Transformer expressivity and learnability · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Mechanistically studies a 2-layer single-head transformer trained on variable assignment plus modular addition, showing it composes an induction-like variable-routing circuit with a Fourier-feature modular addition MLP and gives a one-step gradient analysis of how this composition emerges.

Problem

How transformers learn to compose variable binding (associative recall/induction-style) with modular arithmetic to generalize to unseen combinations, and how this composition emerges during training.

Approach

Uses a controlled synthetic task where addition operands can be variables assigned in context, withholding certain addition pairs and variable positions from training, and reverse-engineers the resulting circuit at the weight level, complemented by progress measures and a one-gradient-step theoretical analysis of the second-layer key-query matrix.

Key results

• Small 2-layer single-head transformer reaches over 98% test accuracy across 0-var, 1-var, and 2-var addition types.
• Isolated OV-MLP circuit W_U(MLP_2(OV_2(a)+OV_2(b))) achieves 99.94% accuracy on all constant pairs.
• Generalization to unseen constant pairs requires training on at least about 25% of addition pairs; two-variable generalization requires relative frequency r of about 0.2.

Surprising findings

• The model learns 0-variable modular addition to near-perfect accuracy before it learns variable assignment routing, and variable-binding accuracy emerges later.

Limitations

• Only a single, very simple task and small model are studied, which limits generality.
• Theoretical analysis relies on simplified disentangled embeddings, orthogonality assumptions, and a small-output-scale approximation.
• Mechanistic analysis focuses on a single training run.

Experimental setup

Datasets: synthetic variable-assignment and modular addition sequences (mod 59, 12 variables)

Benchmarks: held-out synthetic test sets (0-var, 1-var, 2-var, add-restricted, var-restricted)

Metrics: accuracy cosine similarity linear probe accuracy

Builds on: induction heads grokking Fourier features for modular addition TransformerLens AdamW

Closest papers by meaning

• Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic tasks (NeurIPS 2024 · similarity 0.91)
• Teaching Arithmetic to Small Transformers (ICLR 2024 · similarity 0.89)
• Limits of Transformer Language Models on Learning to Compose Algorithms (NeurIPS 2024 · similarity 0.88)
• Task Generalization with Autoregressive Compositional Structure: Can Learning from $D$ Tasks Generalize to $D^T$ Tasks? (ICML 2025 · similarity 0.88)
• Compositional Capabilities of Autoregressive Transformers: A Study on Synthetic, Interpretable Tasks (ICML 2024 · similarity 0.87)
• Extrapolation by Association: Length Generalization Transfer In Transformers (NeurIPS 2025 · similarity 0.87)