Stochastic Runge-Kutta Methods: Provable Acceleration of Diffusion Models

TL;DR

This paper introduces a provably efficient stochastic Runge-Kutta method to accelerate diffusion model sampling, reducing computational cost while maintaining high sample quality, with theoretical guarantees and empirical validation.

Contribution

It presents a training-free acceleration algorithm for diffusion samplers based on stochastic Runge-Kutta methods, with improved theoretical complexity bounds over prior approaches.

Findings

Achieves $ ilde{O}(d^{3/2}/ ext{epsilon})$ score evaluations for small epsilon

Improves theoretical guarantees from $ ilde{O}(d^{3}/ ext{epsilon})$ to $ ilde{O}(d^{3/2}/ ext{epsilon})$

Numerical experiments confirm the efficiency of the proposed method.

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.