不砸掉重来,给自回归 LLM 装一对"扩散翅膀":Uno 的无损加速之路
上周刷到一篇让我眼前一亮的论文。说实话,"扩散语言模型加速推理"这个方向我已经有点审美疲劳了——LLaDA、Mercury、DiffusionGemma 一波接一波,但看来看去总绕不开一个死结:速度是快了,质量掉了;而且 batch 一大,加速比直接蒸发。
这篇论文的思路很不一样。它不跟自回归(AR)模型较劲,而是问了一个特别务实的问题:既然 AR 模型质量已经够好了,能不能不动它一根汗毛,只在旁边加一小组"扩散权重",让它学会一次并行吐出好几个 token,再由原来的 AR 权重来做验证?
答案是可以。这就是 Uno。
核心摘要
自回归解码一次只能出一个 token,这个串行瓶颈在推理 trace 越来越长、RL 后训练 rollout 占大头的今天,已经成了实打实的成本黑洞。这篇来自 Institute of Foundation Models、UIUC、Cornell Tech、Harvard 和 Cerebras 的论文提出了扩散增强 LLM:每个 Transformer 层保留原有 AR 权重(管质量),再挂一组 rank-128 的 LoRA 作为扩散权重(管速度),通过一个叫 Diffusion Distillation 的轻量蒸馏阶段训练后者,配上一族 Ψ-Spec 采样器做"扩散起草 + AR 验证"的无损拒绝采样。结果相当能打:8B 的 Uno 在所有评测 batch size 下都比 EAGLE-3、DFlash 这些投机解码方法快,最大 batch 下仍保 1.5 倍 以上加速,batch size 1 时达 2.2 倍,还在 agentic 工具调用、编程、长上下文推理上全面赢过 26B 的 DiffusionGemma 和闭源的 Mercury 2。我的判断:这不是底层突破,但工程定位极准——它把并行生成"从换模型降级成了加 LoRA",这个门槛的下降比加速比本身更值钱。
论文信息
- 标题: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
- 机构:Institute of Foundation Models、UIUC、Cornell Tech、Harvard University、Cerebras Systems
- 链接:https://arxiv.org/abs/2609.04010 (代码与权重:https://s-sahoo.github.io/uno/)
- 发表:2026 年 9 月 3 日
🎯 问题动机:batch size 1 的加速比,是个美丽的误会
先聊聊为什么我觉得这篇论文的问题意识比大多数同行清醒。
扩散 LLM(d-LLM)这两年的叙事一直是"并行解码,起飞"。但论文里有一句话直接戳破了窗户纸:d-LLM 的加速比在大 batch 下会消失。原话更狠——"batch-size-one latency captures a narrow operating regime and may overstate speedups that diminish under concurrency"。
你想想看,现在真实的 LLM 负载长什么样?是 agentic 工作负载:一个用户请求背后是一群并行 agent、分支轨迹、工具调用和重试。serving 系统会把跨请求的调用 batch 起来榨干 GPU。在这种场景下,拿 batch size 1 测出来的延迟加速比来吹,基本等于在空无一人的高速公路上测油耗。
推理在本质上常常是 memory bound 的——搬运权重和 KV cache 才是瓶颈,GPU 算力大量闲置。一次预测多个 token 可以摊薄这些内存搬运,这是所有并行解码方法的理论基础。问题在于已有的三条路都走得不顺:
| 路线 | 代表方法 | 核心毛病 |
|---|---|---|
| 投机解码 | EAGLE-3、DFlash | 要养一个独立的 draft 模型,还要维护两套 KV cache |
| 扩散 LLM | LLaDA2.1、DiffusionGemma、Mercury 2 | 质量换速度,有损;大 batch 下加速消失 |
| 多头预测 MTP | Medusa 类 | 改架构加预测头,但头与头之间没有真正的并行联合建模 |
Uno 的切入点:把投机解码的"无损验证"和扩散模型的"并行起草"缝在一起,且起草者和验证者共享同一份基座权重。
🧠 方法核心:一个模型,两套权重
一句话直觉
把每个 Transformer 层的权重矩阵 \(\mathbf{W}_{\text{base}}\) 冻结不动,旁边挂一组 LoRA 适配器 \(\mathbf{W}_{\text{diff}}\)。生成时分两步:扩散通路(\(\mathbf{W}_{\text{base}} + \mathbf{W}_{\text{diff}}\))一次性并行起草 \(B\) 个 token,然后纯 AR 通路(只用 \(\mathbf{W}_{\text{base}}\))做拒绝采样验证,保留最长的合法前缀。
因为验证分布就是原 AR 模型的分布、且 \(\mathbf{W}_{\text{base}}\) 全程冻结,输出分布在数学上与原 AR 模型完全一致——这就是"lossless"的底气。它本质上是投机解码,但 draft 模型不是一个独立小模型,而是基座自己加了一层 LoRA 的"第二人格"。
这个设计带来几个很实在的工程红利:不需要单独设计 draft 模型的深度、宽度、头数;draft 和 verify 共享同一份 KV cache,峰值显存更低;训练开销只有 LoRA 级别。
Diffusion Distillation:一步蒸馏 + 块级扩散
训练扩散权重的目标很直白:让扩散通路一步并行生成的 token 块,分布上逼近 AR 模型串行生成的同一块。作者改造了 Discrete Consistency Distillation(DCD),但做了两个关键手术。
手术一:一步蒸馏。 标准 DCD 要模拟 PF-ODE 的中间状态、构造多步去噪轨迹,在 LLM 尺度上反复构造和存储中间高斯隐变量贵得离谱。Uno 直接把整条轨迹压成一步:从完全腐蚀的序列 \(\mathbf{z}_1 \sim \pi^L\) 直接映射到干净序列 \(\mathbf{x}\),teacher 是冻结的 AR 分布,student 是扩散通路,直接对齐。
手术二:块级扩散。 一整条长序列一步还原太难了,所以把序列切成 \(N\) 个大小为 \(B\) 的块,逐块做单步去噪。干净序列 \(\mathbf{x}\) 和全腐蚀序列 \(\mathbf{z}_1\) 拼成 \([\mathbf{x}, \mathbf{z}_1]\) 一次前向搞定:块因果注意力掩码让 \(\mathbf{x}\) 内部因果、每个噪声块 \(\mathbf{z}_1^{(b)}\) 内部因果、同时噪声块能注意到前面所有干净块。
这里有个挺精巧的实现细节,值得单拎出来说:teacher logits 要在 \(\mathbf{x}\) 的位置上只用 \(\mathbf{W}_{\text{base}}\) 算,student logits 要在 \(\mathbf{z}_1\) 的位置上用 \(\mathbf{W}_{\text{base}} + \mathbf{W}_{\text{diff}}\) 算——同一次前向里两拨 logits 用不同权重。作者用了 gated LoRA,在干净位置把适配器"关掉"、噪声位置"打开",一次前向两全。块级 DCD 目标长这样:
再加一个 TV 损失。 采样器做的是拒绝采样,只保留扩散预测里最长连续合法前缀——所以效率取决于这个前缀的期望长度。作者顺着 Leviathan 等人投机解码论文的 Corollary 3.6,最小化扩散分布和 AR 分布之间的逐块 Total Variation 距离:
总损失是 \(\mathcal{L} = \alpha \mathcal{L}_{\text{DCD}} + \beta \mathcal{L}_{\text{TV}}\)。后面的消融会告诉你一个反直觉的结果:只用 TV 损失反而最好。
Ψ-Spec 采样器:起草、验证、两种模式
采样分两半。起草端用 Ψ-sampler 的转移核从噪声块一步采出候选块,其中第一个 token 永远只用 AR 基座权重采样(所以它 100% 被接受),剩下 \(B-1\) 个用扩散通路并行采。验证端就是标准投机解码的拒绝校正,AR 模型并行验完整块,保留最长合法前缀,被拒绝的位置从重整化的残差分布里补采。
候选集合的构造给了两种模式,对应两种部署场景:
| 模式 | 思路 | 适用场景 |
|---|---|---|
| Linear Sampler | 每个位置直接从边缘分布采一条候选 | 高 batch、compute-bound,最大化系统吞吐 |
| Tree Sampler | 借鉴 Medusa 树采样,每位置取 top-K,按 log-prob 剪枝保留 top-V 前缀,树注意力并行验证 | 低 batch、memory-bound,榨闲置算力,最大化单请求吞吐 |
一个值得记住的数字边界:每次迭代要两次前向(起草一次、验证一次),所以每前向 token 数 TPF 满足 \(1 \leq \mathrm{TPF} \leq \frac{B+1}{2}\)。就算起草全被拒,保底也有 2 个 token 输出(永远被接受的第一个 + 验证者补采的一个)。
作者还埋了个彩蛋:Ψ-Spec 允许去噪步数 \(T\) 超过起草 token 数 \(B\),在不增加上下文长度的前提下做 inference-time scaling。论文坦承这块留给未来工作——这个"固定上下文花更多算力换质量"的方向我觉得挺有想象空间。
🔬 实验:两条战线都打赢了
实验分两条线:一条从零训练 8B 的 Uno(AR 权重在约 23T token 内部数据上预训练),另一条直接拿开源 Qwen3-8B 挂载扩散权重(Uno_Qwen,扩散权重在 OpenThoughts3-1.2M 上训 14.7B token,约 32 小时 / 32 张 H200)。
吞吐测试的方法论也值得说一句:作者没沿用"在下游任务上测吞吐"的老做法(那样会偏爱生成更短 trace 的模型),而是固定 1024 token 输入、8192 token 输出的 1K/8K 吞吐测试,按各方法的平均 TPF 折算解码步数,保证大家都在同样的有效输入输出长度下比。

图 1:上方是训练流水线——NTP 训练 AR 权重之后,插入一个可插拔的 Diffusion Distillation 阶段训练扩散权重,再接 RL 后训练;左下角显示 Uno 相对 AR 基座、Mercury 2、DiffusionGemma 的系统吞吐优势;右侧四个 agentic/长上下文基准上,8B 的 Uno 与自家 AR 基座质量持平,同时大幅领先 26B 的 DiffusionGemma 和 Mercury 2。
战线一:从零训练的 Uno vs 扩散 LLM 们
先看主表,Uno(8B)对 DiffusionGemma(26B-A4B)、Nemotron-Labs-Diffusion(14B)和闭源 Mercury 2:
| 基准 | Uno (8B) | Mercury 2 | DiffusionGemma | Nemotron-Diffusion |
|---|---|---|---|---|
| τ² Telecom(工具调用) | 90.1 | 71 | 68.1 | 14.3 |
| Terminal-Bench v2.1 | 39.6 | 27 | 14.7 | 4.5 |
| SWE-bench Verified | 68.4 | -- | 18.7 | 0.8 |
| AA-LCR(长上下文) | 68.0 | 36 | 19.7 | 7.3 |
| AIME-25 | 90.7 | -- | 74.3 | 40.0 |
| 系统吞吐 toks/sec | 5255 | 1197 | 1136 | 2794 |
| 单请求吞吐 toks/sec | 405 | 769 | 836 | 290 |
几个让我停下来多看两眼的点:
- 5255 toks/sec 的系统吞吐,是 Mercury 2 报告值(1197 toks/sec,还是跑在更快的 Blackwell GPU 上)的约 4.4 倍(论文按 1K 输入 batch 10 的口径算约 4.6 倍)。硬件更慢、精度还是 bf16,吞吐反超四倍多。
- 唯一输的是 AA-Omniscience 这个知识类基准(14.3 vs Mercury 2 的 20),作者归因于模型体量和训练数据差异,诚实。
- 单请求吞吐上 DiffusionGemma(836)确实更快,但质量掉得惨不忍睹——SWE-bench 18.7 对 68.4,这不是一个量级的对话。而且论文点出:DG 和 NLD 在大 batch 下比它们各自的 AR 基座还慢。
对比自家 AR 基座:batch size 1 时约 2.2 倍 加速;AR 基座能塞进单卡 H200 的最大 batch 是 64,此时 Uno 仍快 1.5 倍。加速比在所有 batch size 下都在——这正是论文最想证明的事。
战线二:Qwen3-8B 挂载 vs EAGLE-3 / DFlash
这条战线更贴近大多数人的实际处境:手里有个开源模型,不想重训,只想加速。Uno_Qwen 只加了 0.35B LoRA 参数(EAGLE-3 的 drafter 是 0.40B,DFlash 的扩散 drafter 是 1.05B,三倍于 Uno)。

图 2:系统吞吐(纵轴)vs 单请求吞吐(横轴)在不同并发下的 Pareto 前沿。蓝色 Uno_Qwen 曲线在所有 batch size 下都位于 EAGLE-3、DFlash 和 AR 基座的右上方——Pareto 支配,最高 2.5 倍于 AR 基座。
具体数字(每步解码 token 数 τ 与吞吐):
| 配置 | Uno_Qwen | EAGLE-3 | DFlash |
|---|---|---|---|
| 系统吞吐最优 τ | 3.89 (B=4) | 2.08 (B=4) | 2.07 (B=4) |
| 系统吞吐 toks/sec | 5733 | 4944 | 5351 |
| 单请求最优 τ | 5.97 (B16,V32) | 3.48 (B8,V60) | 2.74 (B=16) |
| 单请求吞吐 toks/sec | 445 | 284 | 370 |
| 峰值显存 GiB | 118.0~122.2 | 129.4~130.0 | 129.8~130.1 |
| 额外参数量 B | 0.35 | 0.40 | 1.05 |
全 batch Pareto 支配、参数最少、显存最低(因为共享 KV cache)。还有一处对 DFlash 训练成本的暴击:DFlash 训练需要 \(B \cdot L\) 的上下文长度,Uno 永远只要 \(2 \cdot L\),与块大小无关。
跟有损扩散加速方法(SDAR、TiDAR、OPDLM、I-DLM、Jacobi Forcing、Fast-dLLM v2、FLARE、LLaDA2.1-Flash)的大表对比里,Uno_Qwen 作为唯一无损方法,TPF 还能打过大多数有损选手——比如 MBPP 上 89.0 的准确率配 4.21 的 TPF,而有损的 TiDAR 准确率掉到 65.4。说实话看到 SDAR 在 AIME-24 上从 76.7 掉到 10.0 的时候我愣了一下,有损加速的代价在难任务上是真的惨烈。
RL 加速与漂移问题
论文还验证了一个我本来要打问号的场景:扩散适配器是拿 SFT checkpoint 蒸出来的,RL 后训练会持续更新 \(\mathbf{W}_{\text{base}}\),draft 分布难道不会漂走吗?
实测:用 DAPO 训数学、代码、工具调用、搜索四个专家(数学专家就生成了 28.9B token 的 rollout),扩散权重全程冻结只做加速。结果端到端 RL 训练最高提速四成,且 RL 之后 TPF 只从 2.25 掉到 2.10——衰减幅度仅 6 个点。适配器比我想象中皮实。
消融里藏着反直觉
- 损失项:只用 \(\mathcal{L}_{\text{TV}}\) 训练的 TPF 是 2.39,比 \(\mathcal{L}_{\text{DCD}}\) 单独(2.23)和两者等权(2.23)都好。作者发现 DCD 损失数值上比 TV 大一个量级,把 DCD 权重压到 0.01 后 TPF 微涨到 2.40。KL 蒸馏在这居然是个配角。
- 课程学习:块大小从 4 逐步升到 16(每半 epoch 升一次)比固定 16 训两 epoch 好,TPF 2.71 vs 2.65。
- LoRA 配置:把适配器摊到所有层比集中堆在少数层好;rank 128 → 256 能提 TPF 但推理变贵;\(\alpha_{\text{LoRA}} / r_{\text{LoRA}} = 64\) 的比例表现最佳。
🤔 我的判断
这篇论文最值钱的地方,是把"并行生成"的落地门槛砍到了地板上。
之前想让 LLM 并行出 token,你要么重训一个扩散模型(接受质量损失),要么养一个独立的 draft 模型(接受显存和维护成本)。Uno 说:给我 32 小时、几十张 H200、0.35B 的 LoRA 预算,我还你一个分布无损、全 batch 段都更快、还能顺手加速 RL rollout 的模型。对于已经拿着开源权重做产品的团队,这个账太好算了。
技术上我最欣赏的是 gated LoRA 那个设计——同一次前向里干净位置关适配器、噪声位置开适配器,teacher 和 student 一箭双雕。TV 损失盖过 KL 蒸馏这个结果也很有意思,它说明对拒绝采样来说,"让连续前缀更容易被接受"比"整体分布更接近"更直接。
但几个地方我保留意见。
一,无损的承诺有个隐性前提。 整套无损性建立在"验证分布 = 冻结的 AR 分布"上,采样参数用的是 temp=1、top-p=0.95。对追求确定性的生产环境,无损是相对这个采样分布而言的,不是相对贪心解码。
二,8B 自家模型的评测里,AR 基座是自己训的、数据是内部的。 质量对标的赢面有多大来自"基座本身训得好",外部读者很难复核。Qwen 那条战线反而更有说服力——基座是公开权重,谁都能验。
三,扩散权重与 AR 权重的数据分布不需要一致(OpenThoughts 训的适配器照样加速 Qwen3,哪怕用 OpenThoughts 直接微调 Qwen3 会掉 15 个点)——这个现象论文给了实验但没给解释。为什么分布失配不破坏接受率?我觉得这里还有文章可做。
四,inference-time scaling 那块只开了个头。 \(T > B\) 时用更多去噪步换质量、甚至质量超过 AR 后就该关掉验证——这个设想很大胆,但目前只有一段讨论。我挺期待后续工作,这可能是扩散通路从"加速工具"升格为"质量工具"的转折点。
放到更大的图景里看:投机解码这两年在 drafter 上卷出了花(EAGLE 系列、DFlash),d-LLM 在另一条战线上卷质量。Uno 实际上宣告了这两条路线的合流点——draft 能力不需要是一个独立模型,它可以是基座的一层薄壳。我敢猜接下来半年会有不少团队复现这个配方,毕竟"冻结基座 + LoRA 蒸馏 + 投机验证"这个组合,工程阻力实在是太小了。
觉得有启发的话,欢迎点赞、在看、转发。跟进最新AI前沿,关注我