🤖 AI

xPress: Parallel Refinement for Diffusion Drafters in Speculative Decoding

该论文提出了 xPress，一种轻量级的因果细化器（causal refiner），它通过并行细化来恢复块扩散草拟器（block-diffusion drafters）中缺失的依赖关系，与 dFlash 等现有方法相比，显著提升了投机采样中的接受长度和端到端吞吐量。

原作者： Zheng Wang, Davis Wertheimer, Yu Chin Fabian Lim, Mudhakar Srivatsa, Raghu K. Ganti, Minjia Zhang, Naigang Wang

发布于 2026-08-04

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Zheng Wang, Davis Wertheimer, Yu Chin Fabian Lim, Mudhakar Srivatsa, Raghu K. Ganti, Minjia Zhang, Naigang Wang

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

竞逐语速：为什么 AI 需要更好的草拟系统

想象一下，你正在尝试写一个故事，但你必须遵守一个严格的规则：你一次只能写一个词，而且在写下一个词之前，必须等待一位超级聪明的编辑检查你当前的词。这就是目前大多数强大的 AI 语言模型的工作方式。它们极其准确，但也极其缓慢，因为它们必须一个接一个地检查每一个词。为了提高速度，科学家们发明了一个叫做“投机采样”（speculative decoding）的小技巧。把它想象成一个快速但稍微不太细心的助手，他会为你预猜接下来的几个词。如果这位超级聪明的编辑同意助手的猜测，你就可以一次性写下所有这些词，从而跳过等待时间。

问题在于，这个助手通常过于急躁了。他猜出的词听起来单独很有道理，但放在句子中却无法很好地衔接，就像一位厨师挑选了美味的食材，却忘了检查它们组合在一起是否味道相投。最近，一种被称为“扩散草拟器”（diffusion drafter）的新型助手被创造了出来。它不再是一个接一个地猜测单词，而是尝试一次性猜出一整块单词，就像把一把拼图碎片同时撒在桌子上一样。这非常快，但因为它是同时撒下的，碎片往往无法正确连接。超级聪明的编辑不得不拒绝其中的大部分，导致速度再次变慢。研究人员提出的核心问题是：我们能否保留这种“一次性猜测”者的速度，同时修复错误，让编辑能够真正接受这些词？

xPress 登场：修复拼图的“并行精炼器”

本文介绍了一种被称为 xPress 的巧妙解决方案。作者在研究像 Qwen3-8B 这样的模型时发现，虽然快速的“扩散草拟器”擅长猜测一整块单词，但它忽略了一个至关重要的规则——即单词之间存在相互依赖关系（因果律）。例如，如果第一个词是“she”，那么即使“are”本身是一个常见的词，下一个词也不应该是“are”。扩散草拟器并不知道这一点，因为它是在同时猜测所有内容。

为了解决这个问题，xPress 充当了一个轻量级的因果精炼器。想象一下，扩散草拟器把一把拼图碎片撒在了桌子上，xPress 就是一只快速且聪明的双手，它一眼看去整个堆叠，并在不拆解碎片的情况下，将碎片微调到正确的顺序。它通过一种叫做 Jacobi 解码的技术来实现这一点。x-Press 不是一个接一个地修复拼图碎片（这很慢），而是观察整体画面，同时对每个碎片进行快速调整，然后再次观察。它重复这种“观察并微调”的循环仅几次（通常为 4 到 6 次），直到碎片完美地锁定在一起。

结果令人印象深刻。通过使用 xPress，该系统比原始的快速草拟器平均能多接受约 30% 的猜测词。在某些特定测试（如解决数学问题）中，加速效果高达 56%。在衡量 AI 说话的总速度时，与原始方法相比，xPress 使过程平均快了 1.3 倍，在最佳情况下甚至快了 1.7 倍。

论文明确反对了另外两种试图解决此问题的方法。一种方法涉及构建一个巨大的猜测“树”，这既复杂又昂贵；另一种方法使用“马尔可夫头”（Markov head）按严格的线性顺序逐一修复单词，虽然准确但很慢，因为它失去了同时猜测带来的速度优势。作者证明 xPress 优于两者：它比那种缓慢的、步进式的修复器更快，也比复杂的树状方法更简单。

这项研究证实，xPress 的工作原理是将“因果信息”（即单词依赖于前序单词的规则）重新注入到快速系统中，且不会降低其速度。他们在数学、编程和对话基准测试中进行了测试，发现 xPress 始终优于竞争对手。论文指出，这种方法是一种稳健且适度的改进，它让 AI 能够兼顾速度与准确性，证明了只要有正确的精炼方式，你就无需为了速度而牺牲质量。

技术摘要：xPress：用于扩散草拟（Diffusion Drafters）在投机采样中的并行细化

1. 问题陈述

投机采样（Speculative Decoding, SD）通过使用轻量级草拟模型（drafter）来提议 Token，并由较大的目标模型进行验证，从而加速自回归生成。虽然自回归（AR）草拟模型（如 EAGLE）非常有效，但在生成长块时会面临顺序开销问题。块扩散草拟模型（如 dFlash）通过在单次并行前向传递中生成整个 Token 块来解决这一问题，极大地降低了草拟延迟。

然而，扩散草拟模型存在一个根本性的准确性限制：它们基于边缘分布（marginal distributions）而非联合分布（joint distributions）来预测 Token。在单次传递的扩散过程中，Token $k$ 的采样并未以 Token $1 $到$ k-1$ 的实现值为条件。这种缺乏块内因果关系（intra-block causality）的问题导致生成的序列虽然单个 Token 在局部是高概率的，但在整体上却是不连贯的（例如，主谓不一致）。在验证阶段，目标模型会过早地拒绝这些不连贯的序列，从而严重限制了接受长度（ $\tau$ ），并抵消了并行草拟带来的潜在加速效果。

现有的解决方案（如基于树的草拟，例如 PRESTO、DDTree）试图通过将候选空间扩展为树来缓解这一问题。然而，这些方法计算成本高昂、服务复杂度高，且往往是在绕过问题而非从根本上修复因果不连贯性。

2. 方法论：xPress

作者提出了 xPress，一种轻量级的因果细化器（causal refiner），旨在恢复扩散草拟模型中缺失的因果依赖关系，且无需退回到缓慢的逐 Token 自回归循环。

2.1 核心概念

xPress 作为一种“因果细化器”运行，它接收扩散草拟模型的初始块提案，并通过**并行细化（parallel refinement）**对其进行协调。xPress 并非重新按顺序生成 Token，而是将因果信息注入到扩散草拟模型的 Logits 中，以修正其边缘分布。

2.2 架构

该细化器是一个轻量级模块，它在每个位置 $k$ 向草拟模型的原始 Logits ( $s_k$ ) 中添加一个学习到的逻辑偏置（logit bias） ( $\delta_k$ )。其架构包含四个阶段：

融合（Fuse）： 将三个输入投影到一个低秩空间（ $r=256$ $r = 256$ ）：
- 草拟模型的逐位置隐藏状态 ( $h_k$ )。
- 通过均值池化获得的块全局摘要 ( $g$ )。
- 前一个 Token 预测值的嵌入 ( $x_{k-1}$ )。
混合（Mix）： 在块位置上应用一个严格下三角线性混合器 ( $L$ )。这使得位置 $k$ 能够在单次操作中聚合来自所有前驱节点（ $j \le k$ ）的信息，从而在不进行顺序计算的情况下强制执行因果结构。
变换（Trans）： 低秩空间中的残差 MLP。
读取（Read）： 将结果投影回词表大小，以产生偏置 $\delta_k$ 。

最终的精细化 Logit 为 $\ell_k = s_k + \delta_k$ 。

2.3 通过 Jacobi 解码进行并行细化

一个直接的因果细化器需要进行从左到右的顺序传递，这会重新引入 xPress 旨在避免的延迟。为了解决这个问题，xPress 采用了 Jacobi 解码：

初始化： 使用草拟模型的初始贪婪预测来初始化块。
迭代： 所有位置使用当前块的状态作为下一次迭代的前缀，进行同步并行更新。
收敛： 该过程重复进行 $K$ 次迭代。理论上，对于大小为 $B$ 的块，最多在 $B$ 次迭代内可达到顺序不动点。在实验中，作者发现 $K \approx 6$ 次迭代足以达到与完整顺序解码相匹配的接受长度，因为草拟模型的种子本身已经是一个高质量的猜测。

2.4 训练策略

xPress 与扩散草拟模型进行协同训练，使用复合损失函数：

交叉熵（教师强制/Teacher Forcing）： 最大化给定真实前缀时真实 Token 的概率。
全变分（Total Variation, TV）损失： 最小化精细化分布与目标模型分布之间的距离，直接优化接受率。
一致性损失（Consistency Loss）： 应用于第二次前向传递的辅助损失，其输入为草拟模型自身的预测（自条件化）。这使训练分布与推理时的 Jacobi 过程保持一致，防止细化器过度拟合金标准前缀（gold prefixes）。
草拟锚点（Drafter Anchor）： 应用于基础草拟模型输出的损失项，以确保底层扩散模型保持高质量。

3. 主要贡献

因果依赖注入： xPress 通过一种极简机制（低秩投影和三角混合器）将因果信息注入扩散草拟模型，该机制作用于草拟模型的隐藏状态，避免了参数量的剧增。
并行细化： 通过利用 Jacobi 解码，xPress 以并行方式解决因果依赖。它仅需 $K \ll B$ 次并行迭代即可实现 $B$ 个依赖性自回归步骤的效果，保留了并行草达的效率优势。
一致的端到端收益： 该方法在不修改验证过程或增加草拟模型规模的情况下，显著提升了接受长度和吞吐量。

4. 实验结果

实验在 Qwen3-8B 上使用块大小为 16 的 dFlash 扩散草拟模型进行，涵盖了七个基准测试（数学：GSM8K, MATH-500, AIME25；代码：HumanEval, MBPP, LiveCodeBench；对话：MT-Bench）。

接受长度： 与原始 dFlash 扩散草拟模型相比，xPress 将平均接受长度提升了 ~30%。
- 在贪婪解码（ $T=0$ ）下，增益范围为 18% 到 56%。
- 在采样（ $T=1$ ）下，增益范围为 19% 到 58%。
吞吐量加速：
- 贪婪模式 ( $T=0$ )： xPress 实现比 dFlash 草拟模型平均 1.3× 的端到端加速，峰值可达 1.7×（例如在 GSM8K 上）。
- 采样模式 ( $T=1$ )： 加速提升至平均 1.46×，最高可达 1.6×。
与顺序细化器的对比： 与 Markov head（来自 DSpark 的顺序因果细化器）相比，xPress 在运行速度显著更快的同时，达到了或超过了其接受长度。
- 在 $K=4$ 次 Jacobi 迭代时，xPress 的准确度与 Markov head 持平，但运行速度快了 1.6×（每块 379 µs 对比 601 µs）。
- Markov head 需要固定的 15 步串行解码，而 xPress 在更少的并行步骤内即可收敛。
可扩展性： 在大规模批处理服务场景（vLLM，批大小为 16 和 32）中，xPress 保持了其优势。虽然随着 Batch 增大，相对于 AR 基线的绝对加速比有所压缩（由于计算受限的验证过程），但 xPress 在吞吐量上始终优于纯 dFlash 草拟模型 1.2–1.4×，证明了其细化器在负载下的有效性。

5. 重要性与主张

论文声称 xPress 解决了并行扩散草拟与顺序验证之间的“内在不匹配”问题，即通过修复边缘分布本身而非扩大搜索空间来解决问题。

效率： xPress 证明了注入因果信息的成本相对于草拟开销几乎可以忽略不计。细化成本仅占总草拟时间的极小部分，且其并行性质避免了顺序修正带来的延迟惩罚。
正交性： 该方法与其他改进方案是正交的；它可以与更大的草拟模型或基于树的验证方法结合使用，因为它改进了这些方法所依赖的“骨干” Logits。
泛化性： 作者指出，通过几次并行迭代来协调相互依赖的预测这一理念，可能适用于扩散采样之外的领域，即任何模型同时发射相互依赖输出的场景。

论文总结道，xPress 成功恢复了单次传递扩散中丢失的因果关系，在不牺牲并行效率的前提下，实现了更长的接受链和更高的吞吐量。