Transformer升级之路：逆用Leaky ReRoPE

news2024/11/25 20:20:32

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 科学空间

研究方向 | NLP、神经网络

在《Transformer升级之路：无限外推的ReRoPE？》中，笔者提出了 ReRoPE 和 Leaky ReRoPE，诸多实验结果表明，它们能够在几乎不损失训练效果的情况下免微调地扩展 LLM 的 Context 长度，并且实现了 “longer context, lower loss” 的理想特性，此外跟 NTK-aware Scaled RoPE 不同的是，其中 ReRoPE 似乎还有表现出了无限的 Context 处理能力。

总之，ReRoPE 看起来相当让人满意，但美中不足的是会增加推理成本，具体表现为第一步推理需要算两次 Attention，以及后续每步推理需要重新计算位置编码。本文试图通过在训练中逆用 Leaky ReRoPE 的方法来解决这个问题。

回顾

让我们不厌其烦地重温一下：RoPE 形式上是一种绝对位置编码，但实际达到的效果是相对位置编码，对应的相对位置矩阵是：

为了在保留局域性的同时避免 Long Context 导致位置越界问题，Leaky ReRoPE 将推理阶段的相对位置矩阵改为：

其中w是窗口宽度，大概取训练长度的到，用来调节可处理的最大长度，一般使得不超过训练长度的一半为佳。至于 ReRoPE，则是直接取了的极限：

反转

从上一篇的评测结果来看，作为一种免训练的外推方案，ReRoPE 和 Leaky ReRoPE 的效果都是相当让人满意的，既没有损失训练长度内的效果，又实现了 “Longer Context, Lower Loss”。唯一美中不足的是，它们的推理速度相比原本的 Attention 来说是变慢的，并且目前尚不兼容 Flash Attention 等加速技术。

那么，能否反过来呢？ReRoPE/Leaky ReRoPE 在训练阶段是正常速度的 RoPE，推理阶段则是变慢了，反过来也就是说：能否让训练阶段变慢，让推理阶段变为常规的 RoPE？

可能有读者疑惑：为什么会想要让训练阶段变慢？训练成本不是更高吗？这是因为 ReRoPE/Leaky ReRoPE 是一种长度外推方法，场景是 “Train Short, Test Long”，训练速度的变慢是短期的、可控的，推理速度的变慢才是长期的、难顶的，所以相较之下，如果是同等程度的变慢的话，我们更愿意将变慢的部分放到训练阶段。

让我们再看一下 Leaky ReRoPE，它在训练阶段的相对位置矩阵是步长为 1 的式（1），推理阶段则在 w 的窗口内使用 1 的步长，在窗口外使用的步长，即式（2），换句话说，差别是推理阶段窗口外使用更小的步长。如果我们反过来，在训练阶段使用 Leaky ReRoPE，并让它窗口外的步长大于 1，那么按照“推理阶段窗口外使用更小的步长”的原则，推理阶段窗口外是否就可以使用等于 1 的步长，从而退化为 RoPE了？

笔者将以上想法称之为 “InvLeaky ReRoPE（Inverse Leaky ReRoPE）”。事不宜迟，我们马上做实验测试。

实验

继续之前的 “GAU + Deep Norm + Tiger + 语言模型”实验组合，在训练阶段使用的 Leaky ReRoPE，在推理阶段使用正常的 RoPE，测试结果如下：

其中是指 RoPE 的频率底数从 10000 换成了 80000。可以看到，“Leaky ReRoPE → RoPE”的InvLeaky ReRoPE 虽然效果上不如 “RoPE → ReRoPE/Leaky ReRoPE”，但依然胜过了 HFWA，并且由于推理阶段是常规的RoPE，可以套用现成的加速技术，因此依然是有相当竞争力的。

此外，笔者对等参数做了一些简单的调参，发现最优解基本上就是以上两个组合了，即“ 设置为‘扩展倍数的 2 倍的倒数’、 设置为训练长度的 、 可选乘以扩展倍数”。

那么，InvLeaky ReRoPE 对训练速度有多大影响呢？在上述实验中，模型是 1 亿参数量，训练长度是 512，每 1000 步的训练时间从 330 秒增加到了 350 秒，增加不到 10%，当然这里边有 GAU 的原因，因为 GAU 是单头的注意力，本就比多头注意力快。如果多头注意力或者训练长度更长的话，增加幅度应该会大一些，但目测应该不超过 50% 都是可以接受的。