Stem Je Stappen Af: Optimalisatie van Bemonsteringsschema's in Diffusiemodellen

Align Your Steps: Optimizing Sampling Schedules in Diffusion Models

April 22, 2024

Auteurs: Amirmojtaba Sabour, Sanja Fidler, Karsten Kreis

cs.AI

Samenvatting

Diffusiemodellen (DMs) hebben zich gevestigd als de state-of-the-art benadering voor generatieve modellering in het visuele domein en daarbuiten. Een cruciaal nadeel van DMs is hun trage sampling snelheid, die afhankelijk is van vele sequentiële functie-evaluaties door grote neurale netwerken. Sampling van DMs kan worden gezien als het oplossen van een differentiaalvergelijking via een gediscretiseerde set van ruisniveaus, bekend als het sampling schema. Terwijl eerdere werken zich voornamelijk richtten op het afleiden van efficiënte oplossers, is er weinig aandacht besteed aan het vinden van optimale sampling schema's, en de gehele literatuur vertrouwt op handmatig gemaakte heuristieken. In dit werk stellen we voor het eerst een algemene en principiële benadering voor om de sampling schema's van DMs te optimaliseren voor hoogwaardige outputs, genaamd Align Your Steps. We maken gebruik van methoden uit de stochastische calculus en vinden optimale schema's die specifiek zijn voor verschillende oplossers, getrainde DMs en datasets. We evalueren onze nieuwe benadering op verschillende benchmarks voor beeld-, video- en 2D-toydatasynthese, met behulp van een verscheidenheid aan verschillende samplers, en observeren dat onze geoptimaliseerde schema's in bijna alle experimenten beter presteren dan eerdere handmatig gemaakte schema's. Onze methode toont het onbenutte potentieel van sampling schema-optimalisatie, vooral in het regime van synthese met weinig stappen.

English

Diffusion models (DMs) have established themselves as the state-of-the-art generative modeling approach in the visual domain and beyond. A crucial drawback of DMs is their slow sampling speed, relying on many sequential function evaluations through large neural networks. Sampling from DMs can be seen as solving a differential equation through a discretized set of noise levels known as the sampling schedule. While past works primarily focused on deriving efficient solvers, little attention has been given to finding optimal sampling schedules, and the entire literature relies on hand-crafted heuristics. In this work, for the first time, we propose a general and principled approach to optimizing the sampling schedules of DMs for high-quality outputs, called Align Your Steps. We leverage methods from stochastic calculus and find optimal schedules specific to different solvers, trained DMs and datasets. We evaluate our novel approach on several image, video as well as 2D toy data synthesis benchmarks, using a variety of different samplers, and observe that our optimized schedules outperform previous hand-crafted schedules in almost all experiments. Our method demonstrates the untapped potential of sampling schedule optimization, especially in the few-step synthesis regime.