Unlocking Lossless Speedups in LLMs via Discrete Diffusion
作者: Subham Sekhar Sahoo, Lingjie Chen, Khiem Pham, Jonathan Geuter, Chaitanya Dwivedi, Varad Pimpalkhute, Yash Akhauri, Alexander Moreno, Mikhail Yurochkin, Zhenting Wang, Mostafa Elhoushi, Nolan Dey, Shane Bergsma, Joel Hestness, John Thickstun, Eric Xing, Zhengzhong Liu
分类: cs.LG
发布日期: 2026-09-03
备注: Code and Checkpoints at https://s-sahoo.github.io/uno/
🔗 代码/项目: PROJECT_PAGE
💡 一句话要点
提出扩散增强LLM以解决自回归生成速度慢的问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大型语言模型 扩散增强 自回归生成 并行生成 速度提升 蒸馏训练 推理加速
📋 核心要点
- 现有大型语言模型的自回归结构导致生成速度缓慢，限制了其应用效率。
- 本文提出扩散增强LLM，通过扩散方法并行生成多个标记，显著提高生成速度。
- Uno模型在多个基准测试中表现优异，8B模型在多项任务上超越了26B的DiffusionGemma和Mercury 2。
📝 摘要（中文）
大型语言模型（LLMs）的成功在于其下一个标记预测（NTP）能力，但自回归（AR）结构导致了缓慢的序列生成。为了解决这一瓶颈，本文提出了一种新型模型——扩散增强LLM，通过扩散方法并行生成多个标记。该模型将参数解耦为两组：使用标准NTP目标训练的AR权重和用于同时生成多个标记的轻量级扩散权重。扩散权重通过简单的扩散蒸馏阶段学习，几乎不增加现有LLM训练管道的开销。此外，本文还引入了$Ψ$-Spec采样器，能够在固定上下文长度下实现无损加速和推理时间扩展。与投机解码不同，该方法不需要单独的草稿模型，并且在不牺牲AR模型质量的情况下加速生成。最终模型Uno在每个评估批量大小上都超过了领先的投机解码方法，并在最大批量大小下实现了高达3倍的速度提升。
🔬 方法详解
问题定义：现有大型语言模型（LLMs）依赖自回归结构进行序列生成，导致生成速度较慢，影响了实时应用的效率。
核心思路：本文提出扩散增强LLM，通过引入扩散方法并行生成多个标记，克服了自回归生成的速度瓶颈。该方法将模型参数分为AR权重和轻量级扩散权重，后者通过简单的蒸馏过程学习。
技术框架：整体架构包括两个主要阶段：首先使用标准NTP目标训练AR权重，然后通过扩散蒸馏阶段学习扩散权重。引入的$Ψ$-Spec采样器在推理时实现无损加速。
关键创新：最重要的技术创新在于通过扩散方法实现并行生成，同时保持AR模型的生成质量。这一方法不需要额外的草稿模型，区别于传统的投机解码方法。
关键设计：模型参数的解耦设计使得扩散权重的学习过程几乎不增加训练开销，且在推理时能够实现高效的多标记生成。
🖼️ 关键图片
📊 实验亮点
实验结果表明，Uno模型在所有评估的批量大小上都优于领先的投机解码方法，尤其在最大批量大小下实现了高达3倍的速度提升。此外，8B的Uno模型在代理工具使用、编码和长上下文推理等任务上超越了26B的DiffusionGemma和Mercury 2。
🎯 应用场景
该研究的潜在应用领域包括实时对话系统、自动文本生成、编程辅助工具等。通过提高生成速度，扩散增强LLM能够在需要快速响应的场景中发挥重要作用，提升用户体验和系统效率。
📄 摘要（原文）
Large Language Models (LLMs) owe much of their success to next-token prediction (NTP), but their autoregressive (AR) structure requires slow, sequential token generation. To overcome this bottleneck, we introduce diffusion-augmented LLMs, a new class of models that defines an AR model distribution while using diffusion to draw multiple tokens in parallel from that distribution. We decouple the parameters of these models into two sets: AR weights, trained using the standard NTP objective, and lightweight diffusion weights, trained to generate multiple tokens simultaneously. The diffusion weights are learned through a simple Diffusion Distillation phase that adds negligible overhead to existing LLM training pipelines. We also introduce $Ψ$-Spec, a family of samplers that enables lossless acceleration and inference-time scaling at a fixed context length. Unlike speculative decoding, our method requires no separate draft model. Unlike diffusion LLMs (d-LLMs), it accelerates generation without sacrificing the quality of the underlying AR model. The resulting models, called Uno, can be trained from scratch or built by augmenting existing open-weight AR LLMs. Uno achieves higher throughput than leading speculative-decoding methods at every evaluated batch size and delivers up to $3\times$ speedups over the base AR model, including at the largest batch size supported by the device. Notably, our 8B Uno model outperforms the leading open d-LLM, the 26B DiffusionGemma, and the proprietary Mercury 2 across all evaluated benchmarks in agentic tool use, coding, and long-context reasoning. We release code and checkpoints at: https://s-sahoo.github.io/uno/