PaperlayerView exact revision on arXiv

Source-linked AI summary

DiffusionGemma Technical Report

DiffusionGemma Team, Adrien Ali Taïga, James Assiene, Daniele Calandriello, Rahma Chaabouni, João Gante, Tamara von Glehn, Nate Keating, Chris Knutsen, Martin Kukla, Tianlin Liu, Ivan Lobov, Ofir Nabati, João Gabriel Oliveira, Nicolas Perez-Nieves, Nastasia Prutianova, Bobak Shahriari, Jean Tarbouriech, Pavel Tyletski, Çağlar Ünlü, Cindy Wu, Glenn Cameron, Jerome Connor, Sertan Girgin, Maarten Grootendorst, Alon Levkovitch, Eliya Nachmani, Omar Sanseviero, Piotr Stanczyk, Quentin Berthet, Andrew Campbell, Clément Crepy, Valentin De Bortoli, Arnaud Doucet, Romuald Elie, Alexandre Galashov, Klaus Greff, Alexis Jacq, David Ruhe, Yu-Han Wu, Sebastian Flennerhag, Brendan O'Donoghue, George Scrivener, Shantanu Thakoor

arXiv:2608.00146v1cs.CLcs.AI

TL;DR

Autoregressive language models are limited by sequential decoding, while existing text-diffusion systems face trade-offs among speed, intelligence, and accessibility. DiffusionGemma fine-tunes Gemma 4 for parallel discrete-diffusion generation and reaches roughly 1,500 output tokens per second on a single H100 while retaining competitive reasoning and multimodal capabilities.

  • Problem

    Existing language models face a sequential-decoding bottleneck, while text-diffusion alternatives lack a demonstrated combination of extreme speed, intelligence, and open-weight accessibility.

  • Method

    DiffusionGemma fine-tunes an existing Gemma 4 mixture-of-experts checkpoint into an encoder-decoder discrete-diffusion model for parallel text generation.

  • Results

    Roughly 1,500 output tokens per second on a single NVIDIA H100 establishes a new speed-to-intelligence frontier while retaining competitive reasoning and multimodal capabilities.

  • Takeaways & Limitations

    DiffusionGemma provides an experimental open-weight, ultra-fast text-generation model that preserves key capabilities of its autoregressive starting point.

  • Takeaways & Limitations

    DiffusionGemma has lower absolute performance than its Gemma 4 autoregressive initialization and produces highly concise outputs.

Abstract

from arXiv · show

We introduce DiffusionGemma, an experimental open-weight language model 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) large language models. 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.

1. Introduction

DiffusionGemma addresses the memory-bound, token-by-token bottleneck of autoregressive generation by predicting token blocks simultaneously. Fine-tuned from Gemma 4, it combines high-speed diffusion decoding with retained multimodal, long-context, and reasoning capabilities, while being released with open weights.

  • Motivation: Autoregressive models become fundamentally memory-bound for single or low-concurrency requests because sequential decoding repeatedly transfers weights and context KV caches.Batching can improve throughput for many simultaneous requests, but memory transfer dominates actual computation when concurrency is low.
  • Contribution: DiffusionGemma is a fine-tuned text-diffusion variant of the Gemma 4 26B A4B mixture-of-experts model.It is designed to improve the intelligence-to-speed trade-off relative to prior text-diffusion and Gemma 4 autoregressive models.
  • Performance: 20 TPF and around 1,500 TPS on a single NVIDIA H100 GPU characterize DiffusionGemma’s high-speed generation.The model generates approximately 20 tokens per forward pass, compared with roughly 3–6 TPF for state-of-the-art speculative decoding.
  • Training: The training pipeline first uses supervised fine-tuning for bidirectional denoising of 256-token blocks, then combines sampler distillation with reinforcement learning.The second stage jointly improves generation quality and reduces the number of forward passes.
  • Capabilities and release: DiffusionGemma preserves multimodal understanding, long-context capabilities, and thinking mode despite a performance penalty relative to its autoregressive baseline.It is released with open weights under the permissive Apache 2.0 license.

2. Generative Text Modeling with Discrete Diffusion

Discrete diffusion replaces autoregressive left-to-right factorization with iterative refinement over noise levels, enabling bidirectional token revision during generation. DiffusionGemma formalizes this process with categorical corruption, learned reverse denoising, and parallel canvas refinement, while exposing conditional-independence trade-offs.

  • Autoregressive modeling: Autoregressive language models generate sequentially via single-token conditional factorization, limiting accelerator speed and preventing revision based on future tokens.The limitation arises from memory-bandwidth constraints and strictly left-to-right generation.
  • Discrete diffusion: Diffusion models factorize generation over a Markov chain of noise levels rather than sequence positions, while discrete diffusion applies this paradigm directly to categorical token states.Earlier text methods often operated in continuous embedding spaces and projected or sampled back to discrete tokens.
  • Forward process: The forward process gradually replaces clean tokens with uniformly random vocabulary tokens using a monotonic noise schedule, producing corrupted training data.The schedule decreases from κ0 = 1 to κ1 = 0 as time advances toward noise.
  • Reverse process: The reverse process learns the conditional distribution of clean tokens given noisy states and uses it to step backward through the corruption trajectory.A neural network approximates the clean-token posterior used by the sampling update.
  • Parallel generation: Parallel denoising refines a noisy canvas with bidirectional attention, but conditional independence between simultaneous token updates can cause local inconsistencies and grammatical conflicts.Each position sees the global noisy canvas but not other positions’ simultaneous sampling choices.

3. The DiffusionGemma Architecture

DiffusionGemma uses shared-weight encoder-decoder components initialized from Gemma 4 and generates open-ended text through block-autoregressive canvas denoising. Its architecture combines bidirectional within-canvas refinement, causal KV-cache updates across canvases, and a modular entropy-bounded sampler with adaptive stopping.

  • Architecture: DiffusionGemma is an encoder-decoder transformer with shared weights, initialized from the publicly released Gemma 4 26B A4B MoE checkpoint.This initialization reduces training cost relative to pretraining a diffusion model from scratch, while trading off absolute generation quality.
  • Block-autoregressive decoding: 256-token canvases are denoised sequentially at the block level, with each fully denoised canvas committed before generation begins on the next.This block-autoregressive strategy converts fixed-length diffusion generation into open-ended text generation.
  • Canvas denoising: The decoder denoises each canvas with bidirectional canvas attention and cross-attention to a context KV cache, while continuous self-conditioning feeds back prior predictions.The cache conditions generation on system instructions, user inputs, and past responses; previously generated canvases are frozen, whereas tokens within the current canvas can still be revised.
  • Entropy-bounded sampling: Approximately 12 effective denoising steps replace the maximum N=48 budget on downstream evaluations, yielding a 4× reduction in overall latency without sacrificing generation quality.Adaptive stopping dynamically scales inference-time computation to prompt complexity by bypassing redundant refinement steps.
  • Hybrid generation: The final DiffusionGemma weights can be loaded into the original transformer architecture for standard AR generation using causal attention.The model shares the exact same transformer architecture as Gemma 4 and retains robust AR capabilities.

4. Supervised Finetuning

Supervised finetuning adapts Gemma 4 to bidirectionally denoise 256-token blocks from noisy inputs using block-diagonal attention and discrete multinomial diffusion. Denoising improves rapidly before settling into a log-linear trend, while extended finetuning is especially important for coherent thinking behavior.

  • Finetuning method: 256-token blocks are predicted from noisy inputs using block-diagonal attention, which permits bidirectional attention within each block while isolating denoising blocks.The model conditions on the prompt and previous uncorrupted tokens through an encoder KV cache.
  • Finetuning method: Discrete multinomial diffusion uniformly samples noisy tokens from the vocabulary as the corruption process.
  • Training dynamics: Denoising performance improves rapidly during initial SFT, then settles into a log-linear improvement trend.The reported downstream results use entropy-bounded sampling with adaptive stopping and a maximum of N=192 denoising steps.
  • Thinking behavior: Extended SFT improves thinking performance by helping the model maintain coherent reasoning traces instead of collapsing into stuttering or cycles.Thinking performance initially starts lower but follows a steeper log-linear slope than non-thinking performance.

5. Sampler Distillation & Reinforcement Learning

SD·RL unifies reward maximization and sampler distillation to improve generation quality while compressing diffusion into the few-step regime. It expands the quality-speed frontier through higher reward, fewer effective denoising steps, and more concise outputs.

  • Motivation: Before SD·RL, quality was strong at many denoising steps but collapsed in the few-step regime, while advanced reasoning and coding lagged the baseline AR model.The training therefore targets simultaneous improvements in model intelligence and denoising-trajectory compression.
  • Unified SD·RL objective: SD·RL jointly optimizes reward maximization and sampler distillation in a unified online-learning stage, rather than using decoupled phases.A single gradient update targets both absolute generation quality and mapping high-quality generations to few-step diffusion.
  • Training dynamics: Higher reward and fewer effective denoising steps emerge together during SD·RL training, with adaptive stopping helping the online teacher reach high rewards faster.These dynamics push the model’s quality-speed Pareto frontier.
  • Emergent conciseness: Nearly 2× shorter generations: the final SD·RL checkpoint produces more concise outputs than the SFT checkpoint, multiplying inference-speed gains.This conciseness contrasts with autoregressive RL, which often rewards longer reasoning traces, though it forgoes some gains associated with extended reasoning.

6. Inference Optimizations

DiffusionGemma’s inference optimizations target GPU-level bottlenecks to reduce per-step cost and exploit diffusion’s parallel token generation. It provides the strongest latency and throughput advantages for single-request and low-batch-size serving, while AR models gain an advantage at around 32 concurrent requests.

  • Low batch size serving: Low-batch serving benefits from fewer memory transfers, which can outweigh diffusion’s higher FLOPs per generated token on memory-bound accelerators.The passage attributes this advantage to KV-cache capacity and memory bandwidth constraints and focuses its analysis on single-request inference.
  • GPU time breakdown: 256 tokens per step incur only a 3.2× per-step latency increase versus single-token AR generation.Figure 11 reports H100, FP8 serving with batch size 1, 4096 input tokens, and 1024 output tokens.
  • GPU time breakdown: DiffusionGemma activates approximately 84 unique experts per 256-token canvas per MoE layer, compared with 8 unique experts per token for Gemma 4 AR.For single-request serving, the MoE computation is described as memory-bound because expert weights must be transferred from high-bandwidth memory.
  • Throughput for batch size of 1 serving: 13.56 ms is the average end-to-end time for one DiffusionGemma denoising step on an H100 in the specified single-request setting.The corresponding per-step GPU time is 12.63 ms; the passage also states that TPF varies by task and gives 19.74 as an average assumption.
  • Multi-user throughput: Around 32 concurrent requests mark the regime where AR models begin to gain a throughput advantage over DiffusionGemma.At low batch sizes, DiffusionGemma provides substantially higher TPS per user and higher total throughput than Gemma 4 AR with MTP; the reported runs use H100 FP8 precision on PG-19.
  • Toward real traffic throughput: DiffusionGemma performs TPF× fewer KV-cache transfers in attention while increasing MoE FLOPs proportionally to effective denoising steps.This trade-off may help address memory-bound attention that limits utilization of feed-forward and MoE layers in real-traffic serving.

7. Experimental Results

DiffusionGemma is evaluated across text-diffusion and autoregressive modes, with and without thinking, against its AR initialization and leading open- and closed-weight diffusion models. It substantially improves diffusion-model speed while remaining highly competitive in quality, though TD mode trades some benchmark performance for faster decoding than the AR baseline.

  • Evaluation setup: Evaluations compare text-diffusion and autoregressive generation with and without thinking against Gemma 4 26B A4B, LLaDA 2.1 Flash 100B, Nemotron Diffusion 14B, and Mercury 2.Unless otherwise specified, text diffusion with thinking is enabled.
  • Evaluation setup: The benchmark suite spans mathematical reasoning, coding, general knowledge, multimodal understanding, instruction following, and agentic capabilities.The reported benchmark groups cover reasoning and knowledge, coding, and instruction following and agentic behaviour.
  • Performance and efficiency: DiffusionGemma substantially outperforms existing open-weight diffusion baselines while increasing Tokens Per Forward by approximately an order of magnitude.Table 4 additionally reports TPF, TPS, effective denoising steps, total generated tokens, and end-to-end latency excluding prefill.
  • Performance and efficiency: 1,500 output tokens per second are reached on a single previous-generation H100 GPU, yielding a roughly 2.5× speedup over Mercury 2.Speed measurements exclude prefill time.
  • Performance versus AR: 1,479 tokens per second versus 303 tokens per second gives TD mode nearly 5× the output throughput of the Gemma 4 AR baseline under heavily optimized MTP serving.The conversion reduces performance across the three capability areas while substantially increasing decoding speed.

8. Open-Source Downstream SFT

The authors release an open-source toolkit for adapting DiffusionGemma with LoRA on consumer hardware, supporting a combined causal-encoder and diffusion-decoder training objective. In a Sudoku case study, finetuning substantially improves puzzle-solving accuracy, with full finetuning exceeding 85% on a held-out set of 4096 puzzles.

  • Sudoku puzzle solving: Above 80% Sudoku accuracy is reported after finetuning with LoRA rank 8, whereas the base model fails to produce a correct grid before finetuning.The finetuned model can use significantly fewer denoising steps because of lower predictive entropy.
  • Toolkit release: The open-source finetuning toolkit adapts DiffusionGemma to domain-specific datasets and provides LoRA recipes for consumer hardware.It builds on Hackable Diffusion, a modular open-source research toolbox for generative modeling.
  • Finetuning procedure: The SFT toolkit combines causal-encoder and diffusion-decoder objectives over sequences containing prompt tokens and multiple canvases.The encoder processes the full sequence and contributes a standard cross-entropy loss, while the decoder contributes a denoising loss; the final loss sums both.
  • LoRA finetuning: LoRA is applied to all linear operations, including attention projections, MLP gates, MoE routers, and the self-conditioning feedforward block.This enables downstream finetuning while training only a small fraction of parameters using 2× A100 80GB GPUs.
  • Sudoku puzzle solving: >85% puzzle-level accuracy is achieved by the full-finetuning sampler on a held-out set of 4096 Sudoku puzzles.Sudoku is used as a testbed for discrete diffusion because solving it is non-autoregressive.

9. Practical Advantages of Text Diffusion

Text diffusion offers practical advantages beyond low latency, including bidirectional conditioning, iterative self-correction, adaptive test-time computation, and rapid convergence on highly constrained outputs.

  • Full bidirectional attention lets tokens across the canvas attend simultaneously to past and future representations, unlike causal AR next-token prediction.This non-causal property enables conditioning on tokens that have yet to be generated.
  • Iterative refinement provides built-in self-correction, allowing premature commitments from earlier denoising steps to be revised later.In the arithmetic example, DiffusionGemma produces the correct answer “−25” immediately, while AR Gemma commits to “−1” and later appends a correction.
  • Adaptive stopping allocates different numbers of denoising steps by task difficulty, conserving compute on easy queries and spending more on complex reasoning.The maximum denoising-step count also explicitly controls the latency-quality tradeoff: fewer steps are faster but less precise, while more steps improve quality at higher compute cost.
  • 2 to 3 steps suffice for highly constrained outputs, drastically reducing latency compared with sequential decoding.Structured formats and strong lexical dependence, including JSON schemas, OCR, code editing, and fine-grained syntactic control, provide structural priors that parallel generation can exploit.

10. Limitations & Known Issues

DiffusionGemma’s experimental release has known limitations, including lower absolute performance than its AR initialization and occasional robustness and multimodal-output issues. These limitations reflect compute-efficient training choices and ultra-low-latency decoding.

  • The experimental release has known limitations despite establishing a Pareto frontier for generative text modeling and inference efficiency.The report frames these as current-release limitations rather than eliminating the approach’s reported efficiency advantages.
  • Lower absolute performance than Gemma 4 26B A4B reflects AR warm-starting, short SFT, ultra-low-latency SD·RL, and inherited design constraints.The cited passage attributes the gap to bypassing native diffusion pretraining, compute-limited supervised fine-tuning, latency-focused online learning, and inherited architectural and optimization factors.
  • Aggressively reduced denoising steps can occasionally compromise generation robustness, leaving an uncommon artifact despite SD·RL mitigating most such instances.This issue is described as a direct consequence of operating in an ultra-low-latency regime.
  • In multimodal tasks, the model does not always reliably generate a closing thought tag, even when reasoning is present.The passage identifies this as an occasional omission when processing multimodal prompts.

11. Conclusion

DiffusionGemma offers a practical, compute-efficient path to ultra-fast text generation by fine-tuning Gemma 4 for text diffusion. It reaches around 1,500 tokens per second on a single H100 while retaining competitive reasoning and multimodal capabilities.

  • Approach: DiffusionGemma fine-tunes the existing Gemma 4 26B A4B autoregressive model to perform text diffusion.The approach uses a two-stage training pipeline combining supervised fine-tuning (SFT) and SD·RL.
  • Results: 1,500 tokens per second is achieved on a single H100.This speed contributes to a new Pareto frontier for the speed-to-intelligence tradeoff.
  • Results: DiffusionGemma retains highly competitive reasoning and multimodal capabilities.These capabilities are preserved alongside its ultra-fast generation performance.

Contributions and Acknowledgments (listed alphabetically)

This section lists contributors and acknowledged individuals alphabetically. The names span eight consecutive entries.

  • Contributions and Acknowledgments (listed alphabetically): The first alphabetical entries list Adrien Ali Taïga through Nicolas Perez-Nieves*.Names include Adrien Ali Taïga, James Assiene, Daniele Calandriello*, Rahma Chaabouni, João Gante*, Tamara von Glehn*, Nate Keating, Chris Knutsen, Martin Kukla*, Tianlin Liu, Ivan Lobov*, Ofir Nabati, João Gabriel Oliveira, and Nicolas Perez-Nieves*.
  • Contributions and Acknowledgments (listed alphabetically): The middle entries continue with Omar Sanseviero, Piotr Stanczyk, Arnaud Doucet, Romuald Elie, Alexandre Galashov, and Klaus Greff.These names occupy the next listed entries.
  • Contributions and Acknowledgments (listed alphabetically): The remaining entries list Quentin Berthet, Andrew Campbell, Clément Crepy, Valentin De Bortoli, Alexis Jacq, David Ruhe, and Yu-Han Wu.These names appear in the subsequent alphabetical entries.
  • Contributions and Acknowledgments (listed alphabetically): The final entries name George Scrivener, Shantanu Thakoor, Sebastian Flennerhag, and Brendan O’Donoghue.These names complete the listed acknowledgments.

Sponsors … A.2. Speculative Decoding

The section acknowledges the report’s sponsors and situates DiffusionGemma within continuous-diffusion and speculative-decoding research. It highlights continuous embedding-space and flow-matching alternatives, while describing speculative decoding’s draft–target latency trade-off.

  • Sponsors: The Sponsors section lists Danilla Sinopalnikov and the Gemma Team.
  • Sponsors: It also lists Demis Hassabis, Prateek Jain, Armand Joulin, and Koray Kavukcuoglu.
  • Sponsors: Additional listed sponsors include Olivier Bachem, Jeff Dean, Zoubin Ghahramani, and Raia Hadsell.
  • Sponsors: The sponsors list further includes Marc’Aurelio Ranzato and Oriol Vinyals.
  • A.1. Continuous Diffusion for Text: Recent work indicates a resurgence of continuous and hybrid approaches, although discrete diffusion remains highly effective.
  • A.1. Continuous Diffusion for Text: Flow-matching models such as Embedded Language Flows avoid per-step token supervision and remain in continuous embedding space until final discretization.The passage states that these models can substantially outperform discrete baselines.
  • A.2. Speculative Decoding: Speculative decoding uses a smaller draft model to propose token sequences that a larger target model verifies in parallel, reducing serving latency.
  • A.2. Speculative Decoding: Speculative-decoding latency depends on both draft and target generation, while autoregressive drafters face sequential-generation constraints as draft quality requires more parameters.The passage notes that larger AR drafters increase per-token latency and diminish the latency benefit.

B. Samples Before and After SD·RL Training … G.3. Structured and Constrained Outputs

Across qualitative examples and implementation details, DiffusionGemma improves reasoning, formatting, adaptive computation, and structured generation through denoising. The report also documents Mercury 2 speed estimation, downstream adaptation, prompt serialization, and a seven-step denoising trajectory.

  • B. Samples Before and After SD·RL Training: SD·RL training resolves repetitive looping in the SFT model, enabling complete reasoning traces and correct answers on representative GPQA-Diamond examples.The SFT model initially formed repetitive token loops and empty responses in biology and physics examples; the final checkpoint completed both analyses cleanly.
  • C. Additional Open-Source Downstream Finetuning Results: 8M parameters are finetuned for Sudoku with rank-8 LoRA across all linear layers, while downstream finetuning makes the model output a correctly solved 9x9 grid.The downstream recipes also cover PubMedQA, where the model produces categorical answers and explanatory paragraphs evaluated with BLEU on the long task.
  • D. Prompt Formatting; D.1. BOS and EOS Special Tokens; D.2. Conversational Structuring; D.3. Thinking Channels; D.4. Tool and Function Calling Serialization: A structured Jinja2-based chat template serializes conversation history, instructions, tool schemas, roles, multimodal inputs, thinking channels, and tool responses using specialized control tokens.The template supports system, user, model, and tool roles; image tokens are replaced by image features, and JSON-like schemas use custom delimiters and tool-response wrappers.
  • F. Denoising trajectory sampled from DiffusionGemma: 7 denoising steps produce a coherent sentence describing text diffusion, with iterative refinement from random noise to structured text.The reported final response is “Text diffusion is a generative process that creates high-quality text by iteratively refining random noise until a structured and coherent message emerges.”
  • G.1. Bidirectional Reasoning and Self-Correction: DiffusionGemma self-corrects a frog-puzzle answer from an initial “Yes” to the correct “No” in 6 denoising steps, whereas Gemma 4 commits to a contradictory “Yes.”Bidirectional propagation of logical constraints lets the response and reasoning co-evolve; both models operate in non-thinking mode.
  • G.2. More on Dynamic and Adaptive Computation: 7 denoising steps resolve a structurally hard 40-token sequential sequence, while 4 steps resolve an easy static-input sequence in parallel.The hard task requires predominantly left-to-right reasoning because each token depends on two preceding outputs; removing that dependency enables simultaneous convergence through bidirectional attention.
  • G.3. Structured and Constrained Outputs: 100% per-token confidence is reached by the second denoising step for schema-constrained JSON, while bug-free code editing completes in 3 denoising steps.The JSON output begins with average confidence exceeding 83%, and code editing converges rapidly because most tokens are copied verbatim and only a localized semantic fix is needed.
Loading 2608.00146v1…