Formalizing the Sampling Design Space of Diffusion-Based Generative Models via Adaptive Solvers and Wasserstein-Bounded Timesteps

TL;DR

This paper introduces SDM, a geometric and adaptive framework for diffusion model sampling that optimizes solver choice and scheduling, reducing computational costs while maintaining high-quality results.

Contribution

The paper proposes a novel, principled approach to sampling in diffusion models by aligning solver selection with the diffusion trajectory's properties and formalizing adaptive timesteps via Wasserstein bounds.

Findings

Achieves state-of-the-art FID scores on CIFAR-10, FFHQ, and AFHQv2.

Reduces the number of function evaluations needed for high-quality sampling.

Provides a systematic, training-free method for adaptive solver scheduling.

Abstract

Diffusion-based generative models have achieved remarkable performance across various domains, yet their practical deployment is often limited by high sampling costs. While prior work focuses on training objectives or individual solvers, the holistic design of sampling, specifically solver selection and scheduling, remains dominated by static heuristics. In this work, we revisit this challenge through a geometric lens, proposing SDM, a principled framework that aligns the numerical solver with the intrinsic properties of the diffusion trajectory. By analyzing the ODE dynamics, we show that efficient low-order solvers suffice in early high-noise stages while higher-order solvers can be progressively deployed to handle the increasing non-linearity of later stages. Furthermore, we formalize the scheduling by introducing a Wasserstein-bounded optimization framework. This method systematically derives adaptive timesteps that explicitly bound the local discretization error, ensuring the sampling process remains faithful to the underlying continuous dynamics. Without requiring additional training or architectural modifications, SDM achieves state-of-the-art performance across standard benchmarks, including an FID of 1.93 on CIFAR-10, 2.41 on FFHQ, and 1.98 on AFHQv2, with a reduced number of function evaluations compared to existing samplers. Our code is available at https://github.com/aiimaginglab/sdm.