Assign and Add: A Mechanistic Study of Compositional Arithmetic

• URL: http://arxiv.org/abs/2605.31497v1
• Date: Fri, 29 May 2026 16:19:20 GMT
• Title: Assign and Add: A Mechanistic Study of Compositional Arithmetic
• Abstract summary: We study a mechanism for compositional generalization in transformers.<n>We observe that small transformers are able to generalize to previously unseen combinations of variables and numbers.<n>These results suggest that compositional generalization can be a natural consequence of the compositionality of internal mechanisms intransformers.
• Score: 17.098170202572348
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Understanding the Staged Dynamics of Transformers in Learning Latent Structure [5.944972519558522]
  We train a small decoder-only transformer on three task variants.<n>We show that the model acquires capabilities in discrete stages.<n>We also identify a crucial asymmetry, where the model can compose fundamental rules robustly, but struggles to decompose complex examples to discover the fundamental rules.
  arXiv  Detail & Related papers  (2025-11-24T17:20:42Z)
• Layer Specialization Underlying Compositional Reasoning in Transformers [6.581088182267414]
  Transformers exhibit compositional reasoning on sequences not observed during training.<n>We investigate this phenomenon using the Random Hierarchy Model (RHM), a probabilistic context-free grammar.<n>Mechanistically, we identify a progressive emergence of layer specialization during training that correlates with generalization performance.
  arXiv  Detail & Related papers  (2025-10-20T12:08:22Z)
• How do Transformers Learn Implicit Reasoning? [67.02072851088637]
  We study how implicit multi-hop reasoning emerges by training transformers from scratch in a controlled symbolic environment.<n>We find that training with atomic triples is not necessary but accelerates learning, and that second-hop generalization relies on query-level exposure to specific compositional structures.
  arXiv  Detail & Related papers  (2025-05-29T17:02:49Z)
• Interpreting Affine Recurrence Learning in GPT-style Transformers [54.01174470722201]
  In-context learning allows GPT-style transformers to generalize during inference without modifying their weights. This paper focuses specifically on their ability to learn and predict affine recurrences as an ICL task. We analyze the model's internal operations using both empirical and theoretical approaches.
  arXiv  Detail & Related papers  (2024-10-22T21:30:01Z)
• Principled Understanding of Generalization for Generative Transformer Models in Arithmetic Reasoning Tasks [5.522116934552708]
  Transformer-based models excel in various tasks but their generalization capabilities, especially in arithmetic reasoning, remain incompletely understood.<n>This paper develops a unified theoretical framework for understanding the generalization behaviors of transformers in arithmetic tasks.
  arXiv  Detail & Related papers  (2024-07-25T11:35:22Z)
• Strengthening Structural Inductive Biases by Pre-training to Perform Syntactic Transformations [75.14793516745374]
  We propose to strengthen the structural inductive bias of a Transformer by intermediate pre-training. Our experiments confirm that this helps with few-shot learning of syntactic tasks such as chunking. Our analysis shows that the intermediate pre-training leads to attention heads that keep track of which syntactic transformation needs to be applied to which token.
  arXiv  Detail & Related papers  (2024-07-05T14:29:44Z)
• Initialization is Critical to Whether Transformers Fit Composite Functions by Reasoning or Memorizing [10.206921909332006]
  Transformers have shown impressive capabilities across various tasks, but their performance on compositional problems remains a topic of debate.<n>We discover that the parameter initialization scale plays a critical role in determining whether the model learns inferential (reasoning-based) solutions.<n>We further find that inferential (reasoning-based) solutions exhibit low complexity bias, which we hypothesize is a key factor enabling them to learn individual mappings for single anchors.
  arXiv  Detail & Related papers  (2024-05-08T20:23:24Z)
• Towards Empirical Interpretation of Internal Circuits and Properties in Grokked Transformers on Modular Polynomials [29.09237503747052]
  Grokking on modular addition has been known to implement Fourier representation and its calculation circuits with trigonometric identities in Transformers.<n>We show that the transferability among the models grokked with each operation can be only limited to specific combinations.<n>Some multi-task mixtures may lead to co-grokking, where grokking simultaneously happens for all the tasks.
  arXiv  Detail & Related papers  (2024-02-26T16:48:12Z)
• Discovering modular solutions that generalize compositionally [55.46688816816882]
  We show that identification up to linear transformation purely from demonstrations is possible without having to learn an exponential number of module combinations. We further demonstrate empirically that meta-learning from finite data can discover modular policies that generalize compositionally in a number of complex environments.
  arXiv  Detail & Related papers  (2023-12-22T16:33:50Z)
• Compositional Capabilities of Autoregressive Transformers: A Study on Synthetic, Interpretable Tasks [23.516986266146855]
  We train autoregressive Transformer models on a synthetic data-generating process. We show that autoregressive Transformers can learn compositional structures from small amounts of training data.
  arXiv  Detail & Related papers  (2023-11-21T21:16:54Z)
• How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations [98.7450564309923]
  This paper takes initial steps on understanding in-context learning (ICL) in more complex scenarios, by studying learning with representations. We construct synthetic in-context learning problems with a compositional structure, where the label depends on the input through a possibly complex but fixed representation function. We show theoretically the existence of transformers that approximately implement such algorithms with mild depth and size.
  arXiv  Detail & Related papers  (2023-10-16T17:40:49Z)
• Inducing Transformer's Compositional Generalization Ability via Auxiliary Sequence Prediction Tasks [86.10875837475783]
  Systematic compositionality is an essential mechanism in human language, allowing the recombination of known parts to create novel expressions. Existing neural models have been shown to lack this basic ability in learning symbolic structures. We propose two auxiliary sequence prediction tasks that track the progress of function and argument semantics.
  arXiv  Detail & Related papers  (2021-09-30T16:41:19Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.