Transformer升级之路：Key归一化助力长度外推作者： PaperWeekly 来源： PaperWeekly ©PaperWeekly 原创 · 作者 | 苏剑林单位 | 月之暗面研究方向 | NLP、神经网络大体上，我们可以将目前 Transformer 的长度外推技术分为两类：一类是事后修改，比如 NTK-RoPE、YaRN [1]、ReRoPE 等，这类方法的特点是直接修改

Transformer升级之路：Key归一化助力长度外推

By AiBard123
December 18, 2023 - 2 min read

作者： PaperWeekly 来源： PaperWeekly

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 月之暗面

研究方向 | NLP、神经网络

大体上，我们可以将目前 Transformer 的长度外推技术分为两类：一类是事后修改，比如 NTK-RoPE、YaRN [1]、ReRoPE 等，这类方法的特点是直接修改推理模型，无需微调就能达到一定的长度外推效果，但缺点是它们都无法保持模型在训练长度内的恒等性。

另一类自然是事前修改，如 ALIBI、KERPLE、XPOS 以及 HWFA 等，它们可以不加改动地实现一定的长度外推，但相应的改动需要在训练之前就引入，因此无法不微调地用于现成模型，并且这类方法是否能够 Scale Up 还没得到广泛认可。

在这篇文章中，笔者将介绍一种意外发现的长度外推方案——“KeyNorm”——对 Attention 的 Key 序列做 L2 Normalization，很明显它属于事前修改一类，但对 Attention 机制的修改非常小，因此看上去非常有希望能够 Scale Up。

最初动机

之所以说“意外发现”，是因为该改动的原始动机并不是长度外推，而是尝试替换 Scaled Dot-Product Attention 中的 Scale 方式。我们知道，Attention 的标准定义是（本文主要考虑 Causal 场景）

其中，Scale 因子我们已经多次进行过解释甚至推广，比如《浅谈Transformer的初始化、参数化与标准化》[2]、《从熵不变性看Attention的Scale操作》、《从梯度最大化看Attention的Scale操作》等。标准的推导是在 “ 均独立地采样自“均值为 0、方差为 1”的分布”的假设下进行的，而在该假设之下，我们还有

这是因为

相关推广还可以参考《让人惊叹的Johnson-Lindenstrauss引理：理论篇》。这个近似式意味着，在 Attention 的初始阶段式（1）与下面两个变体有着相同的效果：

因此，就有了验证这两个变体与标准的式（1）哪个更优的想法。为了描述的方便，我们可以相应地称为 “Query/Key-Normalized Dot-Product Attention”，分别简称为 “QNA” 和 “KNA”。

此外，既然可以 QueryNorm 和 KeyNorm，那么自然也可以考虑两者都 Norm一下，所以我们将如下 “Scaled Cosine Attention（CosA）”也一并进行实验：

其中采用《从梯度最大化看Attention的Scale操作》中的结果，即（原文是 3.5，但下面训练长度比较小，改为 4 更精准一些），其中固定为训练长度的一半，或者动态取位置 id 加 1。

先看结果

沿着之前 [3] 做长度外推的实验设置，都是 1 亿参数的小模型，GAU 架构，训练相同的步数（时间有限，这个步数下其实模型还没训练充分），训练长度 512，并考虑外推到 4096 长度，实验结果如下表。其中 Baseline 就是式（1），就是加入《从熵不变性看Attention的Scale操作》介绍的长度相关的缩放因子。评价指标是语言模型的逐 token 准确率，越大越好。

从表格中我们可以看出：1）不管是 QueryNorm 还是 KeyNorm，都在训练长度上取得了更好的效果，虽然这个优势非常微弱，大概率随着训练的进一步推进可以忽略不计，但这个优势非常稳定，暗示着让训练更加平稳的可能性；2）KeyNorm对长度外推的提升非常明显 ，这就是实验结果中的“意外之喜”！

注意，跟 NTK-RoPE、YaRN 需要修改推理阶段的模型不同，这里的 KNA 和 CosA 的长度外推在推理阶段是完全不做改动的。因此，可能有读者想知道，既然 KNA 和 CosA 推理时不加改动外推效果都这么好了，如果配合 NTK-RoPE、YaRN 等外推技巧，效果会不会“更上一层楼”？对此，笔者也进行了测试，结果如下表：

这个表比较啰嗦，主要是为了让大家对主流长度外推技巧的效果差异有一个全面的感知，大家选择自己感兴趣的维度比较就好，但要注意如果看长度外推效果的话，应该以“不重复”一列为主，“重复”一列为辅。

从上表看，结果着实有点让人意外，KeyNorm 似乎“免疫”了已有的 RoPE 外推技巧，NTK、YaRN 等技巧叠加上去并没有明显提升，甚至可能会下降，不过总体来看“重复”一列还是有显著提升的，不显著的是“不重复”一列。这些结果表明，KeyNorm 依然有着无法有效识别超出训练长度的位置（所以“重复”的结果不高）的问题，但有效地避免了 PPL 爆炸问题（所以“不重复”的结果还不错）。

这对做 Long Context 的同学来说可能是个好消息：一方面，KeyNorm 不像 ALIBI、KERPLE 等，它的长度外推不用加 Local 约束，训练完成后也不做任何修改，纯属是“免费的午餐”，甚至看上去加了 KeyNorm 后训练效果都变好了；另一方面，也因为它是非 Local 的，所以可以更长文本继续训练，并且继续训练时再也不用纠结是选 PI [4] 还是 ABF [5] 了，对于 KeyNorm 来说，啥也不改就行。

原理分析

尽管这是个意外发现，但我们仍需要尝试去解释它，不然它就一直只是个意外。所以这一节我们尝试来思考，为什么 KeyNorm 会有助于长度外推。

让我们重新回到式（1），第个 token 与第个token的相关性打分由内积完成：

第二个等号，我们从几何意义出发，将它分解为了各自模长与夹角余弦的乘积。注意力是一个条件概率，只跟当前位置 i 有关，它不改变注意力的相对大小，而只改变稀疏程度；则有能力改变的相对大小，但它不涉及到 i,j 的交互，可以用来表达一些绝对信号，比如 Scissorhands [6] 表明某些绝对位置的 token 的注意力一直都会很高，这就有可能用来表达；剩下的就是用来表达的交互，它是自由度最大的一项。

很明显，为了提高某个位置的相对重要性，模型有两个选择：1、增大模长；2、增大，即缩小的夹角大小。然而，由于“维度灾难”[7] 的存在，在高维空间中显著地改变夹角大小相对来说没有那么容易，所以如果能靠增大模长完成的，模型会优先选择通过增大模长来完成，这导致的直接后果是：的训练可能并不充分。

这里笔者作出一个断言（猜测）：

的训练不充分是 Attention 无法长度外推的主要原因。

的训练不充分，是指被训练过的的夹角只是一个有限的集合，而进行长度外推时，它要面对一个更大的集合，从而无法进行正确的预测。

仔细思考 YaRN [8] 一文的推导就会发现，NTK、YaRN 之所以有效，是因为修改了推理阶段 RoPE 的实现，使得的夹角落到原本训练阶段的有限集合中，避免面对没见过的更大的集合，转外推为内插；ReRoPE 则更加干脆，直接截断 Window 以外的相对位置，这使得推理阶段的位置编码都不会“面生”。这些技巧一定程度上都间接地验证了这个断言。

从这个断言出发，KeyNorm 的长度外推起因就变得简单了。不论是只进行 KeyNorm 的 KNA，还是 QueryNorm、KeyNorm 都做的 CosA，它们都将从 Attention 的定义中排除掉了，于是为了改变 j 的相对重要性，模型就只有“调整 ”这一个选择，这将会使得模型更加充分地训练和利用，从而间接促进了长度外推性。

此外，笔者也实验过 “KeyNorm + NoPE” 的组合，但并没有发现长度外推性，这说明 RoPE 也在 KeyNorm 的长度外推中担任重要角色。事实上这也不难理解，RoPE 对进行旋转，更有助于扩大训练期间的范围，从而使得的训练更为充分。

有没有工作已经尝试过 QueryNorm 和 KeyNorm 了呢？有。2020 年的论文《Query-Key Normalization for Transformers》[9] 曾实验过 CosA，论文还提出了一个类似的长度对数的 Scale 因子，但没有讨论到长度外推问题。

此外，今年初 Google 的论文《Scaling Vision Transformers to 22 Billion Parameters》[10] 也在 Query 和 Key 加了 Norm，但加的是 LayerNorm，LayerNorm 或者 RMSNorm 都带有可学的 gamma 参数，这使得 Norm 之后的向量模长未必为常数，因此并不好说是否能达到本文一样的长度外推效果。

文章小结

本文介绍了笔者意外发现的一种长度外推方案 “KeyNorm”——对 Attention 的 Key 序列进行 L2 归一化，在训练长度上取得了更好的效果，并在长度外推方面表现出显著的提升。它属于“事前修改”方案，跟其他事前修改方案如 ALIBI、KERPLE 等相比，它没有 Local 约束，因此更有希望能够 Scale Up；相比于 NTK-RoPE、YaRN 等“事后修改”方案，它在外推的时候则不会损失训练长度内的性能。