归纳头可能是任何大小的Transformer模型中通用ICL能力的来源

最近读了Anthropic的一篇长文（确实非常长），叫做In-context Learning and Induction Heads，觉得写的非常好，看到中文互联网应该还没有这个文章的中文翻译，也怕自己忘掉，所以读到哪就翻译到哪，做一个记录。大部分内容靠GPT翻译，我在其基础上做术语和逻辑的修正，有些地方中文没法特别好地表达出来，我会附带原文用词。有些地方原文就很抽象，我会注明自己的理解。

Introduction

在此之前，机制可解释性主要关注CNN视觉模型，但最近，我们（指Anthropic，下同）提出了一些关于Transformer语言模型的机制可解释性的非常初步的进展。具体而言，在我们之前的工作中，我们开发了一个数学框架，用于分解Transformer的操作，这使我们能够理解小型模型（仅有单层或两层attention），并几乎完全解释了这些模型是怎么工作的。也许最有趣的发现是归纳头（induction heads），这是网络中的一条circuit，其功能是回顾当前token（记做A）之前的内容，找到上一次紧随A之后出现的token（记做B），然后按照同样的方式来补全当前A后面的内容（即，归纳头倾向于再预测一个B出来，以形成序列[A][B] … [A] → [B]）。换句话说，诱导头通过复制和补全之前发生过的序列来“补全某种模式”。在我们的模型中，诱导头由两个注意力头组成：第一个头是“previous token head”，它将信息从前一个token复制到下一个token，而第二个头（实际的“诱导头”）使用该信息来查找由当前token前导的token。对于2层注意力模型，我们能够精确地表明诱导头实现了这种模式复制行为，并且似乎是in-context learning的主要来源。

注：这段的末位原文写的非常抽象，可能要多读几次才能懂。通俗地讲，假设我们的输入是ABA，在两层的注意力模型中，induction heads是两个attention head的集合，第一层那个head叫previous token head，每一个位置注意当前和之前的位置，通过第一层的previous token head，每个位置已经聚合了之前位置的信息，此时第二层中的induction head（是的，他这里重名了）用第一层送来的输入来查询当前位置的token（假设我们考虑第二个A）之前出没出现过，如果之前出现过，那么这个head则倾向于把上一次A出现的时候的下一个token（即B）直接拿过来用作最后的输出。

然而，我们的最终目标是逆向工程那些先进的语言模型（通常包含数百层和数百亿或数万亿个参数），而不仅仅是2层attention model。不幸的是，多层attention和MLP的存在使得数学上精确确定这些模型的circuit变得更加困难。但是，也可以采用不同的方法：通过经验观察、扰动和研究learning process以及各种结构的形成，我们可以尝试间接推断网络内部可能发生的机制。这在某种程度上类似于神经科学家通过观察神经发展、研究受伤患者的大脑部分、扰动动物的大脑功能或观察与相关神经元有关的一小部分相关神经元来理解大脑的一部分如何工作的方式。

在本文中，我们迈出了构建这种间接研究的第一步。特别地，我们展示了很多初步且间接的证据来证实一个假说：归纳头可能构成了大型Transformer模型中所有in-context learning的主要来源。具体而言，我们的论点是，存在一些circuits，其机制与2层的归纳头相同或相似，并且执行一个“模糊”版或“最近邻”版的模式补全，比如，这些circuits可以完成[A*][B*] … [A] → [B]，其中A*与A，B*与B，在某个空间中相似；而且，这些circuits实现了大模型中的大多数in-context learning。

注：作者主要想说，在大模型中应该也有类似归纳头一样的circuit，并且在补全ABA->B这种模式的时候，不一定要在上文中严格匹配A，可能A*由于与A语义相似，A*B*A->B也是可以完成的，其中B与B*的相似程度与A*与A的相似程度差不多，而A*B*A->B已经是一个in-context learning的形式了。

我们获取这些证据的主要方式是通过发现和研究不同大小的语言模型（只要它们是多层的）在训练初期发生的一个“相变”，这个“相变”的一个表现是training loss的一个突变。在这个相变期间，模型得到了其大多数的ICL能力，并且同时在模型中形成了归纳头，其能够实现相当“抽象”和“模糊”的模式补全。我们详细研究了这种circuit，以试图建立其中的因果关系，包括我们展示了如果我们以一种方式更改transformer结构，使得归纳的突变在训练的另一个阶段发生，那么归纳头，以及ICL能力的形成阶段也会同时发生相应的移动。

具体来说，本文提出了六条互补的证据线，支持归纳头可能是任何大小的Transformer模型中通用ICL能力的来源的论点：

论点1 宏观共存（Macroscopic co-occurence）：Transformer语言模型在训练初期经历了“相变”，在此期间归纳头部形成，并且同时ICL能力得到了显著改善。

论点2 宏观共扰动，（Macroscopic co-perturbation）：当我们以一种改变transformer结构，使归纳头能否形成（以及何时形成）发生变化时，上述ICL能力的那种显著改善精确地随之发生变化。

论点3 直接消融 （Direct ablation）：在小模型中，在测试时直接“去除”归纳头后，ICL大大减少。

论点4 归纳头普遍性的具体示例 （Specific examples of induction head generality）：尽管我们在复制文字序列方面非常狭窄地定义了归纳头，但我们经验性地观察到这些相同的头似乎也能够实现更复杂类型的上下文学习，包括高度抽象的行为，这使得它们可以解释大部分ICL。

论点5 归纳头普遍性的机制合理性 （Mechanistic plausibility of induction head generality）：对于小型模型，我们可以解释归纳头的工作机制，并且可以展示它们对上下文学习的贡献。此外，操作的实际机制（the actual mechanism of operation）表明它可以被重新用于执行更一般的上下文学习。

论点6 从小型到大型模型的连续性（Continuity from small to large models）：在前面的5个论点中，归纳头解释上下文学习的情况对于小型模型比大型模型更有说服力。然而，与归纳头和上下文学习相关的许多行为和数据在小型到大型模型之间都具有平滑的连续性，所以最简单的解释是：他们的机制是相同的。

最后，除了在将归纳头与上下文学习联系起来方面具有实际意义之外，上面提到的相变可能还与安全性有关。神经网络的能力，比如多位数相加，已知有时会在模型训练或规模增加时突然形成或改变，这对于安全性来说是特别令人担忧的，因为它意味着不希望的或危险的行为可能会突然出现。例如，一种安全性问题，即reward hacking，可能会在这种相变中出现。因此，深入研究相变并更好地理解其内部机制可能包含着未来系统解决安全问题的通用教训。特别是，我们观察到的相变形成了微观领域的解释性 与 宏观领域的scaling law以及learning dynmaics之间的有趣潜在桥梁。

关键概念

1. 上下文学习分数（ICL score）

在整个工作中，我们计算了一种简单的上下文学习的启发式度量：

上下文学习分数 = 上下文中第500个token的loss减去上下文中第50个token的loss（全体测试集上平均）。

注：作者设定这个指标是基于Kaplan et al.的概念：ICL中loss随着token index的增长而下降。

我们在选择第500个和第50个token的理由有些随意：第500个token位于长度为512的上下文的末尾附近，而第50个token在上下文中足够远，以便一些文本的基本属性（如语言和文档类型）已经建立，同时又接近开头。我们还将展示，在此处选择不同的数字不会改变我们的结论。

2. 归纳头 （Induction Heads）

我们证实地将归纳头定义为在重复的随机token序列上表现出以下两个特性的头：

1.前缀匹配 (prefix matching)：头会回顾之前的tokens，这些token紧跟在当前和/或最近的token之后。也就是说，它会关注"按照归纳，有可能作为下一个token“的那个token。

2.复制 (copying)：头的输出增加了与被关注token对应的logit值。

换句话说，归纳头是指在给出一系列完全随机的token的重复序列时，根据经验增加了[B]在[A][B]...[A]之后出现的可能性的任何头。下面是一个归纳头行为的示例：

归纳头倾向于在整体上擅长重复序列。例如，给定“The cat sat on the mat. The cat”，归纳头将推荐继续“sat on the mat”。这为归纳头如何与一般上下文学习甚至few-shot learning相连接提供了一个初步的线索：它们学会了重复任意序列，这是一种（简单的）few-shot learning形式。

3. Per-Token Loss Analysis

为了分析模型在训练时是如何演变（evolve）的，我们分析了“per-token loss vectors”。其核心思想来自Erhanet al.。

我们搞了一个模型集合（我们训练了多种不同的模型架构，并在训练过程中保存每个模型的数十个“训练快照”。我们将使用这一系列快照作为我们的模型集合。）接下来，我们收集每个模型对一个不变的测试集（其中包含10,000个随机token，这些token分别来自不同的输入）的对数似然。我们将这些对数似然组合成“per-token loss vectors”，然后应用主成分分析（PCA）：

通过将这种方法应用于多个模型的训练快照，我们可以可视化并比较不同模型的训练轨迹（in terms of outputs）如何演变。

论点1: Transformer语言模型在训练时经历“相变”，在这期间归纳头形成，同时ICL能力极大提升

我们的第一个观察是，如果我们在训练过程中测量transformer的上下文学习能力（基于ICL score），它会在训练早期的狭窄窗口内（大约在25亿到50亿token之间）突然发展，然后在训练的其余部分保持恒定（请参见下面的图表）。在此窗口之前，上下文学习少于0.15 nats，在此之后，上下文学习大约为0.4 nats，这个数量在训练的其余部分保持恒定，并且在许多不同的模型大小上也保持恒定（除了一层模型几乎没有形成太多上下文学习能力）。这似乎是一件令人惊讶的事 - 因为很自然地，人们会期望上下文学习能力会在训练过程中逐渐改善，并且会随着模型大小的增加而改善，正如机器学习中的大多数事情所做的那样。

虽然我们上面只展示了三个模型，但这种模式普遍成立。（在论文后面的模型分析表中展示了许多示例，包括不同模型架构和规模的模型。感兴趣可以看原文附录）

有人可能好奇ICL能力的突然提升是否在某种程度上是将ICL Score定义为第500个和第50个token之间的差导致的。我们稍后会更深入讨论这个问题。但现在，一个简单的方法来确认这种突然提升是一个稳健的现象是查看 the derivative of loss with respect to the logarithm token index in context。你可以将其视为衡量“ICL能力在ε%的context长度增加时的表现”的指标。我们可以在二维图上可视化这个指标，其中一个轴是经过的训练量，另一个轴是被预测的token的index。在相变之前，loss在50th token附近基本停止了改善，但在相变之后，loss继续被改善。

此外，突然出现的ICL的改善并不是这个时间窗口中发生的唯一变化。如果我们对注意力头进行评估，以确定它们是否是归纳头（使用prefix matching score），我们会发现归纳头在正好与上下文学习发展的窗口期突然同时形成（下图所示）

这已经强烈暗示了归纳头和上下文学习之间的某种联系，除此之外，似乎这个窗口对于训练过程总体来说是一个关键点：无论发生了什么，都可以在训练曲线上看到一个突起（下图）。实际上，这是在训练中的loss曲线上唯一一个非凸的地方。

这可能听起来并不重要，但loss曲线是在成百上千个token上进行平均的。要知道，人们发现许多有趣的语言模型行为，例如算术能力的涌现，在loss曲线上都是难以观察到的，如果一些东西在loss曲线上可见，则表明它是模型行为的重大变化。这个转变似乎也是损失曲线与单层模型分歧的第一个点（至少对于小型模型而言） - 单层模型没有体现出这个突起，就像它没体现出其他突变一样。

我们还可以对“per-token loss”应用PCA，如在关键概念一节的‘per-token loss anaylsis‘所述，这使我们能够总结在训练过程中模型的预测如何变化。

下面我们展示了这些模型预测的前两个主成分，黄色轮廓突出显示了与上文所述的“窗口期”，即在上下文学习突然改善的时候。我们可以看到在训练轨迹中，正好在其他变化发生的窗口时，PCA可视化出的轨迹发生了转折。从某种意义上说，当上下文学习突然改善时，这个窗口期内发生的任何事情都是与transformer在训练过程中遵循的基本轨迹的重大偏差。唯一的例外还是单层模型——其中归纳头无法形成，上下文学习也不会改善。

总结一下，在同一突变窗口期间发生了以下事情：

1. 上下文学习能力急剧改善（基于ICL score衡量）。

2. 形成归纳头。

3. loss经历了一个小的“颠簸”（即，loss曲线在之前和之后的部分相比，经历了一个明显更陡的改善期）。

4. 模型的轨迹突然发生改变（在per-token loss空间中，通过PCA可视化分析）。

综合这些结果表明，在训练早期的 2.5e9 到 5e9 tokens窗口期间（对于大模型来说，这可能是训练过程中的大约1-2%），发生了一些重要的转变。我们称这个转变为“相变”，因为它是一种突然的改变，改变了模型的行为，并在宏观层面（loss和上下文学习曲线）和微观层面（归纳头）都有表现。

更近距离观察相变

一个自然的解释是，对于所有这些模型，归纳头实现了上下文学习：它们的形成驱动了观察到的所有其他变化。为了加强这一假设，我们检查了一些事情。首先，相变发生的窗口似乎与学习率、warmup或weight decay不相关。其次，我们尝试在不同的数据集上训练一些小型模型，并观察到相变以相同的方式发展（more details see 原文模型分析表）。

第三，为了加强这种联系，我们定性地来看模型在相变期间的行为。一种方法是查看模型在预测上变得更好和更差的特定token。模型的loss是数十亿token对数似然损失的平均值，通过分析它们，我们可以了解到发生了什么变化。

具体地，我们选择一段文本 - 为了好玩，我们将使用《哈利·波特》的第一段文字 - 并比较相变的开始和结束时的对数似然差异。我们会注意到大多数变化发生在文本中多次重复出现的token时。如果一系列token多次出现，模型在第二次出现时更擅长预测该序列。另一方面，如果一个token后面跟着与先前不同的标记，那么在相位变化后的模型在预测它时效果更差：

我们也可以在模型训练的过程中进行相同的分析。让我们来看看《哈利·波特》第一段中两个token的“per-token loss”轨迹。我们用红色展示了一个在相变期间预测显著改善的token：它是“The Dursleys”的四个token中的最后一个（即 leys），因为这个序列在文本中多次出现。我们用蓝色展示了一个在相变期间loss明显增加的token：即“Mrs Potter”的首次出现，因为之前的“Mrs”都后跟着的是“Dursleys”

对以上证据进行评估

尽管我们有了上述所有证据（以及模型分析表中的更多证据），事实仍然是我们尚未表明归纳头是上下文学习的主要机制。我们只是表明归纳头与上下文学习同时形成，并且此后上下文学习没有进一步改善。在这个故事中存在一些潜在的混淆因素。以下是我们总结的支持和反对认为这种联系是因果关系的原因。支持的论点可能如下所示：

1. 归纳头的形成与各种规模的模型的上下文学习能力大幅增加相关。
2. 这两个明显的转变在这么多模型中同时发生，不太可能纯粹是偶然的，不太可能没有任何因果关系。
3. 几乎肯定存在某种联系，而最简单的可能性是归纳头是导致观察到的上下文学习增加的主要机制。（然而，如下所讨论的，它也可能是一个混淆变量。）
4. 由于超过75%的最终上下文学习能力是在此时形成的，人们可能自然地认为这就是归纳头负责的上下文学习能力的“量”。

然而，以下问题和混淆因素需要谨慎考虑：

1. 在大型模型中，我们对训练过程的分析具有较低的时间分辨率。当只有15个时间点时，co-occurance不太令人惊讶，证据也较弱。
2. 也许在这一时间点上，我们的模型中还形成了其他机制，这些机制不仅对归纳头有贡献，还对上下文学习的其他来源有贡献。（例如，也许相变实际上是模型学习如何通过残差流来组合不同层的时间点，从而使归纳头和潜在需要多头组合的许多其他机制也可能发生共现）换句话说，共现也许主要是由一个共享的潜在变量引起的，而不是由归纳头直接引起全面观察到的上下文学习变化的直接因果关系。
3. 在相变点之后，上下文学习得分大致保持不变（在0.4 nats左右），但这并不一定意味着上下文学习的底层机制在那个点之后是不变的。特别是，我们使用的指标度量了token index 500和50之间的相对损失，并且我们知道模型在50th token处的性能随着训练时间的推移而改善。因为从较低的baseline降低定量的loss是很困难的，因此，这里面可能有随着训练时间的推移而驱动的其他机制的影响。

注：这里第三点是说，我们知道ICL score = 500th token loss - 50th token loss，50th token loss肯定是会被持续优化降低的，但是ICL score竟然保持非常稳定，如果这个稳定只来源于500th token loss一直按“同样的速率”被优化，是很难的一件事，所以这里面可能有其他因素导致这种“稳定”。

值得注意的一点是，在这里，“在相变点时，归纳头导致大部分ICL能力”的论点比“在训练结束时，归纳头导致大部分ICL能力“更为可靠 - 即使我们的度量值ICL score保持不变，后续的训练期间可能会发生许多变化。

论点2: 当改变transformer结构时，相变的出现随之改变

之前已经讨论过，induction heads需要两层模型才能形成，作者想能否修改模型结构使得一层模型即可出现类似的行为？因此作者在第一层attention计算的过程中引入了一个smeared key：

其中alpha是可学习的参数，j下标代表位置，这个修改意为：之前需要两层模型是因为第一层需要prefix matching，即注意自己之前的token，第二层进行copy，在这里通过强制让当前位置引入之前前一个位置的key 信息，即可在单层注意力内完成prefix matching和copy两种行为。

然后作者发现一层的模型也出现相变了，对于两层的模型，icl相变的时间点大幅提前了。

作者评论道：“然而，根据这些证据，我们也许不该对大型模型做出过于强烈的推断。这个实验表明归纳头是在transformer中显著提升ICL能力的最小机制。但在更大的模型中，很容易想象到，这个机制并不是全部，而且这个实验也不能否定在训练过程中上下文学习机制发生变化的想法。”

论点3: 在测试的时候我们关掉induction head，ICL能力急剧下降

这是很直观的现象，这个实验的细节可以看附录。