Transformer升级之路:18、RoPE的底数选择原则
By 苏剑林 | 2024-05-29 | 140740位读者 |我们知道,在RoPE中频率的计算公式为$\theta_i = b^{-2i/d}$,底数$b$默认值为10000。目前Long Context的主流做法之一是,先在$b=10000$上用短文本预训练,然后调大$b$并在长文本微调,其出发点是《Transformer升级之路:10、RoPE是一种β进制编码》里介绍的NTK-RoPE,它本身有较好长度外推性,换用更大的$b$再微调相比不加改动的微调,起始损失更小,收敛也更快。该过程给人的感觉是:调大$b$完全是因为“先短后长”的训练策略,如果一直都用长文本训练似乎就没必要调大$b$了?
上周的论文《Base of RoPE Bounds Context Length》试图回答这个问题,它基于一个期望性质研究了$b$的下界,由此指出更大的训练长度本身就应该选择更大的底数,与训练策略无关。整个分析思路颇有启发性,接下来我们一起来品鉴一番。
期望性质 #
RoPE这里就不再详细介绍了,它本质上是一个分块对角矩阵
\begin{equation}\boldsymbol{\mathcal{R}}_n = \scriptsize{\left(\begin{array}{cc:cc:cc:cc}
\cos n\theta_0 & -\sin n\theta_0 & 0 & 0 & \cdots & \cdots & 0 & 0 \\
\sin n\theta_0 & \cos n\theta_0 & 0 & 0 & \cdots & \cdots & 0 & 0 \\
\hdashline
0 & 0 & \cos n\theta_1 & -\sin n\theta_1 & \cdots & \cdots & 0 & 0 \\
0 & 0 & \sin n\theta_1 & \cos n\theta_1 & \cdots & \cdots & 0 & 0 \\
\hdashline
\vdots & \vdots & \vdots & \vdots & \ddots & \ddots & \vdots & \vdots \\
\vdots & \vdots & \vdots & \vdots & \ddots & \ddots & \vdots & \vdots \\
\hdashline
0 & 0 & 0 & 0 & \cdots & \cdots & \cos n\theta_{d/2-1} & -\sin n\theta_{d/2-1} \\
0 & 0 & 0 & 0 & \cdots & \cdots & \sin n\theta_{d/2-1} & \cos n\theta_{d/2-1} \\
\end{array}\right)}\end{equation}
然后利用恒等式
\begin{equation}(\boldsymbol{\mathcal{R}}_m \boldsymbol{q})^{\top}(\boldsymbol{\mathcal{R}}_n \boldsymbol{k}) = \boldsymbol{q}^{\top} \boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n \boldsymbol{k} = \boldsymbol{q}^{\top} \boldsymbol{\mathcal{R}}_{n-m} \boldsymbol{k}\end{equation}
给$\boldsymbol{q},\boldsymbol{k}$注入绝对位置信息,并自动实现了相对位置的效果。其中$\theta_i = b^{-2i/d}$,这里的$b$的取值就是本文要探讨的问题。
除了给模型注入位置信息外,我们期望RoPE能具备两个理想性质,以达到更好的效果:1、远程衰减,即位置相近的Token平均来说获得更多的注意力;2、语义聚合,即语义相似的Token平均来说获得更多的注意力。其中第一点我们早在《Transformer升级之路:2、博采众长的旋转式位置编码》有过相关讨论,RoPE确实有一定的远程衰减性质。
所以接下来我们来分析第二点。
不等关系 #
所谓语义聚合,指的是当$\boldsymbol{k}$与$\boldsymbol{q}$相近时,不管它们的相对距离$n-m$多大,其注意力$\boldsymbol{q}^{\top} \boldsymbol{\mathcal{R}}_{n-m} \boldsymbol{k}$平均来说都应该更大(至少要比随机的两个Token更大)。为了得到一个量化的结论,我们进一步简化问题,假设$\boldsymbol{q}$的每个分量都是独立同分布的,每个分量的均值为$\mu$,方差为$\sigma^2$。
现在我们考虑两种不同的$\boldsymbol{k}$:一种是在$\boldsymbol{q}$的基础上,加上一个零均值的扰动$\boldsymbol{\varepsilon}$,我们记$\tilde{\boldsymbol{k}} = \boldsymbol{q} + \boldsymbol{\varepsilon}$,代表跟$\boldsymbol{q}$语义相近的Token;另一种则是假设$\boldsymbol{k}$跟$\boldsymbol{q}$独立同分布,这代表两个随机的Token。根据第二点理想性质,我们希望有
\begin{equation}\mathbb{E}_{\boldsymbol{q},\boldsymbol{k},\boldsymbol{\varepsilon}}\big[\boldsymbol{q}^{\top} \boldsymbol{\mathcal{R}}_{n-m} \tilde{\boldsymbol{k}} - \boldsymbol{q}^{\top} \boldsymbol{\mathcal{R}}_{n-m} \boldsymbol{k}\big] \geq 0\end{equation}
注意我们刚才反复强调了“平均来说”,意味着我们只是期望一个平均的趋势,而不是每一点都能严格成立,所以我们在上式加了取数学期望$\mathbb{E}_{\boldsymbol{q},\boldsymbol{k},\boldsymbol{\varepsilon}}$。现在根据假设和RoPE的定义,我们可以把上式具体地算出来:
\begin{equation}\begin{aligned}
&\,\mathbb{E}_{\boldsymbol{q},\boldsymbol{k},\boldsymbol{\varepsilon}}\big[\boldsymbol{q}^{\top} \boldsymbol{\mathcal{R}}_{n-m} (\boldsymbol{q} + \boldsymbol{\varepsilon}) - \boldsymbol{q}^{\top} \boldsymbol{\mathcal{R}}_{n-m} \boldsymbol{k}\big] \\[5pt]
=&\, \mathbb{E}_{\boldsymbol{q}}\big[\boldsymbol{q}^{\top} \boldsymbol{\mathcal{R}}_{n-m} \boldsymbol{q}\big] - \mathbb{E}_{\boldsymbol{q},\boldsymbol{k}}\big[\boldsymbol{q}^{\top} \boldsymbol{\mathcal{R}}_{n-m} \boldsymbol{k}\big] \\[5pt]
=&\, \mathbb{E}_{\boldsymbol{q}}\big[\boldsymbol{q}^{\top} \boldsymbol{\mathcal{R}}_{n-m} \boldsymbol{q}\big] - \mathbb{E}_{\boldsymbol{q}}[\boldsymbol{q}]^{\top}\boldsymbol{\mathcal{R}}_{n-m} \mathbb{E}_{\boldsymbol{k}}[\boldsymbol{k}] \\[5pt]
=&\, \mathbb{E}_{\boldsymbol{q}}\big[\boldsymbol{q}^{\top} \boldsymbol{\mathcal{R}}_{n-m} \boldsymbol{q}\big] - \mu^2\boldsymbol{1}^{\top}\boldsymbol{\mathcal{R}}_{n-m} \boldsymbol{1} \\[5pt]
=& \mathbb{E}_{\boldsymbol{q}}\left[\sum_{i=0}^{d/2-1} (q_{2i}^2 + q_{2i+1}^2)\cos (n-m)\theta_i\right] - \sum_{i=0}^{d/2-1} 2\mu^2\cos (n-m)\theta_i \\[5pt]
=& \sum_{i=0}^{d/2-1} 2(\mu^2 + \sigma^2)\cos (n-m)\theta_i - \sum_{i=0}^{d/2-1} 2\mu^2\cos (n-m)\theta_i \\[5pt]
=& \sum_{i=0}^{d/2-1} 2\sigma^2\cos (n-m)\theta_i \\
\end{aligned}\end{equation}
如果训练长度最大为$L$,那么$n-m\leq L-1$,因此第二点理想性质可以用如下不等式近似描述:
\begin{equation}\sum_{i=0}^{d/2-1} \cos m\theta_i \geq 0,\quad m\in\{0,1,2,\cdots,L-1\}\label{neq:base}\end{equation}
其中$L$是最大长度,是训练前就要选定的超参,而$d$是模型的head_size,按照LLAMA的一般设置是$d=128$,这也就意味着,上式的唯一可调参数就是$\theta_i = b^{-2i/d}$中的$b$。在《Transformer升级之路:1、Sinusoidal位置编码追根溯源》中我们就简单探究过这个函数,它整体趋势是衰减的,$b$越大则衰减速度越慢,对应的连续非负区间就越大,所以存在一个最小的$b$使得上述不等式恒成立,即
\begin{equation}b^* = \inf\left\{\,\,b\,\,\,\left|\,\,\,f_b(m)\triangleq\sum_{i=0}^{d/2-1} \cos m b^{-2i/d} \geq 0,\,\, m\in\{0,1,2,\cdots,L-1\}\right.\right\}\end{equation}
数值求解 #
由于$f_b(m)$涉及到多个三角函数的求和,并且$\theta_i$关于$i$还是非线性的,很难想象上述问题会有解析解,因此只能诉诸数值求解了。然而,$f_b(m)$越到后面震荡越频繁且不规律,因此即便数值求解也不是那么简单的事情。
笔者一开始以为,如果$b_0$使得$f_{b_0}(m)\geq 0$恒成立,那么$\forall b \geq b_0$都恒成立$f_b(m)\geq 0$,所以用二分法就可以了。但事实上这个假设并不成立,所以二分法宣告破产。继续想了一段时间,依然没什么优化思路,期间向原论文作者请教过,他们采用的是逆函数法,即给定$b$求使得$f_b(m)\geq 0$恒成立的最大$L$是比较简单的,于是我们可以得到很多$(b, L)$对,理论上只要枚举的$b$足够多,那么对于任意$L$都可以找出最小的$b$。然而这里有个精度问题,原论文最大的$L$计算到了$10^6$,$b$至少要枚举到$10^8$,如果枚举间隔小,那么计算成本非常大,如果枚举间隔大,那么可能漏掉很多解。
最后,笔者决定还是用“Jax + GPU”进行暴力搜索,以求得到更高精度的结果,大致流程是:
1、初始化$b=1000L$(在$10^6$内$b=1000L$可以使得$f_b(m)\geq 0$恒成立);
2、遍历$k=1,2,3,4,5$,执行以下操作:
2.1)将$[0,b]$等分为$10^k$份,遍历等分点,判断$f_b(m)\geq 0$是否恒成立;
2.2)取最小的使得$f_b(m)\geq 0$恒成立的等分点,更新$b$;
3、返回最终的$b$。
最终结果普遍要比原论文的更紧一些
$$\scriptsize\begin{array}{c|cccccccccc}
\hline
L & 1k & 2k & 4k & 8k & 16k & 32k & 64k & 128k & 256k & 512k & 1M \\
\hline
b^*(\text{原文}) & 4.3e3 & 1.6e4 & 2.7e4 & 8.4e4 & 3.1e5 & 6.4e5 & 2.1e6 & 7.8e6 & 3.6e7 & 6.4e7 & 5.1e8 \\
b^*(\text{本文}) & 4.3e3 & \color{red}{1.2e4} & 2.7e4 & 8.4e4 & \color{red}{2.3e5} & \color{red}{6.3e5} & 2.1e6 & \color{red}{4.9e6} & \color{red}{2.4e7} & \color{red}{5.8e7} & \color{red}{6.5e7} \\
\hline
\end{array}$$
参考代码:
from functools import partial
import numpy as np
import jax.numpy as jnp
import jax
@partial(jax.jit, static_argnums=(2,))
def f(m, b, d=128):
i = jnp.arange(d / 2)
return jnp.cos(m[:, None] * b ** (-2 * i[None] / d)).sum(axis=1)
@np.vectorize
def fmin(L, b):
return f(np.arange(L), b).min()
def bmin(L):
B = 1000 * L
for k in range(1, 6):
bs = np.linspace(0, 1, 10**k + 1)[1:] * B
ys = fmin(L, bs)
for b, y in zip(bs, ys):
if y >= 0:
B = b
break
return B
bmin(1024 * 128)
渐近估计 #
除了数值求解外,我们也可以通过渐近分析来得到一个解析的估计结果,这个估计比数值结果要小,本质上是$d\to\infty$的解,但同样能够得出“$b$应该随着$L$增大而增大”的结论。
渐近估计的思路,是用积分代替求和:
\begin{equation}f_b(m) = \sum_{i=0}^{d/2-1} \cos m b^{-2i/d}\approx \int_0^1 \cos m b^{-s} ds \xlongequal{\text{令}t = mb^{-s}} \int_{mb^{-1}}^m \frac{\cos t}{t \ln b}dt\end{equation}
其中我们记
\begin{equation}\text{Ci}(x) = -\int_x^{\infty} \frac{\cos t}{t} dt\end{equation}
这是被前人研究过的三角积分(参考 Trigonometric integral ),利用这个记号,我们可以写出
\begin{equation}f_b(m) \approx \frac{\text{Ci}(m) - \text{Ci}(mb^{-1})}{\ln b}\end{equation}
$\text{Ci}(x)$的图像长这样:
它的第一个零点是$x_0=0.6165\cdots$,对于$m\geq 1$,可以看出$|\text{Ci}(m)|\leq 1/2$,所以其实$\text{Ci}(m)$相对来说是小项,对于渐近估计来说可以忽略,那么问题近似地变成了$\text{Ci}(mb^{-1})\leq 0$对于$m=1,2,\cdots,L$恒成立,我们只需要让相应的$mb^{-1}$都落在$[0,x_0]$区间内就可以实现,这意味着$Lb^{-1}\leq x_0$,即
\begin{equation}b \geq L / x_0 \approx 2L\end{equation}
或者简单点$b^* = \mathcal{O}(L)$。不出意料这个结果比精确的数值结果要小,因为它对应于$d\to\infty$,无限个三角函数叠加会使得函数图像的震荡更少,看起来更加平稳(相比于有限的$d$),从而对于固定的$b$,$f_b(m)$的连续非负区间更长,或者反过来,对于固定的$L$,保持$m=0,1,2,\cdots,L-1$的$f_b(m)$都非负的$b$更小。
相关思考 #
在《Transformer升级之路:10、RoPE是一种β进制编码》中,我们将RoPE类比为一种$\beta$进制表示,其中$\beta = b^{2/d}$,那么$b - 1= \beta^{d/2} - 1$正好是$d/2$位$\beta$进制编码能够表示的最大数字,于是要表示$0,1,2,\cdots,L-1$这$L$个位置编码,至少有$b \geq L$,这个朴素的类比再次给出了“$b$应该随着$L$增大而增大”的结论,其结果跟上一节的渐近分析结果更为接近。
另一方面,Meta最新发布的LLAMA3,训练长度为8192,但RoPE的底数选择了惊人的500000(5e5),这比前面的数值结果(8.4e4)还要大将近一个数量级,不管从哪个角度看,这个数值笔者都认为是偏大的,可能LLAMA3的这个底数本就是给更大文本长度预留的。但不论如何,更大的文本长度选择更大的RoPE底数,似乎已经成为了很多训练人员的共识。
其实不管是数值结果还是渐近估计,都只是一个参考值,实际上对于给定的$L$,一个相当大范围内的$b$都应该会有相近的效果。所以具体的数值都不重要,关键是原论文通过语义聚合的出发点和一系列推导,澄清了“$b$应该随着$L$增大而增大”的结论及其原理,这是笔者所认为的原论文的核心贡献。
此外,其实语义聚合的出发点和结论也可以用来解释Position Interpolation(PI)。刚才我们说了,同一个$b$,$f_b(m)$的连续非负区间是固定的,如果要使$0,1,2,\cdots,L-1$都落在非负区间内,就需要随着$L$的增大而相应的增加$b$。但反过来,我们也可以不增加$b$,而是减少相邻位置的间隔(即位置ID改成$0,1/k,2/k,\cdots$),那么就可以在同样大小的非负区间内表示$k$倍的位置了,这便是语义聚合视角下的Position Interpolation。
部分旋转 #
RoPE提出于2021年,当时只有一篇中文博客,后来得到了EleutherAI组织的认可和实验,继而才逐渐向学术界推广。当时EleutherAI实验发现,如果只对部分维度加RoPE,会取得稍优的结果,相关内容可以参考这里、这里和这里,后来这个操作用到了它们的GPT-NeoX中。
当然,部分旋转还不是当前LLM的主流选择,但这不妨碍我们研究它,也许它未成为主流选择只是因为我们对它还不够了解。那为什么部分旋转反而可能会更优呢?笔者发现可以用本文的结论来一定程度上解释它。以只旋转一半维度为例,它在数学上等价于选择如下的$\theta_i$:
\begin{equation}\theta_i = \left\{\begin{aligned}&b^{-4i/d},& i < d/4 \\
&0,&i \geq d/4\end{aligned}\right.\end{equation}
此时我们有
\begin{equation}\sum_{i=0}^{d/2-1} \cos m\theta_i = \sum_{i=0}^{d/4-1} (1+\cos mb^{-4i/d})\geq 0\end{equation}
也就是不论$m,b$如何,我们所期望的不等式$\eqref{neq:base}$都自动成立,这意味着从本文的观点来看,部分旋转在赋予位置信息的同时有更好的语义聚合能力,这对模型的效果也许更加有利。同时,部分旋转对模型的长文本能力或许也更有利,因为不等式恒成立,所以按照本文的观点,不论长短文本训练都不用修改$b$。
值得一提的是,DeepSeek提出的MLA也应用了部分旋转,虽然在MLA的原始推导中,部分旋转更多是为了整合RoPE的无奈之举,但结合以往的部分旋转实验结果来看,也许MLA的优异效果有部分旋转的一分功劳。
文章小结 #
本文简单介绍了论文《Base of RoPE Bounds Context Length》,它从语义聚合的期望性质讨论了RoPE的底数下界,由此指出更大的训练长度应该选择更大的底数,而不单单是为了配合“先短后长”的训练策略、继而利用NTK-RoPE来降低初始损失的折中选择。
转载到请包括本文地址:https://www.spaces.ac.cn/archives/10122
更详细的转载事宜请参考:《科学空间FAQ》
如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。
如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!
如果您需要引用本文,请参考:
苏剑林. (May. 29, 2024). 《Transformer升级之路:18、RoPE的底数选择原则 》[Blog post]. Retrieved from https://www.spaces.ac.cn/archives/10122
@online{kexuefm-10122,
title={Transformer升级之路:18、RoPE的底数选择原则},
author={苏剑林},
year={2024},
month={May},
url={\url{https://www.spaces.ac.cn/archives/10122}},
}
May 30th, 2024
有一个困扰挺久的问题想请教一下苏神,每个位置对应的位置编码都是两两成对的形式(从Sinusoidal到RoPE),而且都建立在emb维度是偶数的基础上,偶数维emb与两两成对的PE组合有什么深意吗
非要扯故事的话,就是数域直到复数域才彻底完备(对于代数运算),而复数分实部和虚部,复数向量展开为实数向量的话,自然就是有偶数个分量了。
May 30th, 2024
请教苏神一个问题,想听听你的看法。之前的做法普遍是先在小的b上大规模训练,再用大的b做少量训练(也可以解释为插值后训练)。 但是不知道如果一开始的大规模训练就用特别大的b是否对于长度外推有效?(个人感觉应该是无效的)
按照 https://kexue.fm/archives/9948 的转圈解释,那么应该就是无助于长度外推。
May 31st, 2024
在之前那个ood转圈的视角下,为了使长度内的所有点都遍历整个圆,需要b越小越好;而在这篇文章的视角下需要b越大越好,这两者是不是出现了冲突?
以前的OOD需要b越小越好,这篇文章说的是b存在一个下界,不能无限小。否则的话,b 越小越好,那 b=0.1 好于 b=1000。 另外这篇文章说bound了上下文长度,但是也没说b越大越好,大于下界即可吧
这里有几个细节:
1、在转圈视角下,我们是得出(理论上)$b$越小越利于长度外推,但长度外推好未必其他一切都好;
2、实际测试显示,当$b$小到一定程度,其实连长度外推都不好,所以Attention堆叠起来后,其工作机制是非常复杂的,很多事后解释都很片面;
3、另一个极端是$b\to\infty$,这相当于NoPE,有些文章说它长度外推性很好,我测试起来长度外推也没有,另外就算它真的长度外推好,也只是长度外推好,实际上的效果不如加了位置编码的;
4、而本文说的是,$b$随着长度$L$增大,可能在某些方面(满足两个期望性质)对效果更加好,不同的假设,不同的评估指标,自然有不同的结论。
September 4th, 2024
苏神您好,YaRN的论文中使用了rope_base=10000(静态YaRN),产生了很好的外推效果,想问一下苏神是否尝试过在使用YaRN的同时,将rope_base调为500000,这样是否会有1+1>2的效果,也就是既好于YaRN(rope_base=10000),也好于NTK-aware(rope_base=500000)?还是说结合了两者会有负面效果呀?
两者结合会有负面效果。YaRN是分析了NTK起作用的原理,然后针对提出更好的策略,所以直接使用的话,YaRN基本都优于NTK。YaRN本身就是改进NTK,不存在两者结合能更优的情况。
感谢苏神回复,我当下的认知是,NTK(rope_base=500000)和YaRN(rope_base=10000)效果是好于RoPE(rope_base=10000)的,而YaRN是在rope_base=10000的基础上进行的改进,从而产生了\alpha、\beta、以及t等参数;但是为什么rope_base=10000就一定是YaRN最好的参数呢,如果适当的调整YaRN的其他超参数(\alpha、\beta),会不会在rope_base=500000上做到比NTK更好的效果呀
你可能没搞清楚逻辑。
从一个base(比如base=10000)出发,在相对较短的语料上,已经训练好了一个模型。现在想要将这个模型微调成长文本的模型,在微调之前我们需要先对RoPE做一些修改,提高它的长度外推效果,从而以达到更高的微调效率。这个修改有两种做法:NTK和YaRN,它们是有一定重叠的,YaRN分析了NTK为什么有效,然后提出了更好的解决方法,所以YaRN通常更好。
NTK和YaRN的优劣,说的是它们的免训练长度外推效果,并不是说微调后甚至是从零训练的效果优劣。经过训练的效果已经不能简单判断了。至于从零训练为什么也需要调大base,这就是本文的主题,已经不属于NTK、YaRN的范畴了。