在扩散模型中提炼多样性与控制性

Distilling Diversity and Control in Diffusion Models

March 13, 2025

作者: Rohit Gandikota, David Bau

cs.AI

摘要

蒸馏扩散模型面临一个关键局限：与基础模型相比，样本多样性显著降低。本研究中，我们发现尽管存在这种多样性损失，蒸馏模型仍保留了基础模型的核心概念表征。我们展示了控制蒸馏——在基础模型上训练的控制机制（如概念滑块和LoRAs）能够无缝迁移至蒸馏模型，反之亦然，从而无需重新训练即可有效蒸馏控制能力。这种表征结构的保留促使我们深入探究蒸馏过程中多样性崩溃的机制。为理解蒸馏如何影响多样性，我们引入了扩散目标（DT）可视化，这一分析和调试工具揭示了模型在中间步骤如何预测最终输出。通过DT可视化，我们识别了生成过程中的伪影与不一致性，并证明初始扩散时间步长对输出多样性具有决定性影响，而后续步骤主要精修细节。基于这些洞察，我们提出了多样性蒸馏——一种混合推理策略，仅在关键的第一时间步长策略性地使用基础模型，随后切换至高效的蒸馏模型。实验表明，这一简单调整不仅恢复了从基础模型到蒸馏模型的多样性能力，甚至超越了前者，同时几乎保持了蒸馏推理的计算效率，且无需额外训练或模型修改。我们的代码与数据公开于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