Nemotron-Labs-Diffusion：统一自回归、扩散和自投机解码的三模式语言模型 Nemotron-Labs-Diffusion: A Tri-Mode Language Model Unifying Autoregressive, Diffusion, and Self-Speculation Decoding
提出三模式LM，联合AR和扩散训练，实现AR/扩散/自投机三种解码方式，在精度和速度上超越SOTA开源模型
前置知识
自回归语言模型（AR LM）
AR LM采用严格的从左到右顺序解码，每个token只能基于之前生成的token进行预测。这种方式的推理过程是串行的，每次前向传播只能生成一个token，限制了GPU利用率和推理吞吐量，尤其在低批量部署场景下资源利用率较低。典型的架构包括GPT系列、LLaMA等，它们使用因果注意力掩码确保模型只能看到已生成的token。
本文的核心创新是在保持AR优势的同时扩展到并行解码模式，理解AR LM的串行解码限制是理解为何需要扩散和自投机解码的基础。
扩散语言模型（Diffusion LM）
扩散LM借鉴图像扩散模型的思想，通过逐步去噪过程生成文本。在离散token设置下，通常采用掩码扩散：随机掩码部分token，然后训练模型基于未掩码的上下文恢复被掩码的token。关键优势是可以通过块-wise结构在同一个前向传播中预测多个token，实现并行生成。块内部使用双向注意力，块之间保持因果依赖以便KV缓存重用。
本文将扩散解码作为实现并行token生成的核心机制，理解其双向注意力和去噪过程对于把握三模式架构至关重要。
自投机解码（Self-Speculation）
投机解码是一种推理加速技术，先用小模型快速草拟多个未来token，再用大模型验证这些草拟token，接受匹配的部分从而一次接受多个token。传统的MTP方法如Eagle3需要额外的预测头和独立的草稿模型。本文提出的自投机解码创新性地让同一个模型同时承担草稿和验证任务：扩散模式生成多token草稿，AR模式验证它们，无需额外模型。
自投机解码是本文实现实际部署加速的关键技术，理解其草稿-验证机制和接受率指标对于评估论文的实际贡献非常重要。
KV缓存
KV缓存是Transformer推理加速的关键优化。在生成token时，之前token的Key和Value矩阵可以缓存起来，后续生成只需计算新token的KV并与缓存拼接，避免重复计算。这对AR LM尤其重要，因为每个token都依赖全部历史上下文。在扩散模型中，由于双向注意力的存在，KV缓存的设计更为复杂，需要精心设计注意力模式以支持部分缓存重用。
本文的块-wise扩散设计专门考虑了KV缓存重用（块内双向、块间因果），理解这一点对于把握本文的工程实现细节和实际推理效率很重要。
速度上限（Speed-of-Light, SOL）分析
SOL分析是一种理论分析方法，用于评估并行解码方法的潜在上限。对于扩散模型，SOL通过串行去噪（每次只接受最置信的一个位置）得到收敛后的目标序列，然后用并行方案（如贪心接受或递归动态压缩）尝试在更少的前向传播中重现该序列。SOL接受率代表在理想情况下，每步前向传播能接受多少个token，作为实际方法的性能上限。
本文通过SOL分析证明扩散解码比自投机模式有76.5%的TPF潜力提升，这个分析对于理解扩散解码的长期潜力和当前实现的差距非常关键。
研究动机
自回归语言模型严格的串行token-by-token解码过程从根本上限制了推理并行性，导致资源利用不足和低吞吐量，尤其在低批量部署场景下。例如，在单用户推理时，GPU大部分时间处于空闲状态。扩散语言模型作为有前景的替代方案，能够通过每个前向传播解码多个token实现并行生成。然而，扩散LM在精度和学习效率上往往落后于AR模型，需要更多数据才能达到可比性能，且在实际效率-精度权衡上经常落后于多token预测（MTP）方法。
本文的目标是本文的核心目标是在一个统一架构内同时保留AR和扩散两种范式的优势，实现三模式解码：AR模式用于高并发场景，扩散模式提供灵活的精度-吞吐量权衡，自投机模式结合两者实现低并发场景的极致加速。通过联合AR-扩散训练，使得单个模型能够根据部署需求动态切换解码模式，在不同并发级别和硬件条件下都保持高吞吐量。
与已有工作不同的是，本文的独特切入角度是将AR和扩散视为互补而非竞争的范式。已有工作要么专注于纯AR优化，要么专注于纯扩散改进，或者需要额外的草稿模型。本文通过联合训练让同一个模型同时掌握两种能力，并创新性地提出自投机模式让扩散草稿、AR验证，消除了对额外预测头和独立草稿模型的需求。这种设计使得模型可以作为一个drop-in替代品无缝替换现有AR模型，同时解锁并行解码的潜力。
核心方法
方法整体思路可以概括为：联合训练、三模式推理、部署优化。首先，通过同时优化AR下一token损失和块-wise扩散去噪损失，让模型学习双向表示能力，同时保持AR的强语言先验。训练完成后，模型可以以三种模式推理：(1) AR模式：标准的从左到右生成，适合高并发；(2) 扩散模式：基于置信度的并行块解码，提供灵活的精度-吞吐量权衡；(3) 自投机模式：扩散生成多token草稿，AR验证，适合低并发加速。关键创新是所有三种模式共享同一套参数，无需架构修改。
核心创新点在于三模式统一的架构设计和联合训练策略。与需要额外预测头的MTP方法（如Eagle3）不同，自投机模式复用模型的扩散路径进行草稿、AR路径进行验证，无需独立草稿模型。与纯扩散LM不同，联合训练保留了AR的强语言先验，使扩散模式不会因为学习任意token排列而浪费容量。此外，本文提出训练轻量级采样器、LoRA增强草稿器、递归动态压缩SOL分析等多项技术创新，将理论潜力转化为实际加速效果。
方法步骤详情
方法步骤完整描述如下：(1) 两阶段训练：第一阶段用纯AR目标训练1T token，建立强左到右语言先验；第二阶段联合训练AR和扩散目标300B token，损失权重alpha=0.3平衡两个目标。(2) 联合损失计算：对序列x，AR损失为负对数似然求和，扩散损失采用块-wise设计，将序列分成B个连续块，每次只去噪一个块，总损失为L_AR + alpha * L_diff。(3) 双流注意力模式：训练时输入包含噪声视图和干净视图，噪声视图内块内双向、块间因果，干净视图严格因果，支持在同一前向传播中同时计算AR和扩散损失。(4) 三模式推理：AR模式直接因果采样；扩散模式采用置信度采样或训练的采样器，迭代去噪每个块；自投机模式先让扩散路径并行生成k个草稿token，再用AR路径验证，接受最长匹配前缀。(5) 优化增强：训练轻量级采样器预测每个位置top-1是否正确；用LoRA在扩散草稿路径微调以更好对齐AR验证器，损失结合LK-hybrid分布匹配和token级交叉熵。
技术新颖性
技术新颖性主要体现在四个方面：(1) 三模式统一架构：首次在单一模型内同时支持AR、扩散和自投机三种解码模式，所有模式共享参数，无需模式特定修改。(2) 自投机解码机制：创新性地让扩散模式草稿、AR模式验证，消除了对额外预测头和独立草稿模型的需求，相比Eagle3等MTP方法有更高接受率。(3) 联合训练策略：两阶段训练配合全局损失平均（对批次内所有token全局平均而非序列平均）有效稳定了扩散训练中的高方差问题。(4) SOL分析与采样器优化：通过递归动态压缩方法得到扩散解码的理论上限，证明当前置信度采样远未达到潜力，并设计了基于轻量Transformer的采样器自动预测接受决策，相比固定阈值显著提升精度-吞吐量权衡。
实验结果
核心发现包括：(1) AR和扩散目标互补而非竞争，两者在alpha=0.3时同时达到最佳性能，表明存在共同的最优操作点。添加AR损失使扩散模式平均精度提升7.48%（从57.06%到70.28%），添加扩散损失保持或略微提升AR精度（基础模型+0.14%，指令模型+0.43%）。(2) 8B指令模型在AR模式下平均精度63.61%，超过Qwen3-8B的62.75%；扩散模式实现2.57倍TPF同时保持63.18%精度；LoRA增强的线性自投机模式达到5.99倍TPF，精度62.81%；二次自投机模式达到6.38倍TPF，精度64.04%。(3) 扩散模式的SOL分析显示，在块长度32时平均接受率可达7.60倍，比线性自投机的实际TPF 3.41倍高出76.5%，表明当前采样器仍有巨大提升空间。(4) 在实际部署中，8B模型在GB200 GPU上使用SGLang达到1015 tok/sec（AR模式），线性自投机模式达到3.97倍加速，SOL预测的理论加速可达7.09倍（FP8）和12.36倍（INT4）。(5) 在VLM任务上，扩散模式提供2.46倍到3.15倍TPF，线性自投机达到3.63倍到7.45倍TPF（长文本大于200 token），优势在需要长推理的任务上更明显。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 综合指令 | 平均精度 | 63.61% | Qwen3-8B 62.75% | +0.86% |
| 综合指令 | Tokens-Per-Forward | 5.99x | Qwen3-8B 1.00x | +499% |
| HumanEval | 精度 | 80.49% | Qwen3-8B 81.71% | -1.22% |
| MBPP | 精度 | 85.19% | Qwen3-8B 81.88% | +3.31% |
| Math500 | 精度 | 88.00% | Qwen3-8B 84.80% | +3.20% |
| GSM8K | 精度 | 94.01% | Qwen3-8B 92.42% | +1.59% |
| SPEED-Bench吞吐量 | 加速比(GB200) | 3.97x | Qwen3-8B AR 1.00x | +297% |
| SPEED-Bench吞吐量 | 加速比(RTX 6000) | 3.46x | Qwen3-8B AR 1.00x | +246% |
局限与改进
局限性分析包括作者承认的和观察到的：(1) 当前置信度采样离SOL上限有较大差距，实际TPF约3倍而SOL可达7.6倍，说明采样策略仍需改进。(2) 二次自投机模式虽然理论TPF更高，但由于使用FlexAttention专用注意力掩码的kernel优化不足，实际设备效率落后于线性自投机。(3) 扩散模式在大块长度（如32）时精度有轻微下降，需要在精度和效率间权衡。(4) 不同任务类别的扩散SOL差异很大（3.49倍到11.26倍），开放性生成任务如roleplay的并行性较低。(5) 自投机模式目前只接受前缀token，无法充分利用扩散的非前缀接受模式，限制了理论TPF的完全实现。(6) 训练时双流注意力带来的计算开销，尤其是VLM中视觉token在噪声流中的重复，通过非对称双流设计部分缓解但仍有一定开销。
独立分析的弱点
独立分析的弱点及改进方向：(1) 采样器仍需优化：当前训练的采样器虽比固定阈值有提升（1.3倍TPF或+10.6%精度），但离SOL上限仍有差距。改进方向是设计更复杂的采样器架构，考虑更丰富的特征（如跨位置依赖、局部内容类型自适应），或采用强化学习直接优化接受率。(2) 扩散块长度权衡：大块长度提升并行性但降低精度。改进方向是研究自适应块长度策略，根据局部内容熵动态调整块大小，或在训练中鼓励更激进的并行token采样。(3) 前缀限制：自投机只接受前缀浪费了扩散非前缀接受的优势。改进方向是探索扩散模式验证或使用另一个扩散验证器，验证多个非连续草拟token，进一步提升有效接受率。(4) 生成顺序倾向：扩散解码仍表现出强左到右倾向，主要是token级并行。改进方向是设计训练算法鼓励段级或段落级并行性，更好解锁全局规划能力。(5) 二次自投机kernel优化：当前FlexAttention实现未充分优化。改进方向是开发专用kernel支持二次注意力掩码，或设计其他同时草稿和验证的高效注意力模式。
未来方向
未来研究方向包括作者提出的和基于成果可延伸的：(1) 关闭实际扩散解码与SOL上限的差距：开发更可靠的采样器识别正确token，或更激进的训练方案使条件独立token的并行采样更可行。(2) 改进自投机的草稿-验证对齐：进一步对齐扩散草稿模式和AR验证模式，提升接受率。可用权重共享的子网通过专门训练技术降低草稿成本。(3) 超越前缀AR验证：探索扩散模式验证，可能使用另一个扩散验证器，验证多个非连续草拟token，进一步提升有效接受率。(4) 扩散生成中的更高层并行性：当前主要是token级并行，未来研究段级或段落级并行训练算法，更好解锁全局规划能力和效率潜力。(5) 扩展到更多模态和任务：当前已扩展到VLM，未来可探索音频、视频等多模态场景，或代码生成、推理链等特定任务的优化。(6) 更大规模模型：当前最大14B，未来可探索更大规模（如70B、100B+）的扩展，研究规模对扩散能力和并行性的影响。
复现评估
复现评估：(1) 开源情况：论文承诺在Hugging Face发布Nemotron-Labs-Diffusion模型家族（3B/8B/14B的基础、指令和VLM版本），训练和推理管道通过Megatron Bridge开源。(2) 数据：使用Ministral3的预训练数据集和SFT数据集，具体细节在相关论文中提供。(3) 算力：训练在256个NVIDIA H100 GPU上进行，8B模型训练约1.3T token（1T AR + 300B联合），时间和计算成本较高。(4) 难度：论文提供了详细的训练设置（学习率、权重衰减、批次大小、序列长度）和评估协议，但复现完整训练流程需要大规模计算资源。不过，推理和评估可以在单GPU上进行，论文提供了SPEED-Bench等标准基准的评估细节。(5) 关键超参数：alpha=0.3的损失权重、两阶段训练策略、全局损失平均、块长度选择等都有详细说明，相对容易复现核心实验。
论文图表
图(a)展示了10个任务（HumanEval、MBPP、LiveCodeBench-CPP、GSM8K、Math500、AIME24、AIME25、GPQA、IFEval、MMLU）的平均精度和tokens-per-forward对比。图(b)和(c)分别展示了编码任务（HumanEval、MBPP、LCB-CPP）和数学任务（GSM8K、Math500、AIME24、AIME25）的平均精度对比。图中包括Qwen3-8B、Ministral3-8B、LLaDA-8B、Dream-7B、SDAR-8B等基线。
这张图提供了模型性能的综合评估，证明了三模式模型在精度和效率上的全面优势。特别是在编码和数学任务上，扩散和自投机模式在保持精度的同时实现了显著的吞吐量提升，这对实际应用非常关键。