💬 NLP

How Do Transformers Learn to Associate Tokens: Gradient Leading Terms Bring Mechanistic Interpretability

本文通过基于训练梯度动力学推导出的大词元、词元可互换性及上下文映射的复合闭式表达式，为变换器如何学习语义关联提供了机制性解释。

原作者： Shawn Im, Changdae Oh, Zhen Fang, Sharon Li

发布于 2026-05-14

📖 1 分钟阅读☕ 轻松阅读

原作者： Shawn Im, Changdae Oh, Zhen Fang, Sharon Li

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

想象一下，将 Transformer（现代 AI 聊天机器人背后的“大脑”）视为一本巨大的空白笔记本。当我们开始训练它时，页面是空白的。你所询问的这篇论文提出了一个简单却深刻的问题：这本笔记本如何仅通过阅读数百万个句子，就学会将“鸟”和“飞”这样的词联系起来？

作者没有去观察最终成型的复杂 AI，而是决定放大观察学习的最初几个瞬间。他们使用数学上的“放大镜”，审视 AI 在开始更新其记忆时所采取的最初几步。

以下是他们发现的拆解，使用了简单的类比：

1. “第一步”捷径

训练 AI 就像尝试预测句子中的下一个词。通常，AI 更新其记忆背后的数学极其混乱且复杂。

作者意识到，在最初阶段，AI 并不需要执行复杂的微积分运算。它只需要遵循数据中最强、最明显的信号。他们称之为“主导项”。这就像一名徒步者开始攀登一座高山。在起步阶段，路径清晰且笔直。只有到了后来，随着他们爬得更高，路径才变得曲折复杂。作者证明，在相当长的一段时间内，AI 的“记忆权重”都紧紧跟随那条简单、笔直的初始路径。

2. 三个构建模块

该论文揭示，AI 的记忆并非随机的混乱。相反，它是由堆叠三种特定类型的“乐高积木”（作者称之为基函数）构建而成的。这些积木直接源自 AI 所阅读文本的统计数据：

“下一个词”模块（双词映射）：
- 类比： 想象一个孩子学习到"The"后面几乎总是跟着像"cat"或"dog"这样的名词。
- 作用： 这个模块简单地记录：“如果你看到词 A，下一个最可能出现的词 B 是什么？”它捕捉简单、直接的关联。
“可互换”模块（互换性映射）：
- 类比： 想想同义词词典。如果你能在句子中将"car"替换为"truck"而不破坏语法，它们就是“可互换的”。
- 作用： 这个模块学习某些词扮演相同的角色。如果"red"通常描述"car"，而"fast"也描述"car"，这个模块帮助 AI 意识到"red"和"fast"可能是相关的，因为它们都与相同的词共现。它根据词的功能而非仅仅是紧邻的邻居来对词进行分组。
“上下文”模块（上下文映射）：
- 类比： 想象阅读一个关于“鱼”的故事。即使"pond"这个词没有紧挨着"fish"，故事前面可能会提到"water"、"lake"或"boat"。
- 作用： 这个模块在句子中向后看得更远。它学习到，如果一个词出现在特定的上下文中（例如关于自然的故事），它很可能与该上下文中的其他词相关联，即使它们相距甚远。

3. 它们如何协同工作

该论文表明，AI 内部的“权重”（决定其思考方式的数字）仅仅是这三个模块的组合。

输出层（负责猜测下一个词的部分）主要是下一个词模块。这是 AI 在说：“基于我刚刚看到的，接下来是什么？”
注意力层（负责决定关注什么的部分）是可互换模块和上下文模块的混合。这是 AI 在说：“我看到了‘鱼’这个词。我应该回看‘池塘’这个词，因为它们经常出现在类似的故事中，即使它们没有紧挨着。”

4. 证明：它确实有效

作者不仅写了方程；他们在真实的 AI 模型上测试了这一点。

他们在儿童故事数据集上训练了一个小型 AI。
他们将 AI 实际学到的记忆与他们的数学公式进行了比较。
结果： 匹配度极高（在某些情况下相似度超过 99%）。即使当他们观察一个巨大的现实世界 AI（Pythia-1.4B）时，同样的模式依然成立。AI 确实正在使用这三个简单的统计构建模块来组织其记忆。

核心启示

该论文认为，我们不需要将 AI 视为一个神秘地“理解”语言的“黑盒”。相反，在其核心，它只是通过以下方式在统计上组织世界：

记录什么通常跟随什么。
将执行相同工作的词分组。
链接共享背景故事的词。

通过理解这三条简单的规则，我们终于可以看到机器如何学会将“鸟”与“飞”联系起来，将“国家”与“首都”联系起来。这不是魔法；这是一种非常结构化的方式来总结它所阅读的文本。

技术摘要：Transformer 如何学会关联词元：梯度主导项带来机制可解释性

问题陈述

基于自注意力机制的大语言模型（LLMs）在捕捉事实性知识和生成连贯文本方面展现了卓越能力，这主要归功于其学习语义关联的能力（例如“鸟”与“飞”之间的联系，或“汽车”与“卡车”之间的可互换性）。尽管近期研究已识别出诸如归纳头（induction heads）或线性语义关系等具体机制，但关于这些关联如何在真实世界自然语言数据上基于注意力的 Transformer 训练过程中涌现，仍缺乏原则性且基于理论的解释。

先前的理论工作往往依赖于不切实际的假设，例如合成结构化语言、简化的架构（缺乏位置编码或残差连接）或非标准的训练制度（例如顺序组件式训练）。这些与实际情况的偏离限制了其见解对现代大语言模型的普适性。本文通过提出以下问题来填补这一空白：基于注意力的语言模型在自然语言数据上训练时，语义关联是如何涌现的？

方法论

作者提出了一个基于训练动力学的理论框架，具体分析了训练早期阶段梯度的主导项近似。

模型设置：本研究聚焦于标准的 $L$ 层基于注意力的 Transformer，包含（相对）位置编码、因果掩码和残差流，并在自然语言语料库上通过标准的全批量梯度下降进行训练。
主导项近似：鉴于 Transformer 在训练早期即获得核心行为（包括语义关系），作者通过近似梯度更新，推导出了权重矩阵（输出 $W_O$ 、值 $V^{(l)}$ 、查询 - 键 $W^{(l)}$ 以及位置编码 $P^{(l)}$ ）的闭式表达式。他们表明，对于足够多的步数 $s$ （其中 $s \leq \eta^{-1} \min(\dots)$ ），权重保持接近其梯度的主导项。
基函数分解：理论分析揭示，学习到的权重可以表示为源自语料库统计的三种基本基函数的简单组合：
- 二元组映射（ $\bar{B}$ ）：捕捉下一个词元的依赖关系（ $P(e_j|e_i)$ ）。
- 可互换性映射（ $\Sigma \bar{B} = \bar{B}^\top \bar{B}$ ）：基于前驱词元分布的相似性，捕捉词元之间的功能相似性（如同义词或共享的语法角色）。
- 上下文映射（ $\bar{\Phi}$ ）：编码更长距离的前缀 - 后缀共现，代表词元出现的平均上下文。

主要贡献

本文做出了三项主要贡献：

显式权重表征：本文首次提出了在真实世界文本语料库上、基于下一个词元预测损失训练的基于注意力的 Transformer 中权重的显式、闭式表征。研究表明，权重是上述三种基函数的组合。
语义关联的机制解释：作者将权重中学习到的特征解释为二元组、可互换性和上下文映射的特定组合。他们展示了这些组合如何使模型能够捕捉语义关联（例如，“鱼”如何通过上下文映射与“池塘”关联，或“红色”如何通过二元组映射与“卡车”关联）。
实证验证：理论预测在两个模型上得到了验证：一个是在 TinyStories 数据集上训练的玩具 3 层注意力模型，以及一个在 OpenWebText 上训练的实际大规模大语言模型（Pythia-1.4B）。

结果

理论一致性：在 TinyStories 上训练的 3 层 Transformer 实验表明，即使经过 100 个 epoch 的训练，理论预测的主导项与实际学习到的权重之间的余弦相似度仍然极高（例如，注意力和值矩阵的相似度 >0.99）。
语义结构：定性分析证实，基函数捕捉到了有意义的语言结构。例如，二元组映射 $\bar{B}$ 将形容词与其修饰的名词相关联，而上下文映射 $\bar{\Phi}$ 将实体与其典型环境相关联（例如“鱼”与“池塘”）。
向真实大语言模型的泛化：当应用于 Pythia-1.4B 时，理论主导项特征与学习到的嵌入和注意力相关性显示出强烈的一致性，特别是在训练早期。分析表明，随着训练的推进，权重虽然向更丰富的知识表征漂移，但在相当长的一段时间内仍保留了理论预测的基础关联特征。
组件协作：对完整模型计算（公式 12）的分析说明了残差流如何提供平均的下一个词元预测，而自注意力块（由 $\bar{Q}$ 和 $\bar{\Phi}^\top \Sigma \bar{B}$ 驱动）则通过根据语料库统计选择性地关注最能指示下一个词元的词元，来细化这一预测。

意义

本文声称通过将深度学习训练动力学与语言理论相结合，为大语言模型中的表示学习提供了机制基础。通过将分析扎根于自然主义文本分布和标准训练程序，这项工作显著缩小了形式理论分析与实际大语言模型使用之间的差距。

作者认为，他们的发现提供了一种严谨的方法来定位每个 Transformer 组件的功能。该理论表明，模型的行为并非将模型视为黑盒，而是通过首先建立这些基础关联特征（二元组、可互换性、上下文）来发展，这些特征作为学习更复杂语义关系的锚点。这项工作为可解释性研究开辟了途径，它表明不同组件的权重矩阵可以分解为共享语料库统计的简单函数，从而允许提出关于概念如何在模型中涌现的更广泛假设，而不仅限于特定机制。