拡散モデルにおける多様性と制御の蒸留

Distilling Diversity and Control in Diffusion Models

March 13, 2025

著者: Rohit Gandikota, David Bau

cs.AI

要旨

蒸留拡散モデルは、基盤となるモデルと比較してサンプルの多様性が減少するという重大な制限を抱えています。本研究では、この多様性の損失にもかかわらず、蒸留モデルが基盤モデルの基本的な概念表現を保持していることを明らかにしました。我々は、基盤モデルで訓練されたConcept SlidersやLoRAsなどの制御メカニズムを、再訓練なしで蒸留モデルにシームレスに転移できる「制御蒸留」を実証します。この表現構造の保存は、蒸留中の多様性崩壊のメカニズムを調査するきっかけとなりました。蒸留が多様性にどのように影響するかを理解するために、我々はDiffusion Target (DT) Visualizationを導入しました。これは、モデルが中間ステップで最終出力をどのように予測するかを明らかにする分析・デバッグツールです。DT-Visualizationを通じて、生成アーティファクトや不整合を特定し、初期の拡散タイムステップが出力の多様性を不均衡に決定し、後期のステップは主に詳細を調整することを実証しました。これらの知見に基づき、我々は多様性蒸留を導入しました。これは、最初の重要なタイムステップのみに基盤モデルを戦略的に使用し、その後効率的な蒸留モデルに移行するハイブリッド推論アプローチです。実験により、この単純な修正が、基盤モデルから蒸留モデルへの多様性能力を回復するだけでなく、驚くべきことにそれを上回りながら、蒸留推論の計算効率をほぼ維持することを示しました。これらは、追加の訓練やモデルの変更を必要とせずに達成されます。コードとデータは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