S4S: Solving for a Diffusion Model Solver

TL;DR

This paper introduces S4S, a learned solver for diffusion models that improves sample quality with fewer neural function evaluations, outperforming traditional methods across various datasets and settings.

Contribution

S4S is a novel learned solver that directly optimizes for high-quality diffusion model sampling, outperforming traditional ODE solvers in efficiency and quality.

Findings

S4S improves sample quality across six pre-trained diffusion models.

S4S achieves state-of-the-art results with only 5 NFEs on CIFAR10 and MS-COCO.

S4S-Alt further enhances performance by optimizing both solver and discretization schedule.

Abstract

Diffusion models (DMs) create samples from a data distribution by starting from random noise and iteratively solving a reverse-time ordinary differential equation (ODE). Because each step in the iterative solution requires an expensive neural function evaluation (NFE), there has been significant interest in approximately solving these diffusion ODEs with only a few NFEs without modifying the underlying model. However, in the few NFE regime, we observe that tracking the true ODE evolution is fundamentally impossible using traditional ODE solvers. In this work, we propose a new method that learns a good solver for the DM, which we call Solving for the Solver (S4S). S4S directly optimizes a solver to obtain good generation quality by learning to match the output of a strong teacher solver. We evaluate S4S on six different pre-trained DMs, including pixel-space and latent-space DMs for both conditional and unconditional sampling. In all settings, S4S uniformly improves the sample quality relative to traditional ODE solvers. Moreover, our method is lightweight, data-free, and can be plugged in black-box on top of any discretization schedule or architecture to improve performance. Building on top of this, we also propose S4S-Alt, which optimizes both the solver and the discretization schedule. By exploiting the full design space of DM solvers, with 5 NFEs, we achieve an FID of 3.73 on CIFAR10 and 13.26 on MS-COCO, representing a $1.5\times$ improvement over previous training-free ODE methods.