$\textit{Revelio}$: Interpreting and leveraging semantic information in diffusion models

TL;DR

Revelio investigates how diffusion models encode semantic information across layers and denoising timesteps. It introduces k-Sparse Autoencoders (k-SAE) to extract monosemantic features from diffusion activations and validates these features with a lightweight Diff-C classifier for transfer learning and visual reasoning. The study spans UNet and DiT architectures, timesteps, and pretraining, revealing non-linear granularity of semantic information and architecture-induced biases, with practical benefits in transfer performance and inference speed. Overall, the work advances interpretability of diffusion models and informs design choices for semantic control and efficient deployment.

Abstract

We study $\textit{how}$ rich visual semantic information is represented within various layers and denoising timesteps of different diffusion architectures. We uncover monosemantic interpretable features by leveraging k-sparse autoencoders (k-SAE). We substantiate our mechanistic interpretations via transfer learning using light-weight classifiers on off-the-shelf diffusion models' features. On $4$ datasets, we demonstrate the effectiveness of diffusion features for representation learning. We provide an in-depth analysis of how different diffusion architectures, pre-training datasets, and language model conditioning impacts visual representation granularity, inductive biases, and transfer learning capabilities. Our work is a critical step towards deepening interpretability of black-box diffusion models. Code and visualizations available at: https://github.com/revelio-diffusion/revelio

Figures (14)

• Figure 1: k-sparse autoencoders (k-SAE) trained on complex visual features help identify monosemantic visual properties represented within black-box diffusion models. We show sample k-SAE neurons and top-4 images that yield highest activations when the k-SAE is trained on intermediate diffusion layer's features on Oxford-IIIT Pet oxford dataset. Note how these features encapsulate distinct fine-grained information about different breeds like Keeshond and Samoyed. Best viewed in color.
• Figure 2: k-SAE visualizations across layers of the U-Net in SD 1.5 and sample images from different neurons yielding highest activations when k-SAEs are trained on different layers for $t=25$ on Oxford-IIIT Pet. We note that across $3$ random neurons of k-SAEs, the bottleneck layer captures very coarse-grained information, where foreground objects positioned similarly are activated by the same neuron. up_ft1 captures valuable breed specific domain information while up_ft2 seems to capture high-frequency visual patterns.
• Figure 3: k-SAE visualizations on Oxford-IIIT Pet of bottleneck, up_ft1, and up_ft2 U-Net layers at $t=25$. bottleneck isolates very coarse patterns of objects positioned similarly with respect to the background. For up_ft1, clear class-specific features are observed helping us isolate different fine-grained breeds. up_ft2 captures more global texture information such as that of grass.
• Figure 4: Top-1 accuracy of different SD 1.5 layer features. Features from up_ft1 consistently yield best performance for SD 1.5.
• Figure 5: k-SAE visualizations on Caltech-101 of bottleneck and up_ft1 UNet layers at $t=25$. Unlike for fine-grained dataset (Fig. \ref{['fig:qual_step25_oxfordpet_up1_up2']}), bottleneck captures class information, likely due to distinct object shapes (sailing ships v/s elephants). up_ft1 captures more abstract information such as sketches or objects with white background.