在 Transformer 模型的演进过程中，位置编码一直是关键的一环。最近，笔者引入了一种新的位置编码方法——ReRoPE (Rectified Rotary Position Embeddings)，通过这种方法，可以显著提升模型在长度外推上的性能。本文将详细介绍 ReRoPE 的背景、原理、实现和实验结果。
背景与问题
在 Transformer 模型中，RoPE（旋转位置编码）是一种绝对位置编码方法，但实际上它给 Attention 带来的是相对位置信息。这种编码方式在长度外推（Length Extrapolation）上效果较差，因为超出训练长度的位置编码并没有被训练过。
研究人员提出了位置内插（Position Interpolation）来解决这个问题，通过调整相对位置矩阵，使得最大相对位置不超过训练长度。然而，这种方法虽然避免了外推，但使位置信息更加“拥挤”，仍需一定步数的微调。
后来，NTK-aware Scaled RoPE 进一步改进，通过“高频外推、低频内插”平摊外推压力，效果有所提升，但仍存在上限。
ReRoPE 的提出
笔者提出了 ReRoPE（Rectified RoPE），通过设定一个窗口大小 ( w )，在窗口内使用常规位置间隔，在窗口外使用较大的间隔，从而精确保持局域性，且所有位置编码不超过训练长度。形式上，ReRoPE 与标准 RoPE 的关系类似于 ReLU 与 Linear 的关系。
具体来说，矩阵 ( T ) 和 ( T' ) 的定义如下：
- 当 ( d = 1 ) 时，( T ) 简化为：
[
T =
\begin{cases}
\text{RoPE}(q, k) & \text{if } |q - k| < w \
\text{Leaky RoPE}(q, k) & \text{otherwise}
\end{cases}
] - 这样，无论输入长度如何，位置编码范围都不超过 ( w )。
计算复杂度
实现 ReRoPE 增加了一些计算复杂度，特别是在自回归解码时，需要为每步解码时给整个 Key 序列补上对应的 RoPE。这种改动会增加推理计算量，但对于 token-by-token 解码，仅需计算一次 Attention 矩阵。
实验结果
我们在 1 亿参数的 GAU 模型和 llama2-13b 模型上进行了实验，结果显示 ReRoPE 的效果显著优于 NTK-aware Scaled RoPE，甚至超过了 HFWA。
在 GAU 模型上的实验结果
| 测试长度 | RoPE-4k | NTK-RoPE-16k | ReRoPE-w1024-16k |
|---|---|---|---|
| 4k | 1.4967 | 1.5163 | 1.4001 |
| 8k | 8.8615 | 1.5417 | 1.4267 |
在 llama2-13b 模型上的实验结果
| 方法 | Loss |
|---|---|
| RoPE-4k | 1.4967 |
| RoPE-8k | 8.8615 |
| NTK-RoPE-4k | 1.6081 |
| NTK-RoPE-8k | 1.5417 |
| ReRoPE-w1024-4k | 1.4996 |
| ReRoPE-w1024-8k | 1.4267 |
| ReRoPE-w1024-16k | 1.4001 |
可以看到，ReRoPE 几乎不损伤训练效果，并且在更长的上下文中表现更好。
实现与使用
我们在 transformers 的 LLAMA 模型基础上实现了 ReRoPE 和 Leaky ReRoPE，读者可以在 GitHub 上查看详细代码。
依赖
- transformers 4.31.0
测试方法
python test.py # 测试聊天功能
python eval_loss.py # 计算 llama2 模型的 loss总结
ReRoPE 提供了一种新的位置编码方法，显著提升了长度外推能力。在实验中，ReRoPE 的表现超过了目前最优的 NTK-aware Scaled RoPE 和 HFWA，且在任意长度下都表现良好。希望本文的分享能为大家在模型优化上带来新的思路。
如有兴趣，欢迎访问 GitHub 仓库获取代码并进行测试：GitHub - bojone/rerope。