自省扩散模型:首次追平自回归模型质量,并行解码吞吐量翻三倍

0 阅读10分钟

扩散语言模型 (Diffusion Language Models, DLMs) 一直被寄予厚望,它们通过并行地迭代式修正一个 Token 块,有望打破自回归 (Autoregressive, AR) 模型逐字生成的序列瓶颈。然而,这一潜力长期以来未能完全兑现,DLM 在生成质量上始终落后于同等规模的 AR 模型。最近,来自普林斯顿、斯坦福等顶尖机构的研究者们发表了一篇名为《Introspective Diffusion Language Models》的论文,直指问题的核心,并提出了一个全新的范式,首次在质量上追平了同规模的 AR 模型,同时在实际服务效率上远超以往的 DLM。

ArXiv URL:arxiv.org/abs/2604.11…

这篇工作的核心洞察在于,过去的 DLM 普遍缺乏一种关键特性——内省一致性 (introspective consistency)。简单来说,AR 模型在训练和生成时,其内部机制能确保模型“认同”自己已经生成的内容;而 DLM 在其多步去噪的生成过程中,往往会产生与自身预测不符的 Token,导致逻辑连贯性受损。基于此,研究团队设计了内省扩散语言模型 (Introspective Diffusion Language Model, I-DLM),它不仅保留了扩散模型并行解码的优势,更通过一种新颖的训练和解码机制,继承了 AR 模型的内省一致性。最终,I-DLM-8B 在多个基准测试上达到了与强大的 Qwen3-8B AR 模型相当的水平,同时在高并发服务场景下,实现了比先前最先进的 DLM 高出约 3 倍的吞吐量。

图注:(a) 内省一致性示意图:标准 DLM 生成的 Token 分布 q 与模型自身的下一步预测 p 不一致;I-DLM 则通过训练使两者趋于一致。 (b) MATH-500 任务上的质量与吞吐量对比:I-DLM-8B 在性能上匹配了 AR 模型 Qwen3-8B,同时吞吐量是 LLaDA-2.1-mini 的 3.1 倍,是 SDAR 的 4.0 倍。

内省一致性:解开自回归模型强大之谜的关键

要理解 I-DLM 的创新,必须先理解 AR 模型为何如此强大。AR 模型的核心机制是因果掩码 (causal masking)logit 移位 (logit shifting)。在训练时,模型只能看到当前位置之前的所有 Token,并被要求预测下一个 Token。这个看似简单的设计,却隐式地教会了模型一种强大的能力:在每一步生成中,它都在用相同的预测规则“重新审视”和“验证”它已经生成的所有前文。换言之,AR 模型被训练得与自己的生成历史保持高度一致。

研究者将这种特性形式化为“内省一致性”,并定义了一个衡量指标——内省接受率 (introspective acceptance rate),用于评估模型在多大程度上会“接受”自己先前生成的 Token。通过这个视角审视,现有 DLM 的根本缺陷暴露无遗。无论是早期的连续扩散模型,还是后来的离散扩散模型,其训练和生成过程(例如多步双向去噪)破坏了这种一致性。一个 DLM 在某一步去噪后生成的 Token,如果放回模型中,模型在下一步的最优预测可能并非是这个 Token 本身。这种“自我矛盾”是导致 DLM 在需要长程逻辑推理的任务上表现不佳的根本原因。

图注:现有 DLM 的三大局限性:(a) 低内省一致性;(b) 计算效率低下,训练和推理需要更多 FLOPs;(c) 推理引擎不兼容,无法利用成熟的 AR 服务栈。

以往的 DLM 发展路线,大多是从扩散范式出发,试图通过引入各种类似 AR 的机制(如 blockwise 解码、causal-mask 解码)来弥补质量差距。这条路虽然取得了一些进展,但始终是在“修补”。I-DLM 的作者们转换了思路:我们能否从 AR 模型出发,保留其最核心的内省一致性原则,并将其嫁接到一个并行生成框架中? 这个视角的转变,是通向 I-DLM 的关键一步。

I-DLM 的三大支柱:训练、解码与服务

为了实现这一目标,I-DLM 围绕算法和系统进行了协同设计,构建了三大核心支柱:内省一致性训练、内省式步进解码算法,以及与 AR 兼容的高效服务栈。

1. 内省一致性训练 (Introspective-Consistency Training)

I-DLM 的训练方法出奇地简洁高效。它直接取一个预训练好的 AR 模型(如 Qwen3),通过一个简单的单阶段训练过程,将其转换为 I-DLM。这个过程仅需约 45 亿 Token 的数据,远少于从头训练一个模型的成本。其成功的秘诀在于同时运用了三个关键组件:

  • 严格的因果注意力 (Causal Attention):与 AR 模型一样,确保在任何时候,模型对一个位置的预测都只能依赖于它之前的信息。这在生成过程中维持了上下文的一致性。

  • Logit 移位 (Logit Shifting):这是 AR 训练的固有特性,它将输入序列向左移动一位与输出对齐。这不仅是实现下一个词预测的基础,也巧妙地统一了生成与验证过程中的模型隐藏状态,确保了训练的稳定性。

  • 全掩码目标 (All-Masked Objective):在训练时,输入序列的所有位置都被 [MASK] 替换,模型需要一次性预测出所有原始 Token。这种密集监督 (dense supervision) 的方式极大地提高了训练效率。

论文特别指出,虽然这三个技术单独来看并非全新,但之前的工作未能认识到它们组合起来对于实现真正的内省一致性的决定性作用。消融实验清晰地证明了这一点:如果去掉因果注意力和 logit 移位,退化为传统的块扩散模式,模型在推理任务上的性能会发生断崖式下跌(例如,HumanEval 准确率从 92.7% 骤降至 60.3%)。这表明,正是这种训练方式,让模型学会了在并行解码的框架内保持“自我认同”。

图注:训练方法消融实验。完整的 I-DLM (causal + logit shift) 与退化的块扩散模型 (block-causal, no logit shift) 相比,在各项任务上性能优势明显,尤其是在需要复杂推理的 HumanEval 和 MATH 任务上。

2. 内省式步进解码 (Introspective Strided Decoding, ISD)

拥有了一个具备内省一致性的模型后,如何设计一个高效的并行解码算法是下一个关键。I-DLM 提出了新颖的内省式步进解码 (ISD) 算法。与传统 DLM 需要反复迭代去噪不同,ISD 在单次前向传播中同时完成了生成新 Token验证旧 Token 两项任务。

图注:不同解码范式对比。ISD 采用严格的因果注意力,在一次前传中同时生成草稿 Token 并验证前一个 Token,使其能无缝集成到 AR 服务架构中。

具体来说,假设我们要并行生成 NN 个 Token。在每一步解码中:

  1. 模型输入包含已确认的上下文,以及 N1N-1[MASK] Token。

  2. [MASK] 位置,模型会提出 (propose) 新的候选 Token。

  3. 在紧邻上下文的第一个生成位置,模型会验证 (verify) 上一步生成的 Token 是否仍然是当前上下文下的最优选择。这个验证是基于模型内生的 AR 分布进行的。

  4. 如果验证通过,该 Token 被正式接受。模型向前推进一个位置,并继续生成下一批 N1N-1 个草稿 Token。如果验证失败,则拒绝该 Token,并以 AR 模式生成一个新 Token。

这种“生成与验证”统一的单次前传机制,消除了传统 DLM 反复去噪带来的巨大计算开销。ISD 的输出在严格模式下可以被证明与基础 AR 模型的分布相匹配,这意味着它在提速的同时没有牺牲理论上的正确性。

研究者用 TPFTPF (Tokens Per Forward) 指标来衡量解码效率。在理想情况下(所有提出的 Token 都被接受),TPFTPF 等于步进大小 NN。在实际应用中,I-DLM 的接受率非常高(通常 85%\geq 85\%),使得一个 N=3N=3 的步进就能实现约 2.3-2.4 倍的实际加速,而计算开销仅增加约 2 倍。这意味着大部分并行解码的理论优势都转化为了真实的吞吐量提升。

3. AR 兼容的服务栈 (AR-Compatible Serving)

这是 I-DLM 最具现实意义的贡献之一。以往的 DLM 由于其独特的解码逻辑(如多步去噪、块同步),需要定制化的推理引擎,无法利用社区为 AR 模型积累的大量优化成果,如 PagedAttention、连续批处理 (continuous batching) 等。这导致它们在实际部署中效率低下,尤其是在高并发场景。

由于 I-DLM 的训练和解码都保留了严格的因果注意力结构,它可以作为“即插即用”的扩展,直接集成到现有的 AR 服务系统(如 SGLang)中。这意味着 I-DLM 天然继承了所有先进的 AR 服务优化:

  • 连续批处理:ISD 每一步都保证至少产出一个高质量 Token,所有请求可以均匀向前推进,不会像块扩散那样因为批次内某个请求的收敛速度慢而卡住整个批次。

  • KV 缓存优化:可以复用 PagedAttention 等高效的 KV 缓存管理机制。

  • CUDA 图优化:可以将整个 ISD 步骤捕获为单个 CUDA 图,大幅减少内核启动开销。

此外,团队还开发了静态批次调度器 (stationary-batch scheduler) 等定制优化,进一步减少了 ISD 解码循环中的 CPU 开销。这些系统层面的设计,确保了 I-DLM 的理论并行优势能够真正转化为生产环境中的高吞吐量。

实验验证:首次实现质量比肩与效率超越

I-DLM 的实验结果令人信服,它在质量和效率两个维度上都树立了新的标杆。研究团队基于 Qwen3-8B 和 Qwen3-32B 训练了 I-DLM-8B 和 I-DLM-32B 两个版本。

在质量方面,横跨 15 个基准测试,I-DLM 是第一个在同等规模下性能媲美强大 AR 模型的 DLM。如下表所示,I-DLM-8B/32B 的表现在绝大多数任务上都与基座 AR 模型 Qwen3-8B/32B 持平或非常接近。与此同时,它显著优于之前所有的 DLM,例如在 AIME-24 和 LiveCodeBench-v6 这两个高难度编码任务上,I-DLM 的得分比 LLaDA-2.1-mini (16B) 分别高出 26 和 15 分之多。

图注:端到端质量对比。I-DLM 在 15 个基准测试上的准确率(%)。下划线表示 30B 以下参数规模的最佳非 AR 模型结果。结果表明 I-DLM 的性能与同规模 AR 模型相当,并远超其他 DLM。

在效率方面,I-DLM 的优势更加突出。与依赖辅助小模型进行草稿生成的推测解码 (speculative decoding) 方法(如 EAGLE3)相比,I-DLM 在所有并发水平下都展现出更高的单请求吞吐量。更值得注意的是,即使是产生与基座 AR 模型输出完全一致的“无损”版 I-DLM (R-ISD),其吞吐量在多数情况下也超过了 EAGLE3。随着并发数的增加,I-DLM 凭借其与 AR 服务栈的良好兼容性,吞吐量能持续扩展,而其他并行解码方法则会更快遇到瓶颈。最终,I-DLM 在高并发服务场景下实现了比先前最先进 DLM 高 2.9 至 4.1 倍的吞吐量。

结论

I-DLM 的工作为并行解码领域带来了突破性的进展。它通过回归 AR 模型的设计本源,识别出“内省一致性”是过去 DLM 未能企及 AR 质量的“缺失环节”。基于这一深刻洞察,I-DLM 通过一套简洁而强大的训练、解码和服务协同设计范式,成功地将 AR 模型的质量保证与扩散模型的并行潜力结合起来。

这项工作不仅在理论上提出了一个更优的 DLM 设计方向,更重要的是,它提供了一个可以直接部署、性能卓越且与现有生态系统兼容的解决方案。它证明了我们不必在生成速度和质量之间做出艰难的权衡。对于日益增长的大模型推理需求,尤其是高并发服务场景,I-DLM 展示了一条极具吸引力的技术路径,预示着并行解码模型在未来实际应用中的广阔前景。