探索 ReRoPE：一种崭新的位置编码方法 2024-05-28 作者 C3P00 在 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-4kNTK-RoPE-16kReRoPE-w1024-16k4k1.49671.51631.40018k8.86151.54171.4267 在 llama2-13b 模型上的实验结果 方法LossRoPE-4k1.4967RoPE-8k8.8615NTK-RoPE-4k1.6081NTK-RoPE-8k1.5417ReRoPE-w1024-4k1.4996ReRoPE-w1024-8k1.4267ReRoPE-w1024-16k1.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。
在 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’ ) 的定义如下：✅
[
T =
\begin{cases}
\text{RoPE}(q, k) & \text{if } |q – k| < w \
\text{Leaky RoPE}(q, k) & \text{otherwise}
\end{cases}
]
计算复杂度
实现 ReRoPE 增加了一些计算复杂度，特别是在自回归解码时，需要为每步解码时给整个 Key 序列补上对应的 RoPE。这种改动会增加推理计算量，但对于 token-by-token 解码，仅需计算一次 Attention 矩阵。
实验结果
我们在 1 亿参数的 GAU 模型和 llama2-13b 模型上进行了实验，结果显示 ReRoPE 的效果显著优于 NTK-aware Scaled RoPE，甚至超过了 HFWA。
在 GAU 模型上的实验结果
在 llama2-13b 模型上的实验结果
可以看到，ReRoPE 几乎不损伤训练效果，并且在更长的上下文中表现更好。
实现与使用
我们在 transformers 的 LLAMA 模型基础上实现了 ReRoPE 和 Leaky ReRoPE，读者可以在 GitHub 上查看详细代码。
依赖
测试方法
总结
ReRoPE 提供了一种新的位置编码方法，显著提升了长度外推能力。在实验中，ReRoPE 的表现超过了目前最优的 NTK-aware Scaled RoPE 和 HFWA，且在任意长度下都表现良好。希望本文的分享能为大家在模型优化上带来新的思路。
如有兴趣，欢迎访问 GitHub 仓库获取代码并进行测试：GitHub – bojone/rerope。