RoPE 想扩上下文只有三条路:硬外推(崩,见第 5110 篇)、位置插值 PI(把全部维度一起压,糊,见第 5090 篇)、以及这篇的主角——按维度区别对待。核心观察很朴素:RoPE 的维度对本来就有波长分工(第 5070 篇),插值该只动负责长程的低频维,负责局部精度的高频维一个相位都不该挪。NTK 和 YaRN 都是这个观察的执行方案,只是一个从核函数视角偷巧,一个把它做彻底。
NTK 缩放:换个底数就完了
设要把上下文扩 $s$ 倍。NTK-aware 缩放(最初来自 Reddit 讨论,被 Llama-2 生态广泛采用)的全部操作是把 RoPE 底数换掉:
$$ \mathrm{base}' = \mathrm{base}\cdot s^{\frac{d}{d-2}} \approx \mathrm{base}\cdot s \quad (d \gg 2). $$
表面上和重训一个长 base 一样,效果完全不同,看逐维频率的变化。第 $i$ 对的新角频率:
$$ \theta_i' = (\mathrm{base}')^{-2i/d} = \theta_i \cdot s^{-\frac{2i}{d-2}}. $$
- $i\to0$(最高频):乘子 $s^0 = 1$,频率不动;
- $i\to d/2$(最低频):乘子 $s^{-d/(d-2)}\approx 1/s$,波长恰好拉伸 $s$ 倍,等价于位置刻度整体除以 $s$——正是 PI 想要的效果;
- 中间维度:按 $s^{-2i/(d-2)}$ 幂律平滑过渡。
一个底数改动,同时实现"高频保真、低频插值",而且位置 $m$ 不用动——全部信息塞进谱线里。这就是它受欢迎的原因:改一行配置,不动权重格式,不动 kernel。
"NTK" 这个名字的来历顺带一提:推导者把 RoPE 的内积看作核函数 $K(m,n)$,直接外推相当于在训练支撑点之外求核值,必然发散;缩放底数相当于对核的频谱做局部拉伸,让训练区间外的核值落回已拟合区域。动机不同,结论和上面的逐维分析一致。
Dynamic NTK:边生成边改底数
静态 NTK 按目标长度定死 $\mathrm{base}'$。推理时更细的做法是动态算:当前序列长度 $m$ 超过原训练长度 $L$ 时取
$$ s(m) = \frac{m}{L},\qquad \mathrm{base}'(m) = \mathrm{base}\cdot s(m)^{\frac{d}{d-2}}, $$
每步解码重算角度表。$m \le L$ 时 $s=1$,原行为一字不差;越界后平滑过渡。代价是 cos/sin 表不能静态缓存(但第 5060 篇说过表是确定性函数,重算很便宜),且 batch 内不同长度的序列要分组。HF Transformers 的 dynamic_rope 模式实现的就是这个。
YaRN:把幂律过渡换成显式斜坡
NTK 的幂律过渡是"碰巧够用",YaRN(Peng et al. 2023)指出两个可以更好的地方,并给出修正。
其一,过渡形状。 NTK 的乘子随维度平滑渐变,但高频端其实不该做任何缩放——YaRN 把维度分成三段:波长 $\lambda_i \le L_{\text{keep}}$ 的维度(纯局部用途)完全不插值;$\lambda_i \ge L$ 的维度做满额 $1/s$ 插值;中间按波长线性斜坡过渡:
$$ h_i = \begin{cases} 1, & \lambda_i \le L_{\text{keep}},\\[4pt] 1-\dfrac{\lambda_i - L_{\text{keep}}}{L - L_{\text{keep}}}\cdot\Big(1-\frac{1}{s}\Big), & L_{\text{keep}} < \lambda_i < L,\\[6pt] 1/s, & \lambda_i \ge L,\end{cases} $$
位置缩放为 $m_i' = m\cdot h_i$(逐维度不同——实现上等价于逐维改角度)。$L_{\text{keep}}$ 一般取 1 附近,即"周期只有几个位置的维度不许动"。
其二,插值的熵副作用。 把低频相位压缩 $s$ 倍后,远距离 token 的注意力 logit 方差整体缩小,注意力分布变平(熵增大),长程检索精度反而受损。YaRN 给注意力打分加一个温度因子补偿:
$$ t = 1 + 0.1\ln s,\qquad a_{ij} \leftarrow a_{ij}\,/\,t . $$
效果上等价于把 $q$ 缩放 $1/\sqrt{t}$,让打分分布回到训练时的锐度。
微调成本。 PI 需要在大段长文本上重训才能找回质量;YaRN 由于高频维完全没动、且温度因子抵消了分布漂移,论文报告用预训练成本的约 0.1% 做少量步数微调即达目标长度,甚至部分设置下免微调可用。这套" ramp + 温度"组合后被 Gemma、Qwen 等长上下文版本大量采用。
谱系小结
三条路的对比可以压成一张图:横轴是维度波长(左短右长),纵轴是位置缩放因子 $h$。PI 是一条平线(所有维度一律 $1/s$);NTK 是从 1 滑到 $1/s$ 的幂律曲线;YaRN 是"平台—斜坡—平台"的分段线。三者都在回答同一个问题——该动哪些维度——只是答案的激进程度递增。下一篇回到最朴素的 PI,把"直接压缩位置刻度"的账算清楚,看看它到底糊在哪。