从 GPT-2 到 LLaMA · EP03

旋转位置编码

RoPE

上一期我们把 FFN 里的 GELU 换成了 SwiGLU。这一期换另一个零件：位置编码。

GPT-2 用的是"可学习的绝对位置向量"，加在 Embedding 上。LLaMA、Qwen、Mistral、DeepSeek 这些主流开源模型，基本都换成了 RoPE（Rotary Position Embedding，旋转位置编码）。

这一期只回答三个问题：RoPE 到底做了什么？为什么一个"旋转"就能表达相对位置？代码里怎么写？

一、先回顾：GPT-2 的位置编码有哪里不舒服

GPT-2 的做法（原系列位置编码那一期讲过）：

输入 = token_embedding + position_embedding

position_embedding 是一张可学习的表，形状 [1024, 768]，第 m 行就是"第 m 个位置"的向量，直接加到 Embedding 上。

这个做法能用，但有三个不舒服的地方：

1. 位置信息只在输入层加一次。 之后经过 N 层 Block，位置信息要靠网络自己一路"记住"，没有任何机制保证它不被冲淡。

2. 它是绝对位置，而注意力真正关心的是相对距离。 "鱼"这个 token 去看"吃"，重要的是"它在我前面 1 个位置"，而不是"我是第 2 个、它是第 1 个"。绝对位置的做法，模型得自己从两个绝对位置里把"差值"学出来。

3. 有硬上限。 表只有 1024 行，第 1025 个位置没有对应的向量，直接没法算。

RoPE 的思路是：位置信息不要加在输入上，而是直接作用在注意力的 Q 和 K 上，并且让 Q·K 的结果天然只依赖相对距离。

二、核心想法：给位置 m 的向量，转 m 份角度

先看最简单的二维情况。

二维旋转矩阵，把向量逆时针转 α 角：

R(α) = [ cos α   -sin α ]
       [ sin α    cos α ]

对向量 (x0, x1) 做旋转：

(x0, x1) → (x0·cos α - x1·sin α,  x0·sin α + x1·cos α)

RoPE 的规定很简单：

位置 m 上的 q 和 k，都旋转 m·θ 这么大的角度。

也就是：

q_m = R(mθ) · q
k_n = R(nθ) · k

位置越靠后，转得越多。位置 0 不转。

为什么这样就只剩相对距离了？

注意力分数是 q 和 k 的点积。把旋转代进去：

q_m · k_n = (R(mθ) q)ᵀ (R(nθ) k)
          = qᵀ · R(mθ)ᵀ · R(nθ) · k

旋转矩阵有两个性质：

  • 转置等于反向旋转：R(α)ᵀ = R(-α)
  • 连续旋转，角度相加：R(α)·R(β) = R(α+β)

所以：

R(mθ)ᵀ · R(nθ) = R(-mθ) · R(nθ) = R((n-m)θ)

代回去：

q_m · k_n = qᵀ · R((n-m)θ) · k

结果里只剩 (n−m)，也就是相对距离。m 和 n 各自是多少，已经消失了。

这就是 RoPE 的全部数学核心：用旋转把"绝对位置"编码进去，点积时绝对位置自动抵消，只留下相对位置。

三、推广到 d 维：两两配对，每一对转速不同

真实的 head_dim 不是 2，是 64、128。做法是把向量的维度两两配对，每一对当作一个二维平面，各自旋转。

head_dim = d 时，一共有 d/2 对，第 i 对（i = 0, 1, ..., d/2 − 1）的旋转速度是：

θ_i = 10000^(-2i/d)

i 越小，θ 越大，转得越快；i 越大，θ 越小，转得越慢。

类比：钟表。

  • 秒针转得快，能分辨很近的两个时刻
  • 分针慢一些，分辨中等的间隔
  • 时针最慢，用来分辨很远的间隔

RoPE 里，低维的配对像秒针，负责近距离的位置区分；高维的配对像时针，负责远距离。几对"指针"放在一起，一个位置就有了一个独特的"时间读数"。

10000 这个底数（base）是超参数，来自原始 Transformer 的正弦位置编码，RoPE 沿用了它。后面讲长上下文时它还会出场。

四、手算：猫吃鱼 × RoPE

沿用我们的例子：三个 token，位置分别是

猫 → 位置 0
吃 → 位置 1
鱼 → 位置 2

为了手算方便，这一期取 head_dim = 4（单头，只看一个头内部发生了什么）。

Step 1：算旋转速度

d = 4，所以有 2 对：(第 0, 1 维) 和 (第 2, 3 维)。

θ_0 = 10000^(-0/4) = 1
θ_1 = 10000^(-2/4) = 10000^(-0.5) = 0.01

第一对每走一个位置转 1 弧度（约 57°），第二对每走一个位置只转 0.01 弧度。一快一慢。

Step 2：旋转一个向量

为了把"位置的影响"单独隔离出来，我们故意让三个 token 的 q、k 原始向量都一样：

q = k = [1, 0, 1, 0]

（这是玩具设定，真实模型里不同 token 的 q、k 当然不同。）

把它按 (x0, x1)、(x2, x3) 分成两对，分别旋转 m·θ。

位置 0（猫）： 角度 0，不转。

[1, 0, 1, 0]

位置 1（吃）： 第一对转 1 弧度，第二对转 0.01 弧度。

第一对：(1, 0) → (cos 1, sin 1)       = (0.5403, 0.8415)
第二对：(1, 0) → (cos 0.01, sin 0.01) = (1.0000, 0.0100)

结果：[0.5403, 0.8415, 1.0000, 0.0100]

位置 2（鱼）： 第一对转 2 弧度，第二对转 0.02 弧度。

第一对：(1, 0) → (cos 2, sin 2)       = (-0.4161, 0.9093)
第二对：(1, 0) → (cos 0.02, sin 0.02) = (0.9998, 0.0200)

结果：[-0.4161, 0.9093, 0.9998, 0.0200]

Step 3：算注意力分数（未缩放的 q·k）

没有 RoPE 时： 三个 token 的 q、k 完全一样，任意两个之间的分数都是

[1,0,1,0] · [1,0,1,0] = 2

模型从分数里看不出任何位置信息。

有 RoPE 时：

「鱼」(位置 2) 看「吃」(位置 1)，距离 1：

[-0.4161, 0.9093, 0.9998, 0.0200] · [0.5403, 0.8415, 1.0000, 0.0100]
= -0.2248 + 0.7652 + 0.9998 + 0.0002
≈ 1.5403

「吃」(位置 1) 看「猫」(位置 0)，距离也是 1：

[0.5403, 0.8415, 1.0000, 0.0100] · [1, 0, 1, 0]
= 0.5403 + 0 + 1.0000 + 0
≈ 1.5403

「鱼」(位置 2) 看「猫」(位置 0)，距离 2：

[-0.4161, 0.9093, 0.9998, 0.0200] · [1, 0, 1, 0]
= -0.4161 + 0.9998
≈ 0.5837

Step 4：看结果

谁看谁 绝对位置 (m, n) 相对距离 分数
吃 → 猫 (1, 0) 1 1.5403
鱼 → 吃 (2, 1) 1 1.5403
鱼 → 猫 (2, 0) 2 0.5837

两件事：

  1. (1,0) 和 (2,1) 的绝对位置不同，但分数完全一样。 因为它们的相对距离都是 1。这正是上一节推导出来的性质。
  2. 距离从 1 变成 2，分数从 1.54 掉到 0.58。 位置信息确实进入了注意力分数。

可以用推导验证一下：第一对的贡献是 cos(δ·1)，第二对是 cos(δ·0.01)，δ = 距离。δ=1 时 cos1 + cos0.01 = 0.5403 + 0.99995 ≈ 1.5403；δ=2 时 cos2 + cos0.02 = -0.4161 + 0.9998 ≈ 0.5837。完全吻合。

诚实声明： "距离越远分数越低"在这个例子里成立，是因为我们人为把三个 token 的内容设成了一样。真实模型里，分数同时取决于内容（q、k 学到的向量）和相对距离，不一定随距离单调下降。RoPE 只保证一件事：位置的影响以相对距离的形式进入分数，具体怎么用，由训练决定。

五、几个容易搞混的点

1. 只作用在 Q 和 K 上，不作用在 V 上。 位置信息的作用是影响"谁该关注谁"（分数），而 V 是"被取走的内容"，不需要旋转。

2. 每一层都做，不是只在输入层。 每个 Block 的注意力里，Q、K 算出来之后都要重新旋转一次。这也回应了前面的第一个不舒服：位置信息不会被冲淡。

3. 没有可学习参数。 旋转角度由位置 m 和固定公式 θ_i 直接算出来，不需要训练。GPT-2 里那张 [1024, 768] 的 wpe 表，在 RoPE 的模型里整个删掉。

4. 旋转不改变向量的长度。 旋转只改方向，不改模长，所以不会影响 q、k 本身的数值规模。

5. 配对方式有两种约定，这是实际加载权重时的坑。

  • 相邻配对（原论文）： (0,1)、(2,3)、(4,5)……
  • 前后半配对（HuggingFace LLaMA 的实现）： (0, d/2)、(1, d/2+1)……

两种在数学上等价，但维度排列不同。加载预训练权重时，必须和权重对应的约定保持一致，否则 q、k 的维度对不上，输出会是乱码。这和我们之前遇到的 Conv1D 转置是同一类坑：不是原理问题，是约定问题。

六、PyTorch 实现

先预先算好所有位置、所有频率的 cos 和 sin（只算一次）：

import torch

def precompute_rope(head_dim: int, max_len: int, base: float = 10000.0):
    # θ_i = base^(-2i/d)，i = 0..d/2-1
    i = torch.arange(0, head_dim, 2).float()        # [d/2]
    theta = 1.0 / (base ** (i / head_dim))          # [d/2]
    pos = torch.arange(max_len).float()             # [max_len]
    angles = torch.outer(pos, theta)                # [max_len, d/2]  每个位置×每个频率
    return torch.cos(angles), torch.sin(angles)     # 各 [max_len, d/2]

然后是旋转本身，这里用和手算一致的相邻配对：

def apply_rope(x, cos, sin):
    # x: [B, H, T, d]
    T = x.shape[-2]
    cos, sin = cos[:T], sin[:T]                     # [T, d/2]，自动广播到 [B, H, T, d/2]

    x1 = x[..., 0::2]                               # 偶数维 [B, H, T, d/2]
    x2 = x[..., 1::2]                               # 奇数维 [B, H, T, d/2]

    out1 = x1 * cos - x2 * sin                      # 对应 x0·cosα − x1·sinα
    out2 = x1 * sin + x2 * cos                      # 对应 x0·sinα + x1·cosα

    out = torch.stack([out1, out2], dim=-1)         # [B, H, T, d/2, 2]
    return out.flatten(-2)                          # [B, H, T, d]

如果要对齐 HuggingFace LLaMA 的前后半配对，只需要改切分方式：

def apply_rope_half(x, cos, sin):
    half = x.shape[-1] // 2
    x1, x2 = x[..., :half], x[..., half:]
    return torch.cat([x1 * cos[:x.shape[-2]] - x2 * sin[:x.shape[-2]],
                      x1 * sin[:x.shape[-2]] + x2 * cos[:x.shape[-2]]], dim=-1)

接进注意力

改动只有两处：删掉输入层的位置 Embedding，在 Q、K 算完之后旋转。

class MultiHeadAttention(nn.Module):
    def forward(self, x, cos, sin):
        B, T, _ = x.shape
        q = self.W_q(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.W_k(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.W_v(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)

        q = apply_rope(q, cos, sin)     # 只转 q
        k = apply_rope(k, cos, sin)     # 只转 k，v 不动

        scores = q @ k.transpose(-2, -1) / self.head_dim ** 0.5
        # 后面的 causal mask、softmax、@ v 都和之前完全一样
        ...

GPT 主体里：

# 之前（GPT-2）
x = self.tok_emb(idx) + self.pos_emb(torch.arange(T))

# 现在（RoPE）
x = self.tok_emb(idx)           # 不再加位置向量
# cos, sin 在模型初始化时 precompute_rope 一次，每层 Block 共用

其余部分（Block、残差、归一化、FFN）完全不用动，这也是"零件替换"的意义：换一个零件，其他结构保持原样。

七、一个预告：长上下文

RoPE 在训练长度之外的表现，是后来一系列工作的出发点。

直觉上：训练时模型只见过位置 0 ~ L 对应的那些旋转角度；超出 L 之后，会出现训练时没见过的角度组合，效果会变差。

后续的位置插值（Position Interpolation）、NTK-aware 缩放、YaRN 等方法，基本都是在调整"角度"或"base"，让更长的序列落进模型熟悉的角度范围内。这些我们之后单独讲，这一期只需要知道：RoPE 的 base（10000）不是随便定的常数，它直接决定了模型能"看清"多远的距离。

八、小结

  • GPT-2：位置向量加在输入上，绝对位置，有长度上限
  • RoPE：位置信息变成对 Q、K 的旋转，位置 m 转 mθ
  • 点积后，m 和 n 各自抵消，只剩相对距离 (n−m)
  • d 维向量两两配对，每对不同转速，像钟表的秒针、分针、时针
  • 只作用于 Q、K，每层都做，没有可学习参数
  • 实现时注意配对约定（相邻 vs 前后半），加载权重要对齐

与前后文的关系：

  • 原系列"位置编码"那一期：GPT-2 的绝对位置向量，本期把它整个替换掉
  • 进阶篇「从 GPT-2 到 LLaMA」 EP1 / EP2：RMSNorm、SwiGLU，本期是第三个被换掉的零件

下期预告

位置编码换完了，下一个要换的是注意力本身的结构：GQA（Grouped-Query Attention）。

为什么 LLaMA 2 / 3 这些模型，K 和 V 的头数比 Q 少？它和我们之前讲过的 KV Cache 有什么关系？下期用猫吃鱼手算给你看。

参考：Su et al., "RoFormer: Enhanced Transformer with Rotary Position Embedding", 2021.