Rotary Position Embedding (RoPE)

旋转位置嵌入（RoPE）通过将每个查询向量和键向量旋转一个与 token 位置成比例的角度，为 Transformer 提供有关词元顺序的信息。自注意力本身无法区分不同的词序，因此模型必须以某种方式注入位置信息；使用 RoPE 时，位置 m 的查询与位置 n 的键之间的注意力分数只取决于它们的偏移量 m − n。RoPE 于 2021 年随 RoFormer 模型提出，并由 EleutherAI 的 GPT-J 和 GPT-NeoX 推广；如今，它已成为开源大型语言模型（LLM）的标准位置编码，包括 Llama、Mistral、Qwen、Gemma 和 DeepSeek。

RoPE 的工作原理#

RoPE 会将每个维度为 d 的查询向量和键向量拆分为 d/2 对特征，并将每对特征视为平面上的一个点。在位置 m 时，第 i 对特征会旋转 m·θᵢ 的角度，其中 θᵢ = base^(−2i/d)，原始论文中 base 为 10,000。前面的特征对旋转较快，可分辨细微的局部顺序；后面的特征对旋转较慢，仍能区分相隔数千个位置的 token。将查询旋转 mθ、将键旋转 nθ 后，两者的点积只取决于差值 (m − n)θ，而且旋转不会改变向量长度。旋转会在每个注意力层内部、查询和键投影之后执行；值向量不会旋转，输入时也不会向 token 的嵌入添加任何内容：

下面的代码片段使用复数实现旋转，与 Meta 的 Llama 参考代码相同，并验证相对位置属性：

import torch

torch.manual_seed(0)
dim = 64
freqs = 10000 ** (-torch.arange(0, dim, 2) / dim)  # one rotation frequency per pair of features


def rope(x, pos):
    """Rotate each consecutive feature pair of x by the angle pos * frequency."""
    pairs = torch.view_as_complex(x.reshape(-1, 2))
    return torch.view_as_real(pairs * torch.polar(torch.ones_like(freqs), pos * freqs)).flatten()


q, k = torch.randn(2, dim)
# The same query and key, three tokens apart, at two different places in the sequence
print(f"positions 5 and 2:     {rope(q, 5) @ rope(k, 2):.4f}")
print(f"positions 105 and 102: {rope(q, 105) @ rope(k, 102):.4f}")
print(f"positions 5 and 4:     {rope(q, 5) @ rope(k, 4):.4f}")  # one token apart
print(f"Vector length kept: {rope(q, 5).norm():.4f} vs {q.norm():.4f}")

无论这两对向量位于位置 5 和 2，还是位置 105 和 102，相隔三个 token 时都会输出 -8.2811；相同向量相隔一个 token 时，输出 -11.4563。旋转后的查询向量长度仍为 8.3705。

RoPE 与其他位置编码的比较#

• 正弦绝对位置编码： 原始 Transformer 在第一层之前，仅将固定的正弦和余弦向量添加到输入嵌入中。模型必须从混入的信号中推断相对距离。
• 可学习绝对位置嵌入： BERT 和 GPT-2 会为每个位置学习一个向量，最多分别支持 512 和 1,024 个位置，超过该限制的位置没有对应的嵌入。
• 相对位置偏置： T5 和 ALiBi 会根据距离向注意力分数添加可学习或固定的偏置。RoPE 通过旋转本身编码相对距离，不增加任何参数。

扩展上下文窗口#

使用 RoPE 在长度为 L 的序列上训练的模型从未见过 L 之后位置对应的旋转角度，因此遇到这些角度时，注意力表现会下降。上下文扩展方法会重新缩放角度，而不是从头开始训练。位置插值会将新增位置压缩到训练范围内，并在 1,000 个微调步骤内将 Llama 模型的上下文窗口扩展到 32,768 个 token。NTK 感知缩放和 YaRN 对旋转较慢的低频特征对进行更大幅度的扩展；与早期方法相比，YaRN 用少 10 倍的训练 token 将 Llama 2 扩展到 128K token。Llama 3 在预训练期间将 base 从 10,000 提高到 500,000，因此旋转最慢的特征对每个 token 的旋转幅度更小，远距离位置也能保持可区分。

视觉模型中的 RoPE#

图像需要两个坐标，因此视觉 Transformer会使用 2D 或轴向 RoPE：每个注意力头的一半特征对根据图像块的列旋转，另一半根据其行旋转。Qwen2-VL 在多模态 M-RoPE 中增加了第三个时间轴，用于处理视频。Meta 的 SAM 2 可在 Ultralytics 软件包中运行，并在其记忆注意力的四层中应用轴向 RoPE；当前视频帧的特征会在此关注早期帧的已存储特征。每个已存储帧的键都会使用与当前帧查询相同的网格角度，因此两帧之间的注意力取决于图像块的空间偏移。