One-sentence Summary

Researchers at Google DeepMind introduce DifusionGemma, an open-weight discrete-diffusion language model fine-tuned from the mixture-of-experts Gemma 4 with 3.8B activated and 25.2B total parameters that iteratively refines 256-token blocks in parallel, achieving roughly 1,500 tokens per second on an NVIDIA H100 and establishing a new speed-capability Pareto frontier while retaining autoregressive, multimodal, and long-context abilities.

Key Contributions

• DiffusionGemma is an experimental open-weight discrete diffusion language model obtained by fine-tuning the Gemma 4 mixture-of-experts checkpoint (3.8B activated, 25.2B total parameters) to refine 256-token blocks in parallel with bidirectional attention instead of decoding one token at a time.
• The two-stage training pipeline uses supervised fine-tuning for bidirectional denoising, followed by reinforcement learning combined with sampler distillation, and consumes less than 10% of the original autoregressive model’s training token budget.
• DiffusionGemma establishes a new speed-versus-capability Pareto frontier, reaching about 20 tokens per forward pass and roughly 1,500 output tokens per second on a single NVIDIA H100 GPU. It retains thinking mode, multimodal inputs, long contexts, and autoregressive generation with minor degradation, while reasoning examples show bidirectional self-correction during denoising.

Introduction

Autoregressive language models dominate serving but are memory-bound at low concurrency, because token-by-token generation spends more time transferring weights and KV cache than computing. Speculative decoding improves utilization but remains constrained by sequential drafting or declining acceptance rates in parallel drafters. Text diffusion models can predict blocks of tokens simultaneously, but existing options force tradeoffs among speed, intelligence, and accessibility: proprietary APIs are closed, while open models often underdeliver on latency or reasoning. The authors introduce DiffusionGemma, an open-weight diffusion variant of the Gemma 4 26B A4B model. It is warm-started from a pretrained autoregressive checkpoint, generates 256-token blocks in roughly 12 forward passes, achieves about 20 tokens per forward pass, reaches around 1,500 tokens per second on an H100, and establishes a new speed-to-intelligence Pareto frontier through supervised fine-tuning followed by sampler distillation and reinforcement learning.

Dataset

Dataset sources and composition

• The authors use open-source, domain-specific datasets for downstream supervised finetuning.
• The section explicitly references an open-source Sudoku puzzle dataset and mentions PubMedQA for additional results.
• The Sudoku evaluation uses a held-out set of 4096 puzzles. Training set size, filtering rules, and dataset mixture ratios are not specified in this section.

Example and sequence structure

• Each training sample is represented as a sequence containing:
  • P prompt tokens
  • K canvases
  • C tokens per canvas
• This gives a total sequence length of P + K*C.
• The text does not describe cropping or metadata construction beyond this prompt/canvas token layout.

How the data is used

• The authors finetune DiffusionGemma on the downstream data using an open-source SFT toolkit.
• The encoder processes all tokens in a sequence and is trained with a standard next-token cross-entropy loss.
• The decoder loss is computed on one uniformly sampled canvas per sequence, using the prompt and any earlier canvases as context.
• For 50% of batch examples, the decoder also uses a self-conditioning state from a previous forward pass; for the other 50%, this state is set to zero.
• Full finetuning on the Sudoku dataset achieves over 85% puzzle-level accuracy on the 4096-puzzle held-out set.
• LoRA finetuning is also reported, with accuracy and compute trade-offs controlled by the LoRA rank. Additional PubMedQA results are referenced but not detailed in this section.

Method

The authors adopt a continuous-time Markov chain (CTMC) approach for discrete diffusion. The forward process dictates the transition from clean text tokens to uniformly distributed tokens. Conditioned on a fixed clean starting canvas $x_0$x0​, the forward transition probability path factorizes independently over each token coordinate $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​]

where $\kappa_t$κt​ is a smoothly varying, monotonically decreasing noise schedule from $\kappa_0 = 1$κ0​=1 to $\kappa_1 = 0$κ1​=0, and $\delta$δ represents the Kronecker delta. As time $t$t progresses toward 1, each token is increasingly likely to be replaced by a token sampled uniformly at random from the vocabulary.

To reconstruct data from noise, the model learns a reverse process to undo this categorical corruption. Discrete flow matching theory shows that to perfectly reverse the forward trajectory, the model needs to infer the conditional distribution of the original uncorrupted tokens given a corrupted state. For a given realization $X_t = x_t$Xt​=xt​, stepping backward in time by a small increment $\Delta t$Δt is governed by a transition mapping, denoted Step, which outputs the probability distribution for the next intermediate 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​))

The authors approximate the true posterior distribution of clean tokens given a noisy state, $\mathbb{P}(X_0^i = \nu \mid X_t = x_t)$P(X0i​=ν∣Xt​=xt​), with a neural network $p_\theta(\nu \mid x_t)$pθ​(ν∣xt​). During generation, they use this approximation to sample the next token state.

As shown in the figure above, the model smoothly shifts probability mass away from a uniform noise distribution toward valid data modes as time runs backward. Individual sequence coordinates undergo continuous-time jump transitions, demonstrating how parallel dimensions coordinate over time without requiring sequential left-to-right generation.

The DifusionGemma architecture functions as an encoder-decoder transformer with shared weights. Rather than pretraining a diffusion model from scratch, the authors initialize their model with the publicly released Gemma 4 26B A4B MoE checkpoint. This initialization allows them to inherit the base model's advanced features, such as its extended context window and native multimodal understanding.

To generate open-ended text, the authors employ a block-autoregressive generation strategy. The model denoises a canvas of 256 tokens at a time. Once a canvas is fully denoised, it is committed to the sequence history, and the model begins denoising the next canvas.

As illustrated in the figure above, the generation pipeline consists of three main stages. First, the input prompt is processed by the causal encoder to initialize the Key-Value (KV) cache. Second, 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. Third, 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.

At each denoising step $t$t, the decoder applies the transformer with shared weights to predict the probability distribution of the clean tokens. The decoder takes three inputs: the current noisy canvas $x_t$xt​, the context KV cache $H$H, and a continuous self-conditioning signal $z_t \in \mathbb{R}^{C \times d}$zt​∈RC×d that feeds the model's previous predictions back into itself. Using bidirectional attention across the canvas tokens and cross-attention to the KV cache, the decoder outputs the unnormalized 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

At each iteration, the authors compute the logits $L_t$Lt​ and evaluate the clean token probabilities $\hat{p}_0$p^​0​, update the self-conditioning signal for the next step, and sample the refined 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​).​

Here, $E \in \mathbb{R}^{V \times d}$E∈RV×d is the token embedding matrix and FFW is a standard feedforward network. The time-dependent temperature $\tau_t > 0$τt​>0 sharpens the model's predictions.

The authors employ an entropy-bounded sampler with temperature annealing and adaptive stopping. Tokens are accepted in rank order from lowest to highest entropy, ensuring that their mutual information bound remains strictly below a predefined error tolerance threshold. Once the threshold is attained, all other tokens are renoised uniformly at random. To balance the rate of convergence and linguistic diversity, token probabilities are artificially sharpened via tempering. A temperature $\tau_t < 1$τt​<1 is annealed linearly from an initial value of $\tau_{\max} = 0.8$τmax​=0.8 down to $\tau_{\min} = 0.4$τmin​=0.4 across the fractional denoising timescale.

As shown in the figure above, adaptive stopping enables DifusionGemma to dynamically adjust its number of denoising steps to task complexity and domain. The sampler dynamically halts the denoising process based on the model's step-wise uncertainty. This early termination is triggered when the mean predictive entropy across the entire canvas falls below a predefined threshold and the deterministic sequence predictions from two consecutive denoising steps are identical.

The authors start from the publicly released Gemma 4 26B A4B checkpoint and run an extended finetuning phase where the model adapts to predicting blocks of 256 tokens from noisy inputs. They use a block-diagonal attention mask, enabling bidirectional attention within each block without allowing the model to condition on other denoising blocks. For a given canvas, the model conditions on the prompt and previous tokens via the encoder KV cache. They use discrete multinomial diffusion as their corruption process and uniformly sample noisy tokens from the vocabulary. Given a clean context, a self-conditioning signal $z_t$zt​, and a noisy canvas $x_t$xt​, the model is trained to minimize the cross-entropy loss between its predictions and the ground-truth 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)

As shown in the figures above, denoising performance improves rapidly within the initial steps of training, after which it settles into a log-linear performance improvement trend. Thinking performance benefits from extended SFT, as the model initially struggles with maintaining coherent reasoning traces, often collapsing into stuttering or cycles.

Following the SFT stage, the model achieves strong generation quality when using a high number of denoising steps. However, its performance on advanced reasoning and coding tasks is somewhat poorer than the baseline AR model, and generation quality collapses in the few-step regime required for ultralow latency inference. To address this, the authors introduce a unified online learning stage, coined sampler distillation and reinforcement learning (SD·RL), which optimizes both axes concurrently. Relying on a joint objective, a single gradient update drives reward maximization to elevate absolute generation quality and alignment, and sampler distillation to map this high-quality generation to the few-step regime.

The model acts as an online teacher that generates denoising trajectories to establish a high-quality reference. The SD·RL joint objective uses these trajectories to simultaneously maximize reward and drive sampler distillation. Over the course of SD·RL training, the online teacher's average reward steadily increases, reflecting improved fundamental capabilities. Facilitated by the adaptive stopping mechanism, the online teacher progressively requires fewer effective denoising steps to achieve these high rewards. This acceleration occurs because the SD·RL objective systematically reduces the predictive entropy of the model. The interplay between the reward objective and adaptive stopping induces a curriculum learning effect. Early in training, high predictive entropy delays the adaptive stopping trigger. As the model's confidence improves and entropy drops, adaptive stopping triggers earlier, seamlessly shifting the training distribution toward ever shorter denoising trajectories.

Experiment

The evaluation covers DiffusionGemma in text diffusion and autoregressive modes, with and without thinking, across reasoning, coding, knowledge, multimodal, instruction-following, and agentic benchmarks. The SD·RL stage shows that jointly optimizing reward maximization and sampler distillation compresses high-quality generation into fewer denoising steps, while the inference analysis finds that the model is most advantageous in low-batch serving by trading memory bandwidth for compute. Overall results position DiffusionGemma ahead of open text diffusion baselines and competitive with Mercury 2, with text diffusion trading some accuracy for substantially higher throughput relative to the autoregressive baseline. Constrained-output experiments confirm rapid convergence on structured tasks, and the reported limitations highlight residual issues such as occasional token stuttering and lower throughput at large batch sizes.

DifusionGemma is a mixture-of-experts transformer with a vocabulary of 262k tokens. Its total parameter count is roughly 25 billion, but only a small fraction is activated per token, and the activated count excludes the vision encoder. The embedder accounts for a large share of non-expert parameters, while the self-conditioning block adds comparatively little overhead. The sparse expert design keeps activated parameters far below total parameters by using a small number of active experts among many total experts. The embedder is larger than the vision encoder and much larger than the self-conditioning block, which adds only a minor parameter cost.

The recommended diffusion sampler uses a maximum of 48 denoising steps but averages roughly 12 steps through adaptive stopping, with separate entropy thresholds for adaptive stopping and token selection and a linear temperature schedule from 0.8 to 0.4. Highly constrained outputs can converge in two to three steps, reducing latency relative to sequential decoding. The low-latency configuration is associated with concise outputs and occasional token stuttering. Adaptive stopping lowers typical denoising from a maximum of 48 steps to an average of about 12 steps. Constrained outputs can converge in two to three denoising steps, drastically cutting latency compared to sequential decoding.

The comparison covers token-level speed and efficiency among text diffusion and autoregressive models across several reasoning, coding, and knowledge benchmarks. Speed measurements use different hardware or API conditions, so throughput values reflect mixed serving setups rather than a single controlled environment. The cited text indicates that DifusionGemma retains autoregressive generation, with AR-mode performance falling between its text diffusion mode and the baseline Gemma checkpoint. DifusionGemma weights can be loaded back into the original architecture for autoregressive sampling, and their AR-mode performance sits between the diffusion mode and the Gemma 4 initialization checkpoint. Speed conditions vary by model: DifusionGemma and Gemma 4 are measured on one H100 with FP8, Nemotron 14B on one H100 with bfloat16, LLaDA 2.1 Flash 100B on eight B200s, and Mercury 2 through its public API. Tokens per second and tokens per forward are averaged over seven benchmarks with full coverage, including AIME 2026, GPQA Diamond, LiveCodeBench-v6, MGSM, HumanEval, LBPP, and Natural2Code. Natural2Code and HiddenMath are specifically highlighted as proprietary, unleaked evaluation sets.

In thinking mode, DiffusionGemma TD improves scores over no-thinking mode on all reported benchmarks, with the largest gains on complex coding and reasoning tasks such as Codeforces ELO. This improvement comes with much higher end-to-end latency, driven by substantially more total tokens and total forwards. Tokens per forward and tokens per second remain broadly comparable across modes, while effective denoising steps vary by task. Thinking mode achieves higher benchmark scores on every reported task, with Codeforces ELO showing the largest gain and GSM8K nearly unchanged. No-thinking mode produces far fewer tokens and total forwards, resulting in much shorter end-to-end latency per sample. Throughput metrics such as tokens per forward and tokens per second are similar between modes, so longer thinking-mode generations drive the latency gap.

Before finetuning, DiffusionGemma did not produce any correct Sudoku grids and used substantially more denoising steps. After LoRA rank 8 finetuning, accuracy rose above 80 percent while the number of denoising steps decreased sharply, consistent with lower predictive entropy. Base DiffusionGemma achieved zero Sudoku accuracy. LoRA finetuning improved Sudoku accuracy to above eighty percent. The finetuned model required roughly a quarter of the denoising steps used by the base model.

DiffusionGemma is a sparse mixture-of-experts text diffusion model whose adaptive sampler typically uses about 12 of 48 maximum denoising steps, with highly constrained outputs converging even faster. Benchmark comparisons across reasoning, coding, and knowledge tasks show that its autoregressive mode sits between diffusion-mode performance and the original Gemma checkpoint, while throughput varies with serving setup. Thinking mode improves scores on all reported benchmarks, especially complex coding and reasoning tasks, but increases end-to-end latency through longer generations rather than per-token slowdown. A LoRA rank-8 finetuning experiment on Sudoku raises accuracy from zero to above 80 percent and sharply reduces required denoising steps.