03. RoPE Tutorial | 旋转位置编码教程

难度： Medium | 环境： CPU-first | 标签： 基础实现, 位置编码, RoPE | 目标人群： 基础实现学习者

🚀 云端运行环境

本章节的实战代码可以点击以下链接在免费 GPU 算力平台上直接运行：

(国内推荐：魔搭社区免费实例)

---

本节导读

注意力本身只看 token 之间的相似度，并不知道谁在前、谁在后。没有位置编码，模型看到的更像一袋 token；但如果位置处理得太死，模型又很难泛化到更长的上下文。

RoPE 的做法不是给 token 额外加一个位置向量，而是把位置信息写进查询（Query，Q）和键（Key，K）的旋转里，让注意力点积感知相对距离。本节沿着成对维度、位置频率、相对相位和长度外推展开，帮助你判断位置配置改变了哪一层表示。

关键词： RoPE, 位置编码, 复数张量

---

前置阅读

导语： 先确认 Q/K 的形状和注意力点积如何计算，再观察旋转角度和位置频率如何改变相对位置信息。

• P0: 05. PyTorch Tensor Fundamentals | PyTorch 张量基础操作
• P0: 16. Attention Mechanism Intro | 注意力机制导论

---

Step 1：让注意力感知相对位置

注意力的 Q/K 相似度本身不知道 token 的顺序。RoPE 将位置写成二维通道对上的旋转：位置越远，旋转相位差越大；两个位置的内积因而能够表达相对距离。旋转只作用于 Q/K，不直接改写 V，也不改变因果掩码决定的可见范围。

对象	RoPE 改变什么	RoPE 不负责什么
Q/K 表示	按位置旋转成对通道	不决定哪些历史位置可见
相似度	让内积携带相对相位差	不替代注意力的内容匹配
长度配置	决定频率覆盖和外推条件	不保证训练长度外的任务质量

Step 2：成对维度如何表示旋转

隐藏通道两两组成二维平面。第 $i$ 个通道对使用频率 ${\omega }_i$，位置 $m$ 对应角度 $m{\omega }_i$；同一旋转同时作用于 Q 和 K，因此两者的内积主要由位置差而不是绝对位置决定。维度必须为偶数，频率表长度必须覆盖当前序列位置。

中间量	形状或约束	机制含义
通道对	D/2 组	每组构成一个二维旋转平面
频率	高频到低频	分别感知短距离变化和较慢位置变化
复数旋转因子	[S, D/2]	用单位模复数统一表达旋转
广播关系	同一位置覆盖所有批次与注意力头	不改变批次、序列、注意力头的轴语义

Step 3：频率尺度、长度扩展与验证证据

RoPE 为不同维度对分配不同频率。训练长度之外继续使用原频率时，模型会遇到训练中没有充分覆盖的相位组合；仅把 max_position_embeddings 调大并不会自动获得可靠的长上下文能力。扩展方案本质上在调整“位置如何映射到旋转角度”，但调整对象和副作用不同。

方案类型	改变什么	主要风险	必须补充的证据
调整基频 theta	改变各通道对的旋转速度	短距离位置分辨率发生变化	训练长度内回归、目标长度质量
位置插值	将更长位置压回训练位置范围	相邻位置被压缩	短文本质量、长距离检索
频率相关缩放	对不同频率采用不同缩放	配置与模型训练方式强相关	配置来源、长度阶梯和失败样例
重新训练或继续训练	让模型适应新的位置分布	需要额外数据与训练预算	训练曲线、目标任务和成本

Step 4：代码设计与 RoPE 旋转验证

题目区验证基础旋转不变量：频率表由位置与通道对共同确定，Q/K 经过同一位置旋转后保持形状、数据类型和模长。长上下文扩展仍需在固定模型和任务上独立评测。

TODO / 函数	学习者完成的机制	必须满足的约束	测试证据
TODO 1 · precompute_freqs_cis	生成位置与成对维对应的复数旋转因子	head_dim 为偶数；频率表为 [S, D/2]	首位置不旋转，频率表形状正确
TODO 2 · apply_rotary_emb	将 Q/K 成对实数维转换为复数表示	保留批次、序列与注意力头轴	成对维形状与 dtype 正确
TODO 3 · apply_rotary_emb	旋转并还原为实数张量	Q/K 输出与输入形状一致	旋转生效，模长与输入 dtype 保持

import torch

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> torch.Tensor:
    """生成 `[sequence_length, head_dim // 2]` 的复数旋转因子。"""
    if dim <= 0 or dim % 2 != 0 or end <= 0 or theta <= 0:
        raise ValueError('dim 必须为正偶数，end 和 theta 必须为正数')
    # ==========================================
    # TODO 1：构造成对维的逆频率、位置角度和单位复数旋转因子。
    # ==========================================
    # pair_indices = ???
    # inv_freq = ???
    # positions = ???
    # angles = ???
    raise NotImplementedError('TODO 1：请生成 RoPE 频率表')


def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
    """将 `[S, D/2]` 频率表扩展到 `[1, S, 1, D/2]`。"""
    if x.ndim != 4:
        raise ValueError('复数 Q/K 必须是 [B, S, H, D/2]')
    if freqs_cis.shape != (x.shape[1], x.shape[-1]):
        raise ValueError('频率表必须与 sequence 和成对特征维对齐')
    return freqs_cis.view(1, x.shape[1], 1, x.shape[-1])


def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """将同一位置旋转应用到 Query 与 Key。"""
    if xq.shape != xk.shape or xq.ndim != 4 or xq.shape[-1] % 2 != 0:
        raise ValueError('xq/xk 必须是形状相同且最后一维为偶数的四维张量')
    # ==========================================
    # TODO 2：将成对实数维转为 FP32 复数表示。
    # ==========================================
    # xq_complex = ???
    # xk_complex = ???
    raise NotImplementedError('TODO 2：请完成 Q/K 的成对维表示')

    freqs = reshape_for_broadcast(freqs_cis, xq_complex).to(xq_complex.dtype)
    # ==========================================
    # TODO 3：复数旋转后还原为输入 dtype 的实数张量。
    # ==========================================
    # xq_out = ???
    # xk_out = ???
    raise NotImplementedError('TODO 3：请完成 RoPE 旋转与还原')

# 测试设计：分别验证频率表、成对旋转性质和输入契约。

def test_rope_frequency_table():
    """验证频率表形状、零位置旋转与非法配置拒绝。"""
    freqs = precompute_freqs_cis(dim=8, end=5)
    assert freqs.shape == (5, 4)
    assert torch.allclose(freqs[0], torch.ones(4, dtype=freqs.dtype))
    for dim, end in ((7, 5), (8, 0)):
        try:
            precompute_freqs_cis(dim, end)
        except ValueError:
            continue
        raise AssertionError('非法频率表配置应被拒绝')


def test_rope_rotation_contract():
    """验证旋转前后的形状与向量模长不变量。"""
    xq = torch.randn(2, 4, 3, 8)
    xk = torch.randn(2, 4, 3, 8)
    freqs = precompute_freqs_cis(dim=8, end=4)
    q_out, k_out = apply_rotary_emb(xq, xk, freqs)
    assert q_out.shape == xq.shape and k_out.shape == xk.shape
    assert torch.allclose(torch.linalg.vector_norm(q_out, dim=-1), torch.linalg.vector_norm(xq, dim=-1), atol=1e-5)
    assert torch.allclose(torch.linalg.vector_norm(k_out, dim=-1), torch.linalg.vector_norm(xk, dim=-1), atol=1e-5)


def test_rope_position_and_dtype_contract():
    """验证位置差异、输入数据类型和奇数维拒绝。"""
    repeated = torch.ones(1, 3, 1, 4, dtype=torch.float16)
    freqs = precompute_freqs_cis(dim=4, end=3)
    rotated, _ = apply_rotary_emb(repeated, repeated, freqs)
    assert rotated.dtype == torch.float16
    assert torch.allclose(rotated[:, 0], repeated[:, 0])
    assert not torch.allclose(rotated[:, 1], repeated[:, 1])
    try:
        apply_rotary_emb(torch.randn(1, 2, 1, 3), torch.randn(1, 2, 1, 3), freqs[:2])
    except ValueError:
        pass
    else:
        raise AssertionError('奇数 head_dim 应被拒绝')


def run_rope_tests():
    """运行 RoPE 的机制测试集合。"""
    test_rope_frequency_table()
    test_rope_rotation_contract()
    test_rope_position_and_dtype_contract()
    print('✅ RoPE：频率表、旋转性质与输入契约测试通过')


run_rope_tests()

🛑 STOP HERE 🛑

请先完成题目区并通过测试，再继续查看参考代码与解析。

参考代码与解析

代码

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> torch.Tensor:
    """生成 `[sequence_length, head_dim // 2]` 的复数旋转因子。"""
    if dim <= 0 or dim % 2 != 0 or end <= 0 or theta <= 0:
        raise ValueError('dim 必须为正偶数，end 和 theta 必须为正数')
    # TODO 1：构造成对维的逆频率、位置角度和单位复数旋转因子。
    pair_indices = torch.arange(0, dim, 2, dtype=torch.float32)
    inv_freq = 1.0 / (theta ** (pair_indices / dim))
    positions = torch.arange(end, dtype=torch.float32)
    angles = torch.outer(positions, inv_freq)
    return torch.polar(torch.ones_like(angles), angles)


def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
    """将 `[S, D/2]` 频率表扩展到 `[1, S, 1, D/2]`。"""
    if x.ndim != 4:
        raise ValueError('复数 Q/K 必须是 [B, S, H, D/2]')
    if freqs_cis.shape != (x.shape[1], x.shape[-1]):
        raise ValueError('频率表必须与 sequence 和成对特征维对齐')
    return freqs_cis.view(1, x.shape[1], 1, x.shape[-1])


def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """将同一位置旋转应用到 Query 与 Key。"""
    if xq.shape != xk.shape or xq.ndim != 4 or xq.shape[-1] % 2 != 0:
        raise ValueError('xq/xk 必须是形状相同且最后一维为偶数的四维张量')
    # TODO 2：将成对实数维转为 FP32 复数表示。
    xq_complex = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_complex = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))

    freqs = reshape_for_broadcast(freqs_cis, xq_complex).to(xq_complex.dtype)
    # TODO 3：复数旋转后还原为输入 dtype 的实数张量。
    xq_out = torch.view_as_real(xq_complex * freqs).flatten(-2).to(xq.dtype)
    xk_out = torch.view_as_real(xk_complex * freqs).flatten(-2).to(xk.dtype)
    return xq_out, xk_out

解析

TODO 1：生成位置频率表

• 判断依据：每个通道对有独立逆频率，位置与逆频率的外积给出旋转角。
• 实现与证据：构造单位模复数因子；测试检查 [S, D/2] 形状和位置 0 不旋转。

TODO 2：建立成对维表示

• 判断依据：最后一维的相邻两个分量共同构成一个二维旋转平面。
• 实现与证据：先转为 FP32 复数表示，同时保留批次、序列与注意力头轴。

TODO 3：旋转并还原实数张量

• 判断依据：单位复数乘法改变方向但保持每个二维通道对的模长。
• 实现与证据：Q/K 使用相同位置因子后还原为输入数据类型；测试核对形状、模长和位置差异。

相关阅读

RoPE 解释位置如何写入 Q/K；若继续研究长上下文，应接着比较访问范围如何限制或传播远程信息。

• RoFormer 原论文：Enhanced Transformer with Rotary Position Embedding
• Transformers 中的 LLaMA 模型实现
• 04. 多头注意力与 GQA
• 44. 局部、滑动窗口与稀疏 Attention
• 大模型架构 · RoPE 与位置编码