Desempacotando SDXL Turbo: Interpretando Modelos de Texto-para-Imagem com Autoencoders Esparsos

Unpacking SDXL Turbo: Interpreting Text-to-Image Models with Sparse Autoencoders

October 28, 2024

Autores: Viacheslav Surkov, Chris Wendler, Mikhail Terekhov, Justin Deschenaux, Robert West, Caglar Gulcehre

cs.AI

Resumo

Os autoencoders esparsos (SAEs) tornaram-se um ingrediente essencial na engenharia reversa de grandes modelos de linguagem (LLMs). Para os LLMs, foi demonstrado que eles decompõem representações intermediárias que frequentemente não são diretamente interpretáveis em somas esparsas de características interpretáveis, facilitando um melhor controle e análise subsequente. No entanto, análises e abordagens semelhantes têm faltado para modelos de texto para imagem. Investigamos a possibilidade de usar SAEs para aprender características interpretáveis para modelos de difusão de texto para imagem de poucas etapas, como o SDXL Turbo. Para isso, treinamos SAEs nas atualizações realizadas pelos blocos transformadores dentro do denoising U-net do SDXL Turbo. Descobrimos que suas características aprendidas são interpretáveis, influenciam causalmente o processo de geração e revelam especialização entre os blocos. Em particular, encontramos um bloco que lida principalmente com composição de imagem, um que é principalmente responsável por adicionar detalhes locais, e um para cor, iluminação e estilo. Portanto, nosso trabalho é um importante primeiro passo para melhor compreender os internos de modelos generativos de texto para imagem como o SDXL Turbo e demonstra o potencial das características aprendidas pelos SAEs para o domínio visual. O código está disponível em https://github.com/surkovv/sdxl-unbox

English

Sparse autoencoders (SAEs) have become a core ingredient in the reverse engineering of large-language models (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