🤖 machine learning

Jet-Long: Efficient Long-Context Extension with Dynamic Bifocal RoPE

Jet-Long 是一种无需微调、零样本的方法，它通过动态调整 RoPE 重缩放因子来平衡短上下文保真度与长上下文外推能力，并通过一种高效的双焦注意力机制，以极低的推理开销在长上下文基准测试中实现了最先进的性能。

原作者： Haozhan Tang, Zerui Wang, Yuxian Gu, Song Han, Han Cai

发布于 2026-07-10

📖 1 分钟阅读☕ 轻松阅读

原作者： Haozhan Tang, Zerui Wang, Yuxian Gu, Song Han, Han Cai

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你拥有一个超级聪明的机器人大脑（大型语言模型），它被训练用来阅读长度在 32,000 个单词以内的故事。现在，你想让它阅读一整本小说或一个巨大的代码库，其长度是原来的四倍。如果你直接把这么长的文本交给它，机器人就会感到困惑。这就像是在读一本书时，页码突然从 100 跳到了 1,000,000；机器人的内部指南针（称为 RoPE）会疯狂旋转，导致它开始产生幻觉或忘记故事的中段。

有一段时间，科学家们试图通过选择一个“魔术数字”来拉伸机器人的视野，以此来修复这个问题。但这是一个双输的游戏：如果拉伸得太厉害，为了看清书的结尾，机器人就会忘记开头；如果保持紧凑以记住开头，它就无法看到结尾。

于是，Jet-Long 出现了，它像是一个为机器人的眼睛提供动态双焦点镜片的新方法。

双焦点眼镜的妙招

Jet-Long 不再强迫机器人在整本书中使用单一的拉伸视角，而是同时给它两副眼镜：

近视镜片（局部窗口）： 对于第一部分文本（“局部窗口”），机器人看到的每一处细节都与它受训时看到的一模一样。没有拉伸，没有畸变。这能让机器人的近期记忆保持清晰。
远视镜片（远程窗口）： 对于剩余的文本（“远程窗口”），机器人使用了一个特殊的技巧。它不仅仅是拉伸视野，还会动态地将页面进行分组。随着书本变长，机器人会自动调整将多少页合并为一个“超级页”。

神奇之处在于，这个分组因子并不是固定的。它会根据当前文本的长度实时变化。如果文本很短，机器人会逐字查看；如果文本非常庞大，它就会将单词进行分组，程度恰到好处，使得机器人的内部指南针保持冷静，不会失控旋转。这意味着机器人对于短任务能保持完美的准确性，同时仍能阅读海量文档而不会迷失方向。

“免费午餐”般的速度

通常情况下，同时进行两种不同的计算（既看近又看远）会让机器人慢上一倍。但论文作者发现了一种巧妙的方法来合并这些计算。他们构建了一个特殊的引擎（“融合算子”），可以在一次处理中完成两种镜片的数学运算。

这就像一位厨师通常需要先切菜、再煮汤、最后煎炸，分三个锅进行；而 Jet-Long 就像一个神奇的锅，可以同时完成这三步且不损失任何风味。

结果： 当机器人阅读高达 128,000 个 token 的超长上下文时，在 H100 芯片上，Jet-Long 在起始阶段（“预填充”阶段）比标准方法（FlashAttention-2）快 1.39 倍。
代价： 当机器人逐个单词生成新文本时，它只比标准方法慢了约 4%。对于能够阅读整个图书馆而非仅仅是一本小册子的能力来说，这只是一个微小的代价。

它们测试了什么（以及没测试什么）

团队在三种不同规模的机器人大脑（17 亿、40 亿和 80 亿参数）上进行了测试，发现 Jet-Long 始终优于之前的最佳方法。

在一项名为 RULER 的测试中（该测试检查机器人能否在文本的“大海捞针”中找到隐藏的针），Jet-Long 在最小的模型上比最强的竞争对手高出 4.79 分，在最大的模型上高出 2.03 分。
在一项名为 HELMET-RAG 的测试中（模拟现实世界的搜索任务），它实现了最高的整体准确率。
在 PG-19 测试（阅读古籍）中，它将机器人的困惑度（Perplexity）保持在最低水平，而其他方法会让机器人的困惑度随着文本变长而飙升。

重要提示（关于限制）： 论文明确排除了“需要从头开始重新训练机器人才能使其奏效”的观点。Jet-Long 适用于“零样本”（zero-shot）模型，这意味着你可以拿一个预训练好的机器人，直接插上这个“镜片”。然而，论文也指出，虽然 Jet-Long 解决了“指南针旋转”的问题，但它并没有解决另一个被称为“注意力弥散”（attention diffusion，即机器人被过多的选项淹没）的问题。那个问题通常需要改变机器人的架构本身，而不仅仅是添加一个镜片。

结论

Jet-Long 并不是让机器人变得无限聪明的魔杖，但它是一个高效、无需调优的工具，能让现有的机器人处理长达 128,000 个 token 的文本而不至于崩溃。它在保持机器人短期记忆完美的同时，动态调整其长程视觉，且运行速度几乎与标准方法持平。作者在真实硬件（H100 GPU）上进行了测量，并发现它在不同规模的模型甚至混合了不同类型注意力的混合型机器人大脑上都能奏效。对于任何想要让 AI 阅读更长内容的人来说，这是一个扎实且经过验证的升级方案。

技术摘要：Jet-Long

问题陈述

现代大语言模型（LLM）越来越多地被部署在长上下文应用中，例如检索增强生成（RAG）、仓库级代码编写和智能体工作流。在这些场景下，累积的推理轨迹和工具输出往往会将输入长度推高到超出预训练窗口一个数量级以上（例如，100K+ token 对比 32K 的训练窗口）。

由于二次方计算成本高昂且长上下文数据稀缺，直接在长上下文上训练模型在成本上是难以承受的，而微调往往会降低短上下文性能。因此，零样本上下文扩展（无需微调即可扩展上下文）已成为开源权重检查点的主流部署路径。

然而，使用旋转位置嵌入（RoPE）的原生 Transformer 类 LLM 在超越其训练窗口时会失效，主要归因于两种失败模式：

位置分布外（OOD）： 在训练期间未见过的序列位置，低频 RoPE 组件产生的旋转角度超出了训练分布，导致注意力分数异常。
注意力弥散（Attention Diffusion）： 随着 Key 集合的增长，Softmax 分布变得平坦，将概率质量分散到无关的 token 上，通常伴随着“U型”位置偏差，即对中间上下文信息关注不足。

现有的零样本方法（如 NTK、YaRN、Self-Extend、DCA）通常预先固定一个单一的缩放因子或分组大小。这迫使开发者进行权衡：激进的因子会牺牲短上下文的保真度，而保守的因子则无法在长上下文处实现清晰的外推。

方法论：Jet-Long

Jet-Long 提出了一种无需微调的零样本上下文扩展方法，该方法将一个符合局部 RoPE 特性的窗口与一个缩放因子随当前序列长度动态变化的远程窗口相结合。

1. 动态双焦 RoPE (Dynamic Bifocal RoPE)

其核心创新在于一种动态双焦映射，避免了固定因子的权衡问题：

局部窗口（Local Window）： 大小为 $w_0$ 的窗口保留经典的 RoPE，精确地保留基座模型的预训练行为，适用于短输入。
远程窗口（Remote Window）： 对于超过 $w_0$ 的位置，应用动态位置映射 $f(\cdot)$ 。不同于连续插值频率的方法，Jet-Long 使用离散位置混叠（Discrete Position Aliasing）。
动态分组： 分组大小 $G$ 根据当前序列长度 $L$ 和预训练窗口 $w_{pretrained}$ 动态计算：
$G = \max\left(1, \left\lceil \frac{L}{w_{pretrained}} \right\rceil\right)$
远程映射采用向下取整函数： $f(x) = \lfloor x/G \rfloor$ 。这确保了每个远程旋转角度都保持在分布内（即处于模型的预训练网格上），同时通过应用最小必要的压缩来最大化位置分辨率。

2. 推理效率机制

为了确保该方法具有计算效率，且不需要重新训练或大规模修改 KV Cache，Jet-Long 采用了两种关键技术：

即时修正旋转（On-the-Fly Correction Rotation）： KV Cache 存储其精确绝对位置的未压缩基础 Key。在生成过程中，Jet-Long 并不重写缓存，而是对当前的 Query 和缓存的 Key 即时应用一个修正旋转 $\Delta$ 。由于 RoPE 旋转是加性组合的，这可以在寄存器中重建外推向量而不触碰物理缓存，从而实现跨长度边界的流式生成而不会发生停顿。
包含-排除注意力合并（Inclusion–Exclusion Attention Merge）： 为了避免在 Prefill 阶段实例化完整的二次方注意力矩阵，Jet-Long 使用包含-排除原理结合了三次标准的 FlashAttention Pass：
1. 局部窗口注意力（基础 Query/Key）。
2. 全量因果注意力（远程 Query/Key）。
3. 局部窗口注意力（远程 Query/Key）。
  通过使用 LogSumExp 统计量来稳定合并过程，该方法可以在不使用布尔掩码矩阵的情况下计算出精确的基于距离的路由输出，从而保留了 FlashAttention 的内存效率。

这些组件被融合进单个 CuTe kernel 中，最大限度地减少了 kernel 启动开销。

核心贡献

Jet-Long 算法： 一种无需微调的双焦方法，其远程窗口缩放因子随序列长度动态调整，使远程旋转保持在分布内，同时精确复现基座模型在原生上下文内的表现。
高效实现： 通过包含-排除注意力合并和即时修正旋转，实现了对 KV Cache 的无损操作。通过融合为单个 CuTe kernel，该方法在 H100 长上下文下的 Prefill 吞吐量达到了 1.28–1.39× FlashAttention-2 (FA2)，并在单 Batch 生成期间产生 $\le$ 4% 的开销。
实证验证： 在 Qwen3-1.7B/4B/8B 模型上高达 128K 上下文的广泛评估中，证明了其在 RULER、HELMET-RAG 和 PG-19 基准测试上优于强力的零样本基线（DNTK、YaRN、DCA、Self-Extend）。
泛化能力： 该方法可无缝迁移至混合注意力架构（如 Jet-Nemotron），并且对单一超参数 $w_0$ 的选择保持鲁棒。

实验结果

在 Qwen3 基座模型（1.7B, 4B, 8B）上，将 32K 原生窗口扩展至 128K 进行评估：

RULER 基准测试： Jet-Long 在 1.7B、4B 和 8B 模型上分别领先最强基线 +4.79、+2.18 和 +2.03 个百分点 (pp)。它在 32K 训练窗口内保留了基座模型性能，并在扩展长度处超越了所有基线。
HELMET-RAG： 在该基准测试中取得了最佳的整体准确率，而 HELMET 研究已确定该测试是预测下游长上下文性能最有效的指标。
PG-19 Perplexity（困惑度）： 在所有外推方法中获得了最低的困惑度。当基座模型崩溃时（例如 1.7B 模型在 128K 时困惑度 >70），Jet-Long 仍能保持低困惑度（例如 1.7B 为 11.41）。
混合架构： 应用于 Jet-Nemotron（交替使用 Softmax 注意力和线性注意力），Jet-Long 在 128K 长度下将 RULER 准确率提升了 +10.01 pp (2B) 和 +11.31 pp (4B)，而此时基座混合模型已趋于崩溃。
消融研究：
- 超参数韧性： 局部窗口大小 $w_0$ 是鲁棒的；在 512 到 4096 之间的取值与最优性能的差距保持在 1–2 pp 以内。
- 混叠 vs. 插值： 位置混叠（Jet-Long 的方法）显著优于频率插值（Yaarn 式方法），例如在 64K 处提升了 +6.99 pp，证实了将旋转保持在离散预训练网格上的优势。

意义

论文声称 Jet-Long 在 RoPE 层级解决了位置 OOD 失败模式，且无需长上下文训练的计算成本，也无需固定因子外推带来的保真度损失。通过使缩放因子动态化并利用融合 kernel 实现“本质免费”的推理构建，Jet-Long 为开源权重 LLM 在长上下文场景下的部署提供了一条实用且高性能的路径。它作为一种针对注意力弥散问题的架构替代方案（如稀疏注意力或状态空间模型）的补充，提供了一种无需重新训练且能泛化至各种模型规模及混合架构的零样本解决方案。