引言

想象一下，你试图总结一本内容密集的小说，但你的记忆一次只能容纳十页。当你读到第三章时，第一章的内容就已经完全忘记了。这就是大型语言模型 (LLM) 在处理有限上下文窗口时面临的根本困境。虽然像 GPT-4 和 LLaMA-2 这样的模型彻底改变了自然语言处理 (NLP) ，但它们处理大量输入 (如整本书或法律资料库) 的能力受到其“上下文窗口”的限制。

解决这个问题的标准方法就像是在调节收音机。研究人员“拉伸”模型的内部定位系统 (具体来说是旋转位置编码，即 RoPE) 以适应更长的序列。像位置插值 (PI) 、YaRN 和 LongRoPE 等技术在这方面取得了很大进展。然而，它们都面临一个共同的问题: 找到确切的频率来调节模型是非常困难的。这需要搜索巨大的参数空间。如果频率稍微偏离，模型对位置的理解就会发生“漂移”，导致困惑度 (perplexity) 飙升和幻觉产生。

在这篇文章中，我们将深入探讨一篇题为 “PSC: Extending Context Window of Large Language Models via Phase Shift Calibration” (PSC: 通过相位偏移校准扩展大型语言模型的上下文窗口) 的研究论文。作者引入了一个名为 相位偏移校准 (Phase Shift Calibration，简称 PSC) 的巧妙且轻量级的模块。PSC 不再花费数天时间去搜索完美的频率，而是接受一个不完美的猜测，并在运行过程中通过数学方法对其进行“校准”。

我们将拆解 RoPE 的数学原理，直观地解释为何现有方法会发生漂移，并说明 PSC 如何像矫正镜片一样，让模型能够以高精度阅读高达 64k 的 token (甚至可能更多) 。

背景: 位置的机制

要理解为什么 LLM 会被长文档搞混，我们首先需要了解它们是如何知道句子中单词的位置的。与循环神经网络 (RNN) 不同，Transformer 并行处理所有单词。它们需要为每个单词打上“时间戳”。

旋转位置编码 (RoPE)

这种“时间戳”的行业标准是旋转位置编码 (RoPE) 。RoPE 不是将静态数字添加到词向量中，而是在高维空间中旋转向量。旋转的角度对应于单词的位置。

让我们看看研究人员提供的数学公式。对于位置 \(m\) 和嵌入向量 \(\mathbf{x}_m\)，RoPE 使用复数对查询 (Query，\(\mathbf{q}\)) 和键 (Key，\(\mathbf{k}\)) 向量进行变换。

公式 1: 查询向量的 RoPE 变换。

公式 2: 键向量的 RoPE 变换。

这里，\(i\) 是虚数单位，\(\theta\) 是基频。关键在于，两个单词之间的相互作用 (注意力分数) 仅取决于它们之间的相对距离。当模型计算注意力时，它会计算查询和键的点积。

公式 3: Softmax 注意力分数。

由于复数的旋转特性，该点积产生一个取决于 \((m - n)\) 的函数，即两个单词之间的距离。

公式 4: 显示相对位置依赖性的内积。

对于模型训练时的上下文长度 (例如 4096 个 token) ，这种机制既优雅又有效。

问题: 外推与“拉伸”

当我们想要扩展上下文窗口 (例如从 4k 到 32k) 时，我们不能简单地进一步旋转向量——因为模型在训练期间没有见过那些旋转角度。流行的解决方案, 位置插值 (Position Interpolation, PI) , 涉及将更长的序列“压缩”到原始旋转范围内。如果你想处理 4 倍的长度，你就让所有东西旋转得慢 4 倍。

其他方法如 YaRN 和 LongRoPE 使用更复杂的缩放因子来保留高频信息。然而，它们都依赖于预定义或搜索缩放因子。

问题就在这里: 这些因子的搜索空间是指数级的。要在计算上找到完美的缩放因子是不可能的。因此，研究人员只能满足于“足够好”的因子。但“足够好”会导致 相位偏移 (Phase Shift) ——即编码的位置稍微偏离最佳目标的偏差。

图 1: 相位偏移导致 sin/cos 值偏离其最佳位置。

如上方的 图 1 所示，蓝线代表最佳频率 (\(\theta^*\)) ，绿线代表我们通过估算方法得到的实际频率。那个水平间距就是相位偏移。在数千个 token 中，这个微小的误差会累积，使模型对信息的位置感到困惑。

核心方法: 相位偏移校准 (PSC)

这篇论文的作者提出了一个新颖的解决方案。与其试图找到完美的频率 \(\theta^*\) (这太昂贵了) ，不如接受不完美的频率 \(\hat{\theta}\)，并添加一个小的、可学习的模块来修正误差。

校准的数学原理

研究人员将理想的查询编码 \(f_q^*\) 建模为最佳频率与估计频率之间的关系。

如果我们假设错过了最佳频率 \(\theta^*\)，而我们拥有估计频率 \(\hat{\theta}\)，那么理想编码与实际编码的关系就是它们之间差值的旋转: \((\theta^* - \hat{\theta})\)。

公式 13: 推导理想编码与实际编码之间的关系。

这个等式揭示了，理想的嵌入只是实际嵌入旋转了误差项 \(e^{im(\theta^* - \hat{\theta})}\)。

以矩阵形式表示，这种修正看起来像一个块对角旋转矩阵。这个矩阵 \(\tilde{\mathbf{R}}\) 代表了将嵌入“校准”回焦点所需的调整。

公式 15: 修正的矩阵形式。

为什么 LoRA 还不够

适应 LLM 的一种常用技术是 LoRA (低秩自适应) , 它冻结主模型并训练小的低秩矩阵来调整权重。你可能会问: 难道我们不能直接用 LoRA 来学习这种修正吗？

作者证明了 LoRA 并不适合这项特定任务。

修正矩阵 \(\tilde{\mathbf{R}}\) (理想旋转与实际旋转之间的差异) 是一个 满秩 (full-rank) 的块对角矩阵。如果估计的频率在所有维度上哪怕只是稍微偏离，所需的修正矩阵的秩就会变得非常高 (可能等于注意力头的数量) 。LoRA 依赖于更新是低秩的假设。因此，LoRA 难以有效地逼近这种高秩的旋转变换。

PSC 模块架构

为了解决这个问题，作者引入了 相位偏移校准 (PSC) 模块。它是专门设计用来学习这种块对角旋转的。

PSC 模块是一个注入到注意力机制中的轻量级多层感知机 (MLP) 。它作用于嵌入 \(\mathbf{x}\)。

公式 19: PSC 模块的数学定义。

这里，\(\mathbf{W}_1\) 和 \(\mathbf{W}_2\) 是可学习的块对角矩阵。这种结构与 RoPE 的块状性质完美契合。

架构在下方的 图 2 中可视化。注意 PSC 模块如何与主路径并行，计算出一个偏移量，该偏移量在应用标准 RoPE 之前 (或之后) 被添加到原始嵌入中。

图 2: PSC 架构显示嵌入被校准到理想位置。

实现

PSC 的优势之一是它的简单性。它只需向模型添加不到 1% 的额外参数。该算法在复数空间 (或块对角等效空间) 中计算校准并调整查询和键向量。

算法 1: Pytorch 风格的相位偏移校准伪代码。

作者发现 预校准 (Pre-calibration) (在位置编码之前应用 PSC) 效果最好。这可能是因为位置编码会引入复杂的非线性失真；在信号被失真之前修正它比事后修复更容易。

实验与结果

研究人员在 LLaMA-2 和 Mistral 模型上验证了 PSC，测试了从 16k 到 64k token 的上下文窗口。他们将 PSC 与标准位置插值 (PI) 、YaRN 和 LongRoPE 进行了比较。

1. 困惑度 (Perplexity) 的改善

困惑度衡量模型对新文本的惊讶程度；数值越低越好。

在下表( 表 1 )中，查看 16384 (16k) token 这一列。

  • PIFT (位置插值微调) : 7.32
  • LongRoPEFT: 7.26
  • LongRoPE + PSC (LongRoPEPSCFT): 7.24

表 1: LLaMA-2 7B 上的滑动窗口困惑度比较。

虽然数字上的差异看起来很小，但在 LLM 的世界里，这些边际收益通常转化为显著更好的稳定性。更重要的是，请注意 PSC 改善了它接触的每一种方法 。 无论你使用 PI、YaRN 还是 LongRoPE，添加 PSC 都会降低困惑度。它充当了一个通用的增强器。

2. 密钥检索测试 (Passkey Retrieval Test)

这是“大海捞针”测试。模型被赋予一段充满垃圾数据的长文本，其中隐藏着一个随机的 5 位数字密钥。然后问它: “密钥是什么？”

这个测试揭示了模型是真正在使用其上下文窗口，还是仅仅因为长度而感到困惑。

图 3: 密钥检索准确率。注意粉红色的线。

图 3 是论文中最令人信服的结果。

  • 蓝线 (基座 LLaMA-2) : 超过 4k token 后立即失败。
  • 品红线 (YaRN FT) : 表现良好，但在 32k token 左右跌至 0% 的准确率。
  • 粉红线 (YaRN + PSC) : 一直到 34k token 都保持 100% 的准确率 。

通常导致模型在窗口极限附近性能下降的“相位偏移”，被 PSC 有效地中和了，从而在更长的范围内保持完美的检索准确率。

3. 通用能力 (标准基准测试)

扩展上下文的一个常见担忧是“灾难性遗忘”——即模型在长文本上表现变好，但在短的、标准的任务 (如数学或逻辑) 上变得愚蠢。

作者在 Hugging Face Open LLM 排行榜 (ARC, HellaSwag, MMLU, TruthfulQA) 上测试了 PSC。

表 5: 标准基准测试上的表现。

如 表 5 所示，配备 PSC 的模型 (带有下标 \(_{PSC}\) 的行) 表现与其非 PSC 对应模型相当，有时甚至更好。例如, YaRN\(_{PSC}\) 在 TruthfulQA 上取得了最高分 (39.81) ，击败了标准的 YaRN 模型。这证实了校准位置频率不会损害模型的通用推理能力。

4. 效率

这个新模块会让模型变得臃肿吗？

图 5: GPU 显存消耗比较。

图 5 比较了标准 LoRA (橙色) 与 PSC + LoRA (绿色) 的 GPU 显存使用情况。这两条线几乎重合，直到上下文变得巨大 (64k) 时，PSC 使用的显存才略多一点。

  • 参数数量: 对于 LLaMA-2 7B，PSC 仅增加了 64M 参数 (占总数的 0.095%) 。
  • 计算开销: 推理时间的增加可以忽略不计 (在 16k token 输入上大约只有 5ms 的差异) 。

表 9: 计算开销。

结论与启示

论文 “PSC: Extending Context Window of Large Language Models via Phase Shift Calibration” 指出了我们目前扩展 LLM 方式中一个微妙但关键的缺陷: 频率估计的不准确性。

通过将这种不准确性视为“相位偏移”并构建特定的神经组件来修正它，作者提供了一个稳健的解决方案，该方案:

  1. 增强了现有方法: 它在 PI、YaRN 和 LongRoPE 之上工作。
  2. 解决了秩不匹配问题: 它处理了 LoRA 无法有效学习的满秩旋转修正。
  3. 参数高效: 它只需极少的计算和显存开销。

对于学生和从业者来说，这篇论文强调了一个重要的教训: 有时候，与其寻找一组完美的静态超参数 (如频率因子) ，不如构建一个能够学会自我校正的系统更为有效。随着我们向无限上下文模型迈进，像 PSC 这样的动态校准机制很可能会成为 Transformer 架构中的标准组件。