MaRS: A Fast Sampler for Mean Reverting Diffusion based on ODE and SDE Solvers
Ao Li, Wei Fang, Hongbo Zhao, Le Lu, Ge Yang, Minfeng Xu
TL;DR
MaRS delivers a fast, training-free sampler for Mean Reverting Diffusion by deriving semi-analytical solutions to the reverse-time SDE and the probability flow ODE, coupling an analytic contribution with a neural-network integral estimated via exponential integrators. It supports noise, data, and velocity parameterizations and achieves 10–20× speedups, reducing sampling to 5–10 NFEs while maintaining high sample quality on image restoration tasks. The work highlights data prediction as numerically more stable than noise prediction and provides detailed transformations between parameterizations to ensure broad compatibility. Overall, MaRS makes MR Diffusion more practical for controllable generation, though it does not surpass distillation methods at extremely low NFEs and will release accompanying code for reproducibility.
Abstract
In applications of diffusion models, controllable generation is of practical significance, but is also challenging. Current methods for controllable generation primarily focus on modifying the score function of diffusion models, while Mean Reverting (MR) Diffusion directly modifies the structure of the stochastic differential equation (SDE), making the incorporation of image conditions simpler and more natural. However, current training-free fast samplers are not directly applicable to MR Diffusion. And thus MR Diffusion requires hundreds of NFEs (number of function evaluations) to obtain high-quality samples. In this paper, we propose a new algorithm named MaRS (MR Sampler) to reduce the sampling NFEs of MR Diffusion. We solve the reverse-time SDE and the probability flow ordinary differential equation (PF-ODE) associated with MR Diffusion, and derive semi-analytical solutions. The solutions consist of an analytical function and an integral parameterized by a neural network. Based on this solution, we can generate high-quality samples in fewer steps. Our approach does not require training and supports all mainstream parameterizations, including noise prediction, data prediction and velocity prediction. Extensive experiments demonstrate that MR Sampler maintains high sampling quality with a speedup of 10 to 20 times across ten different image restoration tasks. Our algorithm accelerates the sampling procedure of MR Diffusion, making it more practical in controllable generation.