DiffusionGemma Technical Report

我们推出了DiffusionGemma——一款实验性质的开源权重语言模型，采用离散扩散（discrete diffusion）机制实现超高速文本生成。与传统自回归（AR）大语言模型逐词解码不同，DiffusionGemma以256词为单位，对文本块进行并行迭代优化，从而规避了自回归模型固有的串行解码瓶颈。该模型并非从零训练，而是在混合专家架构的Gemma-4模型（激活参数量3.8B，总参数量25.2B）基础上微调得到。我们设计了一种计算高效的两阶段训练流程，整体训练所用token总量不足原始AR模型初始训练预算的10%：第一阶段采用监督式微调（SFT），教会模型执行双向去噪；第二阶段则融合强化学习与采样器蒸馏（sampler distillation），同步提升生成质量与推理效率。DiffusionGemma在生成速度与模型能力之间确立了全新的帕累托最优前沿（Pareto frontier）。在其全部评估任务上的平均表现显示，每次前向传播可生成约20个词元，单张NVIDIA H100 GPU上输出吞吐量达约1500词元/秒——即便对比当前最先进的推测解码（speculative decoding）技术，其速度仍显著优于各类自回归模型。此外，DiffusionGemma完整保留了原始模型对“思维模式”（thinking mode）、多模态输入以及长上下文的支持能力。值得注意的是，尽管经过扩散式微调，该模型仍能以自回归方式生成文本，且性能仅出现轻微下降，这为未来构建兼具扩散与自回归特性的混合解码范式提供了可行路径。

解决自回归语言模型固有的顺序解码瓶颈问题，即单次前向传播仅生成一个token导致的低吞吐量和高延迟，尤其在高算力硬件（如H100）上无法充分释放并行计算潜力；该问题并非全新，但高效、保能力、可实用的非自回归/迭代并行生成方案仍属前沿挑战。

将离散扩散（discrete diffusion）范式引入大语言模型生成流程，不从头训练，而是对Gemma-4（MoE架构，3.8B激活参数）进行两阶段高效微调：第一阶段监督学习教模型双向去噪（即同时建模上下文依赖），第二阶段结合强化学习与采样器蒸馏，联合优化文本质量与推理速度；核心创新在于以256-token块为单位并行迭代精炼，而非逐token自回归生成。

• 实现约20 tokens/forward pass的极高吞吐（vs AR模型通常1 token/forward）；• 单H100达~1500 tok/s，显著超越带Speculative Decoding的SOTA AR模型；• 训练仅用<10%原始AR模型token预算，计算高效；• 完整保留原Gemma-4的‘thinking mode’、多模态输入支持与长上下文能力；• 支持无缝回退至AR模式，性能下降轻微，验证了扩散-AR混合解码可行性；• 模型开源（open-weight），但论文未明确提及代码是否同步开源；实验覆盖通用语言理解、推理、代码、多轮对话等综合评测套件；值得深入的方向包括：扩散步数-质量-延迟的理论权衡、块大小自适应机制、跨模态扩散生成统一框架。

• 'Mask-Predict: Parallel Decoding of Sequence-to-Sequence Models' (Gu et al., 2019); • 'Non-Autoregressive Transformer' (Lee et al., 2018); • 'Diffusion-LM Improves Controllable Text Generation' (Li et al., 2022); • 'Multinomial Diffusion for Language Modeling' (Austin et al., 2022); • 'Denoising Diffusion Probabilistic Models for Discrete Data' (Hoogeboom et al., 2021); • 'Speculative Decoding' (Leviathan et al., 2023); • 'Flow-Matching for Scalable Simulation of Discrete Sequences' (Liu et al., 2024)