DiffusionGemma Technical Report — part 4/6

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

topsep=0pt Sampling. Beyond AR’s single-token softmax-and-sample, text diffusion sampling requires additional operations, most notably a self-conditioning embedding matmul and softmax over the full canvas of 256 tokens (see Algorithm 1 ). These operations are performed in full precision with a vocabulary dimension of 262k. Consequently, text diffusion sampling takes 3.06ms while AR sampling takes only 0.56ms. We implement sampling using standard PyTorch primitives optimized with torch.compile , rather than hand-written GPU kernels, to facilitate easier extensibility by the community.

topsep=0pt Attention. Unlike AR, DiffusionGemma uses bidirectional attention over a canvas of 256 tokens. As such, we cannot utilize the fast single-token decoding attention available to AR models, but we can take advantage of the highly-optimized FlashAttention-4 kernel ( Dao et al., 2022 ; Zadouri et al., 2026 ) . Under these optimizations, the attention operation is 4 × 4\times slower for the DiffusionGemma model than the Gemma 4 AR model.

##### Eliminating CPU-GPU synchronization.

There is additional complexity when serving DiffusionGemma: denoising steps depend on adaptive stopping and there are two types of forward passes (the denoising step and the KV cache update) that require different attention masks and that can occur within the same batch. To minimize request latency, it is important that this complexity is handled solely by GPU operations without triggering any additional CPU-GPU synchronization. This is achieved by extending asynchronous scheduling to the text diffusion models and introducing a per-sequence causal attention flag (see The vLLM Team and Google DeepMind Team, 2026 for more details).

##### Throughput for batch size of 1 serving.

The decoding throughput for text diffusion models is calculated by the following formula: TPS ≜ TPF t fwd , \text{TPS}\triangleq\dfrac{\text{TPF}}{t_{\text{fwd}}}, (12) where TPS is Tokens Per Second, TPF is Tokens Per Forward (see Equation 10 ), and t fwd t_{\text{fwd}} is the time needed for a single denoising step, which varies with context length (due to attention). On an H100 GPU (FP8 precision), a single denoising step of DiffusionGemma takes t fwd = 13.56 ​ ms t_{\text{fwd}}=13.56\text{ms} on average (end-to-end; the per-step GPU time is 12.63ms, see Figure 11 ) when serving a single request with 4096 input tokens and 1024 output tokens. TPF varies depending on the task. Assuming a TPF of 19.74 (the average TPF across the 7 benchmarks reported in Table 3 ), the average decoding throughput of the model is 1456 TPS—a 7.1 × 7.1\times improvement over the Gemma 4 AR model (204 TPS) and a 4.8 × 4.8\times improvement over the AR model with MTP (303 TPS), on the same device setup.

##### Multi-user throughput.

Figure 12 shows the trade-off between total and per-user throughput of the Gemma 4 AR model (with MTP) and DiffusionGemma depending on the number of concurrent users (i.e., batch size). Importantly, these results are without targeted optimization for batch sizes larger than 1: current kernel selection is suboptimal, and sampling has not been tuned to scale with batch size; for example, applying top- k k truncation to the sampling step is expected to yield significant throughput gains at higher batch sizes with negligible impact on output quality. Despite this, DiffusionGemma offers substantially higher TPS per user and higher total throughput than the Gemma 4 AR (MTP) model in the low batch size regime, with AR models beginning to gain a throughput advantage only at moderate batch sizes (around 32 concurrent requests).

##### Toward real traffic throughput.

Compared to the AR equivalent, DiffusionGemma changes the compute characteristics of two transformer layers: for the attention layer, it performs TPF × \times fewer transfers of KV cache; for the MoE layer, it performs proportionally more FLOPs (scaling with the number of effective denoising steps). This has the potential to address one of the challenges in modern LLM serving: memory-bound attention limiting compute utilization in FFW/MoE layers ( Zhu et al., 2025a ; Tang et al., 2024 ; Zhu et al., 2025b ), which is particularly relevant for agentic workflows with long context. DiffusionGemma effectively trades data movement for compute, which is favorable on modern GPU hardware where compute-to-bandwidth ratios continue to grow; a thorough empirical analysis under realistic traffic conditions is beyond the scope of this work.

## 7 Experimental Results

We evaluate DiffusionGemma’s capabilities and inference efficiency across four operational modes: text diffusion (TD) versus autoregressive (AR) generation, each evaluated with and without thinking enabled; unless otherwise specified, TD with thinking is enabled. We benchmark the model against its AR initialization (Gemma 4 26B A4B), contemporary open-weight text diffusion models (LLaDA 2.1 Flash 100B and Nemotron Diffusion 14B), and the proprietary Mercury 2 API.

We consider a diverse suite of benchmarks spanning core domains such as mathematical reasoning, code generation, general knowledge, multimodal understanding, and instruction following alongside agentic capabilities. Specifically, for mathematical reasoning, we utilize AIME ( Dekoninck et al., 2026 ) , GSM8K ( Cobbe et al., 2021 ) , MGSM ( Shi et al., 2023 ) , Putnam ( Tsoukalas et al., 2024 ) , and HiddenMath (internal). Coding performance is measured against LiveCodeBench-v6 ( Jain et al., 2025 ) , Codeforces ( Quan et al., 2025 ) , HumanEval ( Chen et al., 2021 ) , BigCodeBench ( Zhuo et al., 2025 ) , LBPP (v2) ( Matton et al., 2024 ) , and Natural2Code (internal). Broad and expert-level general knowledge is measured against GPQA-Diamond ( Rein et al., 2024 ) , BIG-Bench ( Suzgun et al., 2023 ) , MMMLU ( OpenAI, 2024 ) , and MMLU-Pro ( Wang et al., 2024 ) , while multimodal reasoning is evaluated on MMMU-Pro ( Yue et al., 2025 ) . Finally, strict instruction following is assessed via IFEval ( Zhou et al., 2023 ) , and agentic task completion is evaluated through the Tau-bench suite, encompassing the Retail, Airline, and Telecom environments ( Yao et al., 2024 ) .

The complete per-benchmark results for all models and inference configurations are reported in Table 3 . Table 4 provides a complementary analysis of DiffusionGemma’s decoding efficiency, reporting Tokens Per Forward (TPF; Equation 10 ), Tokens Per Second (TPS; Equation 12 ), effective denoising steps (Equation 9 ), total generated tokens (Equation 7 ), and end-to-end generation latency, excluding prefill time. For a higher-level comparison, Figure 13 groups the benchmarks into three capability areas—reasoning and knowledge, coding, and instruction following and agentic behaviour—and reports the unweighted mean of the 0–100 benchmark scores within each area, together with output throughput. A model is shown for a capability area only if it completed every constituent benchmark; missing bars therefore indicate incomplete benchmark coverage rather than a score of zero. Hatched bars denote the no-think variants.

DiffusionGemma sets a new performance frontier for text diffusion models, substantially outperforming existing open-weight diffusion baselines while increasing TPF by approximately an order of magnitude. In terms of quality it is highly competitive with Mercury 2, a closed-weight text-diffusion model, while reaching roughly 1,500 output tokens per second on a single previous-generation H100 GPU, a roughly 2.5 × 2.5\times speedup over Mercury 2.

Relative to the AR model used for initialization, DiffusionGemma in text-diffusion (TD) mode trades some absolute benchmark performance for substantially greater decoding speed. Although the conversion reduces performance across the three capability areas, TD mode delivers nearly 5 × 5\times the output throughput of the original Gemma 4 AR baseline under heavily optimized MTP serving: 1,479 tokens per second compared with 303 tokens per second. The two-stage training pipeline nevertheless retains support for autoregressive decoding. When run in standard left-to-right AR mode, DiffusionGemma recovers part of the performance gap observed in TD mode and narrows the capability gap to the original baseline, albeit at lower throughput. This dual-mode capability could enable requests to be routed dynamically according to latency requirements and task complexity.

* Benchmarks are grouped into: reasoning and knowledge ( AIME 2026 , GPQA Diamond , BigBench EH , GSM8K , MGSM , MMMLU , Putnam , MMLU-Pro , and HiddenMath ); coding ( LiveCodeBench V6 , HumanEval , BigCodeBench , LBPP , and Natural2Code ); and instruction following and agentic behaviour ( IFEval , Tau2 Retail , Tau2 Airline , and Tau2 Telecom ).

## 8 Open-Source Downstream SFT

Alongside DiffusionGemma, we release an open-source finetuning toolkit that allows practitioners to adapt the model to their own domain-specific datasets. We build on top of Hackable Diffusion ( Crepy et al., 2026 ) , a modular open-sourced research toolbox for generative modeling. We provide Low-Rank Adaptation (LoRA) recipes ( Hu et al., 2022 ) to enable finetuning on consumer hardware.

##### Finetuning procedure.

Our open-source SFT toolkit includes both a causal encoder and a diffusion decoder objective. Training sequences are of length P + K ​ C P+KC where P P is the number of prompt tokens, K K is the number of canvases and C C is the canvas size. To compute the total loss, the encoder first processes all tokens in the entire sequence, populates a KV cache H H , and provides next token predictions that are fed into a standard cross-entropy loss for the encoder. The decoder loss is computed by uniformly sampling a canvas k k from the K K canvases available. For canvas k k , we evaluate a denoising cross-entropy loss using the decoder’s predicted logits given the current noisy state, x t x_{t} , the KV cache H H for the prompt and any prior canvases in the sequence (i.e., canvas 1 1 , 2 2 , …, k − 1 k-1 ); for 50 % 50\% of the data-points in a batch, the decoder is also conditioned on a self-conditioning state z t z_{t} computed from a previous forward pass. The other 50 % 50\% have z t = 0 z_{t}=\textbf{0} . Formally, we write the losses as L encoder ( θ ) = − 1 P + K ​ C ∑ j = 1 P + K ​ C log p θ ( x j ∣ x 1 : j − 1 ) , L decoder ( θ ) = − 1 C ∑ i = 1 C log p θ ( x 0 i ∣ x t , z t , H ) . L_{\text{encoder}}(\theta)=-\frac{1}{P+KC}\sum_{j=1}^{P+KC}\log p_{\theta}(x^{j}\mid x^{1:j-1}),\quad\quad L_{\text{decoder}}(\theta)=-\frac{1}{C}\sum_{i=1}^{C}\log p_{\theta}(x_{0}^{i}\mid x_{t},z_{t},H). (13) For the encoder loss, j j indexes along all tokens in the sequence (prompt and canvases); for the decoder loss, i i indexes along tokens in the given canvas. The final loss is the sum of the two losses.

##### LoRA for parameter-efficient finetuning.

LoRA is applied to all linear operations (attention projections, MLP gates, MoE routers, and the self-conditioning feedforward block). This allows us to achieve strong downstream performance while training only a small fraction of the model’s parameters, using 2 × \times A100 80GB GPUs. All of our training details are reported in Appendix C .

##### Case study: Sudoku puzzle solving.

Solving Sudoku puzzles is a compelling testbed for discrete diffusion due to the non-autoregressive nature of the task. We finetune DiffusionGemma on an open-source dataset of Sudoku puzzles. 5 5 5 https://www.kaggle.com/datasets/rohanrao/sudoku With full finetuning, the sampler (Algorithm 1 ) achieves >85% puzzle-level accuracy evaluated on a held out set of 4096 4096 puzzles (Figure 14 ). Decreasing the LoRA rank trades-off accuracy and compute. In Table 5 , we report performance of the original model as well as the finetuned model. See Appendix C for additional results on PubMedQA ( Jin et al., 2019 ) .

## 9 Practical Advantages of Text Diffusion

In this manuscript, we have primarily emphasized low latency as the principal advantage of text diffusion over standard AR language modeling ( Ghazvininejad et al., 2019 ; Gong et al., 2023 ; Li et al., 2022 ; Dieleman et al., 2022 ) . However, the architectural paradigm of text diffusion offers several benefits that extend beyond computational efficiency. Here, we demonstrate the practical advantages through concrete examples and qualitative analysis of samples generated by the model. To isolate the intrinsic capabilities of the architecture, we disable explicit thinking modes in both DiffusionGemma and the baseline Gemma AR model. This allows us to observe and evaluate the underlying generative process.

### 9.1 Bidirectional Reasoning and Self-Correction

AR next-token prediction is inherently causal; during the generation of a given token, the model can only attend to the preceding context. It fundamentally lacks the capacity to condition upon tokens that have yet to be generated. By contrast, text diffusion operates with full bidirectional attention across the canvas. This non-causal property allows tokens at arbitrary positions across the canvas to attend simultaneously to both past and future representations—a critical capability for complex planning and reasoning tasks where early decisions depend on eventual outcomes ( Papadopoulos et al., 2024 ; Kitouni et al., 2024 ; Zhang-Li et al., 2024 ; Nagarajan et al., 2025 , e.g.,) . Consequently, locally within a canvas, future tokens can directly influence the formation of earlier tokens.

Because text diffusion employs an iterative refinement process, it possesses a built-in mechanism for self-correction. Any premature commitments generated in earlier denoising steps can be revised during subsequent denoising iterations ( Reid et al., 2022 ; Han et al., 2023 ) . While we have previously demonstrated the efficacy of bidirectional reasoning and self-correction in structured logic puzzles like Sudoku, we demonstrate these mechanics in two reasoning examples.

Consider the multi-step arithmetic problem detailed in Figure 15 . Constrained by causal generation, the AR Gemma model must commit to the first token of the final answer before articulating the intermediate calculations. This forces an incorrect initial prediction ( − 1 -1 ) that the model must later correct. DiffusionGemma, conversely, leverages its bidirectional attention canvas to simultaneously evolve the final answer and its underlying logic. As visualized in the accompanying denoising trace, the parallel diffusion process explores incorrect intermediate states but uses the emerging reasoning tokens to course-correct, seamlessly converging on the correct answer ( − 25 -25 ). See Appendix G.1 for the second example of logical reasoning and correction via bidirectional attention and denoising.

### 9.2 Dynamic and Adaptive Computation

Unlike AR models, which inherently allocate a fixed amount of computation per generated token, text diffusion facilitates dynamic test-time compute. Text diffusion models can autonomously calibrate the computational effort expended on a given prompt, trading increased inference time for enhanced generative performance based on the inherent difficulty of the task ( Ye et al., 2024 ) . DiffusionGemma automatically adapts to task difficult via adaptive stopping. As illustrated in Figure 5 , tasks of varying complexity organically elicit different numbers of effective denoising steps. This allows the model to conserve compute on easier queries while expending more compute on complex reasoning. In Appendix G.2 we concretely illustrate the adaptive behavior of text diffusion by contrasting a structurally “hard” generation task against a structurally “easy” task and show that the easy task requires fewer denoising steps, as well as showing how the information propagates differently across the canvas of tokens in these two cases.