长上下文训练：从 RoPE 到百万 Token 上下文

上下文长度是 LLM 能力的核心维度之一。从最初的 512 token 到现在的百万 token，长上下文技术经历了快速而深刻的演进。本文将系统解析这一技术方向的理论基础和工程实践。

一、为什么长上下文重要

1.1 应用驱动

长上下文能力直接解锁了一系列应用场景：

上下文长度	可以做什么	典型应用
4K	短对话、简单问答	ChatGPT 初期
32K	长文档摘要	论文总结、合同审阅
128K	整本书分析	代码库理解、书籍问答
1M	多文档推理	企业知识库、复杂研究
10M+	视频/长音频理解	Gemini Pro 的目标

1.2 “Needle in a Haystack”

长上下文不仅仅是能 “装下” 更多文本，关键是能在长文本中精确检索和推理。“大海捞针”（Needle in a Haystack）测试就是衡量这一能力的标准基准。

二、位置编码基础

2.1 为什么需要位置编码

Transformer 的自注意力机制本身是位置无关的——“我爱你” 和 “你爱我” 在没有位置信息的情况下会得到相同的表示。位置编码为每个 token 注入其在序列中的位置信息。

2.2 绝对位置编码

Sinusoidal（正弦位置编码）：

Transformer 原始论文提出的方法：

$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d}}\right)$

$PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d}}\right)$

可学习位置编码：

GPT-2/3 使用的方法——位置编码是可训练的参数矩阵。

问题：绝对位置编码在推理时无法泛化到训练时没见过的位置。如果训练时最长序列是 2048，那么位置 2049 的编码完全未定义。

2.3 相对位置编码

相对位置编码不编码绝对位置，而是编码 token 之间的相对距离。

ALiBi（Attention with Linear Biases）：

BLOOM 使用的方法，在注意力分数上加一个与距离成线性关系的偏置：

$\text{Attention}(q_i, k_j) = q_i^T k_j - m \cdot |i - j|$

其中 $m$ 是每个注意力头不同的斜率。

优势：天然支持外推（Extrapolation），训练时用短序列，推理时可以用更长的序列。 劣势：线性衰减假设过于简单，远距离信息衰减太快。

2.4 RoPE：旋转位置编码

RoPE（Rotary Position Embedding） 是目前最主流的位置编码方法，被 LLaMA、Qwen、DeepSeek、Mistral 等几乎所有主流开源模型采用。

核心思想：将位置信息编码为向量的旋转角度。

对于查询向量 $q$ 在位置 $m$ 的编码：

$f(q, m) = q e^{im\theta}$

具体来说，将向量分成二维子空间，每个子空间旋转不同的角度：

$R_\Theta(m) = \begin{pmatrix} \cos m\theta_0 & -\sin m\theta_0 & & \\ \sin m\theta_0 & \cos m\theta_0 & & \\ & & \cos m\theta_1 & -\sin m\theta_1 \\ & & \sin m\theta_1 & \cos m\theta_1 \\ & & & & \ddots \end{pmatrix}$

其中 $\theta_i = 10000^{-2i/d}$（base = 10000）。

RoPE 的优美性质：

1. 相对位置通过内积体现：$q_m^T k_n$ 只依赖于 $m - n$
2. 长距离衰减：自然地对远距离 token 给予更低的注意力权重
3. 无额外参数：不增加模型参数
4. 旋转操作高效：只需要简单的三角函数运算

三、上下文长度扩展方法

3.1 RoPE 外推的挑战

RoPE 的一个关键限制是：训练时如果只见过位置 0 到 $L_{train}$，推理时遇到位置 $> L_{train}$ 时性能会急剧下降。

原因：高频分量（$\theta$ 值大的维度）在超过训练长度后会 “旋转” 到未见过的角度，导致注意力分数异常。

3.2 位置插值（Position Interpolation, PI）

Meta 2023 提出的最直接方法：将长序列的位置压缩到训练范围内。

$f'(x, m) = f\left(x, \frac{mL_{train}}{L_{target}}\right)$

如果训练长度是 4K，目标长度是 32K，那么位置 32000 会被映射到位置 4000（乘以 4K/32K = 1/8）。

优势：简单，只需少量微调（~1000 步）就能扩展 8 倍 劣势：压缩了近距离的位置分辨率

3.3 NTK-aware Interpolation

核心洞察：不应该对所有频率维度做相同的缩放。低频分量（编码长距离关系）需要压缩，高频分量（编码近距离关系）应该保持不变。

做法是修改 RoPE 的 base：

$\text{base}' = \text{base} \times \alpha^{d/(d-2)}$

其中 $\alpha = L_{target} / L_{train}$。

效果：比 PI 更好地保持了近距离精度，同时扩展了远距离覆盖。

3.4 YaRN（Yet another RoPE extensioN）

YaRN 是目前最先进的 RoPE 扩展方法之一，结合了：

1. NTK-aware 缩放：调整不同频率的缩放比例
2. 注意力缩放：补偿长序列的注意力分散
3. 动态 NTK：在推理时根据实际序列长度自适应调整

$s(n) = 1 - \frac{1}{1 + s_n \cdot ((n/n_{train}) - 1)}$

其中不同频率分量有不同的缩放函数。

3.5 ABF（Adjusted Base Frequency）

LLaMA-3 的方法：直接将 RoPE 的 base 从 10000 增大到 500000。

$\theta_i = 500000^{-2i/d}$

直觉：更大的 base 意味着更低的旋转频率，等效于 “拉伸” 了位置编码的周期，可以覆盖更长的序列。

优势：极其简单，无需额外插值或缩放 代价：需要在长序列上做更多的训练

3.6 方法对比

方法	原理	微调需求	扩展倍数	近距离精度
直接外推	什么都不做	无	1.2x	不变
PI	线性压缩	~1K 步	8x	下降
NTK-aware	频率感知缩放	~1K 步	8-16x	较好
YaRN	组合缩放+注意力	~400 步	16-64x	好
ABF	增大 base	大量训练	原生 32x+	好

四、训练策略

4.1 渐进式长度训练

主流方法是分阶段逐步增加训练序列长度：

阶段 1: 4K context  (主预训练，~90% 数据)
  ↓
阶段 2: 32K context (中训练，~8% 数据)
  ↓
阶段 3: 128K context (微调，~2% 数据)

为什么不直接用长序列训练？

1. 计算成本：自注意力的复杂度是 $O(n^2)$，128K 的计算量是 4K 的 1024 倍
2. 数据稀缺：真正高质量的长文本数据远少于短文本
3. 内存限制：长序列的 KV Cache 巨大

4.2 长文本数据准备

长上下文训练对数据有特殊要求：

数据源：

• 长篇论文和报告（arXiv 全文）
• 书籍和长篇小说
• 完整的代码仓库（文件拼接）
• 长对话记录
• 法律文件和合同

数据构造技巧：

1. 文档拼接：将相关的短文档拼接成长序列
2. 上采样长文本：长文本通常数量少，需要适当上采样
3. 保持多样性：不能全是同一类型的长文本
4. 混入短文本：30-50% 的短文本保持基础能力

4.3 注意力机制优化

长序列训练的效率瓶颈在注意力计算。核心优化技术：

Flash Attention（2022, 2023）：

• 不改变数学计算，只优化 GPU 内存访问模式
• 利用 GPU 的 tiling 和 kernel fusion
• 将注意力的内存复杂度从 $O(n^2)$ 降低到 $O(n)$
• FlashAttention-2 进一步优化了并行度

Ring Attention（2023）：

• 将长序列分布到多个 GPU 上
• 每个 GPU 持有序列的一段
• 通过环形通信计算完整的注意力
• 理论上可以支持任意长的序列

稀疏注意力：

• 不计算所有 token 对之间的注意力
• 只关注局部窗口 + 少量全局 token
• 典型方法：Sliding Window Attention（Mistral 使用）

五、KV Cache 管理

5.1 KV Cache 的内存挑战

在推理时，每个 token 需要缓存所有层的 Key 和 Value 向量：

$\text{KV Cache 大小} = 2 \times n_{layers} \times n_{heads} \times d_{head} \times L \times \text{精度}$

以 LLaMA-2 70B 为例（80层, 64头, d_head=128）：

上下文长度	KV Cache (FP16)	KV Cache (INT8)
4K	2.5 GB	1.25 GB
32K	20 GB	10 GB
128K	80 GB	40 GB
1M	625 GB	312 GB

128K 的 KV Cache 就需要 80GB——比模型参数本身还大！

5.2 KV Cache 压缩技术

GQA（Grouped-Query Attention）：

• LLaMA-2 70B、Mistral 使用
• 多个查询头共享一组 KV 头
• KV Cache 直接缩小 4-8 倍

MQA（Multi-Query Attention）：

• 所有查询头共享一组 KV
• KV Cache 缩小到 1/n_heads
• 通常伴随微小的质量损失

KV Cache 量化：

• 将 KV Cache 从 FP16 量化到 INT8 或 INT4
• 缩小 2-4 倍
• 通常配合 per-channel 量化保持精度

KV Cache 剪枝：

• 丢弃注意力权重低的 token 的 KV 缓存
• 如 H2O（Heavy Hitter Oracle）：保留注意力分数最高的 token
• StreamingLLM：保留初始 token + 滑动窗口

5.3 前沿方法

Paged Attention（vLLM）：

• 借鉴操作系统的虚拟内存分页思想
• KV Cache 按需分配，避免内存碎片
• 支持动态长度和请求间共享

SnapKV：

• 观察到注意力模式在不同层之间高度相似
• 利用这种冗余压缩 KV Cache
• 在保持质量的同时大幅减少内存

六、评估长上下文

6.1 常用基准

基准	测试内容	长度范围
NIAH (Needle in a Haystack)	精确检索	1K - 1M
LongBench	多任务长文本	4K - 32K
L-Eval	长文本理解	4K - 100K
∞Bench	超长文本	128K+
RULER	多种长上下文能力	4K - 128K

6.2 NIAH 的局限

“大海捞针” 测试虽然直观，但有严重局限：

1. 只测检索，不测推理：找到一句话 ≠ 理解长文本
2. 与实际使用差距大：真实场景需要综合多段信息推理
3. 容易过拟合：针对性优化可以通过测试但实际无用

更有意义的评估是多跳推理（multi-hop reasoning）——需要结合文档不同部分的信息才能回答的问题。

6.3 有效上下文长度

一个重要概念是有效上下文长度——模型能真正利用的最大长度，通常远小于声称的上下文窗口。

声称支持 128K context
  ↓
NIAH 测试通过 128K
  ↓
但多跳推理能力在 32K 后显著下降
  ↓
有效上下文长度 ≈ 32K

七、工程实践建议

7.1 长上下文训练清单

1. 选择位置编码方案：首选 RoPE + ABF 或 YaRN
2. 使用 Flash Attention：必须，否则训练不现实
3. 渐进式训练：短 → 中 → 长，逐步增加
4. 准备长文本数据：至少需要目标长度的数据
5. 序列并行：超长序列需要跨 GPU 分布
6. KV Cache 优化：GQA 或 MQA，按需启用量化

7.2 上下文长度选择

应用场景	推荐上下文长度	原因
聊天对话	8K - 32K	通常不需要太长
文档 QA	32K - 128K	覆盖大多数单文档
代码理解	64K - 256K	代码仓库可能很长
多文档推理	128K - 1M	需要同时参考多份文档
书籍/长视频	1M+	内容本身就很长

7.3 推理时的性能优化

长上下文推理的关键优化：

1. Prefill 和 Decode 分离：长 Prompt 的 Prefill 是计算密集型，Decode 是带宽密集型
2. Chunked Prefill：将长 Prompt 分块处理，避免 OOM
3. Speculative Decoding：对长上下文特别有效，因为检验比生成快
4. KV Cache 预加载：对固定的系统 Prompt 预计算 KV Cache

八、展望

长上下文技术的发展方向：

1. 10M+ Token：Gemini 已在探索的方向，需要根本性的架构创新
2. 亚线性注意力：将注意力复杂度降到 $O(n \log n)$ 或 $O(n)$
3. 记忆增强：显式记忆模块存储长距离信息
4. 分层理解：不同粒度的信息用不同机制处理
5. 与 RAG 的融合：长上下文处理结构化检索结果

长上下文能力的提升正在深刻改变 LLM 的应用边界——从 “问答工具” 向 “认知引擎” 的转变。