This notebook explores the inner workings of diffusion-based language models (DLMs) like LLaDA using the LogitLens interpretability method. The analysis is inspired by two key sources:
- Alessio de Voto's blog post, which provides an excellent guide to the inner workings of Diffusion Language Models.
- The "Prophet" paper, which investigates verifying emergent properties in diffusion models and suggests that answers often converge early in the decoding timeline.
Unlike traditional left-to-right text generators, DLMs are bidirectional models that learn to fill in masked tokens through an iterative denoising process. LogitLens allows us to explore how a model's predictions evolve across its layers by applying the unembedding matrix (the language modeling head) to the hidden states of all intermediate layers.
Today, we're turning this lens on LLaDA to uncover how it builds its understanding layer by layer, providing insights into DLM functionality and highlighting differences with respect to autoregressive models.
Final Analysis: Two Axes of Convergence
By extending the LogitLens analysis across the full iterative decoding process, we can visually reconcile the findings of the "Prophet" paper (early convergence in time) with the layer-by-layer patterns seen in our initial analysis (late convergence in depth).
The results reveal that Diffusion Language Models exhibit two distinct types of convergence:
1. Convergence in Time (The "Prophet" Paper's Finding)
Examining the sequence of heatmaps from each timestep makes the "early convergence" phenomenon clear.
- In the initial steps (
t=0tot≈7), the model's final prediction (the bottom row) is uncertain and incoherent. - Around the halfway mark (e.g., timestep
t=8), the model's output stabilizes on a complete and correct sentence. - In all subsequent steps (
t=9tot=15), this final prediction remains largely unchanged, with the model only growing more certain (as shown by the brighter colors in the final layers).
This confirms the core finding that the final answer is often determined early in the iterative decoding process.
2. Convergence in Depth (The Logit Lens Finding)
Within any single heatmap—including those where the final answer has already converged—we still observe a "late convergence" pattern across the model's layers.
- The initial layers (top rows) consistently show low certainty (dark colors) and often predict nonsensical or placeholder tokens.
- The correct tokens and high certainty (bright yellow colors) only emerge in the final, deeper layers of the network.
Conclusion
This experiment reveals a more complete picture of how DLMs operate. They first establish a high-level, globally coherent "plan" for the full sentence early in the decoding timeline. Then, within each step of executing that plan, they refine the details from abstract representations (in early layers) to concrete tokens (in the final layers).