Fast ODE-based Sampling for Diffusion Models in Around 5 Steps (2312.00094v3)

Abstract: Sampling from diffusion models can be treated as solving the corresponding ordinary differential equations (ODEs), with the aim of obtaining an accurate solution with as few number of function evaluations (NFE) as possible. Recently, various fast samplers utilizing higher-order ODE solvers have emerged and achieved better performance than the initial first-order one. However, these numerical methods inherently result in certain approximation errors, which significantly degrades sample quality with extremely small NFE (e.g., around 5). In contrast, based on the geometric observation that each sampling trajectory almost lies in a two-dimensional subspace embedded in the ambient space, we propose Approximate MEan-Direction Solver (AMED-Solver) that eliminates truncation errors by directly learning the mean direction for fast diffusion sampling. Besides, our method can be easily used as a plugin to further improve existing ODE-based samplers. Extensive experiments on image synthesis with the resolution ranging from 32 to 512 demonstrate the effectiveness of our method. With only 5 NFE, we achieve 6.61 FID on CIFAR-10, 10.74 FID on ImageNet 64$\times$64, and 13.20 FID on LSUN Bedroom. Our code is available at https://github.com/zju-pi/diff-sampler.

• The paper introduces AMED-Solver, reducing function evaluations to around five steps while maintaining high sample quality.
• It leverages geometric insights to approximate the mean direction in sampling trajectories, effectively mitigating truncation errors in ODE solvers.
• Competitive FID scores on CIFAR-10, ImageNet, and LSUN highlight its efficiency and practical impact on generative modeling.

Fast ODE-based Sampling for Diffusion Models in Around 5 Steps: An Analysis

This paper introduces the Approximate MEan-Direction Solver (AMED-Solver), a novel approach aimed at enhancing the sampling efficiency within diffusion models, specifically targeting the reduction of the number of function evaluations (NFE) required to achieve high-quality sample generation. The AMED-Solver presents a significant advancement by demonstrating the potential to minimize NFE to approximately five while producing visually comparable results to more computationally intensive methods.

Numerical Methods in Diffusion Models

Diffusion models have emerged as a prominent tool in generative tasks, offering stable training and high-quality sample generation. The sampling process in diffusion models is synonymous with solving ordinary differential equations (ODEs). Traditional methods often involve extensive computational efforts due to lower-order ODE solvers that require many iterations (or function evaluations) to converge to a solution with acceptable quality. Recent advancements have introduced higher-order solvers, thereby shortening the NFE. However, these advancements typically introduce approximation errors that degrade sample quality when the function evaluations are overly curtailed.

Geometric Insight and the AMED-Solver

A foundational insight driving the AMED-Solver is the recognition that the sampling trajectories in high-dimensional spaces tend to confine themselves within two-dimensional subspaces. This geometric observation allows the solver to approximate the mean direction in sampling more directly, circumventing issues related to truncation errors prevalent in existing techniques. AMED-Solver leverages this by learning the mean direction, thereby optimizing the sampling process while maintaining sample integrity with minimal NFEs.

Comparative Performance and Implications

The paper offers a comprehensive evaluation of AMED-Solver across multiple datasets, such as CIFAR-10, ImageNet, and LSUN bedroom images. With merely five NFEs, AMED-Solver achieves a Fréchet Inception Distance (FID) of 6.61 on CIFAR-10, 10.74 on ImageNet 64x64, and 13.20 on LSUN Bedroom, highlighting its efficacy across varying resolutions.

Furthermore, the research extends AMED into a versatile plugin that can enhance existing ODE solvers. This AMED-Plugin integrates seamlessly with solvers such as the improved Pseudo Numerical Methods for Diffusion Models (iPNDM), yielding substantial improvements, as evidenced by superior FID scores achieved with minimal NFEs.

Practical and Theoretical Implications

Practically, AMED-Solver's ability to function with significantly reduced computational resources without compromising the sample quality makes it suitable for applications where computational efficiency is paramount. Theoretically, the introduction of a geometrically motivated approach to trajectory simplification may prompt further exploration into low-dimensional subspace approximations in diffusion models and other ODE-driven processes.

Future Directions

This work paves the way for future exploration into adaptive sampling schedules and solvers that dynamically respond to trajectory characteristics in real-time, potentially further reducing computational requirements. This research establishes a robust foundation for continued advancements in efficient ODE solvers, contributing to the broader landscape of generative modeling, and sets a precedent for leveraging geometric insights to tackle complex computational problems effectively.

In summary, AMED-Solver represents a promising approach in the field of diffusion models, effectively addressing a quintessential challenge in generative modeling by achieving rapid and efficient sampling. The insights and methods presented herein are poised to influence subsequent research directions, enriching the toolkit available for researchers and practitioners in artificial intelligence and machine learning.