DiffusionGemma Technical Report — part 3/6

Cleaned-text code points [30937, 46420). Bibliography omitted here; source.md includes it.

By successfully bypassing redundant refinement steps, this mechanism allows the model to dynamically scale its inference-time compute to the complexity of the prompt. As illustrated in Figure 5 , instead of exhausting the maximum N = 48 N=48 budget, the model averages approximately 12 12 effective denoising steps (defined in Section 3.4 ) across the downstream evaluations shown in the figure. This yields a 4 × 4\times reduction in overall latency without sacrificing generation quality. Furthermore, the boxplots reveal distinct convergence profiles across domains: the model tends to deploy fewer steps for structured tasks like code and more for natural language. Moreover, it is not only the domain, but also the task complexity that dictates this behavior; for instance, harder code problems (e.g., LiveCodeBench) naturally require more steps to converge than easier ones (e.g., HumanEval).

### 3.4 Inference Efficiency Metrics

Because adaptive stopping turns the number of denoising steps into a random variable, we introduce a set of metrics to quantify denoising inference efficiency. Let K K be the total number of canvases generated, N k ≤ N N_{k}\leq N be the number of denoising steps executed for the k k -th canvas, and C k ≤ C C_{k}\leq C be the number of valid tokens produced in that canvas (i.e., all C C tokens, or up to the first end-of-sequence token if one is generated). First, we define the Total Tokens as the sum of all valid tokens generated across the entire sequence: Total Tokens ≜ ∑ k = 1 K C k . \text{Total Tokens}\triangleq\sum_{k=1}^{K}C_{k}. (7) We define the Total Denoising Steps as the absolute number of denoising steps executed across the entire generation: Total Denoising Steps ≜ ∑ k = 1 K N k . \text{Total Denoising Steps}\triangleq\sum_{k=1}^{K}N_{k}. (8) Since each denoising step incurs a fixed computational cost, the Total Denoising Steps is the primary driver of end-to-end generation latency. However, because this absolute metric inherently scales with the total sequence length, we also evaluate a normalized measure of model efficiency. We define the Effective Denoising Steps as the token-weighted average of steps across all canvases: Effective Denoising Steps ≜ 1 Total Tokens ​ ∑ k = 1 K N k ​ C k . \text{Effective Denoising Steps}\triangleq\frac{1}{\text{Total Tokens}}\sum_{k=1}^{K}N_{k}C_{k}. (9) We weight by the number of valid tokens to prevent the metric from being downwardly biased by the final canvas, which tends to require fewer denoising steps when it is only partially filled. Note that without adaptive stopping, the Effective Denoising Steps metric always reduces to the denoising budget N N . Finally, we calculate the Tokens Per Forward (TPF) : TPF ≜ Total Tokens Total Denoising Steps + K − 1 , \text{TPF}\triangleq\frac{\text{Total Tokens}}{\text{Total Denoising Steps}+K-1}, (10) where the K − 1 K-1 term accounts for the single additional forward pass required between canvases to encode the newly generated clean tokens and append their key-value pairs to the KV cache.

### 3.5 Retained Autoregressive Capability

Because DiffusionGemma shares the exact same transformer architecture as Gemma 4, the final DiffusionGemma weights can be seamlessly loaded back into the original architecture to perform standard AR generation using causal attention, exactly as the base model does. As demonstrated in Section 7 , Table 3 , the model maintains robust capabilities in this setting; its performance scores in AR mode land squarely between those of DiffusionGemma’s primary text diffusion mode and the baseline Gemma 4 checkpoint from which DiffusionGemma was initialized.

## 4 Supervised Finetuning

We start from the publicly released Gemma 4 26B A4B ( Gemma Team et al., 2026 ) checkpoint and run an extended finetuning phase where the model adapts to predicting blocks of 256 tokens from noisy inputs. We 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 (uncorrupted) tokens via the encoder KV cache. We use discrete multinomial diffusion as our corruption process and uniformly sample noisy tokens from the vocabulary ( Hoogeboom et al., 2021 ; Austin et al., 2021 ) . For a given canvas, we sample a noise level t ∼ Unif ⁡ [ 0 , 1 ] t\sim\operatorname{Unif}[0,1] and noise each token in the canvas with probability t t . Given a clean context of prompt and previous canvases H H (encoded through the KV cache), a self-conditioning signal z t z_{t} , and a noisy canvas x t x_{t} , the model is trained to minimize the cross-entropy loss between its predictions and the ground-truth canvas extracted from the training data ( Gat et al., 2024 ) : L ( θ ) = − ∑ i = 1 C log p θ ( x 0 i ∣ x t , z t , H ) , L(\theta)=-\sum_{i=1}^{C}\log p_{\theta}(x_{0}^{i}\mid x_{t},z_{t},H), (11) where superscript i i denotes indexing along the canvas dimension; p θ p_{\theta} is parameterized via a softmax transformation over the neural network’s output logits. As shown in Figures 6 and 7 , 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—a common challenge for internalizing reasoning in language models ( Zelikman et al., 2024 ) .

## 5 Sampler Distillation & Reinforcement Learning

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. More critically, when operating in the few-step regime required for ultra-low latency inference, generation quality collapses. To address this, we target a dual improvement: pushing the model’s intelligence while simultaneously compressing its denoising trajectory.

Traditionally, achieving these two goals requires a decoupled, multi-stage pipeline that treats reward-driven alignment and sampler distillation as distinct phases. We bypass this with a unified online learning stage, coined sampler distillation & reinforcement learning ( SD ⋅ \cdot RL ), which optimizes both axes concurrently. Relying on a joint objective, a single gradient update drives:

topsep=0pt Reward maximization : Elevating absolute generation quality and alignment, analogous to standard RL for AR and diffusion models ( Black et al., 2024 ; Liu et al., 2025 ; Zheng et al., 2026 ; Zhao et al., 2026 ; Ma et al., 2026 ; Wallace et al., 2024 ; Fan et al., 2023 ; Clark et al., 2024 ; Dong et al., 2023 ; Lee et al., 2023 ) .

topsep=0pt Sampler distillation : Mapping this high-quality generation to the few-step regime, addressing a compression challenge unique to iterative diffusion frameworks ( Salimans and Ho, 2022 ; Song et al., 2023 ; Deschenaux and Gulcehre, 2025 ; Fu et al., 2025 ; Luo et al., 2023 ; Sauer et al., 2024 ; Liu et al., 2024 ; Yin et al., 2024 ; Hoogeboom et al., 2026 ) .

##### Training setup.

We adapt the data distribution used by the Gemma 4 RL recipe ( Gemma Team et al., 2026 ) for our use case. Encompassing both thinking and non-thinking modes, our setup targets improvements across a variety of capabilities such as helpfulness, mathematical reasoning, coding and instruction-following. Initialized from the SFT weights, the model acts as an online teacher that generates denoising trajectories (using a sampler configured with high maximum denoising steps and mild temperature annealing) to establish a high-quality reference. The SD ⋅ \cdot RL joint objective uses these trajectories to simultaneously maximize reward and drive sampler distillation in order to compress the model’s highest-quality output into the few-step regime.

##### Training dynamics & implicit curriculum effect.

Over the course of SD ⋅ \cdot RL training, two synergistic dynamics emerge, as shown in Figure 8 . First, the online teacher’s average reward steadily increases, reflecting improved fundamental capabilities. Second, 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 ⋅ \cdot RL objective systematically reduces the predictive entropy of the model. Crucially, 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. This seamlessly shifts the training distribution toward ever shorter denoising trajectories, allowing the algorithm to dynamically pace its own sampler distillation. Consequently, and in contrast to RL for AR models, prolonging the SD ⋅ \cdot RL phase remains highly beneficial even after the reward metric plateaus; continued entropy reduction translates directly into further inference speedups.

##### Improved speed-to-intelligence Pareto frontier.

As illustrated in Figure 9 , SD ⋅ \cdot RL training significantly expands the Pareto frontier established by the SFT checkpoint. On the quality axis, it yields a 10-point improvement on the combined GPQA-Diamond and LiveCodeBench-v6 score. On the efficiency axis, it quadruples the TPF from 5 to nearly 20, unlocking ultra-low latency inference.

As shown in Figure 9 , evaluating the SFT checkpoint with the default DiffusionGemma sampler (maximum of N = 48 N=48 denoising steps) results in poor downstream accuracies and artificially low effective denoising steps. This counterintuitive behavior stems from the SFT model frequently degenerating into repetitive token loops in this restricted-step regime. Once it falls into a loop, its predictive entropy collapses, which prematurely triggers the adaptive stopping mechanism. Figures 16 and 17 (Appendix B ) provide generated samples on the GPQA-Diamond benchmark that illustrate how these degeneracies manifest: the SFT model begins with a valid, logical reasoning trace but suddenly collapses into a token repetition loop. In contrast, the samples post- SD ⋅ \cdot RL avoid these degenerative traps to sustain coherent reasoning at minimal latency.

##### SD ⋅ \cdot RL specialization in the few-step regime.

For the SFT checkpoint, downstream performance scales consistently with the maximum number of denoising steps N N up to 192, at which point denoising is roughly any-order autoregressive. Following SD ⋅ \cdot RL , the scaling behavior improves faster in N N and plateaus earlier: performance improves steadily up to N = 48 N=48 , but exhibits diminishing returns thereafter (Figure 10 ). This early saturation emerges because the SD ⋅ \cdot RL objective explicitly specializes the model for the few-step regime by aggressively minimizing predictive entropy.

##### Emergent conciseness.

An emergent property of our SD ⋅ \cdot RL optimization is that it encourages the model to produce concise, token-efficient outputs. This contrasts with RL for AR models, which often maximizes reward by inducing longer reasoning traces. Our final checkpoint produces generations nearly 2 × 2\times shorter than the SFT checkpoint (Figure 9 ). While this means the model forgoes some capability gains typically associated with extended reasoning, it acts as a multiplier on inference speed gains. Compounding fewer total tokens with fewer effective denoising steps per canvas directly drives exceptionally low end-to-end latency: DiffusionGemma requires less than 5 % 5\% of the total forward passes used by the Gemma 4 AR baseline across our eval suite.

## 6 Inference Optimizations

The preceding sections focused on maximizing the number of tokens produced per forward pass through our SFT and SD ⋅ \cdot RL pipeline. We now turn to the complementary axis: minimizing the wall-clock cost of each forward pass through targeted GPU-level optimizations. Fundamentally, the throughput advantage of text diffusion over AR decoding is dictated by two competing quantities: on one hand, each denoising step decodes TPF \mathrm{TPF} tokens, and thus requires TPF \mathrm{TPF} times fewer forward passes than an AR baseline. On the other hand, because each forward pass processes a canvas of 256 tokens, each such step is r r times slower, and so the overall throughput relative to the AR baseline becomes TPF / r \mathrm{TPF}/r . The forward pass overhead can be substantially reduced by hardware optimizations.

##### Low batch size serving.

While text diffusion inference requires more floating-point operations (FLOPs) per generated token than AR models, it relies on significantly fewer forward passes. Because LLM serving on modern hardware accelerators is typically memory-bound—largely dictated by KV cache capacity and memory bandwidth ( Kwon et al., 2023 ; Pope et al., 2023 ) —this reduction in memory transfers creates a latency advantage that outweighs the higher computational cost. As a result, text diffusion is highly effective in low-batch-size scenarios, leveraging available compute capacity to minimize per-request latency. For simplicity, our analysis focuses on the single-request inference throughput (a batch size of 1) of DiffusionGemma, comparing it directly to its AR counterpart, Gemma 4 26B A4B. Reference inference implementations are available in HuggingFace Transformers ( Google DeepMind Team, 2026 ) and vLLM ( The vLLM Team and Google DeepMind Team, 2026 ) .

##### GPU time breakdown.

Figure 11 presents the per-step GPU kernel time breakdown for both models; we focus on GPU kernel time rather than end-to-end latency to isolate model-specific bottlenecks. For each model individually, end-to-end latency exceeds GPU kernel time by approximately 1 ms due to detokenization and other CPU-side serving overhead, with this gap being similar for both the Gemma 4 AR and DiffusionGemma models. Each DiffusionGemma step processes 256 × 256\times more tokens, yet is only 3.2 × 3.2\times slower than the single-token AR step. The time difference can be mainly attributed to three operations: mixture-of-experts (MoE), sampling, and attention. Other operations (e.g., shared expert, attention output projection, etc.) are at most 2 × 2\times slower.

topsep=0pt MoE. When serving a single request, the MoE layer computation is memory-bound, with its running time dominated by the transfer of expert weights from high-bandwidth memory, a well-documented bottleneck when serving sparse MoE models ( Rajbhandari et al., 2022 ; Huang et al., 2024 ) . For the Gemma 4 AR model, only 8 unique experts are activated per token per MoE layer. For the DiffusionGemma model, on average, approximately 84 unique experts are activated per canvas of 256 tokens per MoE layer, as measured on the PG-19 benchmark ( Rae et al., 2020 ) . 4 4 4 Dataset available at https://github.com/google-deepmind/pg19 . Activating more experts per forward pass results in a 4.3 × 4.3\times slower MoE kernel. It is the same type of penalty incurred by verification in speculative decoding, but scaled to larger token parallelism. For a dense architecture, this overhead would be eliminated, reducing the per-step feed-forward network slowdown to less than 2 × 2\times .