三十三：Long Context~NTK-aware 插值

来源:http://mp.weixin.qq.com/s?__biz=MzYyNTk3Njg1NA==&mid=2247484394&idx=1&sn=c9ec222c33dfa9ed4b4f50ccc4be4340&chksm=f01eb093c76939857bb8eba1487b90a7cd7e917822179ad0a9e125b243aca72c8b5cb2b86dcc#rd

1. 学习范围

本日主题是 RoPE 长上下文扩展中的插值方法，重点是 NTK-aware 插值。今天接着 Day 32 的 RoPE 复习，讨论当模型训练长度有限、推理时想扩大上下文窗口时，为什么不能只把长度配置改大，以及常见插值方法如何改变 RoPE 的位置到相位映射。 需要掌握： - RoPE 外推失败的原因。

• 直接外推、线性位置插值、进制转换理解的区别。

• RoPE 可以被理解为一种 β 进制位置编码。

• Position Interpolation 的公式和分辨率损失。

• NTK-aware scaling / interpolation 的核心思想。

• 动态 NTK scaling 的工程公式。

• 插值方法的评估、适用场景和常见坑。

2. 参考资料

• 从 ALiBi、PI、NTK-aware 到 YaRN、S2-Attention：https://blog.csdn.net/v_JULY_v/article/details/135072211

• 苏剑林：RoPE 是一种 β 进制编码：https://kexue.fm/archives/9675

• RoFormer 原论文：https://arxiv.org/abs/2104.09864

• Position Interpolation 论文：https://arxiv.org/abs/2306.15595

• Hugging Face RoPE utilities：https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_rope_utils.py

3. 长上下文扩展的问题定义

假设一个 LLM 训练时最大上下文长度为：

L_train = 4096

现在希望推理时支持：

L_target = 32768

直觉上，RoPE 没有学习式位置表，可以为任意位置 m 计算 cos(m theta) 和 sin(m theta)。但这只说明公式可计算，不说明模型能可靠理解。模型训练时只见过 0..L_train-1 的位置、相位组合和长距离依赖模式；直接推到 L_target 会遇到分布外的位置。 长上下文扩展要解决的是：

如何让模型在更长序列中使用 RoPE，
同时尽量保留训练长度内的短上下文能力。

4. RoPE 的频率形式

RoPE 常用频率写法：

inv_freq_i = 1 / base^(2i / d)
angle(m, i) = m * inv_freq_i

其中： - m 是 position id。

• i 是二维旋转子空间编号。

• d 是旋转维度。

• base 通常为 10000，也常叫 rope_theta。

对每个二维子空间：

R(m, i) =
  [ cos(m * inv_freq_i)  -sin(m * inv_freq_i) ]
  [ sin(m * inv_freq_i)   cos(m * inv_freq_i) ]

位置越大，旋转角越大；频率越高，角度变化越快。长上下文扩展的本质，就是修改 m、base 或 inv_freq_i，使更长的位置范围落在模型更容易泛化的相位空间里。

5. 直接外推

直接外推不改 RoPE 公式，只把最大长度配置调大：

angle_new(m, i) = m * inv_freq_i

优点： - 实现最简单。

• 训练长度内的位置完全不变。

• 短上下文能力通常不受影响。

缺点： - 超出训练长度后，模型看到未训练过的相位。

• 高频维度相位快速绕圈，远距离 attention score 可能异常。

• 长文档任务中容易出现“看得到但用不好”的问题。

直接外推适合小幅扩展和初步实验，不适合作为可靠长上下文方案。

6. 线性位置插值

Position Interpolation 的核心是压缩 position：

s = L_target / L_train
m' = m / s
angle_new(m, i) = (m / s) * inv_freq_i

例如从 4k 扩到 32k，s = 8。推理位置 m=32000 会映射到原始位置 4000 附近，因此模型看到的位置范围仍接近训练范围。 优点： - 避免直接进入很远的未见位置。

• 数学和实现都简单。

• 配合少量长上下文微调常有稳定效果。

缺点： - 所有距离都被压缩，局部位置分辨率下降。

• 相邻 token 的相位差变小，短距离模式可能受损。

• 压缩比例越大，越依赖长上下文微调修复。

线性插值的面试表达：

PI 不是让 RoPE 直接外推到更大 position，而是把更长位置缩放回训练窗口内，用牺牲位置分辨率换取更稳定的相位分布。

7. RoPE 的 β 进制理解

苏剑林提出的一个有用理解是：RoPE 可以看成一种 β 进制位置编码。位置 m 被不同频率维度编码，不同维度类似不同“位数”的数字。 类比十进制：

m = digit_0 * 10^0 + digit_1 * 10^1 + digit_2 * 10^2 + ...

RoPE 中的不同频率维度对应不同尺度：

angle_i = m / beta^(2i/d)

低维高频像低位数字，对小距离变化敏感；高维低频像高位数字，对长距离变化更稳定。扩大上下文可以被理解为：原来的 β 进制容量不够时，需要改变进制或者改变位置映射。 这个视角能解释两类方法： - 线性插值：把数字整体缩小后再放入原来的编码系统。

• NTK-aware：改变 β / base，让编码系统本身能覆盖更长范围。

8. NTK-aware scaling 的核心思想

NTK-aware scaling 不直接把所有 position 除以同一个系数，而是调整 RoPE 的 base，改变不同频率维度的波长。直接增大编码系统的进制基底 base！ 原始：

inv_freq_i = 1 / base^(2i / d)

扩展后：

base_new = base * scale_base
inv_freq_new_i = 1 / base_new^(2i / d)

当 base_new 变大时： - i=0 的最高频维度基本不变，因为指数为 0。

• i 越大，频率下降越明显，波长变长。

• 局部位置能力保留得比线性插值更好。

• 长距离维度可以覆盖更长上下文。

这就是 NTK-aware 方法吸引人的地方：它更像“改变进制”，而不是把所有位置统一压缩。

9. 常见 NTK-aware 公式

静态 NTK-aware scaling 常见近似形式：

scale = L_target / L_train
base_new = base * scale^(d / (d - 2))

动态 NTK scaling 常见工程形式：

if seq_len <= L_train:
    base_new = base   # 不超训练长度，完全使用原始RoPE，无性能损失
else:
    base_new = base * ((factor * seq_len / L_train) - (factor - 1))^(d / (d - 2))

其中： - factor 是目标扩展倍率。

• seq_len 是当前实际序列长度。

• d 是 rotary dimension。

• base_new 用来重新计算 inv_freq。

动态版本的好处是：当输入没有超过训练长度时，不改变 RoPE；超过训练长度时，再随着当前长度调整 base。

10. NTK-aware 与线性插值对比

方法 修改对象 局部短距离能力 长距离覆盖效果 实现难度 核心风险 / 缺陷 直接外推 不改动 RoPE 公式、不缩放 position、不修改 base，仅放开最大长度限制 完全保留，无任何损失 极差，超长位置属于训练外相位，高频维度相位缠绕，远距离注意力失效 最低，仅修改长度超参 长文本远端信息无法有效利用，仅适合小幅临时实验 线性 PI 插值 固定缩放位置：m′=m/s，base 保持不变 全局分辨率下降，相邻 token 相位差缩小，短句性能下滑 稳定可控，所有位置映射至训练区间，配合微调效果稳定 低，仅对位置做除法缩放 扩容倍数越大，短文本退化越明显，一般必须长文本微调修复 静态 NTK-aware 固定放大基底：basenew=base⋅scaled−2d，position 不变 保留较好，高频维度几乎不受影响，仅轻微损失短句效果 较好，低频维度波长同步拉长，天然适配长依赖 中等，需要预计算新的逆频率数组 短序列推理仍使用放大后的 base，短句轻微退化；倍率设置不当会相位偏移 动态 NTK 根据输入 seq_len 自适应调整 base，未超训练长度则使用原始 base 最优，≤训练长度时完全原生 RoPE，无退化 灵活自适应，序列越长基底放大越多，适配任意超长输入 中等偏高，推理分支判断，KV 缓存复用逻辑更繁琐 推理时实时更新 inv_freq，增加计算开销；缓存管理逻辑复杂，工程易出 bug 简单记忆：

PI: 缩小 position。
NTK-aware: 调大 base / rope_theta。
Dynamic NTK: 只在实际超长时调大 base。

11. 工程实现流程

典型实现步骤：

def compute_inv_freq(base, dim):
    return 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))

def dynamic_ntk_base(base, seq_len, train_len, factor, dim):
    if seq_len <= train_len:
        return base
    scale = (factor * seq_len / train_len) - (factor - 1)
    return base * (scale ** (dim / (dim - 2)))

base_new = dynamic_ntk_base(
    base=10000,
    seq_len=current_seq_len,
    train_len=4096,
    factor=8,
    dim=rotary_dim,
)
inv_freq = compute_inv_freq(base_new, rotary_dim)

应用 RoPE：

position_ids -> freqs = position_ids * inv_freq
freqs -> cos/sin cache
q/k -> apply_rope(q/k, cos, sin)

如果使用 KV cache，要保证 prefill 和 decode 阶段使用一致的 scaling 策略。尤其是动态 NTK，如果后续 token 导致 seq_len 增长并改变了 base_new，历史 cache 中已旋转的 K 是否需要重算，是实现中必须处理的问题。因此很多部署会在 prefill 时根据目标长度固定 scaling，避免生成中途改变 RoPE 参数。

12. 评估方法

RoPE 插值是否有效，要同时评估： - 短上下文任务是否掉点。

• 长上下文 needle-in-a-haystack 是否能定位。

• 长文档问答是否能引用正确证据。

• 多跳长依赖是否稳定。

• PPL 是否随长度平滑变化。

• 推理显存、KV cache 和延迟是否可接受。

常见实验分桶：

0-4k: 原训练长度内能力
4k-8k: 轻度扩展
8k-32k: 中长上下文
32k+: 极长上下文压力测试

只测最大长度是不够的，因为有些 scaling 方法会牺牲短上下文能力。

13. 常见坑

只改 max_position_embeddings： 配置变大不等于模型学会长上下文。需要同步修改 RoPE scaling，并做长上下文验证。 使用错误的 rotary_dim： 公式中的 d 应该是实际 RoPE 旋转维度，不一定等于 hidden size。 训练和推理 scaling 不一致： 如果模型微调时使用一种 scaling，推理时换另一种，效果可能明显下降。 动态 NTK 与 KV cache 不一致： 生成中途改变 base，会让历史 K 和新 Q/K 使用不同相位体系。 忽视短上下文回归： 长上下文指标上升不代表整体模型更好。上线前必须回测常规短任务。

14. 面试表达模板

可以这样回答 NTK-aware 插值：

RoPE 长上下文扩展的本质是修改 position 到旋转相位的映射。线性位置插值是把 m 变成 m/s，让更长序列压缩回训练长度内，但会降低局部位置分辨率。NTK-aware scaling 则更像改变 RoPE 的 base，也就是改变不同维度的频率分布。base 变大后，高维低频的波长变长，能覆盖更远距离；低维高频受影响较小，所以局部位置能力保留得更好。

工程上常见做法是根据扩展倍率或当前 seq_len 计算 base_new，再用 base_new 重新生成 inv_freq 和 cos/sin cache。要注意 rotary_dim、KV cache、训练推理一致性，以及短上下文回归测试。

15. 紧凑总结

• 直接外推：公式不变，但长位置是分布外。

• 线性 PI：m -> m/s，稳定但牺牲分辨率。

• RoPE 可理解为 β 进制位置编码，不同频率像不同位数。

• NTK-aware：调大 base / rope_theta，改变频率分布。

• 动态 NTK：输入超出训练长度后才按当前长度调整 base。

• 插值不是只改配置，必须配合正确实现和长短上下文评估。

          预览时标签不可点
  

  <div class="