汤不热吧在大语言模型的演进历程中，位置编码（Positional Encoding）始终是一个核心而微妙的问题。Transformer 架构本身是排列等变的，它无法区分“猫追狗”和“狗追猫”——必须通过位置编码注入顺序信息。自 2021 年苏剑林等人在 RoFormer 中提出旋转位置编码（Rotary Position Embedding, RoPE）以来，这一方案迅速成为 LLaMA、Qwen、Mistral、DeepSeek 等几乎所有主流开源大模型的标配。本文将从数学原理出发，深入剖析 RoPE 的工作机制，并重点讨论长上下文外推（length extrapolation）的工程实现方案。
一、为什么需要位置编码：从绝对到相对的范式转变
最初的 Transformer（Vaswani et al., 2017）使用正弦-余弦绝对位置编码，它为每个位置生成一个固定的向量并加到 token 嵌入上。这种方案简单有效，但存在根本缺陷：它只能编码“第几个位置”，而非“两个 token 之间的相对距离”。而注意力机制的核心是 Query 和 Key 的点积——真正影响注意力分数的是 q_i 与 k_j 之间的相对位置关系 (i – j)，而非绝对位置 i 和 j 各自的值。
相对位置编码（RPE）直接将相对距离信息注入到注意力分数的计算中，例如 T5 的相对偏置、ALiBi 的线性偏置等。RoPE 则走了一条更优雅的路线：它通过对 Query 和 Key 施加旋转变换，使得它们的点积自然成为相对位置的函数，同时保持了绝对位置编码“在进入注意力之前就处理”的简洁性。
1.1 RoPE 的核心设计目标
RoPE 的设计目标可以用一个数学条件来概括：我们希望找到一种变换 f，使得对于任意位置 m 和 n，Query q_m 和 Key k_n 经变换后的内积仅依赖于相对位置 (m – n)：
也就是说，无论绝对位置 m 和 n 是多少，只要它们的相对距离 (m – n) 相同，注意力分数就应该一致。这个性质被称为“平移不变性”，是长上下文外推的理论基础。
二、RoPE 的数学推导：二维旋转的自然推广
苏剑林的洞察来自一个简洁的二维情形分析。考虑 q 和 k 为二维向量，我们可以用复数表示：q = q_0 + iq_1，k = k_0 + ik_1。如果我们对 q_m 旋转角度 mtheta，对 k_n 旋转角度 ntheta，那么它们的内积恰好是相对角度 (m-n)*theta 的函数：
2
= |q|*|k| * cos(phi_q - phi_k + (m-n)*theta)
这里 phi_q 和 phi_k 是 q 和 k 各自在复平面上的辐角。可以看到，最终的内积只依赖于 (m – n)，完美满足我们的设计目标。将这个二维旋转推广到 d 维向量，只需将 d 个维度两两分组，对每组 (q_{2i}, q_{2i+1}) 施加不同频率 theta_i 的旋转即可。
2.1 旋转矩阵的正式定义
对于维度为 d 的向量，RoPE 将其视为 d/2 个二维子向量，第 i 组的旋转频率为 theta_i = 10000^{-2i/d}。完整的旋转变换可以表示为分块对角矩阵 R_m：
2
3
4
5
6
R(m, theta_0), R(m, theta_1), ..., R(m, theta_{d/2-1})
)
其中 R(m, theta_i) = [[cos(m*theta_i), -sin(m*theta_i)],
[sin(m*theta_i), cos(m*theta_i)]]
变换后的 Query 和 Key 为 q_tilde_m = R_m * q，k_tilde_n = R_n * k。由于旋转矩阵是正交矩阵，R_m^T * R_n = R_{n-m}，因此：
内积仅通过 R_{n-m} 依赖于相对位置，证明完毕。这个推导的精妙之处在于：我们无需修改注意力分数的公式，只需在进入点积之前对 Q 和 K 做旋转，就自动获得了相对位置编码的效果。
三、工程实现：从朴素到高效的算子优化
直接构造 d x d 的稀疏旋转矩阵再做矩阵乘法在计算上极其低效。实际实现中，我们利用旋转矩阵的稀疏结构，将旋转操作转化为逐元素的乘法和加法。以下是基于 PyTorch 的参考实现：
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
import torch.nn as nn
import math
def precompute_freqs_cis(dim, end, theta=10000.0):
# 预计算旋转角度的 cos 和 sin 值
# dim: 每个注意力头的维度, end: 最大序列长度
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: dim // 2].float() / dim))
t = torch.arange(end, device=freqs.device)
freqs = torch.outer(t, freqs) # (end, dim//2)
freqs_cos = freqs.cos()
freqs_sin = freqs.sin()
return freqs_cos, freqs_sin
def apply_rotary_emb(x, freqs_cos, freqs_sin):
# 将 RoPE 应用到输入张量
# x: (batch, seq_len, n_heads, dim)
d = x.shape[-1]
# 将最后一维拆成两半：(x_even, x_odd)
x1, x2 = x[..., :d//2], x[..., d//2:]
# 广播: (1, seq_len, 1, dim//2)
freqs_cos = freqs_cos.unsqueeze(0).unsqueeze(2)
freqs_sin = freqs_sin.unsqueeze(0).unsqueeze(2)
# 旋转: x1' = x1*cos - x2*sin, x2' = x1*sin + x2*cos
rotated = torch.stack([
x1 * freqs_cos - x2 * freqs_sin,
x1 * freqs_sin + x2 * freqs_cos,
], dim=-1).flatten(-2)
return rotated
class RoPEAttention(nn.Module):
def __init__(self, dim, n_heads, max_seq_len=8192, theta=10000.0):
super().__init__()
self.n_heads = n_heads
self.head_dim = dim // n_heads
self.wq = nn.Linear(dim, dim, bias=False)
self.wk = nn.Linear(dim, dim, bias=False)
self.wv = nn.Linear(dim, dim, bias=False)
self.wo = nn.Linear(dim, dim, bias=False)
# 预计算并缓存 cos/sin 值
freqs_cos, freqs_sin = precompute_freqs_cis(
self.head_dim, max_seq_len, theta)
self.register_buffer('freqs_cos', freqs_cos, persistent=False)
self.register_buffer('freqs_sin', freqs_sin, persistent=False)
def forward(self, x, mask=None):
B, T, C = x.shape
q = self.wq(x).view(B, T, self.n_heads, self.head_dim)
k = self.wk(x).view(B, T, self.n_heads, self.head_dim)
v = self.wv(x).view(B, T, self.n_heads, self.head_dim)
# 仅对 Q 和 K 应用 RoPE，V 不需要
q = apply_rotary_emb(q, self.freqs_cos[:T], self.freqs_sin[:T])
k = apply_rotary_emb(k, self.freqs_cos[:T], self.freqs_sin[:T])
# 标准注意力计算
attn = torch.nn.functional.scaled_dot_product_attention(
q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
attn_mask=mask, is_causal=True if mask is None else False)
return self.wo(attn.transpose(1, 2).reshape(B, T, C))
上述实现中有几个工程细节值得注意：第一，cos 和 sin 值可以预计算并缓存，避免每次 forward 重复计算；第二，RoPE 只作用于 Q 和 K，不作用于 V，因为 V 不参与 Query-Key 匹配的位置敏感计算；第三，在融合实现中（如 FlashAttention），RoPE 通常被融合进 QK 投影之后的 kernel 中，避免额外的显存读写往返。
四、长上下文外推：RoPE 的阿喀琉斯之踵与解法
RoPE 的一个关键问题是：当模型在 L_train 长度上训练后，若推理时序列长度超过 L_train，注意力分数会急剧退化。这是因为预计算的频率值在超出训练范围后，旋转角度变得过大，导致不同位置的 Q-K 内积分布发生剧烈变化，注意力变得无意义。
围绕这一问题，学术界和工业界发展出了多种外推方案。以下是主流方案的对比：
| 方案 | 核心思想 | 训练成本 | 外推倍数 | 代表模型 |
|---|---|---|---|---|
| Position Interpolation (PI) | 将位置索引缩放到训练范围 | 微调（少量） | 4-8x | LLaMA-2-Long |
| NTK-Aware Scaling | 调整基础频率 theta 使高频不缩放、低频缩放 | 零微调 | 2-4x | Code Llama |
| YaRN | 分段缩放：高频不动、中频 NTK、低频 PI | 微调（少量） | 8-16x | Qwen2、Mistral-Nemo |
| Dynamic NTK | 推理时根据当前长度动态调整 theta | 零微调 | 2-4x | vLLM 默认 |
4.1 Position Interpolation：最朴素的缩放
PI 的思路直截了当：如果训练时最大位置是 L_train，推理时序列长度为 L_test，就将所有位置索引乘以 L_train/L_test 进行缩放。相当于把“更长的序列”压缩到模型已经见过的位置范围内：
2
3
4
5
6
# Position Interpolation 缩放因子
# 将 [0, test_len) 映射到 [0, train_len)
scale = train_len / test_len
return positions * scale
# 在 RoPE 中等价于: m * (L_train / L_test) * theta_i
PI 的优点是实现极简，仅需修改位置索引的计算方式。缺点是它对所有频率一刀切地缩放，高频旋转也被压缩了，导致局部位置信息的精度下降，通常需要少量微调来恢复质量。
4.2 NTK-Aware Scaling：保高频、压低频
NTK（Neural Tangent Kernel）感知缩放的核心洞察是：高频分量（小的 theta_i，对应近距离相对位置）应该保持不变，因为局部注意力模式在不同序列长度下是通用的；只有低频分量（大的 theta_i，对应远距离相对位置）需要外推。这通过修改基础频率 theta 来实现：
2
3
4
5
6
7
# NTK-Aware 缩放后的基础频率
# 高频（大指数）几乎不缩放，低频（小指数）被大幅拉伸
scale = test_len / train_len
# 新的 base_theta = base_theta * scale^(dim/(dim-2))
new_theta = base_theta * (scale ** (dim / (dim - 2)))
return new_theta
这个方案最大的优势是可以“零微调”——直接修改 theta 参数即可获得 2-4 倍的外推能力，无需重新训练。Code Llama 便是采用了这一方案将上下文从 4K 扩展到 16K。
4.3 YaRN：分段缩放的最优解
YaRN（Yet another RoPE extensioN）是目前最成熟的外推方案，被 Qwen2、Mistral-Nemo 等模型广泛采用。它将频率维度分为三个区间：
- 高频区（波长 < 训练长度）：完全不缩放，保持局部精度
- 低频区（波长 > 外推长度）：使用 PI 方式线性缩放
- 中频区：使用 NTK 在两者之间平滑过渡
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
# 计算 YaRN 的每个频率维度的缩放因子
def find_correction_dim(num_rot, dim, base=10000, max_len=4096):
return dim * (math.log(max_len) - math.log(num_rot)) / \\
(math.log(base) - 2 * math.log(num_rot) / dim)
def find_correction_range(low, high, dim, base, max_len):
low = math.floor(find_correction_dim(low, dim, base, max_len))
high = math.ceil(find_correction_dim(high, dim, base, max_len))
return max(low, 0), min(high, dim // 2 - 1)
# 频率缩放
scale = test_len / train_len
# 计算 alpha 和 beta 边界
low, high = find_correction_range(
beta_fast, beta_slow, dim, 10000.0, train_len)
factors = []
for i in range(dim // 2):
if i < low:
factors.append(1.0) # 高频：不缩放
elif i > high:
factors.append(1.0 / scale) # 低频：PI 缩放
else:
# 中频：平滑插值
t = (i - low) / max(high - low, 1)
factors.append(1.0 / (scale * (1 - t) + t))
return factors
YaRN 还引入了一个注意力缩放因子 t 来补偿由于旋转角度变化带来的注意力分数分布偏移。实践表明，YaRN 配合少量微调（通常 1000-5000 步），可将上下文从 4K 外推至 128K 甚至更长，且几乎不损失短文本上的表现。
五、生产部署实践：常见陷阱与调优建议
在实际将 RoPE 外推方案部署到生产环境时，有几个容易被忽视的细节往往决定成败。
5.1 注意力缩放因子的补偿
当修改了 RoPE 的缩放参数后，Q-K 点积的统计特性会发生变化。标准 Attention 中的 1/sqrt(d) 缩放假设 Q 和 K 是独立的随机向量，但旋转角度的拉伸破坏了这一假设。YaRN 建议引入一个额外的缩放因子 t：
2
3
attn_scale = 0.1 * math.log(scale) + 1.0 # scale = test_len/train_len
# 在 softmax 前应用: logits = (Q @ K^T) / (sqrt(d) * attn_scale)
忽略这一补偿会导致长序列上注意力分布过于尖锐（或过于平坦），表现为模型在长文本上“幻觉”加重或“走神”。
5.2 KV Cache 中的位置索引管理
在使用 KV Cache 的自回归生成中，新 token 的位置索引必须与缓存中已有 token 的索引连续。当使用 PI 或 YaRN 缩放时，必须确保缓存中的历史 Key 已经按正确的缩放因子旋转过。一个常见错误是在训练时用了原始 RoPE，推理时直接切换到缩放版本，但缓存中的历史 Key 还是按原始角度旋转的——这会导致新旧 token 之间的相对位置关系完全错乱。
5.3 不同层使用不同外推策略
一些研究表明，Transformer 不同层对位置信息的敏感度不同：底层更多依赖局部位置信息（高频），高层更多依赖全局语义（低频）。Qwen2 和某些实验性方案采用了“分层缩放”策略——底层使用较小的缩放因子保留高频精度，高层使用较大缩放因子增强外推能力。这种方案虽然增加了工程复杂度，但在超长上下文（>64K）场景下可以进一步提升质量。
六、总结与展望
RoPE 之所以能成为大模型位置编码的事实标准，在于它巧妙地将绝对位置的旋转变换与相对位置的注意力计算统一起来，兼具理论优雅和工程高效。而围绕 RoPE 的外推方案——从 PI 到 NTK 再到 YaRN——本质上都是在回答一个核心问题：如何在“保持局部精度”和“扩展全局视野”之间取得平衡。
展望未来，以下几个方向值得关注。第一，无训练外推的极限：能否在零微调条件下将 4K 训练的模型直接扩展到 1M 上下文，目前 Blockwise NTK、Self-Extend 等方案正在探索这一方向。第二，与注意力机制本身的协同设计：如 ALiBi、NoPE（无位置编码）等替代方案在某些场景下展现了竞争力。第三，硬件层面的支持：当前 RoPE 的旋转操作在 GPU 上需要额外的逐元素运算，未来若能在 Tensor Core 层面原生支持旋转融合，将进一步降低长上下文推理的延迟开销。
对于工程师而言，掌握 RoPE 的原理和外推方案，是在大模型时代做好长上下文推理部署的基本功。希望本文的数学推导和代码实现能帮助你构建起完整的知识框架。