Google's Diffusion Gemma: Parallel Token Generation Explained

Diffusion-Based Parallel Decoding

Unlike traditional autoregressive models that generate tokens sequentially and cannot revise past output, Diffusion Gemma uses a diffusion process to generate text in fixed-window patches of 256 tokens. The model generates a rough draft of the entire block simultaneously and iteratively refines it. During each denoising step, the model calculates the entropy of its predictions at each position. It locks in positions where it has high confidence, up to a set budget, while discarding lower-confidence tokens back into noise for reconstruction in the next pass. This hybrid approach functions as diffusion within 256-token blocks and autoregressive generation across blocks, enabling the model to correct errors in earlier tokens as it generates.

Architecture and Performance

The model is a 26B parameter Mixture-of-Experts (MoE) architecture with approximately 4B active parameters per token. It utilizes 128 experts, with each token routed to 8 experts plus one shared expert. The core architecture features 30 layers of sliding-window attention with periodic global layers, supporting a context window of up to 256K tokens. While the model offers significant speed advantages—reaching approximately 700 tokens per second on an H100 compared to 300 tokens per second for standard autoregressive models—it generally lags behind equivalent-sized autoregressive models on standard benchmarks, representing a trade-off between inference speed and raw accuracy.

Hardware and Deployment

Memory requirements scale with quantization precision:

• BF16: 52 GB VRAM
• FP8: 27 GB VRAM
• NVFP4: 18 GB VRAM
• GGUF: 17 to 27 GB VRAM

The model is supported out-of-the-box by Transformers, vLLM, MLX, and llama.cpp. For local deployment on Apple Silicon, 4-bit quantization allows for usable generation speeds, though performance is sensitive to the length of the context window.