Circuit explained: How does a transformer perform compositional generalization.

Abstract

Compositional generalization-the systematic combination of known components into novel structures-is fundamental to flexible human cognition, yet the mechanisms that enable it in neural networks remain poorly understood in both machine learning and cognitive science. [1] showed that a compact encoder-decoder transformer can achieve simple forms of compositional generalization in a sequence arithmetic task. In this work, we identify and mechanistically interpret the circuit responsible for this behavior in such a model. Using causal ablations, we isolate the circuit and show that this understanding enables precise activation edits to steer the model's outputs predictably. We find that the circuit performs function composition without encoding the specific semantics of any given function-instead, it leverages a disentangled representation of token position and identity to apply a general token remapping rule across an entire family of functions. Although the circuit mechanism was identified in a limited number of small scale models with a synthetic task, it sheds light to how symbolic compositionality can emerge in transformers and offer testable hypotheses for similar mechanisms in large-scale models. Code for model and analysis is publicly available.