← 最新论文
🤖 AI

PRESTO: Prefix-Aligned Tree Drafting for Diffusion Speculative Decoding

该论文介绍了 PRESTO,这是一个通过实现前缀对齐评分和基于优先级的树搜索来解决扩散边际分布与自回归验证之间不匹配问题的原则性框架,从而增强了基于扩散的投机解码,并显著提高了端到端吞吐量。

原作者: Zheng Wang, Zhifan Ye, Qi Cheng, Yonggan Fu, Ziyan Wang, Feng Zhu, Haozhe Zhao, Jan Kautz, Pavlo Molchanov, Humphrey Shi, Minjia Zhang

发布于 2026-07-28
📖 1 分钟阅读☕ 轻松阅读

原作者: Zheng Wang, Zhifan Ye, Qi Cheng, Yonggan Fu, Ziyan Wang, Feng Zhu, Haozhe Zhao, Jan Kautz, Pavlo Molchanov, Humphrey Shi, Minjia Zhang

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

想象一下你正在试图预测一个故事中的下一个词。长期以来,最聪明的计算机(被称为大语言模型)都是一次只做一个词,就像一个人朗读书籍一样,每读一个词就要停下来思考下一个词是什么。这种方式很准确,但速度很慢。最近,科学家们发现了一种使用“扩散”(diffusion)模型来编写这些故事的新方法。把它们想象成雕塑家,从一块大理石块开始,一次性雕刻出整个雕像,而不是一次只凿掉一小块。这使得计算机可以同时猜测许多个词,这非常快。

然而,这其中有一个陷阱。当你同时猜测许多词时,你可能会猜错几个。为了修复这个问题,有一个聪明的技巧叫做“投机采样”(speculative decoding)。这就像是一个快速的初级助手在猜测接下来的几个词,然后由一位缓慢但超级聪明的上司来检查这些猜测是否正确。如果上司同意,他们就会立即接受整批词汇,从而节省大量时间。问题在于,这个初级助手(扩散模型)擅长猜测单个词,但它并不总是知道这些词如何按特定顺序组合在一起。这就像助手非常擅长挑选蛋糕的各种原料,但并不总是知道哪些原料的组合才会真正好喝。

这就是一篇新论文发挥作用的地方。由 Zheng Wang 及其同事领导的研究人员意识到,目前使用这些快速助手的现有方式浪费了很多速度。他们发现,虽然助手可以生成极其多样化的可能的词汇组合,但当前的方法只检查单一路径,就像走在一条单行道上并希望尽头的门是开着的。作者提出了一个名为 PRESTO(前缀对齐树状草拟,Prefix-Aligned Tree Drafting)的新系统。与其走在一条单行道上,PRESTO 构建了一棵可能性的树,同时探索许多不同的路径。但神奇之处在于:它修复了一个关于衡量助手置信度的根本缺陷。助手的原始置信度是“前缀盲目”的,这意味着它不在乎之前的词是什么。PRESTO 添加了一个“前缀对齐”的分数,这就像一个指南针,确保被选出的用于探索的路径是那些最有可能被上司接受的路径。

结果是一个显著更快的系统。在测试中,PRESTO 帮助计算机在每一轮猜测中接受了更多的词。在一些现有的最佳设置上,它使整个过程快了 1.5 倍。在其他设置上,它提供了 1.12 倍的加速。该论文表明,通过将猜测过程视为一场树状分支的冒险而非一条直线,并确保选择的路径是基于它们与目前故事的契合程度,我们可以获得两全其美:扩散模型的速度和仔细检查的准确性。

问题: “单路径”陷阱

要理解为什么需要 PRESTO,想象你在和一位朋友玩“填词游戏”(Mad Libs),而你的朋友正在试图猜测缺失的词。你的这位朋友是一个扩散模型。他们非常擅长观察一个空白处并说:“我猜这里的词是‘猫’!”或者“也许是‘狗’?”或者“或者是‘火箭’?”他们可以同时喊出所有这些选项。

然而,目前使用这位朋友的方法是非常僵化的。它采取朋友的最佳猜测,将其写下来,然后询问“上司”(目标模型)这个猜测是否正确。如果上司说“不”,那么整个过程都会被丢弃,你必须重新开始。如果上司说“是”,你就移动到下一个词并重复此过程。这被称为线性草拟(linear drafting)。这就像是在森林中行走,并且只看正前方的路径。

作者观察到,这种方法效率低下。因为扩散模型可以同时生成许多选项,所以存在一个巨大的“组合空间”。这就像拥有一张印有上千条小径的地图,但你只被允许走一条路。论文显示,通过坚持只走一条路径,系统错失了许多有效的路线。事实上,在像 GSM8K 这样的数学问题上,当前的方法平均接受了约 6.5 个词,但研究人员计算出,如果他们能够检查所有最佳路径,他们本可以接受近 10 个词。这是一个巨大的差距!

不匹配: “盲目”的指南针

研究人员进行了深入研究,发现仅仅通过增加路径(构建树)并不能完美解决旧方法的问题。他们确定了一个“根本性的不匹配”。

在标准的 AI 世界(自回归模型)中,一个词的置信度分数高度依赖于出现在它之前的词。如果句子是“猫坐在……”,模型知道“垫子”是一个非常可能的后续词,但“比萨”则不是。这是前缀对齐(prefix-aligned)。

但扩散模型的工作方式不同。它们独立地为每个位置生成一个“边缘”概率。这就像模型说:“在位置 5,‘猫’的可能性是 80%”,而不关心位置 4 是“那个”还是“那个快速的棕色”。这是前缀盲目(prefix-blind)。

当你尝试使用这些盲目分数来构建一棵猜测树时,你会遇到排名问题。你可能会选择一条对于第一个词看起来很棒、但对于第二个词来说很糟糕的路径,因为模型没有意识到第一个词改变了语境。这就像一个 GPS,它仅根据当前的街道给出方向,却忽略了你刚刚转弯并进入了一条单行道的事实。论文指出,使用这些盲目分数来构建树会导致“不可靠的路径排名”,这意味着系统探索了错误的路径并浪费了时间。

解决方案: PRESTO

PRESTO(基于前缀对齐评分和优先级树搜索的扩散投机解码)通过为扩散模型的评分添加一个“修正”来解决这个问题。

  1. 前缀对齐评分: 作者意识到我们需要将扩散模型的强“边缘”信号(单词自身的可能性)与“前缀调节”信号(给定前序词后的可能性)结合起来。他们创建了一个新的评分公式,该公式将扩散概率乘以一个源自简单 n-gram 模型(一种观察词组组合的轻量级工具)的修正因子。这创造了一个尊重故事流动的分数。
  2. 基于优先级的树搜索: PRESTO 不仅仅是选择顶层路径,而是构建一棵树。它使用新的、经过修正的分数来决定生长哪些分支。它优先考虑那些最有潜力被上司接受的路径。这就像一个徒步旅行者,不再只是直走,而是看一张地图,并选择那条最有希望通往顶峰的路径,即使这条路径在开始时并不是最显眼的。

论文测试了两种生长树的方法:束搜索(Beam Search,在每一步保留固定数量的顶层路径)和最佳优先搜索(Best-First Search,始终扩展目前发现的单一最佳路径)。他们发现,对于他们的特定设置,束搜索的表现与更复杂的最佳优先搜索一样好,因此他们选择了更简单、更高效的选项。

结果:更快且更聪明

作者在各种任务上对 PRESTO 进行了测试,包括数学问题(GSM8K, Math500)、编程挑战(HumanEval, LiveCodeBench)以及聊天对话。他们使用了两种不同类型的系统:

  • 专用扩散草拟器: 一个小的、快速的扩散模型,为较大的自回归模型进行猜测(如 dFlash)。
  • 自我投机扩散 LLM: 一个单一的扩散模型,既进行猜测又进行检查(如 Nemotron-Labs-Diffusion)。

结果在各方面都是一致的。PRESTO 持续增加了平均接受长度,即上司在一次运行中接受的词数。

  • dFlash 系统(使用 Qwen3-8B)上,平均接受长度从约 6.6 个词跳升至 9.6 个词。这转化为 1.5 倍 的端到端加速。
  • Nemotron-Labs-Diffusion 系统上,接受长度从 8.8 增加到 9.9 个词,实现了 1.12 倍 的加速。

或许最令人印象深刻的是,论文展示了即使在系统处于“随机化”(stochastic)状态时,PRESTO 依然有效,而这通常是最难预测的情况。在这种情况下,加速效果更加显著,某些基准测试的吞吐量几乎翻倍。

作者还检查了这种新方法的“成本”。他们发现,构建树和计算新分数所需的额外工作量微乎其微——不到总时间的 4%。绝大部分时间(超过 90%)仍然花在由上司模型进行的实际验证上。这意味着 PRESTO 是一个高效的升级,不会因开销而拖慢系统速度。

PRESTO 不是什么

需要注意的是,该论文并未声称做出的内容。作者明确指出,仅仅应用朴素的树结构(如果不使用他们的前缀对齐评分)是次优的。如果你只是使用扩散模型的原始分数并构建一棵树,你无法获得完整的收益,因为这些分数的“盲目”特性。PRESTO 特指修复这种评分不匹配的问题。

此外,论文并未声称已经彻底解决了扩散模型的问题。他们承认,他们的方法依赖于一个“可处理的前缀对齐信号”(例如他们使用的 n-gram 模型)来承担主要的修正工作。他们建议未来的工作可以探索更丰富的信号,但目前,他们简单的修正足以看到巨大的收益。

为什么这很重要

在让 AI 变得更快、更高效的竞赛中,每一分速度都至关重要。投机解码一直是一个热门话题,因为它允许我们在不支付全额时间惩罚的情况下使用最好、最准确的模型。然而,当前的方法受限于将扩散模型视为线性机器,忽略了它们同时生成许多选项的独特能力。

PRESTO 通过将扩散模型视为多路径探索者,改变了游戏规则。通过将评分与上司检查工作的方式相对齐,它释放了扩散模型并行生成的全部潜力。结果是,它不仅是一个稍快一点的系统,而是一个显著更高效的系统,它允许我们以此前被认为对于这类模型而言是不可能的的速度来生成文本、解决数学问题和编写代码。正如作者所言,他们将“单路径”行走变成了一场“基于树”的远征,确保每一步都是迈向正确答案的一步。

您所在领域的论文太多了?

获取与您研究关键词匹配的最新论文每日摘要——附技术摘要,使用您的语言。

试用 Digest →