MRS: 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
2025-02-17
Summary
This paper talks about a new method called MRS (MR Sampler) that makes it faster and easier to create high-quality images using a special type of AI called Mean Reverting (MR) Diffusion. It's like finding a shortcut to make a complex art project much quicker without losing any of the details.
What's the problem?
Current AI methods for creating controllable images are either too slow or not flexible enough. MR Diffusion is good at making controllable images, but it takes too many steps (hundreds of calculations) to create a high-quality image, which makes it impractical for real-world use.
What's the solution?
The researchers created MRS, which uses clever math to solve complex equations related to MR Diffusion. They found a way to break down the problem into two parts: one that can be solved exactly and another that uses a neural network (a type of AI) to estimate. This allows MRS to create high-quality images in much fewer steps than before, without needing any extra training.
Why it matters?
This matters because it makes AI-generated images both faster and more controllable. MRS can create images 10 to 20 times faster than previous methods while keeping the same high quality. This speed boost could make AI image generation more practical for real-world applications, like quickly editing photos or creating custom images for various purposes. It opens up new possibilities for using AI in creative and practical ways that were previously too slow or difficult.
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 MRS (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.