Transformer升级之路：RoPE的底数设计原则

2024-06-16 06:06

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 科学空间

研究方向 | NLP、神经网络

我们知道，在 RoPE 中频率的计算公式为，底数默认值为 10000。目前 Long Context 的主流做法之一是，先在上用短文本预训练，然后调大并在长文本微调，其出发点是《Transformer升级之路：RoPE是一种β进制编码》里介绍的 NTK-RoPE，它本身有较好长度外推性，换用更大的再微调相比不加改动的微调，起始损失更小，收敛也更快。

该过程给人的感觉是：调大完全是因为“先短后长”的训练策略，如果一直都用长文本训练似乎就没必要调大了？

近期的论文《Base of RoPE Bounds Context Length》[1] 试图回答这个问题，它基于一个期望性质研究了的下界，由此指出更大的训练长度本身就应该选择更大的底数，与训练策略无关。整个分析思路颇有启发性，接下来我们一起来品鉴一番。

期望性质

RoPE 这里就不再详细介绍了，它本质上是一个分块对角矩阵

然后利用恒等式

给注入绝对位置信息，并自动实现了相对位置的效果。其中，这里的的取值就是本文要探讨的问题。

除了给模型注入位置信息外，我们期望 RoPE 能具备两个理想性质，以达到更好的效果：1）远程衰减，即位置相近的 Token 平均来说获得更多的注意力；2）语义聚合，即语义相似的 Token 平均来说获得更多的注意力。其中第一点我们早在《Transformer升级之路：博采众长的旋转式位置编码》有过相关讨论，RoPE 确实有一定的远程衰减性质。

所以接下来我们来分析第二点。

不等关系

所谓语义聚合，指的是当相近时，不管它们的相对距离多大，其注意力平均来说都应该更大（至少要比随机的两个 Token 更大）。为了得到一个量化的结论，我们进一步简化问题，假设的每个分量都是独立同分布的，每个分量的均值为，方差为。

现在我们考虑两种不同的：一种是在的基础上，加上一个零均值的扰动，我们记，代表跟语义相近的 Token；另一种则是假设跟独立同分布，这代表两个随机的 Token。根据第二点理想性质，我们希望有

注意我们刚才反复强调了“平均来说”，意味着我们只是期望一个平均的趋势，而不是每一点都能严格成立，所以我们在上式加了取数学期望。现在根据假设和 RoPE 的定义，我们可以把上式具体地算出来：

如果训练长度最大为 L，那么，因此第二点理想性质可以用如下不等式近似描述：

其中 L 是最大长度，是训练前就要选定的超参，而是模型的 head_size，按照 LLAMA 的一般设置是，这也就意味着，上式的唯一可调参数就是中的。在《Transformer升级之路：Sinusoidal位置编码追根溯源》中我们就简单探究过这个函数，它整体趋势是衰减的，越大则衰减速度越慢，对应的连续非负区间就越大，所以存在一个最小的 b 使得上述不等式恒成立，即

数值求解

由于涉及到多个三角函数的求和，并且关于还是非线性的，很难想象上述问题会有解析解，因此只能诉诸数值求解了。然而，越到后面震荡越频繁且不规律，因此即便数值求解也不是那么简单的事情。

笔者一开始以为，如果使得恒成立，那么都恒成立，所以用二分法就可以了。但事实上这个假设并不成立，所以二分法宣告破产。

继续想了一段时间，依然没什么优化思路，期间向原论文作者请教过，他们采用的是逆函数法，即给定求使得恒成立的最大 L 是比较简单的，于是我们可以得到很多对，理论上只要枚举的足够多，那么对于任意都可以找出最小的。然而这里有个精度问题，原论文最大的计算到了，至少要枚举到，如果枚举间隔小，那么计算成本非常大，如果枚举间隔大，那么可能漏掉很多解。

最后，笔者决定还是用 “Jax + GPU” 进行暴力搜索，以求得到更高精度的结果，大致流程是：

1. 初始化（在内可以使得恒成立）；

2. 遍历，执行以下操作：

2.1 将等分为份，遍历等分点，判断是否恒成立；

2.2 取最小的使得恒成立的等分点，更新；

3. 返回最终的。

最终结果普遍要比原论文的更紧一些。

参考代码：

 1from functools import partial
 2import numpy as np
 3import jax.numpy as jnp
 4import jax
 5
 6@partial(jax.jit, static_argnums=(2,))
 7def f(m, b, d=128):
 8    i = jnp.arange(d / 2)
 9    return jnp.cos(m[:, None] * b ** (-2 * i[None] / d)).sum(axis=1)
10
11@np.vectorize
12def fmin(L, b):
13    return f(np.arange(L), b).min()
14
15def bmin(L):
16    B = 1000 * L
17    for k in range(1, 6):
18        bs = np.linspace(0, 1, 10**k + 1)[1:] * B  
19        ys = fmin(L, bs)
20        for b, y in zip(bs, ys):
21            if y >= 0:
22                B = b
23                break
24    return B
25
26bmin(1024 * 128)

渐近估计

除了数值求解外，我们也可以通过渐近分析来得到一个解析的估计结果，这个估计比数值结果要小，本质上是的解，但同样能够得出“ 应该随着增大而增大”的结论。

渐近估计的思路，是用积分代替求和：

其中我们记

这是被前人研究过的三角积分（参考 Trigonometric integral [2]），利用这个记号，我们可以写出

的图像长这样：

▲ Ci(x) 的图像【来自维基百科】

它的第一个零点是，对于，可以看出，所以其实相对来说是小项，对于渐近估计来说可以忽略，那么问题近似地变成了对于恒成立，我们只需要让相应的都落在区间内就可以实现，这意味着，即

或者简单点。不出意料这个结果比精确的数值结果要小，因为它对应于，无限个三角函数叠加会使得函数图像的震荡更少，看起来更加平稳（相比于有限的），从而对于固定的的连续非负区间更长，或者反过来，对于固定的，保持的都非负的更小。

相关思考

在《Transformer升级之路：RoPE是一种β进制编码》中，我们将 RoPE 类比为一种进制表示，其中，那么正好是位进制编码能够表示的最大数字，于是要表示这个位置编码，至少有，这个朴素的类比再次给出了“ 应该随着增大而增大”的结论，其结果跟上一节的渐近分析结果更为接近。

另一方面，Meta 最新发布的 LLAMA3，训练长度为 8192，但 RoPE 的底数选择了惊人的 500000（5e5），这比前面的数值结果（8.4e4）还要大将近一个数量级，不管从哪个角度看，这个数值笔者都认为是偏大的，可能 LLAMA3 的这个底数本就是给更大文本长度预留的。但不论如何，更大的文本长度选择更大的 RoPE 底数，似乎已经成为了很多训练人员的共识。

其实不管是数值结果还是渐近估计，都只是一个参考值，实际上对于给定的L，一个相当大范围内的 b 都应该会有相近的效果。所以具体的数值都不重要，关键是原论文通过语义聚合的出发点和一系列推导，澄清了“ 应该随着增大而增大”的结论及其原理，这是笔者所认为的原论文的核心贡献。

此外，其实语义聚合的出发点和结论也可以用来解释 Position Interpolation [3]（PI）。刚才我们说了，同一个的连续非负区间是固定的，如果要使都落在非负区间内，就需要随着的增大而相应的增加。但反过来，我们也可以不增加 b，而是减少相邻位置的间隔（即位置ID改成），那么就可以在同样大小的非负区间内表示 k 倍的位置了，这便是语义聚合视角下的 Position Interpolation。