DiffusionGemma Technical Report — part 2/6

Cleaned-text code points [15828, 30937). Bibliography omitted here; source.md includes it.

The forward process dictates the transition from clean text tokens to uniformly distributed tokens. Conditioned on a fixed clean starting canvas x 0 x_{0} , the forward transition probability path factorizes independently over each token coordinate i i : ℙ ⁡ ( X t = x t ∣ X 0 = x 0 ) = ∏ i = 1 C [ κ t ​ δ ​ ( x t i , x 0 i ) + ( 1 − κ t ) ​ 1 V ] , \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], (1) where κ t ∈ [ 0 , 1 ] \kappa_{t}\in[0,1] is a smoothly varying, monotonically decreasing noise schedule from κ 0 = 1 \kappa_{0}=1 to κ 1 = 0 \kappa_{1}=0 , and δ \delta represents the Kronecker delta. In practical terms, as time t t progresses toward 1 1 , each token is increasingly likely to be replaced by a token sampled uniformly at random from the vocabulary ( Hoogeboom et al., 2021 ; Austin et al., 2021 ) . We use this closed-form forward process to generate training data from clean text.

#### 2.1.2 The Backward Denoising Process

To reconstruct data from noise, we learn a reverse process to undo this categorical corruption. Discrete flow matching theory shows that to perfectly reverse the forward trajectory, we need to infer the conditional distribution of the original uncorrupted tokens given a corrupted state ( Campbell et al., 2024 ; Gat et al., 2024 ; Ou et al., 2025 ) . For a given realization X t = x t X_{t}=x_{t} , stepping backward in time by a small increment Δ ​ t \Delta t is governed by some transition mapping, which we denote Step \operatorname{Step} , which outputs the probability distribution for the next intermediate step 1 1 1 As a typical example, the function Step \operatorname{Step} can be instantiated as a discrete Euler update step introduced in Gat et al. (2024) , in which ℙ ⁡ ( X t − Δ ​ t i = v ∣ X t = x t ) ≈ δ ⁡ ( v , x t i ) − Δ ​ t ​ κ ˙ t 1 − κ t ​ [ ℙ ⁡ ( X 0 i = v ∣ X t = x t ) − δ ⁡ ( v , x t i ) ] \mathbb{P}(X_{t-\Delta t}^{i}=v\mid X_{t}=x_{t})\approx\delta(v,x_{t}^{i})-\Delta t\frac{\dot{\kappa}_{t}}{1-\kappa_{t}}\left[\mathbb{P}(X_{0}^{i}=v\mid X_{t}=x_{t})-\delta(v,x_{t}^{i})\right] with κ t ∈ [ 0 , 1 ] \kappa_{t}\in[0,1] . : ℙ ( X t − Δ ​ t = ⋅ ∣ X t = x t ) ≈ Step ( x t , ℙ ( X 0 = ⋅ ∣ X t = x t ) ) . \mathbb{P}(X_{t-\Delta t}=\cdot\mid X_{t}=x_{t})\approx\operatorname{Step}\Big(x_{t},\mathbb{P}(X_{0}=\cdot\mid X_{t}=x_{t})\Big). (2) To compute this update, we need the true posterior distribution of clean tokens given a noisy state, ℙ ⁡ ( X 0 i = v ∣ X t = x t ) \mathbb{P}(X_{0}^{i}=v\mid X_{t}=x_{t}) . We approximate this posterior with a neural network p θ ​ ( v ∣ x t ) p_{\theta}(v\mid x_{t}) . During generation, we use this approximation to sample the next token state: x t − Δ ​ t ∼ Step ( x t , p θ ( ⋅ ∣ x t ) ) x_{t-\Delta t}\sim\operatorname{Step}(x_{t},p_{\theta}(\cdot\mid x_{t})) .

Figure 3 presents a highly stylized example to illustrate this reverse sampling trajectory within a simplified two-token canvas ( C = 2 C=2 ). As time runs backward from t = 1.0 t=1.0 to t = 0.0 t=0.0 , the model smoothly shifts probability mass away from a uniform noise distribution (e.g., around “blue moon”) toward valid data modes. Concurrently, individual sequence coordinates undergo continuous-time jump transitions—visualized by the step-by-step path from “blue moon” at t = 1.0 t=1.0 , to “dark cloud” at t = 0.66 t=0.66 , and finally landing on a high-probability mode like “red sunset” at t = 0.0 t=0.0 . This highlights how parallel dimensions coordinate over time without requiring sequential left-to-right generation.

Because this denoising process is factorized, it assumes conditional independence during individual reverse steps. A newly updated token at position i i is conditioned globally on the current noisy canvas x t x_{t} , but it cannot see the simultaneous sampling choices made at other positions ( j ≠ i j\neq i ). This structural tradeoff can occasionally introduce local inconsistencies or conflicting grammatical predictions. In the next section, we counteract these uncoordinated mechanics directly through our self-conditioning architecture and entropy-bounded sampling strategies.

## 3 The DiffusionGemma Architecture

The DiffusionGemma architecture functions as an encoder-decoder transformer ( Vaswani et al., 2017 ) with shared weights θ \theta . Rather than pretraining a diffusion model from scratch, we initialize our model with the publicly released Gemma 4 26B A4B MoE checkpoint ( Gemma Team et al., 2026 ) . Tying our architecture to an existing AR backbone is a pragmatic choice that trades off absolute generation quality for a substantial reduction in the computational cost of training. This initialization also allows us to inherit the base model’s advanced features—such as its extended context window and native multimodal understanding. Moreover, our experiments demonstrate that our final model weights retain the ability to be sampled from autoregressively (see Table 3 ).

The remainder of this section details how DiffusionGemma uses this architecture to generate text at inference time. We break this process down into three components: the block-autoregressive decoding strategy for handling long sequences (Section 3.1 ), the general framework for denoising individual canvases (Section 3.2 ), and a specific entropy-bounded sampler (Section 3.3 ).

### 3.1 Block-Autoregressive Generation

The discrete flow matching formulation introduced in Section 2 operates on fixed-length sequences. To generate open-ended text, we use a block-AR generation strategy. The model denoises a canvas of 256 tokens at a time, and once a canvas is fully denoised, it is committed to the sequence history, and the model begins denoising the next canvas.

##### KV cache initialization.

As illustrated in Figure 4 , the first step in block-AR generation is to encode the context, h ∈ 𝒱 L h\in\mathcal{V}^{L} , into a KV cache H H . The context includes system instructions and user prompts of maximum token count L L . 2 2 2 In practice, h h also includes interleaved multimedia embeddings which are not part of the token vocabulary. We write the encoding as H = Encoder θ ⁡ ( h ) , H=\operatorname{Encoder}_{\theta}(h), (3) where Encoder θ \operatorname{Encoder}_{\theta} is the transformer forward pass with weights θ \theta and causal attention masking.

##### Canvas denoising conditioned on KV cache.

By cross-attending to the KV cache H H , canvas generation can be conditioned on system instruction, user inputs and past responses. Canvas generation starts from a canvas of uniformly random tokens and is iteratively denoised using the process described in Sections 3.2 and 3.3 . Once a canvas is fully denoised, 3 3 3 A canvas is fully denoised either by reaching t = 0 t=0 or by the adaptive stopping condition described in Section 3.3 . denoted by x ^ 0 \hat{x}_{0} , its keys and values are appended to the KV cache (depicted in red in Figure 4 ): H ← H ⊕ Encoder θ ⁡ ( x ^ 0 ) . H\leftarrow H\oplus\operatorname{Encoder}_{\theta}(\hat{x}_{0}). (4) We repeat this denoising process for subsequent canvases until the model generates a special token marking the end of the model’s turn, after which all canvases x ^ 0 \hat{x}_{0} are concatenated into a final response. Since we use causal attention in the encoder, it lets the encoder update the KV cache by appending only the newest canvas. This constitutes an architectural inversion of standard encoder-decoder models such as BART ( Lewis et al., 2020 ) or T5 ( Raffel et al., 2020 ) . Whereas those models utilize a bidirectional encoder for context and a causal decoder for generation, our approach utilizes a causal encoder for the sequence history and a bidirectional decoder for the diffusion-based canvas generation. This causal encoding prevents early context tokens from attending to the latest canvas, but it eliminates the need to re-encode the growing context from scratch, making our approach scalable to long reasoning generations and yielding a structural blend of diffusion and AR generation. This block-AR strategy was developed by our group in an unpublished June 2023 manuscript; analogous approaches to blockwise sequence modeling and KV caching were independently developed by Wu et al. (2025) , Arriola et al. (2025) , and Deschenaux and Gulcehre (2026a) .

### 3.2 The Denoising Framework at Inference

##### Canvas initialization.

To generate a canvas, we iteratively refine it over a maximum of N N denoising steps, with step size Δ ​ t = 1 / N \Delta t=1/N . Let x t ∈ 𝒱 C x_{t}\in\mathcal{V}^{C} denote the corrupted canvas at step t t . We initialize the process at t = 1 t=1 with a canvas x 1 x_{1} of uniformly random tokens from the vocabulary 𝒱 \mathcal{V} .

##### The decoder forward pass.

At each step t t , we apply the transformer with the shared weights θ \theta as a decoder, written as Decoder θ \operatorname{Decoder}_{\theta} , to predict the probability distribution of the clean tokens. The decoder takes three inputs: the current noisy canvas x t x_{t} , the context KV cache H H , and a continuous self-conditioning signal z t ∈ ℝ C × d z_{t}\in\mathbb{R}^{C\times d} that feeds the model’s previous predictions back into itself ( Chen et al., 2023c ; Strudel et al., 2022 ; Jo et al., 2026 ) . Using bidirectional attention across the canvas tokens and cross-attention to the KV cache, the decoder outputs the unnormalized logits L t L_{t} : L t = Decoder θ ⁡ ( x t , z t , H ) ∈ ℝ C × V . L_{t}=\operatorname{Decoder}_{\theta}(x_{t},z_{t},H)\in\mathbb{R}^{C\times V}. (5)

##### The denoising iteration.

At each iteration, we compute the logits L t L_{t} and evaluate the clean token probabilities p ^ 0 \hat{p}_{0} , update the self-conditioning signal for the next step, and sample the refined canvas: p ^ 0 \displaystyle\hat{p}_{0} = Softmax ⁡ ( L t / τ t ) ∈ [ 0 , 1 ] C × V , \displaystyle=\operatorname{Softmax}\big(L_{t}/\tau_{t}\big)\in[0,1]^{C\times V}, (6) z t − Δ ​ t \displaystyle z_{t-\Delta t} = FFW ⁡ ( p ^ 0 ​ E ) ∈ ℝ C × d , \displaystyle=\operatorname{FFW}(\hat{p}_{0}E)\in\mathbb{R}^{C\times d}, x t − Δ ​ t \displaystyle x_{t-\Delta t} ∼ Step ⁡ ( x t , p ^ 0 ) . \displaystyle\sim\operatorname{Step}(x_{t},\hat{p}_{0}). Here, E ∈ ℝ V × d E\in\mathbb{R}^{V\times d} is the token embedding matrix and FFW \operatorname{FFW} is a standard feedforward network. The time-dependent temperature τ t > 0 \tau_{t}>0 can sharpen the model’s predictions before we calculate the final transition probabilities. A transition mapping Step \operatorname{Step} , as described in Equation ( 2 ), computes the categorical distribution used to sample the updated canvas x t − Δ ​ t ∈ 𝒱 C x_{t-\Delta t}\in\mathcal{V}^{C} , which, along with the new self-conditioning signal z t − Δ ​ t z_{t-\Delta t} , is then fed into the next denoising step. Section 3.3 details our specific choices for transition mapping Step \operatorname{Step} and the temperature τ t \tau_{t} .

##### Multinomial diffusion.

Our denoising iteration (Equation 6 ) uses multinomial (or uniform) diffusion rather than masked diffusion ( Hoogeboom et al., 2021 ; Austin et al., 2021 ) . Because all tokens can transition between one another, the model can continuously correct its own errors: tokens accepted during earlier denoising steps ( t ′ > t t^{\prime}>t ) within the current canvas can still be revised. However, we note that tokens from previously generated canvases are permanently frozen.

##### Remark on interpretability.

Recent interpretability analyses demonstrate that while the self-conditioning signal z t z_{t} introduces a continuous latent space injection into the sequence, these intermediate self-conditioning vectors map robustly to an interpretable token bottleneck, preserving the model’s algorithmic transparency ( Engels et al., 2026 ; Asaria et al., 2026 ) .

### 3.3 The DiffusionGemma Sampler

In contrast to AR generation, which relies on rigid heuristics like temperature, top- p p , and top- k k , diffusion modeling introduces a vastly richer sampling design space that includes discrete predictor-corrector mechanisms and multi-step solvers ( Lezama et al., 2023 ; Deschenaux et al., 2026 ; Liu et al., 2026a ; Yao et al., 2026 ; Ren et al., 2025 ) . Framing text generation as an iterative temporal process allows us to decouple the inference-time algorithm from the underlying architecture, granting fine-grained control over the trade-off between computational cost and generation quality.

While we described the overall approach of canvas denoising in Section 3.2 , this subsection details a specific instantiation: the entropy-bounded sampler ( Ben-Hamu et al., 2026 ) with temperature annealing and adaptive stopping (Algorithm 1 ). While this is our default and recommended sampler, DiffusionGemma is modular and is not strictly bound to this exact sampling configuration for high-quality inference.

##### Entropy-bounded token refinement.

After producing the marginal distributions p θ ​ ( x 0 ∣ x t , z t , H ) p_{\theta}(x_{0}\mid x_{t},z_{t},H) with the denoiser and applying temperature scaling (detailed below), tokens are sampled from the resulting probability distribution. We employ an entropy-bounded sampler ( Ben-Hamu et al., 2026 ) , whereby tokens are accepted in rank order from lowest to highest entropy (similar to the MaskGIT decoding scheme by Chang et al. (2022) ), ensuring that their mutual information bound remains strictly below a predefined error tolerance threshold ( b = 0.1 b=0.1 ). Once the threshold is attained, all other tokens are renoised uniformly at random, maintaining these uncommitted positions as a uniform prior to force local exploration during the next forward pass.

##### Temperature annealing.

To balance rate of convergence and linguistic diversity, token probabilities are artificially sharpened via tempering. A temperature τ t < 1 \tau_{t}<1 is annealed linearly from an initial value of τ max = 0.8 \tau_{\max}=0.8 down to τ min = 0.4 \tau_{\min}=0.4 across the fractional denoising timescale t ∈ [ 0 , 1 ] t\in[0,1] . While static temperature adjustments are traditionally used to truncate the unreliable tail of token distributions and prevent text degeneration, this dynamic annealing ensures that the model explores diverse token possibilities in early, highly-noised states and aggressively commits to high-confidence sequences as the semantic structure crystallizes ( Zhang et al., 2024 ) .

##### Adaptive stopping heuristic.

To optimize inference efficiency, the sampler dynamically halts the denoising process based on the model’s step-wise uncertainty, strictly capping iterations at N N . This early termination is triggered when two conditions are simultaneously satisfied: topsep=0pt Confident predictions: The mean predictive entropy across the entire canvas falls below a predefined threshold ( e stop = 0.005 e_{\mathrm{stop}}=0.005 ).

topsep=0pt Stable predictions: The deterministic sequence predictions (i.e., the most likely tokens) from two consecutive denoising steps are identical.