「大模型学习」(14)Long Context之线性插值、NTK-aware、Yarn、LongRoPE

Long Context要解决的问题是什么？

你整理得非常清晰，我来为你补充和系统化讲解一下 大语言模型（LLMs）中的 Long Context 插值技术，帮助你建立起一整套认知框架，尤其聚焦在位置编码插值策略。

---

🔧 为什么需要“位置插值”？

大语言模型如 GPT、LLaMA 等在训练时使用的是有限长度的上下文窗口（如 2048 token）。如果想在推理阶段扩展上下文到 4K、16K、100K 甚至百万 token，必须解决一个核心问题：

原始位置编码机制不支持未见过的位置索引！

例如：

• RoPE（旋转位置编码）是 GPT 系列中使用的机制，训练时只学习到 2048 个位置索引。
• 如果你强行外推，比如输入 4096 个 token，它就必须处理未见过的位置，从而性能显著下降。

于是引入插值：把原本“稀疏的”位置信息，插值成“密集的”分布，让模型在训练过的位置区间“内”继续工作，而不是外推“未知位置”。

---

回顾rope

$m$ 是 token 在序列中的位置编号（position index），比如在一句话中：

“The cat sat on the mat”

如果我们将这句话 tokenize 为：

Token	Index $m$
The	0
cat	1
sat	2
on	3
the	4
mat	5

那么：

• 第一个 token “The” 的位置是 $m=0$；
• 第二个 token “cat” 的位置是 $m=1$；
• 以此类推。

在位置编码中，我们通常对每个 token 使用它的 $m$ 值来生成对应的旋转编码。

$i$ 是 embedding 向量维度中的索引（维度编号），特别是 RoPE 中的偶数-奇数维度对。

假设一个 token 的嵌入向量维度 $d=8$，我们可以将它分为 4 个维度对：

$i$	维度对 $(2i,2i+1)$
0	(0, 1)
1	(2, 3)
2	(4, 5)
3	(6, 7)

RoPE 在每个维度对 $(2i,2i+1)$ 上使用不同的频率 ${\theta }_i$ 来控制旋转角度：

$${\theta }_i=\frac{1}{10000^{2i/d}}$$

• $i$ 越大，频率 ${\theta }_i$ 越小；
• 这意味着高维旋转更慢、低维旋转更快；

RoPE 编码每个 token 的嵌入向量时，使用如下公式在维度对 $(2i, 2i+1)$ 上进行旋转：

$$\left[\begin{array}{c}{x}_{2i}^{\text{rot}}\\ x_{2i+1}^{\text{rot}}\end{array}\right]=\left[\begin{array}{cc}\cos ({\theta }_i\cdot m)& -\sin ({\theta }_i\cdot m)\\ \sin ({\theta }_i\cdot m)& \cos ({\theta }_i\cdot m)\end{array}\right]\left[\begin{array}{c}{x}_{2i}\\ x_{2i+1}\end{array}\right]$$

• $m$ 决定位置（输入序列的 token 序号）；
• $i$ 决定频率（某个维度对使用哪一个频率）；
• 二者共同决定了该 token 的旋转角度，从而编码位置信息。

---

📏 常见插值技术的逻辑与演化

1️⃣ 位置插值（Positional Interpolation，PI）🔧

本质：缩放位置索引，而不扩展位置数量

• 训练：L = 2048，索引为 0, 1, …, 2047

• 推理想扩展到 L′ = 4096，只需将每个位置 $m$ 映射为：

  $$m^{\prime }=\frac{m\cdot L}{L^{\prime }}=\frac{m\cdot 2048}{4096}=0.5m$$

实现方式（针对 RoPE）

RoPE 编码本质是：

$$\text{RoPE}(m)={\left[\cos (m/{\theta }_i),\sin (m/{\theta }_i)\right]}_{i=1}^{d/2}$$

我们只需要将 $m$ 替换为：

$$m\leftarrow m\cdot \frac{L}{L^{\prime }}$$

即可完成 PI。

这在代码中通常只涉及 一个位置索引缩放参数 scale = L / L’，插入到 RoPE 位置计算函数中。

RoPE 可以接收非整数位置，于是直接对 $f(q,\frac{mL}{L^{\prime }})$ 进行插值。

✅ 优点：

• 不增加参数
• 不外推，安全、稳定
• 微调很快，仅适应新坐标空间

⚠️ 缺点：

• 高频维度被“挤压”，周期短 → 精度损失

---

2️⃣ NTK-aware 插值

对不同频率做差异化处理：“高频外推、低频内插”

在 NTK-aware 插值中我们对这两个变量做了如下变换：

变量	原始 RoPE	NTK-aware 插值
$m$（位置）	原始 token index	缩放为 $m/\alpha$
${\theta }_i$（频率）	$10000^{-2i/d}$	$(10000^{-2i/d}{)}^{1/\alpha }$

最终计算的是：

$$\cos \left(\frac{m}{\alpha }\cdot {\theta }_i^{1/\alpha }\right)$$

✅ 优点：

• 低频更重要，因此优先保留低频结构
• 保证部分维度高频可外推

⚠️ 缺点：

• 存在“越界”频率，带来训练不稳定
• 理论上的 k 并不总能匹配真实推理窗口 → 微调必不可少

---

3️⃣ NTK-by-parts 插值

观察每个 RoPE 维度的波长，分段做插值，只对长波长（低频）维度进行位置插值，保留短波长（高频）维度的原始结构，对中间波长做平滑过渡。

核心思想：

令：

• $L$ 为训练时支持的最大上下文长度

• $L^{\prime }$ 为推理时目标上下文长度

• $s=\frac{L^{\prime }}{L}$ 为扩展倍数（插值因子）

• 第 $d$ 维的 RoPE 原始角频率为 ${\theta }_d=10000^{-2d/D}$，其中 $D$ 为维度总数

• 对应的 波长 为：

  $${\lambda }_d=\frac{2\pi }{{\theta }_d}=2\pi \cdot 10000^{\frac{2d}{D}}$$

✅ 比率定义

定义维度 $d$ 的相对位置比率：

$$r_d=\frac{L}{{\lambda }_d}$$

也即：当前训练上下文长度 $L$ 中，RoPE 第 $d$ 维经历了多少个周期。

✅ 插值策略

设阈值 $\alpha$、$\beta$：

• 若 $r_d>\beta$（即波长 $\lambda \_d\ll L$，频率高） → 不做插值
• 若 $r_d<\alpha$（即波长 $\lambda \_d\gg L$，频率低） → 完全线性插值PI
• 否则：按线性斜坡函数平滑插值

引入过渡因子 ${\gamma }_d$：

$${\gamma }_d=\{\begin{array}{ll}0,& r_d>\beta \\ 1,& r_d<\alpha \\ \frac{\beta -r_d}{\beta -\alpha },& \text{else}\end{array}$$

然后每个维度的角频率缩放为：

$${\theta }_d^{\prime }=(1-{\gamma }_d)\cdot {\theta }_d+{\gamma }_d\cdot \frac{{\theta }_d}{s}$$

✅ 优点：

• 保留高频维度，避免插值过度损害短距离语义建模能力
• 提高长距离建模能力
• 插值策略更加灵活，对多种任务更稳健

---

4️⃣ Yarn 插值 （PI + 动态温度调整）

在 attention logits 上引入温度参数 t：

$$\text{softmax}\left(\frac{q^{T}k}{t\cdot \sqrt{d}}\right)$$
可以根据上下文长度动态调整温度：

• 短上下文：t 较小，插值强度低，保持原始位置编码特性。
• 长上下文：t 增大，增强低频段的插值强度，避免位置编码碰撞。

结合 NTK-by-parts 插值，得到 Yarn 插值方案

✅ 特点：

• RoPE插值 + 动态 softmax 调节
• 理论与工程效果良好，无需改动模型结构

---

6️⃣ LongRoPE

非均匀插值 + 动态伸缩 + RoPE 维度搜索

在保持兼容性的前提下增强 RoPE 的长距离建模能力，尤其在超过预训练上下文长度时减少性能损失。

• 非均匀位置插值：
• • LongRoPE中发现，有效的位置编码插值应考虑两种非均匀性：不同的RoPE维度（${\theta }_i中的i$）和不同的token位置（$cos(m{\theta }_i)中的m$）。低维和初始token位置存储着关键信息，因此需要进行更少程度的插值。相比之下，高维存储的信息相对较为稀疏，可进行较大程度的插值。通过搜索RoPE每个维度以及不同token位置的旋转角度缩放因子，有效地保留了原始 RoPE 位置编码中的信息。这种方法最大程度地减小了位置插值引起的信息损失，从而为微调提供了更好的初始化。
• 上下文渐进式扩展：
• • LongRoPE首先在预训练的LLM上进行256k长度的微调，然后对微调后的模型进行第二次位置插值，以实现2048k的上下文窗口。
• 短上下文窗口性能恢复：
• • 在扩展到2048k上下文窗口后，LongRoPE通过调整RoPE位置插值因子来恢复短上下文窗口的性能。LongRoPE在扩展后的大模型上对8K长度内的RoPE缩放因子进行了重新搜索，以鼓励在较短上下文长度上进行较少的位置插值。在推理过程中，大模型可根据输入长度动态调整相应的 RoPE 缩放因子。

实现过程：

一、主要计算模块及作用

1. RoPEPositionalEncoding

angles = positions.unsqueeze(-1) * self.theta
pos_embed = torch.stack([angles.cos(), angles.sin()], dim=-1).flatten(-2)

对应公式：

对于 token 位置 $m$，RoPE 原始位置编码为：

$$\text{RoPE}(m{)}_i=\left[\begin{array}{c}\cos (m\cdot {\theta }_i),\sin (m\cdot {\theta }_i)\end{array}\right]$$

其中：

$${\theta }_i={\text{base}}^{-\frac{2(i//2)}{d}}$$

2. non_uniform_interpolation

scale = torch.where(mask, torch.ones_like(...), 1 / lambda_factors[i])

对第 $i$ 维度（偶数维）进行插值调整：

👉 对应插值公式（非均匀插值）：

$$\text{若 }m<\hat{n}:{\theta }_{m,i}=m\cdot {\theta }_i$$

$$\text{若 }m\ge \hat{n}:{\theta }_{m,i}=\frac{m}{{\lambda }_i}\cdot {\theta }_i$$

即在 RoPE 中将 大于阈值 $\hat{n}$ 的位置使用维度特定的缩放 ${\lambda }_i$ 进行调整。

最终插值嵌入为：

$$\text{LongRoPE}(m,i)=\left[\begin{array}{c}\cos \left({\theta }_{m,i}\right),\sin \left({\theta }_{m,i}\right)\end{array}\right]$$

二、进化搜索：寻找最优 ${\lambda }_i$

3. initialize_population 设定三类初始化策略：

• PI 初始化：${\lambda }_i=s$ （线性压缩）
• NTK-aware：${\lambda }_i=s^{2i/d}$ （对数压缩）
• YaRN 分段式插值：前128维为1，中间段为 $s^{1/3}$，后段为 $s$

4. search_lambda_factors

通过进化搜索进行迭代优化，目标是找到：

arg ⁡ min ⁡ { λ i } Perplexity ( LongRoPE ( λ i ) ) \arg\min_{\{\lambda_i\}} \text{Perplexity}(\text{LongRoPE}(\lambda_i)) arg{λi​}min​Perplexity(LongRoPE(λi​))

完整插值流程

1. 初始构建 RoPE：生成角度 ${\theta }_i$（频率向量）

2. 位置编码生成：使用 $\cos (m\cdot {\theta }_i)$ 和 $\sin (m\cdot {\theta }_i)$

3. 应用插值：

   • 前$\hat{n}$ 位置：保持原始 RoPE
   • 后位置：插值为 $\frac{m}{{\lambda }_i}$

4. 与输入相加：RoPE 编码 + token embedding

5. 输入 Transformer 编码器

✅ 优点：

• 精度保持最好
• 支持百万 token 推理
• 与原模型兼容，可微调迁移

---

总结一张图（简化版）

插值策略	插值对象	高频维	低频维	特点
PI	统一缩放所有位置	插值	插值	简单、高效
NTK-aware	高频外推/低频插值	外推	插值	提升低频精度
NTK-by-parts	按波长动态选择策略	保留	插值	减少高频损失
Yarn	+ attention 温度缩放	动态	动态	插值 + logits 调节
LongRoPE	token 维度双非均匀插值	动态	动态	精度最优，支持百万级 context