Ученые Yandex Research и НИУ ВШЭ представили метод Scale-Wise Distillation of Diffusion Models (SwD) для снижения вычислительной нагрузки при генерации изображений с помощью диффузионных моделей. Он сокращает время на получение результата с нескольких секунд до 0,3–0,4 секунды.

Процесс создания изображения в диффузионных моделях состоит из десятков шагов с вычислениями в высоком разрешении. SwD предлагает начинать генерацию с низкого качества, что исключает избыточные вычисления на ранних этапах. В рамках метода также используется дистилляция обученных моделей, таких как FLUX и Stable Diffusion 3.5: более простая модель учится воспроизводить результат сложной. Это сокращает число шагов генерации с десятков до 4–6.

Для обучения исследователи предложили новую функцию потерь — Maximum Mean Discrepancy (MMD). Она сравнивает как модель-учитель видит изображение на внутренних уровнях обработки с тем, как модель-студент представляет это изображение. Как отмечают исследователи, MMD может использоваться в качестве самостоятельного метода дистилляции. В экспериментах он сокращал время одной итерации обучения в семь раз по сравнению с более сложными комбинированными подходами.