通过离散扩散实现 LLM 无损加速
来自 arXiv · 期号 2026-09-03
摘要
自回归 LLM 逐 token 生成，速度是瓶颈。 这篇论文提出 Uno，在 AR 骨干上附加轻量扩散权重，实现多 token 并行生成。扩散权重通过一个简单的蒸馏阶段训练，几乎不增加现有训练管线的开销。配套的 Ψ-Spec 采样器族支持无损加速和推理时扩展。 Uno 在最大 batch size 下仍比基础 AR 模型快 3 倍，8B 版本在 agent 工具调用、编程和长上下文推理上全面超过 26B 的 DiffusionGemma 和闭源的 Mercury 2。代码和权重已开源。
深读
导读 自回归 LLM 逐 token 生成的速度瓶颈，是当前推理优化的核心挑战。这篇论文提出 Uno，通过在标准自回归骨干上附加轻量扩散权重，实现多 token 并行生成，且不牺牲生成质量。Uno 在最大 batch size 下仍比基础自回归模型快 3 倍，8B 版本在 agent 工具调用、编程和长上下文推理上全面超越 26B 的 DiffusionGemma 和闭源的 Mercury 2。 关键信息 - Uno 采用参数解耦设计：自回归权重用标准 next-token prediction 目标训练，轻量扩散权重通过一个名为 Diffusion Distillation 的蒸馏阶段训练，对现有训练管线几乎不增加开销。 - 配套的 Ψ-Spec 采样器族支持无损加速（lossless acceleration）和推理时扩展（inference-time scaling），无需像推测解码那样依赖独立的草稿模型。 - 在每档 batch size 下，Uno 的吞吐量均超过主流推测解码方法；在设备支持的最大 batch size 下，仍比基础自回归模型快 3 倍。 - 8B 参数的 Uno 在 agent 工具调用、编程和长上下文推理的全部评测基准上，优于 26B 的 DiffusionGemma 和闭源的 Mercury 2。 逐节讲解 问题与思路：自回归的串行瓶颈 论文首先指出，LLM 的成功很大程度上归功于 next-token prediction，但其自回归结构要求逐 token 顺序生成，速度成为瓶颈。现有加速方案各有局限：推测解码需要额外训练一个草稿模型；扩散 LLM（d-LLM）虽然支持并行生成，但会牺牲底层自回归模型的质量。Uno 的目标是在不牺牲质量的前提下实现并行加速。 方法：扩散增强 LLM 与参数解耦 Uno 将模型参数解耦为两组：自回归权重（AR weights）用标准 NTP 目标训练，保持原有生成质量；轻量扩散权重（diffusion weights）通过 Diffusion Distillation 阶段训练，学会从自回归分布中并行采样多个 token。扩散权重的训练开销极小，可以附加到现有 LLM 训练管线中。Uno 可以从头训练，也可以基于已有的开源自回归 LLM 改造。 采样器：Ψ-Spec 与无损加速 论文提出 Ψ-Spec 采样器族，支持在固定上下文长度下实现无损加速和推理时扩展。与推测解码不同，Uno 不需要独立的草稿模型；与扩散 LLM 不同，它不会降低底层自回归模型的质量。这使得 Uno 在每档 batch size 下吞吐量都超过主流推测解码方法，包括设备能支持的最大 batch size。 评测结果：小模型超越大模型 实验显示，8B 参数的 Uno 在 agent 工具调用、编程和长上下文推理的全部评测基…