通过离散扩散解锁 LLMs 的无损加速

TL;DR 用离散扩散蒸馏出轻量并行采样器，让 LLM 一次生成多个 token 而不改变原自回归分布，实现无损加速且无需草稿模型。

问题

问题背景

LLM 推理吞吐受限于 自回归解码(autoregressive decoding) 的串行依赖：每生成一个 token 都必须等待前一个 token 完成，导致 GPU 计算资源大量空闲，尤其在 batch 场景下延迟线性累积。

现有方法局限

• 投机解码(Speculative Decoding) 依赖一个独立的 draft model 来猜测后续 token，但 draft model 与 target model 的分布差异会造成频繁拒绝，且需要额外维护和调优一个模型，显存与工程复杂度上升。
• 扩散语言模型(d-LLMs) 可以并行生成多个 token，但通常重新定义生成分布，无法保证与原始 AR 模型分布严格一致，导致生成质量下降，难以直接替换现有 AR 模型。
• 已有并行解码方案如 Medusa 或 Lookahead Decoding 需修改模型结构或引入额外头，且加速比在不同 batch size 下不稳定。

为什么这个问题难/重要

核心难点在于：如何在 不改变 AR 模型输出分布 的前提下，利用并行计算同时采样多个 token？这要求扩散模块学到与 AR 条件分布一致的近似，并通过精确采样器消除近似误差。业界对推理成本极度敏感，无损加速意味着可以在不牺牲质量的情况下提升 serving 吞吐，对 Agent 工具调用、代码生成、长上下文推理等低延迟场景至关重要。

行业类比

类似 投机解码 在实时聊天机器人中的定位：既要降低首 token 延迟，又不能引入 draft model 的维护负担；Uno 相当于把 draft model 压缩进主模型内部的轻量扩散头，实现“自给自足”的并行推测。

核心洞察

• 核心创新在于将 AR 生成分布与采样机制解耦：模型仍由 AR 权重定义 next-token 条件分布，扩散权重只负责从该分布并行采样多个 token。与 speculative decoding 需要额外 draft model 且可能因拒绝采样浪费计算不同，与 diffusion LLMs 直接替换生成过程导致分布偏移不同，该方法在严格保持原始 AR 模型质量的前提下实现并行加速，因此是无损的。
• `Ψ-Spec` 采样器提供无损加速与 inference-time scaling 的统一接口。它在固定上下文长度下支持不同计算预算的采样，从而灵活权衡吞吐与质量；而 speculative decoding 的上下文长度会随验证失败次数变化，导致批次调度复杂，d-LLMs 则难以在不牺牲质量的情况下动态调节。该设计使得同一模型能适配异构部署环境。
• Diffusion Distillation 阶段只训练轻量 diffusion 权重，开销可忽略，并可直接适配已有 open-weight AR LLM。相比需要维护 draft model 和 target model 的 speculative decoding 方案，或从头训练 diffusion LLM 的大规模成本，该路径显著降低落地门槛，为存量 AR 模型提供可插拔的加速方案。

方法

输入与参数解耦

输入为标准 AR LLM 的上下文 token 序列。模型将参数分为两类：AR 权重 负责基础自回归分布，沿用标准 NTP 目标训练；扩散权重 为轻量模块，通过 扩散蒸馏 阶段从已训练 AR 分布中学习并行生成多个 token。该阶段对现有训练管线开销极小。

扩散蒸馏与并行生成

扩散权重并不独立定义新分布，而是蒸馏自 AR 模型的输出分布。训练时，给定前缀 token，让扩散采样器一次性重建后续多个 token，目标是最小化其与 AR 逐步采样结果的分布差异，从而保证并行采样与自回归采样统计一致。

Ψ-Spec 采样器

推理阶段使用 Ψ-Spec 采样器族，在固定上下文长度下实现无损加速与推理时缩放。工作流程：

1. 扩散权重快速生成一组候选 token 序列；
2. AR 权重对候选序列进行验证与修正；
3. 输出分布与纯 AR 解码严格一致，避免质量损失。 该机制不需要独立的草稿模型，也不改变 AR 模型的原始输出分布。

输出与工程启示

输出为与 AR 模型完全相同的 token 序列分布，生成吞吐大幅提升。与 speculative decoding 相比，省去草稿模型训练与通信开销；与 diffusion LLMs 相比，保留了 AR 基础模型的质量与生态兼容性。

该方法首次在保持 AR 模型分布不变的前提下，仅通过蒸馏轻量扩散权重和专用采样器实现无损并行加速，无需替代 AR 骨干。

实验

实验设计

Uno 实验覆盖 生成速度 与 任务质量 两个维度。速度测试在不同 batch size 下对比 base AR 模型与主流 speculative decoding 方法；质量评估则涉及 agentic tool use、coding、long-context reasoning 三类基准（具体数据集未在摘要中披露），并对比 26B DiffusionGemma 与专有模型 Mercury 2。

关键发现

• Uno 在所有评估 batch size 下吞吐量均超过领先 speculative decoding 方法，且相对 base AR 模型最高实现 3x 加速。
• 8B 参数的 Uno 在全部评估基准上超过 26B DiffusionGemma 和 Mercury 2，证明扩散增强不牺牲 AR 模型的原始能力。

对比解读

• vs speculative decoding：无需单独 draft model，避免了 draft-target 分布不匹配问题，且加速是无损的。
• vs diffusion LLMs (d-LLMs)：传统 d-LLM 往往在质量上妥协，Uno 通过解耦 AR 权重与轻量扩散权重保持原 AR 分布，从而在加速同时守住质量底线。
• 工程启示：该方法可从零训练或增强现有开源 AR 模型，仅需轻量 Diffusion Distillation 阶段，训练开销可忽略，易于集成到生产管线。

行业影响

落地场景

Uno 适用于所有对延迟敏感的 AR LLM 推理服务，包括实时对话、代码补全、Agent 工具调用和长上下文推理。其多 token 并行采样在 batch size 较大时仍保持吞吐优势，适合云端高并发推理 API。

商业价值

• 降本：无损加速提升 token/秒，降低单位推理成本。
• 体验：用户感知延迟显著下降，提升产品留存与付费意愿。
• 部署简化：无需独立 draft 模型，减少内存与运维开销。
• 开源友好：可对现有 open-weight 模型做增量增强，避免从头训练成本。

接口集成

• 推理引擎：需在 vLLM / TensorRT-LLM 等框架中实现 Ψ-Spec 采样器，并对扩散权重做序列化。
• 训练流程：Diffusion Distillation 仅增加轻量参数和有限训练步，可插入现有 NTP 后处理阶段。
• 替换 Speculative Decoding：可作为现有 speculative decoding 方案的无缝替代，不需要额外 draft 模型。

具体 use case

• 企业级智能客服：电商平台客服机器人需低延迟回答用户咨询，Uno 加速可使单个 GPU 服务更多并发会话。
• 代码辅助工具：IDE 插件的实时补全对延迟极敏感，Uno 的无损加速直接提升开发者体验。

局限

• **扩散蒸馏的训练开销与收敛性**：论文虽强调 diffusion distillation 阶段对现有 LLM 训练管线增加的开销“可忽略”，但未给出与全量 NTP 训练相比的精确算力 / 时间数据，也未讨论在更大模型（如 70B+）上蒸馏扩散权重的稳定性和超参敏感性。对于资源有限的团队，轻量化扩散权重的初始化、学习率、采样步数等细节直接影响能否复现 lossless 加速。
• **Ψ-Spec 的适用范围与 batch size 上限**：方法与 speculative decoding 的对比显示在设备支持的最大 batch size 下仍有加速，但实验可能受限于可用显存或特定并行策略。当上下文长度显著增加时，并行采样多个 token 的接受率是否会下降、是否需要动态调整采样温度，论文未展开。此外，在极低延迟场景（如单步生成）下，扩散并行采样带来的额外计算是否抵消加速收益，缺乏敏感性分析。
• **与 diffusion LLMs 的权衡边界**：Uno 实现了不牺牲底层 AR 模型质量的加速，但扩散权重仅用于并行采样，并未改变 AR 分布本身。这意味着其生成多样性 / 可控性依旧由 AR 模型决定，可能无法获得 d-LLMs 在非自回归生成中展现的全局一致性和特定结构化约束下的优势。论文未在数学上严格证明 Ψ-Spec 在所有温度、采样策略下均保持分布一致性，仅提供经验验证。