Assign and Add: A Mechanistic Study of Compositional Arithmetic
作者: Brady Exoo, Alberto Bietti, John Sous
分类: cs.LG, stat.ML
发布日期: 2026-05-29
💡 一句话要点
研究Transformer在变量赋值与模块化加法中的组合泛化机制
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: Transformer 组合泛化 机制分析 变量赋值 模块化加法
📋 核心要点
- 大型语言模型在组合技能方面表现出色，但其内部组合机制尚不明确，是当前研究面临的挑战。
- 该论文通过变量赋值和模块化加法的受控实验，深入研究了Transformer的组合泛化能力。
- 实验结果表明，Transformer能够泛化到未见过的变量和数字组合，揭示了模型内部机制的组合性。
📝 摘要（中文）
大型语言模型能够组合技能以执行复杂任务，其中许多任务可能在训练期间未曾见过。这种组合的具体方式仍然难以捉摸。本文通过考虑一个涉及变量赋值和模块化加法的简单受控环境，研究了Transformer中组合泛化的机制。通过将训练数据划分为不相交的集合，我们观察到小型Transformer能够泛化到以前未见过的变量和数字的组合。我们的机制分析表明，无论是直接给出输入还是通过单独的变量赋值机制间接给出输入，都使用相同的“模块化加法”MLP模块。我们还从经验角度分析了训练动态，揭示了学习的三个阶段：首先，学习模块化加法，然后是变量赋值所需的结构，最后是一个细化阶段，模型泛化到训练中未见过的一些困难序列。最后，我们提供了一个理论框架来解释组合性如何从训练动态中出现。这些结果表明，组合泛化可能是Transformer内部机制组合性的自然结果。
🔬 方法详解
问题定义：现有大型语言模型虽然展现出强大的组合能力，但其内部实现机制仍然是一个黑盒。理解模型如何将不同的技能组合起来完成复杂任务，是当前研究的痛点。特别是在算术运算等任务中，模型如何处理变量赋值和运算的组合，仍然缺乏深入的理解。
核心思路：该论文的核心思路是通过构建一个简单的、可控的实验环境，即变量赋值和模块化加法，来研究Transformer的组合泛化能力。通过分析模型在不同阶段的学习动态，以及模型内部模块的使用情况，来揭示模型如何实现组合泛化。
技术框架：该研究的技术框架主要包括以下几个部分：1) 构建一个包含变量赋值和模块化加法的训练数据集，并将数据划分为不相交的集合，以测试模型的泛化能力。2) 使用小型Transformer模型进行训练，并监控模型的训练动态。3) 通过机制分析，研究模型内部的MLP模块在处理不同输入时的行为，特别是“模块化加法”模块的使用情况。4) 提出一个理论框架，解释组合性如何从训练动态中涌现。
关键创新：该论文的关键创新在于，它通过一个简单的受控实验，揭示了Transformer在组合泛化方面的内部机制。通过机制分析和训练动态分析，发现模型能够将变量赋值和模块化加法这两个模块组合起来，实现对未见过的变量和数字组合的泛化。此外，该论文还提出了一个理论框架，解释了组合性如何从训练动态中涌现。
关键设计：在实验设计方面，论文将训练数据划分为不相交的集合，确保模型在训练过程中无法直接学习到所有变量和数字的组合。在模型方面，论文使用了小型Transformer模型，以便于进行机制分析。在损失函数方面，论文使用了标准的交叉熵损失函数。在训练过程中，论文监控了模型的训练动态，包括损失函数的变化、准确率的变化以及模型内部模块的激活情况。
🖼️ 关键图片
📊 实验亮点
实验结果表明，小型Transformer能够泛化到以前未见过的变量和数字的组合。机制分析显示，相同的“模块化加法”MLP模块被用于处理直接输入和通过变量赋值机制间接给出的输入。训练动态分析揭示了学习的三个阶段：模块化加法学习、变量赋值结构学习和泛化到困难序列的细化阶段。
🎯 应用场景
该研究成果有助于理解大型语言模型的组合泛化能力，并为提升模型的可靠性和可解释性提供指导。潜在应用包括：开发更强大的智能体，能够处理复杂的任务；改进模型的训练方法，使其更容易泛化到新的任务；以及设计更可信赖的AI系统，能够解释其决策过程。
📄 摘要（原文）
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.