一句话总结

Google DeepMind 的研究人员介绍了 DifusionGemma，这是一个开放权重的离散扩散语言模型，由混合专家 Gemma 4 微调而来，激活参数为 3.8B，总参数为 25.2B；该模型并行迭代细化 256 个 token 的块，在 NVIDIA H100 上达到约 1,500 tokens/s，并在保留自回归、多模态和长上下文能力的同时建立了新的速度-能力帕累托前沿。

核心贡献

• DiffusionGemma 是一个实验性的开放权重离散扩散语言模型，通过对 Gemma 4 混合专家检查点（激活参数 3.8B，总参数 25.2B）进行微调得到，以双向注意力并行细化 256 个 token 的块，而不是一次解码一个 token。
• 两阶段训练流程先使用监督微调进行双向去噪，随后结合强化学习和采样器蒸馏，且消耗的训练 token 预算不到原始自回归模型的 10%。
• DiffusionGemma 建立了新的速度与能力帕累托前沿，在单块 NVIDIA H100 GPU 上每次前向传递约生成 20 个 token，每秒约输出 1,500 个 token。它保留了思考模式、多模态输入、长上下文和自回归生成能力，仅有轻微退化；推理示例显示去噪过程中存在双向自校正。

引言

自回归语言模型在推理服务中占主导地位，但在低并发下受限于内存带宽，因为逐 token 生成花在权重传输和 KV cache 上的时间多于计算。投机解码提高了利用率，但仍受顺序草稿或并行草稿接受率下降的限制。文本扩散模型可以同时预测 token 块，但现有方案被迫在速度、智能和可获取性之间权衡：专有 API 不开放，而开放模型在延迟或推理方面往往表现不佳。作者介绍了 DiffusionGemma，这是 Gemma 4 26B A4B 模型的一个开放权重扩散变体。它从预训练的自回归检查点热启动，大约 12 次前向传递即可生成 256 个 token 的块，每次前向传递约生成 20 个 token，在 H100 上达到每秒约 1,500 个 token，并通过监督微调以及随后的采样器蒸馏和强化学习，建立了新的速度与智能帕累托前沿。

数据集

数据来源与构成

• 作者使用开源的领域特定数据集进行下游监督微调。
• 本节明确引用了开源的 Sudoku 谜题数据集，并提到 PubMedQA 用于额外结果。
• Sudoku 评估使用了包含 4096 个谜题的留出集。训练集规模、过滤规则和数据集混合比例在本节中未说明。

示例与序列结构

• 每个训练样本表示为一个序列，包含：
  • P 个 prompt token
  • K 个 canvas
  • 每个 canvas 有 C 个 token
• 因此总序列长度为 P + K*C。
• 除这种 prompt/canvas token 布局外，文本未描述裁剪或元数据构建。

数据如何使用

• 作者使用开源 SFT 工具包在下游数据上微调 DiffusionGemma。
• 编码器处理序列中的所有 token，并使用标准 next-token 交叉熵损失进行训练。
• 解码器损失在每条序列中均匀采样一个 canvas 上计算，并使用 prompt 及任何先前的 canvas 作为上下文。
• 对 50% 的批次样本，解码器还会使用前一次前向传递中的自条件状态；对另外 50%，该状态被置零。
• 在 Sudoku 数据集上进行全量微调后，在包含 4096 个谜题的留出集上实现了超过 85% 的谜题级准确率。
• 还报告了 LoRA 微调，准确率与计算开销的权衡由 LoRA rank 控制。额外 PubMedQA 结果被提及，但本节未详细说明。

方法

作者采用连续时间马尔可夫链（CTMC）方法进行离散扩散。前向过程规定了从干净文本 token 到均匀分布 token 的转移。在给定固定干净起始 canvas $x_0$x0​ 的条件下，前向转移概率路径在每个 token 坐标 $i$i 上独立分解：

$$\mathbb{P}(X_t = x_t \mid X_0 = x_0) = \prod_{i=1}^C \left[ \kappa_t \delta(x_t^i, x_0^i) + (1 - \kappa_t) \frac{1}{V} \right]$$P(Xt​=xt​∣X0​=x0​)=i=1∏C​[κt​δ(xti​,x0i​)+(1−κt​)V1​]

其中 $\kappa_t$κt​ 是从 $\kappa_0 = 1$κ0​=1 到 $\kappa_1 = 0$κ1​=0 平滑变化且单调递减的噪声调度，$\delta$δ 表示克罗内克 delta。随着时间 $t$t 向 1 推进，每个 token 越来越可能被从词表中均匀随机采样的 token 替代。

为了从噪声中重建数据，模型学习逆过程来逆转这种离散损坏。离散流匹配理论表明，要完美逆转前向轨迹，模型需要推断在给定损坏状态下原始未损坏 token 的条件分布。对于给定实现 $X_t = x_t$Xt​=xt​，在时间上向后移动一个小增量 $\Delta t$Δt 由记作 Step 的转移映射控制，该映射输出下一中间步骤的概率分布：

$$\mathbb{P}(X_{t-\Delta t} = \cdot \mid X_t = x_t) \approx \operatorname{Step}\left(x_t, \mathbb{P}(X_0 = \cdot \mid X_t = x_t)\right)$$P(Xt−Δt​=⋅∣Xt​=xt​)≈Step(xt​,P(X0​=⋅∣Xt​=xt​))

作者用神经网络 $p_\theta(\nu \mid x_t)$pθ​(ν∣xt​) 近似在给定含噪状态下干净 token 的真实后验分布 $\mathbb{P}(X_0^i = \nu \mid X_t = x_t)$P(X0i​=ν∣Xt​=xt​)。在生成过程中，他们使用该近似来采样下一个 token 状态。

如上图所示，随着时间反向运行，模型将概率质量从均匀噪声分布平滑地移向有效数据模式。序列中各坐标经历连续时间跳跃转移，展示了并行维度如何随时间协调，而无需按顺序从左到右生成。

DifusionGemma 架构是一个权重共享的编码器-解码器 transformer。作者没有从头预训练扩散模型，而是使用公开的 Gemma 4 26B A4B MoE 检查点初始化模型。这种初始化使其能够继承基础模型的高级能力，例如扩展上下文窗口和原生多模态理解。

为了生成开放式文本，作者采用块自回归生成策略。模型一次去噪一个包含 256 个 token 的 canvas。一旦某个 canvas 完全去噪，它就被提交到序列历史中，模型开始去噪下一个 canvas。

如上图所示，生成流水线包括三个主要阶段。首先，输入 prompt 由因果编码器处理，以初始化键值（KV）缓存。其次，解码器利用对 canvas 内的双向注意力以及到 KV 缓存的交叉注意力，迭代细化含噪 canvas，直到文本完全去噪。第三，最终干净的 canvas 通过因果编码器传回并追加到 KV 缓存，为下一个 token 块设置上下文。

在每个去噪步骤 $t$t，解码器应用权重共享的 transformer 来预测干净 token 的概率分布。解码器接收三个输入：当前含噪 canvas $x_t$xt​、上下文 KV 缓存 $H$H，以及连续自条件信号 $z_t \in \mathbb{R}^{C \times d}$zt​∈RC×d，后者将模型先前的预测反馈给自身。通过跨 canvas token 的双向注意力和到 KV 缓存的交叉注意力，解码器输出未归一化的 logits $L_t$Lt​：

$$L_t = \operatorname{Decoder}_\theta(x_t, z_t, H) \in \mathbb{R}^{C \times V}$$Lt​=Decoderθ​(xt​,zt​,H)∈RC×V

在每次迭代中，作者计算 logits $L_t$Lt​ 并评估干净 token 概率 $\hat{p}_0$p^​0​，更新下一步的自条件信号，并采样细化后的 canvas：

$$\begin{array}{r} \hat{p}_0 = \operatorname{Softmax}(L_t / \tau_t) \in \mathbb{R}^{C \times V}, \\ z_{t-\Delta t} = \operatorname{FFW}(\hat{p}_0 E) \in \mathbb{R}^{C \times d}, \\ x_{t-\Delta t} \sim \operatorname{Step}(x_t, \hat{p}_0). \end{array}$$p^​0​=Softmax(Lt​/τt​)∈RC×V,zt−Δt​=FFW(p^​0​E)∈RC×d,xt−Δt​∼Step(xt​,p^​0​).​

其中 $E \in \mathbb{R}^{V \times d}$E∈RV×d 是 token 嵌入矩阵，FFW 是标准前馈网络。与时间相关的温度 $\tau_t > 0$τt​>0 会使模型预测更加锐利。

作者采用带温度退火和自适应停止的熵受限采样器。token 按熵从低到高的排序被接受，确保它们的互信息界严格低于预定义的误差容限阈值。一旦达到阈值，所有其他 token 会被均匀随机地重新加噪。为平衡收敛速度与语言多样性，token 概率通过温度调节被人工锐化。温度 $\tau_t < 1$τt​<1 在整个分数去噪时间尺度上从初始值 $\tau_{\max} = 0.8$τmax​=0.8 线性退火至 $\tau_{\min} = 0.4$τmin​=0.4。

如上图所示，自适应停止使 DifusionGemma 能够根据任务复杂度和领域动态调整去噪步骤数。采样器根据模型逐步不确定性动态停止去噪过程。当整个 canvas 的平均预测熵低于预定义阈值，并且两个连续去噪步骤的确定性序列预测相同时，就会触发提前终止。

作者从公开的 Gemma 4 26B A4B 检查点开始，进行扩展微调阶段，使模型适应从含噪输入中预测 256 个 token 的块。他们使用块对角注意力掩码，在每个块内实现双向注意力，同时不允许模型以其他去噪块为条件。对于给定的 canvas，模型通过编码器 KV 缓存以 prompt 和先前的 token 为条件。他们使用离散多项分布扩散作为损坏过程，并从词表中均匀采样含噪 token。给定干净上下文、自条件信号 $z_t$zt​ 和含噪 canvas $x_t$xt​，模型被训练以最小化其预测与真实 canvas 之间的交叉熵损失：

$$L(\theta) = - \sum_{i=1}^C \log p_\theta(x_0^i \mid x_t, z_t, H)$$L(θ)=−i=1∑C​logpθ​(x0i​∣xt​,zt​,H)

如上图所示，去噪性能在训练初期迅速提升，之后进入对数线性的性能改进趋势。思考性能受益于扩展 SFT，因为模型最初难以维持连贯的推理轨迹，通常会陷入卡顿或循环。

在 SFT 阶段之后，当使用大量去噪步骤时，模型获得了较强的生成质量。然而，它在高级推理和编码任务上的表现略差于基线 AR 模型，并且在超低延迟推理所需的少步区间中生成质量会坍缩。为解决这一问题，作者引入统一的在线学习阶段，称为采样器蒸馏与强化学习（SD·RL），同时优化两个方向。依靠联合目标，单次梯度更新既推动奖励最大化以提升绝对生成质量和对齐，也推动采样器蒸馏以将高质量生成映射到少步区间。

模型充当在线教师，生成去噪轨迹以建立高质量参考。SD·RL 联合目标使用这些轨迹同时最大化奖励并推动采样器蒸馏。在 SD·RL 训练过程中，在线教师的平均奖励稳步增加，反映出基础能力提升。在自适应停止机制的促进下，在线教师逐步需要更少的有效去噪步骤即可获得这些高奖励。这种加速出现的原因是 SD·RL 目标系统性地降低了模型的预测熵。奖励目标与自适应停止之间的相互作用诱导了课程学习效应。训练早期，较高的预测熵会延迟自适应停止触发。随着模型置信度提高、熵下降，自适应停止会更早触发，从而无缝地将训练分布转向越来越短的去噪轨迹。

实验

评估覆盖 DiffusionGemma 在文本扩散和自回归模式下、有思考和无思考的表现，涉及推理、编码、知识、多模态、指令遵循和 agentic 基准。SD·RL 阶段表明，联合优化奖励最大化和采样器蒸馏可将高质量生成压缩到更少的去噪步骤中；推理分析发现，该模型在低批次服务中最具优势，用内存带宽换取计算。总体结果使 DiffusionGemma 领先于开放文本扩散基线，并与 Mercury 2 具有竞争力；相对自回归基线，文本扩散用一定准确率换取大幅提高的吞吐量。受限输出实验证实了在结构化任务上的快速收敛，所报告的局限性则突出了一些残余问题，例如偶发 token 卡顿和大批次规模下吞吐量较低。

DifusionGemma 是一个混合专家 transformer，词表大小为 262k 个 token。其总参数约为 250 亿，但每个 token 只激活其中一小部分，且激活参数数量不包括视觉编码器。嵌入器在非专家参数中占很大比例，而自条件块带来的开销相对较小。稀疏专家设计通过在许多总专家中使用少量活跃专家，使激活参数远低于总参数。嵌入器大于视觉编码器，也远大于自条件块，后者只增加少量参数成本。

推荐的扩散采样器最多使用 48 个去噪步骤，但通过自适应停止平均约为 12 步；自适应停止和 token 选择使用不同的熵阈值，温度按线性调度从 0.8 降至 0.4。高度受限的输出可在两到三步内收敛，从而相对顺序解码降低延迟。低延迟配置与简洁输出和偶发 token 卡顿相关。自适应停止将典型去噪从最多 48 步降低到平均约 12 步。受限输出可在两到三个去噪步骤内收敛，大幅降低相对顺序解码的延迟。

对比覆盖了文本扩散模型和自回归模型在多个推理、编码和知识基准上的 token 级速度和效率。速度测量使用不同的硬件或 API 条件，因此吞吐量数值反映的是混合服务设置，而非单一受控环境。引用文字表明，DifusionGemma 保留了自回归生成，其 AR 模式表现介于文本扩散模式和基线 Gemma 检查点之间。DifusionGemma 权重可以加载回原始架构进行自回归采样，其 AR 模式表现介于扩散模式和 Gemma 4 初始化检查点之间。速度条件因模型而异：DifusionGemma 和 Gemma 4 在一块 H100 上以 FP8 测量，Nemotron 14B 在一块 H100 上以 bfloat16 测量，LLaDA 2.1 Flash 100B 在八块 B200 上测量，Mercury 2 通过其公开 API 测量。每秒 tokens 和每次前向传递 tokens 在七个具有完整覆盖的基准上取平均，包括 AIME 2026、GPQA Diamond、LiveCodeBench-v6、MGSM、HumanEval、LBPP 和 Natural2Code。Natural2Code 和 HiddenMath 被特别强调为专有且未泄漏的评估集。

在思考模式下，DiffusionGemma TD 在所有报告的基准上相比无思考模式均提高了得分，其中在复杂编码和推理任务（如 Codeforces ELO）上提升最大。这种提升伴随着更高的端到端延迟，由大量增加的总 token 数和总前向次数驱动。两种模式下的每次前向传递 tokens 和每秒 tokens 仍大致相当，而有效去噪步骤因任务而异。思考模式在每个报告任务上都取得了更高的基准得分，其中 Codeforces ELO 增益最大，GSM8K 几乎不变。无思考模式生成的 token 和总前向次数少得多，因此每个样本的端到端延迟短得多。每次前向传递 tokens 和每秒 tokens 等吞吐量指标在两种模式之间相似，因此思考模式更长的生成过程导致了延迟差距。

在微调之前，DiffusionGemma 没有生成任何正确的 Sudoku 棋盘，并且使用了明显更多的去噪步骤。在 LoRA rank 8 微调后，准确率提升到 80% 以上，而去噪步骤数大幅下降，这与较低的预测熵一致。基础 DiffusionGemma 在 Sudoku 上的准确率为零。LoRA 微调将 Sudoku 准确率提高到 80% 以上。微调后的模型所需去噪步骤约为基础模型的四分之一。

DiffusionGemma 是一个稀疏混合专家文本扩散模型，其自适应采样器通常只使用 48 个最大去噪步骤中的约 12 步，高度受限输出的收敛甚至更快。推理、编码和知识任务上的基准对比显示，其自回归模式表现介于扩散模式和原始 Gemma 检查点之间，而吞吐量随服务设置而变化。思考模式在所有报告的基准上都提高了得分，尤其是在复杂编码和推理任务上，但通过更长的生成过程而非单个 token 速度下降来增加端到端延迟。Sudoku 上的 LoRA rank-8 微调实验将准确率从零提高到 80% 以上，并大幅减少所需去噪步骤。