Distillation de la diversité et du contrôle dans les modèles de diffusion

Distilling Diversity and Control in Diffusion Models

March 13, 2025

Auteurs: Rohit Gandikota, David Bau

cs.AI

Résumé

Les modèles de diffusion distillés souffrent d'une limitation critique : une diversité d'échantillons réduite par rapport à leurs homologues de base. Dans ce travail, nous découvrons que malgré cette perte de diversité, les modèles distillés conservent les représentations conceptuelles fondamentales des modèles de base. Nous démontrons la distillation de contrôle - où des mécanismes de contrôle comme les Concept Sliders et les LoRAs entraînés sur des modèles de base peuvent être transférés de manière transparente vers des modèles distillés et vice-versa, distillant ainsi le contrôle sans aucun réentraînement. Cette préservation de la structure représentationnelle a motivé notre investigation sur les mécanismes de l'effondrement de la diversité lors de la distillation. Pour comprendre comment la distillation affecte la diversité, nous introduisons la Visualisation des Cibles de Diffusion (DT-Visualization), un outil d'analyse et de débogage qui révèle comment les modèles prédisent les sorties finales à des étapes intermédiaires. Grâce à la DT-Visualization, nous identifions des artefacts de génération, des incohérences, et démontrons que les premières étapes de diffusion déterminent de manière disproportionnée la diversité des sorties, tandis que les étapes ultérieures affinent principalement les détails. Sur la base de ces insights, nous introduisons la distillation de diversité - une approche d'inférence hybride qui utilise stratégiquement le modèle de base uniquement pour la première étape critique avant de passer au modèle distillé efficace. Nos expériences montrent que cette simple modification non seulement restaure les capacités de diversité des modèles de base aux modèles distillés, mais les dépasse de manière surprenante, tout en maintenant presque l'efficacité computationnelle de l'inférence distillée, le tout sans nécessiter d'entraînement supplémentaire ou de modifications du modèle. Notre code et nos données sont disponibles à l'adresse 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