DLM-Scope: Mechanistic Interpretability of Diffusion Language Models via Sparse Autoencoders

TL;DR

This paper introduces DLM-Scope, a novel SAE-based interpretability framework for diffusion language models, revealing unique effects of SAE insertion and enabling improved model interventions and understanding.

Contribution

It is the first to adapt sparse autoencoders for interpretability in diffusion language models, demonstrating their effectiveness and uncovering distinct behaviors compared to autoregressive models.

Findings

SAEs can faithfully extract interpretable features in DLMs

SAE insertion reduces cross-entropy loss in early DLM layers, unlike in LLMs

SAE features enable more effective diffusion-time interventions

Abstract

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.