Align Your Steps: Optimizing Sampling Schedules in Diffusion Models

Align Your Steps: Optimizing Sampling Schedules in Diffusion Models

22 Apr 2024 | Amirmojtaba Sabour, Sanja Fidler, Karsten Kreis
This paper introduces Align Your Steps (AYS), a novel framework for optimizing sampling schedules in diffusion models to enhance output quality, especially in few-step synthesis. Diffusion models (DMs) are powerful generative models but suffer from slow sampling due to sequential function evaluations. Sampling from DMs can be viewed as solving a differential equation through a discretized noise schedule. Previous works focused on efficient solvers, but little attention was given to optimizing the sampling schedule. AYS leverages stochastic calculus to find optimal schedules specific to different solvers, trained models, and datasets. The framework minimizes the Kullback-Leibler divergence between the true and linearized generative SDEs to find optimal schedules. AYS is general and applicable to all DMs regardless of data modality. The method is evaluated on various benchmarks, including image, video, and 2D toy data, showing improved output quality compared to hand-crafted schedules. AYS is also effective with different solvers, including stochastic SDE solvers. The results demonstrate that optimizing the sampling schedule leads to significant improvements in generation quality, especially in low NFE regimes. The method is applicable to a wide range of diffusion models and can be integrated with other generative techniques. The broader impact includes faster synthesis, reduced computational demands, and potential applications in creative content generation. However, the use of deep generative models also raises concerns about the production of deceptive imagery and videos.This paper introduces Align Your Steps (AYS), a novel framework for optimizing sampling schedules in diffusion models to enhance output quality, especially in few-step synthesis. Diffusion models (DMs) are powerful generative models but suffer from slow sampling due to sequential function evaluations. Sampling from DMs can be viewed as solving a differential equation through a discretized noise schedule. Previous works focused on efficient solvers, but little attention was given to optimizing the sampling schedule. AYS leverages stochastic calculus to find optimal schedules specific to different solvers, trained models, and datasets. The framework minimizes the Kullback-Leibler divergence between the true and linearized generative SDEs to find optimal schedules. AYS is general and applicable to all DMs regardless of data modality. The method is evaluated on various benchmarks, including image, video, and 2D toy data, showing improved output quality compared to hand-crafted schedules. AYS is also effective with different solvers, including stochastic SDE solvers. The results demonstrate that optimizing the sampling schedule leads to significant improvements in generation quality, especially in low NFE regimes. The method is applicable to a wide range of diffusion models and can be integrated with other generative techniques. The broader impact includes faster synthesis, reduced computational demands, and potential applications in creative content generation. However, the use of deep generative models also raises concerns about the production of deceptive imagery and videos.
Reach us at info@study.space