DiffusionGemma Technical Report

DiffusionGemma is Google DeepMind's answer to the question "can you take a frontier autoregressive LLM and cheaply convert it into a text diffusion model without giving up the things that make it useful?" The result: a 25.2B-parameter MoE (3.8B active) that denoises 256-token blocks in ~12 forward passes, hits roughly 1,500 tokens/second on a single H100 at batch size 1, and keeps thinking mode, multimodality, long context — and, surprisingly, the ability to still decode autoregressively. The conversion cost less than 10% of the original model's training token budget. This is the strongest open evidence yet that the autoregressive decoding bottleneck is a choice, not a law.

Why single-stream decoding is slow, and why the obvious fixes stall

When an LLM serves one user, generation is memory-bound: each forward pass must stream all the model weights (and the KV cache) from HBM into the compute units to produce one token. The GPU's arithmetic units sit mostly idle. Batching many users amortizes the weight transfer, but per-user latency stays stuck at one token per weight-read.

Speculative decoding attacks this by drafting several tokens with a cheap model and verifying them in one pass of the big model. But in practice this yields 3–6 tokens per forward pass (TPF): autoregressive drafters are themselves sequential, and parallel drafters see acceptance rates collapse at later draft positions because tokens drafted independently don't cohere.

Text diffusion takes a different bet: instead of extending the sequence one token at a time, generate a whole block of tokens as a joint denoising problem. Start from noise, run a handful of refinement passes over the entire block in parallel, and commit. Now each forward pass is heavier (it processes 256 token positions), but you need dramatically fewer of them. You've traded memory bandwidth for compute — exactly the trade modern GPUs want you to make, since compute-to-bandwidth ratios keep growing.

The catch has always been quality. Open diffusion LMs so far were either weak reasoners or failed to deliver the promised speed; the strong ones (Gemini Diffusion, Mercury) were closed. DiffusionGemma's claim is that it's the first model to be simultaneously fast, capable, and open.

The mechanism: an AR model turned inside out

The architectural idea is elegant enough to sketch in one figure. DiffusionGemma is a single transformer (the Gemma 4 26B A4B MoE, weights shared) used in two attention modes, arranged as an inverted encoder-decoder:

• Causal encoder for history. The prompt and all previously committed text are encoded with ordinary causal attention into a KV cache. Because it's causal, appending a new block only requires encoding that block — no re-encoding of the growing context.
• Bidirectional decoder for generation. A 256-token "canvas" of noise is iteratively refined with full bidirectional attention within the canvas, cross-attending to the KV cache.

This flips BART/T5 (bidirectional encoder, causal decoder) on its head, and it's what makes long generations tractable: diffusion happens locally within a block, autoregression happens between blocks. Once a canvas is fully denoised, it's frozen, encoded causally, appended to the cache, and the next canvas begins.

Two less-obvious design choices matter a lot:

Multinomial noise, not masking. Most discrete diffusion LMs use an absorbing [MASK] state: unrevealed positions are masks, and once a token is unmasked it's fixed. DiffusionGemma instead corrupts tokens to uniformly random vocabulary tokens. Because every token can transition to any other, earlier commitments within a canvas remain revisable — the model can correct its own mistakes mid-generation. The paper shows this concretely: on an arithmetic problem, the model initially places high probability on the wrong answer -1 (the same mistake the AR baseline makes and has to walk back in text), then revises through -15 to the correct -25 over subsequent denoising steps, because the co-evolving reasoning tokens feed back into the answer tokens bidirectionally.

Self-conditioning. Alongside the noisy canvas x_t, the decoder receives a continuous signal z_t = \mathrm{FFW}(\hat p_0 E) — the previous step's predicted token distribution \hat p_0, pushed through the embedding matrix E and a small MLP (7.8M parameters). This gives each step a soft summary of what the model believed last step, mitigating the core weakness of factorized denoising: within a single step, position i is sampled independently of the simultaneous choices at positions j \neq i, which can otherwise cause locally inconsistent text.

The sampler ties it together. At each step, temperature-annealed logits (\tau: 0.8 → 0.4) yield per-position entropies; positions are accepted greedily from lowest to highest entropy until a cumulative entropy budget (b = 0.1) is exhausted — accept where you're sure, and re-noise everything else to uniform random to force exploration next pass. Generation halts early when mean canvas entropy drops below 0.005 and the argmax prediction is stable across two steps. This adaptive stopping means the nominal budget of 48 steps collapses to ~12 on average — a free 4× latency win — and, satisfyingly, the model spends more steps on hard problems (LiveCodeBench) than easy ones (HumanEval), a form of automatic test-time compute allocation that AR models simply don't have.

Training: SFT, then a two-birds-one-stone online phase

Stage one is plain: supervised fine-tuning from the released Gemma 4 checkpoint with the diffusion cross-entropy loss — sample a noise level t \sim \mathrm{Unif}[0,1], corrupt each canvas token with probability t, predict the clean canvas given the noisy one plus the causal KV cache. Denoising ability appears fast; thinking takes much longer, with early checkpoints collapsing into stutter loops mid-reasoning-trace. Performance then improves log-linearly with training.

The more interesting stage is what they call SD·RL — sampler distillation and reinforcement learning fused into one objective. The problem it solves: after SFT, the model is decent if you give it many denoising steps (approaching 192, which is nearly any-order autoregressive), but quality collapses in the few-step regime you actually need for speed. Conventionally you'd run RL for quality, then a separate distillation phase to compress the trajectory. Here, the model acts as its own online teacher (sampling with generous steps and mild annealing), and a single gradient update both maximizes reward and distills those high-quality trajectories into fewer steps.

The elegant emergent dynamic: the objective systematically drives down predictive entropy, and low entropy triggers adaptive stopping earlier, which shifts the training distribution toward ever-shorter trajectories — a self-pacing curriculum for compression. Unlike RL on AR models, training keeps paying off after reward plateaus, because continued entropy reduction converts directly into speed.

The results of this stage are large: +10 points on combined GPQA-Diamond/LiveCodeBench, TPF quadrupled from ~5 to ~20, and it fixes the SFT model's failure mode of degenerating into repetition loops (Appendix B shows an SFT checkpoint deriving the correct GPQA answer, then drowning in a thousand repeated ) tokens and returning nothing). One side effect worth noting: SD·RL makes outputs ~2× shorter — great for latency (the paper reports <5% of the AR baseline's total forward passes end-to-end), but it means the model forgoes the gains that longer reasoning traces usually buy.

The numbers, honestly

Here is the speed story first. A DiffusionGemma denoising step processes 256 positions but is only 3.2× slower than a single-token AR step (13.56 ms vs. ~4 ms, H100, FP8, batch 1) — the overheads being MoE (84 unique experts activated per canvas vs. 8 per token, a 4.3× kernel slowdown), bidirectional attention (4×), and canvas-wide sampling over a 262k vocabulary. With ~20 TPF, the arithmetic gives ~1,456 TPS averaged across benchmarks:

Now the quality cost, which the abstract soft-pedals. Against its own AR parent (thinking mode, and against the MTP-served variant, since that's the fair speed-optimized AR comparison):

The pattern is consistent: near-parity on instruction following and easier tasks (IFEval 97.4 vs. 98.7, GSM8K 96.3 vs. 96.7), a real but modest gap on knowledge (GPQA −9), and the largest losses exactly where long, careful reasoning matters most — AIME −19, Putnam −14, Codeforces −289 Elo. That's not a coincidence: SD·RL's emergent conciseness cuts reasoning traces roughly in half, and hard math is where extended traces pay. Against other diffusion models, though, DiffusionGemma is dominant among open weights (LLaDA 2.1 Flash scores 39.4 on LiveCodeBench at 375 TPS on 8×B200; DiffusionGemma scores 69.1 at ~1,479 TPS on one H100) and roughly matches the closed Mercury 2 at ~2.5× its speed.

Two more results deserve attention:

AR mode survives. Load the same weights into the vanilla Gemma architecture and decode causally: scores land between the diffusion mode and the original baseline (e.g., AIME 84.2, GPQA 79.8). Diffusion fine-tuning didn't destroy the causal capability — it's still sitting in the weights. This makes dynamic routing (diffusion for latency-critical requests, AR for quality-critical ones) a single-model deployment option, and hints at hybrid decoding schemes.

The concurrency crossover. At batch size 1, diffusion wins on both per-user and total throughput. AR+MTP only regains the total-throughput advantage around ~32 concurrent requests, where diffusion's higher per-token FLOPs start to bite. So this is a low-concurrency technology: on-device, single-user, agent loops, latency-critical APIs — not high-batch cloud serving. (The authors note batch>1 kernels are untuned, so the crossover may shift.)

What to make of it

The framing worth internalizing: diffusion converts memory traffic into compute. Every architectural trend — FP8, growing FLOPs-to-bandwidth ratios, MoE — makes that trade more favorable over time. And the conversion recipe is cheap enough (<10% of pretraining tokens, LoRA fine-tuning of the released model runs on 2×A100s) that this stops being an exotic research artifact and becomes something teams can actually adopt. Early third-party serving already reports ~2,000 TPS on an RTX 6000, and downstream fine-tunes (ASR, radiology drafting) appeared within weeks.

There are also genuinely new capabilities here, not just speed. Bidirectional generation lets the answer and its justification co-evolve (the AR baseline commits to a wrong answer token before writing the reasoning that refutes it; the diffusion model revises in place). Structured outputs — JSON schemas, code edits that mostly copy the input — converge in 2–3 denoising steps because the model locks in all the boilerplate simultaneously. Adaptive stopping gives per-prompt test-time compute scaling for free. A Sudoku LoRA fine-tune goes from 0% to 84% puzzle accuracy, a task class where left-to-right decoding is structurally handicapped.

Reasons for skepticism, stated plainly. The quality gap on hard reasoning is real and the authors admit multiple contributing compromises: no diffusion pretraining, a compute-constrained SFT phase, an SD·RL objective that explicitly trades asymptotic quality for latency, and enforced brevity that caps reasoning depth. The headline TPS is batch-size-1 on well-covered benchmarks; the Mercury 2 comparison is a black-box regression against an API with acknowledged outlier contamination, and Inception's own numbers (~1,000 TPS on Blackwell) muddy the cross-hardware picture. The MoE architecture is actually unfavorable here (84 experts per canvas vs. 8 per token); a dense model would see less per-step slowdown, which cuts both ways — it means the speed numbers understate diffusion's potential on dense architectures, but also that this particular model pays an avoidable tax. Occasional token stuttering and a known multimodal thinking-tag bug (MMMU-Pro thinking scores below non-thinking, 54.3 vs. 66.0) remind you this is labeled experimental. And "Pareto frontier" claims are averaged over just two benchmarks in the headline figure.

Still, the honest summary is that the Pareto claim survives scrutiny: nothing else, open or closed, delivers this quality at this latency, and nothing at this latency comes close on quality. If the reasoning gap closes — via diffusion-aware pretraining, longer SFT, or hybrid AR/diffusion routing exploiting the retained dual mode — the assumption that serious LLMs must decode left-to-right, one token at a time, starts looking like a historical accident.

Where to spend your reading time: Section 5 (SD·RL) is the most novel contribution — the entropy-driven self-pacing curriculum is a genuinely new training dynamic. Section 6 plus Figure 11 gives you the hardware accounting needed to sanity-check the speed claims yourself. Skim Table 3 for the unvarnished quality picture, and Appendix B for a visceral look at what SD·RL actually fixes.