# Assign and Add: A Mechanistic Study of Compositional Arithmetic

arXiv: https://arxiv.org/abs/2605.31497 · HTML: https://mlsanity.com/paper/2605.31497/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 25 of its 32 references that are in the map.
- First posted on arXiv 29 May 2026.

## Builds on (most cited first, 25 of 25)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [Measuring and Narrowing the Compositionality Gap in Language Models](https://mlsanity.com/paper/2210.03350/) (2022)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [Faith and Fate: Limits of Transformers on Compositionality](https://mlsanity.com/paper/2305.18654/) (2023)
- [Transformers Learn Shortcuts to Automata](https://mlsanity.com/paper/2210.10749/) (2022)
- [COGS: A Compositional Generalization Challenge Based on Semantic Interpretation](https://mlsanity.com/paper/2010.05465/) (2020)
- [How does GPT-2 compute greater-than?: Interpreting mathematical abilities in a pre-trained language model](https://mlsanity.com/paper/2305.00586/) (2023)
- [Birth of a Transformer: A Memory Viewpoint](https://mlsanity.com/paper/2306.00802/) (2023)
- [Analyzing Transformers in Embedding Space](https://mlsanity.com/paper/2209.02535/) (2022)
- [The mechanistic basis of data dependence and abrupt learning in an in-context classification task](https://mlsanity.com/paper/2312.03002/) (2023)
- [How Transformers Learn Causal Structure with Gradient Descent](https://mlsanity.com/paper/2402.14735/) (2024)
- [Unveiling Transformers with LEGO: a synthetic reasoning task](https://mlsanity.com/paper/2206.04301/) (2022)
- [Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic tasks](https://mlsanity.com/paper/2406.02550/) (2024)
- [Learning Compositional Functions with Transformers from Easy-to-Hard Data](https://mlsanity.com/paper/2505.23683/) (2025)
- [Discovering Variable Binding Circuitry with Desiderata](https://mlsanity.com/paper/2307.03637/) (2023)
- [Feature emergence via margin maximization: case studies in algebraic tasks](https://mlsanity.com/paper/2311.07568/) (2023)
- [Iteration Head: A Mechanistic Study of Chain-of-Thought](https://mlsanity.com/paper/2406.02128/) (2024)
- [Compositional Capabilities of Autoregressive Transformers: A Study on Synthetic, Interpretable Tasks](https://mlsanity.com/paper/2311.12997/) (2023)
- [How Do Transformers Learn Variable Binding in Symbolic Programs?](https://mlsanity.com/paper/2505.20896/) (2025)
- [Alternating Gradient Flows: A Theory of Feature Learning in Two-layer Neural Networks](https://mlsanity.com/paper/2506.06489/) (2025)
- [On the Mechanism and Dynamics of Modular Addition: Fourier Features, Lottery Ticket, and Grokking](https://mlsanity.com/paper/2602.16849/) (2026)
- [Composing Global Solutions to Reasoning Tasks via Algebraic Objects in Neural Nets](https://mlsanity.com/paper/2410.01779/) (2024)

## Abstract (arXiv)

Large language models are able to compose skills in order to perform complex tasks, many of which might not have been seen during training. The details of how exactly this composition occurs remain elusive. In this paper, we study a mechanism for compositional generalization in transformers by considering a simple controlled setting involving variable assignment and modular addition. By partitioning our training data into disjoint sets, we observe that small transformers are able to generalize to previously unseen combinations of variables and numbers. Our mechanistic analysis shows that the same ``modular addition'' MLP module is used whether the inputs are given directly or indirectly through a separate variable assignment mechanism. We also analyze the training dynamics from an empirical lens, which reveals three phases of learning: first, modular addition is learned, then the structure required for variable assignment, and finally a refinement phase where the model generalizes to some hard sequences not seen in training. Finally, we provide a theoretical framework to explain how compositionality emerges from training dynamics. These results suggest that compositional generalization can be a natural consequence of the compositionality of internal mechanisms in~transformers.