DLM-Scope: Mechanistic Interpretability of Diffusion Language Models via Sparse Autoencoders
This work provides foundational tools for interpreting DLMs, an emerging alternative to autoregressive LLMs, though it is incremental in applying existing SAE methods to a new model type.
The authors tackled the problem of mechanistic interpretability for diffusion language models (DLMs) by developing DLM-Scope, a sparse autoencoder-based framework, and found that it extracts interpretable features, reduces cross-entropy loss in early layers by up to 0.15 bits per token, and enables more effective interventions than in autoregressive LLMs.
Sparse autoencoders (SAEs) have become a standard tool for mechanistic interpretability in autoregressive large language models (LLMs), enabling researchers to extract sparse, human-interpretable features and intervene on model behavior. Recently, as diffusion language models (DLMs) have become an increasingly promising alternative to the autoregressive LLMs, it is essential to develop tailored mechanistic interpretability tools for this emerging class of models. In this work, we present DLM-Scope, the first SAE-based interpretability framework for DLMs, and demonstrate that trained Top-K SAEs can faithfully extract interpretable features. Notably, we find that inserting SAEs affects DLMs differently than autoregressive LLMs: while SAE insertion in LLMs typically incurs a loss penalty, in DLMs it can reduce cross-entropy loss when applied to early layers, a phenomenon absent or markedly weaker in LLMs. Additionally, SAE features in DLMs enable more effective diffusion-time interventions, often outperforming LLM steering. Moreover, we pioneer certain new SAE-based research directions for DLMs: we show that SAEs can provide useful signals for DLM decoding order; and the SAE features are stable during the post-training phase of DLMs. Our work establishes a foundation for mechanistic interpretability in DLMs and shows a great potential of applying SAEs to DLM-related tasks and algorithms.