Assign and Add: A Mechanistic Study of Compositional Arithmetic
This study provides insights into the mechanistic underpinnings of compositional generalization in transformers, which is a fundamental problem for understanding and improving large language models.
This paper investigates how transformers compose skills by studying variable assignment and modular addition. The authors found that small transformers can generalize to unseen combinations of variables and numbers, using the same modular addition MLP module for both direct and indirect inputs.
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.