Stochastic Runge-Kutta Methods: Provable Acceleration of Diffusion Models
Yuchen Wu, Yuxin Chen, Yuting Wei
TL;DR
This work delivers a training-free, high-order stochastic Runge-Kutta diffusion sampler that provably accelerates SDE-based diffusion models. By leveraging an exponential integrator and RK-inspired score-difference corrections, it achieves $\\mathsf{KL}$-error $\\varepsilon^2$ with $\\widetilde{O}(d^{3/2}/\\varepsilon)$ score evaluations, improving the prior dimension dependency. Theoretical analysis dissects discretization, score-estimation, and initialization errors, and empirical results on CIFAR-10 and ImageNet-64 confirm practical gains over standard DDPM and prior SDE accelerators. The framework sets the stage for higher-order, training-free diffusion samplers and motivates future work on total-variation guarantees and structure-adaptive diffusion.
Abstract
Diffusion models play a pivotal role in contemporary generative modeling, claiming state-of-the-art performance across various domains. Despite their superior sample quality, mainstream diffusion-based stochastic samplers like DDPM often require a large number of score function evaluations, incurring considerably higher computational cost compared to single-step generators like generative adversarial networks. While several acceleration methods have been proposed in practice, the theoretical foundations for accelerating diffusion models remain underexplored. In this paper, we propose and analyze a training-free acceleration algorithm for SDE-style diffusion samplers, based on the stochastic Runge-Kutta method. The proposed sampler provably attains $\varepsilon^2$ error -- measured in KL divergence -- using $\widetilde O(d^{3/2} / \varepsilon)$ score function evaluations (for sufficiently small $\varepsilon$), strengthening the state-of-the-art guarantees $\widetilde O(d^{3} / \varepsilon)$ in terms of dimensional dependency. Numerical experiments validate the efficiency of the proposed method.