TL;DR

Google DeepMind's recent model DiffusionGemma (DG) generates text via diffusion, meaning many diffusion steps happen before generating the final output. In particular, these diffusion steps carry vectors in addition to tokens. If we cannot interpret these tokens and vectors, the model has significant opaque serial depth, potentially harming monitorability. Recently, Engels et al. found that DG nevertheless maintains high monitorability, for instance by showing that projecting the distribution to its top-k items largely retains performance. We strengthen these results by showing that this performance degradation is largely a sampler artifact and good performance can be maintained with only the top item, supporting the case for high monitorability. Still, we also find some rare case studies where the distribution vector is load-bearing computationally, i.e. where top-1 projection would be detrimental. However even in these cases, it just encodes superposition, remaining interpretable.

Apart from model behavior, we also examined how interpretability techniques carry over to DiffusionGemma, including probes, steering, and J-lens. We find that performance is largely retained. This is a positive update on the interpretability of diffusion models that are derived from text-pretrained LLMs (an efficient training method more likely to be deployed), but might not apply for more general paradigms.

Overall, this supports the paper's conclusion that DiffusionGemma remains highly monitorable, while nevertheless showing that there are cases where models can learn to use vector-valued information.

Github

Introduction

Large language models arrive at answers to complex questions through chains-of-thought. These chains are generated token-by-token in a sequence of autoregressive steps. This gives fairly large visibility into the model's process of arriving at the answer, and has been the main pillar for monitorability in recent years (see e.g. Guan et al.). In contrast, latent reasoning models (see here for a survey) pass vectors between diffusion steps, a priori destroying monitorability.

DiffusionGemma is a particular model whose architecture allows latent reasoning, by passing a vector encoding a probability distribution between steps.

Engels et al. recently surveyed DiffusionGemma's behavior, identifying cases where the model will use its distribution to express uncertainty about positioning of tokens. Here, we explore whether there are instances where the model uses the distribution computationally.

Background on DiffusionGemma

DiffusionGemma is a text-generation model, based on the Gemma architecture. Generation proceeds as follows. Let be the prompt, and let be the noise-initialized token canvas comprising positions. Let denote a single forward pass through the transformer stack (a finetune of Gemma). Roughly, the final output then is obtained via

where is the number of diffusion steps. Here, represents the distribution passed between diffusion steps. While attends causally to , it attends bidirectionally to itself. also functions as a confidence score for any token : unless a cumulative entropy bound is crossed, gets replaced with a uniformly chosen random token at every step . The distribution at every step is furthermore sharpened by a temperature rescaling the logits, which itself follows a linear cooling schedule over steps.

For a visual and more detailed introduction to DiffusionGemma, see this post or the technical report.

Thus, generation in DiffusionGemma mainly differs in the following ways:

1. Generation happens as reverse diffusion, not as token-by-token autoregression. This is reminiscent of a looped transformer (see e.g. Giannou et al., 2023).
2. Attention is bidirectional.
3. Between every diffusion step, not only the current text output is sampled, but the output distributions across all positions are also passed to the next diffusion step .

The last point is especially interesting, since it passes the vector between diffusion steps, in addition to the token canvas lacking the axis. A priori, this allows the model to transport drastically more information between diffusion steps in an illegible way, hindering monitorability.

Performance degradation from top-k truncation largely is a sampler artifact

Investigating this risk, Engels et al. surprisingly found that DiffusionGemma scores similar monitorability to Gemma. However, when truncating to just its top-k entries, performance significantly dropped. Thus somehow the information in seemed to have been essential, conflicting with the results of high monitorability.

However, when replicating their experiments, we observed that the model will often fall into a "degenerate loop", outputting the same token over and over, never reaching the final answer. We found that adopting a gentler sampler together with a higher number of diffusion steps largely prevents this failure mode, suggesting that in fact the distribution is not essential to solve these problems. Rather, the truncation might have put the model outside of its training distribution. Still, this does not rule out that there is a functional necessity in other tasks that the paper had not investigated.

A gentler sampler prevents the degenerate loop that caused performance degradation on top-k truncating the distributional state observed in Engels et al.. Specifically, we use more diffusion steps ( = 48 → 96), a wider temperature range (0.8–0.4 → 1.0–0.5), and a lower entropy bound (0.1 → 0.02), all of which contribute to stabilization. Soft denotes the untruncated distribution .

A case study for using the distribution computationally: letter arithmetic

We next investigated whether there may still be some other tasks where in fact is essential. We looked for tasks where the model plausibly would hold several "hypotheses" in superposition, and do computation on them in parallel. We were able to construct a case study with letter arithmetic, where we ask the model to shift a letter 3 down in the alphabet:

Pick any uppercase letter (the operand ) between A and W, write it, then write the letter (the target ) 3 (the increment ) positions later in the alphabet.

DG answers these correctly, consistently choosing its own natural operand and target (e.g. Letters: G, J).

We here find that DG holds a probability distribution on the starting letter. Note however that there is a natural mechanism to implement this computation: embedding letters as points on a circle and rotating by an angle of 3/26. We were not able to easily find similar parallel computation in other settings.

To intervene on this computation, we capture the canvas at some intermediate denoising step , add probability mass on a different operand letter at the operand position. We choose the injection such that is still not the top logit. This is important, because we would like to measure what DG does to states that are not the most probable ones.

Then, we measure the response to that perturbation , where indicates an average over seeds.

Perturbing subleading tokens predominantly triggers a response at corresponding target tokens. x-axis: Injection token, y-axis: possible target tokens. Note the negative response at J, the target of the model's natural response Letters: G, J.

For illustration, here is a single intervention in full:

Example of one intervention. a subleading injection on H (the leader G keeps rank 1) switches the answer slot from J to K=H+3 one step later. Top shows the prompt, and the model's "natural" output. The lower part shows previous and next step vocabulary distributions (minor columns), for the baseline run and after intervention, respectively. Note how the "natural" output J loses mass due to the competing injected mass on H (purple).

Parallel computation

This simple response behavior begs the question whether it is possible to perturb source letters simultaneously (again, keeping them subleading) and see whether the corresponding target images respond. This would be interesting, since it would suggest that the model can carry out simultaneous computation, a key advantage that latent reasoning models have on paper.

To this end, we measured the response on the image of the injected operands and a baseline averaged over images of alternative possible operands, respectively, and measure the effect . We then report the fraction of simultaneous injections that has positive response on all targets, as well as a null hypothesis if tokens would respond randomly.

DiffusionGemma responds to multiple injections in parallel, more than expected by chance. Letters task with : for each injected operand we define the per-member effect , where and is the pool of possible operands, with the response as defined above. An injection set counts as responding if . Dashed: chance level . Error bars are 95% CIs over possible injection operand sets.

We observe that the average target response exceeds the non-target response for all injected members more often than chance, averaged over possible injection sets. This behavior is plausible considering the computation in question: a shift is easily implemented by a linear operation in representation space. Therefore by linearity, superpositions of letters will be transported to superpositions of responses. Overall, this leaves an interpretable picture.

In contrast, we found that for multiplication (e.g. letter × 3) or taking the absolute value, there was no significant effect.

Autonomous computational usage of

In the previous letter arithmetic task, DiffusionGemma did respond to modifications of the distributional state in the way we expected. While suggestive, it is unclear if the model also would make use of autonomously, i.e. without interventions.

To investigate this, we searched for a non-trivial (i.e., not just a binary choice) task where DG will use a hypothesis subleading in its distribution. An instance we found is a word-level palindrome task. We ask the model Please write a word-level palindrome. DG maintains two competing completions of the same canvas: a literal seasonal phrase (All leaves fall when leaves fall all.) and an idiom (All for one and one for all.). We measure how the probability mass of these two typical outcomes, summed over positions, varies over the course of diffusion.

Ablating a nascent contender in the distribution can prevent takeover. Probability mass of the two completions at the contested slots, summed over all tokens of each completion (black: idiom, purple: seasonal). Top row: Base run without interventions. Bottom row: ablation of the idiom's probability mass (red) beyond step .

The canvas will read the seasonal answer for the first few diffusion steps, after which it flips and stays at idiom. Interestingly, this is accompanied by dynamics in : idiom will start at ~0, and progressively gain weight, replacing seasonal.

We validated that the emergence of idiom is indeed causal: when ablating idiom's tokens, the transition can be prevented.

Transfer of interpretability techniques

So far, we have focused on DiffusionGemma's behavior. We here study whether the model's representation supports the behavioral finding of high monitorability.

Representation similarity

We first ask about how the representation changes between Gemma and DiffusionGemma, and then will ask about functional implications.

A simple way is to just measure overlap between pairs of inputs, where each element of the pair is fed through Gemma or DiffusionGemma, respectively. We here compare a simple cosine similarity, and centered kernel analysis (CKA, a measure that is invariant to global rotations of the representation).

Representational similarity falls with layers, and is significantly driven by bare model difference and bidirectional attention. Inputs are concepts from (Kantamneni et al., 2025). Notably, both the continued training from Gemma itself, as well as the switch to bidirectional attention mode, contribute to the similarity decrease.

A priori, this rather large drop suggests a significant change in how the representation is organized, which we then went on to test more specifically.

Probe retention

Probing allows us to study how well a model separates concepts. We use 56 binary concept datasets from the SAE-Probes benchmark (Kantamneni et al., 2025) and train logistic-regression probes on gemma-4's residual stream (≤ 1024 training examples for most concepts, 256 held-out), optimizing the layer via heldout AUC. We then test this probe on DG.

Probes largely transfer from Gemma to DiffusionGemma. Top: mean held-out AUC over the 56 concepts, probe source (trained on) × target (applied to). DG is split by attention mode (last-position read everywhere). Bottom: a held-out positive and negative test text for one concept (clickbait): On this instance, the same Gemma-trained probe scores Gemma-4 and DG activations near-identically. Ticks: <read model> · <attention mode> · <read position>. G = gemma-4, DG = DiffusionGemma, "last" = the probe reads the residual at the last token.

DiffusionGemma's representation is more linearly separable

Interestingly, we observed training and applying probes on DiffusionGemma in bidirectional attention mode yields a somewhat higher AUC. This is suggestive of DG's bidirectional attention yielding a better structured representation, but is confounded by potentially longer training.

Steering retention

To see whether these similarities in representation are causally load-bearing, we consider the steering experiments from the RepE paper (Zou et al., 2023). For each of 11 RepE concept tasks we fit a direction , where denotes the mean last-token residual activation over the task's positive contrastive stimuli (likewise for ), read separately on each stream (Gemma-4, DG causal, DG bidirectional). The direction is then injected additively into the residual stream of the steered model at a fixed strength ( with , layers 9–19) while it completes a neutral carrier prompt, once with and once with . A blinded judge sees the two generations in random order and needs to make a forced choice to select the -steered one.

Steering largely transfers from Gemma to DiffusionGemma. Top: blind judge accuracy, direction source × steered model. Bottom: the same Gemma-fit happiness direction applied to Gemma-4 and DiffusionGemma on one carrier prompt. Ticks: <read model> · <attention mode> · <read position of the direction fit>. "pr80 write" = the direction is added over the last 80% of prompt positions.

J-Lens retention

The Jacobian lens describes how a residual-stream activation affects future output, on average. It linearly transports into the final-layer basis and decodes it with the model's own unembedding, , where is the input–output Jacobian averaged over a generic text corpus (WikiText). We fit separate lenses on gemma-4, DG causal, and DG bidirectional, and read each on both models' residual streams.

J-lens largely transfers from Gemma to DiffusionGemma. Top: Transfer matrix, with scores being the fraction of layers * positions slots where the presumed intermediate is in the top-20 (the eval tasks from the J-Lens paper). Bottom: An example poetry eval task, where the models surface the rhyme already one line in advance. Ticks show the computed Jacobian: is the residual at source layer and position , the final-layer residual at target position . The causal fits average over causal targets , the bidirectional fit over all canvas positions . The expectation is taken over samples of WikiText.

DiffusionGemma represents tokens non-causally

Since DiffusionGemma uses bidirectional attention, it is interesting to ask whether J-space percepts will also be distributed in a non-causal way. In particular, do they occur before the token that triggers them? Indeed, we found such instances on some problems:

A position preceding the operand will hold it in J-space. Shown are the top-5 J-space tokens.

Conclusion

In this post, we have studied whether DiffusionGemma does latent reasoning via its vector-valued state . We found that top-k truncation of largely preserves accuracy, suggesting that it is not significantly being used in typical reasoning-focused tasks. However, for some tasks, we found that the model can make some use of , in terms of carrying out (parallel) computation on it.

Overall, this supports Engels et al.'s conclusion of a highly monitorable DiffusionGemma. In particular, we did not find any evidence of latent reasoning, as in using in a way that carries out nontrivial computation and is opaque.

Going forward, these largely negative results are an update that "true" latent reasoning is somewhat hard to learn. However, there already exist models like CODI that pass a vector-valued state that is not clearly a superposition of tokens, though they so far don't show superior performance. We believe that finding methods to better interpret such models is an important area of research.

Appendix

Here we discuss some additional interesting phenomena in DiffusionGemma relating to its monitorability that do not depend on using its distributional state, but rather arise from bidirectional attention and looped generation per se.

Post-hoc rationalization

An important question in monitorability (Bogdan et al., 2025) is whether the CoT is actually being used. An approach to measure this is to intervene on a fragment of the CoT, and see whether the answer changes.

For autoregressive models, the answer appears after the CoT. This incentivizes actually using CoT positions to compute the answer. In contrast, DiffusionGemma may arrive at an answer throughout multiple steps of computation, and then fill in a CoT post-hoc to match the format of the training distribution. An increased presence of such post-hoc rationalization relative to Gemma would be a blackpill for DiffusionGemma's monitorability.

Problem difficulty, CoT load-bearingness, and commitment time correlate. Problem difficulty is the mean rating of three blind LLM raters (shown only the problem and its correct answer), CoT load-bearingness is measured via a susceptibility : the flip rate of the answer when a -fraction of the CoT canvas positions (, pinned every step) is replaced with random tokens (averaged over , against the own-CoT-pinned baseline), and commitment time is the denoising step after which the answer's canvas positions stop changing (median over 5 seeds). We consider problems from four task families: cognitive-reflection-test problems that have a suggestive but wrong answer (crt), counting/enumeration (count), hard arithmetic/enumeration (hard), and multi-hop compute-then-transform chains (transform), see our code for details. Spearman with two-sided permutation per panel; × = accuracy < 0.5, dashed = least-squares fit.

For illustration, consider an easy and a hard problem side by side:

A pair of easy and hard examples illustrating the previous correlation. Top row: normal rollout, model answer highlighted, with its commitment time. Bottom row: the susceptibility, measured like in the previous figure.

Load-bearing problems commit the answer only after the CoT

In the previous paragraph, we found that commitment time of the answer correlates with load-bearingness. A natural hypothesis is whether the load-bearing problems will correspondingly also commit the CoT similarly late.

Post-hoc answers commit before their CoT, load-bearing answers wait till the CoT is committed. Mean token entropy of the answer positions (solid) and CoT positions (dashed) per denoising step, one rollout per battery problem, grouped by the susceptibility of the answer ( vs , faint = single problems, bold = group mean). is the number of tasks in each panel, taken from the previous figure.

How bidirectional are DiffusionGemma's generations?

One of the differences in DiffusionGemma is that it has bidirectional attention. While some tasks relying on some form of logical induction require causal left-to-right generation, other tasks do not obviously require it. However, being a finetune of Gemma, it is unclear how much it will just inherit a bias towards causal generation versus actually using bidirectional generation.

Commitment order follows the task's logical direction, but has a bias towards causal generation. Center of mass of the canvas positions committed by step (0 = left, 1 = right of the content span) over diffusion progress (faint = single rollouts, bold = mean). = Spearman between a judge's logical ordering of CoT steps and their actual appearance in the text. denotes distinct task instances.

We considered three tasks: GPQA, which involves logic-style problems and therefore should require left-to-right generation, poem writing, which has no clear left-to-right bias but only needs to satisfy a global structure, and reverse_chain, which requires the model to reverse its reasoning order.