This was a delightful collaboration with Sitan Chen, @jerryzli.bsky.social, @pangwei.bsky.social, and my wonderful advisors Lillian J. Ratliff and
@sewoong79.bsky.social. There’s lots of remaining work to be done on learning to sample that we’re excited to explore!
arxiv.org/abs/2502.17423
7/
S4S-Alt achieves a 1.5-2x improvement over a “vanilla” diffusion ODE solver, particularly with few NFEs, and achieves SOTA performance for techniques that do not modify the underlying DM, while being competitive with training-based methods for a fraction of a compute cost!
6/
Furthermore, to fully explore the DM solver design space, we learn both the discretization steps and the solver coefficients with S4S-Alt, which proceeds through alternating minimization, giving even further quality improvements.
5/
In S4S, we directly learn the coefficients for any diffusion solver by minimizing the global error in the final generated sample. S4S consistently outperforms traditional ODE solvers across model architectures and time discretizations, especially in the few-step regime.
4/
Other prior work fixes the ODE solver and focuses on selecting optimal discretization steps for the reverse process. While this leads to impressive gains, it overlooks the full design space of the diffusion model solver.
3/
Previous work on solving diffusion ODEs focuses on minimizing numerical errors at each step of sequentially solving the ODE. This approach works well with many steps, but breaks down when the number of steps is small.
2/
Want to quickly sample high-quality images from diffusion models, but can’t afford the time or compute to distill them? Introducing S4S, or Solving for the Solver, which learns the coefficients and discretization steps for a DM solver to improve few-NFE generation.
Thread 👇 1/