Transformer残差流中的特权基向量

Anthropic:Transformer Circuits(可解释性研究)·2023-03-16 00:00·1263天前·Nelson Elhage∗
AI 导读

研究发现Transformer模型的残差流中存在“特权基向量”现象,即某些坐标方向持续出现异常大的激活值,这与“无特权基”的理论预期相悖。通过实验,研究者将根源指向Adam优化器中的逐维度归一化器,而非层归一化或浮点精度问题。在2亿参数模型中,典型层有20至60个维度的激活绝对值超过6。研究还提出使用峰度作为检测指标,发现激活分布峰度普遍大于3,进一步证实了基向量对称性被破坏。

Anthropic:Transformer Circuits(可解释性研究)
精选
73AI 编辑部评分,满分 100

Transformer残差流中的特权基向量

2023-03-16 00:00· 1263天前· Nelson Elhage∗
AI 导读

研究发现Transformer模型的残差流中存在“特权基向量”现象,即某些坐标方向持续出现异常大的激活值,这与“无特权基”的理论预期相悖。通过实验,研究者将根源指向Adam优化器中的逐维度归一化器,而非层归一化或浮点精度问题。在2亿参数模型中,典型层有20至60个维度的激活绝对值超过6。研究还提出使用峰度作为检测指标,发现激活分布峰度普遍大于3,进一步证实了基向量对称性被破坏。

推荐理由

揭示Transformer内部基依赖性的根源,帮助研究者改进模型可解释性。

正文 · AI 翻译

所属机构

发表时间

2023年3月16日

* 核心研究贡献者;

‡ 通讯作者

colah@anthropic.com

作者贡献声明见下文

摘要

我们对Transformer架构的数学理论表明,残差流中的各个坐标不应具有特殊意义(即基方向在某种意义上应是“任意的”,其编码信息的可能性并不比随机方向更高)。近期研究表明,这一观察在实践中并不成立。我们对此现象进行了研究,并初步得出结论:Adam优化器中的逐维度归一化器是造成该效应的原因。

我们探讨了Transformer中基依赖性的另外两个明显来源:层归一化和有限精度浮点计算。我们确信地排除了这两者作为观察到的基对齐现象的原因。

更详细的故事

Tim Dettmers近期发布了一系列关于他称之为Transformer模型中“涌现异常值”的研究结果和代码:在大型Transformer中,残差流的某些坐标会出现非常大的异常值,其数值可达其他任何坐标的20倍以上。

这项工作提出的一个明显的可解释性问题是:

这些特征是什么?它们代表什么,或者它们起什么作用?

然而,还有第二个问题,只有对Transformer有更深入的数学模型理解才会注意到:

为什么这些特征会与基对齐?

我们通常认为残差流“没有特权基”。这意味着没有理由期望残差流中的各个坐标具有任何特定含义或显著属性。这一信念源于以下观察:所有对残差流进行读取或写入的操作,都是通过任意满秩线性变换进行的。这反过来意味着,我们可以对残差流进行任意满秩线性变换,然后将相同的变换以适当方式乘入Transformer中的每一个其他矩阵,从而得到一个坐标完全不同但功能完全相同的函数。

在假设模型为残差流选择任意基底的条件下,我们预期大型特征会“分散”到多个基底坐标上——从期望值来看,它们对每个坐标的贡献幅度大约为 1/√d。

因此,当我们观察到某些残差流维度上持续出现极端值时,这表明模型或其训练过程中的某些因素正在打破这种对称性!那会是什么呢?

实验

首先,我们将使用 Anthropic 的代码库,在一个 2 亿参数的模型上演示这一行为。Dettmers 在此规模下观察到了异常值,但指出它们在更大的模型中才会稳定出现;我们发现,这些异常值在我们的实验中已经足够频繁,使我们能够在相对较小的模型上进行实验。

测量异常值

为了证明我们观察到的现象与 Dettmers 类似,我们将使用他的初始定义来探索我们的模型:将残差流中绝对值大于 6 的单个标量值定义为“异常值”(我们已验证,对于我们的模型,此阈值能筛选出极端的尾部数据,并且在一系列不同的阈值下,我们观察到了性质相似的结果)。然后,我们可以绘制出在模型各层中,曾出现过异常值的残差流激活数量,数据基于一个(128 条序列 × 1024 个 token)的测试批次:

我们看到,随着模型层数的加深,异常值数量逐渐增加,在典型层中,会出现 20 到 60 个异常值维度(该模型的总隐藏维度 d_model=1280)。

激活峰度

虽然这些(部分)与基底对齐的异常值是我们发现异常现象的第一个线索,但我们希望找到一个更通用的指标来检测特权基底的存在——该指标的特点之一是无需调整任意阈值。

我们提出,如果将单个 token 的激活向量视为来自某个概率分布的独立样本,那么若模型不处于特权基(privileged basis)中,该分布的峰度应为 3。请注意,峰度 >3 意味着存在特权基,但峰度等于 3 并不一定意味着不存在特权基——如果特征本身的激活值呈高斯分布,也可能出现峰度为 3 的情况。然而,我们通常直觉认为许多特征是稀疏的,在这种情况下上述结论就不成立。以下是支持这一论点的理由。

如果我们相信某个表示不存在特权基,那么可以预期每个特征都由一个随机方向来表示。高维空间中的随机方向具有什么性质呢?事实证明,存在一种标准的技巧来采样随机单位向量:先从各向同性高斯分布中采样,再将得到的向量缩放至单位范数。请注意,这种方法对大多数分布并不适用——关键在于各向同性高斯分布在任何标准正交基变换下保持不变。所有这些意味着,我们可以将随机方向理解为(经过重新缩放后)来自高斯分布的独立样本构成的向量。需要指出的是,这并非说 n 维球面上的点整体分布是高斯分布,而是说 n 维球面上的任意一点都可以被理解为某个高斯分布采样序列的缩放结果。

如果每个特征都以此方式表示,那么所得激活向量的各分量应呈高斯分布。这是因为激活向量将是每个特征对应的基方向分布的叠加。对高斯分布进行缩放仍得到高斯分布,而高斯变量相加也产生高斯分布。此时,我们可以通过多种方式表征“高斯性”,但我们选择聚焦于峰度。峰度是衡量尾部分布程度的指标;高斯分布的峰度为3;任何更大的值都表明存在重尾。因此,在期望上,这些“高斯样本”的峰度应为三。为准确估计这一期望值,我们计算多个token激活值的峰度并取均值。

将这一指标按层绘制,结果显示我们的分布严重偏离高斯分布(因此并非处于任意基中):

为便于说明并验证我们指标的行为,我们可以观察在施加固定随机旋转后(等价于在某个其他随机选择的标准正交基中观察)的相同残差流数值。如果对所得激活值采用相同的峰度指标,我们会发现数值几乎恰好等于3,与预测一致。

关于基对齐的假设

层归一化

标准Transformer前向传播中一个值得注意的基相关操作是层归一化。层归一化存在两个潜在的基依赖:

  • 它减去均值,这等价于与 \(\frac{1}{d_\text{model}}(1, 1, 1, \ldots{})\) 做点积
  • 它有一个逐通道的可学习权重。应用该权重仍是一个线性操作,但该权重是在标准基中学习并应用的,因此可以认为它在某种程度上赋予了该基以特权。

为检验这一假设,我们修改层归一化以消除基依赖。修改后的操作类似于RMSNorm,后者有时也用于Transformer训练。我们可以将修改后的归一化视为“层归一化,但不减去均值,并使用单个可学习的缩放参数”。即:

\text{RMSNorm}(x_i) = \alpha\cdot\frac{x_i}{\text{RMS}(\bf{x})} ~~~\text{其中}~~~\text{RMS}(\bf{x})=\sqrt{\frac{1}{d}\sum_{i=1}^dx_i^2}

该运算在任何标准正交基下都是相同的。

我们发现,采用这种归一化策略的模型,在我们的训练设置下表现与基线 LayerNorm 完全相同。这能解决基依赖性问题吗?

结果与参考模型大致相似,甚至可以说尾部更重。由此我们得出结论:标准 LayerNorm 中(确实存在的微小)基依赖性并非导致异常值的原因。

有限精度

Neel Nanda 和 Alex Silverstein 推测,基偏好源于 Transformer 实现中使用了有限精度(通常为 16 位或 32 位)浮点数。据我们理解,该假设是:当混合不同量级的特征时,最好将它们放入不同的坐标中,因为浮点数在跨不同尺度求和时会迅速损失精度。

验证模型与基无关

通过修改后的 RMSNorm,我们的模型现在应完全具有旋转不变性,因此我们可以通过实际旋转模型来测试浮点精度假设——至少在前向传播阶段可以做到。

我们生成一个随机标准正交矩阵 R,然后将模型中的每个矩阵乘以 {R}(如果该矩阵从残差流中读取)或 R^\intercal(如果该矩阵向残差流写入)。

随机旋转平均而言确实会带来损失,但数值极小(作为参考,基线模型的测试损失约为 3.02 nats),而且我们观察到来自其他因素(例如模型初始化时不同的随机种子)的噪声要大得多。

由此我们得出结论:该模型真正具有旋转不变性,并且即使在考虑浮点误差的情况下,这一事实仍然成立。

在任意基下进行优化

上述实验基本上排除了 Transformer 前向传播过程中对浮点数精度细节的微妙依赖。然而,仍存在一种可能性:训练过程中的模型梯度(其坐标也可能跨越多个数量级)会以某种方式与浮点数精度相互作用,并且这种相互作用导致了标准基的优先选择。

为了探究这种可能性,我们在训练过程中使用类似的旋转操作来训练一个 Transformer。

具体来说,我们可以在初始化时生成固定的随机旋转矩阵,并在每次从残差流读取或写入激活值时,将这些矩阵乘入其中。由于我们是在全秩乘法操作之前/之后立即执行此操作,因此这不会影响 Transformer 可表达的函数类别,但它确实解耦了模型计算过程中不同阶段所使用的基。

我们以这种方式训练了两个模型,包含两种变体:

一个“共享旋转”模型,在该模型中,我们固定一个单一的随机旋转,并在每次从残差流读取(写入)时应用该旋转(或其转置)。这与前向传播实验的设置类似,区别在于这里我们旋转的是激活值,而不是参数矩阵。在这个模型中,我们实质上解耦了残差流基与计算基;所有计算(注意力层和 MLP 层)都在一个单一的共享基内进行;但信息通过残差流时则经由一个不同的基传递。

一个“独立旋转”模型,在该模型中,每次从残差流读取或写入都使用不同的随机旋转。在这个模型中,每个计算层都在其自身的基中进行,该基与任何其他层或残差流无关。

我们发现,这两个模型的表现与基线模型基本相同(我们在本文档末尾附上了损失曲线)。

由此我们得出结论:即便考虑到浮点精度,Transformer 架构在训练和正常运行时也并不依赖某个特权基。训练过程中的某些动态确实会偏向标准基,但这种效应似乎更多是一种附带结果,而非必要条件。

仅凭这一观察结果,我们将其视为反对"浮点精度"假说的中等强度证据;我们预期,在大多数因浮点精度原因而使得标准基至关重要的场景中,强制模型在多个无关基上运行都会导致显著的性能损失。

不过,我们现在可以利用峰度指标对这些模型进行更深入的研究。正如我们所预期的,两个模型的激活值峰度都几乎恰好等于 3:

此外,我们还可以观察旋转后的激活值,即它们被送入 MLP 上投影层之前的状态。对于"独立旋转"模型,每一层都在不同的基中运行,因此我们不一定预期会看到异常现象。而对于"共享旋转"模型,这让我们得以观察"计算基"——模型实际进行计算的空间,该空间现已与用于层间通信的残差流基分离开来。

在计算基内部,"共享旋转"模型在尾部分布的重度方面与基线模型非常相似。

由此我们得出结论:基对齐现象更多是 Transformer 层内部计算的产物,而非残差流中表征的结果。我们认为这是反对浮点精度假说的有力证据。

有趣的是,尽管"独立旋转"模型显示出一些较小的尾部特征(尤其是在早期层中),但这种效应非常微弱。这表明基对齐现象某种程度上源于所有层在同一个基中运行,并"共谋"产生了这些离群值;单个层只能产生微小影响,而主要效应来自所有层的共同作用。

结论

我们认为这些实验提供了相当有力的证据,表明数值精度并非我们所观察到的异常离群值的驱动因素。虽然论证并非完全无懈可击,但我们认为其说服力很强。

当我们使用不同的基(basis)来分别表示残差流和每一层内部的计算时,观察到在计算基中出现了重尾激活,但在残差流中并未出现。

Adam 优化器会追踪动量,并以逐点方式归一化梯度更新,因此它更偏向于权重存储所在的基,而非参数空间中的任意方向。在排除了本文中讨论的其他效应后,这仍然是我们在 Transformer 模型中所知的、与基相关性最强的操作。因此,这些实验的结论倾向于怀疑优化器是导致离群值近似基对齐的原因。

话虽如此,我们并不能断言已将责任完全归咎于 Adam;有可能还存在另一种尚未被识别的机制在起作用。我们还可以进行其他一些实验来进一步巩固这一假设,但我们决定不再继续,不过这些实验可以作为自然的后续研究方向:

  • 我们可以尝试完全使用 SGD 或其他不依赖基的优化器来训练 Transformer。我们的直觉是,虽然调整学习率会是一个挑战,但使用非常低的学习率和大量的训练步数,或许能够成功训练小规模模型。
  • 我们可以使用 Adam 训练模型,但将 Adam 的动量存储在不同于权重和计算的基中。这将与我们之前“在随机基中进行计算”的实验类似,但能进一步将 Adam 的影响分离出来。

损失曲线

以下是本文档所述模型的训练损失曲线:

我们观察到,在训练初期,不同模型之间存在轻微差异,但在大约十亿个 token 之后,它们收敛到了我们认为处于噪声范围内的水平。

作者贡献

基础结果与实验——Nelson Elhage 设计并执行了实验,包括实现和训练 Transformer 的各种变体。

分析与理论——Nelson Elhage 提出,峰度可能是一种有效的基无关度量指标。Chris Olah 完成了理论工作,论证并选定了我们最终采用的具体公式。

撰写——本文主要由 Nelson Elhage 起草,部分章节由 Chris Olah 撰写。Robert Lasenby 参与了编辑工作并提供了有益的意见。

插图——Shan Carter 设计了本文中的插图,Nelson Elhage 和 Chris Olah 提供了协助。

来源:Anthropic:Transformer Circuits(可解释性研究)· transformer-circuits.pub