DiffusionGemma Technical Report

Abstract: We introduce DiffusionGemma, an experimental open-weight LLM that uses discrete diffusion to generate text at exceptionally high speed. Rather than decoding one token at a time, DiffusionGemma iteratively refines blocks of 256 tokens in parallel, avoiding the sequential decoding bottleneck of conventional autoregressive (AR) LLMs. Instead of training from scratch, we obtain DiffusionGemma by fine-tuning the mixture-of-experts Gemma 4 model with 3.8B activated and 25.2B total parameters. Our compute-efficient two-stage training pipeline uses fewer than 10% of the starting AR model's total training token budget. The first stage uses supervised fine-tuning to teach bidirectional denoising, while the second stage combines reinforcement learning with sampler distillation to jointly improve generation quality and inference efficiency. DiffusionGemma establishes a new Pareto frontier for the trade-off between generation speed and model capability. Averaged across our full evaluation suite, it generates around 20 tokens per forward pass and achieves roughly 1,500 output tokens per second on a single NVIDIA H100 GPU, which is substantially faster than AR models even with state-of-the-art speculative decoding. DiffusionGemma also retains the starting model's support for thinking mode, multimodal inputs, and long contexts. Despite diffusion fine-tuning, it remains capable of AR generation with only minor performance degradation, suggesting a path toward hybrid diffusion-AR decoding.

• The paper introduces DiffusionGemma, an open-weight discrete diffusion LLM that denoises 256-token blocks in parallel and reaches 1,479 tokens per second at 19.74 tokens per forward pass on one NVIDIA H100.
• Its joint reinforcement-learning and sampler-distillation pipeline reduces denoising steps to roughly 12–16 while maintaining competitive results on GPQA-Diamond, LiveCodeBench, GSM8K, and MMLU-Pro.
• The model enables bidirectional self-correction, adaptive compute allocation, and rapid structured generation, while retaining AR usability and supporting LoRA adaptation such as 84% Sudoku accuracy with four times fewer denoising steps.

DiffusionGemma: Accelerating Open-Weight LLMs with Discrete Diffusion

Introduction and Motivation

DiffusionGemma defines an experimental open-weight LLM instantiating discrete diffusion for text generation, utilizing a block-parallel denoising mechanism instead of traditional autoregressive (AR) decoding. The model is derived by finetuning the Gemma 4 26B mixture-of-experts (MoE) AR LLM and offers a compute-efficient generation pipeline that outputs entire blocks of 256 tokens in parallel. This reduces the memory-bound bottleneck of AR decoding and introduces ultra-fast inference, moving the hardware utilization from memory to compute-bound regimes. Empirically, DiffusionGemma achieves an average of 20 tokens per forward pass (TPF) and can produce approximately 1,500 output tokens per second (TPS) on a single NVIDIA H100 GPU, surpassing major AR and open diffusion baselines Figure 1.

Figure 1: DiffusionGemma establishes a new Pareto frontier for the trade-off between generation quality and decoding speed, outperforming both AR and previous diffusion-based models.

Notably, DiffusionGemma exposes open, permissively licensed weights and reference implementations, targeting both research transparency and practical downstream adaptation.

Discrete Diffusion Framework

Instead of AR factorization, DiffusionGemma adopts a continuous-time Markov chain over categorical token spaces, parameterized by a monotonic noise schedule. The generative process is formally rooted in recent discrete flow-matching and Markov score-based modeling literature: data is progressively corrupted to noise during training; denoising is modeled as a blockwise, bidirectional process. At inference, the model reconstructs text from noise by iteratively refining a canvas of token states in parallel. Each denoising step stochastically transitions tokens based on their marginal posterior distributions, leveraging self-conditioning.

Figure 2: Stylized illustration of discrete diffusion with parallel sampling and probability paths across a 2D canvas.

Compared to continuous diffusion over projected embeddings, the strict categorical formalism avoids projection mismatch and rounding artifacts, preserving theoretical and statistical integrity in the generative trajectory.

Model Architecture and Blockwise Generation

DiffusionGemma utilizes an encoder-decoder Transformer architecture, with weights initialized from the public Gemma 4 AR checkpoint. The encoder performs causal context encoding into a Key-Value (KV) cache, while the decoder leverages bidirectional attention for parallel denoising. Generation proceeds in block-autoregressive (block-AR) fashion: canvases of 256 tokens are iteratively denoised, with the completed block appended to the context, supporting efficient long-sequence generation.

Figure 3: The inference pipeline: context encoding, parallel denoising via bidirectional decoder, and context update for next block.

Sampling is governed by an entropy-bounded strategy with temperature annealing and adaptive stopping, balancing quality, diversity, and compute. Hard convergence is enforced when confidence thresholds are met, further optimizing inference latency Figure 4.

Figure 4: Adaptive stopping modulates the number of denoising steps according to task complexity and domain, reducing unnecessary computation.

Training Pipeline: Supervised Finetuning and Joint RL/Sampler Distillation

Two training stages are employed:

1. Supervised Fine-Tuning (SFT): The model adapts to bidirectional, blockwise discrete diffusion via cross-entropy minimization using noise-corrupted inputs sampled at random time schedules. Downstream performance grows log-linearly, with significant gains for extended SFT in reasoning-oriented (“thinking”) mode (Figure 5, Figure 6).

Figure 5: SFT performance evolution differentiating thinking and non-thinking behaviors.

2. Joint RL and Sampler Distillation: To simultaneously raise generation quality and compress the denoising trajectory for ultra-low latency inference, an online objective combines reward maximization and sampler distillation. This curriculum results in reduced effective denoising steps and higher reward over time Figure 7. The resulting model advances the speed-to-quality Pareto frontier notably compared to the SFT baseline, especially in the few-step regime (Figure 8, Figure 9).

Figure 8: RL- and distillation-optimized checkpoints deliver significantly improved quality-speed tradeoffs over SFT baselines as measured by TPF and benchmark score.

Inference Efficiency and Hardware Optimization

The practical efficiency of DiffusionGemma arises from substantially fewer overall forward passes with moderately heavier each-pass compute. Per-step, DiffusionGemma is only 3.2× slower than AR for single-token generation, yet processes 256 tokens in parallel. Careful GPU kernel and MoE optimizations, including use of FlashAttention and efficient memory management, ensure minimal throughput loss per step Figure 10. In low-batch-size regimes, both per-user and aggregate throughput exceed AR baselines and even AR models using multi-token prediction (MTP), only becoming less efficient at batch sizes well beyond standard user-facing scenarios Figure 11.

Figure 10: Per-step GPU breakdown shows blockwise MoE and softmax dominate overhead, but the overall gain in TPF outweighs these factors.

Figure 11: DiffusionGemma outperforms AR and AR+MTP models in both total and per-user throughput at typical concurrency levels.

Experimental Results and Strong Numerical Claims

• Quality and Speed: DiffusionGemma achieves 1,479 TPS at 19.74 TPF (reflecting 5–10× acceleration over AR and at least 2× over proprietary Mercury 2 diffusion LLMs), with average effective denoising steps around 12–16.
• Benchmark Coverage: On complex eval suites (GPQA-Diamond, LiveCodeBench-v6, GSM8K, MMLU-Pro, etc.), DiffusionGemma delivers competitive performance versus AR and proprietary diffusion LLMs, while substantially outperforming open-weight diffusion competitors.
• Retention of AR Capability: Importantly, the model remains usable in AR mode with weights shared from the diffusion checkpoint, with only minor loss in benchmark accuracy.
• Practical Adaptation: The open-source downstream SFT toolkit enables rapid finetuning with LoRA on consumer hardware. For example, domain adaptation on Sudoku achieves 84% accuracy and a 4× reduction in denoising steps post-finetuning Figure 12.

Figure 12: Parameter-efficient LoRA downstream adaption enables strong accuracy for structure-constrained tasks at reduced compute cost.

Analysis: Architectural and Practical Implications

• Bidirectional Reasoning and Self-Correction: DiffusionGemma leverages parallel, bidirectional attention to revise earlier tokens based on later context during denoising, a property not possible in AR models. Denoising traces for arithmetic reasoning and logical puzzles exhibit iterative self-correction, preventing typical AR failure modes (Figures 15, 23).

Figure 13: Denoising trace for arithmetic reasoning: the final answer and supporting steps converge together in 5 refinement steps, illustrating non-causal revision.

• Dynamic Compute Allocation: Through adaptive stopping, DiffusionGemma allocates more or fewer refinement steps according to task difficulty, balancing latency and performance dynamically (Figures 24, 25).
• Structured Generation Acceleration: For tasks with rigid syntactic targets (e.g., JSON extraction, code editing), convergence often occurs in only 2–3 steps due to immediately lockable structure, providing dramatic speedups over AR decoding.
• Conciseness and Output Efficiency: RL-distilled checkpoints tend to output more concise, token-efficient generations, effectively compounding inference speedups.

Limitations

Despite establishing a new Pareto frontier, DiffusionGemma introduces certain trade-offs:

• Its absolute performance on some reasoning tasks lags the AR initialization, attributable to non-pretrained initialization, shorter SFT, and aggressive RL-driven step compression.
• Emergent concise outputs, while beneficial for speed, limit long-form, detailed responses common in AR-optimized LLMs.
• Occasional error modes include token stuttering or degeneration under ultra-aggressive stopping, and throughput falls below AR at high batch sizes unless further optimized.

Conclusion and Future Directions

DiffusionGemma demonstrates the feasibility of compute-efficient, blockwise discrete diffusion for ultra-fast text generation at scale. The release provides a practical baseline for further advances in inference optimization, bidirectional reasoning, and hybrid AR-diffusion frameworks. Directions for future work include larger-scale SFT and RL alignment, specialized sampling algorithms for distinct prompt varieties, further MoE and batch-size optimization, and deeper exploration of discrete-continuous hybrid flows and flow matching for text. Critically, the open-weight release sets a platform for empirical study of non-AR generative mechanisms under realistic LLM workloads and paves the way for domain-specific, efficient adaptation across industry and academic applications.

---

Reference:

DiffusionGemma Technical Report (2608.00146)