自从在《Transformer升级之路:11、将β进制位置进行到底》中引入混合进制的思路进一步推广了NTK-aware Scaled RoPE后,笔者感觉类似思路的效果已经达到了上限,想要更大幅度的提升就必须另辟蹊径了。这时候笔者想起了此前构思过的一个思路,该思路由于复杂度较高所以被搁置下了,既然现在已经遇到了瓶颈,那么“唯一的办法就是最好的办法”,于是便将它重拾起来。

万万没想到的是,尽管该方法增加了一些推理复杂度,但它的实验效果却惊人地好——甚至隐约有无限的长度外推能力!因此,笔者迫不及待地撰写了本文来分享该方法。由于形式上跟ReLU激活函数的相似性,所以笔者将该方法命名为“ReRoPE (Rectified Rotary Position Embeddings)”。

重温 #

我们知道,RoPE形式上是一种绝对位置编码,但实际上给Attention带来的是相对位置信息,即如下的Toeplitz矩阵
\begin{equation}\begin{pmatrix}0 & \\
1 & 0 & \\
2 & 1 & 0 &\\
3 & 2 & 1 & 0 & \\
\ddots & 3 & 2 & 1 & 0 & \\
\ddots & \ddots & 3 & 2 & 1 & 0 & \\
\ddots & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots \\
\small{L - 2} & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots \\
\small{L - 1} & \small{L - 2} & \ddots & \ddots & \ddots & 3 & 2 & 1 & 0 & \\
\end{pmatrix}\label{eq:rope}\end{equation}
这里的$L$是当前样本长度。当$L$明显超出了训练长度时,多出来的位置由于没有被训练过,所以无法保证效果,这就是直接外推(Length Extrapolation)表现通常比较差的原因。

后来,研究人员提出了位置内插(Position Interpolation),它相当于将相对位置矩阵改为:
\begin{equation}\begin{pmatrix}0 & \\
\frac{1}{k} & 0 & \\
\frac{2}{k} & \frac{1}{k} & 0 &\\
\frac{3}{k} & \frac{2}{k} & \frac{1}{k} & 0 & \\
\ddots & \frac{3}{k} & \frac{2}{k} & \frac{1}{k} & 0 & \\
\ddots & \ddots & \frac{3}{k} & \frac{2}{k} & \frac{1}{k} & 0 & \\
\ddots & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots \\
\small{\frac{L-2}{k}} & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots \\
\small{\frac{L-1}{k}} & \small{\frac{L-1}{k}} & \ddots & \ddots & \ddots & \frac{3}{k} & \frac{2}{k} & \frac{1}{k} & 0 & \\
\end{pmatrix}\end{equation}
这样一来,只要调整$k$,就可以保证最大的相对位置也不超过训练长度,因此避免了外推。然而,它使得位置信息更加“拥挤”了,所以还需要进行一定步数的微调才能让模型重新工作。而也正因为避免了外推,所以它所需要的微调步数相比直接外推要少得多(神经网络往往更擅长内插而不是外推)。

至于后面提出的NTK-aware Scaled RoPE,则是“剑走偏锋”,巧妙地将外推压力平摊到每一个维度上,所以它不微调也能有不错的效果,但它终究还是依赖外推,这是神经网络不擅长的事情,所以效果存在上限,在笔者的实验中,它的Long Context表现还无法很接近训练效果。

融合 #

我们也可以从语言模型的局域性来考察这些方法。所谓局域性,是指语言模型在推断下一个token时,明显更依赖于邻近的token。直接外推保持了局域性(0附近位置编码不变),效果差是因为引入了超出训练长度的位置编码;位置内插虽然没有外推位置编码,但扰乱了局域性(0附近位置编码被压缩为$1/k$),所以不微调效果也不好;而NTK-aware Scaled RoPE通过“高频外推、低频内插”隐含了两者优点,保证了局域性,又没有明显外推位置编码,所以不微调也有不错的效果。

有没有能更直接地结合外推和内插的方法呢?有,我们可以设定一个窗口大小$w$,在窗口内我们使用大小为$1$的位置间隔,在窗口外我们使用大小为$1/k$的位置间隔,整个相对位置矩阵如下:
\begin{equation}\begin{pmatrix}
\color{red}{0} & \\
\color{red}{1} & \color{red}{0} & \\
\color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{\small{w + \frac{1}{k}}} & \color{green}{w} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{\small{w + \frac{2}{k}}} & \color{green}{\small{w + \frac{1}{k}}} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{\ddots} & \color{green}{\small{w + \frac{2}{k}}} & \color{green}{\ddots} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \\
\color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\small{w + \frac{2}{k}}} & \color{green}{\small{w + \frac{1}{k}}} & \color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{\small{w + \frac{L-1-w}{k}}} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\small{w + \frac{2}{k}}} & \color{green}{\small{w + \frac{1}{k}}} & \color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\end{pmatrix}\label{eq:leaky-rerope}\end{equation}
只要$w$小于训练长度,那么通过控制$k$,我们就可以在精确保持了局域性的前提下,使得所有位置编码不超过训练长度,简单直接地结合了直接外推和位置内插。

特别地,矩阵$\eqref{eq:leaky-rerope}$还有一个特别的case:当$k\to\infty$时,它简化为
\begin{equation}\begin{pmatrix}
\color{red}{0} & \\
\color{red}{1} & \color{red}{0} & \\
\color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{w} & \color{green}{w} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{w} & \color{green}{w} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{\ddots} & \color{green}{w} & \color{green}{\ddots} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \\
\color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{w} & \color{green}{w} & \color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{w} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{w} & \color{green}{w} & \color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\end{pmatrix}\label{eq:rerope}\end{equation}
在这个case下,不管输入长度是多少,它的位置编码范围都不超过$w$,所以这是一种有可能支持任意长度的Context的方案!

形式上,矩阵$\eqref{eq:rerope}$、$\eqref{eq:leaky-rerope}$与标准RoPE矩阵$\eqref{eq:rope}$的关系,就相当于ReLU、Leaky ReLU与Linear的关系,所以笔者将$\eqref{eq:rerope}$称为“ReRoPE(Rectified RoPE)”,将$\eqref{eq:leaky-rerope}$称为“Leaky ReRoPE”。

计算 #

其实,类似的思路并不难想到,以往基于Attention Bias的相对位置编码(比如经典相对位置编码T5位置编码)经常会出现这样的分块运算。然而跟这些相对位置编码不同,在RoPE中实现这样的分块运算会明显增加计算量,这也是该思路会被笔者搁置的主要原因。

怎么理解增加计算量呢?我们知道RoPE是“通过绝对位置实现相对位置”,这样只能得到线性的相对位置,而矩阵$\eqref{eq:leaky-rerope}$、$\eqref{eq:rerope}$是非线性的(或者说分段线性的),要实现它只能算两次Attention矩阵,然后组合起来。具体来说,首先用标准的RoPE计算一次Attention矩阵(Softmax之前)
\begin{equation}a_{i,j}^{(1)} = \left(\boldsymbol{\mathcal{R}}^i\boldsymbol{q}_i\right)^{\top}\left(\boldsymbol{\mathcal{R}}^j\boldsymbol{k}_j\right) = \boldsymbol{q}_i^{\top}\boldsymbol{\mathcal{R}}^{j-i}\boldsymbol{k}_j\end{equation}
这里第一个等号是实现方式,第二个等号是等效结果,其中$\boldsymbol{\mathcal{R}}$就是RoPE的旋转矩阵,简单起见我们省略了Attention的scale因子。接着,我们需要计算间隔为$1/k$的RoPE的Attention矩阵(Leaky ReRoPE):
\begin{equation}a_{i,j}^{(2)} = \left(\boldsymbol{\mathcal{R}}^{(i-w)/k+w}\boldsymbol{q}_i\right)^{\top}\left(\boldsymbol{\mathcal{R}}^{j/k}\boldsymbol{k}_j\right) = \boldsymbol{q}_i^{\top}\boldsymbol{\mathcal{R}}^{(j-i+w)/k-w}\boldsymbol{k}_j\end{equation}
如果是ReRoPE,那么简单一些:
\begin{equation}a_{i,j}^{(2)} = \left(\boldsymbol{\mathcal{R}}^w\boldsymbol{q}_i\right)^{\top}\boldsymbol{k}_j = \boldsymbol{q}_i^{\top}\boldsymbol{\mathcal{R}}^w\boldsymbol{k}_j\end{equation}
最后,根据$i - j < w$这个条件,将它们合并起来:
\begin{equation}a_{i,j} = \left\{\begin{aligned}
&a_{i,j}^{(1)},\quad (i - j < w) \\[8pt] &a_{i,j}^{(2)}, \quad (i - j \geq w)
\end{aligned}\right.\end{equation}
不管是ReRoPE还是Leaky ReRoPE,都不可避免地计算两次Attention矩阵(如果有更高效的实现方法,请赐教),这便是增加的计算量之一。此外,需要自定义计算Attention矩阵也导致了不能直接套用现成的flash attention实现,因此相对之下又增加了计算成本。

另一方面,同样是由于非线性的相对位置,所以在自回归解码时,Key序列的cache只能存RoPE之前的,然后在每步解码时给整个Key序列补上对应的RoPE,这样的改动也会增加推理计算量。唯一的好消息是,在token by token解码时,从第二步开始Query序列的长度就为1,此时只需要为Key序列定制RoPE,那么可以只算一次Attention矩阵:
\begin{equation}a_{i,j} = \left\{\begin{aligned}
&\boldsymbol{q}_i^{\top}\left(\boldsymbol{\mathcal{R}}^{\max(j-i,-w)}\boldsymbol{k}_j\right), \quad(\text{ReRoPE})\\[8pt]
&\boldsymbol{q}_i^{\top}\left(\boldsymbol{\mathcal{R}}^{\max(j-i,(j-i+w)/k-w)}\boldsymbol{k}_j\right), \quad(\text{Leaky ReRoPE})
\end{aligned}\right.\end{equation}

实验 #

继续沿着《Transformer升级之路:11、将β进制位置进行到底》的设置,我们对ReRoPE进行了实验,效果如下表:
\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^{\color{red}{\dagger}}\text{-fixed} & 49.41\% & 55.94\% & 41.11\% \\
\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^{\color{red}{\dagger}}\text{-mixed} & 49.41\% & 59.11\% & 42.38\% \\
\text{NTK-RoPE-}\log n\text{-mixed} & 49.40\% & 68.91\% & 45.41\% \\
\hline
\text{ReRoPE-w256} & 49.41\% & 77.90\% & 48.48\% \\
\text{ReRoPE-w256-}\log n^{\color{red}{\dagger}} & 49.41\% & 82.40\% & 48.85\% \\
\text{ReRoPE-w256-}\log n & 49.40\% & \boldsymbol{85.12\%} & \boldsymbol{49.07\%} \\
\hline
\text{HFWA} & 48.70\% & 80.84\% & 48.15\% \\
\hline
\end{array}
正如文章开头所说,ReRoPE不微调外推的效果可谓出奇地好,不仅明显超越了此前最优的NTK-RoPE-mixed,还明显超过了从零预训练的HFWA!这里的$\text{w256}$指的$w=256$,$\log n^{\color{red}{\dagger}}$是指预训练没有加入$\log n$缩放(比如LLAMA),测试阶段每个$\boldsymbol{q}_n$都乘上$\max(1, \log_{\text{maxlen}} n)$,$\log n$则是指预训练就加入了$\log n$缩放因子。

以下是一些消融实验,显示出ReRoPE关于$w$还是很鲁棒的,最优值大致是训练长度的$1/4\sim 1/2$左右:
\begin{array}{c|cc}
\hline
\text{测试长度} & 512(\text{训练}) & 4096(\text{重复}) & 4096(\text{不重复})\\
\hline
\text{ReRoPE-w64} & 49.41\% & 69.39\% & 45.19\% \\
\text{ReRoPE-w64-}\log n^{\color{red}{\dagger}} & 49.41\% & 78.58\% & 47.42\% \\
\text{ReRoPE-w64-}\log n & 49.40\% & 84.38\% & 48.14\% \\
\hline
\text{ReRoPE-w128} & 49.41\% & 76.11\% & 47.82\% \\
\text{ReRoPE-w128-}\log n^{\color{red}{\dagger}} & 49.41\% & 82.28\% & 48.78\% \\
\text{ReRoPE-w128-}\log n & 49.40\% & \boldsymbol{85.47\%} & 48.87\% \\
\hline
\text{ReRoPE-w256} & 49.41\% & 77.90\% & 48.48\% \\
\text{ReRoPE-w256-}\log n^{\color{red}{\dagger}} & 49.41\% & 82.40\% & 48.85\% \\
\text{ReRoPE-w256-}\log n & 49.40\% & 85.12\% & \boldsymbol{49.07\%} \\
\hline
\text{ReRoPE-w384} & 49.41\% & 70.72\% & 48.15\% \\
\text{ReRoPE-w384-}\log n^{\color{red}{\dagger}} & 49.41\% & 76.42\% & 48.31\% \\
\text{ReRoPE-w384-}\log n & 49.40\% & 83.24\% & 48.62\% \\
\hline
\text{ReRoPE-w512} & 49.41\% & 7.09\% & 8.25\% \\
\text{ReRoPE-w512-}\log n^{\color{red}{\dagger}} & 49.41\% & 7.08\% & 8.25\% \\
\text{ReRoPE-w512-}\log n & 49.40\% & 15.84\% & 10.83\% \\
\hline
\end{array}

下表则对比了ReRoPE和Leaky ReRoPE:
\begin{array}{c|cc}
\hline
\text{测试长度} & 512(\text{训练}) & 4096(\text{重复}) & 4096(\text{不重复})\\
\hline
\text{ReRoPE-w128-}\log n & 49.40\% & \boldsymbol{85.47\%} & 48.87\% \\
\text{Leaky ReRoPE-w128-k64-}\log n & 49.40\% & 85.29\% & 48.96\% \\
\text{Leaky ReRoPE-w128-k32-}\log n & 49.40\% & 85.31\% & 49.03\% \\
\text{Leaky ReRoPE-w128-k16-}\log n & 49.40\% & 85.15\% & \boldsymbol{49.10\%} \\
\text{Leaky ReRoPE-w128-k8-}\log n & 49.40\% & 80.00\% & 48.11\% \\
\hline
\text{ReRoPE-w256-}\log n & 49.40\% & 85.12\% & 49.07\% \\
\text{Leaky ReRoPE-w256-k64-}\log n & 49.40\% & 84.60\% & 49.03\% \\
\text{Leaky ReRoPE-w256-k32-}\log n & 49.40\% & 84.30\% & 48.97\% \\
\text{Leaky ReRoPE-w256-k16-}\log n & 49.40\% & 83.59\% & 48.87\% \\
\text{Leaky ReRoPE-w256-k8-}\log n & 49.40\% & 69.80\% & 45.72\% \\
\hline
\end{array}

作为ReRoPE的一般化,经过精调的Leaky ReRoPE是有机会超过ReRoPE的,但提升很微弱。此外,当$k$取有限值时,能处理的最大长度也是有限的,因为我们不能提前知道要生成的总长度,所以只能预设一个足够大的$k$,但设定为有限值之后,当输入足够长时,就会因为位置编码超出训练长度而效果大幅下降,相比之下ReRoPE则不会有这个风险。总的来说,精调Leaky ReRoPE相比ReRoPE的价值似乎不大。

以上实验结果都只是在1亿参数的GAU模型上测试的,下面给出基于llama2-13b的测试结果(指标是loss,越小越好),它代表了在真正的LLM表现:
\begin{array}{c|cc}
\hline
\text{测试长度} & 4096(\text{训练}) & 8192 & 16384\\
\hline
\text{RoPE} & 1.4967 & 8.8615 & \text{-} \\
\text{NTK-RoPE} & 1.6081 & 1.5417 & 1.5163 \\
\text{ReRoPE} & 1.4996 & 1.4267 & 1.4001 \\
\hline
\end{array}
可以看到,ReRoPE真正做到了几乎不损训练效果(RoPE-4096代表训练效果),并且满足“longer context, lower loss”的理想特点(更多的context应该更加有助于预测)。此外,笔者也在OpenBuddy开源的LLAMA2-13b微调模型上测试了chat的效果,自我感觉还不错(最多测试过20k tokens的Context)。

最后,分享笔者在transformers的LLAMA模型基础上实现ReRoPE和Leaky ReRoPE的代码,读者也可以自行加载LLAMA系列模型进行测试:

小结 #

在这篇文章中,笔者提出了ReRoPE (Rectified RoPE),它同样是一种RoPE的后处理方案,实验结果显示它的不微调长度外推能力不仅明显超过了此前的NTK-aware Scaled RoPE,甚至还超过了之前专门设计的需要从零训练的HFWA。此外,不同于NTK-aware Scaled RoPE在超过某个长度后能力会大幅下降,ReRoPE似乎在任意长度下都表现良好。除了对比实验外,文章还给出了基于transformers-llama的参考实现,有兴趣的读者可以自行测试。

转载到请包括本文地址:https://spaces.ac.cn/archives/9708

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

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

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

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

苏剑林. (Aug. 07, 2023). 《Transformer升级之路:12、无限外推的ReRoPE? 》[Blog post]. Retrieved from https://spaces.ac.cn/archives/9708

@online{kexuefm-9708,
        title={Transformer升级之路:12、无限外推的ReRoPE?},
        author={苏剑林},
        year={2023},
        month={Aug},
        url={\url{https://spaces.ac.cn/archives/9708}},
}