Transformer升级之路：逆用Leaky ReRoPE解决推理成本问题 2024-05-29 作者 C3P00 在上周的文章《Transformer升级之路：12、无限外推的ReRoPE？》中，我们介绍了ReRoPE和Leaky ReRoPE的创新方法。实验结果表明，它们在几乎不损失训练效果的情况下能够扩展LLM的Context长度，并且实现了“longer context, lower loss”的理想特性。尤其是ReRoPE，似乎表现出了无限的Context处理能力。然而，这些方法也带来了推理阶段的成本增加问题。本文将探讨一种新的解决方案：逆用Leaky ReRoPE。 回顾RoPE与Leaky ReRoPE RoPE与位置编码 RoPE（Rotary Position Embedding）形式上是一种绝对位置编码，但实际上它实现了相对位置编码。其对应的相对位置矩阵为： [latex][\left(\begin{array}{cccccccccc}0 & 1 & 2 & 3 & \cdots & L-2 & L-1\-1 & 0 & 1 & 2 & \cdots & L-3 & L-2\-2 & -1 & 0 & 1 & \cdots & L-4 & L-3\\vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots\2-L & 3-L & 4-L & \cdots & -2 & -1 & 0\1-L & 2-L & 3-L & \cdots & -3 & -2 & -1\\end{array}\right)][/latex] Leaky ReRoPE与窗口化处理 为了在保留局域性的同时避免Long Context导致位置越界问题，Leaky ReRoPE将推理阶段的相对位置矩阵改为： [latex][\left(\begin{array}{cccccccccc}0 & 1 & 2 & \cdots & w-1 & w & \cdots & w+k & \cdots & w+L-1-wk\-1 & 0 & 1 & \cdots & w-2 & w-1 & \cdots & w+k-1 & \cdots & w+L-2-wk\\vdots & \vdots & \vdots & \ddots & \vdots & \vdots & \ddots & \vdots & \vdots & \vdots\-w+1 & -w+2 & -w+3 & \cdots & -1 & 0 & \cdots & k-1 & \cdots & L-2-wk\\vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots\\end{array}\right)][/latex] 其中，( w ) 是窗口宽度，( k ) 用来调节可处理的最大长度。 ReRoPE的无限扩展能力 ReRoPE直接取了 ( k \to \infty ) 的极限： [latex][\left(\begin{array}{cccccccccc}0 & 1 & 2 & \cdots & w-1 & w & \cdots & w & \cdots & w\-1 & 0 & 1 & \cdots & w-2 & w-1 & \cdots & w & \cdots & w\\vdots & \vdots & \vdots & \ddots & \vdots & \vdots & \ddots & \vdots & \vdots & \vdots\-w+1 & -w+2 & -w+3 & \cdots & -1 & 0 & \cdots & w & \cdots & w\\vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots\\end{array}\right)][/latex] 逆用Leaky ReRoPE：让训练阶段变慢，推理阶段变快 为什么逆用？ 原本的ReRoPE和Leaky ReRoPE在推理阶段增加了计算成本。如果我们反过来在训练阶段使用Leaky ReRoPE，而在推理阶段使用常规的RoPE，能否解决这一问题呢？ 实验与结果 我们进行了以下实验组合：“GAU + Deep Norm + Tiger + 语言模型”。在训练阶段使用 ( k=1/16, w=128 ) 的Leaky ReRoPE，在推理阶段使用正常的RoPE。测试结果如下： 测试长度BaselineBaseline-lognNTK-RoPE-fixedNTK-RoPE-logn†-fixedNTK-RoPE-logn-fixedNTK-RoPE-mixedNTK-RoPE-logn†-mixedNTK-RoPE-logn-mixedReRoPE-w256ReRoPE-w256-logn†ReRoPE-w256-lognInvLeaky ReRoPE-w128-lognInvLeaky ReRoPE-w128-b8-lognHFWA512 (训练)训练49.41%49.40%49.41%49.41%49.40%49.41%49.41%49.40%49.41%49.41%49.40%49.38%49.62%48.70%4096（重复）24.17%24.60%51.86%55.94%62.85%53.09%59.11%68.91%77.90%82.40%85.12%82.25%81.15%80.84%4096（不重复）23.16%24.02%39.61%41.11%44.14%40.12%42.38%45.41%48.48%48.85%49.07%48.32%48.85%48.15% 其中，( b8 ) 是指RoPE的频率底数从10000换成了80000。可以看到，“Leaky ReRoPE → RoPE”的InvLeaky ReRoPE虽然效果上不如“RoPE → ReRoPE/Leaky ReRoPE”，但依然胜过了HFWA，并且由于推理阶段是常规的RoPE，可以套用现成的加速技术，因此依然是有相当竞争力的。 我们对 [latex]( k, w, b )[/latex] 等参数做了一些简单的调参，发现最优解基本上就是上述两个组合。具体来说： 调参与训练速度 ( k ) 设置为“扩展倍数的2倍的倒数” ( w ) 设置为训练长度的四分之一 ( b ) 可选乘以扩展倍数 在上述实验中，模型参数量为1亿，训练长度为512，每1000步的训练时间从330秒增加到了350秒，增加不到10%。虽然这里有GAU的原因，因为GAU是单头的注意力，本身比多头注意力快。如果是多头注意力或者训练长度更长，增加幅度可能会更大，但估计不会超过50%，依然在可接受范围内。 小结 本文提出了Leaky ReRoPE的“逆用”方法，通过在训练阶段使用更大步长的Leaky ReRoPE，使得推理阶段可以退回常规的RoPE，从而保持推理速度不变。实验结果表明，这种方法依然具有相当的竞争力。未来的工作可以进一步优化参数设置，提升模型性能。 希望这篇文章能够帮助您更好地理解逆用Leaky ReRoPE的方法及其优势。如果有任何疑问或建议，欢迎在评论区留言讨论。
在上周的文章《Transformer升级之路：12、无限外推的ReRoPE？》中，我们介绍了ReRoPE和Leaky ReRoPE的创新方法。实验结果表明，它们在几乎不损失训练效果的情况下能够扩展LLM的Context长度，并且实现了“longer context, lower loss”的理想特性。尤其是ReRoPE，似乎表现出了无限的Context处理能力。然而，这些方法也带来了推理阶段的成本增加问题。本文将探讨一种新的解决方案：逆用Leaky ReRoPE。
回顾RoPE与Leaky ReRoPE
RoPE与位置编码
RoPE（Rotary Position Embedding）形式上是一种绝对位置编码，但实际上它实现了相对位置编码。其对应的相对位置矩阵为：
[latex][
\left(\begin{array}{cccccccccc}
0 & 1 & 2 & 3 & \cdots & L-2 & L-1\
-1 & 0 & 1 & 2 & \cdots & L-3 & L-2\
-2 & -1 & 0 & 1 & \cdots & L-4 & L-3\
\vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots\
2-L & 3-L & 4-L & \cdots & -2 & -1 & 0\
1-L & 2-L & 3-L & \cdots & -3 & -2 & -1\
\end{array}\right)
][/latex]
Leaky ReRoPE与窗口化处理
为了在保留局域性的同时避免Long Context导致位置越界问题，Leaky ReRoPE将推理阶段的相对位置矩阵改为：
[latex][
\left(\begin{array}{cccccccccc}
0 & 1 & 2 & \cdots & w-1 & w & \cdots & w+k & \cdots & w+L-1-wk\
-1 & 0 & 1 & \cdots & w-2 & w-1 & \cdots & w+k-1 & \cdots & w+L-2-wk\
\vdots & \vdots & \vdots & \ddots & \vdots & \vdots & \ddots & \vdots & \vdots & \vdots\
-w+1 & -w+2 & -w+3 & \cdots & -1 & 0 & \cdots & k-1 & \cdots & L-2-wk\
\vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots\
\end{array}\right)
][/latex]
其中，( w ) 是窗口宽度，( k ) 用来调节可处理的最大长度。
ReRoPE的无限扩展能力
ReRoPE直接取了 ( k \to \infty ) 的极限：
[latex][
\left(\begin{array}{cccccccccc}
0 & 1 & 2 & \cdots & w-1 & w & \cdots & w & \cdots & w\
-1 & 0 & 1 & \cdots & w-2 & w-1 & \cdots & w & \cdots & w\
\vdots & \vdots & \vdots & \ddots & \vdots & \vdots & \ddots & \vdots & \vdots & \vdots\
-w+1 & -w+2 & -w+3 & \cdots & -1 & 0 & \cdots & w & \cdots & w\
\vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots\
\end{array}\right)
][/latex]
逆用Leaky ReRoPE：让训练阶段变慢，推理阶段变快
为什么逆用？
原本的ReRoPE和Leaky ReRoPE在推理阶段增加了计算成本。如果我们反过来在训练阶段使用Leaky ReRoPE，而在推理阶段使用常规的RoPE，能否解决这一问题呢？
实验与结果
我们进行了以下实验组合：“GAU + Deep Norm + Tiger + 语言模型”。在训练阶段使用 ( k=1/16, w=128 ) 的Leaky ReRoPE，在推理阶段使用正常的RoPE。测试结果如下：
其中，( b8 ) 是指RoPE的频率底数从10000换成了80000。可以看到，“Leaky ReRoPE → RoPE”的InvLeaky ReRoPE虽然效果上不如“RoPE → ReRoPE/Leaky ReRoPE”，但依然胜过了HFWA，并且由于推理阶段是常规的RoPE，可以套用现成的加速技术，因此依然是有相当竞争力的。
我们对 [latex]( k, w, b )[/latex] 等参数做了一些简单的调参，发现最优解基本上就是上述两个组合。具体来说：
调参与训练速度
在上述实验中，模型参数量为1亿，训练长度为512，每1000步的训练时间从330秒增加到了350秒，增加不到10%。虽然这里有GAU的原因，因为GAU是单头的注意力，本身比多头注意力快。如果是多头注意力或者训练长度更长，增加幅度可能会更大，但估计不会超过50%，依然在可接受范围内。
小结
本文提出了Leaky ReRoPE的“逆用”方法，通过在训练阶段使用更大步长的Leaky ReRoPE，使得推理阶段可以退回常规的RoPE，从而保持推理速度不变。实验结果表明，这种方法依然具有相当的竞争力。未来的工作可以进一步优化参数设置，提升模型性能。
希望这篇文章能够帮助您更好地理解逆用Leaky ReRoPE的方法及其优势。如果有任何疑问或建议，欢迎在评论区留言讨论。