三十二：Long Context

来源:http://mp.weixin.qq.com/s?__biz=MzYyNTk3Njg1NA==&mid=2247484378&idx=1&sn=3fa3a6caf3366eca0c3ff219038ea3c1&chksm=f01eb0a3c76939b5d173157f2ae0d561715c34608a695f50a766524a80d03a4fb1c8e856ac71#rd

1. 学习范围

本日主题是 Long Context 中的位置编码复习，重点是 1.4.3 旋转位置编码 RoPE。今天不是重新泛泛学习所有位置编码，而是把 RoPE 讲到面试可推导、工程可实现、长上下文可排错的程度。 需要掌握： - Transformer 为什么需要位置编码。

• RoPE 为什么作用在 Q/K，而不是 V。

• 二维旋转矩阵和高维成对旋转公式。

• RoPE 为什么让 attention score 依赖相对位置。

• cos/sin cache、position_ids、rotate_half 的实现形状。

• RoPE 在推理 KV cache 中如何处理位置偏移。

• RoPE 长上下文外推为什么会退化。

• position interpolation、NTK-aware scaling、YaRN 等方法的直觉。

2. 参考资料

• 课程截图推荐 RoPE 帖：https://blog.csdn.net/v_JULY_v/article/details/134085503

• 苏剑林 RoPE 原始介绍：https://spaces.ac.cn/archives/8265

• RoFormer 原论文：https://arxiv.org/abs/2104.09864

• Position Interpolation：https://arxiv.org/abs/2306.15595

• YaRN：https://arxiv.org/abs/2309.00071

• LongRoPE：https://arxiv.org/abs/2402.13753

3. 位置编码的定位

Self-attention 的核心计算是：

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k))V

如果没有位置机制，attention 只知道 token 内容之间的相似性，不天然知道 token 的顺序和距离。位置编码的目标是让模型区分：

我 爱 你
你 爱 我

常见位置机制：

绝对位置编码:
  x_m + p_m
  每个绝对位置 m 有一个位置向量。

相对位置编码:
  attention score 中显式加入 f(i-j)。

RoPE:
  对 Q/K 按位置做旋转。
  Q/K 内积天然包含相对距离。

ALiBi:
  attention score 加与距离相关的线性 bias。

RoPE 的特别之处在于：它不把位置向量加到 embedding 上，而是在 attention 计算前，对每个位置的 Q 和 K 做同角度体系下的旋转。

4. RoPE 的核心直觉

RoPE 把向量的每两个维度看成一个二维平面，在位置 m 上旋转角度 m * theta_i。 对第 i 个二维子空间：

q_i = [q_{2i}, q_{2i+1}]
k_i = [k_{2i}, k_{2i+1}]
theta_i = base^(-2i / d)

位置 m 的旋转矩阵：

R(m, theta_i) =
  [ cos(m theta_i)  -sin(m theta_i) ]
  [ sin(m theta_i)   cos(m theta_i) ]

旋转后的向量：

q'_{2i}   = q_{2i}   * cos(m theta_i) - q_{2i+1} * sin(m theta_i)
q'_{2i+1} = q_{2i}   * sin(m theta_i) + q_{2i+1} * cos(m theta_i)

K 在位置 n 上同理旋转：

k' = R(n, theta_i) k

这就是“旋转位置编码”的名字来源。

5. RoPE 为什么体现相对位置

关键性质来自旋转矩阵：

R(m)^T R(n) = R(n - m)

因此，对位置 m 的 query 和位置 n 的 key：

(R(m)q)^T (R(n)k)
  = q^T R(m)^T R(n) k
  = q^T R(n - m) k

attention score 依赖的是 n - m，也就是相对距离，而不是只依赖绝对位置。实际高维 RoPE 是很多二维旋转块的组合，每一对维度有不同频率。 这说明 RoPE 兼具两点： - 用绝对位置 m 和 n 参与编码。

• 进入 QK 内积后，天然表现出相对位置性质。

这也是 RoPE 在现代 decoder-only LLM 中广泛使用的原因。

6. 高频与低频维度

RoPE 通常使用类似 sinusoidal position encoding 的频率设计：

theta_i = 1 / base^(2i / d)
base 通常为 10000
i = 0, 1, ..., d/2 - 1

小 i 对应更高频，位置变化一点角度就变化明显；大 i 对应更低频，能表达更长距离的缓慢变化。 直觉： - 高频维度对局部顺序敏感。

• 低频维度对长距离关系更稳定。

• 多个频率叠加后，模型能同时感知近距离和远距离相对位置。

长上下文问题也和频率有关：当位置超过训练长度很多时，高频维度的角度分布可能落到训练时很少见的区域，导致 attention score 泛化变差。

7. RoPE 在 attention 中的张量形状

假设：

B = batch size
S = sequence length
H = num heads
D = hidden size
Dh = head_dim
D = H * Dh

线性投影后：

hidden_states: [B, S, D]
q: [B, H, S, Dh]
k: [B, H, S, Dh]
v: [B, H, S, Dh]

RoPE 通常作用在 q 和 k 上：

q_rot = apply_rope(q, cos, sin)
k_rot = apply_rope(k, cos, sin)

然后计算 attention：

scores = q_rot @ k_rot.transpose(-1, -2) / sqrt(Dh)
scores shape: [B, H, S_q, S_k]

cos/sin cache 常见形状：

inv_freq: [Dh / 2]
position_ids: [B, S] 或 [S]
freqs = position_ids[:, None] * inv_freq[None, :]
freqs shape: [S, Dh / 2]
cos/sin after duplicate: [S, Dh]
broadcast to q/k: [1, 1, S, Dh] 或 [B, 1, S, Dh]

实际模型中还可能只对一部分 head_dim 做 RoPE：

rotary_dim < Dh
q_rot_part = q[..., :rotary_dim]
q_pass_part = q[..., rotary_dim:]

这种 partial rotary 需要特别注意拼接顺序和维度一致。

8. rotate_half 的实现

一种常见实现方式：

def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

def apply_rope(x, cos, sin):
    return x * cos + rotate_half(x) * sin

另一种实现按偶数维和奇数维成对旋转：

def rotate_every_two(x):
    x_even = x[..., 0::2]
    x_odd = x[..., 1::2]
    y_even = -x_odd
    y_odd = x_even
    return torch.stack((y_even, y_odd), dim=-1).flatten(-2)

两种实现对应不同的维度排列约定。只要训练和推理一致，数学上都可以；但如果把不同模型的权重、cache 或实现混用，排列约定不一致会直接导致效果崩掉。

9. RoPE 为什么不作用在 V 上

Attention score 由 Q 和 K 的内积决定：

score_{m,n} = q_m^T k_n / sqrt(Dh)

RoPE 的目标是让 score 感知位置关系，所以作用在 Q/K 上最直接。V 是被加权汇聚的内容向量，不参与 score 的位置匹配。如果把 RoPE 也作用在 V 上，会把内容表示也按位置旋转，通常不是必要设计，还可能干扰值信息。 因此现代 LLM 通常：

apply RoPE to Q and K
do not apply RoPE to V

10. 推理 KV cache 中的位置处理

自回归生成时，不会每一步都重新计算所有历史 token 的 K/V，而是缓存历史：

prefill:
  input length = S
  compute q/k/v for all positions 0..S-1
  cache k/v

decode step t:
  new token position = S + t
  compute q/k/v only for new token
  append k/v to cache

RoPE 在 KV cache 中最容易出错的是 position_ids： - prefill 阶段，token 位置通常是 0..S-1。

• decode 阶段，新 token 位置必须接着历史长度走。

• 如果每步都把新 token 当成位置 0，模型会把所有新 token 看成同一位置，生成质量会异常。

正确逻辑：

past_key_values_length = len(cache)
position_ids = past_key_values_length + arange(new_tokens)

还要注意 padding 侧： - left padding 时，position_ids 通常不能简单等于列下标。

• batch 中不同样本 cache 长度不同，需要正确对齐。

• attention mask 和 position_ids 必须一致。

11. RoPE 与长上下文

RoPE 没有固定长度的位置表，所以理论上可以计算任意位置的 cos/sin。但这不等于模型自然具备无限长上下文能力。 退化原因： - 训练时只见过有限长度的位置分布。

• 超出训练长度后，旋转角度组合进入未见区域。

• 高频维度相位变化过快，attention score 分布异常。

• 远距离 token 增多后，注意力稀释和检索干扰增加。

• 模型的训练数据和任务未覆盖长距离依赖。

因此 RoPE 具备一定外推潜力，但长上下文扩展通常还需要 scaling、继续训练或专门长上下文数据。

12. Position Interpolation

Position Interpolation 的核心思想是：不要让模型直接看到训练长度外的新位置，而是把长位置压缩回训练长度范围内。 假设原训练长度为 L_train，目标长度为 L_target，缩放因子：

s = L_target / L_train

把真实位置 m 映射为：

m' = m / s

然后用 m' 计算 RoPE：

cos(m' theta), sin(m' theta)

直觉是把更长序列“压缩”到模型熟悉的位置范围。优点是稳定，缺点是位置分辨率下降，近距离 token 的相对间隔也被压缩。通常需要少量长上下文微调来恢复性能。

13. NTK-aware Scaling 与动态 RoPE Scaling

NTK-aware scaling 的直觉是调整 RoPE 的 base 或频率分布，让低频维度覆盖更长距离，同时尽量保留局部位置信息。 粗略理解：

原始 theta_i = 1 / base^(2i / d)
扩展上下文时增大或动态调整 base
低频维度变得更适合长距离

动态 RoPE scaling 会根据实际序列长度决定缩放程度。它比固定插值更灵活，但实现细节依赖具体模型，不能随意把一个模型的 scaling 配置搬到另一个模型上。 面试中不必背复杂公式，但要能说清： - Position Interpolation 是压缩 position。

• NTK-aware scaling 是调整 frequency/base。

• YaRN 进一步组合局部保真和长距离扩展，并常配合少量微调。

14. YaRN 与 LongRoPE 的直觉

痛点总结：PI 全局压缩位置，近距离 token 区分度暴跌；单纯 NTK 统一修改 base，长短维度无法差异化调节。 YaRN 关注高效扩展 RoPE 上下文窗口。它的关键思想是对不同频率维度采用更细的缩放策略（分频率维度差异化缩放），尽量保留短距离能力，同时扩展长距离能力，并用较少训练 token 达到较好长上下文效果。 - 高频维度（负责局部、近距离语义）：缩放系数接近 1，几乎不压缩，保住短句、相邻字词分辨能力；

• 低频维度（负责长距离依赖）：大幅缩放，拉长周期，适配超长上下文远距离关联。

LongRoPE 进一步关注更长窗口扩展中的非均匀位置插值和渐进式扩展。它的核心问题意识是：不同维度、不同位置区间对扩展的敏感度不同，统一缩放可能不是最优。 LongRoPE 在 YaRN 分层频率缩放基础上，再加一层位置区间非均匀插值，解决两大观察： - 1、不同位置区间敏感程度不一样：文本靠前的位置、末尾远端位置对缩放的耐受度不同；

• 2、统一线性缩放（PI / 普通 NTK / 基础 YaRN）无法适配全部位置区间，存在分布偏移。

• 两大关键设计：

• 1、非均匀位置插值：不使用固定全局缩放 s，不同位置段采用不同缩放比例；

• 2、渐进式窗口扩展训练：训练时逐步提升上下文长度，让模型平滑适应更长位置分布，避免一次性拉满长度带来的分布断层。

面试表达可以保持在直觉层：

RoPE 长上下文扩展不是简单把 max_position_embeddings 改大。
需要处理频率外推、局部分辨率、长距离覆盖和训练分布偏移。
Position Interpolation、NTK scaling、YaRN、LongRoPE 都是在调 RoPE 的位置到角度映射。

15. 工程实现中的常见坑

位置从 0 还是 1 开始： 训练和推理必须一致。绝大多数实现从 0 开始。 Q/K 维度顺序： 有的实现是前半/后半旋转，有的实现是偶/奇维成对旋转。加载权重时必须匹配。 cos/sin cache dtype： 长上下文位置下，低精度计算可能带来数值误差。很多实现用 FP32 生成 cos/sin，再转成模型 dtype。 KV cache offset： decode 阶段 position_ids 必须从 cache 长度继续。这个 bug 很隐蔽，常表现为短 prompt 正常，长生成混乱。 padding 和 attention mask： left padding 时，列下标不等于真实 position。batch 内不同长度样本要仔细构造 position_ids。 partial rotary： 只旋转部分维度时，rotary_dim 必须为偶数，拼接后维度必须回到 Dh。 scaling 配置： 不同模型的 rope_theta、rope_scaling、训练长度不同。错误修改可能让模型短上下文也变差。

16. 面试紧凑表达

可以这样回答 RoPE：

RoPE 是 Rotary Position Embedding。它不把位置向量加到 token embedding 上，而是在 attention 前对 Q 和 K 做位置相关的二维旋转。每两个维度构成一个平面，位置 m 对应旋转角 m * theta_i。由于旋转矩阵满足 R(m)^T R(n)=R(n-m)，所以旋转后的 Q 和 K 做内积时，attention score 会自然依赖相对距离 n-m。

工程上，Q/K 形状通常是 [B,H,S,Dh]，先根据 position_ids 和 inv_freq 生成 cos/sin cache，再对 Q/K 应用 x*cos + rotate_half(x)*sin。RoPE 不作用在 V，因为 V 不参与 attention score 的位置匹配。

长上下文时，RoPE 可以计算训练长度之外的位置，但模型未必见过这些相位和长距离分布，所以会退化。常见扩展方法包括 position interpolation、NTK-aware scaling、YaRN 和 LongRoPE，本质上都是调整位置到旋转角度的映射，并通常需要长上下文验证或微调。

17. 紧凑总结

• RoPE 对 Q/K 做位置相关旋转，而不是向 embedding 加位置向量。

• 每两个维度构成一个二维旋转平面。

• R(m)^T R(n)=R(n-m) 是 RoPE 表达相对位置的核心。

• RoPE 的 attention score 同时使用绝对位置编码形式和相对位置效果。

• 工程实现要关注 cos/sin cache、position_ids、rotate_half、KV cache offset。

• RoPE 不保证无限长上下文，超出训练长度仍会退化。

• 长上下文 RoPE scaling 的本质是调整 position 到 angle 的映射。

          预览时标签不可点
  

  <div class="