Align Your Steps: Optimizing Sampling Schedules in Diffusion Models (2404.14507v1)

Abstract: 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 $\textit{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.

• The paper presents the 'Align Your Steps' framework that optimizes sampling schedules by minimizing the KL divergence between true SDE trajectories and their discretized approximations.
• The methodology significantly reduces numerical function evaluations while enhancing output quality across multiple datasets and generative tasks.
• The framework demonstrates broad applicability, efficiently adapting to various stochastic solvers and data modalities in diffusion models.

Align Your Steps: A Novel Framework for Optimizing Sampling Schedules in Diffusion Models

Introduction to Diffusion Model Sampling

Diffusion models, particularly in visual tasks, operate by progressively converting data distributions into Gaussian noise and then learning to reverse this process. However, they require numerous neural network evaluations, making them computationally costly. Conventional acceleration strategies either focus on new training methods or on enhanced SDE/ODE solvers, which doesn't directly address the inefficiencies in existing sampling schedules that most models rely on. These schedules, surprisingly, are often chosen heuristically and have not been systematically optimized relative to the specific traits of different solvers, training configurations, or datasets.

Introducing the Align Your Steps Framework

The paper presents a unified and rigorous approach titled "Align Your Steps" (AYS) for optimizing sampling schedules across various stochastic SDE solvers. The key insight is that all stochastic SDE solvers, by nature, can be seen as solving an "approximation" of the complex, data-derived stochastic differential equation (SDE) used during the generative process. The method involves minimizing the divergence between the trajectory distributions of this theoretical SDE and its solver-specific discretized approximation.

1. Core Mechanism:
   • By leveraging stochastic calculus, it quantifies the "distance" between the true and approximated generation process using a calculated Kullback-Leibler divergence upper bound (KLUB). This calculated distance guides the optimization of the sampling schedule to ensure faithful reproduction of the generative process with fewer computational steps.
2. Evaluation and Results:
   • Extensively tested on 2D toy data, standard image datasets (CIFAR10, FFHQ, ImageNet), large-scale text-to-image models (Stable Diffusion, SDXL), and video generative models (Stable Video Diffusion). Outperforms traditional, heuristically designed schedules by generating higher-quality outputs with notably fewer numerical function evaluations (NFEs).

Theoretical Contributions and Practical Implications

• Dependency on Dataset Characteristics: Initially demonstrated on straightforward Gaussian data, this method analytically shows the variability and dependency of optimal schedules based on dataset characteristics.
• Framework Generality: AYS is broadly applicable to any diffusion model, independent of the underlying data modality, making it a universally applicable optimization technique in the generative model toolkit.
• Optimization Efficiency: It is computationally efficient to optimize the schedules, typically requiring fewer than 300 iterations to converge.

Future Work and Speculations

The framework opens several avenues for further exploration:

• Extension to Conditional Generative Models: Applying the AYS framework to optimize schedules in label- or text-conditional diffusion models could potentially refine generation further in discriminative settings.
• Adaptation to Higher-Order Solvers: Investigating the applicability of AYS to higher-order deterministic solvers (like adaptive step-size Runge-Kutta methods) could enhance its utility in scenarios where such solvers are preferred.
• Exploring Broader Applicability: The theory underpinning AYS might extend to other generative paradigms that inherently contain a transition from structured data to a simplified noise model and back, inviting cross-pollination with newer generative approaches such as score-based generative models and energy-based models.

Conclusion

The "Align Your Steps" framework represents a significant stride towards more efficient generative modeling. By rigorously optimizing the sampling schedule, researchers and practitioners can achieve faster and potentially more cost-effective model training and inference cycles, enhancing the practical deployment of diffusion models in real-world applications. The provided schedules and the accompanying empirical results underscore the method's effectiveness and utility across various domains and datatypes.