Differentiable Solver Search for Fast Diffusion Sampling

2025-05-27

AI Summary

Diffusion model sampling requires numerous function evaluations, incurring substantial computational overhead; existing ODE-based solvers predominantly rely on time-$t$-dependent Lagrange interpolation, which has been shown to be suboptimal. This paper proposes a differentiable solver search framework that jointly optimizes ODE solver architecture end-to-end within a compact spatio-temporal differentiable space, thereby uncovering and circumventing the inherent limitations of $t$-dependent interpolation for the first time. Our method integrates differentiable architecture search, ODE parameterization modeling, and gradient-driven coefficient optimization, enabling plug-and-play deployment across diverse diffusion models—including Rectified Flow and DDPM. On ImageNet256, our solvers achieve FID scores of 2.40, 2.35, and 2.33 using only 10 sampling steps for SiT-XL/2, FlowDCN-XL/2, and DiT-XL/2, respectively—significantly outperforming conventional solvers. Moreover, the approach demonstrates strong generalization across model architectures, image resolutions, and model scales.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Deep Neural Architectures and Foundation Models

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Diffusion models have demonstrated remarkable generation quality but at the cost of numerous function evaluations. Recently, advanced ODE-based solvers have been developed to mitigate the substantial computational demands of reverse-diffusion solving under limited sampling steps. However, these solvers, heavily inspired by Adams-like multistep methods, rely solely on t-related Lagrange interpolation. We show that t-related Lagrange interpolation is suboptimal for diffusion model and reveal a compact search space comprised of time steps and solver coefficients. Building on our analysis, we propose a novel differentiable solver search algorithm to identify more optimal solver. Equipped with the searched solver, rectified-flow models, e.g., SiT-XL/2 and FlowDCN-XL/2, achieve FID scores of 2.40 and 2.35, respectively, on ImageNet256 with only 10 steps. Meanwhile, DDPM model, DiT-XL/2, reaches a FID score of 2.33 with only 10 steps. Notably, our searched solver outperforms traditional solvers by a significant margin. Moreover, our searched solver demonstrates generality across various model architectures, resolutions, and model sizes.

Problem

• Optimizing ODE-based solvers for faster diffusion sampling
• Addressing suboptimal t-related Lagrange interpolation in solvers
• Improving computational efficiency without sacrificing generation quality

Innovation

• Differentiable solver search algorithm
• Compact search space optimization
• Generality across model architectures