Destilando Diversidad y Control en Modelos de Difusión

Distilling Diversity and Control in Diffusion Models

March 13, 2025

Autores: Rohit Gandikota, David Bau

cs.AI

Resumen

Los modelos de difusión destilados presentan una limitación crítica: una reducción en la diversidad de muestras en comparación con sus contrapartes base. En este trabajo, descubrimos que, a pesar de esta pérdida de diversidad, los modelos destilados conservan las representaciones fundamentales de conceptos de los modelos base. Demostramos la destilación de control, donde mecanismos de control como Concept Sliders y LoRAs entrenados en modelos base pueden transferirse sin problemas a modelos destilados y viceversa, efectivamente destilando control sin necesidad de reentrenamiento. Esta preservación de la estructura representativa motivó nuestra investigación sobre los mecanismos de colapso de diversidad durante la destilación. Para comprender cómo la destilación afecta la diversidad, introducimos la Visualización de Objetivo de Difusión (DT-Visualization), una herramienta de análisis y depuración que revela cómo los modelos predicen los resultados finales en pasos intermedios. A través de DT-Visualization, identificamos artefactos de generación, inconsistencias, y demostramos que los pasos iniciales de difusión determinan de manera desproporcionada la diversidad de salida, mientras que los pasos posteriores principalmente refinan los detalles. Basándonos en estas ideas, introducimos la destilación de diversidad, un enfoque híbrido de inferencia que emplea estratégicamente el modelo base solo para el primer paso crítico antes de transicionar al modelo destilado eficiente. Nuestros experimentos demuestran que esta simple modificación no solo restaura las capacidades de diversidad de los modelos base a los destilados, sino que sorprendentemente las supera, manteniendo casi la eficiencia computacional de la inferencia destilada, todo sin requerir entrenamiento adicional o modificaciones del modelo. Nuestro código y datos están disponibles en 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