Масштабная дистилляция диффузионных моделей

Scale-wise Distillation of Diffusion Models

March 20, 2025

Авторы: Nikita Starodubcev, Denis Kuznedelev, Artem Babenko, Dmitry Baranchuk

cs.AI

Аннотация

Мы представляем SwD, масштабно-ориентированную фреймворк дистилляции для диффузионных моделей (DMs), которая эффективно использует идеи предсказания следующего масштаба для генераторов на основе диффузии с малым количеством шагов. В частности, SwD вдохновлен недавними исследованиями, связывающими процессы диффузии с неявной спектральной авторегрессией. Мы предполагаем, что DMs могут начинать генерацию на более низких разрешениях данных и постепенно увеличивать масштаб выборок на каждом шаге шумоподавления без потери производительности, при этом значительно снижая вычислительные затраты. SwD естественным образом интегрирует эту идею в существующие методы дистилляции диффузии, основанные на согласовании распределений. Кроме того, мы расширяем семейство подходов к согласованию распределений, вводя новый патч-лосс, который обеспечивает более детальное сходство с целевым распределением. Примененный к современным диффузионным моделям для генерации изображений по тексту, SwD приближается к времени вывода, эквивалентному двум шагам на полном разрешении, и значительно превосходит аналогичные методы при одинаковом вычислительном бюджете, что подтверждается автоматическими метриками и исследованиями предпочтений пользователей.

English

We present SwD, a scale-wise distillation framework for diffusion models (DMs), which effectively employs next-scale prediction ideas for diffusion-based few-step generators. In more detail, SwD is inspired by the recent insights relating diffusion processes to the implicit spectral autoregression. We suppose that DMs can initiate generation at lower data resolutions and gradually upscale the samples at each denoising step without loss in performance while significantly reducing computational costs. SwD naturally integrates this idea into existing diffusion distillation methods based on distribution matching. Also, we enrich the family of distribution matching approaches by introducing a novel patch loss enforcing finer-grained similarity to the target distribution. When applied to state-of-the-art text-to-image diffusion models, SwD approaches the inference times of two full resolution steps and significantly outperforms the counterparts under the same computation budget, as evidenced by automated metrics and human preference studies.