这篇论文用简洁的实验揭示了 Transformer 组合泛化的内部机制，对理解大模型如何组合技能有启发意义，做可解释性或模型架构研究的读者值得一看。
这篇论文通过变量赋值和模加法任务，研究了 Transformer 如何实现组合泛化。作者将训练数据划分为不相交集合，发现小型 Transformer 能泛化到未见过的变量与数字组合。机制分析显示，无论输入是直接给出还是通过变量赋值间接给出，模型都使用相同的“模加法”MLP 模块。训练动力学分析揭示了三个阶段：先学习模加法，再学习变量赋值结构，最后精炼阶段泛化到困难序列。理论框架解释了组合性如何从训练动力学中自然涌现。
Assign and Add: A Mechanistic Study of Compositional Arithmetic
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.