Unpacking SDXL Turbo: Interpreting Text-to-Image Models with Sparse Autoencoders (2410.22366v2)

Abstract: Sparse autoencoders (SAEs) have become a core ingredient in the reverse engineering of large-LLMs. For LLMs, they have been shown to decompose intermediate representations that often are not interpretable directly into sparse sums of interpretable features, facilitating better control and subsequent analysis. However, similar analyses and approaches have been lacking for text-to-image models. We investigated the possibility of using SAEs to learn interpretable features for a few-step text-to-image diffusion models, such as SDXL Turbo. To this end, we train SAEs on the updates performed by transformer blocks within SDXL Turbo's denoising U-net. We find that their learned features are interpretable, causally influence the generation process, and reveal specialization among the blocks. In particular, we find one block that deals mainly with image composition, one that is mainly responsible for adding local details, and one for color, illumination, and style. Therefore, our work is an important first step towards better understanding the internals of generative text-to-image models like SDXL Turbo and showcases the potential of features learned by SAEs for the visual domain. Code is available at https://github.com/surkovv/sdxl-unbox

• The paper employs Sparse Autoencoders to analyze SDXL Turbo, unveiling how distinct transformer blocks specialize in image composition, detail enhancement, and style infusion.
• The study rigorously combines qualitative and quantitative methods with over 1.5 million LAION-COCO prompts to validate the interpretability of the model's internal features.
• The paper’s methodology offers a practical framework for refining text-to-image models and advancing mechanistic interpretability in diffusion-based generative systems.

Interpreting SDXL Turbo Using Sparse Autoencoders: Insights on Text-to-Image Models

The paper presents an innovative paper on understanding the intermediate representations of modern text-to-image generative models, specifically focusing on SDXL Turbo, a recent open-source few-step text-to-image diffusion model. This research employs Sparse Autoencoders (SAEs) to gain insight into the operations of SDXL Turbo's denoising U-net, with a particular emphasis on interpreting feature learning within the model's transformer blocks.

Methodology and Analysis

To explore whether SAEs can elucidate the computation performed during SDXL Turbo's generation process, the paper trains SAEs on transformer block updates within the model. Using the SDLens library, the authors cache and manipulate SDXL Turbo’s intermediate results, creating a dataset with over 1.5 million prompts from the LAION-COCO dataset. Each transformer block's dense feature maps are collected and used to train multiple SAEs. The paper reports a detailed analysis of the learned features, employing both qualitative and quantitative methods.

The empirical analysis demonstrates that SAEs have the potential to learn interpretable features within diffusion-based text-to-image models. Visualization techniques are developed to showcase the interpretability and causal effects of the SAE-learned features across various transformer blocks. Notably, different blocks in the SDXL Turbo pipeline specialize in discrete aspects of image generation: image composition, local detail addition, and style, among others.

Quantitative Validation

Quantitative experiments were performed to confirm the qualitative findings on a larger dataset, demonstrating the robustness of the hypotheses. An automatic feature annotation pipeline was developed for the transformer block deemed responsible for image compositions. This approach highlighted the efficacy of SAEs in endowing researchers with a tool for understanding the computational intricacies of SDXL Turbo’s forward pass.

Theoretical and Practical Implications

From a theoretical standpoint, this work advances the field of mechanistic interpretability by exploring the less-explored domain of diffusion models. The successful application of SAEs, originally a tool developed for LLMs to decompose internal representations into interpretable features, to image generation models marks an essential step forward. Practically, the insights gained from this research can aid in refining text-to-image pipelines for various applications, potentially enhancing the precision and control over the generated images.

Future Directions

The open-sourcing of both the SAEs and the SDLens library provides a solid foundation for further research. Future studies might benefit from exploring deeper interactions between features within and across blocks, as well as leveraging advanced interpretability techniques, such as circuit discovery, to unravel the higher-order relations within the computational process.

Complex visual features that necessitate particular contexts for effect realization add a layer of challenge that contemporary visual LLMs struggle to annotate adequately. Thus, research aimed at improving annotation techniques, possibly through extended prompt and feature space exploration, could provide additional benefits in understanding and controlling text-to-image generative models.

In conclusion, the paper makes significant strides toward demystifying the operation of text-to-image models, employing SAEs to extract interpretable and causally relevant features, thereby offering the research community a pathway to deeper understanding and innovation in AI-driven generative technologies.