跳到主要内容

位置编码与外推

核心要点

  • 位置编码给 attention 注入序列位置信息,RoPE 是主流
  • 外推指超训练长度仍保持质量,直接外推会退化
  • 三类外推:训练友好 / 推理调整 / 短训长推
  • 短训长推性价比最高:少量微调扩到 128K-2M
  • 工业主流:RoPE + YaRN / LongRoPE 路线

本文回答 02-第一性挑战 中"位置外推困难"这一第一性问题,不涉及训练侧的长数据合成(→ 06-训练侧)。

位置编码做什么

Transformer 的 attention 本质是集合运算 —— 把 token 序列打乱顺序,attention 输出完全一致。位置信息必须显式注入,否则模型分不清 "猫追狗" 和 "狗追猫"。

注入方式代表方法核心特征
绝对位置学习式位置 embedding (GPT-2 / BERT)、正弦余弦 (Vanilla Transformer)给每个位置一个独立向量加到 token embedding 上
相对位置T5 相对位置 bias、RoPE注入 token 对之间的相对距离
注意力偏置ALiBi在 attention score 上加距离衰减项,不改 embedding

@tbl-longctx-pe-injection-modes 位置编码的三类注入方式

长上下文场景下,绝对位置编码的劣势暴露:训练时只见过 $[0, n_{\text{train}})$ 区间的 embedding,超过这个范围模型无法处理。相对位置方案外推性更好,是当前长上下文模型的标准选择。

RoPE:当前主流方案

核心问题:为什么相对位置编码 RoPE 取代了绝对位置编码成为主流?

RoFormer[1] 提出旋转位置编码(Rotary Position Embedding),把位置 $m$ 编码为对 query/key 向量做复平面旋转

$$\begin{equation} \tilde q_m = R_{\Theta, m} \, q_m, \quad \tilde k_n = R_{\Theta, n} \, k_n \label{eq:lc-pe-rope-rotation} \end{equation}$$

其中 $R_{\Theta, m}$ 是按位置 $m$ 构造的分块对角旋转矩阵,每个 $2 \times 2$ 块旋转角度 $m \theta_i$,频率序列 $\theta_i = 10000^{-2i/d}$

关键性质:旋转后的 query-key 内积仅依赖相对位置 $m - n$

$$\begin{equation} \tilde q_m^\top \tilde k_n = q_m^\top R_{\Theta, n-m} k_n \label{eq:lc-pe-rope-relative} \end{equation}$$

RoPE 因为如下三点成为主流:

  1. 相对位置:内积只依赖距离,天然适合 attention 的位置语义
  2. 简洁高效:实现成本低,与 FlashAttention 等 kernel 兼容
  3. 频率维度独立:不同维度有不同旋转频率(高频对近距离敏感、低频对远距离敏感),是后续外推方法的发力点

LLaMA、Qwen、DeepSeek、Mistral 等几乎所有现代开源模型用 RoPE。

RoPE 的外推问题

直接把 4K 训练的 RoPE 模型用在 32K 上,perplexity 会显著上升。原因有二:

失败模式机制
未见频率组合:长位置的低频维度旋转角 $m \theta_i$ 超出训练分布高频维度(小 $\theta_i$ 对应快旋转)训练时已周期性重复见过;低频维度(大 $\theta_i$)训练时只见过窄角度范围,外推到大位置进入未见区域
注意力 logit 失控:长距离 logit 偏小或偏大,softmax 退化RoPE 不像 ALiBi 自带距离衰减,长序列下偶发的 logit 异常会被 softmax 放大成"指针漂移"

@tbl-longctx-pe-rope-extrap-failure RoPE 直接外推的两类失败模式

ALiBi:训练时就为外推设计

核心问题:在训练阶段就内置外推能力,能否替代 RoPE 加外推方法?

Press et al., 2021[2] 提出 ALiBi(Attention with Linear Biases),不在 embedding 注入位置,而在 attention score 上加一个线性距离惩罚

$$\begin{equation} \text{score}(q_m, k_n) = q_m^\top k_n - s_h \cdot |m - n| \label{eq:lc-pe-alibi} \end{equation}$$

其中 $s_h$ 是每个 head 的斜率(不同 head 不同,等比递减;ALiBi 原论文用 $m$ 表示斜率,本文为避免与位置变量 $m$ 撞名改写为 $s_h$)。距离越远,得分被压得越低,自然有"近优先"的归纳偏置。

优势劣势
训练时即具备一定外推性(4K 训练 → 推理可用到 16K)远距离信息几乎被 ALiBi 压没,真正长依赖任务退化
实现简单,无频率维度缺乏 RoPE 的"相对位置丰富表示"

@tbl-longctx-pe-alibi ALiBi 的优劣

实际案例

  • 原论文数据[2]:1.3B 模型训练长度 1024 token,外推到 2048 时 PPL 为 8.92,低于训练时在 1024 上的 PPL 9.16;与在 2048 上全量训练的 sinusoidal 模型质量相当,且训练快 11%。有效外推约 6–10× 训练长度。
  • MPT-7B[3]:基础模型训 2048 token;StoryWriter 变体结合 ALiBi 微调到 65k,实测可外推到 84k token
  • 边界:过度预训练(预训练 token 量极大)的模型外推上限收窄到约 1.12×;结合 Position Interpolation 可达约 2×。

业界现状:早期 MPT、BLOOM 用 ALiBi;后续主流(LLaMA 系、Qwen、DeepSeek)回到 RoPE 路线。原因是 ALiBi 在真正长依赖任务上比不过 RoPE + 外推方法。

RoPE 外推的三类方法

核心问题:如何在不重新训练或仅少量微调下,把 RoPE 模型扩到更长上下文?

推理时调整:NTK-aware Scaling

/u/bloc97 在 Reddit 提出 NTK-aware Scaling[4],核心想法:不在低频维度上线性压缩,而是在高频维度上微调频率基

Position Interpolation (PI)NTK-awareYaRN
把位置 $m$ 等比缩放到 $m / s$$s$ 为扩展倍数)修改 $\theta$ 基数:$\theta_i' = \theta_i \cdot s^{-2i/(d-2)}$(仅低频维度受影响)分段:高频不变,低频插值,再加 attention 温度修正
零训练,简单零训练,效果优于 PI极少量微调,效果最佳

@tbl-longctx-pe-rope-extension-methods RoPE 外推三类方法的核心机制

Position Interpolation (PI)[5] 的思路最直接:把推理时的位置 $m$ 缩到训练范围 $[0, n_{\text{train}})$ 内:

$$\begin{equation} \tilde m = m \cdot \frac{n_{\text{train}}}{n_{\text{infer}}} \label{eq:lc-pe-pi} \end{equation}$$

零训练即可把 LLaMA 1 从 2K 扩到 8K。但所有维度等比压缩,精细分辨率下降——分不清相邻几个 token 的位置。

NTK-aware[4] 改进点:只在低频维度做插值,高频维度保持原频率。物理直观:高频对应"近距离细节",不需要扩;低频对应"远距离粗粒度",扩这里就够了。零训练效果优于 PI。

短训长推:YaRN

YaRN[6] 把 NTK-aware 思路进一步精化,分三段处理频率

频率区间处理
高频(已见周期足够多)保持原值,不插值
低频(未见周期)等比缩放(PI)
中频(部分见过)NTK 风格插值(平滑过渡)

加上 attention 温度修正(针对长上下文 logit 失控),YaRN 在仅 0.1% 训练 token 微调下就能把 4K 模型扩到 128K。

实际影响:YaRN 是 2024 年开源社区主流外推方案。LLaMA 3.1 把上下文从 8K 扩到 128K 用了类似思路;Qwen2.5-Turbo 1M 上下文也基于这类方法。

短训长推:LongRoPE

LongRoPE[7] 进一步把上限推到 2M+。核心想法:

  1. 非均匀插值搜索:不假设所有低频维度等比缩放,用搜索算法给每个维度找最优插值因子
  2. 渐进扩展:先扩到 256K 微调,再扩到 2M 微调,分阶段降低训练显存
  3. 短上下文恢复:在最终阶段加少量短上下文 SFT,避免模型只擅长长上下文

LongRoPE 把 LLaMA 2-7B 扩到 2M 上下文,是开源界第一个公开报告 2M 的方案。

外推方法选型对照

核心问题:给定目标长度与训练预算,应该选哪种外推方法?

方法训练成本上限工业采用情况
直接外推0$\le 1.5 \times$ 训练长度不推荐
Position Interpolation0$4 \times$(高于则失真)已被 NTK / YaRN 取代
NTK-aware0$8 \times$推理快速验证
ALiBi0(训练即用)约 6–10×(有效),可达 40×(微调,MPT-7B 84k)早期 MPT / BLOOM;现已少用
YaRN0.1% 训练 token$32 \times$+LLaMA 3.1 / Qwen2.5 / DeepSeek 风格
LongRoPE多阶段微调$500 \times$+(2M+)学术领先,工程门槛较高

@tbl-longctx-pe-method-comparison 主流外推方法的训练成本与上限对比

当前实践共识

  • 训练强模型选 RoPE + YaRN 路线,性价比最高
  • 极长上下文(2M+)研究用 LongRoPE 系思路 + 多阶段微调
  • ALiBi 与原版 PI 已逐步退出主流
  • 外推方法只解决"能看长",看不"看得懂长" —— 见 02-第一性挑战 的中段遗忘部分

Takeaway

知识点核心结论
位置编码作用attention 是集合运算,必须显式注入位置信息
主流方案RoPE 用相对位置内积,外推性优于绝对位置编码
直接外推超训练长度 1.5× 即退化,不可用
推理时调整NTK-aware 零训练扩到约 8×,用于快速验证
短训长推(≤128K)RoPE + YaRN,0.1% token 微调,性价比最高
极长上下文(2M+)LongRoPE 非均匀插值 + 多阶段微调
ALiBi训练即外推,但长依赖任务退化,已少用

@tbl-longctx-pe-takeaway 全文要点

延伸阅读

参考资料

  1. Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding, arXiv 2021. https://arxiv.org/abs/2104.09864
  2. Press et al., Train Short, Test Long: Attention with Linear Biases (ALiBi), arXiv 2021. https://arxiv.org/abs/2108.12409
  3. MosaicML, Introducing MPT-7B: A New Standard for Open-Source, Commercially Usable LLMs, 2023. https://www.databricks.com/blog/mpt-7b
  4. bloc97, NTK-Aware Scaled RoPE (Reddit post), 2023. Reddit 受反爬限制可能 403;技术细节亦见 EleutherAI Blog: Extending the RoPE 的回顾段。https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
  5. Chen et al., Extending Context Window of Large Language Models via Positional Interpolation, arXiv 2023. https://arxiv.org/abs/2306.15595
  6. Peng et al., YaRN: Efficient Context Window Extension of Large Language Models, arXiv 2023. https://arxiv.org/abs/2309.00071
  7. Ding et al., LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens, arXiv 2024. https://arxiv.org/abs/2402.13753