Дистилляция разнообразия и контроля в диффузионных моделях

Distilling Diversity and Control in Diffusion Models

March 13, 2025

Авторы: Rohit Gandikota, David Bau

cs.AI

Аннотация

Дистиллированные диффузионные модели страдают от критического ограничения: сниженное разнообразие выборок по сравнению с их базовыми аналогами. В данной работе мы обнаруживаем, что, несмотря на эту потерю разнообразия, дистиллированные модели сохраняют фундаментальные концептуальные представления базовых моделей. Мы демонстрируем дистилляцию управления — когда механизмы управления, такие как Concept Sliders и LoRAs, обученные на базовых моделях, могут быть бесшовно перенесены на дистиллированные модели и наоборот, эффективно дистиллируя управление без необходимости повторного обучения. Это сохранение структуры представлений побудило нас исследовать механизмы коллапса разнообразия в процессе дистилляции. Чтобы понять, как дистилляция влияет на разнообразие, мы представляем визуализацию целевых диффузий (Diffusion Target Visualization, 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