破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

378 0 1

随着大

语言模型

的快速发展，其长度外推能力（length extrapolating）正日益受到研究者的关注。尽管这在 Transformer 诞生之初，被视为天然具备的能力，但随着相关研究的深入，现实远非如此。传统的 Transformer 架构在训练长度之外无一例外表现出糟糕的推理性能。

研究人员逐渐意识到这一缺陷可能与位置编码（position encoding）有关，由此展开了绝对位置编码到相对位置编码的过渡，并产生了一系列相关的优化工作，其中较为代表性的，例如：旋转位置编码（RoPE）(Su et al., 2021)、Alibi (Press et al., 2021)、Xpos (Sun et al., 2022) 等，以及近期 meta 研发的位置

插值

（PI）(Chen et al., 2023)，reddit 网友给出的 NTK-aware Scaled RoPE (bloc97, 2023)，都在试图让模型真正具备理想中的外推能力。

然而，当研究人员全力将目光放在位置编码这一众矢之的上时，却忽视了 Transformer 中另一个重量级角色 –self-attention 本身。蚂蚁

人工智能

团队最新研究表明，这一被忽视的角色，极有可能成为扭转局势的关键。Transformer 糟糕的外推性能，除了位置编码外，self-attention 本身仍有诸多未解之谜。

基于此发现，蚂蚁

人工智能

团队自研了新一代

注意力机制

，在实现长度外推的同时，模型在具体任务上的表现同样出色。

破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

论文地址：https://arxiv.org/abs/2309.08646
Github仓库：https://github.com/codefuse-ai/Collinear-Constrained-Attention
ModelScope：https://modelscope.cn/models/codefuse-ai/Collinear-Constrained-Attention/summary
HuggingFace：敬请期待

背景知识

在深入探讨之前，我们快速回顾一些核心的背景知识。

长度外推 (Length Extrapolating)

长度外推是指大

语言模型

在处理比其训练数据中更长的文本时的能力。在训练大型

语言模型

时，通常有一个最大的序列长度，超过这个长度的文本需要被截断或分割。但在实际应用中，用户可能会给模型提供比训练时更长的文本作为输入，如果模型欠缺长度外推能力或者外推能力不佳，这将导致模型产生无法预期的输出，进而影响模型实际应用效果。

自注意力

(Self-Attention)

(Vaswani et al., 2017) 于 2017 年提出的 multi-head self-attention，作为如今大

语言模型

的内核，对于推动

人工智能

领域的发展起到了举足轻重的作用。这里以下图 1 给出形象化的描述，这项工作本身已经被广泛认可，这里不再进行赘述。初次接触大

语言模型

，对这项工作不甚了解的读者可以前往原论文获取更多细节 (Vaswani et al., 2017)。

破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

^{图1. 多头

注意力机制

示意图，引自（Vaswani, et al., 2017)。}

位置编码 (Position encoding)

由于 self-attention 机制本身并不直接处理序列中的位置信息，因此引入位置编码成为必要。由于传统的 Transformer 中的位置编码方式由于其外推能力不佳，如今已经很少使用，本文不再深入探讨传统的 Transformer 中的编码方法，对于需要了解更多相关知识的读者，可以前往原论文查阅详情 (Vaswani et al., 2017)。在这里，我们将重点介绍目前非常流行的旋转位置编码（RoPE）(Su et al., 2021)，值得一提的是，Meta 的 LLaMa 系列模型 (Touvron et al., 2023a) 均采用了此种编码方式。

RoPE 从建模美学的角度来说，是一种十分优雅的结构，通过将位置信息融入 query 和 key 的旋转之中，来实现相对位置的表达。

^{图2. 旋转位置编码结构，引自(}

^{Su et al., 2021}

^)。

位置

插值

(Position Interpolation)

尽管 RoPE 相比绝对位置编码的外推性能要优秀不少，但仍然无法达到日新月异的应用需求。为此研究人员相继提出了各种改进措施，以 PI (Chen et al., 2023) 和 NTK-aware Scaled RoPE (bloc97, 2023) 为典型代表。但要想取得理想效果，位置

插值

仍然离不开微调，实验表明，即使是宣称无需微调便可外推的 NTK-aware Scaled RoPE，在传统 attention 架构下，至多只能达到 4～8 倍的外推长度，且很难保障良好的语言建模性能和长程依赖能力。

破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

^{图3. 位置

插值

示意图，引自(}

^{Chen et al., 2023}

^)。

CoCA

过去的研究主要集中在位置编码上，所有相关研究工作均默认 self-attention 机制已经被完美实现。然而，蚂蚁

人工智能

团队近期发现了一个久被忽视的关键：要从根本上解决 Transformer 模型的外推性能问题，self-attention 机制同样需要重新考量。

^{图4. CoCA 模型架构，引自(}

^{Zhu et al., 2023}

^)。

RoPE 与 self-attention 的异常行为

在 Transformer 模型中，self-attention 的核心思想是计算 query（q）和 key（k）之间的关系。

注意力机制

使用这些关系来决定模型应该 “关注” 输入序列中的哪些部分。

考虑输入

和

，分别代表输入序列中的m个和第n个位置索引对应的 token，其 query 和 key 分别为
破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

和

。它们之间的注意力可以表示为一个函数
破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

，如果应用 RoPE，则可以进一步简化为仅依赖于m和n相对位置的函数
破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

。

从数学角度看，

可以解释为

组复数（

，其中d为 hidden dimension 维度，省略位置索引m,n）经过旋转后的内积之和。这在直觉上是有意义的，因为位置距离可以建模为一种序，并且两个复数的内积随着旋转角度
破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

的变化而变化，以图 5 为例，其中
破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

，

为

和

的初始夹角。

然而，它隐藏着一个此前一直被忽略的技术缺陷。

为了便于理解，我们首先考虑双向注意力模型，例如 Bert (Devlin et al., 2019) 和 GLM (Du et al., 2021) 等。如图 5 所示，对于
破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

组复数中的任意一组

，

，它们分别具有位置索引m和n。

不失一般性，我们假设复平面上有一个小于
破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

的角度

逆时针从

旋转到

，那么它们的位置关系有两种可能的情况（不考虑 =，因为它是平凡的）。

正常保序关系

：当

时，如图 5 右侧所示。注意力分数随着位置距离的增加而降低（直到它们相对角度超出

，超出

的这部分在原论文附录中进行了讨论 (Zhu et al., 2023)）。
异常行为

：然而，当

时，如图 5 左侧所示，异常行为打乱了

个最邻近的 token 的序。当

时，

和

之间的相对角度将随着

的增大而减小，这意味着最接近的 token 可能会获得较小的注意力分数。（我们在这里使用 “可能”，因为注意力分数是

个内积的总和，也许其中一个是微不足道的。但是，后续实验证实了这一重要性。）并且，无论应用 PI 还是 NTK-aware Scaled RoPE，均无法消除这一影响。

破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

图5.双向模型中的序被破坏，引自(Zhu et al., 2023)。

对于因果模型来说，虽然m总是大于n，但问题同样存在。如图 6 所示，对于某些j，当存在小于
破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

的角度

从

逆时针旋转到

时，而不是从

到

。

破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

^{图6.因果模型中的序被破坏，引自(Zhu et al., 2023)。}

共线约束注意力（CoCA）

基于以上 RoPE 与 self-attention 的异常行为分析，要从根本上解决这个问题，仅从位置编码入手很显然是药不对症，根本的解决方法是让 self-attention 中的 query 和 key 初始夹角为 0，这是论文中共线约束（Collinear Constrained Attention）的由来。

详细的推导和公式，这里不进行一一展开，读者可以阅读原文进行深入理解。

需要特别指出的是，论文给出的理论解释中，表明：

稳定的远程衰减特性

：CoCA 相对于 RoPE 显示出了更为稳定的远程衰减特性。
显存瓶颈与解决方案

：CoCA 有引入显存瓶颈的风险，但论文给出了十分高效的解决方案，使得 CoCA 的计算和空间复杂度几乎与最初版本的 self-attention 无异，这是十分重要的特性，使得 CoCA 的实用性十分优异。
无缝集成

：CoCA 可以与当前已知的

插值

手段（论文中实验了 NTK-aware Scaled RoPE）无缝集成，且在无需微调的情况下取得了远超原始 attention 结构的性能，这意味着使用 CoCA 训练的模型，天然就拥有近乎无限的外推能力，这是大

语言模型

梦寐以求的特性。

实验结论

论文比较了 CoCA 与 RoPE (Su et al., 2021)、 ALibi (Press et al., 2021) 在外推性能上的差异，取得了令人振奋的结果。相应模型记为：

Origin：原始 attention 结构，位置编码方式为 RoPE
ALibi：原始 attention 结构，位置编码方式为 ALibi
CoCA：论文模型结构，位置编码方式为 RoPE

详细的实验背景，可以查阅原文。

长文本建模能力

论文评估了 CoCA 和 Origin、ALibi 模型的长文本语言建模能力。该评估使用了 100 个文档，每个文档至少拥有 8,192 个 token，这些文档随机采样于 PG-19 数据集 (Rae et al., 2019)。所有 3 个模型的训练长度均为 512，模型大小为 350M。

图 7 说明了一个值得注意的趋势：当推理长度超出其训练长度，Origin 模型的困惑度迅速偏离（>1000）。相比之下，CoCA 模型能够维持低困惑度，即使是在其训练长度的 16 倍时，困惑度没有出现发散趋势。

NTK-aware Scaled RoPE (bloc97, 2023) 作为一种无需微调的外推方法，论文中允许在实验中应用该方法，但即使在 Origin 模型上应用了动态 NTK 方法，其困惑度仍然远高于 CoCA。

ALibi 在困惑度上的表现最好，而 CoCA 在应用动态 NTK 方法后，可以取得与 ALibi 相近的结果。

破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

^{图7.滑动窗口困惑度测试结果，引自(}

^{Zhu et al., 2023}

^)。

捕捉长程依赖

困惑度是衡量

语言模型

预测下一个 token 的熟练程度的指标。然而，它并不能完全代表一个理想的模型。因为虽然局部注意力在困惑度上表现出色，但它在捕获长程依赖性方面往往表现不佳。

为了深入评估这一问题，论文采用 (Mohtashami & Jaggi, 2023) 提出的密钥检索综合评估任务评估了 CoCA 和 Origin、ALibi 模型。在此任务中，有一个隐藏在长文档中的随机密钥需要被识别和检索。

如图 8 所示，像 ALibi 这一类具有一定局部假设的模型，尽管在困惑度任务上表现良好，但在捕捉长程依赖时具备无法弥补的劣势，在外推 1 倍长度时，

准确率

就开始迅速下降，最终下跌至 10% 以下。相比之下，即使测试序列长度扩展到原始训练长度的 16 倍，CoCA 也始终表现出很高的准确性，在 16 倍外推长度时仍然超过 60%。比 Origin 模型整体高出 20%，而比 ALibi 整体高出 50% 以上。

破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

^{图8.随机密钥识别检索性能曲线，引自(Zhu et al., 2023)。}

超参数

稳定性

由于在实验中应用了动态 NTK 方法，论文针对 Origin 和 CoCA 模型在动态 NTK 方法下的 scaling factor

超参数

稳定性进行了深入探讨。

如图 9 所示，Origin 模型在不同 scaling factor 下存在剧烈波动（200～800），而 CoCA 模型则处于相对稳定的区间（60～70）。更进一步，从 Table 4 中的明细数据可以看出，CoCA 模型最差的困惑度表现仍然比 Origin 模型最好的困惑度表现要好 50% 以上。

在 Passkey 实验中，Origin 和 CoCA 模型表现出了与困惑度实验中类似的特性，CoCA 模型在不同的 scaling factor 均有较高的

准确率

，而 Origin 模型在 scaling factor=8 时，

准确率

下跌至了 20% 以下。更进一步，从 Table 5 的明细数据可以看出，Origin 模型即使在表现最好的 scaling factor=2 时，仍然与 CoCA 模型有 5%～10% 的

准确率

差距。

同时，Origin 模型在 scaling factor=2 时困惑度表现却是糟糕的，这从侧面反映了原始的 attention 结构在长度外推时很难同时保证困惑度和捕捉长程依赖两方面的性能，而 CoCA 做到了这一点。

破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

^{图9. Origin model 和 CoCA 在不同 scaling factor 的困惑度，引自(Zhu et al., 2023)}

破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

^{图10. Origin 模型和 CoCA 在不同 scaling factor 的通行密钥准确性，引自(Zhu et al., 2023)}

破解自注意力推理缺陷的奥秘，蚂蚁自研新一代Transformer或实现无损外推

长度外推中的注意力得分

如同 PI (Chen et al., 2023) 论文所探究的一样，大

语言模型

在长度外推中的失败与注意力得分的异常值（通常是非常大的值）直接相关。论文进一步探讨了这一现象，这一现象也从侧面说明了为什么 CoCA 模型在长度外推中的表现要优于传统 attention 结构。

实验使用了 PG-19 数据集 (Rae et al., 2019) 中的一个随机片段，长度为 1951 个 token，大约是模型训练长度的 4 倍。

如图 11 所示，其中 (a1) 是 Origin 和 CoCA 模型在不使用动态 NTK 方法时的各层注意力得分，(b1) 是使用动态 NTK 方法后的得分，low layers 表示模型的第 6、12、18 层，last layer 表示第 24 层，(a2) 是 (a1) 在最后 500 个 token 的放大版本，(b2) 同理。

从 (a1) 和 (b1) 可以发现，Origin 模型的注意力得分存在少量异常值，数值比 CoCA 模型注意力得分大 10～20 倍。
由于这些异常值影响了观察效果，(a2) 局部放大了最后 500 个 token，可以看到 Origin 模型的 last layer 注意力得分几乎为 0，这说明 Origin 模型在长度外推时，发生了关注邻近 token 的失效。
从 (b2) 可以看出，当应用动态 NTK 方法后，Origin 模型在邻近 token 处的注意力得分变得异常大，这一异常现象与前文论证的 RoPE 与 self-attention 的异常行为息息相关，Origin 模型在邻近 token 处可能存在着严重的

过拟合

。