在文章《Transformer升级之路:10、RoPE是一种β进制编码》中,我们给出了RoPE的β进制诠释,并基于进制转化的思路推导了能够在不微调的情况下就可以扩展Context长度的NTK-aware Scaled RoPE。不得不说,通过类比β进制的方式来理解位置编码,确实是一个非常美妙且富有启发性的视角,以至于笔者每次深入思考和回味之时,似乎总能从中得到新的领悟和收获。

本文将重新回顾RoPE的β进制诠释,并尝试将已有的NTK-aware Scaled RoPE一般化,以期望找到一种更优的策略来不微调地扩展LLM的Context长度。

进制类比 #

我们知道,RoPE的参数化沿用了Sinusoidal位置编码的形式。而不知道是巧合还是故意为之,整数n的Sinusoidal位置编码,与它的β进制编码,有很多相通之处。

具体来说,整数nβ进制表示的(从右往左数)第m位数字是:
nβm1mod
而它的Sinusoidal位置编码是
\begin{equation}\boldsymbol{p}_n=\big[\cos\theta_1,\sin\theta_1,\cos\theta_2,\sin\theta_2,\cdots,\cos\theta_{d/2},\sin\theta_{d/2}\big]\\[5pt] \theta_m = \frac{n}{\beta^{m-1}},\quad \beta=10000^{2/d} \label{eq:sinu}\end{equation}
可以看到,两者都有相同的\frac{n}{\beta^{m-1}},并且\bmod\cos,\sin同为周期函数,所以两者的唯一差距,只是无关紧要的取整\lfloor\cdot\rfloor了。所以说,将RoPE/Sinusoidal位置编码类比为它\beta进制表示,是非常直观且合理的结果。

修正NTK #

沿着《Transformer升级之路:10、RoPE是一种β进制编码》的思路,直接外推会将外推压力集中在“高位(m较大)”上,而位置内插则会将“低位(m较小)”的表示变得更加稠密,不利于区分相对距离。而NTK-aware Scaled RoPE其实就是进制转换,它将外推压力平摊到每一位上,并且保持相邻间隔不变,这些特性对明显更倾向于依赖相对位置的LLM来说是非常友好和关键的,所以它可以不微调也能实现一定的效果。

仔细看式\eqref{eq:sinu}\cos,\sin事实上是一个整体,所以它实际只有d/2位,也就是说它相当于nd/2\beta进制编码。如果我们要扩展到k倍Context,将\beta进制转换为\beta\lambda进制,那么至少应该有
\begin{equation}\lambda^{d/2}=k\quad\Rightarrow\quad\lambda = k^{2/d}\end{equation}
于是新的RoPE变为
\begin{equation}\boldsymbol{p}_n=\big[\cos\theta_1,\sin\theta_1,\cos\theta_2,\sin\theta_2,\cdots,\cos\theta_{d/2},\sin\theta_{d/2}\big]\\[5pt] \theta_m = \frac{n}{(\beta\lambda)^{m-1}},\quad \beta=10000^{2/d},\quad \lambda = k^{2/d}\label{eq:ntk-old}\end{equation}
这就是上一篇文章我们提出的NTK-RoPE。

然而,后来笔者仔细思考后,发现这其实还不够合理。回到式\eqref{eq:mod},如果要计算\beta\lambda进制的第m位数字,那么应该是
\begin{equation}\left\lfloor\frac{n}{(\beta\lambda)^{m-1}}\right\rfloor\bmod(\beta\lambda)\end{equation}
也就是说,除了\frac{n}{\beta^{m-1}}要换成\frac{n}{(\beta\lambda)^{m-1}}之外,求\bmod的周期也要扩大\lambda倍,这等价于求\cos,\sin之前,要多除以一个\lambda
\begin{equation}\boldsymbol{p}_n=\big[\cos\theta_1,\sin\theta_1,\cos\theta_2,\sin\theta_2,\cdots,\cos\theta_{d/2},\sin\theta_{d/2}\big]\\[5pt] \theta_m = \frac{n}{\lambda(\beta\lambda)^{m-1}},\quad \beta=10000^{2/d},\quad \lambda = k^{2/d}\label{eq:ntk-fixed}\end{equation}
在后面的实验中,我们把上一篇文章提出的式\eqref{eq:ntk-old}称为“NTK-RoPE-old”,而式\eqref{eq:ntk-fixed}称为“NTK-RoPE-fixed”。

混合进制 #

现在,不妨让我们更加“天马行空”一些——既然我们可以用\beta进制来表示位置,那么为何不干脆使用更一般化的“混合进制”呢?这里的混合进制,指的是每一位数字所使用的进位基数不尽相同,这对于我们来说并不鲜见,比如60秒是1分钟、60分是1小时,但24小时是1天、7天是1周,这里的60、60、24、7就是不同进制基数,换句话说秒、分、时、天、周就是一个使用混合进制的例子。

假设从右往左数,第1位使用\beta_1进制、第2位使用\beta_2进制、第3位使用\beta_3进制、...,那么求n的第m位数字,结果是
\begin{equation}\left\lfloor\frac{n}{\beta_1\beta_2\cdots\beta_{m-1}}\right\rfloor\bmod\beta_m\label{eq:mod2}\end{equation}
为什么会考虑到混合进制呢?这是因为某天笔者发现了一个有趣的事实:RoPE本质上是一种相对位置编码,相对位置是Toeplitz矩阵的一个特例,它长这个样(由于本文主要关心语言模型,所以右上角部分就没写出来了)
\begin{equation}\begin{pmatrix}0 & \\ 1 & 0 & \\ 2 & 1 & 0 &\\ 3 & 2 & 1 & 0 & \\ 4 & 3 & 2 & 1 & 0 & \\ 5 & 4 & 3 & 2 & 1 & 0 & \\ 6 & 5 & 4 & 3 & 2 & 1 & 0 & \\ \end{pmatrix}\end{equation}
从上式我们可以发现,相对位置编码的位置分布是不均衡的!0的出现次数最多、1次之、2再次之,以此类推,即n越大出现次数越少。这就意味着,作为一种\beta进制编码的RoPE,它的“高位”很可能是训练不充分的,换言之高位的泛化能力很可能不如低位。刚才我们说了,NTK-RoPE将外推压力平摊到每一位上,如果这里的猜测合理的话,那么“平摊”就不是最优的,应该是低位要分摊更多,高位分摊更少,这就导致了混合进制。

分摊优化 #

具体来说,我们通过将\beta进制转换为\beta_1,\beta_2,\cdots,\beta_{d/2}混合进制的方式来扩展到k倍Context,这里\beta_m = \beta \lambda_m。此时式\eqref{eq:mod2}变为
\begin{equation}\left\lfloor\frac{n}{\beta^{m-1}(\lambda_1\lambda_2\cdots\lambda_{m-1})}\right\rfloor\bmod(\beta\lambda_m)\end{equation}
\eqref{eq:ntk-fixed}也相应地变成
\begin{equation}\boldsymbol{p}_n=\big[\cos\theta_1,\sin\theta_1,\cos\theta_2,\sin\theta_2,\cdots,\cos\theta_{d/2},\sin\theta_{d/2}\big]\\[5pt] \theta_m = \frac{n}{\beta^{m-1}(\lambda_1\lambda_2\cdots\lambda_m)},\quad \beta=10000^{2/d}\end{equation}
根据“扩展k倍”和“低位要分摊更多”的原则,约束条件是
\begin{equation}\lambda_1\lambda_2\cdots\lambda_{d/2} = k,\quad \lambda_1 \geq \lambda_2 \geq \cdots \geq \lambda_{d/2} \geq 1\end{equation}
我们讨论如下形式的解(有兴趣的读者也可以试探别的形式的解,这里自由度本身就很大)
\begin{equation}\lambda_1\lambda_2\cdots\lambda_m = \exp(am^b)\end{equation}
a > 0, b\leq 1时,它满足\lambda_1 \geq \lambda_2 \geq \cdots \geq \lambda_{d/2} \geq 1的条件,当b=1时,实际上就是前面的“NTK-RoPE-fixed”,当b=0时,就是“Positional Interpolation(PI)”。\lambda_1\lambda_2\cdots\lambda_{d/2} = k给出了约束
\begin{equation}a\left(\frac{d}{2}\right)^b = \log k\end{equation}
所以只有一个自由度可以调。经过简单的二分法搜索,笔者发现在自己的实验中,b=0.625能取得平均来说比较好的扩展效果(不同的模型可能会有不同的最优解,请自行调试),这个版本被称为“NTK-RoPE-mixed”。

实验结果 #

《Transformer升级之路:10、RoPE是一种β进制编码》的实验基础上,笔者补做了“NTK-RoPE-fixed”和“NTK-RoPE-mixed”的实验,对比如下:
\begin{array}{c|cc} \hline \text{测试长度} & 512(\text{训练}) & 4096(\text{重复}) & 4096(\text{不重复})\\ \hline \text{Baseline} & 49.41\% & 24.17\% & 23.16\% \\ \text{Baseline-}\log n & 49.40\% & 24.60\% & 24.02\% \\ \hline \text{PI-RoPE} & 49.41\% & 15.04\% & 13.54\% \\ \text{PI-RoPE-}\log n & 49.40\% & 14.99\% & 16.51\% \\ \hline \text{NTK-RoPE-old} & 49.41\% & 51.28\% & 39.27\% \\ \text{NTK-RoPE-}\log n\text{-old} & 49.40\% & 61.71\% & 43.75\% \\ \hline \text{NTK-RoPE-fixed} & 49.41\% & 51.86\% & 39.61\% \\ \text{NTK-RoPE-}\log n\text{-fixed} & 49.40\% & 62.85\% & 44.14\% \\ \text{NTK-RoPE-mixed} & 49.41\% & 53.09\% & 40.12\% \\ \text{NTK-RoPE-}\log n\text{-mixed} & 49.40\% & \boldsymbol{68.91\%} & \boldsymbol{45.41\%} \\ \hline \end{array}

可以看到,相比等进制的“NTK-RoPE-old”和“NTK-RoPE-fixed”,混合进制推导出来的“NTK-RoPE-mixed”所带来的提升还是很明显的,而且不用微调,可谓是“免费午餐”了。此外,可以看到\log n版的外扩性能确实更好,但是\log n技巧需要在预训练阶段就加入,之前就有读者问过像LLAMA这种在预训练阶段并没有加入\log n技巧的模型,可否享受到\log n的“红利”呢?经过笔者测试,发现它可以通过加入如下scale因子来提升效果:
\begin{equation}\max(1, \log_{\text{maxlen}} n)\label{eq:plogn}\end{equation}
这里的\text{maxlen}是预训练的最大长度,在本文的实验中是512,在LLAMA中是2048,LLAMA2则是4096,实现时可以直接给每个\boldsymbol{q}_n乘上相应的因子。这样一来,在\text{maxlen}之内的部分不受影响,之外的部分则按\log n缩放,算是一种简单的过渡,效果如下(加个\color{red}{\dagger}区别原来的\log n):
\begin{array}{c|cc} \hline \text{测试长度} & 512(\text{训练}) & 4096(\text{重复}) & 4096(\text{不重复})\\ \hline \text{NTK-RoPE-fixed} & 49.41\% & 51.86\% & 39.61\% \\ \text{NTK-RoPE-}\log n^{\color{red}{\dagger}}\text{-fixed} & 49.41\% & 55.94\% & 41.11\% \\ \text{NTK-RoPE-mixed} & 49.41\% & 53.09\% & 40.12\% \\ \text{NTK-RoPE-}\log n^{\color{red}{\dagger}}\text{-mixed} & 49.41\% & 59.11\% & 42.38\% \\ \hline \end{array}
可以看到,这个\log n^{\color{red}{\dagger}}也算得上免费的午餐了。总之,如果你打算进行从零预训练,不妨事先就加入\log n技巧,如果已经训练完成,那么可以使用式\eqref{eq:plogn}替代,最后再加上NTK-RoPE-mixed,能够取得较优的拓展Context效果。

文章小结 #

在这篇文章中,我们重温了RoPE的\beta进制视角,并尝试对NTK-aware Scaled RoPE进行推广,在混合进制的启发下,我们得到了一个更优的不微调扩展Context长度的策略,最后通过实验表明了它的有效性。

转载到请包括本文地址:https://kexue.fm/archives/9706

更详细的转载事宜请参考:《科学空间FAQ》

如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。

如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!

如果您需要引用本文,请参考:

苏剑林. (Jul. 31, 2023). 《Transformer升级之路:11、将β进制位置进行到底 》[Blog post]. Retrieved from https://kexue.fm/archives/9706

@online{kexuefm-9706,
        title={Transformer升级之路:11、将β进制位置进行到底},
        author={苏剑林},
        year={2023},
        month={Jul},
        url={\url{https://kexue.fm/archives/9706}},
}