Summary (Overview)

• Diffusion Language Models (DLMs) internally encode a latent signal related to denoising progress (the fraction of unmasked tokens, $\tau$), despite not being explicitly conditioned on timestep.
• This signal can be reliably decoded using simple MLP probes from individual token hidden states across all layers, achieving $R^2 > 0.5$ for both LLaDA and Dream.
• Steering the model's internal representation of $\tau$ by swapping mean activation vectors produces predictable changes in model confidence, entropy, and KL divergence, demonstrating causal relevance.
• The $\tau$ representation lives in a low-dimensional subspace (≈2 principal components) that follows a parabolic manifold structure, shared across layers in LLaDA and organized into blocks in Dream.
• Models can partially correct early-layer perturbations to the $\tau$ signal, indicating the representation is repeatedly recomputed throughout the depth.

Introduction and Theoretical Foundation

Large Language Models (LLMs) have driven a paradigm shift, and understanding their internal representations is an active research area. Recently, Masked Diffusion Language Models (DLMs) have emerged as an alternative to autoregressive models. Instead of left-to-right generation, they progressively denoise [MASK] tokens in a BERT-style manner. While prior work has focused on improving efficiency and generation quality, the internal mechanisms of DLMs remain largely unexplored. In particular, it is unclear whether DLMs internally represent denoising progress during generation.

Theoretical work suggests that the strong performance of absorbing-state DLMs may be attributed to the special role of the [MASK] token. The reverse process only needs to denoise masked positions, and the score can be expressed as a clean-data conditional distribution up to an analytic time-dependent factor (Ou et al., 2024). The fraction of masked tokens is a natural proxy for denoising progress, but it is unknown whether and how such a signal is internally represented.

The paper investigates three research questions:

• RQ1: Do DLMs internally represent a signal related to the denoising step?
• RQ2: Is the identified signal important for generation? How does modifying it affect downstream computations?
• RQ3: What characteristics does the signal exhibit? Is there a unified model-level representation?

The analysis is conducted on two representative large-scale masked diffusion language models: LLaDA-1.5 (Zhu et al., 2026) and Dream (Ye et al., 2025).

Methodology

Background on DLMs

DLMs are trained to recover a clean sequence $\mathbf{x}_0$ from a corrupted version $\mathbf{x}_s$, where corruption replaces a subset of positions with the [MASK] token. The corruption level $s \in [0,1]$ controls the fraction of masked tokens. The training loss is:

$$-\mathbb{E}_{\mathbf{x}_0, s, \mathbf{x}_s} \left[ \frac{1}{s} \sum_{j \in \mathcal{M}(\mathbf{x}_s)} \log p_\theta(x_0^j | \mathbf{x}_s) \right] \tag{1}$$

where $\mathcal{M}(\mathbf{x}_s)$ is the set of masked positions. Inference proceeds through $T$ denoising steps. At each step $t$, the fraction of unmasked tokens is:

$$\tau_t := 1 - \frac{|\mathcal{M}(\mathbf{x}_t)|}{L} \tag{2}$$

where $L$ is the sequence length. $\tau_t$ serves as a proxy for denoising progress.

Probing for $\tau$

To test whether $\tau$ is recoverable from hidden states, MLP probes $\phi_l(\cdot): \mathbb{R}^d \to (0,1)$ are trained for each layer $l$ to predict $\tau_t$ from a single residual-stream hidden state $h_{t,l,n}^j$. The loss is:

$$\mathcal{L}_{\text{MLP}} = \mathbb{E}_{t,n,j} \left[ \left( \tau_t - \phi_l(h_{t,l,n}^j) \right)^2 \right]$$

Probes are trained separately for [MASK] tokens, non-[MASK] tokens, and all tokens.

Mean Activation Vectors and Steering

To characterize the internal representation, mean activation vectors are computed for each denoising step $t$ and layer $l$:

$$\mu_{t,l} := \mathbb{E}_{n,j} \left[ h_{t,l,n}^j \right] \tag{3}$$

The denoising process is discretized into 100 bins. Steering is performed by replacing the mean vector of the current step with that of a target step $\hat{t}$:

$$\tilde{h}_{t,l}^j := h_{t,l}^j - \mu_{t,l} + \mu_{\hat{t},l} \tag{4}$$

As a control, random perturbations of matched norm are sampled from the empirical covariance matrix:

$$\tilde{h}_{t,l}^j := h_{t,l}^j + a_{t \to \hat{t}}, \quad a \sim \mathcal{N}(0, C_{t,l}) \tag{5}$$

Effects are measured via entropy change $\Delta \bar{S}_t$, confidence change $\Delta \bar{c}_t$, and KL divergence $\text{KL}_t$ between clean and steered distributions.

Subspace Analysis

Principal Component Analysis (PCA) is applied to the mean vectors $\mu_{t,l}$ per layer. Steering is repeated using only the projection onto the top $k$ principal components and the orthogonal complement, rescaled to match the original norm:

$$\Delta_{t \to \hat{t}}^{\parallel, l} := \frac{P_{\parallel,l} \Delta_{t \to \hat{t}}^l}{\|P_{\parallel,l} \Delta_{t \to \hat{t}}^l\|} \|\Delta_{t \to \hat{t}}^l\|, \quad \Delta_{t \to \hat{t}}^{\perp, l} := \frac{P_{\perp,l} \Delta_{t \to \hat{t}}^l}{\|P_{\perp,l} \Delta_{t \to \hat{t}}^l\|} \|\Delta_{t \to \hat{t}}^l\| \tag{7}$$

Empirical Validation / Results

Probing Results (RQ1)

Figure 2 shows the $R^2$ coefficients for MLP probes across layers of LLaDA. Probes maintain $R^2 > 0.5$ across all layers, with slight degradation at early and late layers. Both [MASK] and non-[MASK] tokens support accurate recovery, indicating the signal is distributed across the residual stream.

Table: MLP probe $R^2$ coefficients for LLaDA (higher is better). Probes maintain high accuracy across layers and token types.

The correlation between probe predictions on mean vectors $\phi_l(\mu_{t,l})$ and normalized timestep $t/100$ is very high: Pearson $0.976$, Spearman $0.980$ for LLaDA; $0.962$ and $0.974$ for Dream.

Steering Results (RQ2)

Steering layer 29 of LLaDA (Figure 4) shows systematic effects: steering towards larger $\hat{t}$ (more denoised) increases confidence and decreases entropy; steering towards smaller $\hat{t}$ has the opposite effect. KL divergence grows proportionally to $|\hat{t} - t|$. Random perturbations of matched norm produce weaker and incoherent effects.

When steering early layers (e.g., layer 6), the model is able to correct the perturbation in deeper layers, as measured by probe-predicted $\tau$ drifting back towards clean values (Figure 5). Extreme targets ($\hat{t}=100$) are harder to compensate for.

Characterization of the Signal (RQ3)

Low-dimensional structure: PCA on mean vectors per layer shows that a single principal component explains >90% of variance on average (Figure 6a). Projections reveal a parabolic manifold (Figure 6b, Figure 1). Steering within the top 2 principal components recovers the full steering effect, while the orthogonal component produces incoherent effects (Figure 7).

Shared representation across layers: Standardizing 2D projections across layers reveals a common parabolic trajectory (Figure 8 for LLaDA), with error bars within ~0.1, indicating a unified model-level geometry.

Cross-layer alignment: Cosine similarity of same-indexed mean vectors across layers (Figure 9) shows high alignment for most layers in LLaDA, except layer 32 which is nearly orthogonal. Dream shows block-wise alignment.

Internal layer representation: Within a layer, self-attention and MLP outputs have opposite representations (highly anti-correlated), suggesting a mechanism for correction (Figure 10, Figure 11).

Theoretical and Practical Implications

• Theoretical significance: The paper provides the first mechanistic evidence that DLMs internally track denoising progress via a structured, low-dimensional representation. This connects DLM behavior to broader findings in LLM interpretability (e.g., "clock" or "time" representations in transformers).
• Practical implications: The ability to steer denoising progress via mean vectors offers a new tool for controlling generation dynamics (e.g., adjusting confidence/entropy at inference). The low-dimensional subspace suggests efficient steering is possible with minimal computation.
• Causal relevance: The signal is not merely correlational; steering it produces predictable changes in output distributions, confirming its functional role.
• Architecture insights: The repeated recomputation and correction of the signal across layers suggests that DLMs have built-in mechanisms to maintain and refine this representation, which could inform future model designs.

Conclusion

The paper demonstrates that Diffusion Language Models encode a latent representation of denoising progress ($\tau$) in their residual streams. This signal is:

• Decodable via simple probes across layers and token types.
• Causally steerable via mean activation vectors, producing predictable changes in entropy, confidence, and KL divergence.
• Low-dimensional (≈2D parabolic manifold) and shared across layers.
• Repeatedly recomputed across depth, allowing correction of early perturbations.

Future directions include identifying the circuit that computes this signal, exploring its use for efficient decoding/remasking strategies, and extending the analysis to other DLM architectures (e.g., block-diffusion models). Limitations include the restriction to LLaDA and Dream, and the focus on sequence-level statistics rather than token-level effects of steering.