DiffusionGemma Explained
by Timothy Gao
Google recently released DiffusionGemma, a 26B open-weight uniform state diffusion language model ~4× faster than Gemma 4 for certain workloads.1
As an experimental model, it doesn’t have a whitepaper, and the officially released docs are pretty brief and high-level. This blog post, based on the colab notebook, attempts to fill in the gaps. It walks through an annotated from-scratch reimplementation derived from its open weights, following the style of the awesome Annotated Transformer.
Along the way, we explain some of the more opaque details — behind the model architecture, self-conditioning, the sampling procedure, how encode/decode weight sharing is actually implemented. Through closely examining the model config and weights, we also uncover some interesting design choices — the layer scalar, partial RoPE, and Google’s Scalar QK Norm — and offer possible explanations for them. Finally, we conclude by making some visualizations and empirical observations about the denoising process.
Thank you Lucas Gu, Arshia Nayebnazar, Henry Ko, Rishi Athavale, and Tejas Prabhune for help proofreading and improving this blog post!
Note: prior knowledge of vanilla autoregressive LLM implementation (Llama 3.1, MOE) is assumed.
Table of contents
Setup + Load Model
Install, import, load configs and weights
!pip install -q -U huggingface_hub
!hf download google/diffusiongemma-26B-A4B-it \
--local-dir /content/diffusiongemma
import glob
import json
import torch
import torch.nn.functional as F
from einops import rearrange, einsum
from safetensors.torch import load_file
from tokenizers import Tokenizer
from torch.distributions import Categorical
from tqdm import trange
torch.set_default_device("cuda")
max_tot_tokens = 2048
checkpoint = "diffusiongemma"
canvas_len = 256
model_config = json.load(open(f"{checkpoint}/config.json"))['text_config']
gen_config = json.load(open(f"{checkpoint}/generation_config.json"))
V = model_config['vocab_size']
sd = {}
for safetensor_path in glob.glob(f"{checkpoint}/model-*.safetensors"):
sd |= {k: v for k, v in load_file(safetensor_path, device="cuda").items() if "vision" not in k}
W_vocab = sd['model.decoder.embed_tokens.weight']
sdGoogle’s Scalar QK Normalization
Printing out the weights, we find that interestingly, all entries within each learned QK-normalization weight vector have the same value. For example:
'model.decoder.layers.17.self_attn.q_norm.weight': tensor(
[0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883,
0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, 0.9883, ...]
)The QK-norm weight vectors have the form
where γ_q and γ_k are the learned scalars.
RoPE preserves Euclidean norm, so the same result holds after applying rotate:
For a query q and key k_i separated by an angle θ_i,
Thus, the attention score for k_i is
where
is a fixed, learned effective attention temperature per layer.
In regular QK norm models (and non-QK norm models), the magnitudes of q and k can vary. This 1) produces a per-query-dependent effective temperature and 2) allows key magnitude to encode a form of query-independent “global importance”:
This may be undesirable in long-context settings, where a high-magnitude early key could remain disproportionately influential even after its relevance has faded. One interpretation is that scalar QK norm effectively avoids this, and makes the inner product truly capture similarity between a q and k without being confounded with anything else.
Verify all QK norm weights are like this
for layer_id in range(model_config["num_hidden_layers"]):
for kind in ("q", "k"):
w = sd[f"model.decoder.layers.{layer_id}.self_attn.{kind}_norm.weight"]
assert (w == w[0]).all(), "QK Norm entries are not all equal"We do RMSNorm in higher precision like usual because precision errors are accumulated through the reduction, potentially amplified through the division, and propagated to every element.
RMSNorm
def rms(x, w = 1):
return (w * x * (x.shape[-1] ** 0.5) / (torch.norm(x, dim=-1, keepdim=True, dtype=torch.float32) + model_config['rms_norm_eps'])).to(x.dtype)Rotary position embeddings
DiffusionGemma has two RoPE configs:
Standard RoPE assigns each of the d/2 complex / rotary pairs one frequency (head dim d):
This is what the sliding-attention layers use, but interestingly, DiffusionGemma’s full-attention layers use partial RoPE: after constructing the normal frequency vector, we then set all but the first 25% of its pair frequencies to zero, which will make rotate a no-op on those dimensions. Notably, we are doing truncating the full frequency spectrum, not the more intuitive option of rescaling it like we do in YaRN etc.
Partial RoPE enables a head to carry both explicitly position-sensitive features and features whose representation is not rotated as position changes.
Interestingly, the rows of W_k and W_q are stored arranged so that the projected vectors use split-half rotary pairs, (t, t + d/2), rather than adjacent pairs, (t, t+1), as in the original RoFormer paper. This makes the rotation easier to express and slightly more cache-friendly.
From equation 34 in RoFormer,
where rot(x) rotates each two-dimensional pair by 90 degrees counterclockwise, thus for a token at position p, the i-th pair of its k/q vector is rotated ccw by ω_i * p. Note our code’s rotate(x) is RoPE(x), not to be confused with rot(x).
In other words, instead of something like
rot_x = np.empty_like(x)
rot_x[..., 0::2] = -x[..., 1::2]
rot_x[..., 1::2] = x[..., 0::2]for (t, t+1) pairs,
we have the more cache friendly
rot_x = rearrange([-x[..., head_dim // 2 :], x[..., : head_dim // 2]], 'z ... d -> ... (z d)')Precompute RoPE frequencies and rotate
# precompute frequencies
freq = {}
# sliding window - regular rope
freq['sliding_attention'] = model_config['rope_parameters']['sliding_attention']['rope_theta'] ** -(torch.arange(0, 1, 2 / model_config['head_dim'], dtype=torch.float32))
# full attention - partial rope
freq['full_attention'] = model_config['rope_parameters']['full_attention']['rope_theta'] ** -(torch.arange(0, 1, 2 / model_config['global_head_dim'], dtype=torch.float32))
freq['full_attention'][int(model_config['rope_parameters']['full_attention']['partial_rotary_factor'] * len(freq['full_attention'])) : ] = 0
def rotate(x, layer_type, start_idx=0): # x [..., seq, head_dim]; rotate each (t, t + hd/2) pair ccw
head_dim = x.shape[-1]
rot_x = rearrange([-x[..., head_dim // 2 :], x[..., : head_dim // 2]], 'z ... d -> ... (z d)')
pos = (torch.arange(x.shape[-2])[:, None] + start_idx) * torch.cat([freq[layer_type], freq[layer_type]]) # pos * (t, t + hd/2 pairs)
return torch.cos(pos).to(x.dtype) * x + torch.sin(pos).to(x.dtype) * rot_xAttention
Decode vs Encode
In the decode stage, each canvas query attends to [committed history | current canvas]. It is:
Similar to autoregressive prefill in the sense that it writes KV’s for several tokens at a time, autoregressive decode in the sense that it reads a context history of past KV’s
Similar to cross attention in that it attends to KV’s from the encoder, whisper-style, similar to self-attention in that it is non-casual, ViT-style
In the encode stage, the attention behaves like an ordinary causal prefill. It can be thought of as the “verify” pass for speculative decoding with a 256-token draft, with all tokens accepted.
The same exact attention weights is used for both.
This code can be easily modified to support batch size greater than 1, but it does not support mixed encode/decode modes within the same batch. Dynamic per-sequence attention mode is supported in vLLM:
Here, “denoise” means decode, “prefill” means encode on the prompt, and “accept” means encode on a denoised canvas.
For sliding layers, we use the approach in Google DeepMind’s JAX sampler: every canvas token sees all other tokens inside the canvas plus the same prefix of KVs immediately before the canvas. Note that vLLM uses a different symmetric, 2K+1, per-token sliding window.
Gemma-specific details
Some notable differences from a standard LLM attention:
No explicit 1/√d multiplier. The inner product is done directly. This can absorbed into the scalar Q/K normalization weights, which already does a fixed inverse temperature.
Gemma also applies an embedding scale, multiplying token embeddings by √D at the start of the residual stream. I’m not really sure why this is done.
Value normalization. Weightless RMS-norm (i.e., unit norm) is applied to every value vector before it is used or stored in cache
Hybrid architecture. 5 : 1 ratio of sliding window layers to global attention layer
Shared global K/V projection. In full-attention layers,
W_v=W_k. K and V are only different due to post-projection processing: K receives scalar-weight RMS normalization plus RoPE, while V receives unit RMS normalization and no RoPE. This is strange since K/V vectors are typically believed to live in different subspaces (see 1, 2 while K/Q live in the same subspace).
Optimizations and Implementation
One potential optimization is to cache only the shared pre-normalization projection and reconstruct K and V when reading it, potentially reducing global-layer cache storage at the cost of additional computation. Combined with the hybrid global/sliding schedule, this could be attractive in long-context, KV-cache-bandwidth-bound regimes. The readable implementation below keeps separate K and V caches instead.
One implemented optimization is that, on encode, we can skip all computation after the K/V projection in the final layer, skipping most of the final layer and unembed work. Encode passes do not use the model’s final logits; they only need to incur the minimum set of computation required to obtain (and commit) correct K/V states in every layer.
In the current implementation with fixed-size statically-shaped cache, full-attention is only differentiated from sliding window through kv_len, thus behaves exactly like sliding-window attention with a max_tot_tokens-sized window:
class AttentionBlock(torch.nn.Module):
def __init__(self, layer_id):
super().__init__()
self.layer_id = layer_id
self.layer_type = model_config['layer_types'][layer_id]
self.W_q, self.W_k, self.W_o = [sd[f'model.decoder.layers.{layer_id}.self_attn.{item}_proj.weight'] for item in ['q', 'k', 'o']]
self.q_norm, self.k_norm = [sd[f'model.decoder.layers.{layer_id}.self_attn.{item}_norm.weight'] for item in ['q', 'k']]
self.pre_norm = sd[f'model.decoder.layers.{layer_id}.input_layernorm.weight']
self.post_norm = sd[f'model.decoder.layers.{layer_id}.post_attention_layernorm.weight']
self.q_heads = model_config['num_attention_heads']
if self.layer_type == 'full_attention':
self.W_v = self.W_k # global layers share K = V
self.kv_heads = model_config['num_global_key_value_heads']
self.head_dim = model_config['global_head_dim']
self.kv_len = max_tot_tokens
else:
assert(self.layer_type == 'sliding_attention')
self.W_v = sd[f'model.decoder.layers.{layer_id}.self_attn.v_proj.weight']
self.kv_heads = model_config['num_key_value_heads']
self.head_dim = model_config['head_dim']
self.kv_len = model_config['sliding_window']
self.k_cache = torch.empty(self.kv_heads, self.kv_len, self.head_dim) # use statically shaped KV buffer
self.v_cache = torch.empty(self.kv_heads, self.kv_len, self.head_dim) # KV's flow into here from left to right, FIFO, latest element is rightmost
def forward(self, x, pos_idx, mode):
assert mode in ['encode', 'decode']
L, D = x.shape
resid_x = x.clone()
x = rms(x, w=self.pre_norm)
q, k, v = x @ self.W_q.T, x @ self.W_k.T, x @ self.W_v.T
q, k, v = [rearrange(z, 'l (n h) -> n l h', h = self.head_dim) for z in [q, k, v]]
q, k = rms(q, self.q_norm), rms(k, self.k_norm) # QK-norm per head: weight is [head_dim], normalize over each head's dims
v = rms(v, 1) # v_norm: weightless, no rope
q, k = rotate(q, self.layer_type, pos_idx), rotate(k, self.layer_type, pos_idx) # absolute positions pos_idx .. pos_idx + L
k = torch.concat([self.k_cache[:, : pos_idx, :], k], axis=1) # attend to [committed history | current block]
v = torch.concat([self.v_cache[:, : pos_idx, :], v], axis=1) # Note python automatically clips on the left to 0, on the right to shape[1] = kv_len
if mode == "encode": # difference #1: writes/updates the kv cache
# Actually, we don't have to put this in a branch, can also just do this on decode too, ok since we'll override with an encode at the end anyways
self.k_cache[:, : pos_idx + L, :] = k[:, -self.kv_len :, :] # automatically clips
self.v_cache[:, : pos_idx + L, :] = v[:, -self.kv_len :, :]
if self.layer_id == model_config['num_hidden_layers'] - 1:
return # encode optimization: notice we don't need to do the remaining computation after this
q = rearrange(q, '(n gqa) qt h -> n gqa qt h', gqa = self.q_heads // self.kv_heads) # Fold GQA into an outer dim
scores = einsum(q, k, 'n gqa qt h, n kt h -> n gqa qt kt').float() # no divide by sqrt(head dim); softmax in fp32
if mode == "encode":
scores += torch.triu(torch.full(scores.shape, -torch.inf), diagonal = scores.shape[-1] - scores.shape[-2] + 1) # this applies a mask that looks like R2 in the vllm figure
scores = torch.exp(scores - torch.amax(scores, axis=-1, keepdims=True))
scores = (scores / torch.sum(scores, axis=-1, keepdims=True)).to(x.dtype)
x = einsum(scores, v, 'n gqa qt kt, n kt h -> n gqa qt h')
x = rearrange(x, 'n gqa qt h -> qt (n gqa h)')
res = x @ self.W_o.T
res = rms(res, self.post_norm)
return res + resid_xMOE
DiffusionGemma routes each token to 8 of 128 experts + one larger shared expert per forward pass.
Each expert is a standard gated GELU MLP:
HuggingFace stores the routed experts’ gate and up projections as a single matrix, which we explicitly split in the MLP module.
MOE Forward
operates independently across the token axis. Each (of B × L many tokens) is routed to its own 8 experts. We could do this naively by looping over each token, but a better approach is to iterate over the experts instead: gather all tokens routed to each expert, do a single forward pass, then send results back.
Across multiple devices / for expert parallelism, this is usually implemented with an AllToAll-style collective rather than the slow explicit Python loop we have here (see the JAX Scaling Book).
Routing Score
DiffusionGemma performs routing in the standard way. Each expert receives a score proportional to
where the expert keys are the rows of W_router. The top 8 scores are renormalized and multiplied by learned per-expert scales.
RMSNorm Fusion
There is an optimization we can do here to save a couple RMS-norm computations. Notice that we apply a pre-norm 1 (shared expert) + 8 (routed experts) times at the start of each MLP. Notice RMSNorm can be factored into two operations
, unit-normalize then multiply by w.
Instead of doing this naively, we can instead apply a single weightless RMS_1(x) at the start, and absorb each expert’s learned RMSNorm weight w into the columns of that expert’s gate and up projection matrix.
MLP and mixture-of-experts block
class MLP(torch.nn.Module):
def __init__(self, layer_id, expert_num, conditioning_mlp : bool = False):
super().__init__()
if(conditioning_mlp):
self.pre_norm = sd['model.decoder.self_conditioning.pre_norm.weight']
self.W_up = sd['model.decoder.self_conditioning.up_proj.weight']
self.W_gate = sd['model.decoder.self_conditioning.gate_proj.weight']
self.W_down = sd['model.decoder.self_conditioning.down_proj.weight']
return
if expert_num is None:
self.pre_norm = sd[f'model.decoder.layers.{layer_id}.pre_feedforward_layernorm.weight']
self.W_up, self.W_gate, self.W_down = [sd[f'model.decoder.layers.{layer_id}.mlp.{item}_proj.weight'] for item in ['up', 'gate', 'down']]
else:
self.pre_norm = sd[f'model.decoder.layers.{layer_id}.pre_feedforward_layernorm_2.weight']
self.W_gate, self.W_up = rearrange(sd[f'model.decoder.layers.{layer_id}.experts.gate_up_proj'][expert_num], '(z intermed) D -> z intermed D', z=2)
self.W_down = sd[f'model.decoder.layers.{layer_id}.experts.down_proj'][expert_num]
def forward(self, x): # every MLP has a prenorm. Also, note since the MLP is a compsition of functions that maps the 0 vector to itself, the entire MLP also maps all 0s to all 0s, this is relevant for the self-conditioning MLP
L, D = x.shape
x = rms(x, self.pre_norm)
a = x @ self.W_up.T
b = F.gelu(x @ self.W_gate.T)
res = (a * b) @ self.W_down.T
assert(res.shape == (L, D)) # share expert, routed expert, and self-conditioning MLP all map (_, D) -> (_, D)
return res
class MOEBlock(torch.nn.Module):
def __init__(self, layer_id):
super().__init__()
self.k_experts = model_config['top_k_experts']
self.num_experts = model_config['num_experts']
self.W_router = sd[f'model.decoder.layers.{layer_id}.router.proj.weight']
self.expert_scale = sd[f'model.decoder.layers.{layer_id}.router.per_expert_scale']
self.scale = sd[f'model.decoder.layers.{layer_id}.router.scale']
self.experts = [MLP(layer_id, e) for e in range(self.num_experts)]
self.shared_expert = MLP(layer_id, None)
self.post_norm_1 = sd[f'model.decoder.layers.{layer_id}.post_feedforward_layernorm_1.weight'] # applied on shared expert output
self.post_norm_2 = sd[f'model.decoder.layers.{layer_id}.post_feedforward_layernorm_2.weight'] # applied on summed contribution from experts
self.post_norm = sd[f'model.decoder.layers.{layer_id}.post_feedforward_layernorm.weight'] # the sum h1 (shared ) + h2 (routed sum), before residual add
def forward(self, x):
L, D = x.shape
resid_x = x.clone()
route_x = rms(x) * self.scale / (D ** 0.5)
expert_scores = F.softmax((route_x @ self.W_router.T).float(), dim=-1) # (L, num_experts)
top_k_scores, top_k_idx = torch.topk(expert_scores, self.k_experts, dim=-1) # (L, k_experts), (L, k_experts)
top_k_scores = (top_k_scores / torch.sum(top_k_scores, dim=-1, keepdim=True) * self.expert_scale[top_k_idx]).to(x.dtype)
res = rms(self.shared_expert(x), self.post_norm_1) # h1: dense branch, shape (L, D)
moe_out = torch.zeros_like(x)
for id, expert in zip(range(self.num_experts), self.experts):
mask = torch.any(top_k_idx == id, dim = -1) # boolean mask (L, ) which tokens routed to expert_id
mult = top_k_scores[mask][top_k_idx[mask] == id] # shape (L', ) where L' <= L is the number of tokens expert_id routed to
moe_out[mask] += mult[:, None] * expert(x[mask]) # (L', D) += (L', 1) * (L', D)
res = res + rms(moe_out, self.post_norm_2) # h2 normed once, then h1 + h2
return rms(res, self.post_norm) + resid_x
Putting Them together
We now combine all our previous components into a single module. Unlike Llama, DiffusionGemma does a couple things different beyond merely interleaving attention and MoE blocks.
Logit Softcapping
Logit softcapping applies the following function to the final logits:
where z is an uncapped logit and c is the softcap value.
For |z| ≪ c, tanh(z/c) ≈ z/c, so small logits are essentially unchanged. As z → ±∞, z_capped → ± c, so large logits are “soft-capped”.
This is a smooth alternative to torch.clip or torch.clamp - unlike a hard cap, it remains differentiable everywhere.
The standard tanh function is bounded between -1 and 1; multiplying by c changes the bounds to [−c,c].
Note this is part of the model, and not the sampling process — not to be confused with the temperature schedule, which is applied later.
Denoise vs Prefill
The following are function signatures of encode vs decode:
# this commits canvas, (last time we) write KV cache
def enc(self, pos_idx, logits) -> None:
self._forward(pos_idx, logits, 0, mode="encode")
# this maps (canvas_i, canvas_prob_i) -> (canvas_prob_i+1)
def dec(self, pos_idx, logits, logit_probs) -> torch.tensor:
return self._forward(pos_idx, logits, logit_probs, mode="decode")On encode, the embedded tokens are directly passed into the first layer.
In decode mode, the embedded tokens additionally incorporates a self-conditioning signal from the previous denoising step. Let P∈ℝ^(L×V) be the previous step’s post-softmax probability distribution over vocabs at each of L canvas positions, let E∈ℝ^(V×D) be the tied embedding table. Then, P @ E computes the expected token embeddings from the previous iteration, L independent convex sums with the weights stored in P. This is then passed through a small conditioning MLP, added to the canvas token embedding, and RMSNormed. Self-conditioning is helpful because for example,
Knowing how confident the previous pass was in this token can inform us about how confident we should be in this pass
During sampling, the highest entropy, least confident tokens are replaced by a random token ID. Thus, self-conditioning allows the model to deduce which tokens are effectively
[MASK]tokens by comparing the probability distribution input to the token input. 2Possibly gives a way for gradients to flow across denoising steps
On the first decode step, the self-conditioning input is all 0s. Note that since the self-conditioning MLP is composed of functions that map all 0s to all 0s, the added contribution from self-conditioning to the input is 0 - this is equivalent to skipping self-conditioning for that step (no previous “self”).
Layer Scalar
Notice each block of computation in DiffusionGemma is additive to the residual stream:
where F is the core computation in that block, and RMSNorm(x) is that block’s pre-norm (both attention and MOE have this).
Thus due to the skip connections, the residual stream remains unnormalized through all 30 layers. This is perfectly fine if the activations truly live in ℝ^n, but unfortunately, they are bound to the subset of representation values of its datatype, which has a fixed dynamic range - upper and lower bounds of the feasible set, and anything outside will over/underflow to infinities. To ensure this doesn’t happen, quantization methods typically multiply by a scalar S before casting, for example:
where
to ensure values lie within the dynamic range [−V_max, V_max].
A similar worry arises for the unnormalized residual stream, the magnitude of the activations may gradually grow through the layers (note however, the inputs to blocks remain well-conditioned due to the RMSNorm). To improve training stability, avoid overflow, and enable stable low-precision inference, DiffusionGemma adds a layer scalar which rescales the completed residual stream by a learned scalar s_l at the end of every layer:
Printing out the layer scalars’ actual values, we see that they are all < 1, which makes sense according to our interpretation. Interestingly, it is a learned parameter — the model could learn to adaptively regulate the magnitude of its own residual stream throughout training.
Notice the RMSNorm is invariant to multiplying its input by a scalar (RMSNorm(cx) = RMSNorm(x)). Since every block receives input through an RMSNorm, the layer scaler has little benefit on the normalized input presented to each block. However, note it will affect the direction of the input after the first layer.
The released Hugging Face checkpoint stores separate encoder and decoder layer-scalar entries,
sd[f"model.encoder.language_model.layers.{i}.layer_scalar"]
and
sd[f"model.decoder.layers.{i}.layer_scalar"]but their values are actually identical. In the code below, we have a single self.layer_scalar.
Complete DiffusionGemma model
class DiffusionGemma(torch.nn.Module): # Any computation that utilizes parameters passes through here
def __init__(self):
super().__init__()
self.W_embed = sd['model.decoder.embed_tokens.weight'] # also used as the unembedding matrix (tie_word_embeddings = True)
self.attn_blocks = [AttentionBlock(i) for i in range(model_config['num_hidden_layers'])]
self.moe_blocks = [MOEBlock(i) for i in range(model_config['num_hidden_layers'])]
self.layer_scalar = [sd[f'model.encoder.language_model.layers.{i}.layer_scalar'] for i in range(model_config['num_hidden_layers'])]
self.model_norm = sd['model.decoder.norm.weight'] # final / model norm
self.embed_scale = torch.tensor(model_config['hidden_size'] ** 0.5)
self.conditioning_MLP = MLP(layer_id=None, expert_num=None, conditioning_mlp=True)
def _forward(self, pos_idx, logits, logit_probs, mode): # logit_probs = 0 <=> skip this path (no bias term anywhere)
x = self.W_embed[logits] * self.embed_scale # (L, ) -> (L, D)
# do self conditioning if decode
if mode == "decode":
condition_x = (logit_probs.to(x.dtype) @ self.W_embed) * self.embed_scale # (L, V) x (V, D) -> convex combination of vocab embeddings
condition_x = self.conditioning_MLP(condition_x)
x = rms(x + condition_x)
# pass through all layers
for i, (attn, moe) in enumerate(zip(self.attn_blocks, self.moe_blocks)):
x = attn(x, pos_idx, mode)
if x is None: # last encode layer wrote its KV cache and returned early; nothing else is needed
return None
x = moe(x)
x = x * (self.layer_scalar)[i] # per-layer encoder/decoder scalar
x = rms(x, self.model_norm)
final_logits = (x @ self.W_embed.T).float() # (L, D) x (D, V) -> (L, V)
return torch.tanh(final_logits / model_config['final_logit_softcapping']) * model_config['final_logit_softcapping']
# this commits canvas, (last time we) write KV cache
def enc(self, pos_idx, logits) -> None:
self._forward(pos_idx, logits, 0, mode="encode")
# this maps (canvas_i, canvas_prob_i) -> (canvas_prob_i+1)
def dec(self, pos_idx, logits, logit_probs) -> torch.tensor:
return self._forward(pos_idx, logits, logit_probs, mode="decode")The following figure is a nice summary of what we’ve put together:
Sampling:
After the first encode on the input prompt, every 256 canvas of tokens incurs one encode pass and at most max_denoising_steps decode passes.
After obtaining a vocabulary distribution for every canvas position, sampling is more involved than in an autoregressive model.
We follow the sampling procedure explained here, implemented here. The meanings of the sampling parameters are documented here:
Apply temperature
For decode pass i out of N = max_denoising_steps,
\(t_i = t_{\max} + \frac{i}{N}\left(t_{\min}-t_{\max}\right), \qquad i=0,\ldots,N-1\)Thus the temperature starts at t_max and decreases toward t_min, increasingly sharpening the distribution. In the i-th iteration, we work with the categorical distribution formed by
logits / t_i.Compute entropy
For every canvas position, we compute from its categorical distribution p,
\(H(p) = -\sum_{v=1}^{V} p_v\log p_v\)Used as a measure of how uncertain the model is about the token at this position.
Keep a low-entropy prefix and renoise the rest
Sort position entropies so that
\(H_{(1)} \le H_{(2)} \le \cdots \le H_{(L)}\)Accept the largest prefix ending at k such that
\(\sum_{j=1}^{k-1} H_{(j)} \le \texttt{entropy_bound}\)Then, replace every unaccepted position by a token drawn uniformly from the vocabulary.
Check for early stopping
Convergence is decided when the canvas is confident and stable:
the argmax canvas has remained unchanged for
stability_threshold(equals 1 here) previous canvases, andthe current canvas mean per-position entropy is below
confidence_threshold.
On convergence or after N steps, we commit the latest canvas.
Tokenizer and sampling configuration
tok = Tokenizer.from_file(f"{checkpoint}/tokenizer.json")
prompt = """
What is the meaning of 67?
"""
chat = f"<bos><|turn>user\n{prompt}<turn|>\n<|turn>model\n"
max_denoising_steps = gen_config['max_denoising_steps'] # decoder forward passes per canvas
entropy_bound = gen_config['sampler_config']['entropy_bound'] # see formula in figure
t_max, t_min = gen_config['t_max'], gen_config['t_min'] # linear schedule of temperatures t_max -> t_min across the steps
confidence_threshold = gen_config['confidence_threshold'] # early-stop a canvas when argmax is stable (equal previous argmax canvas) and mean entropy < this
tokens = torch.tensor(tok.encode(chat, add_special_tokens=False).ids)
model = DiffusionGemma()
Lazy Sampling
The usual way a denoising iteration is implemented:
Model forward → retain distribution for self-conditioning → sample and renoise tokens
Instead, we do sampling lazily, passing only the categorical distribution between steps and sampling only once we need it:
Stored distribution → sample and renoise input tokens → model forward → new distribution
The input into the first iteration is just all identical logits, which will generate a random canvas for us.
This is nice since:
The only information we need to pass between steps is the self-conditioning input
Initialization falls out naturally from a uniform categorical distribution, can be folded into the first step
We can think of the decoder input as the previous canvas probs plus a source of randomness (from the sampling). Under this view, preparing the inputs (sampling) lazily is optimal. Additionally, it suggests an alternative interpretation of self-conditioning as the primary rather than auxiliary input:
One way to think of the decoder is that it repeatedly transports our current distribution toward the target distribution, reminiscent of a flow-like process on the probability simplex. So one can think about an alternate DiffusionGemma design that instead starts from a random position in the simplex at initialization, Categorical(probs=torch.rand(L, V)) or Categorical(logits=torch.rand(L, V)), rather than “Categorical(logits=torch.zeros(L, V))”.
Stage Then Commit
We also intentionally give general functions for denoise which generates staged tokens; and commit, which takes in some staged tokens, generates their encoded KVs, and writes them to cache.
denoise(pos_idx) repeatedly calls decode against the already committed cache[:pos_idx] on a fresh canvas at positions [pos_idx : pos_idx + 256], and returns staged tokens sampled from the final canvas of logits. This function has no side effects on the KV cache, thus we can for example denoise multiple canvases at the same position and choose the best one. We can also place the canvas anywhere, as long as pos_idx ≤ len(tokens) - canvas_len, since the correct slice of the KV cache will automatically be read.
commit(l, r, staged_tokens) applies the casual encode and writes (or overwrites) the slice of KV cache from l to r, with the staged tokens’ KVs. This function allows arbitrary-lengthed blocks, thus we can for example commit only a confident prefix of the current canvas. We can also commit tokens at any positions, but if an earlier region is overwritten, all later tokens are invalidated because their cached states depended on the old prefix.
This design is to make alternative schedules easier to study, including overlapping canvases, revising an earlier block, selecting among several denoised candidates, or advancing only part of a canvas.
Denoise, stage, commit, and generate
def denoise(pos_idx): # returns staged_tokens
assert pos_idx + canvas_len <= len(tokens) # must be length canvas_len (what if it wasn't fixed? analyze how casual the self attention is)
plotting_data = []
t = t_max
t_step = (t_min - t_max) / max_denoising_steps
last_canvas = Categorical(logits = torch.ones((canvas_len, V)))
for step in trange(max_denoising_steps):
# renoise last_canvas
sH, sidx = last_canvas.entropy().sort(-1)
accepted = torch.zeros_like(sH, dtype=torch.bool).scatter(-1, sidx, sH.cumsum(-1) - sH <= entropy_bound)
last_canvas_noised = torch.where(accepted, last_canvas.sample(), torch.randint(0, V, (canvas_len,)))
# pass in the noised tokens, but un-noised normalized probs (all-zero probs on step 0: the conditioning path maps 0 to 0)
canvas = model.dec(pos_idx, last_canvas_noised, last_canvas.probs if step != 0 else torch.zeros(canvas_len, V)) / t; plotting_data.append(canvas.detach().cpu())
canvas = Categorical(logits = canvas)
if torch.mean(canvas.entropy()) < confidence_threshold and (canvas.logits.argmax(dim=-1) == last_canvas.logits.argmax(dim=-1)).all():
return canvas.sample(), plotting_data
t += t_step
last_canvas = canvas
assert False, f"Denoising not finished after {max_denoising_steps} steps"
def commit(l, r, staged_tokens):
global tokens
assert len(staged_tokens) == r - l + 1
if(r + 1 < len(tokens)): # commits staged_tokens
print(f"Invalidating {len(tokens) - (r+1)} tokens")
tokens = tokens[:r+1]
tokens[l:] = staged_tokens
model.enc(l, staged_tokens)
def new_canvas():
global tokens
if len(tokens) + canvas_len > max_tot_tokens:
return False
nxt = torch.randint(V, (canvas_len,))
tokens = torch.concat([tokens, nxt])
return True
commit(0, len(tokens) - 1, tokens) # prefill: encode the chat prompt into the KV cache
pos_idx = len(tokens)
plotting_data = []
while new_canvas():
staged_tokens, cur_data = denoise(pos_idx)
commit(pos_idx, pos_idx + canvas_len - 1, staged_tokens)
pos_idx = len(tokens)
print('=' * 50)
print('Final Canvas')
print(tok.decode(tokens.tolist()))
print('=' * 50)
plotting_data.append(cur_data)
if torch.isin(staged_tokens, torch.tensor(gen_config['eos_token_id'])).any():
breakOutput:
...
==================================================
Final Canvas
user
What is the meaning of 67?
model
thought
The meaning of **67** depends entirely on the context in which it is used (science, mathematics, pop culture, etc.). Here are common interpretations:
### 1. Mathematics
* **Prime Number:** 67 is a prime number, meaning it can only be divided by 1 and itself.
* **Lucky Prime:** It is considered a "lucky prime."
* **Sum of Primes:** It is the sum of five consecutive prime numbers.
### 2. Science and Astronomy
* **Atomic Number:** 67 is the atomic number of **Holmium (Ho)**, a rare earth element belonging to the lanthanide series.
* **Astronomy:** Messier object 67 (M67) is an open star cluster in the constellation of Virgo.
### 3. Culture and Slang
* **The "67" Connection:** In the UK, "67" is a well-known drill music group from Brixton, London.
* **Age:** In many countries, 67 is considered the standard age for full retirement eligibility or social security.
### 4. Numerology and Spirituality
* In numerology, the number 67 is often associated with combining the energies of **6** (home, stability, and responsibility) and **7** (spirituality, intuition, and inner wisdom). It is often interpreted as a sign of building practical foundations through spiritual growth.
### 5. Other Uses
* **Country Code:** +67 is not a complete country code, but codes starting with +67 are used in various regions (like +670 for East Timor or +679 for American Samoa).
**Is there a specific area (like a dream, a song, or a math problem) where you saw this number?** Providing more context can help me give you a more specific answer.
==================================================Preliminary Analysis
Using the saved plotting_data from denoise, we can make some figures of the denoising process over time:
Note the cells are laid out in reading order.
Without being explicitly trained to do so, it appears the model is approximately finalizing canvas tokens from left to right. Plotting two more figures:
Further confirms our suspicions.
At every step, the model seems to mostly focus its efforts on figuring out the next causal rolling window of tokens — it “thinks” casually. Is this a general behavior the model has learned? Let’s try a different prompt, solving a magic square puzzle:
Canvas 2
Canvas 3, More visualizations available here.
Interestingly, the prose tokens seem to exhibit the same causal pattern, but the magic square tokens do not. This makes intuitive sense. When humans solve these puzzles, the next value that’s easiest to deduce is not necessarily the next one in autoregressive order. This motivates perhaps a better explanation for how DiffusionGemma “thinks” during denoising — the “easiest”, lowest-entropy tokens are fixed / finalized / accepted first, which in turn unlocks / makes other tokens “easier” for the next denoising step. 3
Note this mirrors what we do during sampling: all but a prefix of lowest entropy tokens is renoised. For prose and chain-of-thought reasoning, the easiest tokens to resolve next often coincide with the earliest tokens in autoregressive order.
However for both prompts, once a token is accepted (i.e. low-entropy and not renoised), it’s extremely unlikely to be renoised in the future:
For 67 Prompt
For Magic Square Prompt
Looking at the top right graph, almost no tokens are “re-masked” for both prompts. Google claims that one of the benefits of DiffusionGemma’s uniform state diffusion over traditional masked diffusion (e.g., LLaDA, Dream) is error correction via re-noising, but from our analysis — how often does this actually happen in practice?
References
O’Donoghue, Brendan, and Sebastian Flennerhag. “DiffusionGemma: 4x faster text generation.” Google, 2026.
Ballantyne, Ian, and Omar Sanseviero. “DiffusionGemma: The Developer Guide.” Google Developers Blog, 2026.
Google AI for Developers. “DiffusionGemma model overview.” 2026.
Google AI for Developers. “Diffusion in Text Generation Explained.” 2026.
Google DeepMind. “DiffusionGemma 26B-A4B-IT model card.” Hugging Face, 2026.
Google DeepMind. “DiffusionGemma reference sampler.” Gemma repository, 2026.
vLLM Team and Google DeepMind Team. “DiffusionGemma: The First Diffusion LLM (dLLM) Natively Supported in vLLM.” vLLM Blog, 2026.
Grootendorst, Maarten. “A Visual Guide to DiffusionGemma.” Exploring Language Models, 2026.
Huang, Austin, Suraj Subramanian, Jonathan Sum, Khalid Almubarak, and Stella Biderman. “The Annotated Transformer.” Harvard NLP, 2022. Original version by Alexander M. Rush.
Peng, Bowen, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. “YaRN: Efficient Context Window Extension of Large Language Models.” arXiv:2309.00071, 2023.
Su, Jianlin, et al. “RoFormer: Enhanced Transformer with Rotary Position Embedding.” arXiv:2104.09864, 2021.
Kayyam, Ali, Anusha Madan Gopal, and M. Anthony Lewis. “Do Transformers Need Three Projections? Systematic Study of QKV Variants.” arXiv:2606.04032, 2026.
Elhage, Nelson, et al. “A Mathematical Framework for Transformer Circuits.” Transformer Circuits Thread, 2021.
Austin, Jacob, et al. “Sharded Matrices and How to Multiply Them.” How To Scale Your Model, 2025.
Roos, Daan, et al. “Categorical Flow Maps.” arXiv:2602.12233, 2026.
Nie, Shen, et al. “Large Language Diffusion Models.” arXiv:2502.09992, 2025.
Ye, Jiacheng, et al. “Dream 7B: Diffusion Large Language Models.” arXiv:2508.15487, 2025.
Its main comparative advantage is higher arithmetic intensity from being able to decode a 256-token canvas in parallel, amortizing KV cache and model weight load across positions, which shines for local, low-concurrency, high-interactivity workloads
However, unlike traditional Masked Language Diffusion models, DiffusionGemma can still mask out a previously un-masked token if it’s no longer confident in it (i.e., the same token position is now assigned high entropy).
Reminiscent of Kahn’s algorithm, this Codeforces problem, and successive interference cancellation