Does DiffusionGemma Do Latent Reasoning? Researchers Probe Its Hidden State

New analysis of Google's DiffusionGemma finds the model can carry parallel computation in its vector-valued hidden state, though most tasks remain interpretable via top-token projection.

AATMA Team
AATMA Team
·3 min read
An analysis of Google's DiffusionGemma found that while the model's vector-valued hidden state can carry parallel computation in specific tasks like letter shifting, most reasoning remains interpretable via top-token projection. Standard interpretability techniques including probes, steering, and J-lens largely transfer from autoregressive Gemma to the diffusion variant.

DiffusionGemma generates text through diffusion rather than autoregression. Instead of producing one token at a time, it starts from noise and iteratively denoises a full canvas. Between each step, the model passes not just tokens but a full vector-valued distribution across all positions. That vector could, in principle, carry hidden computations that are invisible to anyone watching only the token stream. A new analysis asks whether this constitutes latent reasoning, and whether it undermines the model's monitorability.

The Monitorability Baseline

Earlier work on DiffusionGemma concluded that the model remains highly monitorable. One key finding was that projecting the distribution to its top-k entries largely preserves performance, suggesting that the full vector is not load-bearing for most tasks. The new analysis strengthens this result: with a gentler sampler (more steps, wider temperature range, lower entropy bound), even top-1 projection maintains accuracy. The performance drop observed in prior work appears to have been a sampler artifact, not evidence of essential latent information.

Where the Vector Does Matter

Despite the general finding, the researchers identified specific case studies where the distribution vector carries computationally load-bearing information. In a letter-shifting task (move a letter three positions down the alphabet), DiffusionGemma maintains a probability distribution over possible starting letters. Injecting probability mass on a subleading operand letter causes the answer to shift accordingly one step later. The model is effectively performing parallel computation on multiple hypotheses encoded in superposition.

This behavior is interpretable because it encodes superposition: a linear combination of known tokens. It is not opaque latent reasoning in the sense of carrying information in a form unrelated to the vocabulary. The researchers tested multiplication and absolute value tasks and found no significant superposition effects, suggesting that parallel computation via distribution vectors is task-dependent and not universal.

Interpretability Transfer

The analysis also tested whether standard interpretability techniques transfer from autoregressive Gemma to DiffusionGemma. Three methods were evaluated:

Probes: Linear probes trained on Gemma-4's residual stream were applied to DiffusionGemma. Held-out AUC remained high, with bidirectional attention mode in DiffusionGemma actually yielding slightly better probe performance than causal mode.

Steering: Directions fitted on Gemma-4's residual stream were injected into DiffusionGemma. A blinded judge could still identify the steered generation at rates well above chance, confirming that representational directions remain causally active.

J-Lens: The linear lens, which maps residual activations to vocabulary space via the model's own unembedding, transferred between models with only moderate degradation. Some J-space percepts appeared before their triggering tokens, consistent with DiffusionGemma's bidirectional attention.

The Verdict

The overall picture is reassuring for monitorability. DiffusionGemma does not appear to engage in opaque latent reasoning on typical tasks. The cases where the vector-valued state is computationally load-bearing involve superposition of interpretable tokens, not illegible hidden states. This supports the view that diffusion models derived from text-pretrained LLMs retain much of the interpretability of their autoregressive counterparts, though the finding may not generalize to diffusion architectures trained from scratch without that inheritance.