DiffusionGemma Technical Report

Authors: DiffusionGemma Team
Affiliations: Google DeepMind
Paper: https://arxiv.org/abs/2608.00146
Code: https://github.com/google/hackable_diffusion
Model: https://deepmind.google/models/gemma/diffusiongemma/

TL;DR

WHAT was done? The authors introduce DiffusionGemma, an experimental open-weight text diffusion model with 25.2B total parameters and 3.85B activated parameters based on the Gemma 4 26B A4B mixture-of-experts backbone. Rather than generating text strictly left-to-right one token at a time, DiffusionGemma iteratively refines canvases of 256 tokens in parallel. Bypassing native diffusion pretraining, the model is warm-started from post-trained autoregressive weights and fine-tuned using a compute-efficient two-stage pipeline consisting of Supervised Fine-Tuning (SFT) and a novel joint Sampler Distillation & Reinforcement Learning (SD-RL) phase.

WHY it matters? Standard autoregressive language models serving single-request or low-concurrency workloads are fundamentally memory-bandwidth bound, spending vastly more time moving weights and KV caches between high-bandwidth memory and compute units than performing floating-point operations. DiffusionGemma shifts generation from a memory-bound regime to a compute-bound regime, generating approximately 20 tokens per forward pass and achieving roughly 1,500 output tokens per second on a single NVIDIA H100 GPU. This represents a 7.1x speedup over its autoregressive baseline while maintaining reasoning, long-context support, and multimodal capabilities.

Executive summary: For technical leaders and systems architects, DiffusionGemma marks a practical breakthrough in non-autoregressive text generation. By converting a strong, pre-trained mixture-of-experts model into a discrete text diffusion engine using less than 10% of the original training token budget, Google DeepMind demonstrates that high-speed parallel generation does not require sacrificing complex reasoning or multimodal understanding. By releasing the model weights under an Apache 2.0 license alongside reference implementations in vLLM and HuggingFace Transformers, this work provides the broader community with a high-throughput substrate for real-time applications, agentic workflows, and low-latency deployments.

Details

The Memory-Bandwidth Bottleneck in Autoregressive Serving

Large language model inference is constrained by the architectural properties of modern hardware accelerators. Standard autoregressive generation factorizes sequence joint probability into an exact product of conditional probabilities

\(p(x) = \prod_{i=1}^L p(x^i \mid x^{<i})\)

necessitating a distinct forward pass for every single token generated. When serving single-request or low-batch-size workloads, execution time is dominated by memory transfers, specifically loading multi-billion parameter model weights and past key-value (KV) states into chip SRAM. This memory-bandwidth wall leaves tensor processing units under-utilized.

While techniques like speculative decoding and Multi-Token Prediction (MTP) attempt to increase memory efficiency by using draft models to propose candidate sequences, they remain constrained by draft-verification overheads and declining acceptance rates at later token positions. Text diffusion provides an alternative by converting single-token serial generation into a parallel canvas refinement process. By processing blocks of 256 tokens simultaneously, diffusion models execute dramatically fewer total forward passes per sequence, exchanging memory bandwidth overhead for floating-point calculations that fully utilize modern GPU compute units.

Discrete Flow Matching on Categorical Canvases

To formalize discrete text diffusion, the authors build upon continuous-time Markov chains (CTMC) and discrete flow matching theory. Let V denote a categorical token vocabulary of size V=262,000, and let C=256 represent the sequence canvas length. A realization of the canvas at time t∈[0,1] is represented as a vector x_t​∈V^C. The model defines a marginal probability path that smoothly interpolates between an uncorrupted data distribution x₀​∼p(⋅) at t=0 and a fully corrupted source canvas x₁∼Unif(V^C) at t=1.

The continuous-time forward process corrupts clean text tokens into uniformly distributed categorical noise. Conditioned on a clean starting canvas x₀​, the transition path factorizes independently across token coordinates i∈{1,…,C} according to the probability rule:

\(\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]\)

In this formulation, κ_t​∈[0,1] is a monotonically decreasing schedule running from κ₀​=1 down to κ₁​=0, and δ denotes the Kronecker delta. To generate text, the model learns a reverse denoising process parameterized by a neural network p_θ​(v∣x_t​) that approximates the posterior distribution of clean tokens given a noisy canvas state x_t​. The transition update stepping backward in time by an increment Δt is expressed as P(X_t−Δt​=⋅∣X_t​=x_t​)≈Step(x_t​,p_θ​(⋅∣x_t​)). Operating directly over discrete categorical distributions avoids the rounding errors and spatial drift that historically hindered continuous embedding space text diffusion.

The Block-Autoregressive Denoising Loop

DiffusionGemma operates via a hybrid block-autoregressive generation pipeline that reconciles fixed-length discrete flow matching with open-ended text generation, as illustrated in Figure 4. The architecture utilizes an inverted transformer setup: context history is processed using a causal encoder to construct a persistent Key-Value (KV) cache H, while generation across the current 256-token canvas is handled by a bidirectional decoder with shared weights θ.

Consider a specific execution trajectory during inference. Given an input prompt h∈V^L, the system populates the initial KV cache using H=Encoder_θ​(h). Canvas generation starts by initializing a sequence x₁​ of 256 uniformly random tokens from vocabulary V, alongside an initial self-conditioning signal z₁​=0. At each denoising iteration step t, the decoder takes the corrupted canvas x_t​, the context KV cache H, and the self-conditioning state z_t​∈R^C×d to produce unnormalized logits:

L_t​=Decoder_θ​(x_t​,z_t​,H)∈R^C×V

The clean token probabilities 𝑝̂₀​=Softmax(L_t​/τ_t​) are evaluated using a time-dependent temperature schedule τ_t​ that anneals linearly from τ_max​=0.8 down to τ_min​=0.4. An updated self-conditioning vector z_t−Δt​=FFW(𝑝̂₀​​E)∈R^C×d is constructed by passing the expected token embeddings through a feedforward network, feeding the model’s intermediate structural predictions directly back into the next forward pass.

To update the canvas x_t−Δt​, DiffusionGemma uses an entropy-bounded sampler (Algorithm 1). Positions across the canvas are ordered by predictive entropy eⁱ=Entropy(𝑝̂₀ⁱ​). Tokens at positions meeting a strict cumulative error tolerance threshold b=0.1 are updated with candidate samples, whereas higher-uncertainty positions are re-noised uniformly at random to force continued local exploration. Denoising is dynamically halted via an adaptive stopping heuristic when the mean predictive entropy across the canvas drops below e_stop​=0.005 and the deterministic predictions between consecutive steps remain identical. Once fully denoised, the clean canvas x̂₀​ is appended to the KV cache via H←H⊕Encoder_θ​(x̂₀​), and the model advances to the next 256-token block.

Optimization, Losses, and Distillation

Rather than executing full pretraining from scratch, DiffusionGemma warm-starts directly from the post-trained weights of the Gemma 4 26B A4B Mixture-of-Experts checkpoint (Gemma Team et al., 2026). This MoE backbone activates 3.85B parameters out of 25.2B total per token, utilizing 8 active experts out of 128 alongside 1 shared expert. As detailed in Figure 2, the model undergoes a compute-efficient two-stage training pipeline consuming under 10% of the base AR model’s training token budget.

The first stage is Supervised Fine-Tuning (SFT), which adapts the pre-trained causal representations to bidirectional block-level denoising. A block-diagonal attention mask enables full non-causal attention within each 256-token canvas while conditioning on prior context via cross-attention to the KV cache. Noise levels t∼Unif[0,1] are sampled to corrupt training tokens, and the model optimizes a cross-entropy loss over uncorrupted targets:

\(\mathcal{L}(\theta) = - \sum_{i=1}^{C} \log p_\theta(x_0^i \mid x_t, z_t, H)\)

Following SFT, the model achieves high output quality when allocated a large step budget (N=192), but experiences performance degradation under few-step low-latency regimes. To address this, the authors introduce a unified second stage: Sampler Distillation & Reinforcement Learning (SD-RL). Rather than separating preference alignment and step distillation, SD-RL applies a joint online gradient objective. An online teacher model generates high-step denoising trajectories to compute reward estimates across mathematical reasoning, coding, and instruction-following environments.

The SD-RL objective optimizes reward generation while aggressively driving down predictive entropy across the reverse process trajectory. As entropy decreases, the adaptive stopping mechanism triggers earlier in the reverse chain (Figure 8), creating an implicit curriculum that compresses high-reward generations into fewer denoising steps.

For practical community adaptations, parameter-efficient fine-tuning recipes were developed using Low-Rank Adaptation (LoRA) (Hu et al., 2022). LoRA matrices are injected across linear projections, including attention heads, MLP gates, MoE routers, and the self-conditioning feedforward block. Training can be executed using open-source JAX utilities like Hackable Diffusion on modest hardware setups (such as 2x NVIDIA A100 80GB GPUs).

Hardware-level inference execution relies on customized lower-level kernels. Bidirectional attention across 256-token blocks is accelerated via FlashAttention-4 (Dao et al., 2022), while sampling routines are optimized using PyTorch torch.compile primitives. Serving efficiency is further enhanced by eliminating CPU-GPU synchronization bottlenecks through asynchronous sequence scheduling within native engine environments like vLLM.

Empirical Pareto Frontiers and Speedup Sources

DiffusionGemma establishes a new Pareto frontier balancing output speed and generation quality, outperforming both open-weight and proprietary baselines as shown in Figure 1.

On a single NVIDIA H100 GPU operating in FP8 precision at batch size 1, DiffusionGemma reaches an average output speed of 1,479 Tokens Per Second (TPS) with a Tokens Per Forward (TPF) metric of 19.74 across benchmark suites (Table 3). This represents a 7.1x speed improvement over the standard Gemma 4 AR baseline (204 TPS) and a 4.8x speedup over Gemma 4 equipped with Multi-Token Prediction (303 TPS). On GPQA-Diamond, DiffusionGemma achieves 73.2% accuracy in text diffusion mode, outperforming LLaDA 2.1 Flash 100B (68.7% at 375 TPS) and Nemotron Diffusion 14B (47.0% at 49 TPS), while remaining highly competitive with closed-weight APIs like Mercury 2 (75.2% at 600 TPS).

The source of this efficiency boost is analyzed through latency breakdowns in Figure 11. Evaluating a single forward pass over a 256-token canvas requires more compute operations than generating a single token autoregressively. Specifically, processing a canvas incurs a 3.2x increase in per-step GPU kernel execution time (12.63 ms vs 4.01 ms for single-token AR). This increase is driven by broader MoE expert activation (4.66 ms vs 1.08 ms, a 4.3x slowdown due to 84 active experts per canvas versus 8 per single token), denoising sampling overhead (3.06 ms vs 0.56 ms, a 5.5x increase), and bidirectional attention computation (1.84 ms vs 0.45 ms, a 4.1x increase).

However, because each step processes 256 tokens simultaneously and converges in roughly 12 effective denoising steps (Table 4), the overall reduction in total forward passes far outweighs the kernel execution overhead.

Ablation analysis confirms the importance of the SD-RL training phase (Figure 9).

The SFT-only baseline evaluated with restricted denoising steps (N=48) suffers from predictive entropy collapse and token repetition loops (Figure 16). SD-RL training resolves these loop artifacts, providing a +10 point accuracy gain on combined GPQA-Diamond and LiveCodeBench-v6 benchmarks while quadrupling generation TPF from 5 to nearly 20.

Qualitative probing reveals that bidirectional attention enables in-canvas self-correction (Figure 15 and Figure 23). When presented with multi-step arithmetic or logical traps, causal AR models frequently commit to incorrect early tokens and must subsequently generate explicit textual corrections. DiffusionGemma leverages non-causal attention across the canvas to co-evolve intermediate reasoning tokens and final outputs simultaneously, self-correcting flawed logic paths across initial denoising iterations before finalizing discrete token commitments.

Contextualizing Text Diffusion and Speculative Baselines

DiffusionGemma advances non-autoregressive text modeling beyond prior open-weights and proprietary approaches. Earlier text diffusion systems faced theoretical trade-offs: continuous diffusion models like Diffusion-LM (Li et al., 2022) required rounding mechanisms to map continuous vectors back to discrete tokens, leading to spatial drift and ungrammatical outputs. Conversely, discrete diffusion architectures like SEDD (Lou et al., 2024) established discrete transition models, but often suffered from slow inference sampling.

In relation to contemporary diffusion baselines, DiffusionGemma demonstrates notable speed and architectural advantages. Models such as LLaDA 2.1 Flash 100B (Bie et al., 2025) achieve 375 TPS across 8x NVIDIA B200 GPUs, while Nemotron Diffusion 14B yields 49 TPS on an H100. Proprietary APIs like Mercury 2 (Inception Labs, 2025) deliver ~600 TPS. DiffusionGemma reaches roughly 1,500 TPS on a single H100 GPU while providing an open-weight implementation that retains extended context processing, multimodal inputs, and an explicit internal reasoning channel.

Compared to speculative decoding frameworks such as Medusa (Cai et al., 2024), EAGLE-3 (Li et al., 2026), or TiDAR (Liu et al., 2026b), DiffusionGemma does not require separate draft-then-verify loops or target model verification overheads. By formulating generation as block-autoregressive canvas denoising, the model maintains high joint token acceptance across full 256-token canvases without suffering from suffix decay at later sequence positions.

Architectural Trade-offs and Generative Artifacts

Despite its decoding efficiency, DiffusionGemma exhibits theoretical and operational trade-offs. The model shows a modest absolute quality delta compared to its starting Gemma 4 26B AR baseline across dense reasoning benchmarks (Table 3). This gap stems from practical design choices: bypassing native diffusion pretraining from scratch in favor of warm-started AR weights, constraining SFT compute budgets, and employing an SD-RL trajectory compression objective that explicitly optimizes for ultra-low latency.

Another emergent characteristic is output conciseness. Post-SD-RL checkpoints generate responses that are roughly 2x shorter than those produced by the SFT baseline (Figure 9). While concise answers minimize total sequence length and accelerate end-to-end latency, they restrict the model’s ability to utilize long, highly detailed Chain-of-Thought reasoning traces that can benefit complex mathematical proofs.

In addition, the model can occasionally exhibit low-latency generation artifacts. Under tightly restricted step bounds, outputs may rarely fall into localized token stuttering. In multimodal settings, the model can occasionally omit closing thought tags (<|channel|>), which can disrupt parsing logic during automated benchmarking. Finally, serving throughput advantages are concentrated in low-to-moderate batch size regimes (N≤32 concurrent requests, as shown in Figure 12); at higher batch sizes, traditional autoregressive models regain throughput parity as workload characteristics shift back toward being compute-bound across larger parallel batches.

Strategic Outlook: Beyond Left-to-Right Generation

DiffusionGemma presents a viable path for deploying non-autoregressive language models in real-time, latency-critical production environments. By converting a standard Mixture-of-Experts transformer backbone into a discrete flow matching engine, the authors demonstrate that text diffusion can overcome the sequential decoding memory-bandwidth bottleneck without requiring custom pretraining infrastructure.

The dual-mode capacity of DiffusionGemma—retaining standard autoregressive execution while enabling discrete block diffusion—points toward dynamic inference routing systems. Future serving infrastructures could route low-latency, highly constrained, or structured tasks through ultra-fast diffusion decoding while reserving sequential autoregressive sampling for open-ended, extended reasoning queries. Releasing the model weights under an Apache 2.0 open license alongside ecosystem integration in HuggingFace and vLLM offers a strong foundation for future research in parallel text generation algorithms, advanced sampler designs, and optimized hybrid decoding architectures.