Destilando Diversidade e Controle em Modelos de Difusão

Distilling Diversity and Control in Diffusion Models

March 13, 2025

Autores: Rohit Gandikota, David Bau

cs.AI

Resumo

Modelos de difusão destilados sofrem de uma limitação crítica: a redução da diversidade de amostras em comparação com seus equivalentes base. Neste trabalho, descobrimos que, apesar dessa perda de diversidade, os modelos destilados retêm as representações fundamentais de conceitos dos modelos base. Demonstramos a destilação de controle - onde mecanismos de controle como Concept Sliders e LoRAs treinados em modelos base podem ser transferidos diretamente para modelos destilados e vice-versa, efetivamente destilando o controle sem necessidade de retreinamento. Essa preservação da estrutura representacional motivou nossa investigação sobre os mecanismos de colapso de diversidade durante a destilação. Para entender como a destilação afeta a diversidade, introduzimos a Visualização de Alvo de Difusão (DT-Visualization), uma ferramenta de análise e depuração que revela como os modelos preveem os resultados finais em etapas intermediárias. Através da DT-Visualization, identificamos artefatos de geração, inconsistências e demonstramos que os passos iniciais da difusão determinam desproporcionalmente a diversidade da saída, enquanto os passos posteriores refinam principalmente os detalhes. Com base nessas percepções, introduzimos a destilação de diversidade - uma abordagem híbrida de inferência que emprega estrategicamente o modelo base apenas para o primeiro passo crítico antes de transicionar para o modelo destilado eficiente. Nossos experimentos demonstram que essa simples modificação não apenas restaura as capacidades de diversidade dos modelos base para os destilados, mas surpreendentemente as supera, mantendo quase a eficiência computacional da inferência destilada, tudo isso sem exigir treinamento adicional ou modificações no modelo. Nosso código e dados estão disponíveis em https://distillation.baulab.info.

English

Distilled diffusion models suffer from a critical limitation: reduced sample diversity compared to their base counterparts. In this work, we uncover that despite this diversity loss, distilled models retain the fundamental concept representations of base models. We demonstrate control distillation - where control mechanisms like Concept Sliders and LoRAs trained on base models can be seamlessly transferred to distilled models and vice-versa, effectively distilling control without any retraining. This preservation of representational structure prompted our investigation into the mechanisms of diversity collapse during distillation. To understand how distillation affects diversity, we introduce Diffusion Target (DT) Visualization, an analysis and debugging tool that reveals how models predict final outputs at intermediate steps. Through DT-Visualization, we identify generation artifacts, inconsistencies, and demonstrate that initial diffusion timesteps disproportionately determine output diversity, while later steps primarily refine details. Based on these insights, we introduce diversity distillation - a hybrid inference approach that strategically employs the base model for only the first critical timestep before transitioning to the efficient distilled model. Our experiments demonstrate that this simple modification not only restores the diversity capabilities from base to distilled models but surprisingly exceeds it, while maintaining nearly the computational efficiency of distilled inference, all without requiring additional training or model modifications. Our code and data are available at https://distillation.baulab.info