HyperAIHyperAI

Command Palette

Search for a command to run...

通过离散扩散实现 LLM 的无损加速

摘要

大型语言模型(LLM)的成功在很大程度上归功于下一词元预测(NTP),但其自回归(AR)结构要求缓慢的、顺序的词元生成。为了克服这一瓶颈,我们引入了扩散增强型 LLM,这是一类新模型,它在定义 AR 模型分布的同时,利用扩散从该分布中并行抽取多个词元。我们将这些模型的参数解耦为两组:使用标准 NTP 目标训练的 AR 权重,以及用于同时生成多个词元的轻量级扩散权重。扩散权重通过一个简单的扩散蒸馏阶段学习,该阶段对现有 LLM 训练流程的额外开销可忽略不计。我们还引入了 Ψ-Spec,一个采样器家族,它能够在固定上下文长度下实现无损加速和推理时扩展。与推测解码不同,我们的方法不需要单独的草稿模型。与扩散 LLM(d-LLM)不同,它在加速生成的同时不牺牲底层 AR 模型的质量。由此产生的模型称为 Uno,可以从头训练,也可以通过增强现有的开放权重 AR LLM 来构建。在评估的每个批处理大小下,Uno 的吞吐量都高于领先的推测解码方法,并且相对于基础 AR 模型提供了高达 3 倍的加速,包括在设备支持的最大批处理大小下。值得注意的是,我们的 8B Uno 模型在智能体工具使用、编码和长上下文推理的所有评估基准上,都优于领先的开放 d-LLM(26B DiffusionGemma)和专有的 Mercury 2。我们在 https://s-sahoo.com/uno 上发布了代码和检查点。

一句话总结

来自伊利诺伊大学厄巴纳-香槟分校、康奈尔科技、哈佛大学和Cerebras Systems的研究人员提出了Uno,一种扩散增强型LLM类别,将自回归权重与轻量级扩散权重解耦,以并行生成多个token,在不依赖草稿模型或损失质量的情况下,相比基础自回归模型实现最高3×3\times3×加速,并在智能体工具使用、编码和长上下文推理方面超越更大的d-LLM。

核心贡献

  • 提出了扩散增强型LLM,一种模型类别,在定义自回归分布的同时使用轻量级扩散权重并行生成多个token,这些权重通过简单的扩散蒸馏阶段训练,对现有LLM训练流程增加的开销可忽略不计。
  • 提出了Ψ-Spec,一系列采样器,支持在固定上下文长度下实现无损加速和推理时扩展,无需单独的草稿模型,并保持基础自回归模型的质量,不同于投机解码或有损扩散LLM。
  • 由此产生的Uno模型,可从零训练或通过增强开放权重自回归LLM获得,在评估的每个批次大小上相比基础自回归模型实现最高3倍加速,在吞吐量上优于领先的投机解码方法,且8B规模的Uno模型在智能体工具使用、编码和长上下文推理基准上超越26B的DiffusionGemma和专有的Mercury 2。

引言

大型语言模型(LLM)通过下一个token预测(NTP)实现强性能,但这一目标迫使采用自回归(AR)解码,即每步仅生成一个token。这种顺序过程成为长推理轨迹的瓶颈,增加延迟并拖慢强化学习(RL)后训练,其中rollout生成主导运行时间。此外,解码通常受内存限制,尤其在长上下文长度下,导致GPU利用率不足,因为移动模型权重和键值状态限制了速度。现有加速方法存在明显缺陷:投机解码需要单独的草稿模型,离散扩散模型提供有损加速但在大批次大小下失效,多token预测方法则通过额外头修改架构。

作者提出了Uno,一种扩散增强型LLM,在单一架构中统一了自回归和扩散权重。核心思想是定义高质量的自回归分布,并学习从该同一分布并行采样多个token。他们通过在每层标准自回归权重旁增加轻量级扩散权重来实现这一点,在冻结自回归权重后通过扩散蒸馏进行训练。这种设计将响应质量与生成速度解耦,避免了对单独草稿模型或有损自回归到扩散转换的需求。所提出的Ψ-Spec采样器支持无损、经自回归验证的加速和推理时扩展。Uno在基础自回归模型支持的最大批次大小下实现最高2倍加速,同时加速推理和端到端RL训练,并在所有评估的批次大小上优于EAGLE-3和DFlash等投机解码方法,以及开放权重和专有的d-LLM。

方法

作者提出了一种扩散增强型LLM框架,将生成质量与生成速度解耦。该架构在标准自回归模型的每一层上增加一组独立的扩散权重,专门用于并行token生成。这种设计引入了两条不同的路径:自回归权重决定响应质量,通过标准下一个token预测训练;扩散权重通过并行生成token加速推理。在生成时,两组权重同时草拟token,随后仅由自回归权重验证草稿。这种分离保留了既有的自回归训练流程,同时在不牺牲响应质量的情况下实现无损加速。

基础自回归权重记为θAR\theta_{AR}θAR,遵循标准因果掩码范式,其中第\ell个token的分布建模为:

pθAR(xx<)=xθAR1(x)p_{\theta_{AR}}(x^\ell \mid x^{<\ell}) = x_{\theta_{AR}}^{\ell-1}(x)pθAR(xx<)=xθAR1(x)

扩散权重记为θΔ\theta_{\Delta}θΔ,参数化为附加到每个自回归权重矩阵上的低秩适配(LoRA)适配器。这种参数化确保草稿分布与验证分布保持紧密耦合,同时增加最小的内存开销。在草拟过程中,模型使用组合权重θAR+θΔ\theta_{AR} + \theta_{\Delta}θAR+θΔ,而验证路径仅依赖θAR\theta_{AR}θAR。扩散路径保留下一个token预测的参数化形式,在给定噪声序列的情况下预测下一个干净token。

为训练扩散参数,作者引入了扩散蒸馏阶段,将单步扩散分布与token块上的自回归分布对齐。他们改编离散一致性蒸馏,通过单步块去噪逼近自回归分布。训练目标结合了蒸馏损失和总变差损失:

L(θΔ;θAR,α,β)=ExD,z1πL[αLDCD(θΔ;θAR,x,z1)+βLTV(θΔ;θAR,x,z1)]\mathcal{L}(\theta_{\Delta}; \theta_{AR}, \alpha, \beta) = \mathbb{E}_{x \sim \mathcal{D}, z_1 \sim \pi^L} \left[ \alpha \mathcal{L}_{DCD}(\theta_{\Delta}; \theta_{AR}, x, z_1) + \beta \mathcal{L}_{TV}(\theta_{\Delta}; \theta_{AR}, x, z_1) \right]L(θΔ;θAR,α,β)=ExD,z1πL[αLDCD(θΔ;θAR,x,z1)+βLTV(θΔ;θAR,x,z1)]

为使单步去噪可处理,序列被划分为块。作者在拼接的干净和噪声序列的单次前向传播中采用门控LoRA机制。该掩码在干净序列位置禁用适配器,以仅使用θAR\theta_{AR}θAR计算教师logits,并在噪声序列位置启用适配器,以使用θAR\theta_{AR}θARθΔ\theta_{\Delta}θΔ计算学生logits。块级蒸馏损失最小化学生分布与教师分布之间的Kullback-Leibler散度。此外,总变差损失最小化扩散分布与自回归分布之间的距离,以增加拒绝采样期间接受草稿前缀的期望长度。

自回归权重和扩散适配器的训练顺序针对不同部署目标进行结构化设计。作者概述了两种主要训练模式。如果唯一目标是更快的推理,则先完成自回归预训练和后训练,之后冻结自回归权重以训练扩散适配器。或者,为同时加速强化学习rollout和推理,扩散蒸馏在监督微调之后、强化学习后训练之前应用。

如下图所示:

这种课程安排使得所得适配器能够在后续强化学习阶段加速rollout生成。标准策略优化方案更新基础自回归权重,同时保持扩散适配器冻结。尽管强化学习期间基础权重发生变化,低秩适配参数化所维持的紧密耦合确保草稿分布不会与验证分布显著偏离,从而保留推理加速效果。

在推理期间,作者引入了Ψ-投机采样器,以并行生成多个token,同时严格从自回归分布中采样。该采样器使用扩散路径提议一个token块,并针对基础自回归分布执行拒绝采样,以接受最长的有效前缀。为平衡接受长度和验证成本,框架支持两种候选采样策略。线性采样器直接从边际分布生成单个候选序列,在高批次大小下优化整体系统吞吐量。树采样器在低批次大小下利用未充分利用的计算资源,通过采样多个候选并使用树注意力将其作为前缀树并发验证。在两种情况下,第一个token均使用基础自回归权重生成,以保证与验证分布精确匹配,确保无损投机加速。

实验

实验评估了Uno,一种扩散增强型LLM,结合自回归权重与秩为128的LoRA扩散适配器,在两种设置下进行:在专有数据上从零训练,以及在开放权重Qwen3-8B模型上使用不同数据分布(OpenThoughts)训练的适配器进行增强。使用标准化的1K/8K吞吐量测试来控制上下文长度和输出大小,Uno在匹配基础自回归模型质量的同时,在各批次大小下实现1.5倍至2.2倍的更高吞吐量,在所有基准上优于开放权重扩散基线(DiffusionGemma、Nemotron-Labs-Diffusion),并在系统吞吐量上以约4.6倍的优势超过专有的Mercury 2。当应用于Qwen3-8B时,Uno在所有批次大小下均以更少的附加参数和共享KV缓存超越投机解码基线EAGLE-3和DFlash的速度,而消融实验表明,仅使用总变差损失、块大小课程以及分布在各层的适配器可获得最佳的每前向传播token数。

Uno在智能体和编码基准上持续优于开放权重扩散模型和专有的Mercury 2,在智能体任务上增益最大。Uno还实现了比开放权重基线更高的系统吞吐量,尽管使用全注意力,而DiffusionGemma仅在批次大小为1时更快,但准确率显著更低。Uno在所有双方均评估的智能体和编码基准上超过Mercury 2。Uno在每项任务上均优于DiffusionGemma和Nemotron-Labs-Diffusion,尤其在智能体任务上差距显著。Uno实现了比两个开放权重基线更高的系统吞吐量,尽管DiffusionGemma使用跨步注意力。DiffusionGemma仅在批次大小为1时比Uno更快,但其准确率远低。

Uno_Qwen在所有测试的数学基准上,在系统最优和每请求最优设置下均实现比EAGLE-3和DFlash更高的接受长度。它还提供更优的系统吞吐量和每请求吞吐量,同时使用更少的附加参数和共享KV缓存,相比基线降低峰值内存使用。Uno_Qwen在所有数学基准上,对于吞吐量最优配置,均报告比EAGLE-3和DFlash更大的接受长度。最大系统吞吐量在批次大小为4时实现,Uno_Qwen超过每秒5700个token,优于两个基线,并相比基础自回归模型实现1.6倍加速。对于每请求吞吐量,Uno_Qwen和EAGLE-3表现最佳,但Uno_Qwen相比基础模型实现2.5倍加速,超越基线。Uno_Qwen使用共享的草稿-验证KV缓存,不同于EAGLE-3和DFlash维护单独缓存,导致更低的峰值内存使用。Uno_Qwen在所有批次大小下均严格快于DFlash和EAGLE-3,同时需要更少的附加参数。

该表将无损的Uno方法与多种有损的基于扩散的加速方法在数学基准上进行比较,报告准确率和每前向传播token数。Uno在大多数有损方法中持续实现更高的准确率和TPF,尽管优先考虑质量而非速度。有损方法通常显示相对于其父模型的准确率下降,而Uno保持无损性能。Uno在评估的基准上实现比大多数有损扩散方法更高的TPF。TiDAR和Jacobi Forcing等有损方法在AIME-24和AIME-25上相比Uno显示出显著的准确率退化。Uno在GSM8K和MATH500上保持超过96%的准确率,而多种有损方法在这些任务上低于90%。Fast-dLLM v2和LLaDA2.1-Flash在某些基准上显示竞争性准确率,但并非始终优于Uno。

Uno在智能体和编码基准上优于开放权重扩散模型和Mercury 2,在智能体任务上增益最大,同时在使用全注意力的前提下实现比开放权重基线更高的系统吞吐量。Uno_Qwen还在数学基准上,在系统最优和每请求最优设置下,在接受长度和吞吐量上超越EAGLE-3和DFlash,使用更少的参数和共享KV缓存以降低峰值内存。与有损的基于扩散的加速方法相比,Uno在数学基准上保持无损准确率和更高的每前向传播token数,而多种有损方法在AIME-24和AIME-25上显示出显著的准确率下降。


用 AI 构建 AI

从创意到上线——通过免费 AI 协同编码、开箱即用的环境和最优惠的 GPU 价格,加速您的 AI 开发。

AI 协同编码
开箱即用的 GPU
最优定价

HyperAI Newsletters

订阅我们的最新资讯
我们会在北京时间 每周一的上午九点 向您的邮箱投递本周内的最新更新
邮件发送服务由 MailChimp 提供