慢读世界

基础符号定义 (Notation)#

假设我们有一个 Query 或 Key 向量 $x$，维度为 $d$。 RoPE (Rotary Positional Embedding) 将向量两两分组，在复平面上旋转。 对于第 $i$ 组分量（其中 $i \in [0, d/2)$），位置为 $m$，基础频率为 $\theta_i$。

• 位置 (Position): $m$ (当前 token 的位置索引)。
• 维度索引: $i \in [0, 1, ..., d/2 - 1]$。
• 基数 (Base): $b = 10000$ (LLaMA 等模型的默认值)。
• 频率 (Frequency): $\theta_i = b^{-2i/d}$。
• 旋转角度: $m \cdot \theta_i$。
• 训练最大长度: $L_{train}$ (例如 2048 或 4096)。
• 推理/目标长度: $L_{new}$。
• 扩展倍率 (Scale): $s = \frac{L_{new}}{L_{train}}$。

原始 RoPE (Baseline)#

在讲外推之前，必须先看原始 RoPE 做了什么。

数学实现： RoPE 的核心是对向量 $x$ 应用一个旋转矩阵。对于第 $i$ 对分量 $(x_{2i}, x_{2i+1})$，旋转角度为 $m \theta_i$：

注意，这里的 $\theta_i$ 公式为：

• 高频分量（$i$ 较小）：旋转速度极快，关注局部位置信息。
• 低频分量（$i$ 较大）：旋转速度极慢，关注长距离全局信息。

外推问题： 当 $m > L_{train}$ 时，$\cos(m\theta_i)$ 会出现训练时从未见过的数值分布，导致注意力机制崩塌（PPL 爆炸）。

PPL 是什么？

PPL (Perplexity)：困惑度，是衡量语言模型预测能力的指标。PPL 越低，模型预测下一个词的能力越强。数学上，PPL 是交叉熵损失的指数形式：

$$\text{PPL} = \exp\left(-\frac{1}{N} \sum_{i=1}^{N} \log P(w_i | w_{<i})\right)$$

其中 $N$ 是词的数量，$P(w_i | w_{<i})$ 是模型预测第 $i$ 个词的概率。PPL可以理解为模型在预测时的“平均分支数”，PPL 越高，表示模型越不确定。

Position Interpolation (PI, 线性插值)#

核心思想： 既然模型没见过大于 $L_{train}$ 的 $m$，那我们把新的长位置 $L_{new}$ 硬塞回 $[0, L_{train}]$ 的范围内。也就是“把尺子刻度变密”。

数学实现： 将位置 $m$ 替换为 $m' = \frac{m}{s}$。 新的频率项变为：

工程实现： 在代码中，不需要改变模型权重，只需要在预计算 cos 和 sin 表（Frequency Buffer）时修改即可。

1
# 伪代码
2
scale = L_new / L_train  # 例如 4096->8192, scale=2
3
t = torch.arange(L_new)
4
t = t / scale  # 将 0~8191 压缩回 0~4095.5
5
freqs = torch.outer(t, inv_freq)

优缺点：

• 优点： 极其简单，微调(Fine-tune)几百步就能让模型适应长文本。
• 缺点： “高频分辨率丢失”。
  • 对于高频部分（$i$ 很小），相邻 token 的旋转角度差值变小了（被除以了 $s$），导致模型无法精确区分相邻的词（例如 “cats” 和 “cat” 的区别）。这就是所谓的“分辨率（Resolution）分辨率”。

NTK-Aware Interpolation#

核心思想： PI 简单粗暴地对所有频率都除以 $s$，这伤害了高频部分。 根据 Neural Tangent Kernel (NTK) 理论，高频特征学习快，低频特征学习慢。 我们不应该改变位置 $m$（保持分辨率），而应该改变“基数”$b$。 通过改变 $b$，我们可以让高频部分保持不变（不插值），只对低频部分进行插值。

数学实现： 我们要寻找一个新的基数 $b'$，使得最低频（$i=d/2-1$）的效果等同于 PI 插值，而高频变化最小。 PI 的低频项是：$\frac{m}{s} \cdot b^{-2i/d}$。 NTK-Aware 保持 $m$ 不变，改变 $b$ 为 $b'$：$m \cdot (b')^{-2i/d}$。

令两者在最低频处相等（$i$ 最大时）：

注：推导细节为了简化忽略了$2i/d$，实际公式考虑所有维度，推导出的修正基数如下：

最终公式：

通常 $\frac{d}{d-2} \approx 1$，所以很多实现里直接近似 $b' \approx b \cdot s$。

工程实现： 不用改 $m$（位置索引），只改计算频率时的 base。

1
# 伪代码
2
alpha = L_new / L_train
3
# 计算新的 base
4
new_base = 10000 * (alpha ** (dim / (dim - 2)))
5
# 重新计算 inv_freq
6
inv_freq = 1.0 / (new_base ** (torch.arange(0, dim, 2) / dim))
7
# 生成 cos/sin 表

优缺点：

• 优点： 解决了 PI 在高频部分分辨率丢失的问题。无需微调也能获得比 PI 更好的 PPL，微调后效果更佳。

Dynamic NTK Interpolation#

核心思想： 前面两种方法 $s$ 是固定的。如果你设 $L_{new}=16k$，但在推理时只输入了 1k 个词，模型性能反而会下降，因为它被不必要地“压缩”了。 Dynamic NTK 根据当前生成的序列长度 $L_{curr}$ 动态计算 $s$。

数学实现： 在每一步推理（或者每隔一段）：

1. 获取当前序列长度 $L_{curr}$。
2. 如果 $L_{curr} \le L_{train}$，使用原始 RoPE ($scale=1$)。
3. 如果 $L_{curr} > L_{train}$，计算 $s = \frac{L_{curr}}{L_{train}}$。
4. 使用 NTK-Aware 重新计算 $b'$ 和频率表。

工程实现： 这对 KV Cache 系统提出了挑战，因为 RoPE 以前是静态的，现在变成了动态的。 如果 RoPE 是在 Attention 之前做的（通常是），意味着以前存进 KV Cache 的 Key 向量是基于旧的 $s$ 旋转的。当 $s$ 变化时，之前的 KV Cache 失效了？

• 实际上： 大多数实现为了性能，在推理时如果采用 Dynamic NTK，通常不重算 KV Cache，这会导致一定的误差，但在 RoPE 的语境下，模型往往能容忍这种渐变。或者采用“不旋转存 KV，取出时旋转”的策略（但计算量大）。

优缺点：

• 优点： 在短文本上保持原始性能，随着长度增加平滑过渡到外推模式。

NTK-by-parts Interpolation (分段 NTK)#

核心思想： NTK-Aware 即使改变了 base，对高频部分还是有微小的影响。 “By-parts” 提出一个更激进的假设：波长（Wavelength，$\lambda = 2\pi / \theta_i$）。

1. 如果波长 $\lambda$ 很短（远小于训练长度）：这是高频，包含局部信息。绝对不要插值（保持原始 RoPE）。
2. 如果波长 $\lambda$ 很长（大于训练长度）：这是低频，包含长距离信息。必须插值（类似 PI）。
3. 中间部分：过渡。

数学实现： 定义波长 $\lambda_i = 2\pi \cdot b^{2i/d}$。 定义两个阈值 $\alpha, \beta$（例如 $\alpha=1, \beta=32$）。 对于每个维度 $i$：

• 如果 $\lambda_i / L_{train} < \alpha$：不插值（$m \to m$）。
• 如果 $\lambda_i / L_{train} > \beta$：全量插值（$m \to m/s$）。
• 中间：线性混合。

由于直接对 $m$ 混合比较麻烦，通常会将这个逻辑转换为改变频率 $\theta_i$：

其中 $\gamma_i$ 是根据波长计算出的 ramp 函数（0 到 1）。

工程实现： 需要为每个维度 $i$ 计算一个缩放因子。

1
# CodeLLaMA 采用的就是类似思路
2
# 计算每个维度的波长
3
wavelen = 2 * math.pi / inv_freq
4
# 这种方法不改 base，而是直接改 freq
5
# 对 wavelen 小的不变，大的除以 scale

YaRN (Yet another RoPE extensioN)#

这是目前的集大成者（CodeLlama 等模型深受其影响）。

核心思想： YaRN 结合了 NTK-by-parts 的分段思想，并解决了一个隐藏问题：熵（Entropy）的变化。 当我们将所有 token 的 RoPE 距离拉开（或压缩）时，Attention 矩阵中的点积数值分布会发生变化（通常变得更“平”或更“尖”），这会导致 Softmax 后的分布熵变大，模型困惑度增加。 YaRN 引入了 temperature scaling 来修正注意力分布。

数学实现（三大步骤）：

1. NTK-by-parts: 使用分段策略（高频不转，低频转，中间插值）。
   • 不像原始 NTK 改 base，YaRN 直接通过修正 frequency 来实现。
2. Length Scaling (注意力修正): 在计算 Attention Score $\frac{q \cdot k^T}{\sqrt{d}}$ 时，乘上一个额外的因子 $\sqrt{t}$。 或者简单地理解为，在 RoPE 之后，把 $q, k$ 向量本身乘以 $\sqrt{\sqrt{t}}$（即 $t^{0.25}$）。 因子 $t$ 通常取值： $$\sqrt{t} \approx 0.1 \ln(s) + 1$$ 如果不加这个修正，长文本外推时，注意力分布会因为上下文变长而变得过于“平坦”，模型不知道该关注谁。

工程实现： YaRN 的实现通常包含两个部分：

1. Get_Rotary_Embedding: 生成修改后的 cos/sin 表（基于分段逻辑）。
2. Model Forward: 在应用 FlashAttention 之前或之后，对 Logits 进行 scaling，或者直接对 RoPE 后的 output 进行 scaling。

1
# YaRN 简易逻辑
2
# 1. 计算 ramp function (ramp) 基于波长
3
# 2. inv_freq_new = inv_freq * (1 - ramp) + (inv_freq / scale) * ramp
4
# 3. 还有一个 scale factor
5
sqrt_t = 0.1 * math.log(scale) + 1
6
# 在应用 RoPE 后，将 q, k 乘以 sqrt(sqrt_t) (或者在 attention socre 乘 1/t)
7
q = apply_rope(q, inv_freq_new) * t_correction
8
k = apply_rope(k, inv_freq_new) * t_correction

总结与对比 (Cheat Sheet)#