Distilling Diversity and Control in Diffusion Models
Authors: name: Rohit Gandikota; name: David Bau
Published: 2025-03-13 17:59:56

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