Transformer升级之路：无限外推的ReRoPE？

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 追一科技

研究方向 | NLP、神经网络

自从在《Transformer升级之路：将β进制位置进行到底》中引入混合进制的思路进一步推广了 NTK-aware Scaled RoPE 后，笔者感觉类似思路的效果已经达到了上限，想要更大幅度的提升就必须另辟蹊径了。这时候笔者想起了此前构思过的一个思路，该思路由于复杂度较高所以被搁置下了，既然现在已经遇到了瓶颈，那么“唯一的办法就是最好的办法”，于是便将它重拾起来。

万万没想到的是，尽管该方法增加了一些推理复杂度，但它的实验效果却惊人地好——甚至隐约有无限的长度外推能力！因此，笔者迫不及待地撰写了本文来分享该方法。由于形式上跟 ReLU 激活函数的相似性，所以笔者将该方法命名为 “ReRoPE（Rectified Rotary Position Embeddings）”。

重温

我们知道，RoPE 形式上是一种绝对位置编码，但实际上给 Attention 带来的是相对位置信息，即如下的 Toeplitz 矩阵 [1]：

这里的是当前样本长度。当明显超出了训练长度时，多出来的位置由于没有被训练过，所以无法保证效果，这就是直接外推（Length Extrapolation）表现通常比较差的原因。

后来，研究人员提出了位置内插（Position Interpolation），它相当于将相对位置矩阵改为：

这样一来，只要调整，就可以保证最大的相对位置也不超过训练长度，因此避免了外推。然而，它使得位置信息更加“拥挤”了，所以还需要进行一定步数的微调才能让模型重新工作。而也正因为避免了外推，所以它所需要的微调步数相比直接外推要少得多（神经网络往往更擅长内插而不是外推）。

至于后面提出的 NTK-aware Scaled RoPE，则是“剑走偏锋”，巧妙地将外推压力平摊到每一个维度上，所以它不微调也能有不错的效果，但它终究还是依赖外推，这是神经网络不擅长的事情，所以效果存在上限，在笔者的实验中，它的 Long Context 表现还无法很接近训练效果。

融合

我们也可以从语言模型的局域性来考察这些方法。所谓局域性，是指语言模型在推断下一个 token 时，明显更依赖于邻近的 token。直接外推保持了局域性（0 附近位置编码不变），效果差是因为引入了超出训练长度的位置编码；位置内插虽然没有外推位置编码，但扰乱了局域性（0 附近位置编码被压缩为），所以不微调效果也不好；而 NTK-aware Scaled RoPE 通过“高频外推、低频内插”隐含了两者优点，保证了局域性，又没有明显外推位置编码，所以不微调也有不错的效果。

有没有能更直接地结合外推和内插的方法呢？有，我们可以设定一个窗口大小，在窗口内我们使用大小为 1 的位置间隔，在窗口外我们使用大小为的位置间隔，整个相对位置矩阵如下：

只要小于训练长度，那么通过控制，我们就可以在精确保持了局域性的前提下，使得所有位置编码不超过训练长度，简单直接地结合了直接外推和位置内插。

特别地，矩阵（3）还有一个特别的 case：当时，它简化为

在这个 case 下，不管输入长度是多少，它的位置编码范围都不超过 w，所以这是一种有可能支持任意长度的 Context 的方案！

形式上，矩阵（4）、（3）与标准 RoPE 矩阵（1）的关系，就相当于 ReLU、Leaky ReLU 与 Linear 的关系，所以笔者将（4）称为 “ReRoPE（Rectified RoPE）”，将（1）称为 “Leaky ReRoPE”。

计算

其实，类似的思路并不难想到，以往基于 Attention Bias 的相对位置编码（比如经典相对位置编码、T5 位置编码）经常会出现这样的分块运算。然而跟这些相对位置编码不同，在 RoPE 中实现这样的分块运算会明显增加计算量，这也是该思路会被笔者搁置的主要原因。

怎么理解增加计算量呢？我们知道 RoPE 是“通过绝对位置实现相对位置”，这样只能得到线性的相对位置，而矩阵（3）、（4）是非线性的（或者说分段线性的），要实现它只能算两次 Attention 矩阵，然后组合起来。具体来说，首先用标准的 RoPE 计算一次 Attention 矩阵

这里第一个等号是实现方式，第二个等号是等效结果，其中就是 RoPE 的旋转矩阵，简单起见我们省略了 Attention 的 scale 因子。接着，我们需要计算间隔为的 RoPE 的 Attention 矩阵（Leaky ReRoPE）：

如果是 ReRoPE，那么简单一些：

最后，根据这个条件，将它们合并起来：

不管是 ReRoPE 还是 Leaky ReRoPE，都不可避免地计算两次 Attention 矩阵（如果有更高效的实现方法，请赐教），这便是增加的计算量之一。此外，需要自定义计算 Attention 矩阵也导致了不能直接套用现成的 flash attention 实现，因此相对之下又增加了计算成本。

另一方面，同样是由于非线性的相对位置，所以在自回归解码时，Key 序列的 cache 只能存 RoPE 之前的，然后在每步解码时给整个 Key 序列补上对应的RoPE，这样的改动也会增加推理计算量。唯一的好消息是，在 token by token 解码时，从第二步开始 Query 序列的长度就为 1，此时只需要为 Key 序列定制 RoPE，那么可以只算一次 Attention 矩阵：

实验

继续沿着《Transformer升级之路：将β进制位置进行到底》的设置，我们对 ReRoPE 进行了实验，效果如下表：

正如文章开头所说，ReRoPE 不微调外推的效果可谓出奇地好，不仅明显超越了此前最优的 NTK-RoPE-mixed，还明显超过了从零预训练的 HFWA！这里的指的，是指预训练没有加入缩放（比如 LLAMA），测试阶段每个都乘上，则是指预训练就加入了缩放因子。

以下是一些消融实验，显示出 ReRoPE 关于还是很鲁棒的，最优值大致是训练长度的左右：

下表则对比了 ReRoPE 和 Leaky ReRoPE：

作为 ReRoPE 的一般化，经过精调的 Leaky ReRoPE 是有机会超过 ReRoPE 的，但提升很微弱。此外，当取有限值时，能处理的最大长度也是有限的，因为我们不能提前知道要生成的总长度，所以只能预设一个足够大的，但设定为有限值之后，当输入足够长时，就会因为位置编码超出训练长度而效果大幅下降，相比之下 ReRoPE 则不会有这个风险。总的来说，精调 Leaky ReRoPE 相比 ReRoPE 的价值似乎不大。

最后，笔者在 transformers 的 LLAMA 模型基础上给出了的 ReRoPE 和 Leaky ReRoPE 的实现，在 OpenBuddy 开源的 LLAMA2-13b 微调模型上测试自我感觉还不错（最多测试过 20k tokens 的 Context），读者也可以自行加载 LLAMA 系列模型进行测试：

Github：

https://github.com/bojone/rerope

小结

在这篇文章中，笔者提出了 ReRoPE（Rectified RoPE），它同样是一种 RoPE 的后处理方案，实验结果显示它的不微调长度外推能力不仅明显超过了此前的 NTK-aware Scaled RoPE，甚至还超过了之前专门设计的需要从零训练的 HFWA。此外，不同于 NTK-aware Scaled RoPE 在超过某个长度后能力会大幅下降，ReRoPE 似乎在任意长度下都表现良好。除了对比实验外，文章还给出了基于 transformers-llama 的参考实现，有兴趣的读者可以自行测试。

更多阅读