位置编码改进方案 - RoPE外推与YaRN

RoPE外推问题的本质

旋转位置编码回顾

RoPE（Rotary Position Embedding）通过旋转矩阵将位置信息注入注意力计算：

$f(q, m) = q \cdot e^{i \cdot m \cdot \theta_j} = \begin{pmatrix} q_{2j} \\ q_{2j+1} \end{pmatrix} \cdot \begin{pmatrix} \cos(m\theta_j) & -\sin(m\theta_j) \\ \sin(m\theta_j) & \cos(m\theta_j) \end{pmatrix}$

其中$\theta_j = 10000^{-2j/d}$为旋转基频，$j$为维度索引。

外推失效的数学解释

当序列长度从训练时的$L_{\text{train}}$外推到$L_{\text{test}}$时，最大旋转角度：

• 训练时：$\phi_{\text{max}}^{\text{train}} = L_{\text{train}} \cdot \theta_0 = L_{\text{train}} \cdot 1$
• 测试时：$\phi_{\text{max}}^{\text{test}} = L_{\text{test}} \cdot \theta_0$

对于4K训练→128K测试： $\phi_{\text{max}}^{\text{test}} = 32 \cdot \phi_{\text{max}}^{\text{train}}$

这意味着高频分量（低$j$）经历了32倍的旋转角度，完全超出了训练分布。

xychart-beta
    title "不同维度的旋转角度随序列长度变化"
    x-axis "位置索引" 0 --> 128000
    y-axis "旋转角度 (弧度)" 0 --> 100
    line [0, 0.001, 0.002, 0.004, 0.008, 0.016, 0.032, 0.064]
    line [0, 0.01, 0.02, 0.04, 0.08, 0.16, 0.32, 0.64]
    line [0, 0.1, 0.2, 0.4, 0.8, 1.6, 3.2, 6.4]
    line [0, 1, 2, 4, 8, 16, 32, 64]
    
    annotation "dim=0" at 128000, 64
    annotation "dim=64" at 128000, 6.4
    annotation "dim=128" at 128000, 0.64

注意力分数分布偏移

位置编码外推导致query-key点积的分布发生偏移：

• 训练时：$QK^T \sim \mathcal{N}(0, \sigma^2_{\text{train}})$
• 外推时：$QK^T \sim \mathcal{N}(\mu_{\text{shift}}, \sigma^2_{\text{test}})$

Softmax对分布偏移极其敏感，轻微的均值偏移会导致注意力权重完全崩溃。

位置插值（Position Interpolation）

线性插值原理

位置插值（PI）通过缩小位置索引来适配训练分布：

$m' = m \cdot \frac{L_{\text{train}}}{L_{\text{test}}}$

对于4K→16K外推： $m' = m \cdot \frac{4096}{16384} = m \cdot 0.25$

这样，位置16384在RoPE中被视为位置4096，落在训练分布内。

实验结果

在LLaMA-2-7B上的测试（困惑度，越低越好）：

关键发现：

• 位置插值无需微调即可实现2-4倍外推
• 仅需1K步微调即可实现8倍外推且性能损失<3%
• 微调后的外推能力可稳定扩展到训练长度的8-16倍

局限性与问题

位置插值存在两个主要问题：

1. 高频信息损失：缩小位置索引导致高频分量（小$j$）的相对位置区分度下降
2. 长距离相对位置模糊：远距离token的位置差异被压缩，难以建模长程依赖

实验显示，在需要精确长程依赖的任务（如长文档推理）上，PI的表现不如其他外推方案。

NTK-aware 缩放

神经切线核理论

NTK（Neural Tangent Kernel）理论描述了神经网络在初始化附近的核函数行为。bloc97发现，RoPE的旋转频率与NTK的频率参数存在对应关系。

动态频率调整

NTK-aware方案通过调整基频$\theta_j$来适应长序列：

$\theta'_j = \theta_j \cdot \alpha^{-2j/d}$

其中$\alpha = \frac{L_{\text{test}}}{L_{\text{train}}}$为外推倍数。

关键洞察：

• 高频分量（小$j$）：$\theta'_j \approx \theta_j \cdot \alpha^{-1}$，频率降低以适应大范围旋转
• 低频分量（大$j$）：$\theta'_j \approx \theta_j$，几乎不变以保持细粒度区分

与位置插值的对比

flowchart LR
    subgraph PI["位置插值 (PI)"]
        P1["缩小位置索引"]
        P2["保持频率不变"]
        P3["高频区分度↓"]
        P4["所有位置均匀压缩"]
    end
    
    subgraph NTK["NTK-aware"]
        N1["保持位置索引"]
        N2["调整频率分布"]
        N3["高频适应大范围"]
        N4["低频保持精度"]
    end

性能评估

在PG-19数据集上的困惑度对比：

**NTK-aware+**是改进版本，通过额外的缩放因子优化远距离注意力：

$\theta''_j = \theta_j \cdot \alpha^{-2j/d} \cdot \beta$

YaRN：最优外推方案

核心思想

YaRN（Yet another RoPE extensioN）结合了NTK-aware频率调整和注意力温度缩放，是当前最有效的RoPE外推方案。

两个关键改进：

1. 频率调整：与NTK-aware相同的动态频率调整
2. 注意力温度缩放：引入缩放因子$t$调整attention logits

$\text{Attention}(Q, K) = \text{softmax}\left(\frac{QK^T}{t \cdot \sqrt{d_k}}\right)$

温度缩放原理

外推时query-key点积的方差增大： $\text{Var}(QK^T)_{\text{test}} > \text{Var}(QK^T)_{\text{train}}$

温度缩放$t$补偿这一变化： $t = \sqrt{\frac{1}{s} \cdot \frac{\log(\alpha \cdot s)}{\log(s)}}$

其中$s$为注意力softmax的温度超参数，LLaMA-2中$s \approx 1$。

计算流程

flowchart TD
    A[输入序列<br/>长度L_test] --> B{计算缩放因子}
    B --> C[α = L_test / L_train]
    B --> D[t = sqrt(log(α*s)/log(s))]
    
    C --> E[调整RoPE频率<br/>θ' = θ * α^(-2j/d)]
    D --> F[缩放Attention Logits<br/>/ (t * sqrt(d_k))]
    
    E --> G[标准Attention计算]
    F --> G
    
    G --> H[输出]

实测性能

在LongChat基准测试上的准确率（%）：

关键结论：

• YaRN在32K长度下保持80%+准确率，而基线在16K时已崩溃
• 无需任何微调即可实现16倍外推
• 微调后（1K steps）可进一步扩展到32倍

超参数调优

YaRN有三个关键超参数：

• $\alpha$：基础外推倍数（通常设为$L_{\text{test}}/L_{\text{train}}$）
• $s$：温度缩放参数（LLaMA-2推荐$s=1$，其他模型可能需要调整）
• 频率混合比例：控制高频和低频的调整程度

网格搜索显示，对于大多数模型，$s \in [0.8, 1.2]$范围内性能最佳。

xPos与ALiBi：替代位置编码方案

xPos：指数位置编码

xPos通过引入指数衰减项改进RoPE：

$f(q, m) = q \cdot e^{i \cdot m \cdot \theta_j} \cdot \gamma^{-j/d}$

其中$\gamma$为衰减系数（通常$\gamma \approx 0.99$）。

优势：

• 天然抑制远距离噪声
• 外推能力优于标准RoPE

局限：

• 需要训练时使用xPos，无法直接应用于预训练模型
• 长距离依赖建模能力略弱

ALiBi：线性偏置注意力

ALiBi（Attention with Linear Biases）完全不使用位置编码，而是在attention score中直接加入位置偏置：

$\text{Attention}(Q, K)_{ij} = \frac{Q_i K_j^T}{\sqrt{d_k}} - m \cdot |i - j|$

其中$m$为负斜率，随head变化。

独特优势：

• 训练稳定：无需学习位置编码参数
• 外推能力强：线性偏置天然支持任意长度
• 实现简单：仅需修改attention mask

MPT-7B的应用： MosaicML的MPT-7B采用ALiBi位置编码，在训练时仅使用8K长度，但可直接推理到65K+长度而无需任何调整。

三种方案对比

实际部署指南

方案选择决策树

flowchart TD
    A[已有预训练模型?] -->|是| B[使用RoPE?]
    A -->|否| C[从头训练]
    
    B -->|是| D[YaRN推荐]
    B -->|否| E[ALiBi或xPos]
    
    C --> F[ALiBi: 稳定+外推强]
    C --> G[xPos: 精度优先]
    C --> H[YaRN+RoPE: 生态兼容]
    
    D --> I[无需重训<br/>直接应用]
    F --> J[重新训练<br/>长期收益]

代码实现示例

YaRN应用（Hugging Face Transformers）：

from transformers import LlamaForCausalLM, LlamaConfig
import torch

# 加载模型
model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b")

# 应用YaRN配置
original_max_position = 4096
target_max_position = 32768
scaling_factor = target_max_position / original_max_position

# 修改RoPE缩放
for layer in model.model.layers:
    layer.self_attn.rotary_emb.scaling_factor = scaling_factor
    layer.self_attn.rotary_emb.attn_factor = 1.0  # YaRN温度缩放
    
# 更新模型配置
model.config.max_position_embeddings = target_max_position

关键参数：

• scaling_factor：外推倍数（32K/4K = 8）
• attn_factor：温度缩放因子（通常1.0-1.2）
• 无需修改模型权重，纯配置即可生效

参考资料

1. NTK-Aware Scaled RoPE (bloc97, 2023)

   • NTK-aware原始实现与讨论

2. YaRN: Efficient Context Window Extension of Large Language Models (Peng et al., 2023)

   • YaRN完整论文，包含理论分析和实验验证

3. Extending Context Window of Large Language Models via Position Interpolation (Chen et al., 2023)

   • 位置插值原始论文

4. xPos: Extrapolatable Position Embeddings (Sun et al., 2022)

   • xPos设计与实验

5. Train Short, Test Long: Attention with Linear Biases (Press et al., 2021)

   • ALiBi原始论文

6. MPT-7B: A New Standard for Open-Source, Commercially Usable LLMs (MosaicML, 2023)

   • ALiBi的工程实践与性能数据