DiffusionGemma: Parallel Denoising From an Autoregressive Checkpoint

Google DeepMind released DiffusionGemma (arXiv:2608.00146, 31 July 2026), an experimental open-weight language model that generates text by discrete diffusion rather than left-to-right decoding. The 55-page technical report describes a finetuning recipe that converts the mixture-of-experts Gemma 4 26B A4B model — 3.85 billion activated and 25.2 billion total parameters — into a block-parallel generator, with code in HuggingFace Transformers (PR #46540) and vLLM. Weights are Apache 2.0.

The paper is accompanied by a finetuning toolkit and a Sudoku case study demonstrating domain adaptation on consumer hardware.

Finding

DiffusionGemma does not generate one token at a time. It denoises a fixed canvas of 256 tokens in parallel and commits canvases block-autoregressively to build open-ended outputs.

The architecture reuses the Gemma 4 transformer with shared weights but inverts the usual encoder-decoder orientation. A causal encoder processes the prompt and previously committed canvases into a KV cache that is appended canvas by canvas, without re-encoding. A bidirectional decoder then refines the current noisy canvas by attending across the canvas and cross-attending to the cache. A 7.8 million parameter self-conditioning feedforward block feeds the previous step prediction back in, which the authors report remains interpretable as a token bottleneck.

Training is a two-stage conversion that uses fewer than 10 percent of the original autoregressive training token budget, starting from the released post-trained Gemma 4 checkpoint:

1. Supervised finetuning. The model is trained to denoise uniform (multinomial) corruption. For each canvas a noise level t is sampled uniformly in [0,1] and each token is replaced with a uniform vocabulary sample with probability t. The loss is cross-entropy to the clean canvas conditioned on the KV cache and self-conditioning state. Performance improves rapidly at first, then log-linearly. Thinking traces — the model generating reasoning before its answer — require the longer end of this phase and initially collapse into repetition.

2. Sampler distillation and reinforcement learning (SD·RL). An online phase combines reward maximization with compression of the denoising trajectory. An online teacher generates with a high step budget and mild temperature annealing; a joint objective pushes quality upward while distilling that quality into few-step sampling. The report shows about 10 points of combined improvement on GPQA-Diamond and LiveCodeBench-v6, tokens per forward pass rising from about 5 to nearly 20, and output length shrinking by roughly half as an emergent effect. Continued training past reward plateaus still reduces latency because predictive entropy keeps falling, which triggers earlier stopping.

Sampling uses multinomial diffusion rather than masked diffusion, so any token in the current canvas can be revised until the canvas is frozen. The default sampler (Algorithm 1) has three mechanisms:

• Entropy-bounded acceptance. Tokens are sorted low to high entropy. A prefix is accepted while cumulative entropy stays below b = 0.1; the remainder is renoised uniformly. This bounds mutual information in the joint update.
• Temperature annealing. A linear schedule from 0.8 to 0.4 across the denoising timescale sharpens predictions from exploratory to committed.
• Adaptive stopping. Denoising halts when mean canvas entropy falls below 0.005 and the argmax prediction is stable for two steps. The maximum budget is N = 48, but the mean effective steps across the evaluation suite is about 12. Harder tasks take more steps: structured code tasks converge faster when constrained, open-ended natural language takes longer.

On hardware the report compares on a single NVIDIA H100 at FP8, batch size 1, with 4096 input tokens and 1024 output tokens. A single DiffusionGemma denoising step processes 256 tokens and is 3.2 times slower than a single-token autoregressive step. The breakdown is instructive. The mixture-of-experts layer is 4.3 times slower because a canvas activates about 84 distinct experts per layer versus 8 for a single token — the same verification penalty seen in speculative decoding, scaled to wider parallelism. Sampling is 3.06 ms versus 0.56 ms because it runs a full-vocabulary 262k softmax and self-conditioning matmul in full precision over 256 positions. Attention is 4 times slower because bidirectional attention over 256 tokens cannot use single-token decode kernels, even with FlashAttention-4.

With tokens per forward pass (TPF) defined to include the extra forward pass for cache updates, and throughput as TPS = TPF / t_fwd, the model averages 19.74 TPF and t_fwd ≈ 13.56 ms end-to-end. That yields around 1,456 tokens per second averaged over seven benchmarks where all models have coverage, reported as roughly 1,500. Third-party measurement by Unsloth is cited at 2,000 tokens per second on an RTX 6000. For reference on the same H100 setup: Gemma 4 autoregressive at 204 TPS, and with multi-token prediction draft length 4 — the current speculative decoding baseline used in the report — at 303 TPS. That is about 7.1 times and 4.8 times respectively.

Capability comparisons in Table 3 group 19 benchmarks. Against its own starting point, diffusion mode trades some accuracy for speed: AIME 2026 69.1 versus 84.2 autoregressive (88.3 with multi-token prediction), GPQA-Diamond 73.2 versus 79.8, LiveCodeBench-v6 69.1 versus 71.4, GSM8K essentially flat at 96.3 versus 96.6. The same final weights run autoregressively recover part of the gap, which the authors present as evidence for dynamic routing or hybrid decoding by latency requirement.

Against other diffusion models the report places DiffusionGemma substantially ahead of open baselines on both quality and speed: LLaDA 2.1 Flash 100B at 4.63 TPF and 375 TPS on 8 times B200, Nemotron Diffusion 14B at 49 TPS on a single H100, both in bfloat16. Against the closed Mercury 2 API measured at roughly 600 TPS, the report claims near parity on quality at about 2.5 times the speed, though precisions and devices differ and the comparison is not hardware-matched.

What was not demonstrated in this report is equally important. Batched serving is not optimized; total throughput advantage over autoregressive with multi-token prediction holds until about 32 concurrent requests, after which autoregressive pulls ahead (Figure 12). Long-context agentic behavior at scale, rare extremes, and real traffic patterns are left for future work. The model generates concise outputs by design after SD·RL — useful for latency, but it forgoes the extra quality that longer thinking traces provide. Rare stuttering loops and a known failure to close thinking tags on some multimodal prompts (MMMU-Pro thinking 54.3 versus non-thinking 66.0) remain.

Meaning

The bottleneck DiffusionGemma addresses is not computation but memory movement at low concurrency.

Single-request autoregressive serving is memory bound. The time spent transferring weights and KV cache from high-bandwidth memory exceeds the time spent computing. Batching hides the cost by amortizing transfers across requests; speculative decoding helps by verifying a draft of typically 8 tokens at once, reaching 3 to 6 TPF. DiffusionGemma widens that parallelism to 256 and keeps adaptive acceptance conservative enough to average about 20 TPF — a step change in forward passes required.

Each forward pass is heavier because it touches more experts, runs a larger softmax, and uses bidirectional attention. That cost is fixed per pass. The net win appears when the reduction in number of passes outweighs the increase in cost per pass. At batch 1 on current accelerators the reduction dominates. At larger batches the per-pass overhead matters more because memory bandwidth is better utilized and compute per token becomes the limiting factor. That crossover around 32 concurrent users in Figure 12 is the point where the tradeoff reverses.

Two engineering choices make the conversion practical without pretraining from scratch.

First, the inverted encoder-decoder. Conventional encoder-decoder models use a bidirectional encoder and causal decoder. DiffusionGemma uses a causal encoder and bidirectional decoder. The causal side allows append-only KV cache updates, so history does not need re-encoding. The bidirectional side allows the whole canvas to revise in light of future tokens. Weight sharing means the checkpoint can still be loaded back into the original architecture and run autoregressively with only minor degradation — a property the authors suggest for latency-aware routing.

Second, distillation in the time domain combined with reinforcement learning. Rather than separate alignment and acceleration phases, SD·RL lets adaptive stopping create its own curriculum. Early in training predictive entropy is high, so stopping triggers late and the model sees long trajectories. As entropy falls, stopping triggers earlier and the training distribution shifts toward short trajectories without an external schedule. That is why extra optimization steps continue to improve latency after quality has leveled off.

The consequence of that optimization is conciseness. The final checkpoint generates roughly half as many tokens as the supervised finetuned checkpoint on the same prompts. Fewer tokens times fewer steps per canvas compounds into latency. It also caps the model on tasks where longer chains of thought help. The report is explicit about this: SD·RL targets ultra-low latency and trades away asymptotic performance that more steps and longer traces would recover.

The bidirectional property is visible in the paper examples beyond speed. An arithmetic prompt that asks for the answer before the reasoning forces the autoregressive Gemma 4 to commit to -1 before working through √121 = 11, (11-3)=8, 7×8=56, 9²=81, 56-81=-25, then patch itself with a correction. DiffusionGemma traverses -1 → -15 → -25 during denoising as reasoning tokens and answer tokens co-evolve, converging in 5 steps. For highly constrained outputs like strict JSON extraction or verbatim code edits, convergence takes 2 to 3 steps because predictable structure can be locked in across the canvas at once. Sequential decoding must pay O(N) even for boilerplate; parallel refinement does not.

Connection

The pattern here recurs in systems that moved from serial to iterative-parallel execution.

Weather forecasting offers a nearby comparison. Last cycle, GenCast showed that machine learning could outperform the best physics ensemble but relied on diffusion sampling over dense grids — powerful, iterative, slow. WeatherNext 2 moved in the opposite direction, from diffusion back to a single forward pass perturbed by a low-dimensional noise vector that acted globally through conditional normalization. Coherence emerged from how variation was introduced, not from an explicit joint objective. DiffusionGemma moves language generation in the analogous direction for a different reason — from serial decoding to a perturbed parallel pass where variation is bounded by entropy — and observes the same consequence: spatially (or sequentially) coherent outputs without having trained a joint likelihood, because the perturbation acts on shared computation.

In hardware terms the tradeoff is data movement versus compute. DiffusionGemma performs TPF times fewer KV cache transfers at the cost of proportionally more FLOPs in the sparse feedforward layers. That trade is favorable on modern GPUs where compute-to-bandwidth ratios keep rising, and increasingly favorable for agentic workloads with long contexts where KV traffic dominates. The crossover in batched throughput is the reminder that no single operating point is optimal everywhere. What is useful here is that one set of weights provides two operating points — diffusion for latency and structure-exploiting tasks, autoregressive for quality — without retraining.

If the low-step regime holds under replication and the per-step overhead falls with denser architectures and batch-aware kernels, the wider implication is not that diffusion replaces autoregressive generation. It is that inference can be offered as a choice along a latency-quality curve negotiated per request, with parallel refinement handling the time-critical and template-bound cases and sequential decoding handling the reasoning-intensive ones.

Primitive implementation. Correct vector.