С помощью диффузионных моделей можно генерировать изображения отличного качества по текстовому описанию, но без дистилляции они требуют десятки прогонов модели, что снижает их эффективность эффективными по сравнению с альтернативными парадигмами. Авторы сегодняшней статьи разработали Distribution Matching Distillation (DMD) — процедуру дистилляции диффузионной модели (учителя) в одношаговый генератор изображений (студента), который хорошо сохраняет качество картинок и их соответствие текстовому описанию.
Авторы предлагают учить одношаговый генератор изображений
p_\theta(x) за счет минимизации обратной KL-дивергенции между распределением студента и распределением данных: KL(p_\theta(x) || p_{data}). Напомним, что в таком KL-е мат-ожидание берется по распределению студента и следовательно при оценке методом Монте-Карло примеры сэмплируются из текущей версии модели-студента. Тем не менее, посчитать обратный KL явно все еще не получится. Однако главная идея работы заключается в том, что мы можем оценить производную обратного KL-я, используя прямой проход уже обученных диффузионных моделей: одна модель оценивает скор реального распределения
(\nabla_x \log p_data(x)), а другая скор распределения студента (\nabla_x \log p_\theta(x)). Отметим, что в самом начале у нас есть только одна модель, которая оценивает \nabla_x \log p_data(x), а вторую нужно учить на текущих генерациях студентов тем самым адаптируя её под изменяющееся распределение студента по ходу обучения. Процедура напоминает обучение GAN-a, но с другим обджективом. Внимательный читатель обратит внимание, что диффузия обычно предсказывает скоры для зашумленных распределений, а не для чистых, как описано выше. Но в данном методе это не баг, а фича, о которой лучше почитать в самой статье.
Важный нюанс при минимизации обратного KL — он склонен к потере мод распределения, что часто выражается в генерации однотипных картинок для одного текстового описания. Чтобы побороть проблему, авторы предлагают добавить дополнительный дистилляционный лосс, который заставляет модель не забывать куски распределения учителя. Дополнительный обджектив представляет собой регрессионный лосс, когда сначала собираются пары (шум — картинка), сгенерированные учителем, и затем по MSE многошаговый выход учителя сближается с одношаговым студентом по тому же входному шуму.
Таким образом, итоговая процедура обучения состоит из двух лоссов: обратный KL, отвечающий за качество генерации студента, и регрессионный лосс, стабилизирующий процедуру обучения и не дающий потерять разнообразие.
По словам авторов статьи, такой подход позволяет добиться лучших результатов, чем предыдущие методы (на момент декабря 2023 года), достигая 2.62 FID на ImageNet 64x64 и 11.49 FID на zero-shot COCO-30k, что сравнимо с результатами Stable Diffusion, но при этом DMD работает гораздо быстрее, генерируя изображения со скоростью до 20 FPS. Однако, стоит обратить внимание, что метрики для генерации картинок — ненадежный показатель качества, поэтому в идеале нужно проводить оценку на людях, чего авторы не сделали. Модели и кода в открытом доступе тоже не было, поэтому о реальном качестве результатов остается только догадываться.
В оправдание авторов: спустя полгода команда выпустила улучшенную версию метода, к которой приложили и код, и модели, и полноценную оценку качества генерации. Почитать о новом методе можно здесь.
Разбор подготовил
CV Time