旋转位置编码 RoPE 的原理与实现

所属方向：注意力与核心零件 ｜ 难度：核心组件 ｜ 前置知识：MHA
关键词：位置编码、RoPE、旋转位置编码、相对位置、复数旋转、长度外推

---

01. 为什么需要它

先说一个具体的失败场景。

你训了一个视频 DiT，训练时用 480p、49 帧。推理时想出 720p、97 帧，结果画面从第 50 帧开始漂——人物五官慢慢挪位，背景纹理开始流动。你换了采样器、加了 CFG、调了 shift，都没用。

先把这次改动的量级算清楚。按常见配置（VAE 空间 8 倍、时间 4 倍压缩，再做 2×2 patchify）：

• 480p/49 帧 → 13 个 latent 帧 × 1560 tokens/帧 = 20280 tokens
• 720p/97 帧 → 25 个 latent 帧 × 3600 tokens/帧 = 90000 tokens

总量涨了 4.44 倍，其中空间方向 2.31 倍、时间方向 1.92 倍。也就是说，模型要给四倍多的位置编号赋予意义，新增的较远坐标与注意力关系超出了训练范围（已有坐标仍然见过）。

这是一个待诊断的假设场景：位置编码、训练尺度、VAE、注意力规模和运动分布都可能造成失败，不能仅凭症状归因给 RoPE。

注意力机制本身是置换等变的：把输入序列的顺序打乱，输出也只是跟着打乱，每个 token 看到的上下文完全不变。换句话说，纯注意力分不清「第 3 帧」和「第 40 帧」。位置信息必须额外注入。

早期做法是可学习的绝对位置嵌入（learned absolute PE）：给每个位置分配一个可训练向量，加到 token 上。这在训练长度内工作良好，但有个硬伤——没见过的位置没有对应向量。你训练时只见过 49 个位置，推理要 97 个，后面 48 个位置的嵌入是凭空来的，模型自然就崩了。

RoPE 解决的正是这件事：它不给位置分配「向量」，而是给位置分配「旋转角度」。角度公式可计算到新位置；生成质量能否外推仍需训练或扩展策略验证。

02. 最小可用理解

三句话：

1. 把 query 和 key 的特征维度两两配对看成复数，每个位置 $m$ 把它们旋转 $m\theta$ 弧度。
2. 两个向量做内积时，指数项相减，绝对位置 $m$、$n$ 自动消掉，只剩相对距离 $(m-n)$。
3. 所以对固定的 q、k，位置项只依赖 token 间的相对距离；完整分数仍依赖内容特征，而旋转角度可以外推到训练时没见过的位置。

如果只记一件事：RoPE 是乘上去的旋转，不是加上去的向量。加法会污染特征的模长，旋转不会——这是它比绝对位置嵌入更「干净」的地方（3.4 节会把这句话量化）。

还有一个位置要记准：它作用在每一层注意力的 $q$ 和 $k$ 上，不是输入层加一次，通常不旋转 $v$。这是本文及常用实现的约定，某些变体也会旋转 value，不能一概判错。

03. 数学推导

3.1 我们到底想要什么性质

先把目标写成式子。设 $x_m$ 是位置 $m$ 上的 token 特征，我们要找一个变换 $f$，使得变换后的 query 和 key 的内积满足：

$$\langle f(q, m), f(k, n) \rangle = g(q, k, m-n)$$

左边依赖两个绝对位置 $m$ 和 $n$，右边只依赖它们的差值 $m-n$。这就是「相对位置不变性」的精确表述。注意 $g$ 的具体形式我们不关心，只要求它不含 $m$、$n$ 的单独出现。

3.2 复数域上的现成答案

二维情况下有个现成的东西满足这个条件：旋转。

把二维向量 $q = (q_0, q_1)$ 看成复数 $q_0 + iq_1$。复数乘法 $e^{im\theta}$ 的几何意义就是逆时针旋转 $m\theta$ 弧度，且不改变模长。

现在令 $f(q, m) = q e^{im\theta}$，算一下内积。复数域上的内积要取共轭：

$$\langle f(q,m), f(k,n) \rangle = \mathrm{Re}\left[(q e^{im\theta}) \overline{(k e^{in\theta})}\right] = \mathrm{Re}\left[q \bar{k} e^{i(m-n)\theta}\right]$$

关键一步在指数：$e^{im\theta} \cdot \overline{e^{in\theta}} = e^{im\theta} e^{-in\theta} = e^{i(m-n)\theta}$。绝对位置就这么消掉了，只留下 $(m-n)$。这就是 RoPE 的全部动机，后面所有工程细节都只是把这个二维结论铺到高维。

3.3 铺到高维

实际的 head_dim 是 64 或 128，不是 2。做法是把 $d$ 维切成 $d/2$ 组，每组两个相邻维度构成一个复数平面，第 $i$ 组用自己的旋转频率：

$$\theta_i = \text{base}^{-2i/d}, \quad i = 0, 1, \dots, d/2 - 1$$

这里的 base 就是代码里的 rope_theta，通常取 10000。

这个频率设计是理解 RoPE 行为的关键。展开看两端：

• $i = 0$ 时 $\theta_0 = 1$，位置每 +1 就转 1 弧度，转得最快
• $i = d/2 - 1$ 时 $\theta_i \approx \text{base}^{-1} = 10^{-4}$，位置每 +1 只转 0.0001 弧度，转得最慢

也就是说，高频维度负责区分近距离，低频维度负责区分远距离。一个位置差 1 的 token 对，主要靠前面几维的相位差来分辨；位置差较大时，高频相位已多次绕圈，而低频相位仍提供变化较慢的尺度；各频段并非互斥地负责某一距离。

这张图要看什么：左图是 $\cos(m\theta_i)$ 的热力图（横轴位置、纵轴维度组）——256 长度内转满一整圈的只有前 13 组；这些高频组提供较大的近距离相位变化，维度索引越大，相位变化越慢，但并非不携带近距离位置信息；右图把位置拉到 32768，高频组早就转了几万弧度，低频组才刚积累出可分辨的相位。这是 32 组频率提供的多尺度结构，不是互斥的距离分工；后面所有外推方法（PI、NTK、YaRN）做的都是重新分配这个分工（图由 code/make_figures.py 生成）。

这个结构直接解释了长度外推的两种主流做法：位置插值（PI）把所有位置除以一个缩放因子，等价于整体降低旋转速度；NTK-aware 缩放只调 base，让低频维度转得更慢而尽量保住高频维度的分辨率。它们改的都是同一组 $\theta_i$。

写成矩阵形式，RoPE 就是一个分块对角的旋转矩阵 $R_m$，第 $i$ 个 $2\times2$ 块为：

$$\begin{pmatrix} \cos m\theta_i & -\sin m\theta_i \\ \sin m\theta_i & \cos m\theta_i \end{pmatrix}$$

于是注意力分数变成 $q_m^\top R_{n-m} k_n$——形式上是把相对位置信息塞进了双线性形式的中间，而没有动 $q$、$k$ 本身的模长。

3.4 「不改模长」为什么要紧

02 节说过 RoPE 是乘上去的旋转而不是加上去的向量，这里补上为什么这件事重要。

$R_m$ 是正交矩阵（$R_m^\top R_m = I$），所以 $\|R_m x\| = \|x\|$——位置编码前后，token 的模长分毫不变。加法式绝对位置嵌入做不到这点：$\|x + p_m\| \neq \|x\|$，而且改变量取决于 $p_m$ 的方向，是随位置乱跳的。

这有实际后果。注意力分数里 $q$、$k$ 的模长直接决定 softmax 前的量级，而模长会影响打分量级，但不能普遍解释成 token 的重要性或置信度。加法式 PE 会把位置信息和这个量级搅在一起：同一个词放在句首和句中，模长不同，softmax 后的分布也跟着偏移。旋转则是干净的——位置只影响方向（相位），不碰长度。

实测一下（code/rope_decay.py 第 1 段）：

原始向量模长                7.315344
RoPE 后（位置 0/1/100/511） 7.315344  7.315344  7.315344  7.315344
跨位置模长极差              8.88e-16
加法式 PE 后                9.843323  10.080425  10.808198  10.190631
跨位置模长极差              0.9649  ← 同一个 token 的模长被位置改写了

RoPE 那行四个位置的模长在浮点精度内完全一致（偏差 8.88e-16 是机器 epsilon 量级）；加法式则在 9.84 到 10.81 之间晃，同一个 token 的「强度」被它所在的位置改写了将近 10%。

04. 代码实现

下面是完整可跑的最小实现。刻意用 numpy 而不是 PyTorch：讲原理时框架的抽象反而是噪声，而且 numpy 谁都能跑。

import numpy as np


def build_rope_cache(seq_len: int, dim: int, base: float = 10000.0):
    """预计算 RoPE 的 cos/sin 表。

    dim 必须是偶数：每两个相邻维度配成一个复数参与旋转。
    base 就是常说的 rope_theta，调它会按维度改变旋转速度；最高频 i=0 始终为 1。
    """
    assert dim % 2 == 0, f"dim 必须为偶数，收到 {dim}"
    # theta_i = base^(-2i/dim)：i 越大转得越慢，对应越低的频率
    inv_freq = 1.0 / (base ** (np.arange(0, dim, 2) / dim))
    pos = np.arange(seq_len)
    freqs = np.outer(pos, inv_freq)              # [seq_len, dim/2]
    return np.cos(freqs), np.sin(freqs)


def apply_rope(x: np.ndarray, cos: np.ndarray, sin: np.ndarray) -> np.ndarray:
    """把旋转作用到 x 上。x: [seq_len, dim]

    最后一维两两配对看成复数 (x_even + i*x_odd)，乘上 e^{i*m*theta}：
        实部 = x_even * cos - x_odd * sin
        虚部 = x_even * sin + x_odd * cos
    """
    x_even, x_odd = x[..., 0::2], x[..., 1::2]
    out = np.empty_like(x)
    out[..., 0::2] = x_even * cos - x_odd * sin
    out[..., 1::2] = x_even * sin + x_odd * cos
    return out

验证核心性质：固定一对向量，只改赋予它们的绝对位置，间距始终为 1。如果推导是对的，内积应该完全不变。

rng = np.random.default_rng(0)
seq_len, dim = 8, 64
cos, sin = build_rope_cache(seq_len, dim)
q = rng.standard_normal(dim)
k = rng.standard_normal(dim)

for shift in range(4):
    qr = apply_rope(q[None, :], cos[shift:shift + 1], sin[shift:shift + 1])
    kr = apply_rope(k[None, :], cos[shift + 1:shift + 2], sin[shift + 1:shift + 2])
    print(f"位置 ({shift}, {shift + 1}) → 内积 {float(qr[0] @ kr[0]):.6f}")

真实输出：

cos/sin 形状: (8, 32) (8, 32)
位置 1 的前 4 维 cos: [0.5403, 0.7318, 0.846, 0.9124]
位置 1 的末 2 维 cos: [1.0, 1.0]

同一对向量，平移绝对位置（间距恒为 1）:
  位置 (0, 1) → 内积 -11.297869
  位置 (1, 2) → 内积 -11.297869
  位置 (2, 3) → 内积 -11.297869
  位置 (3, 4) → 内积 -11.297869

对照组，改变间距:
  间距 1      → 内积 -11.297869
  间距 2      → 内积 -12.637566
  间距 3      → 内积 -13.049657

未加 RoPE 的原始内积: -9.701179

三个数字值得盯一下。

第一，四行内积完全相同（-11.297869），小数点后六位都不差。这就是 3.1 节那个式子在数值上的体现：绝对位置从 (0,1) 挪到 (3,4)，分数一点没变。而间距一变（1→2→3），内积立刻变化。

第二，位置 1 的前 4 维 cos 是 0.5403、0.7318、0.846、0.9124，递增。对应 $\theta_i$ 递减：第 0 维转了整整 1 弧度（cos 1 ≈ 0.5403），第 3 维只转了约 0.42 弧度。这就是「高频维度在前」的直接观测。

第三，末 2 维 cos 打印出来是 1.0。不是真的等于 1，而是 $\theta_{31} \approx 10^{-4}$，位置 1 只转了 0.0001 弧度，$\cos(0.0001) = 0.999999995$，六位小数下显示成 1.0。这些维度在近距离上几乎不携带位置信息——它们是留给远距离的。也正因为如此，NTK-aware 缩放对这些低频维度施加更大压缩；实际质量代价仍取决于训练后的频段使用方式。

上面这张图（code/make_figures.py 图 3）把本节的两件事画在了一起。左图是平移不变性的可视化：同一对 $q$、$k$，把绝对起点放在 0 / 500 / 2000，逐点画出内积随相对距离的曲线——三条曲线完全重合（最大偏差 5.2e-13），这就是上面输出里「四行内积完全相同」的图形版。右图是另一个常被混着说的现象：距离越远，加了 RoPE 的分数与不加位置编码的原始内积的相关性在所采样的区间整体降低，但并不单调（4000 对随机向量实测，$\Delta=1$ 时 0.97、$\Delta=4096$ 时接近 0；对独立各向同性输入，虚线为 $\frac{2}{d}\sum_i\cos(\Delta\theta_i)$，d 是特征维度、Δ 是相对距离）。注意衰减的是这个相关性，不是内积本身的大小——旋转保模长，本文独立各向同性高斯向量的内积分布与距离无关（实验三会验证这一点）。

完整脚本见 code/rope_minimal.py。

4.1 顺手把两种外推方法也实现掉

3.3 节说过 PI 和 NTK-aware「改的都是同一组 $\theta_i$」。这句话用代码写出来就是两个函数，各三行——把它们并排放一起，两种方法的取舍立刻就清楚了。

def theta_original(dim=64, base=10000.0):
    return 1.0 / (base ** (np.arange(0, dim, 2) / dim))


def theta_pi(dim=64, base=10000.0, scale=8):
    """Position Interpolation：位置索引整体除以 scale。
    等价于把每个 theta_i 都乘上 1/scale——所有频率被同等减速。"""
    return theta_original(dim, base) / scale


def theta_ntk(dim=64, base=10000.0, scale=8):
    """NTK-aware：把 base 放大，让减速量随维度索引递增（频率越低减速越多）。"""
    new_base = base * scale ** (dim / (dim - 2))
    return theta_original(dim, new_base)

跑 code/rope_extrapolate.py 看它们对各维度做了什么（外推 8 倍）：

  维度 i      原始         PI(÷8)      NTK-aware    PI减速比  NTK减速比
     0   1.000e+00   1.250e-01   1.000e+00      8.00x     1.00x
     4   3.162e-01   3.953e-02   2.418e-01      8.00x     1.31x
     8   1.000e-01   1.250e-02   5.847e-02      8.00x     1.71x
    16   1.000e-02   1.250e-03   3.419e-03      8.00x     2.92x
    31   1.334e-04   1.667e-05   1.667e-05      8.00x     8.00x

PI 那一列减速比是齐刷刷的 8.00x，从最高频到最低频一视同仁。NTK 那一列从 1.00x 单调爬到 8.00x——最高频维度完全没动，最低频维度承担了全部的 8 倍压缩。

这张图要看什么：左图把上表画成曲线——PI 是一条平线（每维都慢 8 倍），NTK 从最高频的 1.00x 单调爬到最低频的 8.00x，两条线夹出的面积就是「NTK 少减的部分」；右图是代价，PI 把每一维的近距离分辨率都砍到 12.5%，NTK 保住高频 100% 的代价是中低频一路掉到 34.2%。两种方法改的是同一组 $\theta_i$，区别只在把这次外推的代价放在哪个频段。

代价体现在近距离分辨率上：

  维度 i     原始        PI          NTK        PI保留   NTK保留
     0   1.000000  0.125000  1.000000    12.5%  100.0%
     4   0.316228  0.039528  0.241809    12.5%   76.5%
     8   0.100000  0.012500  0.058472    12.5%   58.5%
    16   0.010000  0.001250  0.003419    12.5%   34.2%

PI 让每维相邻位置的相位增量降到 12.5%；这可能影响局部分辨，但不等于任务指标必然下降，微调与模型适应会改变实际结果。NTK 在第 0 维保住了 100%，代价是第 16 维只剩 34.2%。

还有一个容易看漏的点。NTK 外推后，最高频维度在位置 32767 处的相位是 32767 弧度，仍然远超训练时见过的 4095：

  训练时见过的最大相位（位置 4095）     4095.0 rad
  原始（直接外推）           位置 32767    32767.0 rad   ← 超出 8.0 倍
  PI                 位置 32767     4095.9 rad   ← 接近但略超训练端点 4095
  NTK-aware          位置 32767    32767.0 rad   ← 超出 8.0 倍

这是该缩放的设计取舍：最高频保留原始步长，低频逐渐压缩。未展开相位超过训练最大值不等于必然失效，相位本身具有周期性；只有最低频获得完整 8 倍压缩，不能据此保证所有频段已回到训练范围或质量安全。PI 则把所有维度的相位范围压到接近训练区间，代价是相邻位置相位步长整体下降；精确端点对齐还需使用端点跨度比。YaRN 的贡献就是把「哪些频段该压、压多少」变成一组可调策略，而不是这两种极端。

05. 工业级实现对照

参考实现（以 2026-08 的 main 分支为准，上游重构频繁）：

• huggingface/transformers → modeling_llama.py：LlamaRotaryEmbedding、rotate_half、apply_rotary_pos_emb
• huggingface/diffusers → embeddings.py：get_3d_rotary_pos_embed

生产代码和上面的最小实现有四处不一样，每一处都有理由。

5.1 配对方式不同，而且不能混用

这是最容易踩的坑。本文把第 $2i$ 和 $2i+1$ 维配成一个复数（奇偶交错），而 HF 的写法是把第 $i$ 和第 $i + d/2$ 维配成一对（前后半切分）：

def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

q_embed = (q * cos) + (rotate_half(q) * sin)

配合 emb = torch.cat((freqs, freqs), dim=-1)（cos/sin 表沿最后一维复制两遍），展开就是 $x_i \cos - x_{i+d/2} \sin$ 和 $x_{i+d/2} \cos + x_i \sin$——和 3.3 节的旋转矩阵一模一样，只是谁跟谁配对换了。

那两者是等价的吗？我写了个脚本实测（code/rope_pairing.py）：

=== 1. 两种配对都满足相对位置不变性 ===
  奇偶交错: 间距恒为 3 时内积 -13.049657，四个绝对位置间的最大波动 3.55e-15
  前后半切分: 间距恒为 3 时内积 -7.706177，四个绝对位置间的最大波动 1.78e-15

=== 2. 同一份权重下，两种配对给出的分数不同 ===
  奇偶交错   -13.049657
  前后半切分 -7.706177
  差值       5.343480  ← 不为 0，实现和权重必须配套

=== 3. 重排维度后两者完全一致 ===
  奇偶交错(原始排布)      -13.049657
  前后半切分(重排后权重)  -13.049657
  差值                    0.00e+00  ← 只差一个固定的维度置换

三段结论连起来才是完整答案：作为位置编码方案两者等价（都满足相对位置不变性，波动 1e-15 是纯浮点误差），但对同一份权重不可互换（差了 5.34），因为它们只差一个固定的维度置换（重排后差值精确为 0）。

实践意义：移植权重时如果配对方式对不上，模型不会报错，只会变笨。这种 bug 极难查——loss 能降，输出也通顺，就是长距离能力莫名其妙地差。判断办法很简单：看它的 cos/sin 表是 cat((freqs, freqs)) 还是 repeat_interleave(freqs, 2)，前者是前后半切分，后者是奇偶交错。

5.2 inv_freq 存成 buffer，而且不进 checkpoint

self.inv_freq = nn.Buffer(inv_freq, persistent=False)
self.original_inv_freq = nn.Buffer(inv_freq.clone(), persistent=False)

缓存是为了省掉每步的幂运算。而 persistent=False 这个细节更值得说：它让 inv_freq 不写进 state_dict。好处是改 rope_theta 做长度外推时，不必担心 checkpoint 里存着旧频率把新配置覆盖掉——频率永远由 config 现算。留一份 original_inv_freq 是给 dynamic NTK 用的：那类方法会随序列长度临时改频率，改完得能退回原值。

5.3 接受 position_ids 而不是内部 arange

def forward(self, x, position_ids):
    inv_freq_expanded = self.inv_freq[None, :, None].expand(position_ids.shape[0], -1, 1)
    position_ids_expanded = position_ids[:, None, :].float()

最小实现里位置就是 arange(seq_len)，生产上不行，至少三个场景要显式传：

• KV cache 增量解码：第 101 个 token 单独前向，它的位置是 100 而不是 0
• 序列打包（packing）：多条短样本拼成一条长序列喂进去，每条的位置要各自从 0 起算
• 多模态与视频：图像 token 和文本 token 的位置编号规则不同，位置得由外部算好

5.4 强制 fp32 算 cos/sin

with maybe_autocast(device_type=device_type, enabled=False):
    freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
    emb = torch.cat((freqs, freqs), dim=-1)
    cos = emb.cos() * self.attention_scaling

注意它显式关掉了 autocast。原因是 $m\theta_i$ 在长序列下会长到几万，bf16 只有 8 位有效尾数，在这个量级上间隔已经大于 1——相邻位置会算出同一个角度，位置信息直接丢失。所以三角函数必须在 fp32 里算完，再转回低精度参与矩阵乘。

5.5 读陌生代码库时的四个问题

接手一个没见过的模型，想搞清它的 RoPE 怎么配的，按这个顺序问四个问题最快，也刚好对应上面四小节：

1. cos/sin 表是怎么铺开的？ 搜 cat 还是 repeat_interleave——torch.cat((freqs, freqs)) 是前后半切分，repeat_interleave(freqs, 2) 是奇偶交错。这决定了权重能不能和别的实现互换。
2. inv_freq 从哪来、会不会变？ 搜 inv_freq 的赋值点。如果只在 __init__ 里算一次，就是静态 RoPE；如果 forward 里还会重算，大概是某种 dynamic NTK，那么推理时的行为会随序列长度变。
3. position_ids 谁给的？ 如果是内部 arange，还要核对是否加了 cache offset 或传入各段位置；从零重启才会破坏 KV cache 增量解码，多模态场景基本要改。
4. 三角函数在什么精度下算的？ 搜附近有没有 float()、autocast、enabled=False。没有强制 fp32 的话，长序列下有精度隐患（5.4 节）。

这四个问题都能靠 grep 在几分钟内回答，比通读整个 attention 实现快得多。

5.6 视频侧：三个轴分走维度

视频 DiT 的 token 有时间、高、宽三个坐标，做法是把 head_dim 切成三份，每份跑一遍 1D RoPE 再拼起来。diffusers 里的分配是：

dim_t = embed_dim // 4        # 时间轴，占 25%
dim_h = embed_dim // 8 * 3    # 高，占 37.5%
dim_w = embed_dim // 8 * 3    # 宽，占 37.5%

这个 1:1.5:1.5 的分配是 CogVideoX 相关实现的一种维度预算。不能从空间占 75%、时间占 25% 推断哪个轴更先失败或失败症状；需结合训练范围、每轴频率和实际消融。详见3D RoPE。

06. 代价与边界

RoPE 不是免费的。

它对固定内容向量只保证位置项的相对位置性质，不保证外推可用。 3.1 节那个性质对任意 $m$、$n$ 都成立，包括训练时没见过的位置。但「公式成立」和「效果好」是两件事：模型只在 $m\theta_i$ 落在见过的区间里学过怎么解释这些相位，超出去就是在外推一个没有监督过的函数。这也是为什么直接放长序列会崩，得配 PI 或 YaRN 这类手段。

它引入了远程衰减，而且衰减偏偏打在你最需要的地方。 这条值得展开，因为流传的说法不太准确。

「RoPE 自带远程衰减」几乎成了共识，但实测下来它只对一部分输入成立。我分两种情况量了一遍（code/rope_decay.py 第 2 段）。

先看彼此无关的随机 $q$、$k$：

不加 RoPE            6.2076   (基线)
相对距离     1       6.3182   (101.8% of 基线)
相对距离    64       6.2415   (100.5% of 基线)
相对距离  1024       6.2583   (100.8% of 基线)
相对距离  4000       6.4830   (104.4% of 基线)

完全没有衰减。 道理也简单：$R_m$ 是正交变换，对本文独立、各向同性高斯向量，不改变其内积的统计分布；任意非各向同性输入并无这个保证，转多少度都一样。

再看 $q = k$（一对高度相似的 token，比如相邻帧里的同一个主体）：

全 1 向量（gap=0 时内积 64.00）:
  相对距离     1      +61.8337   ( 96.6% of gap=0)
  相对距离    16      +38.7456   ( 60.5% of gap=0)
  相对距离    64      +27.8122   ( 43.5% of gap=0)
  相对距离  1024      +16.8272   ( 26.3% of gap=0)
  相对距离  4000       +0.0989   (  0.2% of gap=0)

衰减掉了 99.8%。 换个随机向量做同样的实验，结论一致（50.96 → 0.97）。

这两组实验只展示了各向同性独立输入与所选相似输入的差别。RoPE 的相位加权和依赖内容在各频段的分布；既不能推出“只有相似 token 才受影响”，也不能据此将真实视频的身份漂移归因于 RoPE，后者需要固定模型的消融实验。

顺带一个细节：衰减不是单调的。上面那个随机向量在 gap=256 处回升到 51.3%，比 gap=64 的 40.9% 还高。它是震荡衰减，因为各频率分量的相位在不同距离上会偶然重新对齐。所以拿单个距离点的数值下结论是不可靠的。

它无法表达绝对位置。 「这是第一帧」这类信息 RoPE 表达不了，因为它被设计成只看差值。需要绝对锚点的任务（比如首帧条件生成、指定时间戳编辑）得靠别的机制补，通常是额外的条件注入而不是位置编码。

算力和显存几乎不是代价。 这点顺手算清楚，免得担错心。拿 01 节那个 90000 token 的 720p 配置、假设 40 层 24 头、head_dim 64：

• 旋转本身约 33.2 GFLOPs，而同规模的注意力矩阵乘是 995.3 TFLOPs——RoPE 占 0.0033%，四舍五入就是零
• cos/sin 表按半维存 fp32 是 22.0 MB，铺满 dim 是 43.9 MB（bf16 减半）。这个表在各层频率配置相同时可共享，算一次就够，相比激活值的开销可以忽略

所以选不选 RoPE 从来不是性能问题，纯粹是「要不要相对位置性和外推能力」的取舍。真正吃资源的是 5.4 节那个 fp32 约束——它强迫这一小段计算不能跟着整个网络一起降精度。

什么时候不该用：序列长度固定且很短、任务本身依赖绝对位置、或者模型是纯双向编码器且已有成熟的绝对位置方案时，换 RoPE 收益不大。RoPE 的价值集中在「训练短、推理长」这个场景上。

07. 经典论文脉络

五篇串起来就是位置编码近几年的主线：

• ALiBi（arXiv:2108.12409，2021-08）不做旋转，直接给注意力分数加一个随距离线性递减的偏置。外推能力很强，但把「远处不重要」写死成了先验，表达力受限。它是 RoPE 的主要对照组。
• RoFormer（arXiv:2104.09864，2021-04）RoPE 的出处。贡献是用复数旋转同时拿到相对位置性和乘法式注入；留下的问题是没解决外推，那时也还没人需要 128K 上下文。
• A Length-Extrapolatable Transformer（arXiv:2212.10554，2022-12）把「外推」这件事本身形式化，提出用注意力分辨率来度量，给后续方法提供了评价标准而不只是刷分。
• Position Interpolation（arXiv:2306.15595，2023-06）最简单有效的一招：把位置索引整体除以缩放因子，让新长度挤回训练区间。代价是所有频率一起被压缩，近距离分辨率跟着下降。
• YaRN（arXiv:2309.00071，2023-08）按频率分段处理——高频维度少动或不动以保住近距离分辨率，低频维度多压缩以覆盖新长度，再配一个注意力温度修正。这是目前长度外推的实用基线。

主线很清楚：RoPE 定下了「相对位置靠旋转」的框架，后面几年都在调同一组 $\theta_i$。理解了 3.3 节的频率分配，这些方法的差别只是一句话。

7.1 视频侧这条支线

上面五篇都出自语言模型。视觉生成这边是另一条时间线，起点晚但收敛很快：

• DiT（arXiv:2212.09748，2022-12）把扩散模型的骨干从 U-Net 换成 Transformer，位置信息用的还是固定的 sincos 绝对嵌入。它证明了 Transformer 在扩散上能 scale，但也把绝对位置嵌入那套外推问题一起继承了过来。
• CogVideoX（arXiv:2408.06072，2024-08）视频 DiT 转向 3D RoPE 的代表作。diffusers 里那个 get_3d_rotary_pos_embed（5.6 节引的那段维度分配代码）就是为它写的。
• HunyuanVideo（arXiv:2412.03603，2024-12）与 Wan（arXiv:2503.20314，2025-03）两个开源大模型延续了 3D RoPE 的路线。到这一步，「时间/高/宽三轴各分走一段 head_dim」基本成了视频 DiT 的默认选择。

值得注意的是这条支线继承了语言侧的全部外推工具：PI 和 NTK 的思路可以直接套到某一个轴上——比如只对空间两轴做插值来支持更高分辨率，时间轴不动。反过来，语言侧没有的新问题也在这里出现：三个轴的维度预算怎么分、变分辨率训练时坐标怎么归一化。这些留给「视频 DiT 里的 3D RoPE 与分辨率外推」。

08. 常见误解

「RoPE 是加在 embedding 上的」。 不是。它作用在每一层注意力的 $q$、$k$ 上，而且是乘法（旋转）。$v$ 在本文与常用实现里不旋转；存在其他变体，不能把所有旋转 value 的方案都判为错误。把它当成输入层的一次性操作，是最常见的理解错误。

「调大 rope_theta 就能延长上下文」。 只对了一半。调大 base 会让除最高频 $\theta_0=1$ 之外的频率变慢，能塞进更长序列，但不重新训练的话，高频维度的近距离分辨率会一起掉，表现为长文本能读了、短距离推理反而变差。YaRN 之所以要分频段处理，就是为了避免这个副作用。

「相对位置不变性意味着模型对绝对位置完全无感」。 不成立。注意力有 causal mask，第 0 个 token 只能看到自己，第 100 个能看到 100 个——这种「可见范围」本身就泄露了绝对位置。研究里管这叫隐式位置信息。

「间距相同分数就一定相同」。 还必须固定同一对内容向量 q、k；真实模型里的它们也会随上下文变化。真实模型里中间还夹着 softmax 归一化和 causal mask，同一对 token 在不同上下文里拿到的注意力权重是不同的。04 节验证的是旋转这一步的数学性质，不是整个注意力层的行为。

「RoPE 自带远程衰减」需要限定输入与统计口径。 独立各向同性高斯输入的内积分布不随旋转改变；所选相似向量可能在某些距离明显降低、在另一些距离回升。示例不是普遍衰减定律，也不能单独证明长视频质量的因果关系。

09. 动手验证

四个五分钟能跑完的小实验，用到的脚本全文都在文末附录，只依赖 NumPy，复制存成同名文件就能跑。另外跑一遍 make_figures.py 可以把本文三张配图重新生成（这个需要 matplotlib，其余脚本只要 NumPy）。

实验一：看外推在哪一步崩。 跑 rope_minimal.py，把 seq_len 从 8 改成 4096，然后对比位置 (0, 1) 和位置 (4000, 4001) 的内积。预期结果是两者仍然完全相等——因为数学性质对任意位置都成立。这恰好说明：外推失败不是公式失效，而是模型没在那个相位区间学过。这是本文最容易被误解的一点，值得亲手确认一次。

实验二：验证配对方式不能混用。 跑 rope_pairing.py，观察第 2、3 段输出。预期看到同一份权重下两种配对差 5.34，而重排维度后差值精确为 0。把 dim 改成 128 再跑一次，结论不变——这是结构性的，与维度大小无关。

实验三：亲手推翻一个常识。 跑 rope_decay.py，重点看 2a 和 2b 的差别。预期看到随机 $q$、$k$ 不衰减、$q = k$ 时衰减 99% 以上。然后把 2b 里的向量换成两个弱相关的向量试试（比如 k = 0.5 * q + 0.5 * rng.standard_normal(dim)），看衰减幅度怎么随相似度变化——这是理解 06 节那笔账的最快路径。

实验四：给自己的模型算外推预算。 跑 rope_extrapolate.py，把 SCALE 从 8 改成你实际需要的倍数（比如 480p→720p 大约是 2.25 倍 token），看 PI 和 NTK 各自在近距离分辨率上要付多少。预期规律是 PI 的保留率恒等于 $1/\text{scale}$，而 NTK 的保留率在高频维度接近 1、随维度下降。这些频率变化只能说明位置编码的几何差异，不能直接预测生成质量。如果任务依赖纹理、口型等近距离细节，应在目标分辨率上比较 PI、NTK 与必要的适配训练，结合质量和长距离一致性决定方案。

10. 延伸阅读

读这篇之前建议先看：

• 自注意力机制的计算与显存账本——RoPE 作用在 $q$、$k$ 上，得先清楚它们在注意力里的位置

读完这篇可以继续看：

• 视频 DiT 里的 3D RoPE 与分辨率外推——5.6 节那个 1:1.5:1.5 的示例维度分配，以及分辨率和帧数变化如何分别影响各轴相位
• FlashAttention 为什么不需要存下注意力矩阵——RoPE 之后紧接着的那一步计算

附录：完整代码

09 节用到的脚本全文如下（make_figures.py、rope_decay.py、rope_minimal.py、rope_extrapolate.py、rope_pairing.py）。复制到本地存成同名文件，按各脚本开头的依赖说明准备环境后即可运行。

make_figures.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""生成「旋转位置编码 RoPE」的三张解释图。

数值全部来自同目录的三个脚本（rope_minimal / rope_extrapolate / rope_decay），
这里只负责画——改了那边这里要重跑，免得图和正文数字打架。

三张图分别回答：
  1. 32 个频率是怎么分层铺在 64 维上的（高频相位在近距离变化更快）
  2. PI 与 NTK-aware 把这 8 倍外推的代价分别摊到了哪些频段
  3. 「相对位置不变」和「随距离衰减」在数值上长什么样

只依赖 numpy + matplotlib。跑法：python make_figures.py
"""

import textwrap
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

import rope_decay as RD
import rope_extrapolate as RE

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "figures"
OUT.mkdir(parents=True, exist_ok=True)

plt.rcParams.update({
    "font.sans-serif": ["PingFang SC", "Arial Unicode MS", "DejaVu Sans"],
    "axes.unicode_minus": False,
    "figure.dpi": 160,
    "savefig.bbox": "tight",
})

INK = "#1f2937"
MUTE = "#6b7280"
C_ORIG = "#2f6fb0"    # 原始 RoPE：蓝
C_PI = "#d1495b"      # PI：红
C_NTK = "#2f9e6f"     # NTK-aware：绿
C_BAD = "#d1495b"
C_OK = "#2f9e6f"


def style(ax, title, xlabel=None, ylabel=None):
    ax.set_title(title, fontsize=11.5, color=INK, pad=10, loc="left")
    if xlabel:
        ax.set_xlabel(xlabel, fontsize=10, color=MUTE)
    if ylabel:
        ax.set_ylabel(ylabel, fontsize=10, color=MUTE)
    ax.tick_params(colors=MUTE, labelsize=9)
    for s in ("top", "right"):
        ax.spines[s].set_visible(False)
    for s in ("left", "bottom"):
        ax.spines[s].set_color("#d1d5db")
    ax.grid(alpha=0.25, linewidth=0.6)
    ax.set_axisbelow(True)


def footer(fig, text, width=118):
    """把「这张图要看什么」放到坐标轴下方。

    必须放在 y<0 的位置：bbox_inches="tight" 会把负坐标的 artist 一起收进来，
    放在 0~0.05 之间的话会和 x 轴标签叠在一起。
    """
    wrapped = "\n".join(textwrap.wrap(text, width=width))
    fig.text(0.012, -0.13, wrapped, fontsize=8.5, color=MUTE,
             va="top", ha="left", linespacing=1.6)


# ──────────────────────────────────────────────────────────────────────
# 图 1：32 个频率怎么分层铺在 64 维上
# ──────────────────────────────────────────────────────────────────────
def fig_freq_layers():
    dim = RE.DIM
    th = RE.theta_original()                      # [dim/2]
    pos_short = np.arange(0, 256)
    pos_long = np.arange(0, 32768, 8)

    cos_short = np.cos(np.outer(pos_short, th))   # [256, dim/2]
    phase_long = np.outer(pos_long, th)           # 累积相位（不取模）

    fig, axes = plt.subplots(1, 2, figsize=(12.8, 4.4))
    fig.subplots_adjust(wspace=0.42)

    ax = axes[0]
    im = ax.imshow(cos_short.T, aspect="auto", origin="lower",
                   cmap="RdBu_r", vmin=-1, vmax=1,
                   extent=[0, len(pos_short), -0.5, len(th) - 0.5])
    ax.set_yticks([0, 8, 16, 24, 31])
    cb = fig.colorbar(im, ax=ax, fraction=0.035, pad=0.02)
    cb.ax.tick_params(labelsize=8, colors=MUTE)
    cb.set_label(r"$\cos(m\theta_i)$", fontsize=9, color=MUTE)
    style(ax, "近距离（0~255）：前 13 组完成至少一圈", "位置 m", "维度组 i")
    # 在 256 个位置内转满一整圈（2π）的维度组数
    n_moving = int((pos_short[-1] * th / (2 * np.pi) >= 1).sum())
    ax.annotate(f"256 长度内转满一整圈的只有前 {n_moving} 组",
                xy=(128, n_moving - 1), xytext=(12, n_moving + 7),
                fontsize=8.5, color="white", fontweight="bold",
                arrowprops=dict(arrowstyle="->", color=MUTE, lw=0.9))

    ax = axes[1]
    for i, c in zip([0, 4, 8, 16, 24, 31],
                    ["#1f6feb", "#2f9e6f", "#e0a03c",
                     "#d1495b", "#8b5cf6", "#6b7280"]):
        ax.plot(pos_long, phase_long[:, i], lw=1.8, color=c, label=f"i={i}")
    trained = (RE.TRAIN_LEN - 1) * th[0]
    ax.axhline(trained, color=INK, ls="--", lw=1.2)
    ax.text(600, trained * 1.35, f"最高频训练最大相位 {trained:.0f} rad",
            fontsize=8.5, color=INK)
    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.legend(fontsize=8.5, frameon=False, loc="upper left", ncol=2)
    style(ax, "远距离（到 32768）：各频段的累积相位",
          "位置 m（对数轴）", r"累积相位 $m\theta_i$（rad，对数轴）")

    fig.suptitle("图 1　RoPE 把 64 维切成 32 组，从「每步 1 弧度」一路排到「每步 1e-4 弧度」",
                 fontsize=12, color=INK, x=0.012, ha="left", y=1.02)
    footer(fig, "左图：0 到 255 的相位差中，只有前 13 组超过一整圈；低频也有非零相位变化。"
                "右图：不同频率提供不同尺度，不是互斥的距离分工。虚线只代表最高频的训练相位上限；"
                "每个频段的训练相位上限都不同。")
    fig.savefig(OUT / "rope_freq_layers.png")
    plt.close(fig)
    print(f"[图1] 相位变化>=2π 的维度组数（0→255）: {n_moving}/{len(th)}")


# ──────────────────────────────────────────────────────────────────────
# 图 2：PI 与 NTK-aware 把 8 倍外推的代价摊在哪
# ──────────────────────────────────────────────────────────────────────
def fig_extrapolation():
    t_orig, t_pi, t_ntk = (RE.theta_original(), RE.theta_pi(), RE.theta_ntk())
    idx = np.arange(len(t_orig))
    slow_pi = t_orig / t_pi
    slow_ntk = t_orig / t_ntk
    keep_pi = t_pi / t_orig * 100
    keep_ntk = t_ntk / t_orig * 100

    fig, axes = plt.subplots(1, 2, figsize=(11.6, 4.3))

    ax = axes[0]
    ax.plot(idx, slow_pi, color=C_PI, lw=2.4, label=r"PI：各维一律减速 8.00x")
    ax.plot(idx, slow_ntk, color=C_NTK, lw=2.0, label=r"NTK-aware：1.00x $\to$ 8.00x")
    ax.fill_between(idx, slow_ntk, slow_pi, color=C_PI, alpha=0.10)
    ax.set_yscale("log")
    ax.set_yticks([1, 2, 4, 8])
    ax.yaxis.set_minor_formatter(plt.NullFormatter())
    ax.get_yaxis().set_major_formatter(plt.FuncFormatter(lambda v, _: f"{v:g}x"))
    ax.legend(fontsize=9, frameon=False, loc="center right")
    style(ax, "谁被减速了多少（对数轴）", "维度组 i",
          r"$\theta_i^{\mathrm{orig}} / \theta_i^{\mathrm{new}}$")

    ax = axes[1]
    ax.plot(idx, keep_pi, color=C_PI, lw=2.4, label="PI：全线只剩 12.5%")
    ax.plot(idx, keep_ntk, color=C_NTK, lw=2.0,
            label=rf"NTK：{keep_ntk[0]:.0f}% $\to$ {keep_ntk[16]:.1f}%（i=16）")
    ax.fill_between(idx, keep_pi, keep_ntk, color=C_NTK, alpha=0.10)
    ax.axhline(100, color=MUTE, ls=":", lw=1.0)
    ax.set_yscale("log")
    ax.set_yticks([10, 20, 50, 100])
    ax.get_yaxis().set_major_formatter(plt.FuncFormatter(lambda v, _: f"{v:g}%"))
    ax.legend(fontsize=9, frameon=False, loc="lower left")
    style(ax, "代价：相邻 token 的相位分辨率还剩多少", "维度组 i",
          r"$\theta_i^{\mathrm{new}} / \theta_i^{\mathrm{orig}}$")

    fig.suptitle("图 2　PI 把减速平均分给每一维，NTK 从高频不压缩逐渐过渡到低频压缩 8 倍",
                 fontsize=12, color=INK, x=0.012, ha="left", y=1.02)
    footer(fig, "要看什么：左图 PI 是一条平线（每维都慢 8 倍），NTK 从最高频的 1.00x 单调爬到"
                "最低频的 8.00x。右图是这个分配的代价——PI 把每一维的近距离分辨率都砍到 12.5%，"
                "这只是相位步长，不能直接预测任务质量；NTK 保住高频（100%），最低频与 PI 同为 12.5%。"
                "两者改的是同一组 theta，区别只在把代价放在哪个频段。")
    fig.savefig(OUT / "rope_extrapolation.png")
    plt.close(fig)
    print(f"[图2] NTK 减速比: i=0 {slow_ntk[0]:.2f}x -> i=31 {slow_ntk[-1]:.2f}x；"
          f"分辨率保留 i=16: PI {keep_pi[16]:.1f}% / NTK {keep_ntk[16]:.1f}%")


# ──────────────────────────────────────────────────────────────────────
# 图 3：相对位置不变 + 随距离衰减
# ──────────────────────────────────────────────────────────────────────
def fig_relativity():
    dim, seq_len = 64, 8192
    cos, sin = RD.build_rope_cache(seq_len, dim)
    rng = np.random.default_rng(0)

    # ── 左：同一对 q,k，只平移绝对位置，内积应当逐点不变 ──
    q = rng.standard_normal(dim)
    k = rng.standard_normal(dim)
    gaps = np.arange(0, 512)
    starts = [0, 500, 2000]
    curves = {}
    for s in starts:
        qr = RD.apply_rope(q[None, :], cos[s:s + 1], sin[s:s + 1])[0]
        kr_all = RD.apply_rope(
            np.repeat(k[None, :], len(gaps), 0),
            cos[s:s + len(gaps)], sin[s:s + len(gaps)])
        curves[s] = kr_all @ qr

    # ── 右：距离越远，RoPE 把「无位置编码的原始内积」改写得多彻底 ──
    # 独立各向同性高斯输入：corr(q^T R_gap k, q^T k) = (2/dim) Σ_i cos(gap θ_i)。
    # 注意：这是「与原始内积的相关性」衰减，不是内积的绝对值变小——
    # 独立各向同性高斯输入的内积分布与距离无关（RMS 为 sqrt(dim)）。
    Q = rng.standard_normal((4000, dim))
    K = rng.standard_normal((4000, dim))
    s0 = np.einsum("ij,ij->i", Q, K)
    th = 1.0 / (10000.0 ** (np.arange(0, dim, 2) / dim))
    probe = np.unique(np.round(np.geomspace(1, 4096, 60)).astype(int))
    emp, theo = [], []
    for g in probe:
        qr = RD.apply_rope(Q, cos[0:1], sin[0:1])
        kr = RD.apply_rope(K, cos[g:g + 1], sin[g:g + 1])
        sg = np.einsum("ij,ij->i", qr, kr)
        emp.append(float(np.corrcoef(sg, s0)[0, 1]))
        theo.append(float(2.0 / dim * np.sum(np.cos(g * th))))

    fig, axes = plt.subplots(1, 2, figsize=(11.6, 4.3))

    ax = axes[0]
    cols = {0: C_ORIG, 500: C_NTK, 2000: "#e0a03c"}
    for s in starts:
        ax.plot(gaps, curves[s], lw=1.6, color=cols[s], label=f"绝对起点 {s}")
    ax.legend(fontsize=9, frameon=False, loc="upper right")
    style(ax, "同一对向量、同一组间距，换绝对位置后曲线完全重合",
          "相对距离 (n-m)", r"$q_m^\top R_{n-m} k_n$")

    ax = axes[1]
    ax.plot(probe, emp, marker="o", ms=3.5, lw=1.8, color=C_ORIG,
            label="实测（4000 对随机向量）")
    ax.plot(probe, theo, lw=1.4, ls="--", color=C_BAD,
            label=r"理论 $\frac{2}{D}\sum_i\cos(\Delta\theta_i)$")
    ax.axhline(0, color=MUTE, ls=":", lw=1.0)
    ax.set_xscale("log")
    ax.legend(fontsize=9, frameon=False, loc="upper right")
    style(ax, "独立高斯输入：与原始内积的相关性随距离振荡",
          "相对距离（对数轴）", r"$\mathrm{corr}(q^\top R_d k, q^\top k)$")

    fig.suptitle("图 3　固定内容的相对位置不变性，与高斯输入下的分数相关性",
                 fontsize=12, color=INK, x=0.012, ha="left", y=1.02)
    max_dev = max(float(np.abs(curves[s] - curves[0]).max()) for s in starts)
    footer(fig, f"左图：固定 q/k 的三条曲线在浮点误差内重合，最大偏差 {max_dev:.1e}，不是逐比特相等。"
                "右图：D=64、Δ 为间距，理论相关性是有限余弦和；它会振荡，不保证单调或在无穷远趋于零。"
                "本图的独立各向同性高斯输入下，内积分布不随旋转改变。")
    fig.savefig(OUT / "rope_relativity.png")
    plt.close(fig)
    print(f"[图3] 不同绝对起点的最大逐点偏差: {max_dev:.2e}；"
          f"corr: d=1 {emp[0]:.3f} → d={probe[-1]} {emp[-1]:.3f}")


def main():
    fig_freq_layers()
    fig_extrapolation()
    fig_relativity()
    print(f"[OK] 三张图已写入 {OUT}")
    for f in sorted(OUT.glob("*.png")):
        print(f"  {f.name}  {f.stat().st_size / 1024:.0f} KB")


if __name__ == "__main__":
    main()

rope_decay.py

"""观测 RoPE 的两个常被口头断言、但很少有人实测的性质。

1. 旋转不改变模长（对比加法式绝对位置嵌入）
2. 所选输入的内积随距离变化；并非普遍衰减定律

运行 `python rope_decay.py` 复现。
"""

import numpy as np


def build_rope_cache(seq_len, dim, base=10000.0):
    inv_freq = 1.0 / (base ** (np.arange(0, dim, 2) / dim))
    freqs = np.outer(np.arange(seq_len), inv_freq)
    return np.cos(freqs), np.sin(freqs)


def apply_rope(x, cos, sin):
    out = np.empty_like(x)
    out[..., 0::2] = x[..., 0::2] * cos - x[..., 1::2] * sin
    out[..., 1::2] = x[..., 0::2] * sin + x[..., 1::2] * cos
    return out


def demo_norm_preserved(rng, dim=64, seq_len=512):
    """旋转保模长，加法不保。"""
    cos, sin = build_rope_cache(seq_len, dim)
    x = rng.standard_normal(dim)
    print("=== 1. 模长是否被位置编码改变 ===")
    print(f"  原始向量模长                {np.linalg.norm(x):.6f}")

    norms = [np.linalg.norm(apply_rope(x[None, :], cos[p:p + 1], sin[p:p + 1]))
             for p in (0, 1, 100, 511)]
    print(f"  RoPE 后（位置 0/1/100/511） "
          f"{'  '.join(f'{n:.6f}' for n in norms)}")
    print(f"  跨位置模长极差              {max(norms) - min(norms):.2e}")

    # 加法式绝对位置嵌入：同尺度的可学习向量直接相加
    pe = rng.standard_normal((seq_len, dim))
    add_norms = [np.linalg.norm(x + pe[p]) for p in (0, 1, 100, 511)]
    print(f"  加法式 PE 后                "
          f"{'  '.join(f'{n:.6f}' for n in add_norms)}")
    print(f"  跨位置模长极差              "
          f"{max(add_norms) - min(add_norms):.4f}  ← 同一个 token 的模长被位置改写了")


def _score(cos, sin, q, k, gap):
    qr = apply_rope(q[None, :], cos[0:1], sin[0:1])
    kr = apply_rope(k[None, :], cos[gap:gap + 1], sin[gap:gap + 1])
    return float(qr[0] @ kr[0])


GAPS = (1, 4, 16, 64, 256, 1024, 4000)


def demo_long_range_decay(rng, dim=64, seq_len=4096, trials=2000):
    """「RoPE 自带远程衰减」到底成不成立——分两种输入分布看。

    结论：本文独立各向同性高斯 q/k 的内积分布不变；所选 q=k 示例会振荡。
    """
    cos, sin = build_rope_cache(seq_len, dim)

    print("\n=== 2a. 随机无关的 q/k：取 |内积| 均值 ===")
    baseline = float(np.mean([abs(float(rng.standard_normal(dim) @ rng.standard_normal(dim)))
                              for _ in range(trials)]))
    print(f"  不加 RoPE            {baseline:.4f}   (基线)")
    for gap in GAPS:
        m = float(np.mean([abs(_score(cos, sin, rng.standard_normal(dim),
                                      rng.standard_normal(dim), gap))
                           for _ in range(trials)]))
        print(f"  相对距离 {gap:>5}       {m:.4f}   ({m / baseline * 100:5.1f}% of 基线)")
    print("  → 没有衰减趋势。独立各向同性高斯输入在正交旋转下分布不变。")

    print("\n=== 2b. q = k（高度相似的一对 token）===")
    for name, vec in (("全 1 向量", np.ones(dim)),
                      ("某个随机向量", np.random.default_rng(7).standard_normal(dim))):
        # gap=0 时位置 0 的旋转角为 0，内积就是未加 RoPE 的原值
        head = _score(cos, sin, vec, vec, 0)
        print(f"  {name}（gap=0 时内积 {head:.2f}）:")
        for gap in GAPS:
            s = _score(cos, sin, vec, vec, gap)
            print(f"    相对距离 {gap:>5}     {s:+9.4f}   ({s / head * 100:5.1f}% of gap=0)")
    print("  → 这些采样点多低于 gap=0，也有回升；不代表任意输入或距离都衰减。")


if __name__ == "__main__":
    demo_norm_preserved(np.random.default_rng(0))
    demo_long_range_decay(np.random.default_rng(1))

rope_minimal.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""RoPE 的最小可运行实现，验证「注意力分数只依赖相对位置」这一核心性质。

只依赖 numpy，直接 python rope_minimal.py 即可运行。
刻意不用 PyTorch：讲原理时框架的抽象反而是噪声，而且 numpy 谁都能跑。
生产实现见文章第 05 节对 transformers 的引用。
"""

import numpy as np


def build_rope_cache(seq_len: int, dim: int, base: float = 10000.0):
    """预计算 RoPE 的 cos/sin 表。

    dim 必须是偶数：每两个相邻维度配成一个复数参与旋转。
    base 就是常说的 rope_theta，调它会按维度改变旋转速度；最高频 i=0 始终为 1。
    """
    assert dim % 2 == 0, f"dim 必须为偶数，收到 {dim}"
    # theta_i = base^(-2i/dim)：i 越大转得越慢，对应越低的频率
    inv_freq = 1.0 / (base ** (np.arange(0, dim, 2) / dim))
    pos = np.arange(seq_len)
    freqs = np.outer(pos, inv_freq)              # [seq_len, dim/2]
    return np.cos(freqs), np.sin(freqs)


def apply_rope(x: np.ndarray, cos: np.ndarray, sin: np.ndarray) -> np.ndarray:
    """把旋转作用到 x 上。x: [seq_len, dim]

    最后一维两两配对看成复数 (x_even + i*x_odd)，乘上 e^{i*m*theta}：
        实部 = x_even * cos - x_odd * sin
        虚部 = x_even * sin + x_odd * cos
    """
    x_even, x_odd = x[..., 0::2], x[..., 1::2]
    out = np.empty_like(x)
    out[..., 0::2] = x_even * cos - x_odd * sin
    out[..., 1::2] = x_even * sin + x_odd * cos
    return out


if __name__ == "__main__":
    rng = np.random.default_rng(0)
    seq_len, dim = 8, 64
    cos, sin = build_rope_cache(seq_len, dim)

    print(f"cos/sin 形状: {cos.shape} {sin.shape}")
    print(f"位置 1 的前 4 维 cos: {np.round(cos[1, :4], 4).tolist()}")
    print(f"位置 1 的末 2 维 cos: {np.round(cos[1, -2:], 6).tolist()}")

    # 固定一对向量，只改赋予它们的绝对位置，间距始终为 1
    q = rng.standard_normal(dim)
    k = rng.standard_normal(dim)

    print("\n同一对向量，平移绝对位置（间距恒为 1）:")
    for shift in range(4):
        qr = apply_rope(q[None, :], cos[shift:shift + 1], sin[shift:shift + 1])
        kr = apply_rope(k[None, :], cos[shift + 1:shift + 2],
                        sin[shift + 1:shift + 2])
        print(f"  位置 ({shift}, {shift + 1}) → 内积 {float(qr[0] @ kr[0]):.6f}")

    print("\n对照组，改变间距:")
    for gap in (1, 2, 3):
        qr = apply_rope(q[None, :], cos[0:1], sin[0:1])
        kr = apply_rope(k[None, :], cos[gap:gap + 1], sin[gap:gap + 1])
        print(f"  间距 {gap}      → 内积 {float(qr[0] @ kr[0]):.6f}")

    print(f"\n未加 RoPE 的原始内积: {float(q @ k):.6f}")

rope_extrapolate.py

"""两种长度外推方法（PI / NTK-aware）到底改了什么。

结论先说：它们改的是同一组 theta_i，区别只在「怎么分配这次减速」。
运行 `python rope_extrapolate.py` 复现。
"""

import numpy as np

DIM = 64
BASE = 10000.0
TRAIN_LEN = 4096          # 训练时见过的长度
SCALE = 8                 # 想外推到 8 倍，即 32768


def theta_original(dim=DIM, base=BASE):
    return 1.0 / (base ** (np.arange(0, dim, 2) / dim))


def theta_pi(dim=DIM, base=BASE, scale=SCALE):
    """Position Interpolation：位置索引整体除以 scale。

    等价于把每个 theta_i 都乘上 1/scale——所有频率被同等减速。
    """
    return theta_original(dim, base) / scale


def theta_ntk(dim=DIM, base=BASE, scale=SCALE):
    """NTK-aware：把 base 放大，让减速量随维度索引递增（频率越低减速越多）。

    base' = base * scale^(d/(d-2))，于是高频维度几乎不受影响，
    最低频维度与 PI 同样减速 scale 倍，其余维度的减速更少。
    """
    new_base = base * scale ** (dim / (dim - 2))
    return theta_original(dim, new_base)


def main():
    t_orig, t_pi, t_ntk = theta_original(), theta_pi(), theta_ntk()

    print("=== 1. 三种方案下各维度的旋转速度 theta_i ===")
    print("  维度 i      原始         PI(÷8)      NTK-aware    PI减速比  NTK减速比")
    for i in (0, 1, 4, 8, 16, 24, 31):
        print(f"  {i:>4}   {t_orig[i]:.3e}   {t_pi[i]:.3e}   {t_ntk[i]:.3e}"
              f"   {t_orig[i] / t_pi[i]:>7.2f}x  {t_orig[i] / t_ntk[i]:>7.2f}x")
    print(f"  → PI 对每一维都减速 {t_orig[0] / t_pi[0]:.2f}x（均匀）；"
          f"NTK 从最高频的 {t_orig[0] / t_ntk[0]:.2f}x 递增到最低频的 "
          f"{t_orig[-1] / t_ntk[-1]:.2f}x（越低频压得越狠）")

    print("\n=== 2. 外推到 8 倍长度后，各方案在第 0 维（最高频）的相位 ===")
    target = TRAIN_LEN * SCALE
    trained_max = (TRAIN_LEN - 1) * t_orig[0]
    print(f"  训练时见过的最大相位（位置 {TRAIN_LEN - 1}）  {trained_max:9.1f} rad")
    for name, th in (("原始（直接外推）", t_orig), ("PI", t_pi), ("NTK-aware", t_ntk)):
        v = (target - 1) * th[0]
        flag = ("落在训练区间内" if v <= trained_max else
                f"接近但略超训练端点 {trained_max:g}" if v <= trained_max * 1.01 else
                f"超出 {v / trained_max:.1f} 倍")
        print(f"  {name:<18} 位置 {target - 1}  {v:9.1f} rad   ← {flag}")
    print("  → NTK 在最高频维度上「不作为」是刻意的，它保留高频相位步长，逐渐压缩更低频的相位步长")

    print("\n=== 3. 代价：近距离分辨率还剩多少 ===")
    print("  (相邻位置差 1 时转过的角度，按维度看。越小＝越难区分相邻 token)")
    print("  维度 i     原始        PI          NTK        PI保留   NTK保留")
    for i in (0, 4, 8, 16):
        print(f"  {i:>4}   {t_orig[i]:.6f}  {t_pi[i]:.6f}  {t_ntk[i]:.6f}"
              f"   {t_pi[i] / t_orig[i] * 100:5.1f}%  {t_ntk[i] / t_orig[i] * 100:5.1f}%")
    print(f"  → PI 在每一维都只保留 {t_pi[0] / t_orig[0] * 100:.1f}% 的分辨率；"
          f"NTK 在第 0 维保留 {t_ntk[0] / t_orig[0] * 100:.1f}%、"
          f"第 16 维只剩 {t_ntk[16] / t_orig[16] * 100:.1f}%。")
    print("     两者都在同一组 theta 上做手脚，只是把代价放在了不同频段——"
          "\n     YaRN 进一步把这个分配变成可调的分段策略。")


if __name__ == "__main__":
    main()

rope_pairing.py

"""对照 RoPE 的两种维度配对方式，并验证它们只差一个固定置换。

奇偶交错（本文最小实现）：第 2i 维和第 2i+1 维配成一个复数
前后半切分（HuggingFace）：第 i 维和第 i+d/2 维配成一个复数

结论先说：两者都满足相对位置不变性，但对同一份权重不可互换。
运行 `python rope_pairing.py` 复现下面的数值。
"""

import numpy as np


def rope_interleaved(x, pos, dim, base=10000.0):
    """奇偶交错配对：(x_0,x_1), (x_2,x_3), ... 每对是一个复平面。"""
    inv_freq = 1.0 / (base ** (np.arange(0, dim, 2) / dim))
    ang = pos * inv_freq                       # [dim/2]
    cos, sin = np.cos(ang), np.sin(ang)
    out = np.empty_like(x)
    out[0::2] = x[0::2] * cos - x[1::2] * sin
    out[1::2] = x[0::2] * sin + x[1::2] * cos
    return out


def rope_half_split(x, pos, dim, base=10000.0):
    """前后半切分配对：(x_i, x_{i+d/2})。等价于 HF 的 rotate_half 写法。

    HF 把 cos/sin 表沿最后一维复制两遍（cat((freqs, freqs))），
    于是 q*cos + rotate_half(q)*sin 展开后就是下面两行。
    """
    half = dim // 2
    inv_freq = 1.0 / (base ** (np.arange(0, dim, 2) / dim))
    ang = pos * inv_freq                       # [dim/2]
    cos, sin = np.cos(ang), np.sin(ang)
    x1, x2 = x[:half], x[half:]
    return np.concatenate([x1 * cos - x2 * sin,
                           x2 * cos + x1 * sin])


def interleave_to_half(x):
    """把奇偶交错的排布重排成前后半切分的排布。

    交错版的第 2i 维（复数实部）搬到前半的第 i 位，
    第 2i+1 维（虚部）搬到后半的第 i 位。
    """
    return np.concatenate([x[0::2], x[1::2]])


def attn_score(rope_fn, q, k, m, n, dim):
    return float(rope_fn(q, m, dim) @ rope_fn(k, n, dim))


def main():
    rng = np.random.default_rng(0)
    dim = 64
    q = rng.standard_normal(dim)
    k = rng.standard_normal(dim)

    print("=== 1. 两种配对都满足相对位置不变性 ===")
    for name, fn in (("奇偶交错", rope_interleaved),
                     ("前后半切分", rope_half_split)):
        scores = [attn_score(fn, q, k, m, m + 3, dim) for m in range(4)]
        spread = max(scores) - min(scores)
        print(f"  {name}: 间距恒为 3 时内积 {scores[0]:.6f}，"
              f"四个绝对位置间的最大波动 {spread:.2e}")

    print("\n=== 2. 同一份权重下，两种配对给出的分数不同 ===")
    s_i = attn_score(rope_interleaved, q, k, 0, 3, dim)
    s_h = attn_score(rope_half_split, q, k, 0, 3, dim)
    print(f"  奇偶交错   {s_i:.6f}")
    print(f"  前后半切分 {s_h:.6f}")
    print(f"  差值       {abs(s_i - s_h):.6f}  ← 不为 0，实现和权重必须配套")

    print("\n=== 3. 重排维度后两者完全一致 ===")
    qp, kp = interleave_to_half(q), interleave_to_half(k)
    s_h_perm = attn_score(rope_half_split, qp, kp, 0, 3, dim)
    print(f"  奇偶交错(原始排布)      {s_i:.6f}")
    print(f"  前后半切分(重排后权重)  {s_h_perm:.6f}")
    print(f"  差值                    {abs(s_i - s_h_perm):.2e}"
          f"  ← 只差一个固定的维度置换")


if __name__ == "__main__":
    main()

---