📊 statistics

Assign and Add: A Mechanistic Study of Compositional Arithmetic

本文研究了小型 Transformer 如何在变量赋值和模运算加法的受控设置中实现组合泛化，揭示了模型如何对直接输入和间接输入复用内部机制，并经历不同的学习阶段，从而自然地泛化到未见的组合。

原作者： Brady Exoo, Alberto Bietti, John Sous

发布于 2026-06-01

📖 1 分钟阅读☕ 轻松阅读

原作者： Brady Exoo, Alberto Bietti, John Sous

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在教一个非常聪明但非常刻板的机器人做数学。你想让它解决这样一个谜题：“如果 'A' 是 5，'B' 是 10，那么 A + B 等于多少？”

棘手之处在于，机器人以前从未见过以这种特定顺序出现的字母“A”和“B”。它只见过数字相加（比如 5 + 10）或者在其他方式下使用的其他字母。核心问题是研究人员提出的：机器人是如何学习将这两个不同的技能——理解一个字母代表一个数字，以及随后进行实际数学运算——结合起来，从而解决一个它从未见过的全新谜题的？

以下是论文《Assign and Add》（赋值与相加）的发现，其简单解释如下：

1. 设置：两步舞曲

研究人员构建了一个微小且简单的“大脑”（神经网络），并给了它一个特定的任务。他们给它输入这样的句子：

c=17, b=42, b+19=?

机器人必须弄清楚 b 实际上是 42，所以数学运算其实是 42 + 19，然后给出答案。

为了观察机器人的学习过程，他们将训练分成了不同的“口味”的谜题：

简单模式： 只有数字（例如：17 + 19 = ?）。
中等模式： 一个字母和一个数字（例如：b + 19 = ?）。
困难模式： 两个字母（例如：b + c = ?）。

2. 学习的三阶段

机器人并不是一次性学会所有东西的。它经历了三个截然不同的阶段，就像学生循序渐进地掌握学科知识一样：

第一阶段：数学大师。 首先，机器人学会了如何进行实际的加法。它变得非常擅长把数字加在一起。在这个阶段，它擅长做 17 + 19，但会被字母完全搞糊涂。它就像一个不知道什么是变量的计算器。
第二阶段：翻译官。 接下来，机器人学会了“赋值”技能。它弄明白了如何通过观察 c=17 来记住 c 代表 17。它学会了把字母替换成数字。
第三阶段：指挥家。 最后，机器人学会了将这两个技能组合起来。它意识到：“哦，我可以利用我的‘翻译官’技能来寻找数字，然后使用我的‘数学大师’技能来进行加法。”

令人惊讶的是： 机器人先学会了数学，才完全学会了翻译。即使在它还在努力理解字母的时候，它也能完美地完成加法运算！

3. 机器人的“大脑”是如何工作的（机制）

研究人员不仅看了最终得分，还深入观察了机器人的“大脑”内部，看它是如何思考的。他们发现有两个特定的工具在协同工作：

工具 A：“前一个标记”头（翻译官）。
想象一个机器人的手臂，它会向后方抓取当前项之前的那个项目。在机器人的第一层中，它学会了使用这个手臂去抓取之前被分配给某个字母的数字。如果它看到 b，它就会向回寻找在句子前半部分 b 被分配到的数字。
工具 B：“傅里叶”数学模块（计算器）。
在第二层中，机器人使用一种特殊的数学技巧（涉及波形和模式）来进行加法。这是机器人进行模运算（即像时钟一样循环的数学运算）的一种已知方法。

神奇时刻： 机器人学会了连接这两个工具。 “翻译官”手臂抓取正确的数字，并将其传递给“计算器”模块。一旦连接成功，即使它从未见过 b 和 c 相加，它也能解决 b + c 的问题。这就像是一个会说“字母语”的翻译员和一个会说“数字语”的计算器，最后建立了一座桥梁。

4. 为什么这很重要

这篇论文表明，泛化（解决新问题）并非魔法。它是在模型学习了微小的、原子化的技能（如加法和寻找变量）之后，再通过理解如何将这些技能编织在一起而自然发生的。

“顿悟”（Grokking）效应： 研究人员注意到了一种现象，称为“顿悟”。有时机器人能完美记忆训练数据，但在面对新测试时却失败了。然后，突然间，它“开窍”了，开始理解背后的“规则”而非仅仅是记忆答案。
结果： 因为机器人分别学习了内部机制，然后又将它们结合起来，所以它可以处理那些在训练期间从未出现过的位置形式的字母谜题。

总结类比

把机器人想象成一名厨师。

首先，它学习如何切菜（数学）。
然后，它学习如何阅读一张写着“使用蓝色碗里的洋葱”的食谱卡（变量赋值）。
最后，它学会了将两者结合：阅读卡片，找到洋葱，然后切菜。

论文证明，如果你将这两项技能分开教授，它们最终会学会如何结合这些技能，去烹饪一道你从未要求它们制作过的菜肴。这些技能的“组合”正是让机器人变得聪明的关键。

技术摘要：赋值与相加：关于组合算术的机制研究

问题陈述

大型语言模型（LLMs）展示了在执行训练期间未见的复杂任务时组合技能的能力，然而实现这种“组合泛化”（compositional generalization）的具体机制仍然难以捉摸。虽然泛化通常被归因于原子机制的组合，但不同机制之间的相互作用——特别是变量绑定（通过归纳头实现的关联召回）与模运算（通过傅里叶特征实现的结构化表示）——尚未在受控设置中得到系统性的研究。

本文旨在解决以下问题：Transformer 如何学习组合变量绑定与模加法机制，从而实现超越训练分布的泛化，以及这些能力在训练过程中是如何涌现的？

方法论

作者采用了一个受控的实验设置，使用一个小型且具有机械可解释性的 Transformer，用于执行结合了变量赋值与模加法的合成任务。

模型架构： 一个 2 层、单头的 Transformer，其参数为 $d_{model} = d_{head} = 128$ 且 $d_{mlp} = 512$ 。该架构刻意省略了偏置（biases）、归一化层（normalization layers）和第 1 层 MLP，以简化电路发现过程。它使用 GPT-2 式的绝对位置嵌入和因果掩码（causal masking）。
任务定义： 模型被训练执行模加法（ $N=59$ ），其中操作数可以是常量，也可以是在上下文中动态赋值的变量（例如：c=1, b=17, a=42, b+19=?）。
数据分布： 序列通过随机填充进行分词，以防止对绝对位置的依赖。训练集被划分为用于评估泛化的不同部分：
- 0 变量（0-var）： 纯常量加法。
- 1 变量（1-var）与 2 变量（2-var）： 涉及已赋值变量的加法。
- 留出集（Held-out Sets）： 约 30% 的加法对从训练中移除（例如，如果 1+3 未见，则排除所有包含该序列的序列）。此外，特定的变量被限制在训练集中的某些位置，以测试位置泛化。
分析技术： 研究利用机械解释性工具（通过 TransformerLens 库）来分析注意力模式、残差流向量和权重矩阵。同时，研究还采用了线性探测（linear probes）和理论梯度分析来追踪训练动态。

核心贡献

可处理的组合算术模型： 作者提出了一个简化的模型，允许对变量赋值如何与模加法进行组合进行机械解释。
解耦的训练动态： 分析表明，模加法是在变量赋值之前且独立于变量赋值被学习的。模型在 0-var 任务上实现了近乎完美的准确率，同时仍在学习如何进行变量绑定。
组合泛化的机械解释： 本文确定了泛化源于两个特定电路的组合：
- 第 1 层中一个修改后的归纳头，负责检索已赋值变量并路由其值。
- 第 2 层的模加法电路（具体为 Output-Value 矩阵和第 2 层 MLP），负责对检索到的值进行算术运算。
理论框架： 作者提供了理论分析，证明了梯度动态如何促进变量绑定与模加法电路的组合，表明在第 2 层 Query-Key 矩阵上的负梯度自然地导致了关联存储块（associative memory block）的形成。

结果

训练动态与阶段

训练过程呈现出三个不同的阶段：

模加法学习阶段： 模型首先学习执行不含变量的序列（0-var）的加法，在变量绑定完全建立之前就达到了高准确率。
变量赋值学习阶段： 模型随后学习变量赋值所需的结构。在准确率出现一次剧烈“峰值”（大约在第 5,300–5,400 步）时，变量赋值路由电路随之涌现，此时第 1 层注意力能够正确地将常量路由到其对应的变量。
精炼阶段： 第二阶段（第 14,500–16,000 步）涉及对“困难”序列（例如，处于留出位置的特定变量）的路由进行精炼，以修正最初对某些变量位置关注失败的问题。

泛化能力

位置泛化： 模型能够泛化到新的变量顺序（例如，在看到 14+b 后预测 b+14）。
分布外（OOD）泛化： 模型成功预测了训练期间从未以任何形式出现过的留出加法对（例如 1+3），在精选测试集上达到了 99.4% 的准确率。
数据需求： 要实现对未见常量对的泛化，需要在训练中覆盖至少约 25% 的总分布（ $f \approx 0.25$ ）。要实现对 2 变量序列的泛化，需要 2-var 序列在训练数据中的相对频率约为 20%（ $r \approx 0.2$ ）。

机械发现

共享表示： 在第 2 层 MLP 之前，变量赋值加法与纯常量加法的残差流向量几乎完全一致（余弦相似度 $\approx 1$ ），表明存在一种共享的计算机制。
傅里叶特征： 第 2 层 MLP 神经元的预激活值符合理论形式的傅里叶特征（ $\cos(\frac{2\pi k}{N}n + \alpha_k) + \cos(\frac{2\pi k}{N}m + \beta_k)$ ），证实了使用既定的周期性机制进行模运算。
归纳头修改： 第 1 层注意力充当常量的“前一 Token 头”，关注变量紧邻的前一个 Token。第 2 层的 Query-Key 矩阵在变量嵌入上形成一个强恒等模式，使得 = Token 可以直接关注到与变量相关的正确常量。

意义

本文声称，Transformer 中的组合泛化可以作为内部机制组合性的自然结果。通过证明不同的机制（归纳头与基于傅里叶的算术）可以被独立学习并通过注意力路由进行组合，这项工作为理解 Transformer 如何泛化到复杂推理任务提供了机械基础。

作者对研究范围保持谦逊，指出该研究受限于一个小型且受控的设置以便于分析。他们并不声称这些发现能直接解释所有 LLM 的行为，而是建议它们提供了一个理解 Transformer 如何产生泛化能力的框架，并可能为未来研究更复杂的任务及缩放法则（scaling laws）提供启发。