工作日早上 7 点读 AI，周日早上 8 点读周报订阅邮件

METAL LAB

DiffusionGemma Technical Report

arXiv:2608.001462026-07-30

谷歌的DiffusionGemma是一款实验性开放权重模型,通过并行打磨256个token组成的区块,让文本生成速度远超逐词生成的传统AR模型

与传统自回归(AR)模型一次生成一个token不同,DiffusionGemma使用离散扩散技术,一次性并行迭代打磨256个token组成的区块来生成文本。它并非从零训练,而是在混合专家(MoE)架构的Gemma 4模型基础上微调而成,该模型激活参数38亿、总参数252亿,微调仅使用了原AR模型训练token预算的不到10%。在完整评测集上平均,它每次前向传播生成约20个token,在单张H100 GPU上可达到约每秒1500个输出token,比使用最先进推测解码的AR模型还要快得多。

METAL LAB 解读图

DiffusionGemma两阶段流水线如何实现快速文本生成

证据状态已报告实测结果

  1. 起点:Gemma 4 AR模型从预训练的Gemma 4 26B A4B MoE权重初始化(激活参数38亿,总参数252亿),而非从零训练
  2. 第一阶段:SFT监督微调教模型在256-token画布上学习双向去噪,仅使用原AR模型训练token预算的不到10%
  3. 第二阶段:SD·RL采样器蒸馏结合强化学习,同时提升基于奖励的生成质量并压缩所需的去噪步数
  4. 推理:区块式扩散解码从随机噪声画布出发,熵约束采样器配合自适应停止平均约12步完成一个256-token区块,再拼接进KV缓存
  5. 结果:全新的速度-质量帕累托前沿单张H100上每秒约1500个token、每次前向传播约20个token,超越了配备最先进推测解码的AR模型
这是 METAL LAB 制作的解读图，并非论文作者提供的原图。

他们做了什么

  1. 传统AR模型一次只能生成一个token,在请求量较低时GPU算力被浪费,因为时间主要花在把模型权重和KV缓存从显存搬运到计算单元上(受限于显存带宽)。
  2. DiffusionGemma用离散扩散绕开了这个瓶颈:从一个由随机噪声组成的256-token画布出发,用双向注意力同时并行地逐步打磨整个区块。
  3. 训练分两个阶段:第一阶段用监督微调(SFT)教模型学会双向去噪;第二阶段称为SD·RL,结合采样器蒸馏与强化学习,同时提升生成质量并压缩所需的去噪步数。
  4. 熵约束采样器搭配自适应停止机制,能在模型预测足够自信且稳定时提前结束去噪,最多可将延迟降低约4倍而不牺牲质量。
  5. 由于与Gemma 4共享完全相同的transformer架构,微调后的权重仍可切换回标准自回归模式运行,为扩散与AR的混合解码留下了空间。
Figure 1: Pareto plot of quality versus output decoding speed comparing DiffusionGemma to the Gemma 4 model family and other diffusion models. Quality and output speed are calculated as the average across GPQA-Diamond and LiveCodeBench-v6 for all models. Gemma 4 and DiffusionGemma output speeds are measured on a single H100 (FP8, batch size 1).2 Nemotron 14B is measured on a single H100 (bfloat16, batch size 1). LLaDA 2.1 Flash 100B is measured on 8× NVIDIA B200 GPUs (bfloat16, batch size 1).3 Mercury 2 was measured via the public OpenRouter API with high reasoning effort (see Appendix E for speed estimation). The light blue dots and dashed line represent the Gemma 4 family of models without MTP.
Figure 1: Pareto plot of quality versus output decoding speed comparing DiffusionGemma to the Gemma 4 model family and other diffusion models. Quality and output speed are calculated as the average across GPQA-Diamond and LiveCodeBench-v6 for all models. Gemma 4 and DiffusionGemma output speeds are measured on a single H100 (FP8, batch size 1).2 Nemotron 14B is measured on a single H100 (bfloat16, batch size 1). LLaDA 2.1 Flash 100B is measured on 8× NVIDIA B200 GPUs (bfloat16, batch size 1).3 Mercury 2 was measured via the public OpenRouter API with high reasoning effort (see Appendix E for speed estimation). The light blue dots and dashed line represent the Gemma 4 family of models without MTP.
Figure 2: Overview of our two-stage training pipeline that converts an autoregressive model (Gemma 4 26B A4B) into a text diffusion model (DiffusionGemma). Initialized from the AR model weights, the model first undergoes SFT that adapts it to discrete text diffusion and bidirectional attention across 256-token canvases. This is followed by an online sampler distillation and reinforcement learning phase, which jointly maximizes reward-driven generation quality and compresses the denoising steps to unlock ultra-low latency.
Figure 2: Overview of our two-stage training pipeline that converts an autoregressive model (Gemma 4 26B A4B) into a text diffusion model (DiffusionGemma). Initialized from the AR model weights, the model first undergoes SFT that adapts it to discrete text diffusion and bidirectional attention across 256-token canvases. This is followed by an online sampler distillation and reinforcement learning phase, which jointly maximizes reward-driven generation quality and compresses the denoising steps to unlock ultra-low latency.
Table 1: Parameter counts. DiffusionGemma’s architecture is a mixture-of-experts transformer with a vocabulary of 262k tokens. The total number of activated parameters does not include the vision encoder. It includes an additional MLP block for the purpose of self-conditioning.
Total25.2B
Activated3.85B
Vision Encoder550M
Embedder740M
Self-Conditioning7.8M
Active / Total Experts8 / 128
+ 1 shared
Figure 3: Stylized example of discrete diffusion probability paths and parallel sampling trajectory. For illustrative purposes, the state space uses distinct, token-specific vocabularies: adjectives on the horizontal axis and nouns on the vertical axis. As time moves backward from t=1.0 (noise) to t=0.0 (data), the marginal distribution smoothly interpolates, concentrating mass away from the uniform noise distribution and onto valid data modes. Black circles track the discrete jump transitions of an individual sequence realization (from “blue moon” at t=1.0 to “red sunset” at t=0.0), demonstrating how parallel canvas dimensions coordinate non-autoregressively over time.
Figure 3: Stylized example of discrete diffusion probability paths and parallel sampling trajectory. For illustrative purposes, the state space uses distinct, token-specific vocabularies: adjectives on the horizontal axis and nouns on the vertical axis. As time moves backward from t=1.0 (noise) to t=0.0 (data), the marginal distribution smoothly interpolates, concentrating mass away from the uniform noise distribution and onto valid data modes. Black circles track the discrete jump transitions of an individual sequence realization (from “blue moon” at t=1.0 to “red sunset” at t=0.0), demonstrating how parallel canvas dimensions coordinate non-autoregressively over time.
Figure 4: The DiffusionGemma generation pipeline. The process consists of three main stages: 1) Context encoding: The input prompt is processed by the causal encoder to initialize the Key-Value (KV) cache. 2) Denoising loop: A noisy canvas is iteratively refined by the decoder, using bidirectional attention across the canvas and cross-attention to the KV cache, until the text is fully denoised. 3) Encode & append: The finalized clean canvas is passed back through the causal encoder and appended to the KV cache, setting the context for the next block of tokens.
Figure 4: The DiffusionGemma generation pipeline. The process consists of three main stages: 1) Context encoding: The input prompt is processed by the causal encoder to initialize the Key-Value (KV) cache. 2) Denoising loop: A noisy canvas is iteratively refined by the decoder, using bidirectional attention across the canvas and cross-attention to the KV cache, until the text is fully denoised. 3) Encode & append: The finalized clean canvas is passed back through the causal encoder and appended to the KV cache, setting the context for the next block of tokens.
Table 3: Comparison of model performance, TPS and TPF across various benchmarks. TPS excludes prefill time; “−” denotes missing data. For speed measurements: DiffusionGemma and Gemma 4 are measured on 1× H100 (FP8, batch size 1); Nemotron 14B on 1× H100 (bfloat16, batch size 1); LLaDA 2.1 Flash 100B on 8× B200 (bfloat16, batch size 1); Mercury 2 via its public API (see Appendix E for speed estimation). TPS, TPF and total tokens are averaged over the 7 benchmarks for which we have full coverage: AIME 2026, GPQA Diamond, LiveCodeBench-v6, MGSM, HumanEval, LBPP, and Natural2Code. TPS and TPF for Gemma 4 (MTP) are measured using SPEED-Bench (1). The Natural2Code and HiddenMath rows are highlighted as they are proprietary, unleaked evals.
Open-weight ModelsClosed-weight Model
DiffusionGemmaGemma 4LLaDA 2.1 FlashNemotron DiffusionMercury 2
26B A4B26B A4B100B14BUnknown
ModeTDTD (No-think)ARAR (No-think)AR (MTP)AR (MTP, No-think)TD (S Mode)TD (Diffusion Mode)HighMedium
AIME 202669.150.884.257.588.380.080.040.091.782.5
GPQA Diamond73.264.679.867.282.373.768.747.075.266.7
LiveCodeBench-V669.160.671.458.377.172.639.428.679.474.9
Codeforces ELO14299591569105917181529718-19861629
BigBench EH47.640.059.142.264.856.2--48.943.8
GSM8K96.395.896.696.196.796.445.0-96.595.8
MGSM84.880.787.984.392.991.56.869.391.991.2
MMMLU81.576.382.278.086.378.0--81.980.6
MMMU Pro54.366.063.366.773.872.5----
Putnam67.457.174.759.781.072.9-45.873.673.6
HumanEval94.592.798.297.698.897.690.286.098.298.2
BigCodeBench46.041.947.745.950.248.1-33.547.645.3
LBPP81.068.986.374.189.577.345.740.789.285.0
IFEval97.494.597.295.798.797.8-72.197.094.5
Tau2 Retail71.557.575.461.085.579.0----
Tau2 Airline69.049.072.050.076.051.0----
Tau2 Telecom28.132.033.832.043.034.2----
MMLU-Pro77.677.978.879.182.682.6--77.675.5
Natural2Code94.090.196.292.396.394.786.973.379.171.3
HiddenMath80.674.385.477.587.281.6-44.382.782.3
Output Speed (TPS)1479151220420430330337549600547
Tokens Per Forward (TPF)19.7418.761.001.001.401.404.631.79--
Average Total Tokens4,0018295,1841,0257,2071,8164,3719413,8821,222
Figure 5: Adaptive stopping enables DiffusionGemma to dynamically adjust its number of denoising steps to task complexity and domain. We report median, first and third quartiles of the effective denoising steps (Equation 9). See Table 4 for full benchmarks and latency metrics.
Figure 5: Adaptive stopping enables DiffusionGemma to dynamically adjust its number of denoising steps to task complexity and domain. We report median, first and third quartiles of the effective denoising steps (Equation 9). See Table 4 for full benchmarks and latency metrics.
Figure 6: Evolution of downstream performance during SFT. Prior to SFT, the model is incapable of denoising text. Only a moderate amount of SFT is needed to achieve good performance in non-thinking mode, however extended SFT is crucial for the model to learn thinking behaviour. Results use entropy-bounded sampling with adaptive stopping and a maximum of N=192 denoising steps.
Figure 6: Evolution of downstream performance during SFT. Prior to SFT, the model is incapable of denoising text. Only a moderate amount of SFT is needed to achieve good performance in non-thinking mode, however extended SFT is crucial for the model to learn thinking behaviour. Results use entropy-bounded sampling with adaptive stopping and a maximum of N=192 denoising steps.
Table 4: Performance and latency metrics of DiffusionGemma TD in Thinking vs. No-Thinking mode. We report accuracy/score, Tokens Per Forward (TPF, higher is faster), Tokens Per Second (TPS), Effective Denoising Steps (DNS, lower is faster), Total Forwards, Total Tokens, and End-to-End Latency per sample in seconds (excluding prefill time). The latency metrics are averaged across all samples within each benchmark.
Score (↑)TPF (↑)TPS (↑)Effective DNS (↓)Total Forwards (↓)Total TokensE2E Time (s) (↓)
BenchmarkThinkNo-ThinkThinkNo-ThinkThinkNo-ThinkThinkNo-ThinkThinkNo-ThinkThinkNo-ThinkThinkNo-Think
AIME 202669.150.819.316.71365.41333.012.614.1390.691.16,4451,3094.720.98
GPQA Diamond73.264.616.716.51207.81330.215.113.4443.448.15,6477264.680.55
LiveCodeBench-V669.160.618.516.91278.31333.413.814.0581.8195.77,5341,8475.891.39
Codeforces ELO142995915.114.0950.51040.317.118.5959.6521.111,6224,27912.234.11
BigBench EH47.640.020.817.71390.21415.111.912.8434.968.19,0621,2336.520.87
GSM8K96.395.823.224.11866.21966.49.17.543.413.28832980.470.15
MGSM84.880.719.016.81526.71367.511.511.963.619.51,0852910.710.21
MMMU Pro54.366.017.715.41351.31255.113.815.3191.631.63,1784722.350.38
Putnam67.457.118.116.01330.81282.413.414.4303.877.44,7251,1033.550.86
HumanEval94.592.723.024.31838.21981.29.48.055.614.31,1743050.640.15
BigCodeBench46.041.919.619.41560.11579.011.310.177.321.81,4103940.900.25
LBPP81.068.920.519.21509.81545.411.610.9264.379.44,7308593.130.56
IFEval97.494.517.29.01368.3732.113.014.4100.824.01,4642391.070.33
Natural2Code94.090.121.121.01682.71706.510.59.670.524.71,3914650.830.27
HiddenMath80.674.321.019.41591.21564.711.311.5206.449.73,5018442.200.54
Figure 7: Downstream performance during SFT improves log-linearly with training progress. The log-linear trend for thinking performance starts off at a lower point yet exhibits a steeper slope than non-thinking mode. Results use the EntropyBounded sampler with adaptive stopping and a maximum of N=192 denoising steps.
Figure 7: Downstream performance during SFT improves log-linearly with training progress. The log-linear trend for thinking performance starts off at a lower point yet exhibits a steeper slope than non-thinking mode. Results use the EntropyBounded sampler with adaptive stopping and a maximum of N=192 denoising steps.
Figure 8: SD⋅RL training simultaneously increases average reward and reduces the effective denoising steps of the online teacher (training metrics shown in buckets of 200 steps). Both effects together push the quality-speed Pareto frontier of the model.
Figure 8: SD⋅RL training simultaneously increases average reward and reduces the effective denoising steps of the online teacher (training metrics shown in buckets of 200 steps). Both effects together push the quality-speed Pareto frontier of the model.
Table 6: PubMedQA performance of the finetuned model with LoRA rank 4. Finetuning leads to a slight increase in accuracy on a model with a good base performance.
ModelEffective Denoising StepsAccuracy (%)BLEU
DiffusionGemma18.0975.610.76
+ LoRA finetuning31.5776.6220.67
Figure 9: SD⋅RL significantly advances the quality-speed Pareto frontier. The SD⋅RL configuration uses the DiffusionGemma sampler with a maximum of N=48 denoising steps. The SFT frontier is derived by sweeping N from 48 to 192 (note that higher N yields milder temperature annealing). Both quality (y-axis) and inference efficiency (x-axes, see Section 3.4) are calculated as the average between GPQA-Diamond and LiveCodeBench-v6 in thinking mode, averaged over 3 seeds.
Figure 9: SD⋅RL significantly advances the quality-speed Pareto frontier. The SD⋅RL configuration uses the DiffusionGemma sampler with a maximum of N=48 denoising steps. The SFT frontier is derived by sweeping N from 48 to 192 (note that higher N yields milder temperature annealing). Both quality (y-axis) and inference efficiency (x-axes, see Section 3.4) are calculated as the average between GPQA-Diamond and LiveCodeBench-v6 in thinking mode, averaged over 3 seeds.
Figure 10: Performance vs. number of denoising steps N, without adaptive stopping or temperature annealing. Performance is measured as the average score over GPQA-Diamond, LiveCodeBench-v6 (3 seeds). To remove confounding effects of early stopping and temperature annealing, we disabled adaptive stopping and use temperature τt=1.
Figure 10: Performance vs. number of denoising steps N, without adaptive stopping or temperature annealing. Performance is measured as the average score over GPQA-Diamond, LiveCodeBench-v6 (3 seeds). To remove confounding effects of early stopping and temperature annealing, we disabled adaptive stopping and use temperature τt=1.
Table 7: Summary of downstream finetuning hyperparameters for the Sudoku and PubMedQA recipes. Both LoRA recipes use adapters applied to all linear layers and train for 2,000 steps. A full-weight alternative is also provided for Sudoku, using Adafactor to manage memory on 8 GPUs.
HyperparameterSudoku (LoRA)Sudoku (Full)PubMedQA
LoRA rank8—4
Canvas size256256128
Number of canvases112
Prompt length2562561024
Batch size282
Peak learning rate3×10−41.125×10−41.0×10−4
End learning rate3×10−51.125×10−51.0×10−5
Training steps8,0002,0002,000
OptimizerAdamAdafactorAdam
LR scheduleCosine with warmupCosine with warmupCosine with warmup
Warmup iterations400100100
Weight decay10−410−410−4
Min. hardware2× A100 80GB8× A100 80GB2× A100 80GB
Figure 11: Per-step GPU time breakdown: DiffusionGemma processes 256 tokens per step with only a 3.2× increase in per-step latency compared to single-token AR generation. Serving a single request (batch size 1) on a H100, FP8 precision, with 4096 input tokens, 1024 output tokens.
Figure 11: Per-step GPU time breakdown: DiffusionGemma processes 256 tokens per step with only a 3.2× increase in per-step latency compared to single-token AR generation. Serving a single request (batch size 1) on a H100, FP8 precision, with 4096 input tokens, 1024 output tokens.
Figure 12: Trade-off between total and per-user throughput of the Gemma 4 AR model (with and without MTP) and DiffusionGemma. In the low batch size regime, DiffusionGemma offers substantially higher TPS per user and higher total throughput. It is only at moderate batch sizes (around 32 concurrent requests) that AR models begin to have a throughput advantage. All models are run on an H100 with FP8 precision using the PG-19 benchmark (4096 input tokens, 1024 output tokens); the Gemma 4 AR (MTP) model uses a draft length of 4.
Figure 12: Trade-off between total and per-user throughput of the Gemma 4 AR model (with and without MTP) and DiffusionGemma. In the low batch size regime, DiffusionGemma offers substantially higher TPS per user and higher total throughput. It is only at moderate batch sizes (around 32 concurrent requests) that AR models begin to have a throughput advantage. All models are run on an H100 with FP8 precision using the PG-19 benchmark (4096 input tokens, 1024 output tokens); the Gemma 4 AR (MTP) model uses a draft length of 4.

研究结果

  • 在完整评测集上平均,DiffusionGemma每次前向传播生成约20个token,在单张NVIDIA H100 GPU上达到约每秒1500个输出token(TPS)。
  • 这比使用最先进推测解码的AR模型(通常每次前向传播仅生成约3到6个token)快得多。
  • 启用自适应停止后,模型在最大48步预算中平均只需约12个有效去噪步,延迟降低约4倍且不牺牲生成质量。
  • 尽管每步处理256个token,但每步延迟相比单token的AR生成仅增加3.2倍(在H100、FP8精度、批大小为1的条件下测得)。
  • 微调后的权重仍可用于标准AR生成,性能仅有轻微下降。

可应用场景

  • 为对响应速度敏感的聊天机器人或代码助手提供低延迟服务
  • 在并发请求量较低(小批大小)、GPU算力容易被浪费的推理场景中使用
  • 借助Apache 2.0的宽松许可,针对特定领域(如医疗问答、语音识别)进行轻量级微调
  • 根据延迟约束和任务复杂度,在扩散解码与AR解码之间动态路由请求的混合服务方案

局限与待验证事项

  • 当并发请求量达到中等规模(约32个并发请求)时,AR模型在吞吐量上会重新占据优势。
  • 扩散微调相比原始AR基线会带来一定的性能下降。
  • 对Mercury 2等闭源竞品的速度评估是通过OpenRouter API间接估算得出,而非直接测量,存在一定误差。
  • 论文正文中间大段内容及部分附录(如附录F、G.3)在本次材料中被省略,未能查看全部细节结果。
  • 由于离散扩散在更新每个位置时无法看到其他位置的同步决定,偶尔会出现局部语法不一致的问题。

为什么重要

对于聊天机器人、代码助手等对响应速度敏感的服务来说,这意味着单张GPU就有可能实现比当前AR服务方案快得多的文本生成。该模型以Apache 2.0协议开放权重发布,研究者和开发者可以直接检视离散扩散的运作机制,并以较低算力成本针对自己的场景进行微调。

本文术语

  • 离散扩散(discrete diffusion) · 直接在词表中的实际token状态之间添加和去除噪声的生成方法,不同于图像常用的连续空间扩散
  • 自回归(AR)模型 · 从左到右逐个生成token的传统语言模型方式
  • 混合专家(MoE) · 每次输入只激活部分专门子网络以降低计算量的模型架构
  • 自适应停止(adaptive stopping) · 当模型预测足够自信且稳定时提前结束去噪过程,而不是每次都跑满最大步数
  • 采样器蒸馏(sampler distillation) · 训练模型用更少的去噪步数就能复现高质量但缓慢的生成结果的技术

论文原文摘要（英文）

We introduce DiffusionGemma, an experimental open-weight language model that uses discrete diffusion to generate text at exceptionally high speed. Rather than decoding one token at a time, DiffusionGemma iteratively refines blocks of 256 tokens in parallel, avoiding the sequential decoding bottleneck of conventional autoregressive (AR) large language models. Instead of training from scratch, we obtain DiffusionGemma by fine-tuning the mixture-of-experts Gemma 4 model with 3.8B activated and 25.2B total parameters. Our compute-efficient two-stage training pipeline uses fewer than 10% of the starting AR model's total training token budget. The first stage uses supervised fine-tuning to teach bidirectional denoising, while the second stage combines reinforcement learning with sampler distillation to jointly improve generation quality and inference efficiency. DiffusionGemma establishes a new Pareto frontier for the trade-off between generation speed and model capability. Averaged across our full evaluation suite, it generates around 20 tokens per forward pass and achieves roughly 1,500 output tokens per second on a single NVIDIA H100 GPU, which is substantially faster than AR models even with state-of-the-art speculative decoding. DiffusionGemma also retains the starting model's support for thinking mode, multimodal inputs, and long contexts. Despite diffusion fine-tuning, it remains capable of AR generation with only minor performance degradation, suggesting a path toward hybrid diffusion-AR decoding.

作者 · DiffusionGemma Team

在 arXiv 阅读

最新论文

全部论文 →

METAL LAB 最新报道

图片来源： DiffusionGemma Team et al., arXiv:2608.00146, CC BY 4.0