YaRN位置插值实现128K长上下文方案

可视化演示

展示RoPE频率插值、分段缩放策略与注意力温度调整，切换模式查看YaRN核心技术。

RoPE 旋转编码

老王的工具箱里有一份祖传的“齿轮转速表”：64 对齿轮塔，第 i 对的转速是 10000 的 -2i/d 次幂。编号 0 的塔一转就是一圈，编号 63 的塔要转上万圈才走一格——从秒针到时针，一整套“速度频谱”就靠这一条公式铺开。大语言模型里的 RoPE（旋转位置编码）用的正是同一套逻辑：位置不是贴标签，而是让向量旋转一个角度。

具体来说，RoPE 把每个 token 的 query（查询）和 key（键）向量按维度两两配对，第 i 对分量在位置 p 处旋转 p·θ_i 弧度。i 越小旋转越快，i 越大旋转越慢，这样不同维度各自负责不同粒度的距离感：快齿轮管“挨得近不近”，慢齿轮管“离得多远”。

$$\displaystyle \theta_i = b^{-\frac{2i}{d}}$$

1. 全局基数 b：b=10000 是钟楼的总动力刻度，决定所有齿轮转速的量级，也决定位置编码的整体“分辨率”。

$$\displaystyle b = 10000$$

2. 维度对编号 i：i 从 0 到 d/2−1 标记第几对齿轮塔，配对的两个分量共享同一转速，协同完成二维旋转。

$$\displaystyle i \in \left\{0, 1, \dots, \frac{d}{2}-1\right\}$$

3. 指数式转速：转速随 i 按基数 b 的 −2i/d 次幂指数衰减，形成从“秒针”到“时针”的完整速度频谱，token 的位置 p 就编码为累计旋转角 p·θ_i。

$$\displaystyle \theta_i = b^{-\frac{2i}{d}}$$

理解了转速表，就明白了为什么“位置”在 RoPE 里不是一张查询表，而是一组会旋转的指针。接下来要回答的问题是：当指针转出钟面，会发生什么？

外推成败关键

镇长说“我要 128K 小时的钟面”，第一个冒出来的馊主意是：什么都不改，让指针接着转。毕竟齿轮没坏，为什么不能多走几圈？可钟面根本没有那些刻度——指针指向 3000 小时的刻度时，整个钟楼的人都不知道那是什么时间。大模型外推时的情况一模一样：位置 p 超过训练长度，旋转角 p·θ_i 进入训练时从未见过的数值区间。

注意力机制要计算两个 token 的相关性，RoPE 下这个相关性只取决于相对位置 p−q：把 key 旋转 (p−q) 的角度再和 query 点积。训练时，模型只见过 (p−q) 在 2048 以内的旋转组合；一旦相对距离超过窗口，高频维度已经转了成百上千圈，旋转角落在哪一象限完全不可预测，点积结果剧烈抖动，softmax 后的注意力分数要么爆炸、要么错乱。

$$\displaystyle a_{pq} = \mathrm{softmax}\left( \frac{q_p^\top R(p-q)\, k_q}{\sqrt{d}} \right)$$

1. 相对位置旋转：R(p−q) 把 key 旋转到与 query 对齐的坐标系，让注意力只依赖两个 token 的距离，而不是绝对位置。

$$\displaystyle R(p-q) = \prod_{i} R_i\big((p-q)\theta_i\big)$$

2. 点积相似度：旋转后的 query 与 key 做点积，数值越大表示该位置的 token 与当前查询越相关。

$$\displaystyle q_p^\top R(p-q)\, k_q$$

3. softmax 归一：先除以 √d 稳定方差，再做 softmax，得到当前 query 对所有 key 的注意力权重分布。

$$\displaystyle a_{pq} = \mathrm{softmax}\Big(\frac{q_p^\top R(p-q)\, k_q}{\sqrt{d}}\Big)$$

外推行不通，指针在空白刻度上乱摆，钟楼发出刺耳噪音。于是有人想到：与其让指针乱转，不如把整个 128K 小时压缩进 2048 的钟面——这就是朴素位置插值。

朴素位置插值

第二个方案粗暴但直观：把所有齿轮统一减速 64 倍。原来走 64 个小时的路程，现在只走 1 个小时，这样 128K 个小时就能全部映射进 2048 的钟面。训练时见过的每个刻度都能对上，模型不至于完全失智。但代价立刻显现：秒针几乎不转，12:01 和 12:02 的夹角肉眼不可见。

朴素位置插值（naive Position Interpolation, PI）在数学上就是把每个位置的旋转角乘以 N_old/N_new = 1/s。对高频维度（小 i）来说，它们原本负责区分相邻 token，被压缩后相邻位置的旋转角差趋近于零——短距离的“分钟精度”彻底丢失。想恢复精度，只能靠大规模重训让模型重新适应被抹平的角度差异，成本极高。

$$\displaystyle \theta_i^{PI} = \theta_i \cdot \frac{1}{s} = \frac{\theta_i}{s}$$

1. 缩放比 s：s 是新窗口与旧窗口的比值，2K→128K 时 s=64，代表所有刻度被压缩到原来的 1/64。

$$\displaystyle s = \frac{N_{new}}{N_{old}} = 64$$

2. 统一缩放：所有维度的旋转角都乘以 1/s，等价于把位置 p 压缩成 p/s，再映射回训练过的窗口。

$$\displaystyle \theta_i^{PI} = \theta_i \cdot \frac{1}{s}$$

3. 高频代价：相邻位置的角度差同样被压缩 64 倍，高频维度几乎无法区分“这一步”和“下一步”。

$$\displaystyle (p+1)\theta_i^{PI} - p\,\theta_i^{PI} = \frac{\theta_i}{s}$$

统一减速保住了“能走”，却牺牲了“精准”。老王的直觉是：秒针和时针的使命本来就不一样，凭什么让秒针也跟着变慢？这个直觉催生了按频率区别对待的思路。

NTK 感知缩放

老王围着齿轮箱转了三圈，提出第三个方案：只给慢齿轮减速，快齿轮保持原速！秒针还是秒针，分钟精度一点不丢；时针慢下来，钟面就能多显示几十个小时。这正是 NTK 感知缩放（NTK-aware scaling）的朴素图景：高频维度原样保留，低频维度按需压缩。

NTK-aware 的巧妙之处在于，它不需要逐维手调转速，只要做一件事——把基数 base 从 b 换成更大的 b'。基数变大后，所有维度的转速按指数规律整体下移，但高频端（小 i）受到的影响微乎其微，低频端（大 i）被显著压缩。一石二鸟：短距离精度保住了，长距离覆盖也打开了。

$$\displaystyle b' = b \cdot s^{\frac{d}{d-2}}$$

1. 目标倍率 s：s 仍然是新窗口对旧窗口的倍率，它决定整体要“多走多远”。

$$\displaystyle s = \frac{N_{new}}{N_{old}}$$

2. 换基指数：d/(d−2) 略大于 1，由头维度 d 决定，保证最高频维度几乎不变、最低频维度约缩小 s 倍。

$$\displaystyle \frac{d}{d-2} \approx 1 + \frac{2}{d}$$

3. 新基座：把新基数代回转速公式，各维度自动获得与自身频率匹配的缩放量，无需手调。

$$\displaystyle b' = b \cdot s^{\frac{d}{d-2}}, \qquad \theta_i^{NTK} = (b')^{-\frac{2i}{d}}$$

NTK 方案让钟楼安静了不少，可老王发现两个新问题：齿轮箱转速过渡处有“硬台阶”，咬合时咔咔作响；而且钟摆的摆动幅度明显不对——注意力分布失真了。要解决这两个隐患，需要真正的 YaRN。

YaRN 对数插值

老王拿出“波长卷尺”——波长 λ_i = 2π/θ_i，也就是这对齿轮走一圈所对应的 token 距离。他给齿轮分了三档：波长比旧窗口还短的（高频），保持原速；波长比新窗口还长的（低频），完全压缩；中间的在“对数调校台”上平滑过渡。这就是 YaRN（Yet another RoPE extensioN）的核心思想：按波长分带，频率在对数空间做线性插值。

为什么必须在对数空间？因为频率跨度高达几十倍，线性空间的混合会在中间维度留下“台阶”。而对数空间的线性混合，相当于对频率做等比过渡——过渡带里每一个维度都平滑衔接，没有扭结。γ_i 是 0 到 1 之间随波长爬升的过渡系数，高频端 γ=0（原速），低频端 γ=1（全压缩）。

$$\displaystyle \theta_i^{Y} = \theta_i \cdot s^{-\gamma_i}, \qquad \gamma_i = \mathrm{clamp}\!\left(\frac{\log\lambda_i - \log\lambda_{old}}{\log\lambda_{new} - \log\lambda_{old}},\;0,\;1\right)$$

1. 波长标尺：波长 λ_i 是第 i 对齿轮转一圈走过的 token 数，波长越长的齿轮转得越慢，管得越远。

$$\displaystyle \lambda_i = \frac{2\pi}{\theta_i} = 2\pi\, b^{\frac{2i}{d}}$$

2. 过渡带系数：γ_i 在“旧窗口波长”与“新窗口波长”之间线性爬升，范围被 clamp 在 0 到 1。

$$\displaystyle \gamma_i = \mathrm{clamp}\!\left(\frac{\log\lambda_i - \log\lambda_{old}}{\log\lambda_{new} - \log\lambda_{old}},\;0,\;1\right)$$

3. 对数线性混合：旋转角的对数在“原值”与“压缩值”之间线性内插，等价于让角速度乘以 s 的 −γ 次幂。

$$\displaystyle \log\theta_i^{Y} = (1-\gamma_i)\log\theta_i + \gamma_i\log\frac{\theta_i}{s}$$

齿轮过渡带终于平滑了，但老王的耳朵还是竖着——钟摆的摆动幅度不对劲。压缩频率改变的不只是齿轮转速，还有注意力分数的“方差”，温度缩放正是为这个问题准备的。

注意力温度缩放

插值之后，所有相对旋转角都变小了，cos 值几乎都贴在 1 附近。这意味着每个 key 和 query 的点积差异被压扁——注意力变成“谁都一样重要”，softmax 之后的分布几乎均匀。对模型来说，这和“失忆”没什么区别：它不再能聚焦关键位置。

YaRN 的解法是在钟摆上加一个“助力器”：把注意力 logits 统一乘以温度系数 t。t 由缩放比 s 和头维度 d 决定，s=64、d=128 时 t≈8。乘上 t 之后，被压扁的差异被重新放大，softmax 的锐度恢复到了预训练阶段的水平——注意力分布重新变得有高有低、层次分明。

$$\displaystyle t = s^{\frac{d}{2(d-2)}}, \qquad a_{pq} = \mathrm{softmax}\left( \frac{t \cdot q_p^\top R'(p-q)\, k_q}{\sqrt{d}} \right)$$

1. 方差压缩：插值后相对旋转角整体变小，logits 的方差远小于预训练时的水平，分布趋近均匀。

$$\displaystyle \mathrm{Var}\big(q_p^\top R'(\Delta)\, k\big) \ll \mathrm{Var}\big(q_p^\top R(\Delta)\, k\big)$$

2. 温度系数：t 由 s 和 d 共同决定，量化了“需要放大多少倍”才能找回原来的方差量级。

$$\displaystyle t = s^{\frac{d}{2(d-2)}}$$

3. 补偿后的分布：logits 乘以 t 再做 softmax，注意力分布的锐度与预训练阶段基本一致。

$$\displaystyle a_{pq} = \mathrm{softmax}\Big( \frac{t \cdot q_p^\top R'(p-q)\, k_q}{\sqrt{d}} \Big)$$

钟摆恢复了正常摆动，大钟看起来已经能挂上 128K 的钟面。但老王知道，新齿轮之间还没磨合过——正式营业前，需要一次轻量的试运行。

轻量微调适配

新钟面和新转速装好了，但齿轮咬合的手感需要适应。好在 YaRN 的插值已经把大部分工作做完了，模型不需要从头学习，只需要“重新校准手感”：用极低的学习率微调约 1000 步，让注意力头在新频率下重新找到熟悉的分布。

微调目标仍然是标准的自回归语言建模——逐 token 预测下一个词，最小化负对数似然。配合“截断序列”技巧，可以在有限的显存里模拟超长距离的注意力依赖，让千步微调真正生效。相比朴素插值动辄需要大量数据的重训，YaRN 的微调成本几乎可以忽略。

$$\displaystyle \mathcal{L} = -\frac{1}{T}\sum_{t=1}^{T}\log P_\theta(x_t \mid x_{<t})$$

1. 自回归目标：模型逐 token 预测下一个词，预测越准，单个 token 的对数似然越高。

$$\displaystyle P_\theta(x_t \mid x_{<t})$$

2. 平均损失：在序列长度 T 上取平均，得到整个序列的平均负对数似然，作为微调的优化目标。

$$\displaystyle \mathcal{L} = -\frac{1}{T}\sum_{t=1}^{T}\log P_\theta(x_t \mid x_{<t})$$

3. 千步收敛：在低学习率下，约 1000 步即可完成频率感知的重新校准，无需大规模预训练。

$$\displaystyle \theta \leftarrow \theta - \eta\,\nabla_\theta \mathcal{L}, \qquad \eta = 2\times10^{-5}$$

试运行结束，镇长亲自验收。大家最关心的问题来了：这块 128K 的钟面，到底走得稳不稳、准不准？

128K 应用效果

验收的核心指标是困惑度（PPL）——它衡量模型对文本的“惊讶”程度，数值越低说明模型预测越稳。朴素插值在窗口超过训练长度后 PPL 迅速飙升，而 YaRN 在 128K 甚至更长的窗口下 PPL 保持平稳：这是“走得稳”。

“走得准”则体现在短距离精度上。因为高频维度几乎没动，模型对相邻 token 的敏感度和预训练时一致——短距离困惑度明显低于朴素插值；低频维度的平滑压缩又让远距离信息变得可用。实验表明，YaRN 在检索、长文档问答等下游任务上稳定提升，而微调使用的 token 量只有预训练语料的约千分之一。

$$\displaystyle PPL = \exp\left( -\frac{1}{N}\sum_{n=1}^{N}\log P_\theta(x_n \mid x_{<n}) \right)$$

1. 困惑度定义：对一段文本计算平均负对数似然再取指数，得到“惊讶程度”的度量。

$$\displaystyle PPL = \exp\Big(-\frac{1}{N}\sum_{n=1}^{N}\log P_\theta(x_n \mid x_{<n})\Big)$$

2. 超窗验证：在 128K 以上窗口，YaRN 的 PPL 不随长度飙升，验证频率插值的稳定性。

$$\displaystyle PPL_{128K} \approx PPL_{2K}$$

3. 短距保真：高频维未动使得短距离困惑度显著低于朴素插值，证明“分钟精度”没有丢失。

$$\displaystyle PPL_{yarn}^{short} < PPL_{PI}^{short}$$

大钟在 128K 的钟面上平稳运行，分钟读数分秒不差。从齿轮转速表到对数调校台，从摆锤助力器到千步试运行，YaRN 用一套完整的工程方案回答了“如何让预训练模型看见更长的世界”。

AI 应用场景

Python 演示

import mathimport torchimport torch.nn.functional as F  def yarn_inv_freq(dim, s, base=10000.0):    """Compute YaRN frequencies: log-space linear interpolation."""    half = dim // 2    idx = torch.arange(half, dtype=torch.float32)    inv_orig = 1.0 / (base ** (2 * idx / dim))    new_base = base * (s ** (dim / (dim - 2)))    inv_new = 1.0 / (new_base ** (2 * idx / dim))    low = 0.25 * half    high = 0.75 * half    gamma = torch.clamp((idx - low) / (high - low), 0.0, 1.0)    inv_freq = torch.exp((1 - gamma) * torch.log(inv_orig) + gamma * torch.log(inv_new))    return inv_freq, gamma  def rope_rotate(x, pos, inv_freq):    """Apply rotary embedding to x: [n, dim]."""    half = inv_freq.shape[0]    angles = torch.outer(pos, inv_freq)    cos_c = torch.cos(angles)    sin_c = torch.sin(angles)    x1 = x[:, :half]    x2 = x[:, half:]    out1 = x1 * cos_c - x2 * sin_c    out2 = x1 * sin_c + x2 * cos_c    return torch.cat([out1, out2], dim=-1)  def main():    torch.manual_seed(0)    dim, s, n = 64, 16, 512    inv_freq, gamma = yarn_inv_freq(dim, s)    pos = torch.arange(n, dtype=torch.float32)    q = torch.randn(n, dim)    k = torch.randn(n, dim)    q_r = rope_rotate(q, pos, inv_freq)    k_r = rope_rotate(k, pos, inv_freq)    scores = (q_r[0] @ k_r.T) / math.sqrt(dim)    attn = F.softmax(scores, dim=-1)    ent_plain = -(attn * torch.log(attn + 1e-9)).sum().item()    t = s ** (dim / (2 * (dim - 2)))    attn_t = F.softmax(scores * t, dim=-1)    ent_temp = -(attn_t * torch.log(attn_t + 1e-9)).sum().item()    print("gamma range: %.3f - %.3f" % (gamma.min().item(), gamma.max().item()))    print("attention entropy without temp: %.3f" % ent_plain)    print("temperature t: %.3f" % t)    print("attention entropy with temp: %.3f" % ent_temp)  if __name__ == "__main__":    main()

1. new_base = base * (s ** (dim / (dim - 2)))NTK 感知缩放的“换基座”，基数变大后各维度自动获得不同缩放。
2. gamma = torch.clamp((idx - low) / (high - low), 0.0, 1.0)过渡带斜坡：高低频之间平滑爬升的 γ 系数。
3. inv_freq = torch.exp((1 - gamma) * torch.log(inv_orig) + gamma * torch.log(inv_new))对数空间线性插值，YaRN 的全部精髓。
4. angles = torch.outer(pos, inv_freq)位置向量与频率向量做外积，得到每个 token 每个维度的累计旋转角。
5. scores = (q_r[0] @ k_r.T) / math.sqrt(dim)位置 0 的 query 与所有 key 点积，得到原始注意力 logits。
6. t = s ** (dim / (2 * (dim - 2)))注意力温度系数，s=16、d=64 时约为 3.36。
7. attn_t = F.softmax(scores * t, dim=-1)温度补偿后的注意力分布，熵显著降低，锐度恢复。

常见误区

面试题