Выровняйте свои шаги: оптимизация расписаний выборки в моделях диффузии

Align Your Steps: Optimizing Sampling Schedules in Diffusion Models

April 22, 2024

Авторы: Amirmojtaba Sabour, Sanja Fidler, Karsten Kreis

cs.AI

Аннотация

Модели диффузии (DMs) установили себя как передовой подход к генеративному моделированию визуальной области и за ее пределами. Одним из критических недостатков DMs является их медленная скорость выборки, основанная на множестве последовательных оценок функций через большие нейронные сети. Выборка из DMs может быть рассмотрена как решение дифференциального уравнения через дискретизированный набор уровней шума, известный как расписание выборки. В то время как прошлые работы в основном сосредотачивались на выводе эффективных решателей, мало внимания уделялось поиску оптимальных расписаний выборки, и вся литература опирается на ручные эвристики. В этой работе мы впервые предлагаем общий и принципиальный подход к оптимизации расписаний выборки DMs для получения высококачественных результатов, названный "Align Your Steps". Мы используем методы стохастического исчисления и находим оптимальные расписания, специфичные для различных решателей, обученных DMs и наборов данных. Мы оцениваем наш новаторский подход на нескольких бенчмарках синтеза изображений, видео, а также 2D игрушечных данных, используя различные сэмплеры, и замечаем, что наши оптимизированные расписания превосходят предыдущие ручные расписания практически во всех экспериментах. Наш метод демонстрирует неиспользованный потенциал оптимизации расписания выборки, особенно в режиме синтеза на небольшом числе шагов.

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.