RoPE 不该一刀切:为每个注意力头配备可学习的旋转频率与长度缩放,长上下文外推更强。
摘要
旋转位置编码(RoPE)已被 Transformer 广泛采用以编码位置信息,但标准实现对所有注意力头施加统一的频率调度与缩放策略。通过简化的检索任务和长度泛化实验,我们从实验和理论两方面证明:承担不同功能的注意力头需要不同的频率范围和缩放因子才能高效工作。忽视这一结构将导致嵌入维度利用率不足和性能下降,在长上下文场景下尤为明显。为此,我们提出 AdaRoPE,为每个注意力头配备可学习的旋转频率和缩放因子。使用 AdaRoPE 预训练的 LLM 在各项评测中一致优于已有的 RoPE 变体,包括 partial-RoPE 和 NoPE 基线。在上下文扩展方面,我们进一步证明 YaRN 等方法采用的统一频率和注意力缩放并非最优。通过逐头缩放,AdaRoPE 在外推设定和长上下文续训设定中均能实现更好的上下文扩展,同时更好地保持短上下文性能。这些结果突显了在单个注意力头层面优化旋转位置编码的重要性。