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

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

进制类比 #

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

具体来说,整数$n$的$\beta$进制表示的(从右往左数)第$m$位数字是:
\begin{equation}\left\lfloor\frac{n}{\beta^{m-1}}\right\rfloor\bmod\beta\label{eq:mod}\end{equation}
而它的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$位,也就是说它相当于$n$的$d/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://www.kexue.fm/archives/9706

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

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

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

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

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

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