Distillation multi-échelle des modèles de diffusion

Scale-wise Distillation of Diffusion Models

March 20, 2025

Auteurs: Nikita Starodubcev, Denis Kuznedelev, Artem Babenko, Dmitry Baranchuk

cs.AI

Résumé

Nous présentons SwD, un cadre de distillation multi-échelle pour les modèles de diffusion (DMs), qui exploite efficacement les idées de prédiction à l'échelle suivante pour les générateurs basés sur la diffusion à faible nombre d'étapes. Plus précisément, SwD s'inspire des récentes perspectives reliant les processus de diffusion à l'autorégression spectrale implicite. Nous supposons que les DMs peuvent initier la génération à des résolutions de données plus faibles et augmenter progressivement la résolution des échantillons à chaque étape de débruitage sans perte de performance, tout en réduisant significativement les coûts de calcul. SwD intègre naturellement cette idée dans les méthodes existantes de distillation par diffusion basées sur l'appariement de distributions. De plus, nous enrichissons la famille des approches d'appariement de distributions en introduisant une nouvelle fonction de perte par patch qui impose une similarité plus fine avec la distribution cible. Lorsqu'il est appliqué aux modèles de diffusion texte-image de pointe, SwD approche les temps d'inférence de deux étapes à pleine résolution et surpasse significativement les alternatives sous le même budget de calcul, comme en témoignent les métriques automatisées et les études de préférence humaine.

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.