多token残差预测
Multi-token Residual Prediction
Modal Research与上海纽约大学HeavyBall Research合作提出多Token残差预测（MRP），将多Token预测（MTP）思想引入扩散语言模型（DLM）。MRP是一个小型3层transformer模块，通过预测相邻去噪步骤间的logit残差而非完整分布，使学习更易收敛。在静态模式下，MRP在SGLang中实现最高1.56倍无损加速（推测模式）或1.89倍有损加速（直接模式）；在动态模式下，通过重新掩码恢复因低阈值解码（τ=0.5）损失的准确率，在SDAR-1.7B/4B/8B上平均最高提升+16个百分点，在GSM8K、MATH500、HumanEval和MBPP上验证。
研究
2026年7月1日·7分钟阅读
编者按：这篇客座博文描述了 Modal Research 与上海纽约大学 HeavyBall Research 之间一项研究合作的结果。
MRP 是一个具有两种用途的小型模块。左图：在静态模式下，它可以无损地加速解码（推测模式，lookahead steps K = 3），或者以较小的质量代价进一步加速（直接模式，lookahead steps K = 1），在 SGLang 中吞吐量最高可达 1.56 倍。右图：在动态模式下，它可以恢复因激进的低阈值解码（阈值 τ = 0.5）而损失的准确率，最高可达 +16 个百分点。结果基于 SDAR-1.7B/4B/8B 在 GSM8K、MATH500、HumanEval 和 MBPP 上的平均值。TL;DR： 多 Token 预测（MTP）通过单次前向传播预测多个 token 来加速自回归模型 (Gloeckle et al., 2024)。我们将这一思想引入扩散语言模型，并做了一项关键改动使其生效：我们不训练一个小型 head 来预测下一个去噪步骤的完整分布，而是训练它预测相邻步骤之间的_残差_。残差是一个更容易学习的目标，因此一个小型模块就能准确预测它，并将其应用于多个步骤。同一个模块随后服务于 DLM 推理中通常需要相互权衡的两种模式：在静态模式下，当需要保持输出质量时，它提供（近乎）无损的加速（在 SGLang 中平均可达 1.56 倍）；在动态模式下，它能够恢复因激进的吞吐量设置而损失的大部分质量（平均最高 +16 个百分点）。
论文：https://arxiv.org/abs/2605.18817
代码：https://github.com/heavyball-research/multi-token-residual-prediction
SGLang 实现：https://github.com/heavyball-research/sglang
模型：https://huggingface.co/collections/heavyball/sdar-mrp
从 MTP 开始
如果你接触过快 LLM 推理，你应该了解多 token 预测（MTP）的故事。自回归模型每次前向传播生成一个 token，这很昂贵。因此，人们会附加轻量级的 head，例如 Medusa、EAGLE、DeepSeek 的 MTP，它们窥视主干网络的隐藏状态，并在一次前向传播中猜测接下来的几个 token。结合推测性验证，我们就能获得真正的加速。
这是一个美妙的想法，并且行之有效。所以我们自然要问：我们能否让它超越自回归模型发挥作用？
扩散语言模型（DLM）是一个自然的选择，因为它们不是从左到右解码的。DLM 从一个完全掩码的序列开始，然后对其进行去噪，一次解除少量高置信度位置的掩码。并行性内置于该过程中，但存在一个权衡：如果在单步中解除过多 token 的掩码，质量就会下降，因为每个 token 的解码没有考虑到同时被提交的其他 token。大多数 DLM 加速文献都处在这条帕累托曲线上，用质量换取速度。
这正是我们想要打破的权衡，而 MTP 看起来是合适的工具。如果一个轻量级 head 能够从主干网络的隐藏状态预测额外的 token，那么它可以在每次前向传播中解码多个位置，同时让每个位置都感知到其他位置，而不是简单地解除更多掩码并为此付出质量代价。问题在于 MTP 的方案能否迁移到扩散设置中。我们发现它不能直接迁移，而理解其原因正是我们方法的起点。
一次天真的尝试
我们首先取一个小型 head，并将其蒸馏，使其直接从当前步骤的隐藏状态预测下一步的完整 log-density。运行几次，以便在每次主干网络传播中解除多个 token 的掩码。
问题在我们要求 head 执行超过一步时立刻显现。蒸馏整个分布意味着每一步都必须从头开始重现一个大的、高动态范围的目标，并且每步误差会累积。我们发现，以这种方式训练的 head 在一步时表现尚可，但到三或四步时，在我们尝试的所有模型规模上都崩溃了。在 SDAR-4B 主干网络上，我们的直接蒸馏 head 在 GSM8K 上的准确率从一步时的 84.8% 下降到四步时的个位数。它根本无法在多次迭代中保持分布。
| 方法 | K=1 | K=2 | K=3 | K=4 |
|---|---|---|---|---|
| 朴素 MTP | 84.8 | 16.9 | 5.9 | 1.9 |
SDAR-4B 上的 GSM8K（0-shot, CoT）准确率。我们在此直接应用 MTP 方案，准确率在第一个 MTP 步骤之后急剧下降。这里 K 表示预测步数。
核心洞察
上述失败也是一个线索。如果完整的下一步分布难以跨多步预测，那么问题在于是否存在一个更容易预测的目标。答案是肯定的，这来自于观察主干网络输出在两个相邻去噪步骤之间实际变化有多小。
这就是全部洞察，它改变了目标。我们不需要一个小型模块从头开始重现下一步的完整分布。我们需要它对一个已经相当不错的预测进行小的修正。
这不仅仅是一个幸运的经验事实。它源于去噪的马尔可夫结构：每一步只扰动少数几个位置，因此根据 Lipschitz 论证，未扰动位置的预测分布只能移动有限的距离，并且随着去噪的进行和模型变得更加自信，这个界限会收紧。我们要求一个小型模块学习的信号本质上是低复杂度的。这正是小型模块能够学习它的原因。
多 Token 残差预测（MRP）
多 Token 残差预测（MRP）是一个附加在冻结的 DLM 主干网络上的小型 transformer（在我们的主要配置中为 3 层）。它读取主干网络的隐藏状态，预测步骤间的 logit 残差，并将其加到主干网络自身的 logits 上。主干网络、其 LM head 和 token embeddings 保持冻结；只有 MRP 模块被训练。
训练目标是 MTP 蒸馏损失的残差版本。我们运行冻结的主干网络两次，一次在揭示一组 token 之前，一次在之后，并在仍然被掩码的位置上使用 KL 散度训练 MRP，以最小化两次输出之间的差异。由于 softmax 下的归一化常数会抵消，这等价于匹配真实的条件分布，而模块本身只表示修正。
| 方法 | K=1 | K=2 | K=3 | K=4 |
|---|---|---|---|---|
| 朴素 MTP | 84.8 | 16.9 | 5.9 | 1.9 |
| MRP | 88.6 | 84.9 | 70.9 | 57.2 |
SDAR-4B 上的 GSM8K（0-shot, CoT）准确率。使用辅助模块预测残差而非整个分布使得学习变得容易得多。这里 K 表示预测步数。残差框架的优势随着 K 的增长而显现。直接建模分布（朴素 MTP）与残差学习（MRP）之间的差距在 K = 1 时只有几个百分点，但在此之后急剧扩大：在 K = 2 时，残差学习已经在 GSM8K 上领先 +65 个百分点，而直接变体在 K = 4 时完全崩溃。
推理中的应用
在冻结的主干网络上使用一个训练好的模块，我们可以廉价地近似下一步去噪步骤将产生的结果。如何最好地利用这个近似取决于主干网络是如何解码的。
DLM 通常在两种模式之一中运行。在静态去噪中，每一步解除少量、固定数量的位置的掩码；这保持了高质量但吞吐量低，是当你需要输出正确时想要的模式。在动态去噪中，每个置信度超过阈值的位置会立即被解除掩码；这提高了吞吐量，但在低阈值下，主干网络每步提交许多 token，质量下降。
MRP 适用于两种模式，但在每种模式中扮演不同的角色：在静态模式下，它_增加_揭示以加快速度；在动态模式下，它_撤销_过于急切的揭示以恢复质量。
应用 I：静态去噪中的无损加速
在静态模式下，MRP 将推理变成了一个可调节的旋钮。同一个训练好的模块为你提供了一系列操作点，从_与主干网络输出完全相同_到_显著更快且质量代价小且可衡量_，你可以根据应用的实际需求选择位置。
一端是推测解码，用于输出必须与主干网络产生的输出完全匹配的情况。这里 MRP 充当起草者：它廉价地提出下一批 token，主干网络在单次传播中验证它们。主干网络同意草案的位置被接受；不同意草案的位置被重新掩码并重做。由于扩散模型在一次前向传播中为每个位置打分，这种逐位置验证是自然的。验证传播并非浪费：其隐藏状态和 logits 为下一次迭代提供种子，因此当接受率高时，它兼作下一步的主干网络传播，其成本被摊销。
| 主干网络 | GSM8K | MATH500 | HumanEvall | MBPP |
|---|---|---|---|---|
| SDAR-4B | 90.0 / 1.36x | 68.0 / 1.26x | 67.7 / 1.35x | 66.5 / 1.27x |
| SDAR-8B | 90.4 / 1.40x | 74.8 / 1.39x | 72.6 / 1.34x | 67.3 / 1.34x |
SGLang 中的推测模式。准确率（%）后跟相对于仅主干网络基线的吞吐量加速。质量通过构造与主干网络匹配。在单个 H100 上测量；我们在 SGLang 中提供了一个实现。
另一端是直接解码，它跳过验证并直接提交 MRP 修正后的 logits。验证需要一次完整的主干网络前向传播，因此放弃它提高了加速的上限，并且残差预测本身足够准确，以至于在推理任务上，质量代价很小：
| 设置 | GSM8K | MATH500 | HumanEval | MBPP |
|---|---|---|---|---|
| 基线 | 90.9 / 1x | 72.2 / 1x | 73.8 / 1x | 67.7 / 1x |
| 直接（MRP Step 1） | 90.1 / 1.59x | 71.4 / 1.61x | 67.1 / 1.53x | 63.8 / 1.51x |
| 直接（MRP Step 2） | 89.2 / 1.89x | 70.8 / 1.91x | 64.0 / 1.78x | 59.9 / 1.75x |
SDAR-8B 上的直接解码。准确率（%）后跟相对于主干网络的吞吐量加速。Token 在无验证的情况下提交。K 设置每次主干网络前向传播运行多少 MRP 步骤，因此它本身就是一个旋钮：K = 1 在推理任务上保持在主干网络一个百分点以内，速度为 1.6 倍；K = 2 将速度推至 1.8 倍以上，但在代码任务上下降更大。较小规模遵循相同模式（完整表格见论文）。
关键在于你可以选择操作点。当正确性不可妥协时无损，当延迟占主导且可接受小的质量代价时更快，可根据任务甚至每个请求进行调整。这种控制只有当你拥有应用程序的推理栈时才存在。在封闭的 API 背后，这种权衡是为你做出的：提供商固定了解码策略，你接受他们提供的任何成本-质量点。使用 MRP 运行你自己的主干网络将旋钮交还到你手中。
应用 II：质量恢复
现在考虑相反的模式：一个交互式设置，其中延迟最重要，因此解除掩码的阈值设置得很低，每步揭示许多 token。在激进的阈值下，这会带来损害，因为主干网络在一步中提交一批 token，每个 token 在选择时无法考虑到其他 token 同时被提交。
这里 MRP 朝另一个方向运行。在主干网络以低阈值过度揭示后，一次 MRP 传播根据这些新揭示的条件预测残差，修正后的 logits 重新评估刚刚提交的 token。任何置信度现在低于阈值的 token 将被重新掩码并推迟到具有更多上下文的后续步骤。由于残差编码了每个预测在考虑到其新邻居后如何变化，MRP 识别出那些仅在孤立情况下才自信的揭示。相同的阈值 τ 同时控制揭示和重新掩码，因此不引入额外的调优。
| 模型 | τ | GSM8K | MATH500 | HumanEval | MBPP |
|---|---|---|---|---|---|
| 1.7B | 0.5 | 41.6 → 59.1 (+17.5) | 26.0 → 37.4 (+11.4) | 17.7 → 28.7 (+11.0) | 26.9 → 41.3 (+14.4) |
| 1.7B | 0.6 | 56.3 → 67.0 (+10.7) | 33.4 → 40.4 (+7.0) | 31.7 → 43.3 (+11.6) | 42.4 → 49.0 (+6.6) |
| 1.7B | 0.7 | 65.4 → 71.8 (+6.4) | 39.4 → 48.6 (+9.2) | 40.9 → 45.1 (+4.2) | 49.8 → 51.0 (+1.2) |
| 1.7B | 0.8 | 70.6 → 75.4 (+4.8) | 47.4 → 52.0 (+4.6) | 45.7 → 48.8 (+3.1) | 51.8 → 51.8 (0.0) |
| 1.7B | 0.9 | 76.2 → 77.3 (+1.1) | 51.2 → 57.0 (+5.8) | 49.4 → 52.4 (+3.0) | 53.7 → 54.1 (+0.4) |
| 4B | 0.5 | 63.4 → 81.1 (+17.7) | 44.2 → 58.4 (+14.2) | 32.3 → 53.1 (+20.8) | 38.5 → 50.6 (+12.1) |
| 4B | 0.6 | 76.4 → 85.5 (+9.1) | 53.6 → 61.4 (+7.8) | 49.4 → 57.9 (+8.5) | 49.8 → 57.2 (+7.4) |
| 4B | 0.7 | 84.6 → 88.5 (+3.9) | 60.4 → 65.6 (+5.2) | 60.4 → 62.2 (+1.8) | 61.1 → 63.4 (+2.3) |
| 4B | 0.8 | 87.9 → 90.1 (+2.2) | 66.8 → 70.6 (+3.8) | 64.6 → 62.8 (−1.8) | 63.8 → 64.2 (+0.4) |
| 4B | 0.9 | 88.5 → 90.1 (+1.6) | 69.0 → 70.6 (+1.6) | 67.1 → 65.9 (−1.2) | 65.4 → 64.6 (−0.8) |
| 8B | 0.5 | 67.9 → 82.3 (+14.4) | 45.2 → 58.0 (+12.8) | 32.3 → 54.9 (+22.6) | 34.6 → 49.4 (+14.8) |
| 8B | 0.6 | 79.6 → 86.8 (+7.2) | 54.8 → 63.8 (+9.0) | 48.8 → 63.4 (+14.6) | 48.3 → 59.9 (+11.6) |
| 8B | 0.7 | 85.9 → 89.0 (+3.1) | 60.8 → 69.0 (+8.2) | 64.6 → 72.6 (+8.0) | 54.9 → 60.3 (+5.4) |
| 8B | 0.8 | 89.3 → 91.0 (+1.7) | 68.0 → 70.2 (+2.2) | 74.4 → 75.0 (+0.6) | 62.3 → 66.5 (+4.2) |
| 8B | 0.9 | 90.8 → 91.4 (+0.6) | 70.0 → 72.0 (+2.0) | 75.0 → 75.0 (0.0) | 66.9 → 68.5 (+1.6) |
MRP 重新掩码提高了低阈值动态解码的准确率。对于每个 SDAR 主干网络（1.7B / 4B / 8B）和解除掩码阈值 τ，每个单元格显示基于纯阈值的动态解码 → MRP 重新掩码。相同的 τ 同时控制揭示和重新掩码，因此不引入额外的调优。增益在激进（低）阈值时最大，此时主干网络过度提交，重新掩码有最多的内容可以撤销（在 τ = 0.5 时，8B HumanEval 上最高 +22.6），并随着 τ 升高和主干网络已经保守地解除掩码而趋近于零。推理（GSM8K, MATH500）在每个操作点都有改善；罕见的小幅退化仅限于高 τ 下的代码。准确率以 % 表示。
重新审视刚刚揭示的 token 并重新掩码那些不再成立的 token 的想法出现在先前的工作中，例如 DMax、RCD 和 WINO。这里的新颖之处在于 MRP 用于做出该决策的信号。MRP 不是花费第二次主干网络传播来重新评估揭示，而是读取步骤间残差，该残差已经编码了每个预测在考虑到新 token 后如何变化，并重新掩码那些修正后置信度回落到阈值以下的位置。那些方法需要为重新检查支付一次完整主干网络前向传播的代价，而 MRP 通过一次轻量级的残差传播就获得了。
我们学到的东西
在此过程中，有几件事很突出：
- 深度在 2-3 层有一个最佳点。 准确率从 1 层到 3 层稳步提高，然后趋于平稳，而吞吐量持续下降，因为额外的层只增加了每步延迟；到 8 层时，权衡严格变差。其他推测性起草模块的深度更深。例如，DFlash 大约在 5 层，差异似乎来自于目标模型信息如何到达起草模块。那些设计使用逐层 KV 注入，将主干网络的隐藏状态馈送到每个起草层，而不仅仅是在输入时，因此更深的模块会持续受益于增加的深度。
- 最佳深度取决于模式。 在直接解码中，MRP 的输出被直接使用，更深的模块帮助更大。在推测解码中，验证无论如何都会纠正 MRP 的错误，因此更浅、更快的模块以更高的吞吐量达到相同的最终质量。
- MRP 可以干净地组合。 在其核心，MRP 只是对下一步去噪步骤将产生的结果进行廉价、准确的预测。这个预测是一个其他方法可以构建的原语，这就是为什么同一个模块支持本文中的两种方案，例如推测解码和重新掩码。它为生态系统增加了一种能力，而不是与之竞争。
如果你使用扩散语言模型，MRP 很容易附加，并为你提供了一个以前没有的选择：无损且更快，或者更快同时恢复你否则会损失的质量。
为什么我们选择 Modal
在项目过程中，最让我们印象深刻的事情之一是 Modal 基础设施的质量。它让我们能够将想法作为模块化实验进行测试，并专注于想法而非基础设施，这就是我们如此迅速地从假设到实现 SGLang 实现的原因。
我们还构建了一个轻量级工具 modal-ssh，用于在我们的团队中共享 Modal 工作流。该工具在一个小的 YAML 配置中捕获研究环境（包、卷、密钥和 Git 仓库），然后使用它启动交互式 VM 或长时间运行的后台作业。这显著加速了我们项目的迭代，因为它使我们能够在团队中共享配置，并帮助新成员以最小的学习成本开始使用 Modal。源代码可在此处获取。
我们将这个项目视为我们联合探索快速推理的第一步，还有更多内核和服务层面的方向我们渴望共同推进。