Transformer升级之路：RoPE是一种β进制编码

2023-07-22 15:07

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 追一科技

研究方向 | NLP、神经网络

对关心如何扩展 LLM 的 Context 长度的读者来说，上周无疑是激动人心的一周，开源社区接连不断地出现令人振奋的成果。首先，网友 @kaiokendev 在他的项目 SuperHOT [1] 中实验了“位置线性内插”的方案，显示通过非常少的长文本微调，就可以让已有的 LLM 处理 Long Context。

几乎同时，Meta 也提出了同样的思路，带着丰富的实验结果发表在论文《Extending Context Window of Large Language Models via Positional Interpolation》[2] 上。惊喜还远不止此，随后网友 @bloc97 提出了 NTK-Aware Scaled RoPE [3]，实现了不用微调就可以扩展 Context 长度的效果！

以上种种进展，尤其是 NTK-Aware Scaled RoPE，迫使笔者去重新思考 RoPE 的含义。经过分析，笔者发现 RoPE 的构造可以视为一种进制编码，在这个视角之下，开源社区的这些进展可以理解为对进制编码编码的不同扩增方式。

进制表示

假设我们有一个 1000 以内（不包含 1000）的整数要作为条件输入到模型中，那么要以哪种方式比较好呢？

最朴素的想法是直接作为一维浮点向量输入，然而 0～999 这涉及到近千的跨度，对基于梯度的优化器来说并不容易优化得动。那缩放到 0～1 之间呢？也不大好，因为此时相邻的差距从 1 变成了 0.001，模型和优化器都不容易分辨相邻的数字。总的来说，基于梯度的优化器都有点“矫情”，它只能处理好不大不小的输入，太大太小都容易出问题。

所以，为了避免这个问题，我们还需要继续构思新的输入方式。在不知道如何让机器来处理时，我们不妨想想人是怎么处理呢。对于一个整数，比如 759，这是一个 10 进制的三位数，每位数字是 0～9。既然我们自己都是用 10 进制来表示数字的，为什么不直接将 10 进制表示直接输入模型呢？也就是说，我们将整数以一个三维向量来输入，分别是的百位、十位、个位。这样，我们既缩小了数字的跨度，又没有缩小相邻数字的差距，代价了增加了输入的维度——刚好，神经网络擅长处理高维数据。

如果想要进一步缩小数字的跨度，我们还可以进一步缩小进制的基数，如使用 8 进制、6 进制甚至 2 进制，代价是进一步增加输入的维度。

直接外推

假设我们还是用三维 10 进制表示训练了模型，模型效果还不错。然后突然来了个新需求，将上限增加到 2000 以内，那么该如何处理呢？

如果还是用 10 进制表示的向量输入到模型，那么此时的输入就是一个四维向量了。然而，原本的模型是针对三维向量设计和训练的，所以新增一个维度后，模型就无法处理了。可能有读者想说，为什么不能提前预留好足够多的维度呢？没错，是可以提前预留多几维，训练阶段设为 0，推理阶段直接改为其他数字，这就是外推（Extrapolation）。

▲ 直接外推

然而，训练阶段预留的维度一直是 0，如果推理阶段改为其他数字，效果不见得会好，因为模型对没被训练过的情况不一定具有适应能力。也就是说，由于某些维度的训练数据不充分，所以直接进行外推通常会导致模型的性能严重下降。

线性内插

于是，有人想到了将外推改为内插（Interpolation），简单来说就是将 2000 以内压缩到 1000 以内，比如通过除以 2，1749 就变成了 874.5，然后转为三维向量 [8,7,4.5] 输入到原来的模型中。从绝对数值来看，新的实际上对应的是 1498，是原本对应的 2 倍，映射方式不一致；从相对数值来看，原本相邻数字的差距为 1，现在是 0.5，最后一个维度更加“拥挤”。所以，做了内插修改后，通常都需要微调训练，以便模型重新适应拥挤的映射关系。

▲ 线性内插

当然，有读者会说外推方案也可以微调。是的，但内插方案微调所需要的步数要少得多，因为很多场景（比如位置编码）下，相对大小（或许说序信息）更加重要，换句话说模型只需要知道 874.5 比 874 大就行了，不需要知道它实际代表什么多大的数字。而原本模型已经学会了 875 比 874 大，加之模型本身有一定的泛化能力，所以再多学一个 874.5 比 874 大不会太难。

不过，内插方案也不尽完美，当处理范围进一步增大时，相邻差异则更小，并且这个相邻差异变小集中在个位数，剩下的百位、十位，还是保留了相邻差异为1。换句话说，内插方法使得不同维度的分布情况不一样，每个维度变得不对等起来，模型进一步学习难度也更大。

进制转换

有没有不用新增维度，又能保持相邻差距的方案呢？有，我们也许很熟悉，那就是进制转换！三个数字的 10 进制编码可以表示 0～999，如果是 16 进制呢？它最大可以表示。所以，只需要转到 16 进制，如 1749 变为，那么三维向量就可以覆盖目标范围，代价是每个维度的数字从 0～9 变为 0～15。

▲ 进制转换

仔细想想，就会发现这真是一个绝妙的想法。刚才说到，我们关心的场景主要利用序信息，原来训练好的模型已经学会了，而在 16 进制下同样有，比较规则是一模一样的（模型根本不知道你输入的是多少进制）。唯一担心的是每个维度超过 9 之后（10～15）模型还能不能正常比较，但事实上一般模型也有一定的泛化能力，所以每个维度稍微往外推一些是没问题的。所以，这个转换进制的思路，甚至可能不微调原来模型也有效！另外，为了进一步缩窄外推范围，我们还可以换用更小的进制而不是 16 进制。

接下来我们将会看到，这个进制转换的思想，实际上就对应着文章开头提到的 NTK-aware scaled RoPE！

位置编码

为了建立起它们的联系，我们先要建立如下结果：

位置的旋转位置编码（RoPE），本质上就是数字的进制编码！

看上去可能让人意外，因为两者表面上迥然不同。但事实上，两者的运算有着相同的关键性质。为了理解这一点，我们首先回忆一个 10 进制的数字，我们想要求它的进制表示的（从右往左数）第 m 位数字，方法是

也就是先除以次方，然后求模（余数）。然后再来回忆 RoPE，它的构造基础是 Sinusoidal 位置编码，可以改写为

其中，。现在，对比式（1），式（2）是不是也有一模一样的？至于模运算，它的最重要特性是周期性，式（2）的是不是刚好也是周期函数？所以，除掉取整函数这个无关紧要的差异外，RoPE（或者说 Sinusoidal位置编码）其实就是数字的进制编码！

建立起这个联系后，前面几节讨论的整数的扩增方案，就可以对应到文章开头的各种进展上了。其中，直接外推方案就是啥也不改，内插方案就是将换成，其中是要扩大的倍数，这就是 Meta 的论文所实验的 Positional Interpolation，里边的实验结果也证明了外推比内插确实需要更多的微调步数。

至于进制转换，就是要扩大倍表示范围，那么原本的进制至少要扩大成进制（式（2）虽然是维向量，但是成对出现的，所以相当于位进制表示，因此要开次方而不是次方），或者等价地原来的底数换成，这基本上就是 NTK-Aware Scaled RoPE [3]。跟前面讨论的一样，由于位置编码更依赖于序信息，而进制转换基本不改变序的比较规则，所以 NTK-Aware Scaled RoPE 在不微调的情况下，也能在更长 Context 上取得不错的效果。

追根溯源

可能有读者好奇，这跟 NTK 有什么关系呢？NTK 全称是 “Neural Tangent Kernel”，我们之前在《从动力学角度看优化算法：SGD ≈ SVM？》也稍微涉及过。要说上述结果跟NTK的关系，更多的是提出者的学术背景缘故，提出者对《Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains》[4] 等结果比较熟悉，里边利用 NTK 相关结果证明了神经网络无法直接学习高频信号，解决办法是要将它转化为 Fourier 特征——其形式就跟式（1）的 Sinusoidal 位置编码差不多。

所以，提出者基于 NTK 相关结果的直觉，推导了 NTK-Aware Scaled RoPE。笔者向提出者请教过他的推导，其实他的推导很简单，就是把外推和内插结合起来——高频外推、低频内插。具体来说，式（2）最低频是项，引入参数变为，让它跟内插一致，即

那么解得。至于最高频是项，引入后变为，由于通常很大，很接近 1，所以它还是接近于，即等价于外推。

所以这样的方案简单巧妙地将外推和内插结合了起来。另外，由于比较大（BERT 是 64，LLAMA 是 128），跟差别不大，所以它跟笔者基于进制思想提出的解是基本一致的。还有，从提出者这个思想来看，任意能实现“高频外推、低频内插”的方案都是可以的，并非只有上述引入的方案，这个读者可以亲自尝试一下。

个人测试

作为号称不用微调就可以增加 LLM 的 Context 长度的方案，笔者第一次看到 NTK-Aware Scaled RoPE 时，也感到很震惊，并且迫不及待地去测试它。毕竟根据《Transformer升级之路：一种全局长度外推的新思路》的经验，在笔者所偏爱的 “GAU+Post Norm” 组合上，很多主流的方案都失效了，那么这个方法又如何？

当取 8 时，对比结果如下（关于“重复”与“不重复”的区别，可以参考这里）

以上报告的都是没有经过长文本微调的结果，其中 Baseline 就是外推，PI（Positional Interpolation）就是 Baseline 基础上改内插，NTK-RoPE 就是 Baseline 基础上改 NTK-Aware Scaled RoPE。带的选项，是指预训练时加入了《从熵不变性看Attention的Scale操作》中的 scale，考虑这个变体是因为笔者觉得 NTK-RoPE 虽然解决了 RoPE 的长度泛化问题，但没有解决注意力不集中问题。

表格的实验结果完全符合预期：

1. 直接外推的效果不大行；

2. 内插如果不微调，效果也很差；

3. NTK-RoPE 不微调就取得了非平凡（但有所下降）的外推结果；

4. 加入来集中注意力确实有帮助。

所以，NTK-RoPE 成功地成为目前第二种笔者测试有效的不用微调就可以扩展 LLM 的 Context 长度的方案（第一种自然是 NBCE），再次为提出者的卓越洞察力点赞！更加值得高兴的是，NTK-RoPE 在“重复”外推上比“不重复”外推效果明显好，表明这样修改之后是保留了全局依赖，而不是单纯将注意力局部化。