基于离散扩散解锁大语言模型的无损加速 Unlocking Lossless Speedups in LLMs via Discrete Diffusion
冻结AR权重、外挂LoRA扩散权重并行起草，AR验证实现全批量无损提速
前置知识
自回归语言模型（AR LM）与下一个词预测
AR 语言模型用链式法则分解序列联合分布 $\log p_\theta(x)=\sum_{\ell=1}^{L}\log p_\theta(x^\ell|x^{<\ell})$，以因果 Transformer 实现，训练目标是下一个词预测（NTP）。推理借助 KV 缓存逐 token 生成，质量高但天生串行：生成长度 $L$ 的序列需要 $L$ 次解码决策。
Uno 的质量上限与验证器就是这个 AR 分布，论文的一切加速都以'不改变该分布'为前提，理解无损性必须先懂 AR 因子分解。
离散扩散语言模型（d-LLM）
用前向过程把干净序列按插值分布 $q_t^\ell(\cdot|x;\pi)=\mathrm{Cat}(\cdot;\alpha_t x^\ell+(1-\alpha_t)\pi)$ 腐蚀成先验 $\pi$（掩码或均匀态），再训练去噪网络 $x_\theta(z_t,t)$ 反向恢复，天然支持并行多 token 生成。均匀态先验（$\pi$ 为均匀分布）还能自我纠错、少步生成。
本文的扩散权重本质是一个均匀态离散扩散去噪器，只是换了 NTP 参数化并借 LoRA 挂在 AR 模型上，核心训练损失也来自离散扩散文献。
投机解码与拒绝采样
小草稿模型自回归起草一段候选 token，大目标模型在一次前向中并行计算其似然，按拒绝采样规则接受最长有效前缀，被拒处由残差分布重采样。数学上保证输出分布与目标模型完全一致（无损），加速取决于草稿被接受的长度。
Ψ-Spec 采样器就是把草稿模型换成扩散通路的投机解码变体，TV 损失的理论依据正是 Leviathan 等人关于接受概率的推论 3.6。
LoRA 与门控 LoRA
LoRA 在冻结权重旁加低秩矩阵 $W+BA$（秩 $r\ll d$），只训练增量。门控 LoRA 进一步按 token 位置动态启停适配器：本文在干净序列位置禁用、噪声位置启用，从而一次前向中同一模型既能给出教师 logits（仅 $\theta_{AR}$）又能给出学生 logits（$\theta_{AR}+\theta_\Delta$）。
扩散权重就是一组 LoRA 适配器（rank-128、$\alpha=256$），门控机制是训练与推理能否单次前向完成的关键工程实现。
离散一致性蒸馏（DCD）
把多步离散扩散去噪轨迹压缩成少步生成器：沿确定性离散轨迹取相邻状态 $z_t,z_s$（$s<t$），训练学生分布 $x_\theta^\ell(z_t,t)$ 逼近教师分布 $x_{\theta_0}^\ell(z_s,s)$，损失为逐位置 KL 散度，逐步拉大 $s,t$ 间距让学生学会大步去噪。
Diffusion Distillation 相位把 DCD 改造成单步、按块的形式来蒸馏 AR 教师分布，是本文训练框架的直接基础。
研究动机
LLM 的下一个词预测（NTP）目标造就了自回归生成：每个解码步只产出一个 token，生成 $L$ 个 token 需要 $L$ 次串行前向。作者指出这在两条线上越来越贵：其一，推理链条变长，长思维链和智能体工作流让延迟不断累积，RL 后训练中 rollout 生成本身就主导运行时间；其二，解码常常是访存受限的，权重与 KV 状态的搬运限制了速度并让 GPU 算力闲置。现有方案只解决一半：投机解码（如 EAGLE-3）无损但需要单独训练、对齐良好的草稿模型，并维护两套 KV 缓存；离散扩散 LLM（如 26B-A4B 的 DiffusionGemma、14B 的 Nemotron-Labs-Diffusion）原生并行，但相对 AR 模型存在质量-速度权衡，且加速在大批量下消失——DiffusionGemma 在大批量时甚至慢于自己的基座。而当代智能体负载（并行 agent、分支轨迹、工具调用、重试）恰恰运行在高并发、大 batch 的服务场景，batch=1 的延迟数字会系统性高估实际收益。
本文的目标是本文目标可以概括为：定义一个高质量的 AR 分布，然后学会从同一个分布中并行抽取多个 token，实现无损（输出分布严格不变）的多 token 生成，且不引入独立草稿模型。具体拆解为四点：(1) 提出'扩散增强 LLM'这一新模型类别，把决定回答质量的参数与优化生成速度的轻量参数解耦在同一架构内；(2) 设计 Ψ-Spec 采样器家族，在固定上下文长度下同时支持 AR 验证的无损加速与推理时扩展（去噪步数换质量）；(3) 给出两条落地路径——从零预训练一个 8B 模型，或给现有开源权重 LLM（Qwen3-8B）外挂扩散权重；(4) 在所有评测批量下都取得对基座 AR 模型的加速，包括设备能容纳的最大 batch，并顺带加速 RL 后训练的 rollout 生成。
与已有工作不同的是，独特切入：不做更小的草稿模型，也不把 AR 模型整体改造成扩散模型，而是'一个模型、两套权重'。与多 token 预测（MTP，如 Medusa）修改架构加预测头不同，作者给每一层的每个权重矩阵旁挂一组 LoRA 形式的扩散权重，与前向配对的 AR 权重共享基底；与自投机/混合方法（TiDAR 等）修改基座权重导致分布漂移、只在 batch=1 有效不同，这里 AR 权重全程冻结，验证分布严格不变，因而严格无损且批量无关。训练上复用离散一致性蒸馏（DCD），但把多步去噪轨迹一步化并按块进行，用门控 LoRA 在一次前向中同时算出教师和学生 logits；再辅以总变差（TV）损失直接优化投机解码的接受长度。评估上也纠正了领域惯例：用固定 1024 输入/8192 输出的'1K/8K 吞吐测试'，在真实批量而非 batch=1 下报告系统吞吐。
核心方法
直觉：语言里大量可预测的固定搭配本可整块生成，于是让'加速版'的自己（AR 权重加 LoRA 扩散权重）一口气起草一小块 token，再让'原版'的自己（纯 AR 权重）逐 token 验收；验收沿用投机解码的拒绝采样，输出分布与原 AR 模型完全一致。技术路线：Uno 在因果 Transformer 每个权重矩阵旁加 rank-128、LoRA-$\alpha$=256 的适配器作为扩散权重 $\theta_\Delta$，AR 权重 $\theta_{AR}$ 冻结。扩散通路用均匀态离散扩散腐蚀序列，保留 NTP 参数化（位置 $\ell$ 的 logits 预测下一位置的干净 token），单步把全噪声块 $z_1$ 直接去噪成干净块，一次前向起草 $B$ 个 token。生成时首个 token 仅用 $\theta_{AR}$ 采样（与验证器同分布、必被接受），其余用联合分布 $p_{draft}=\prod_{\ell=2}^{B}x_\theta^\ell$ 采样，再由 $\theta_{AR}$ 以树注意力并行验证，保留最长被接受前缀。每轮需草稿与验证两段前向，TPF 介于 1 与 $(B+1)/2$ 之间。
核心创新是'解耦 + 耦合'的辩证设计：质量参数与速度参数解耦（各自独立的训练目标和训练阶段），但通过共享基座保持分布耦合（LoRA 使草稿分布贴近验证分布），这正是拒绝采样高效的前提。与已有方法的本质区别有三：(1) 不需要独立草稿模型——EAGLE-3 要 0.40B 参数的草稿器、DFlash 要 1.05B 的扩散草稿器，Uno 只加 0.35B 的 LoRA，且草稿与验证共享一份 KV 缓存，峰值显存 122.2 GiB 对比基线的约 130 GiB；(2) 严格无损——d-LLM 微调会改动基座权重造成质量损失，而这里验证器就是未改动的原始 AR 模型，任何批量下加速都不缩水；(3) 训练开销可忽略——Diffusion Distillation 用门控 LoRA 技巧（干净位置禁用适配器、噪声位置启用）在块级因果注意力掩码下对拼接序列 $[x,z_1]$ 一次前向同时得到教师与学生 logits，无需模拟或存储 PF-ODE 中间高斯隐变量；TV 损失按 Leviathan 等人的推论 3.6 直接最小化块级总变差距离，显式提升连续接受长度。
方法步骤详情
训练与推理分四步。第一步：按标准流程（预训练→SFT→RL）训练或直接加载 $\theta_{AR}$（自研 8B 用 23T token、上下文 8K→512K 分阶段；或用开源 Qwen3-8B）。第二步 Diffusion Distillation：冻结 $\theta_{AR}$，把干净序列 $x$ 与全噪声 $z_1$ 各切成长 $B$ 的 $N$ 个块，在拼接序列 $[x,z_1]$ 上用块级因果掩码一次前向得到教师分布 $x_{\theta_{AR}}$ 与学生分布 $x_{\theta_{AR},\theta_\Delta}$（门控 LoRA：干净位置禁用适配器、噪声位置启用），优化 $\mathcal{L}=\alpha\,\mathcal{L}_{DCD}+\beta\,\mathcal{L}_{TV}$（逐块 KL 蒸馏 + 逐块总变差距离），块大小按 2→4→8（自研 8B，7B token，64 张 H200 约 60 小时）或 2→16（UnoQwen，14.7B token，32 张 H200 约 32 小时）课程递增。第三步 Ψ-Spec 采样：在已生成序列后追加 $B-1$ 个先验噪声 token 一步去噪；首 token 由 $\theta_{AR}$ 采样必被接受，其余从 $p_{draft}$ 采样，可选 Linear（单候选，系统吞吐最优，$B=4$）或 Tree（$(B,K,V)=(16,32,32)$：每位置取 top-K、按对数概率留 top-V 前缀、树注意力并行验证）采样器，拒绝采样保证无损。第四步（可选）RL：仅更新 $\theta_{AR}$，冻结适配器继续加速 rollout，RL 后 TPF 仅降约 6%。
技术新颖性
新颖性体现在四个层面。(1) 架构：把扩散去噪从独立模型降格为冻结 AR 模型上的 LoRA 适配器，使扩散成为 AR 的'加速外挂'而非替代品，首次同时做到无损、免独立草稿模型、跨批量有效；(2) 训练：把 DCD 多步一致性蒸馏一步化、块化，借助门控 LoRA 避免存储中间高斯隐变量，开销压到对现有流水线可忽略——自研 8B 只蒸馏 7B token（约 60 小时/64 张 H200），UnoQwen 只需 14.7B token（32 小时/32 张 H200），比基座数万亿 token 的预训练少三个数量级；(3) 目标函数：首次用 TV 距离显式优化投机接受率，消融显示仅 $\mathcal{L}_{TV}$ 即达 TPF 2.39，优于组合损失的 2.23，且 $\mathcal{L}_{DCD}$ 数值大一个量级、需把权重降到 0.01 才微升至 2.40；(4) 采样器：Ψ-Spec 把 dLLM 的预测-校正采样（$\kappa_t$ 控制校正强度）与 AR 验证结合，并揭示新的推理时扩展轴——去噪步数 $T$ 可超过草稿长度 $B$，在不增加上下文长度的前提下用算力换质量，这是传统 test-time scaling 做不到的。
实验结果
五组核心实验。(1) 对基座 AR：1K/8K 吞吐测试下 batch=1 提速约 2.2×，batch=64（H200 最大批量）仍 1.5×；Qwen3-8B 增强版 batch=1 达 2.5×。(2) 对开源 d-LLM（表 1）：8B Uno 全面超过 26B-A4B DiffusionGemma 与 14B Nemotron——Terminal-Bench v2.1 上 39.6 对比 14.7/4.5，SWE-bench Verified 68.4 对比 18.7/0.8，τ²-Telecom 90.1 对比 68.1/14.3，AIME-25 90.7 对比 74.3/40.0；系统吞吐 5255 tok/s 对比 1136/2794。(3) 对闭源 Mercury 2：智能体、编码、长上下文基准全胜（AA-LCR 68.0 对比 36，Terminal-Bench 39.6 对比 27），系统吞吐约 4.6×（5255 对比 1197），仅 AA-Omniscience 落后（14.3 对比 20）。(4) 对无损投机解码（表 2）：temp=1 系统最优配置接受长度 τ=3.89 对比 EAGLE-3 的 2.08、DFlash 的 2.07，吞吐 5733 对比 4944/5351 tok/s；每请求最优 τ=5.97 对比 3.48/2.74，batch=1 吞吐 445 对比 284/370 tok/s，且显存最低（118 GiB）、附加参数最少（0.35B）。(5) RL 加速：DAPO 训练四个专家端到端提速至多 40%，RL 后 TPF 仅降 6%（表 8）。消融（表 12）确认 TV 损失（TPF 2.39 对比 2.23）、块大小课程（2.65→2.71）、全层布置适配器是关键。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| SWE-bench Verified（智能体编码） | pass@1 | 68.4 | DiffusionGemma 18.7；Nemotron-Labs-Diffusion 0.8 | +49.7 pp（对比最强基线） |
| Terminal-Bench v2.1（智能体终端操作） | pass@1 | 39.6 | Mercury 2 为 27；DiffusionGemma 14.7；Nemotron 4.5 | +12.6 pp（对比 Mercury 2） |
| τ²-Telecom（智能体工具调用） | pass@1 | 90.1 | DiffusionGemma 68.1；Nemotron 14.3 | +22.0 pp |
| AA-LCR（长上下文推理） | pass@1 | 68.0 | Mercury 2 为 36；DiffusionGemma 19.7；Nemotron 7.3 | +32.0 pp |
| 系统吞吐（1K/8K 测试，H200 最大批量） | tokens/s | 5255 | Mercury 2 报告 1197；DiffusionGemma 1136；Nemotron 2794 | 约 4.6×（对比 Mercury 2） |
| 平均接受长度（对比投机解码，temp=1，系统最优配置） | τ（token/步） | UnoQwen 3.89 | EAGLE-3 2.08；DFlash 2.07 | 约 +87%（对比 EAGLE-3） |
| 每请求吞吐（batch=1，对比投机解码） | tokens/s | 445（UnoQwen） | EAGLE-3 284；DFlash 370 | 对基座 AR 提速 2.5×，且帕累托支配全部批量 |
| AIME-25（数学，对比有损扩散方法） | pass@1 | 76.7 | SDAR 10.0；Fast-dLLM v2 为 0.0；I-DLM 12.4 | +66.7 pp（同时 TPF 达 3.96） |
局限与改进
作者承认的：(1) 从零训练的 Uno 8B 依赖 23T token 专有数据，不可复现；(2) RL 加速只在四个专家上初步验证（端到端至多 40%），作者明言详细结果下一版给出，ISO-Merger 合并一笔带过；(3) 推理时扩展（去噪步数 $T>B$ 换质量、是否可关闭验证器）只提出概念；(4) 每步需草稿+验证两次前向，TPF 上界 $(B+1)/2$，Quadratic 采样可合并为单次前向但缺高效 kernel。我的补充观察：(5) 草稿质量依赖训练数据与目标域匹配——UnoQwen 用 OpenThoughts 训练适配器可行，但同一数据 SFT 基座反而降质（表 9），适配器在域外提示上的接受率待检验；(6) 评测集中于单张 H200 + bfloat16、32K 上下文，未测量化、超长上下文与其他硬件；(7) 与并行工作 I-DLM 的对比中作者无法复现其无损声明，需社区复核；(8) 首 token 恒被接受意味着收益下限是每两段前向 2 个 token，高熵内容（数学推导）接受率低时加速会向 1× 收缩。
独立分析的弱点
(1) 两段式前向（草稿+验证）在高批量下浪费算力——计算受限时验证额外候选的空间小，Linear 采样器 TPF 从 B=4 的 1.8 到 B=16 仅 2.3 即是例证；改进：实现 Quadratic 采样合并单次前向。(2) 门控 LoRA 依赖按 token 状态启停适配器的路由能力，在不支持自定义 kernel 的推理栈上会打折；改进：把门控行为蒸馏进统一权重或用 MoE 路由近似。(3) 适配器与基座检查点绑定——RL 更新 $\theta_{AR}$ 后草稿分布理论上漂移，论文只测得 TPF 降 6%，更长训练下的漂移动力学未知；改进：RL 循环中周期性轻量重蒸馏。(4) 扩散草稿对高熵、依赖长程推理的 token（数学关键步、代码分支）接受率天然偏低，可能拉大尾延迟方差；改进：按熵或滑动接受率自适应调节块大小与采样器切换。(5) 吞吐评测固定 1024 输入/8192 输出，短输出对话场景与真实负载分布下的收益未报告；改进：补充按生产流量采样的端到端基准。
未来方向
作者明确提出的方向：(1) 系统探索推理时扩展——增加去噪步数 $T$ 使其超过草稿长度 $B$、调节校正强度 $\kappa_t$，在固定上下文长度下用更多计算换更高质量；若质量最终超越 AR 生成，甚至可以关闭验证器以释放上限；(2) 实现 Quadratic 采样的高效 kernel，把草稿与验证合并为一次前向；(3) 与 MTP 预测头结合进一步提升接受率（作者明言二者互补）。基于成果可延伸的方向：(4) 把框架推广到更大规模与 MoE 架构（如 100B 级稀疏模型），检验 LoRA 扩散权重在专家路由下的扩展性；(5) 应用于多模态 AR 模型的并行视觉 token 生成；(6) 研究在线 RL（GRPO/DAPO）中策略更新与适配器漂移的动力学，设计自适应重蒸馏调度；(7) 利用拒绝采样的分布不变性，把加速安全地用于对分布严格敏感的场景（形式化证明、代码执行反馈回路）；(8) 复制 UnoQwen 的'开源基座+领域数据蒸馏'配方到垂直领域（医疗、法律），可能形成低成本的领域加速适配器生态。
复现评估
作者在 s-sahoo.com/uno 承诺开源代码与检查点，采样器已集成进 Nano-vLLM 与 SGLang（论文实验用 Nano-vLLM）。UnoQwen 路线完全可复现：基座为公开 Qwen3-8B，适配器用公开 OpenThoughts3-1.2M 训练 3 个 epoch（14.7B token，rank-128、LoRA-$\alpha$=256，全局 batch 64，恒定学习率 $10^{-5}$、2% warmup，块大小每半 epoch 从 2 增到 16），约 32 小时、32 张 H200——中等规模实验室可负担。从零训练路线不可复现（23T 专有 token），但蒸馏本身不贵（7B token、约 60 小时/64 卡）。评测协议很细：基准套件、采样参数（temp=1、top-p=0.95、top-k=50、262K 上下文）、1K/8K 吞吐测试定义、各基线确切配置均有说明；EAGLE-3/DFlash/Fast-dLLM v2/SDAR 用开源检查点重跑，DiffusionGemma/Nemotron 由公开权重计算，但 Mercury 2 数据取自 Artificial Analysis 追踪器、闭源不可复验。总体：算法与 Qwen 路线复现难度中等，门槛是 32 张 H200 级算力与智能体评测基础设施（Terminal-Bench、SWE-bench 环境）的搭建。
论文图表
