Distilling Diversity and Control in Diffusion Models
Rohit Gandikota, David Bau
2025-03-14
Summary
This paper explores how to make diffusion models, which are used for generating images, more diverse and controllable, while also making them more efficient.
What's the problem?
Diffusion models, after being 'distilled' to make them faster, often lose some of their ability to create a wide variety of images. They become less diverse, even though they still understand the basic concepts.
What's the solution?
The researchers found that control mechanisms can be transferred between the original and distilled models without retraining. They also discovered that the initial steps in generating an image are most important for diversity, so they created a hybrid approach that uses the original model for the first step and then switches to the faster, distilled model.
Why it matters?
This work matters because it allows for faster image generation without sacrificing diversity. It makes these models more practical and versatile for various applications, while maintaining efficiency.
Abstract
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