Alignez vos étapes : Optimisation des calendriers d'échantillonnage dans les modèles de diffusion

Align Your Steps: Optimizing Sampling Schedules in Diffusion Models

April 22, 2024

Auteurs: Amirmojtaba Sabour, Sanja Fidler, Karsten Kreis

cs.AI

Résumé

Les modèles de diffusion (DMs) se sont imposés comme l'approche de modélisation générative de pointe dans le domaine visuel et au-delà. Un inconvénient majeur des DMs réside dans leur lenteur d'échantillonnage, qui repose sur de nombreuses évaluations séquentielles de fonctions via de grands réseaux de neurones. L'échantillonnage à partir des DMs peut être vu comme la résolution d'une équation différentielle à travers un ensemble discrétisé de niveaux de bruit, connu sous le nom de planification d'échantillonnage. Alors que les travaux précédents se sont principalement concentrés sur la dérivation de solveurs efficaces, peu d'attention a été accordée à la recherche de planifications d'échantillonnage optimales, et toute la littérature repose sur des heuristiques artisanales. Dans ce travail, nous proposons pour la première fois une approche générale et rigoureuse pour optimiser les planifications d'échantillonnage des DMs afin d'obtenir des résultats de haute qualité, appelée Align Your Steps. Nous exploitons des méthodes issues du calcul stochastique et trouvons des planifications optimales spécifiques à différents solveurs, DMs entraînés et jeux de données. Nous évaluons notre nouvelle approche sur plusieurs benchmarks de synthèse d'images, de vidéos ainsi que de données jouets en 2D, en utilisant une variété de différents échantillonneurs, et observons que nos planifications optimisées surpassent les planifications artisanales précédentes dans presque toutes les expériences. Notre méthode démontre le potentiel inexploité de l'optimisation des planifications d'échantillonnage, en particulier dans le régime de synthèse à faible nombre d'étapes.

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.