MaRS: A Fast Sampler for Mean Reverting Diffusion based on ODE and SDE Solvers
Authors: Ao Li
Wei Fang
Hongbo Zhao
Le Lu
Ge Yang
Minfeng Xu
Venue: ICLR 2025 Spotlight (2025)
Keywords: Fast Sampler
Mean Reverting Diffusion
Primary area: generative models
TL;DR: We propose a fast sampler for Mean Reverting Diffusion based on both ODE and SDE solvers.

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.

Published as a conference paper at ICLR 2025
MARS: A FAST SAMPLER FOR MEAN REVERTING DIFFUSION BASED ON ODE AND SDE SOLVERS
Ao Li 1,2s
Wei Fang 3,1 Hongbo Zhao 1.2* Le 3
Ge Yang 1.2f Minfeng Xu 3,4f
School of Artificial Intelligence, University of Chinese Academy of Sciences, Beijing. China "Institute of Automation, Chinese Academy of Sciences (CASIA) 3 DAMO Academy, Alibaba Group 'Hupan Laboratory. Hangzhou, China
f11ao2022, zhaohongbo2 C22, g yang)eia ac. en
lucas. fwwal Mataaaalaaaa-incc com minfengxue163.. com tiger,
ABSTRACT
In applications of diffusion models, controllable generation of practical sig- nificance, but is also challenging. Current methods for controllable generation primarily foeus on modifying the score function of diffusion models, while Mean Reverting (MR) Diffusion directly modifies the structure of the stochastic differ- ential equation (SDE), making the incorporation of image conditions simpler and more natural. However, current training-free fast samplers are not directly ap- plicable 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 de- rive semi-analytical solutions. The solutions consist of an analytical function and an integral parameterized by neural network. Based on this solution, we can gen- erate 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 speedup of 10 to 20 times across ten different image restoration tasks. Our algorithm accelerates the sampling pro- cedure of MR Diffusion, making it more practical in controllable generation,
INTRODUCTION
Diffusion models have emerged a powerful class of generative models, demonstrating remarkable capabilities across a variety of applications, including image synthesis (Dhariwal & Nichol [001]] Ruiz al., 2023 Rombach al. 2022)) and video generation (Ho al., 2022a b). In these applica- tions, controllable generation is very important in practice, but it also poses considerable challenges. Various methods have been proposed incorporate text image conditions into the score function of diffusion models(Ho & Salimans 2022, Ye et al. 2023 Zhang et al., 2023), whereas Mean Re- verting (MR) Diffusion offers a new avenue of control in the generation process (Luo et al. 2023b). Previous diffusion models (such as DDPM (Ho et al.] 2020)) simulate diffusion process that grad- ually transforms data into pure Gaussian noise, followed by learning to reverse this process for sample generation (Song & Ermon 2020 Song et al. (021). In contrast, MR Diffusion is designed to produce final states that follow a Gaussian distribution with a non-zero mean, which provides simple and natural way to introduce image conditions. This characteristic makes MR Diffusion particularly suitable for solving inverse problems and potentially extensible to multi-modal condi- tions. However, the sampling process of MR Diffusion requires hundreds of iterative steps, which is time-consuming.
$ "Thee authors contributed equally to this work. 1 CCoresponding authors. Code is available at http::///gtttbb..co/grrreeeeeeee

---

Published as a conference paper at ICLR 2025
Dehazing
NFE=5
NFE=20
NFE=50
NFE=00
NFE=100
Ground Truth
Posterior Sampling
MR Sampler-2
Inpainting
NFE=5
NFE=20
NFE=50
NFE=80
NFE=100
Ground Truth
Posterior Sampling
MR Sampler-2
Figure 1: Qualitative comparisons between MR Sampler and Posterior Sampling, All images generated by sampling from pre-trained MR Diffusion (Luo et al. 2024a) the RESIDE-6k
(Qin et al. 2020b) dataset and the CelebA-HQ Karras 2017) dataset.
To improve the sampling efficiency of diffusion models, various acceleration strategies have been proposed, which can be divided into two categories. The first explores methods that establish di- rect mappings between starting and ending points on the sampling trajectory, enabling acceleration through knowledge distillation (Salimans & Ho 2022, Song et al. 2023 Liu et. (0226). However, such algorithms often come with trade-offs. such as the need for extensive training and limitations in their adaptability across different tasks and datasets. The second category involves the design of fast numerical solvers that increase step sizes while controlling truncation errors, thus allowing for
faster convergence to solutions (Lu al. 2022a Zhang & Chen 2022] Song et al. 2020a).
Notably, fast sampling solvers mentioned above are designed for common SDEs such as VPSDE and VESDE (Song et al. 2020b). Due to the difference between these SDEs and MRSDE, exist- ing training-free fast samplers cannot be directly applied to Mean Reverting (MR) Diffusion. In this paper, we propose novel algorithm named MaRS (MR Sampler) that improves the sampling efficiency of MR Diffusion. Specifically, solve the reverse- time stochastic differential equation (SDE) and probability flow ordinary differential equation (PF-ODE) (Song et al. 2020b) derived from MRSDE, and obtain a semi-analytical solution, which consists of an analytical function and an integral parameterized by neural networks. We prove that the difference of MRSDE only leads to change analytical part of solution, which can be calculated precisely. And the integral part can be estimated by discretization methods developed in several previous works (Lu et al., 2022a Zhang Chen 2022; Zhao et al. 2024). We derive sampling formulas for two types of neural network parameterizations: noise prediction (Ho et al.) 22200] [ong et al.) (020b) and data prediction (Sali-) mans & Ho 2022). Through theoretical analysis and experimental validation, we demonstrate that data prediction exhibits superior numerical stability compared to noise prediction. Additionally,
propose transformation methods for velocity prediction networks (Salimans & Ho 2002) so that our algorithm supports all common training objectives. Extensive experiments show that our fast sam- pler converges in 5 or 10 NFEs with high sampling quality. As illustrated in Figure T our algorithm achieves stable performance with speedup factors ranging from 10 to 20.
In summary, our main contributions are as follows:
We propose MR Sampler, a fast sampling algorithm for MR Diffusion, based on solving the PF-ODE and SDE derived from MRSDE. Our algorithm is plug-and-play and can adapt to all common training objectives.
2

---

Published as a conference paper at ICLR 2025
We demonstrate that posterior sampling (uo et al. 2024b) for MR Diffusion is equiva- lent to Sarer-Maruyama discretization, whereas MR Sampler computes a amii-nnayttiaa solution, thereby eliminating part of approximation enrors.
Through extensive experiments ten image restoration tasks, we demonstrate that MR Sampler can reduce the required sampling time by factor of 10 to 20 with comparable sampling quality. Morcover, reveal that data prediction exhibits superior numerical stability compared to noise prediction.
BACKGROUND
In this section, we briefly review the basic definitions and characteristics of diffusion probabilistic models and mean-reverting diffusion models.
2.1 DIFFUSION PROBABILISTIC MODELS
According to Song et al. (020b), Diffusion Probabilistic Models (DPMs) can be defined as the solution of the following Itō stochastic differential equation (SDE), which is stochastic process (t with T > 0, called forward process, where I e IR is D-dimensional random variable.
da f(x, f) + g(t)dw.
The forward process performs adding noise to the data Io. while there exists a corresponding reverse process that gradually removes the noise and recovers Zp- Anderson (1982) shows that the reverse of the forward process is also a solution of an Itō SDE:
da [f(x, ( log (x))dd + g())dw,
where f and the drift and diffusion coefficients respectively, t is standard Wiener process running backwards in time, and time flows from 7 to 0, which means dl. < 0. The score function v log Pe (ar) is generally intractable and thus neural network sp () is used to estimate it by optimizing the following objective (Song et al. 2020b Hyvärinen & Dayan 2005):
0* arg min E. Fan Te lea
[100 (2:, t) Va,, log logp(ar, (20)
(3)
where X(t) [0, 77 R is positive weighting function, is uniformly sampled over T1, 20 Po (a) and x; p(2:/200) To facilitate the computation of p(æ, (a)) the drift coefficient f(x,, L) typically defined linear function of æ, as presented in Eq.44). Based on the inference by Särkkä & Solin (2019) in Section 5.5, the transition probability p(æ;, (2)) corresponding to Eq. follows Gaussian distribution, as shown 55.
dæ f(t)adt g(f)dw,
(4)
S 2 S(ErdE seraa (T))) 1):
p( N o;; wh ((r)dr
Krraar
(5)
Song al. (020b) proved that Denoising Diffusion Probabilistic Models (Ho al. 2020 and Noise Conditional Score Networks (Song & Ermon, 2019) can be regarded as discretizations of Variance Preserving SDE (VPSDE) and Variance Exploding SDE (VESDE), respectively. As shown in Table [ the SDEs corresponding to the two most commonly used diffusion models both follow the form of Eq.
Table 1: Two popular SDEs, Variance Preserving SDE (VPSDE) and Variance Exploding SDE (VESDE), m(t) and vtt) refer to mean and variance of the transition probability p(æ)
SDE
S())
g(()
ml)
v(1)
VPSDE Ho et al. 2020
-)t())
/ [tt) d[o*(0)
The ((7)dr
Ie 53 ((r)dr
VESDE(Song & Ermon 2019
0
[ o0) I
3

---

Published as a conference paper at ICLR 2025
2.2
MEAN REVERTING DIFFUSION MODELS
Luo et al. (2033)) proposed a special case of Itō SDE named Mean Reverting SDE (MRSDE), as follows:
dæ = S(() (m - æ) di + s(()dw,
(6)
where 1a is parameter vector that has the same shape of variable æ, and s(), g(t) are time- dependent non-negative parameters that control the speed of the mean reversion and stochastic volatility, respectively. To prevent potential confusion, we have substituted the notation used in
the original paper (Luo et al., 2023b). For further details, please refer to Appendix [B] Under the assumption that (t) 20 for any E [0, 7] with T O, Eq. has a closed-form solution,
given by
3t.- Ene f(r)dr + 1(1 - (o))dr + o V e- /(T)rr
where doe is a positive hyper-parameter that determines the standard deviation of x, when - 6o and z ~ ((, J). Note that æt starts from Fy. and converges to μ + 02 2 / - According to Anderson 1982 's result, we can derive the following reverse-time SDE:
dæ = [f(t) (m x) - (t) pog) dt + g(t)dw.
(8)
Similar to DPMs, the score function in (q.(8) can also be estimated by score matching methods Song & Ermon 2019); Song et al.] (221)) Once the score function is known, we can generate wo from noisy state ær. In summary, MRSDE illustrates the conversion between two distinct types
of data and has demonstrated promising results in image restoration tasks (Luo et al. (023c).
Various algorithms have been developed to accelerate sampling of VPSDE, including methods like CCDF (Chung et al., (202)), DDIM (Song et al. (0200)). PNDM (Liu et al., 2022a)) DPM-Solver (Lu al. 2022a) and UniPC (Zhao al., 2024). Additionally, Karras et al. (2022) and Zhou et al. (2024) have introduced techniques for accelerating sampling of VESDE. However, the drift coefficient of VPSDE and VESDE is linear function of æ. while the drift coefficient in MRSDE is an affine function w.r.t T. adding intercept fe (see (q.(4) and (6) Therefore, current sampling acceleration algorithms cannot be applied MR Diffusion. To the best of our knowledge, MR Sampler has been the first sampling acceleration algorithm for MR Diffusion far.
3. FAST SAMPLERS FOR MEAN REVERTING DIFFUSION WITH NOISE PREDICTION
According to Song et al. (2020b), the states e; the sampling procedure of diffusion models cor- respond to solutions of reverse-time SDE and PF-ODE. Therefore, look for ways accelerate sampling by studying these solutions. In this section, we solve the noise- prisseeredittion- prediction-basee based reverse- time SDE and PF-ODE, and we numerically estimate the non-closed-foo tion, which serves to accelerate the sampling process of MR diffusion models. Next, we analyze the sampling method currently used by MR Diffusion and demonstrate that this method corresponds to a variant of discretization for the reverse-time MRSDE.
component of the solu-
3.1 SOLUTIONS TO MEAN REVERTING SDES WITH NOISE PREDICTION
Ho al. (2020) reported that score matching can be simplified to predicting noise, and (Song et al.) 2020b revealed the connection between score function and noise prediction models, which is
E ca( 11,
V sy log pl (2))))))) (o)) =
where o, Gz V 2 - f(rr
is the standard deviation of the transition distribution p(a,|=0) Because μ is independent of t and æ, we substitute eœ( μr, t) with ea(x,, t) for notation simplicity. According to Eq. we can rewrite (q. as
dæ = f(1) (y x) (2())
at co 1) dt g(t)dw.
(10)
Using Itô's formula (in the differential form), obtain the following semi-analytical solution:
4

---

Published as a conference paper at ICLR 2025
Proposition 1. Given an initial value ag at time [0,T], the solution 2t at time t [0,s] of Eq.I00 T0,11O
is
Og + (1 11 at 02,
p at J: (²) co(2T: dr 0,07
ai o of
(T))
(11)
where we denote ae == e %o (())dr and x NO,, I). The proof in Appendix A.1
However, the integral with respect to neural network output is still complicated. There have been several methods (Lu et al. 2022a Zhang & Chen 2022) Zhao et al., 2024) to estimate the integral nu- merically. We follow Lu et al. (2022b)ss method and introduce the half log-SNR R A; log(as/od Since both (t) and g(t) are deliberately designed ensure that a₁ is monotonically decreasing over and o is monotonically increasing over t. Thus, X, is strictly decreasing function of and there exists an inverse function ((A). Then we can rewrite g(r) Eq. 11) as
g²( = (Tr) = 2f(T)(o² + az) = (0 f
07
(12)
do? - 20 ddr
= =202( + 20? dr
dr
By substituting Eq. 12 into Eq. we obtain
Ar
Ct
1
at
20h
e e EO AJdA + ony (c2(de-do) - 1)2, (13)
Ox
where a ) ( A) :- 6 (0( ((),)) According to the methods of exponential in- tegrators (Hochbruck & Ostermann 2010 2005), the (k- 1)-th order Taylor expansion of co(xx, and integration- cictiiooo--y-rrrree by-parts of the integral part in Eq. yields
]
h e co A)dA -2
22
(h) m
-204
(n) Co
(x), Ao)
e
O(h+t),, (14)
m!!
where h :- - We drop the discretization error term O(h++1) and estimate the derivatives with backward difference method. We name this algorithm MR Sampler-SDE-n-k, where means noise prediction and k is the order. We present details in Algorithm anddd
3.2 SOLUTIONS TO MEAN REVERTING ODES WITH NOISE PREDICTION
Song et al. (2020b) have illustrated that for any Itô SDE, there exists probability flow ODE, sharing the same marginal distribution (r(x) as a reverse-time SDE. Therefore, the solutions of PF-ODEs
are also helpful in acceleration of sampling. Specifically, the PF-ODE corresponding to Eq.
dæ dt
f(t) (μ x) g²()) co)
(15)
20,
The aforementioned equation exhibits a semi-linear structure with respect to thus permitting resolution through the method of "variation of constants". We can draw the following conclusions:
Proposition 2. Given an initial value as at time 15) is
[0, the solution x at time t [0,s] of
Ot x.
(1-0t Ct
1
(²(r) o, co co T)dT,
Et
1 at
(16)
as
ag
where a; := e 15 f(r)dr The proof is in Appendix A.1
Then we follow the variable substitution and Eq.
14) 12
in Section 3.1 and we obtain
La +
0g (1-
a1
Σ n=0
(n) 60
(h)m + +(h+t), (17)
1
oe
0i.
m!
m==
(n) (xx, A) (= is the n-th order total derivatives of co with respect to X. By
where Co
dropping the discretization error term Oh*++)) and estimating the derivatives of cœ As) with backward difference method, we design the sampling algorithm from the perspective of ODE (see
Algorithm 4
5

---

Published as a conference paper at ICLR 2025
3.3 POSTERIOR SAMPLING FOR MEAN REVERTING DIFFUSION MODELS
In order to improve the sampling process of Mean Reverting Diffusson, Luo et al. (2024b) proposed the posterior sampling algorithm. They define a monotonically increasing time series (Li)-- and the reverse process as a Markov chain:
T
p( p(Z1r Øa) p(ær xD) I
p(x i-2
Is 2o) and ær N),,,
(18)
where we denote a; == 21, for simplicity. They obtain an optimal posterior distribution by mini- mizing the negative log- ligliielihhod, which is Gaussian distribution given by
p(a,-11 2i x0) = (N (x, 20), B,1),
(1 - ag 1)a
(xi 14)
1 (10 11) ,s,
- 03)
o
(19)
8. az-1)(1
a²
0
1 oz
where a; -c- - f(r)dr and To (x; 1 o;co(Z;; μ4, (t))//a; + μ.. Actually, the reparame- terization of posterior distribution Eq. (9) is equivalent to a variant of the Euler-Maruyama dis- cretization of the reverse-time SDE (see details in Appendix(A.2) Specifically, the Maruyaaa method computes the solution in the following form:
m7 F. + ['[(o)(-=.))+ g²(T) ( (o dr + S. g(T))d
x.)
(20)
or
which introduces approximation errors from both the analytical term and the non non-linear component associated with neural network predictions. In contrast, our approach delivers exact solution for the analytical part, leading to reduced approximation errors and a higher order of convergence.
FAST SAMPLERS FOR MEAN REVERTING DIFFUSION WITH DATA PREDICTION
Unfortunately, the sampler based noise prediction can exhibit substantial instability, particularly with small NFEs, and may perform even worse than posterior sampling. It is well recognized that the Taylor expansion has limited convergence domain, primarily influenced by the derivatives of the neural networks. In fact, higher-order derivatives often result smaller convergence radii. Dur- ing the training phase, the noise prediction neural network is designed to fit normally distributed Gaussian noise. When the standard deviation of this Gaussian noise is set to 1, the values of sam- ples can fall outside the range of [-1, 1] with probability of 34.74%. This discrepancy results in numerical instability in the output of the neural network, causing its derivatives to exhibit more pronounced fluctuations (refer to the experimental results in Section 5 for further details). Conse- quently, the numerical instability leads to very narrow convergence domains, or in extreme cases, no convergence all, which ultimately yields awful sampling results.
Lu et al. (022b) have identified that the choice of parameterization for either ODEs or SDEs is critical the boundedness of the convergent solution. In contrast noise prediction, the data prediction model (Salimans & Ho 2022) focuses on fitting æp, ensuring that its output remains strictly confined within the bounds 1-1,1 1]. thereby achieving high numerical stability.
4.1 SOLUTIONS TO MEAN REVERTING SDES WITH DATA PREDICTION
According to Eq. we can parameterize To as follows:
0g ( 1) (1 oz)
co(x),,)
(21)
Ø:

---

Published as a conference paper at ICLR 2025
By substituting Eq.(21 (1
into Eq.(0)
we derive the following SDE that incorporates data prediction:
(10()) of f(t) æ -(t)
(1)
(2( (1 (1-ai) F of a 8) dt (t) (22) of
dæ
f(t)
This equation remains remi-linear with respect to a and thus we can employ Itô's formula (in the differential form) to obtain the solution to Eq. (2))
Proposition Given an initial value 2a at time [0, the solution 2, at time [0, s) of
Eq. 22 is
Tt -de--ee Et
x, 1
2(k - at a,e 2(A--A))))
0s
a₂
(23)
dr e -(2 x (x
+ +20g
10 Ald1 + or -o.VI -00--e2200 -((11-A))
where x (). The proof is in Appendix A.I
Then we apply Taylor expansion and integration- barrtttss---- to estimate the integral part in Eq 2
and
obtain the stochastic sampling algorithm for data prediction (see details in Algorithm 5 and
4.2 SOLUTIONS TO MEAN REVERTING ODES WITH DATA PREDICTION
By substituting Eq. diction.
21 into Eq.
15
we can obtain the following ODE parameterized by data pre-
dæ dt =
(1) 20
æ [0
p0()) (1 (1-aa) 20?
on) (1- μe 20?
g²(1))
1(t)
a,zo( t).
(24)
The incorporation of the parameter does not disrupt the semi-linear structure of the equation with respect to æ; and y is not coupled to the neural network. This implies that analytical part of solutions can still be derived concerning both x and μ.. We present the solution below (see Appendix A. for a detailed derivation).
Proposition 4. Given an initial value Da at time Eq. 24) is
[0, the solution D; at time t
[0, s] of
+ az S, é xx (xx,, A)d2.
C F. + p1
05 a0
(25)
Similarly, only the neural network component requires approximation through the exponential inte- grator method (Hochbruck & Ostermann) 20055 2010). And we can obtain the deterministic sam- pling algorithm for data prediction (see Algorithm and for details).
4.3 TRANSFORMATION BETWEEN THREE KINDS PARAMETERIZATIONS
There are three mainstream parameterization methods. Ho et al. (2020) introduced a training objec- tive based on noise prediction, While|Saliman Ho (2022) proposed parameterization strategies for data and velocity prediction to keep network outputs stable under the variation of time or log-SNR. All three methods can be regarded as score matching approaches (Song et al. 2020bb Hyvärinen & Dayan 2005) with weighted coefficients. To ensure our proposed algorithm is compatible with these parameterization strategies, is necessary to provide transformation formulas for each pairs among the three strategies.
The transformation formula between noise prediction and data prediction can be easily derived from Eq.
x s. o1-
(26)
(t)))) 2t-0;-0(t--
d:
a. which is slightly different from the def- For velocity prediction, we define de := arctan( o.n.
inition of Salimans & Ho (2022). Then we have at cos cos the at as sin oi and hence at cos 0; + 11(1 cos Φt) + Tea sin (dt).. And the definition of vt)) is
dx: siin Φ. mo sin or + ax cos(
(27)
do,

---

Published as a conference paper at ICLR 2025
If we have score function model (t) trained with velocity prediction, obtain (t and Ea(t) by (see Appendix A.3 for detailed derivations)
(t) 2y cos Φ1 + fa cos Φ1) vo sin Φ2,
(28)
co(t) = (t)) cos Φe + æ sin de - pesin @t)
(29)
EXPERIMENTS
In this section, we conduct extensive experiments show that MR Sampler can significantly speed up the sampling of existing MR Diffusion. To rigorously validate the effectiveness our method, we follow the settings and checkpoints from Luo et al. (2024a) and only modify the sampling part. Our experiment is divided into three parts. Section 5.] compares the sampling results for different NFE cases. Section 5.2 studies the effects of different parameter settings on our algorithm, including network parameterizations and solver types. In Section 5.3) we visualize the sampling trajectories to show the speedup achieved by MR Sampler and analyze why noise prediction gets obviously worse when NFE is less than 20.
5.1
MAIN RESULTS
Following Luo 2024a). we conduct experiments with ten different types of image degradation:
blurry, hazy. (see Appendix D.I for details). We adopt LPIPS (Zhang et al. 2018) and FID (Heusel et al.) (017) as main metrics for perceptual evaluation, and also report PSNR and SSIM (Wang et al.) (2004) for reference. We compare MR Sampler with other sampling methods, including posterior sampling (Luo et al. 2024b) and Buler-Maruyama discretization (Kloeden al. 1992). We take tasks as examples and the metrics are shown in Figure 2 Unless explicitly mentioned, we always use MR Sampler based SDE solver, with data prediction and uniform X. The complete experimental results can be found in Appendix DD33 The results demonstrate that MR Sampler converges in a few (5 10) steps and produces samples with stable quality. Our algorithm significantly reduces the time cost without compromising sampling performance, which is of great practical value for MR Diffusion.
low-light, noisy, raindrop, rainy, shadowed, snowy, and inpainting
5.2 EFFECTS OF PARAMETER CHOICE
In Table 2 we compare the results of two network parameterizations. The data prediction shows stable performance across different NFEs. The noise prediction performs similarly to data prediction with large NFEs, but its performance deteriorates significantly with smaller NFEs. The detailed analysis can be found in Section 5.3 In Table 3 we compare MR Sample--ODE-d- and MR Sampler-SDEEdd22 on the impainting task, which are derived from PF-ODE and reverse-time SDE respectively. SDE-based solver works better with a large NFE, whereas ODE-based solver is more effective with small NFE. In general, neither solver type is inherently better.
Table 2: Ablation study of network parameteriza- tions on the Rain100H dataset.
Table 3: Ablation study of solver types on the CelebA-HQ dataset.
NFE
Parameterization
LPIPSI
FIDI
PSNRT
SSIM
NFB
Solver Type
LPIPSU
FID
PSNR?
SSIMT
50
Noise Prediction
Data Prediction
0.0606
0.0620
27.28
27.65
28.85
28.89
0.8602
0.8615
50
ODE SDE
0.0402
U.0499
22.91
19.09
29.15
28.49
0.9046
0,8921
20
Noise Prediction
0.1929
47.31
27.68
0.7954
20
ODE
0.0475
21.35
28.51
0.8940
Data Prediction
0.0635
27.79
28.60
0.8559
SDE
0.0408
19.13
28.98
0.9032
10
Noise Prediction
Data Prediction
0.0678
1.376
402.3
29.54
28.09
6.623
0.8483
0.0114
10
ODE
SDB
0.0437
0.0417
19.29
19.44
28.45
28.94
0.8996
0,9048
Noise Prediction
1.416
447.0
5.755
0.0051
ODE
0.0526
27.44
31.02
0.9335
5
Data Prediction
0.0637
26.92
28.82
0.8685
SDE
0.0529
24.02
28.35
0.8930
5.3
ANALYSIS
Sampling trajectory. Inspired by the design idea of NCSN (Song & Ermon 2019), we provide
new perspective of diffusion sampling process. Song & Ermon (019) consider each data point
8

---

Published as a conference paper at ICLR 2025
13337
44 1666
421.17
Postrefor Sumpling
Fatee Disrrvebratinn
1.0
1.297
390.1
Mil Samplee-1
260
MB Samplee-2
0.9
Posterier Sampling
0,752
17.00
Buler Discetization
MiI Samplard
MM Sumplar-2
F10
120
131,32
0.3
50
41.51
35
3603
41.00
35.28
34:65
$2.91
34.70
37.40
0.009
0.078
0.003
02854
0.00a
0.079
0.07
0.078
0,077
0.101
0,000
10
20
50
100
10
20
50
100
NFE
NEE
(a) FID on low-light dataset
(b) LPPPS on low-light dataset
200
Posisciur lanpling
velor
15 2m
141.00
120.35
Euler Discretiaation
Somplee-
1.54
Euler Disceelicatio
Sompledd
100
Samplee-2
0.8
1.69
Sampler-
61.21
LPIPSI
FTD:
0.4
0.51
45.13
15.65
14,92
14:47
14.61
16,00
0.042
0.329
0.140
14220
0.027
10
0.1
NFE
(c) FID on motion-blle notion-bburry datasct
(d) LPIPS on motion-blurry datasct
Figure 2: Perceptual evaluations on low-light and motion-blurry datasets.
Motion luny
Postarior Sampling (10)
MR Sampler1 (10)
Pestarior Sumpliagr
Postertor Samplingri0:
MB Suep 2210
Ground Truth
Euler Discretization (10)
MR Sampler-2 (10)
-100
(a) Sampling results.
(b) Trajectory.
Figure 3: Sampling trajectories. (a), we compare our method (with order and order 2) and previous sampling methods (i.e., posterior sampling and Euler discretization) on motion blurry im- age. The numbers parentheses indicate the NFE. (b), we illustrate trajectories of each sampling method. Previous methods need to take many unnecessary paths to converge. With few NFEs, they fail to reach the ground truth (i.e., the location of zo). Our methods follow a more direct trajectory.
(e.g., an image) as a point in high- nigh-dimensional space. During the diffusion process, noise is added to each point &p, causing to spread throughout the space, while the score function (a neural net- work) remembers the direction towards In the sampling process, we start from random point by sampling a Gaussian distribution and follow the guidance of the reverse-time SDE (or PF-ODE) obtain sam-
and the score function to locate Zo. By connecting each intermediate state space, making it difficult to
pling trajectory. However, this trajectory exists in a high-dimensione Principal Component Analysis (PCA) reduce two dimensions,
visualize. Therefore, obtaining the projection of the sampling trajectory in 2D space. As shown Figure 3 we present example. Previous sampling methods (Luo al. 2024b) often require long path find æn.

---

Published as a conference paper at ICLR 2025
and reducing NFE can lead cumulative errors, making it impossible to locate Tg. In contrast, algorithm produces more direct trajectories, allowing us to find 2o with fewer NFEs.
Low light
Noise prediction (15)
1.0
Noise prediction
0.8
Data prediction
0.6
Ground prediction (155)
0.4
0.2
6 8
10
12
Index timestep
(a) Sampling results.
Ratio of convergence.
Figure 4: Convergence of noise prediction and data prediction. In (a), we choose a low-light image for example. The numbers in parentheses indicate the NFE. (b), we illustrate the ratio of components of neural network output that satisfy the Taylor expansion convergence requirement.
Numerical stability of parameterizations. From Table noise prediction in the case of few NFEs. The reason may be that the neural network parameterized by noise prediction numerically unstable. Recall that we used Taylor expansion in Eq. (14,, and the condition for the equality to hold is J Aal R(s). And the radius of convergence R( can be calculated by
observe poor sampling results for
Cn+1 (t)
lim
(30)
R(I)
C, (1)
where C, is the coefficient of the n-th term in Taylor expansion. We are unable to compute this limit and can only compute the 0 case approximation. The output of the neural network can be viewed as a vector, with each component corresponding to a radius of convergence. At each time step, count the ratio of components that satisfy (s(s) |> As criterion for judging the convergence, where i denotes the i-th component. As shown Figure 4 the neural network parameterized by data prediction meets the convergence criteria almost every step. However, the neural network parameterized by noise prediction always has components that cannot converge, which will lead large errors and failed sampling. Therefore, data prediction has better numerical stability and is more recommended choice.
CONCLUSION
We have developed the fast sampling algorithm of MR Diffusion. Compared with DPMs, MR Diffusion is different in SDE and thus not adaptable to existing training-free free fast samplers. We propose MR Sampler for acceleration of sampling of MR Diffusion. We solve the reverse-time SDE and PF-ODE derived from MRSDE and find semi-analytical solution. We adopt the methods of exponential integrators estimate the non-linear integral part. Abundant experiments demonstrate that algorithm achieves small errors and fast convergence. Additionally, visualize sampling trajectories and explain why the parameterization of noise prediction does not perform well in the case of small NFEs.
Limitations and broader impact. Despite the effectiveness of MR Sampler, our method is still inferior to distillation methods (Song et al., 2023 Luo et al., 2023a) within less than 5 NFEs. Addi- tionally, our method can only accelerate sampling, but cannot improve the upper limit sampling quality.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
Our codes are based on the official code of MR Diffusion (Luo et al. 2023b) and DPM-Solver (Lu al. 2022b). And use the checkpoints and datasets provided by MR Diffusion (Luo et al.
2023b), We will release them after the blind review.
ACKNOWLEDGEMENT
This work was supported in part by the National Natural Science Foundation of China (grant 92354307), the National Key Research and Development Program of China (grant 2024YFF0729202), the Strategic Priority Research Program of the Chinese Academy of Sciences (grant XDA0460305), and the Fundamental Research Funds for the Central Universities (grant E3E45201X2). This work was also supported by Alibaba Group through Alibaba Research Intern Program.

---

Published as a conference paper at ICLR 2025
REFERENCES
Eirikur Agustsson and Radu Timofte. Ntire 2017 challenge on single image super-resolution; dataset and study. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, pp. 126-135, 2017.
Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Ap- plications, 12(3):313-322 1982.
Hyungjin Chung, Byeongsu Sim, and Jong Chul Ye. Come-cllser-ddifus--fass Accelerating con- ditional diffusion models for inverse problems through stochastic contraction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Paltern Recognition, pp. 12413-12422, 2022.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-877 8794, 2021.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho and Tim Salimans. arXiv:2207. 12598, 2022.
Classifier-fre diffusion guidance.
arXiv preprint
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Jonathan Ho. William Chan, Chitwan Saharia, Jay Whang. Ruiqi Gao, Alexey Gritsenko, Diederik Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, et al. Imagen video: High definition video generation with diffusion models. arXiv preprint 1rX:v:2210.02303, 2022a.
Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. Advances Neural Information Processing Systems, 35:8633- 8646, 2022b.
Marlis Hochbruck and Alexander Ostermann. Explicit exponential runge-kutta methods for semi- linear parabolic problems. SIAM Journal on Numerical Analysis, 43(3):1069-1099 2005.
Marlis Hochbruck and Alexander Ostermann. Exponential integrators. Acta Numerica, 19:209-286, 2010.
Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score match- ing. Journal of Machine Learning Research, 6(4), 2005.
Tero Karras. Progressive growing of gans for improved quality, stability, and variation. arXiv preprint arXiv:1710. 10196, 2017.
Tero Karras, Miik Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in neural information processing systems, 35:26565-2633 2022.
Peter E Kloeden, Eckhard Platen, Peter Kloeden, and Eckhard Platen. Stochastic differential equations. Springer, 1992.
Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. arXiv preprint arXiv:2202.09775 097778 2022a.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. arXiv preprint Xiiv:2209.0300) 2022b.
Yun-Fu Liu, Wei Jaw, Shih-Chia Huang, and Jenq-Neng Hwang. Desnownet: Context-aware deep network for snow removal. IEEE Transactions on Image Processing, 27(6):3064-3077 2018.
12

---

Published as a conference paper at ICLR 2025
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 1755775-5787, 2022a.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint arX00:221..01955, 2022b.
Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11461-11471, 2022.
Simian Luo, Yiqin Tan, Longbo Huang, Jian Li, and Hang Zhao. Latent consistency models: Synthe- sizing high-resolution images with few-step inference. arXiv preprint arXiv:2310.04378. 2023a.
Ziwei Luo. Fredrik K Gustafsson, Zheng Zhao, Jens Sjölund, and Thomas B Schön. Image restora- tion with mean-reverting stochastic differential equations. arXiv preprint arXiv:2301.11699 2023b.
Ziwei Luo, Fredrik K Gustafsson, Zheng Zhao, Jens Sjölund, and Thomas Schön. Refusion: Enabling large-size realistic image restoration with latent-space diffusion models. In Proceedings of the JEEE/CVE conference on computer vision and paltern recognition, pp. 1680-1691, 2023c.
Ziwei Luo, Fredrik K Gustafsson, Zheng Zhao, Jens Sjölund, and Thomas B Schön. Controlling vision-language models for multi-task image restoration. In The Twelfth International Conference on Learning Representations, 2024a.
Ziwei Luo, Fredrik K Gustafsson, Zheng Zhao, Jens Sjölund, and Thomas B Schön. Photo-realistic image restoration in the wild with controlled vision-language models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6641-6651 2024b.
David Martin, Charless Fowlkes, Doron Tal, and Jitendra Malik. A database of human segmented natural images and its application to evaluating segmentation algorithms and measuring ecological statistics. In Proceedings of the 18th IEEE International Conference on Computer Vision (ICCV), volume 2. pp. 416 416-423. 423. IEEE, 2001.
Seungjun Nah, Tae Hyun Kim, and Kyoung Mu Lee. Deep multi-scale convolutional neural network for dynamic scene deblurring. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 3883-3891, 2017.
Rui Qian, Robby T Tan, Wenhan Yang, Jiajun Su, and Jiaying Liu. Attentive generative adversarial network for raindrop removal from single image. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 2482-2491, 2018
Xu Qin, Zhilin Wang, Yuanchao Bai, Xiaodong Xie, and Huizhu Jia. FEA-Net: Feature fusion attention network for single image dehazing. In Proceedings of the AAAI Conference Artificial Intelligence. pp. 11955-11915, 2020a.
Xu Qin, Zhilin Wang, Yuanchao Bai, Xiaodong Xie, and Huizhu Jia, Ffa-net: Feature fusion at- tention network for single image dehazing. In Proceedings of the AAAI conference artificial intelligence, volume 34, pp. 11908-11915, 2020b.
Liangqiong Qu, Jiandong Tian, Shengfeng He, Yandong Tang, and Rynson WH Lau. Deshadownet: A multi-context embedding deep network for shadow removal. In Proceedings of the IEEE Con- ference on Computer Vision and Pattern Recognition, pp. 4067-4075, 2017.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman. Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation. In Pro- ceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 22500- 22510, 2023.
13

---

Published as a conference paper at ICLR 2025
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00512, 2022.
Simo Särkkä and Arno Solin. Applied stochastic differential equations, volume 10. Cambridge University Press, 2019.
H Sheikh. Live image quality assessment database release 2. http://live. ece. utexas.
edlureeearcchqualltyy 2005.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020a.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems, 33: 1243--1248,, 2020.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben
Poole. bcore-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020b.
Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in neural information processing systems, 34:1415- 1428, 2021.
Yang Song, Prafulla Dhariwal, Mark Chen, and llya Sutskever. Consistency models. arXiv preprint 2rX.v22303.01469, 2023.
Radu Timofte, Eirikur Agustsson, Luc Van Gool, Ming-Hsuan Yang, and Lei Zhang. NTIRE 2017 challenge on single image super-resolution: methods and results. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition (CVPR) Workshops, pp. 114-125, 2017.
Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero P Simoncelli. Image quality assessment: from error visibility structural similarity. IEEE transactions on image processing, 13(4):600- 612, 2004.
Chen Wei, Wenjing Wang, Wenhan Yang, and Jiaying Liu. Deep retinex decomposition for low-light enhancement. arXiv preprint arXiv: 04560, 2018
Wenhan Yang. Robby T Tan, Jiashi Feng, Jiaying Liu, Zongming Guo, and Shuicheng Yan. Deep joint rain detection and removal from a single image. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1357-1366, 2017.
Hu Ye, Jun Zhang, Sibo Liu, Xiao Han, and Wei Yang. Ip-adapter: Text compatible image prompt adapter for -mage diffusion models. arXiv preprins arXiv:2308.06721, 2023.
Lvmin Zhang, Anyi Rao, and Mancesh Agrawala. Adding conditional control to text-to-image diffuuson models. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 3836-3847, 2023.
Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. arXiv preprint arXiv:2204.13900 2022.
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and puttern recognition, pp. 586-595, 2018
Wenliang Zhao, Lujia Bai, Yongming Rao, Jie Zhou, and Jiwen Lu. Unipe: A unified predictor- corrector framework for fast sampling of diffusion models. Advances Neural Information Processing Systems, 36, 2024.
Zhenyu Zhou, Defang Chen, Can Wang, and Chun Chen. Fast bde-based sampling for diffusion models in around 5 steps. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 7777-7786, 2024.
14

---

Published as a conference paper at ICLR 2025
APPENDIX
We include several appendices with derivations, additional details and results. Appendix provide derivations of propositions in Section and equivalence between posterior sampling and
wc
Maler--aruyama discretization, and velocity prediction, respectively. In Appendix B we compare the notations used in this paper and MRSDE (Luo al., 2023b). Appendix
list detailed
algorithms of MR Sampler with various orders and parameterizations. In Appendix D
we present we provide in-depth
details about datasets, settings and results in experiments. In Appendix E discussion on determining the optimal NFE.
A DERIVATION DETAILS
A.1 PROOFS OF PROPOSITIONS
Proposition 1. Given initial value a, time [0, TJ. the solution æ, time / [0 of Eq. TO
is
at 4 + at g (o(2,,T) dr +
[ ai a
Dt = Eg +
(31)
0x
where ar = e 3/ f(r)dr and N).
Proof. For SDEs in the form of Eq. Itt's formula gives the following conclusion:
dv(( t) 0rb(ar, 1)
dt (x) dt
[f t)dt g(t)dw] + 1b² () (t)
(32)
dx
2 Ox"
where (xx 1) is differentiable function. And we define
(<x () = æeloo f(r)dr
By substituting f(x, t) and g(t) with the corresponding drift and diffusion coefficients in Eq we obtain
10
g0(1) co( t)dt + 9(t)dw
d (f(()) f(r)dr dt + ele f(r)dr
ar
And we integrate both sides of the above equation from to (z, t) (xx (2 s) pa(elt ((r)dr -et (67)dr) )ff dl
(2)) co(x,, dr+ S clo ((S)dE g(r)
or
Note that v is standard Wiener process running backwards in time and we have the quadratic variation (dm)) dr. According to the definition of vb( 1) and at, we have
a4 x.s 0 a
1 ( a -d))))== [[ ((()) co(ar, dr + v
(P(()) a?
drz,
a.)
a,o,
which is equivalent to E1.
Proposition Given an initial value as at time [0,T], the solution æ; at time t [0, s] of
15. is
Iu at S g2(r) 20,0, co( r)d7,
aj v. + (1
2t
(33)
a2
as
where a; : c- f f(r)dr
Proof. For ODEs which have semi-linear structure as follows:
dæ
dt P(t)æ + Q(x, t),
(34)
the method of "variation of constants" gives the following solution:
[ Q(x, r)E -5 ((r)dr +
a=()) 2I ((r)dr
15

---

Published as a conference paper at ICLR 2025
By simultaneously considering the following two equations
xxt) = elo P(r)dr
St Q(x T)e to Ptr) dr C
a=()) elo P(r)dr
So Q(x, r)e to Pr C
and eliminating C, we obtain
<xt) x(s)eli P(r)dr + S( Q(x,))ef PTT)
(35)
Now compare Eq. (5) with Eq. (3) and let
P(t) = f(t)
and Q(x, t) = f(t) g2(t)
co co(
20,
Therefore, we can rewrite (5) as
2t = age E. ((r)d + [ e s1S)a f(T)H+ g2(T)
20, co(o,, dr
= .zC H(ridr + u(1 - e (())dr + [" - -Ssss (2( 20,
co( T)dr,
which is equivalent (3).
Proposition 3. Given an initial value a. at time
[0, the solution æ, at time t [0, of
Eq.22. 22 is
Gr -Ar--ee - Dt
x. + (1- at 2(Ae--A) -a; + are a,e- -(Ac--a))))
(1-
Os
as
(36)
+20, Mr -2(( x0 A)da + +0v1 -2(1, -((1--1))
where z NNOO,)).
Proof. According to Eq.
32 we define
g²(() oî
u(t)
f(t)
and (xx Se u(>)dr
We substitute f(x,t) and g(t) in E3. with the corresponding drift and diffusion coefficients in Eq. 22). and integrate both sides of the equation from s to d:
t I" eli u(E)dE
(2( (1 -a)) dr of
21 = a.e. Ca(rddr
f(r)
(37)
S cl u(E)dE (T( a.xo dr + J dc u(E)dE g())
We can rewrite g(T) as Eq. 12 and obtain
(((()) exp S ( ddr
dr [(T) dr
0: e -(,,-,,)
-e -Dr-oo)
0:
(38)
0,
16

---

Published as a conference paper at ICLR 2025
Next, we consider each term in Eq.(77) by employing Eq. 12 and Eq 38 second term:
Firstly, we simplify the
f(7) (((T)
elt u(s)de
(1 ar) dr
f(T)+ + +2(11 -a,) dht
dr
Of
dr
ext
more de
[f(r]dr + +2(1111--1111 - ar)dd
0T
pore At
ar
[f(r)dT + 2dd, - -2,
o
dda
20r di, -2 2 dd
a²
-aoc
(39)
o?
o
o7
Note that
d = d log
Cx1
da
+
a,da
dot
(40)
1 03
0r - a?)
Substitute Eq. 40) into (3) and we obtain
(T)
"
elt. ((EdEE
f(T)
(1 - ar) dr
o²
day
2da,
2
p0.e de
2 ddr
o7
o2 a?)
07
1 + o
-o0ee de
da, 2c 2c dd
2
22
-ooe
d
ar
21
1 - a²?
Ot 2()x-x))) at + afe 2(Ac--a))))
(41)
Secondly, rewrite the third term in Eq.
37 by employing Eq. 12 and (8).
- [" st u((id
(T(
To a xy (2,, dr ((x--a))
- 2 ddr
a 10, (2,, dr
o²
O+
dr
d 2r (Beys y dx )dd
Ade
2a
c c-2(Ai-N) 20 A)dd.
(42)
Thirdly, consider the fourth term in E7. (note that (dw) = -dr):
S. cSt ule)dE g(r)dl
ESS u(E)de (T
of e 2((x
(-203 ddr
drz
dr
of
2e2 -xilddrz
CxV 1 - - 200
(43)
Lastly, we substitute Eq. (36
38 and Eq. 41 43
43
into Eq.377 and obtain the solution as presented in
17

---

Published as a conference paper at ICLR 2025
Proposition 4. Given an initial value Dg at time [O,T], the solution 2t at time t [0, of Eq. 24
is
dt x, + 1o (1-25 - a,
Ar + oz
exx0(20,, A)dA.
(44)
A
Proof. Note that Eq. 24 shares the same structure as Eq.344 Let
g²(t)
P(t)
f(t),
20
and (x, f(t) (t)
(2)
(1 ar)
11
axso f).
207
207
According to Eq. 12 first consider
g²(T) 20?
ett P(r)dr = exp
-[(T) dr exp exp
-ddr + dllog o+
lo dlog at O+
Tt dlog
exp
exp
(45)
0x
Then, rewrite Eq.
35)
(r(
0z g2(
2t = a: Is + H
at or
f(r)
(1 - a+)
dr
a 20 20 (I,, T)dr.
(46)
20?
o, 20²
Firstly, we consider the second term in (6)
(7(
11
σt T
f(T)
(1 a dr
20?
S(T) (²(
+ (²(T)
μ
20 0t
dr
207
(²(
(2(T) 20
μ
f(T)
dr +
dr
20
d or
a dl,
μ
(refer to Eq.
12 and (4.)
OT
Or
μ
d
d
Gt
Ft
0s
Ot
(47)
Secondly, we rewrite the third term in Eq. 46
(T)
20 ar a a (Br, r)dr = of S" e dx mo (ax,, 2)dd,.
(48)
0r
By substituting Eq. 447 47 and Eq. (88 into Eq. (46), we can obtain the solution shown Eq. (4).
A.2 EQUIVALENCE BETWEEN POSTERIOR SAMPLING AND DISCRETIZATION
The posterior sampling (Luo et al. 2024b) algorithm utilizes the reparameterization of Gaussian distribution in Eq. (9) and computes æ;- from æ; iteratively as follows:
F;-1 = pi( xo) + Bizi,
a²
(1 a, 1)
ja,(:;; va)
(x: (1)
04 1200 14) + n.
(49)
(1 a )aj-1
o
az (1
a²
B
a._1
1 1-a a/
18

---

Published as a conference paper at ICLR 2025
where 2; -N( a; = c So f(r)dr and Vo (æ) tuting æp into f4g, we arrange the equation and obtain
μ o;co(x;; (4) /a; μ. By substi-
ai-1
ac
pi( 20)
a
I; + (1
a411
|
a=11 0fCg(
Di u1
ai
1 a?
(50)
P111
a
0:-1 ag
Ij + (1
0/-1
|
a:
Co (Zi 1,t.).
ai
aj
We note that
=1+ [ f(r)dr +o to (f" f(r)dr ~ f(t))ttt
0i-1 / f(1)d7
(51)
02
where the high-order error term is omitted and A4e 4--4-11 By substituring Eq. into Eq. (50) and Eq.(49) we obtain
(,(2,, = f(t)) f(t))Ati 21(1)
E p,,,,,
(52)
oz
(1 - af (1
2f(t.) az-1)
B,
(53)
0
1-
On the other hand, the reverse-time SDE has been presented in Eq. TO). Combining the assumption (t) = 20 in Section 2.2 and the definition of or in Section 3.1 the Euler-Maruyama
descretization this SDE is
2;-1 I2 = - At) (24))
Eo( μ. ti) +
20 201)
1.,0--1 = (1+ ((.)AA)) - 1(4)A.1
En (ii, (1 4)44 + g(4))/(441 Ane
JZ.
ag
2f(1)) A400
(1 + f(t))AA) f(L)AtH
Eo(2;; + 2oo
-1
1- Bizi.
(1; + Jx
(54)
Thus, the posterior sampling algorithm is special Euler-Maruyama descretization of reverse-time SDE with different coefficient of Gaussian noise.
A.3
DERIVATIONS ABOUT VELOCITY PREDICTION
Following Eq. 27
We can define the velocity prediction as (t) = μl sin Φ. zt(t) sin d; + x (os(
(55)
And we have the relationship between (t(() and ca(t) as follows:
2t = (t cos de + ra - cos Φ1) + ax sin (o)))
(56)
In order to get ag from Vp we rewrite Eq. 55 as
=o(() sin? Φi μ sin Φ (t) sin de + co sin de cos d..
(57)
Then replace (t according to Eq. 56
ze(() sin? Φ. μ sin? 2 o, vo(t) sin d, [x, - zt( cos dr 1(1 cos di cos de
(1 cos oe μ (t) sin Φi + 2, cos di =t cos² di.
(58)
Arranging the above equation, Similarly, we can also rewrite (55) and replace (t) as follows:
obtain the transformation from vy to æo, shown in Eq.
28
on cos 2 (et)) = ve(t) cos di fa sin Φt cos Φt + xe () sin di cos oz
(t( cos di pu sin de cos o; sin [x p(1 - cos cos - do sin (dr)
(t) cos di - pa sin di + mg (t) sin di aue sin? ()) Thus obtain the transformation from Vy to eg, as presented in Eq. 29
(59)
19

---

Published as a conference paper at ICLR 2025
B NOTATION COMPARISON TABLE
This paper
f(t)
g(t)
MRSDE (Luo et al. 2023b)
02
e So 0,dr
ox. V - 210 2 0.
Table 4: The correspondence between the notations used in this paper (left column) and notations used by MRSDE (right column).
C DETAILED SAMPLING ALGORITHM OF MR SAMPLER
We list the detailed MR Sampler algorithm with different solvers, parameterizations and orders as follows.
Algorithm MR Sampler-SDE-n-
Require: initial value ØT a 6 Gaussian noise sequence {zi|z, N/O, I) time steps [t data prediction model æg. Denote h; := X dt_,, for i = I.....M, M.
1: Pto <- 2T. Initialize an empty buffer Q.
2: Q 2r (Dtto ln) to M do
buller
3: for - 1
mie 20 (o)olelt-))e(((.L_utt + 02, ve 1z; a1,-1 5: If < M, then Q (oo(at,, li)
ote It... + (1
4: Lt. o
6: end for
7: return Zts
Algorithm MR Sampler-SDEEnn2
Require: initial value Fy d dwc, Gaussian noise sequence {2,/2, N(0, I)M time steps [4:),1 data prediction model Wg. Denote h; := do. dr, for i 1, 1,..., M.
1: Eto - IT. Initialize an empty buffer Q.
2: Q baffer xz xe( lu)
σe Dto + (1
alg ) μ 20,, (e 1)co(2 to) + oi, Vo2n1 121
3: Ft, oio
4: Q buffer 20 #0(Dt, (Et, ,t2)
5: for - 2 M do
D, 60(201 1 1!t1-1) 1 211-2)
6:
h1-1
7: Et. acr.
at; Itin + (1
cr;
20,, [(eht - 1)co( ti-1) + (c - h:)
Tt. 2e 12;
8:
if < M, then Q buller
Rel tx)
9: end for
10:
return ZtM
20

---

Published as a conference paper at ICLR 2025
Algorithm 3 MR Sampler-ODE-n-l.
Require: initial value ar lr + 6xe, time steps (4)) data prediction model ag. Denote h: dt; At=- for ...... M.
1: Eto - ΦT. Initialize an empty buffer Q.
2: Q huffer
3: for - to M do
4: Et. 04-1 It.-1 + (1- - on,
at;
and 1 on(c"t - :ii
5: if M, then Q butter 0 ti)
6: end for
7:
return Dta
Algorithm 4 MR Sampler-ODE--n-2.
Require: initial value ør = l + axe, time steps (4)) data prediction model æg. Denote h; λt, - d for i I,..., M.
1: Eto - ZT. Initialize an empty buffer Q.
2: Q batter xg( xp( to)
3: Pt are0
de Fito + (1-
as
1 - o4, (e 1)e0 (Fto to)
0sg
4: Q buller
Izl (Iz,, ,t1)
5: for - 2 to M do
6: D; en (0/0) -00l))
bi
7: Kt. 0:-11
at Iti- + (1
an,
H on, [[e 1)eo(= + (e - 1 - h.)
8: Ifi < M, then Q batfer
x0(I1;,
9: end for
10: return ZrM
Algorithm MR Sampler-SDE-d-1
Require: initial value æ + 6, Gaussian noise sequence {zi|zz ~ NoO,, I) time steps (t data prediction model æz. Denote h; == de, di_,, for i 1...,,M M.
1: Eto ZT. Initialize an empty buffer Q.
2: Q Σg( (te( tu) buller
3: for 1 M
4: It, = o1;1
ou his Dt..1 + l
01 C 2h: a1, + a2 22) + or, VI 2h + crr
Qt, (1 (1-e) -2hu x 10 14-1)
5: If < M, then Q Bo( tx)
buller
6: end for
7: return ØrM
21

---

Published as a conference paper at ICLR 2025
Algorithm 6 MR Sampler-SDE-d-2.
Require: initial value æT la + oxe, Gaussian noise sequence (2i|2, ~ NoO, D) time steps {ti data prediction model æg. Denote h4 == dr. Ar,- for i ,,..., M.
1: Eto - 2r. Initialize an empty buffer Q.
2: Q Hzl (Fto to)
batter
del e hi Tto
Tto + pa (1 aey e 221 ot, + Cc2 e 2221 + ar, (1 e no (it 1o) +
3: St
ot vI e- 2h11 batter
e-2h1
4: Q
Igl (It, ti)
5: for 2 to M do
6: Di ( sti-1
1) (
14-11
Tu De.- p (1 an e-2h ar, + aq a 2h) oz, v1 -e 2-2hi zi +
7: Et.
0::-1
anp-
01, e1-e-2h)) anlaa 194111) On (h-- م 24 D.
8: if M, then Q buffer xe x0( bs)
9: end for
return Ft
Algorithm 7 MR SDE-ler-ODE-dII
Require: initial value Fr = μ + oze, time steps (t data prediction model zz. Denote h4 := dt, de, for i =,. 1,..., M.
1: Eto < tT. Initialize an empty buffer Q.
2: Q buller Ztpl (Fpo la)
3: for i - to M do
4: Ft. 0
mi -1.. + 4 (1- ou + Bie at as.) + as, (1- e hhv xs (E vo( titi) a=;-1
ax
5: If M, then Q ar (00 li)
buffer
6: end for
7: return Frar
Algorithm 8 MR Sanpler-ODE-d-2. ODE-d-2.
Require: initial value xr = 4 + oxe, time steps (t data prediction model æg. Denote h; dt, - >t.,, for 1,..., M.
1: Eto - ØT. Initialize an empty buffer Q.
2: Q bufter
2o(F tq)
3: St, deg
de Tooo 11
Otoo
0t (1 e-hs) xr (0o( to) orh)
aeg
0rg
4: Q buffer
20 (Itr, (1)
5: for - 2 to M do
6: Et. 0:.
01 It... + 1
oy + Ti ath, - av.) + ar, (1 (1-E
e hi)
+
hh) wo( dn-1 1)- (000)))))))-)
O11 (h:
him
7: if < M, then Q buller co( tx)
8: end for
9: return ØtM
22

---

Published as a conference paper at ICLR 2025
D DETAILS ABOUT EXPERIMENTS
D.1 DETAILS ABOUT DATASETS
We list details about the used datesets in 10 image restoration tasks in Table
Task name
Dotiset naie
Reference
Number of testing images
Blury
Goltu
Nuh al.
2017
1111
Hazy
RESIDE-6K
Qia
20200u
1000
DIVCK, FlickrāK LIVEI
Agustsson Timatte
Timooec
Sholkh 2005
Low-light
2015
Nossy
DIV2K, Plikk2K and CUSD68
Agustsson Timodde
2017
Timofie et
2017
Martin al.
2001
68
Raindrup
RainDrop
Quan et al.
2018
58
Rainy
Raint0001
Yang al,
2017
100
Shadlwwd
Qu
Snowy
1D0KL
Liu 2018
Inpainting
CelebaHQ)
Lugmayy et al.
2022
100
Table Details about the used datasets in 10 image restoration tasks
D.2 DETAILS ON THE NEURAL NETWORK ARCHITECTURE
In this section, we describe the neural network architecture used in experiments. We follow the framework of
Luo et al. (2024a), an image restoration model designed to address multiple degrada- tion problems simul Itnnousyy without requiring prior knowledge of the degradation. The diffusion model in Luo et al. (2024a) is derived from Luo al. (2023c), and its neural network architecture based on NAFNet. NAFNet builds upon the U-Net architecture by replacing traditional activa- tion functions with SimpleGate and incorporating an additional multi-layer perceptron to manage channel scaling and offset parameters for embedding temporal information into the attention and feedforward layers. For further details, please refer to Section 4.2 in Luo et al. (202cc).
D.3 DETAILED METRICS ON ALL TASKS
We list results on four metrics for ten image restoration tasks in Table 15
NFE
Method
LPIPS
FIDJ
PSNRT
SSIM
NFF
Method
LPPPS
FIDI
PSNR1
SSIM
Posterior
0.0385
18.35
29.49
0,9102
Posterior
0.0614
21.42
27.43
0.8763
100
Sampler-1
Euler
0,0426
20.71
18.16
29.34
29.45
0.9086
0.8981
100
MR Sampler. I
Fuler
0.0608
0.0683
21.30
23.27
27.45
27.09
0.8577
0,8754
MR Sampler-2
0.0397
18.81
29.30
0.9055
MR Sampler-2
0.0626
21.47
28.18
0.8691
Posterior
0.6048
Posterior
0.2374
21.35
0.7037
Euler
0.4449
249.5
12.68
0.5800
Euler
0,2730
76.02
21.02
0.6676
50
Sampler-1
50
MR Sampler-1
0.0602
20.91
27.63
MR Sampler-2
0.0402
19.09
29.15
0.9046
MR Sumpler 2
0.0628
21.85
27.08
0,8685
Posterior
0.7130
347.4
10.19
0.2171
Posterior
U.6622
123.8
16.42
0.3546
Euler
0.7257
344.3
10.16
0.2073
Euler
0.6861
126.0
16.23
0.3431
Sampler-1
18.29
30.05
0.9172
20
MR Sampler-1
0.0601
21.32
28.07
MR Samples-2
0.0408
19.13
28.98
0,9032
MR Sampler-2
0.0650
22.34
26.89
0.8645
Posterior
0.8097
374.1
9.802
0.1333
Posterior 0.8013
138.5
14.76
0,2694
Fuler
0.8154
381.1
9.786
0.1305
Fuler
0.8164
140.4
14.64
0,2640
MR Sampler-
0.0401
18.46
30.61
0.9229
MR Samplor-1
0.0608
22.26
28.50
0.8992
MR Sampler-2
0.0437
19.29
28.48
0.8996
MR Sampler-2
0.0698
23.92
26.49
0,8573
Posterior
0,8489
385.6
9.599
0.1057
Posterior
0,8590
145.8
13.992
0,2318
Euler
0.8525
384.7
9.587
0.1042
Euler
13.85
MR Samples-1
0.0428
20.00
31.03
0.9262
5
MR Sampler- I
0.0611
23.29
28.89
0.9065
Sampler-2
0.0529
28.35
0.8930
MR Sampler-2
0.0628
21.95
27.06
Table 6: Image inpainting.
Table Snowy image restoration.
23

---

Published as a conference paper at ICLR 2025
NFE
Method
LPIPS)
FIDI
PSNRT
SSIMT
NFE
Method
LPIPSL
FIDI
PSNRT
SSIMT
Posterior
0,0970
20.14
27.84
0.8391
Posterior
0.0594
25.58
29.14
0.8704
Euler
0.1129
20.30
27.59
0.8112
Euler
0,0725
28.80
28.77
0.8473
100
MR Sampler-
0,0984
20.03
27.73
0.8370
100
MR Sumpler-1 I
0.0594
30.53
29.15
0,8679
MR Sampler-2
0,0989
20.69
27.44
0,8329
MR Sampler-2
0.0616
27.35
28.92
0.8614
Posterior
0.6119
101.8
18:29
0.4720
Posterior
0.4418
183.1 1
16.41
0,4903
50
MR Sampler-1
Euler
0.6985
0.0978
20.90
114.2
27.92
17.80
0.4123
0.8409
50
MR Sampler-1
Euler
0.0586
0,4560
185.2
30.73
16.24
29.34
0.8730
Sampler-2
0.1006
20.74
27.20
0.8310
MR Sampler-2
0.0620
27.65
28.85
Posterior
1,043
187.2
14.73
0,2049
Posterior
0.6865
293.6
12.54
0,2464
20
MR Sampler.1
Baler
0.0954
1,065
19.79
192.9
28.31
14.56
0.8505
0.1955
20
MR Sampler- I
Euler
0.0604
0.6945
299.0
31.19
29.81
12.49
0,2402
0,8845
MR Sampler-2
0.1014
20.93
27.17
0.8299
MR Sampler-2
0.0635
27.79
28.60
0.8559
Posterior
1,122
208.4
13:67
0.1525
Posterior
0,7972
323.0
11.50
0.17
10
MR Sampler-1
Fuler
0.0956
1,133
209.7
19.87
28.67
13.57
0.1484
0.8554
10
MR Sampler-1
Fuler
0.0699
0.8043
31
330.8
30.28
11.46
0.8943
0.1724
MR Sampler-
0.1044
21.85
26.99
0.8276
MR Sampler-2
0.0678
29.54
28.09
0.8483
Posterior
1.155
13.12
0.1298
Posterior
0.8663
332.4
10.96
0.1450
5
Buler
1,161
221.4
13.06
0.1278
5
Euler
0.8714
332.5
10.94
0,1435
Sampler-1
20.16
0.8601
MR Sampler-1 I
0.0729
32.06
30.68
MR Sampler-2
0.2203
36.69
23.73
0.6690
MR Sampler-2
0.0637
26.92
28.82
0.8685
Table 8: Shadowed image restoration.
Table 9: Rainy image restoration.
NFE
Method
LPIPS]
FIDI
PSNRT
SSIMT
NPE
Method
LPIPS,
FIDI
PSNR
SSIM
Posterior
0.0443
19.70
30.05
0.8910
Posterior
0.1694
65,79
25.97
0.7267
Baler
0.0634
24.19
29.31
0.8438
Euler
0,2719
68. 69
24.28
0.5686
100
MR Sampler-1
0.0437
19.94
29.91
0.8852
100
MR Sampler-1
0.1629
59.12
25.97
0.7244
MR Sampler-2
0.0454
21.35
29.55
0.8768
MR Sampler-2
0,1586
64.02
25.67
0,7126
Posterior
0.4289
100.1
23. 19
0.4663
Posterior
0.7713
135.7
18.19
0,2763
Euler
0.5011
111.6
22.35
0.4106
Euler
0.8060
143.5
17.74
0.2615
MR Sampler-
0.0428
19.14
30.07
0.8914
50
MR Sampler-1
0.1680
65.14
26.20
0.7330
MR Sampler-
0,0459
20.50
29.40
0.8764
MR Sampler 2
0,1615
64,24
25.61
0.7127
Posterior
0.8873
190,8
17.37
0.1925
Posterior
0,9941
181.6
14.76
0,1821
Baler
0.9082
194.1
0.1839
Euler
1.006
188.3
14.61
0,1781
20
MR Sampler-1
0,0439
19.31
30.44
0.9025
20
MR Sumplee-1 1
0,1872
71.31
26.68
0.7494
Sampler-2
0.0470
0.8745
MR Sampler-2
0.1695
64,90
25.46
Postetior
0.9884
215.5
15.67
0,1419
Posterior
1.057
202.6
13.59
0.1550
10
Huler
0,9993
213.1
15.52
0.1381
10
Euler
1.063
207.0
13.51
0,1529
MR Sampler-1
0.0466
20.60
30.77
0,9114
MR Samplee-1
0,2043
79.28
27.13
0.7628
MR Sampler-2
0.0485
22.17
29.37
0.8779
MR Sampler-2
0.1853
70.19
25.09
0.6984
Posterior
1,300
226.3
14.82
0.1
Posterior
1.087
213.5
13.00
0.1419
5
MR Sampler-1
Baler
0.0497
1.037
226.4
21.18
31,04
14.74
0.1
0.9175
5
MR Sampler-1 I
Euler
0,2046
1.091
218.6
80.45
27.51
12.96
0.7743
0. 1404
MR Sampler-2
0.0733
28.26
28.03
0,8369
MR Sampler-2
0,3178
73.93
24.32
0,5485
Table 10: Raindrop image restoration.
Table 11: Noisy image restoration.
24

---

We choose one Published as a conference paper at ICLR 2025
NFE
Method
LPIPSI
FIDI
PSNRT
SSIMT
NFE
Method
LPIPS
FIDI
PSNRT
SSIMT
Posterior
0.0796
34.70
23.84
0.8496
Posterior
0.1702
45.77
25.78
0.7380
Euler
0.1014
37,46
23.27
0,8027
Euler
0,2949
56.99
23.84
0.5398
100
MR Sampler-
0.0774
32.91
23,80
0.8451
100
MR Sumpler-1 I
0.1636
43.67
25.81
0,7362
MR Sampler-2
0,0789
32.91
23.59
0.8394
MR Sampler-2
0.1585
44.58
25.46
0.7220
Posterior
0.6572
151.3
9,490
0.2746
Posterior
0.6494
103.7
19.34
0,2908
50
MR Sampler-1
Euler
0.7517
0.0784
34.45
176.7
23.51
0.8476
0.2340
50
MR Sampler-1
Euler
0.7035
0,1734
47.09
113.7
26.06
18.62
0.7470
Sampler-2
0,0786
32.85
23.52
0.8382
MR Sampler-2
0.1567
45.56
25.41
Posterior
1,288
390.1
8.211
0,0648
Posterior
0.8252
140.6
16.44
0,1988
20
MR Sampler.1
Baler
0.0791
1.297
35.28
396.8
24.22
8.212
0.8586
0.0625
20
MR Sampler- I
Euler
0.1993
0.8477
51.43
144.5
26.58
16.19
0.7649
0.1921
MR Sampler-2
0.0792
32.80
23.63
0.8399
MR Sampler-2
0.1675
47.36
25.33
0.7235
Posterior
1.351
432.2
8.130
0.0476
Posterior
0.8723
158.9
15.49
0.1738
10
MR Sampler-1
Fuler
0.0831
1.354
41.10
424.2
8.136
24.04
0.8619
0.0467
10
MR Sampler-1
Fuler
0,2183
0.8859
57.83
159.9
26.98
15.34
0.7771
0.1701
MR Sampler-
0.0841
36.53
23.22
0.8398
MR Sampler-2
0.1871
51.25
25.08
0.7197
Posterior
1.371
453.0
0.0408
Posterior
0.8941
163.6
14.99
0.1615
5
Sampler-
Buler
0.0860
1.372
447.2
24.02
8,118
0,0405
0.8676
5
MR Sampler. I
Euler
0,9013
0.2281
59.17
162.0
27.28
14.91
0,1596
MR Sampler-2
0.0782
33.98
24.13
0.8507
MR Sumpler-2
0,3751
62.60
23.15
0.4797
Table 12: Low-light image restoration.
Table 13: JPEG image restoration.
NFF
Method
LPIPSS
FIDJ
PSNR1
SSIM
NFF
Method
LPIPS.
FIDJ
PSNR1
SSIM
Posterior
0.0211
4.755
30.37
0.9485
Posterior
0.1249
14.48
27.48
0.8442
100
MR Sampler-
Baler
0.0358
0,0219
6.182
4,826
29.95
30.42
0.9319
0.9462
100
MR Sampler- I
Euler
0,1404
0.1239
14.61
16.06
27.16
27.46
0.8179
0.8419
MR Sampler-2
0,0230
4,978
30.26
0.9431
MR Sampler-2
0.1248
14.61
27.30
0,8365
Postetior
0.3994
35.47
15.34
0.5808
Posterior
0,5112
45.13
24.41
0,5571
50
MR Sampler-1
Fuler
0.0211
4.737
30.44
0.9484
0.5160
50
MR Samplee-1 1
Euler
0.5739
0.244
61.21
14.47
27.59
23.79
0.8461
Sampler
0.0233
4.993
0.9427
MR Sampler-2 2
0.1251
14.62
27.24
Posterior
0.9911
114.0
12.81
0.1832
Posterior
1.069
117.8
18.20
0.1717
20
MR Sampler-1
Buler
0.0200
1.012
4,682
118.5
30.63
12.71
0,9534
0.1752
20
MR Sampler-1 I
Euler
0,1287
1.089
120.8
14.92
27.85
17.94
0.1637
0.8544
MR Sampler-2
0.0240
5.077
30.06
0.9409
MR Sampler-2
0.1266
14.82
27.13
0,8337
Posterior
1,116
144.0
11.94
0.1261
Posterior
1.187
141.9
16.11
0.1136
Baler
1.128
147.3
11.87
0.1228
Euler
1.197
143.7
15.96
1.4
10
MR Sampler-1
0.0197
4.785
30.80
0,9579
10
MR Sampler-1
0.1356
15.65
28.10
0.8613
MR Sampler-2
0.0246
5.228
29.65
0.9372
MR Sampler-2
0.1300
15,51
26.88
0.8295
Posterior
1.162
159.1
11.47
0.1042
Posterior
1,228
155.3
0.0922
5
Sampler-1
Buler
0.0205
1,168
4.926
161.2
30.56
11.43
0,1026
0.9604
5
MR Sampler- I
Euler
0.1422
1,234
157.3
16.32
28.31
15.00
0.0907
0.8668
MR Sampler-2
0.0228
5.174
29.65
0,9416
MR Sampler-2
0.1248
14,20
26.92
0.8354
Table 14: Hazy image restoration.
Table 15: Motion-blurry image restoration.
D.4 DETAILS ON NUMERICAL STABILITY AT LOW NFES
We have included further details regarding numerical stability at 5 NFEs to complement the ex- periments presented in Section 5.3 As illustrated in Figure 5 when the NFE is relatively low, the convergence rate of noise prediction at each step does not exceed 40%, which results in sampling collapse. In contrast, during the final 1 or 2 steps, the convergence rate of data prediction approaches nearly 100%.
D.5 WALL CLOCK TIME
We measured the wall clock time in our experiments on a single NVIDIA A800 GPU. The average wall clock time per image of two representative tasks (low-light and motion-blurry image restora- tion) are reported in Table T6 and 17
D.6 PRESENTATION OF SAMPLING RESULTS
We present the sampling results for all image restoration tasks in Figure 6 and image for each task.
7
25

---

Published as a conference paper at ICLR 2025
Staned
Noise predictoon (5)
1.0
Noise prediction
0.8
Data prediction
0.6
Ground Truth
Data predctioonn (5)
0.4
0.2
0.0
Index of timestep
(a) Sampling results.
(b) Ratio of convergence.
Figure 5: Convergence of noise prediction and data prediction at 5 NFEs. In (a), choose a stained image for example. The numbers in parentheses indicate the NFE. In (b), we illustrate the ratio of components of neural network output that satisfy the Taylor expansion convergence requirement.
Table 16: Wall clock time the Low-light dataset.
Table 17: Wall clock time the Motion- blurry dataset.
NFE
Method
T
NFE
Method
Time(s)
Posterior Sampling
17.19
Posterior Sampling
82.04
100
100
MR Sampler-2
17.83
MR Sampler-2
81.16
Posterior Sampling
8.605
Posterior Sampling
41.23
50
MR Sampler-2
8.439
MR Sampler-2
40.18
Posterior Sampling
3,445
Posterior Sampling
16.44
20
MR Sampler-2
3.285
MR Sampler-2
15.59
Posterior Sampling
1.727
Posterior Sampling
8.212
10
10
MR Sampler-2
1.569
MR Sampler-2
7.413
Posterior Sampling
0.8696
Posterior Sampling
4.133
5
MR Sampler-2
0.7112
MR Sampler-2
3.294
SELECTION OF THE OPTIMAL NFE
In practice, sampling quality expected to degrade as the NFE decreases, which requires us make trade-off between sampling quality and efficiency. From the perspective of reverse-time SDE and PF-ODE, the estimation error of the solution plays a critical role determining the sampling quality. This estimation error is primarily influenced by hmaz maxi gig (A; > O where M represents the NFE. During sampling, the noise schedule must designed in advance, which also determines the X schedule. Since X; is monotonically increasing, a larger NFE results in smaller kmax- thereby reducing the estimation error and improving sampling quality. However, different sampling algorithms exhibit varying convergence rates. Experimental results show that the MR Sampler (our method) can achieve good score with as few as NFEs, whereas posterior sampling and Euler discretization fail to converge even with 50 NFEs.
Specifically, we conducted experiments on the snowy dataset with low NFE settings, and the results presented in Table [8] The sampling quality remains largely stable for NFE values larger than 10, gradually deteriorates when the NFE is below 10, and collapses entirely when NFE is reduced 2. Based on our experience, we recommend using 10-20 NFEs, which provides a reasonable trade-off between efficiency and performance.
26

---

Published as a conference paper at ICLR 2025
Snowy
NFC-5
NFB-20
NFE-50
NFE-60
NFE-100
Cround Truth
Postarior Sampling
MR Sampler- SDE-d-2
Stadowed
NFC-5
NFE-20
NFE-50
NFE-60
NFE-100
Cround Truch
Posterior
Sampling
MR Sampler- SDE-d-2
Rainy
NFE=20
Ground
Sampling
MR Sampler- SDE-d-2
Raindrop
NFE-6
NFE-20
NFE-50
NFE=60
NEE-100
Ground Truth
Posterior
Sampling
MR Sampler- SDE-d-2
Noisy
NFC-5
NFE-20
NFE-50
NFE-80
NFE=100
Ground Truth
Posterior
Sempling
MR Sampler- SDE-d-2
Low-ligha
NFE-5
NFE-20
NFE-50
NFE-80
NFE-100
Ground Truth
Posturior Sampling
MR Sampler- SDE-d-2
Figure 6: Comparisons between posterior sampling and MR Sampler-SDE-4-2 on snowy, shad- owed, rainy, raindrop, noisy and low-light datasets.
Table 18: Results of MR Sampler-SDE- with data prediction and uniform on the snowy dataset.
NFE
100
50
20
10
8
6
5
2
LPIPS
0.0626
0.0628
0.0650
0.0698
0.0725
0.0744
0.0628
0.1063
1.422
FIDL
21.47
21.85
22.34
23.92
24.81
25.60
21.95
30.18
421.1
PSNR1
27.18
27.08
26.89
26.49
26.61
26.40
27.06
25.38
6.753
SSIM1
0.8691
0.8685
0.8645
0.8573
0.8462
0.8407
0.8718
0.7640
0.0311
27

---

Published as a conference paper at ICLR 2025
NFE-6
NFB-20
NEE-50
NFE-60
NEE-100
Cround Truth
Posterior
Sampling
MR Sampler-
SDE-d-2
Huzy
NFE-5
NFE-20
NFE-50
NFE=60
NFE-100
Cround Truth
Posterior
Sampling
MR Sampler-
SDE-d-2
Stained
NFE-6
NFR=20
NFR=100
Ground Truth
Posterior
Sampling
MR Sampler- SDF-d-2
Motion blury
NFF=5
NFR=20
NFF=50
NFB=80
NFR=100
Ground Truth
Posterior Sampling
MR Sampler SDE-d-2
Figure 7: Comparisons between posterior sampling and textitMR Sampler-SDE-d-2 on JPEG- compressed, hazy, inpainting and motion-blurry datasets.
28


Reviews:
[{"review_id": "TO9ZGf2HGZ", "reviewer": "Reviewer_q9n1", "summary": "Standard MR Diffusion models typically require a high number of function evaluations to obtain high-quality samples. To address this limitation, MR Sampler introduces a novel approach that combines an analytical function with an integral parameterized by a neural network. The proposed method demonstrates a 10 to 20-fold speedup across various image restoration tasks, obtaining efficient sampling without compromising quality.", "questions": "Please see the weaknesses section.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1). The paper is very well written, and easy to understand especially very coherent even with the math part.\n\n2). The core idea of the paper is very interesting. Especially how the authors have formulated  the core idea.\n\n3). The shown results are very impressive and competitive.\n\n4). Even though the concept behind the \"sampling trajectory\" part is simple, the visualizations obtained through it really helps to understand the core concept and contrast with existing methods. \n\n5). The MRS is plug and play.", "weaknesses": "I believe this is a great paper with solid reasoning and a well-thought-out approach. However, I have some questions regarding the numerical implementation and the comparison methods used.\n\n1). Can the authors elaborate on why only the backward difference method is used? Are there specific benefits to this approach over others? \n\n2). The proposed MR Sampler is only compared with posterior sampling, and Euler-Maruyama discretization. What are the current SOTA methods? How does the proposed approach compare with these, beyond posterior sampling and Euler-Maruyama?\n\n3). I would like to understand why the authors say \"frequently fall outside\" in this part (line304)\n \"Although the standard deviation of this Gaussian noise is set to 1, the values of samples can frequently fall outside the range of [-1,1]\"\n\n4). Could the authors specify the neural network architecture they used? \n\nPlease address these, and I'm leaning towards acceptance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3XONVwB7x4", "reviewer": "Reviewer_hx57", "summary": "This paper proposes a fast sampling algorithm for mean-reverting diffusion. This addresses a gap for mean-reverting diffusion SDE solvers, as current fast samplers for SDEs do not readily apply to mean-reverting SDEs. Two flavors of solvers are proposed, one based on noise prediction, and the other on data prediction. Results show both perform similarly for larger NFEs, but the latter outperforms the former for fewer NFEs.", "questions": "- Why wasn't an ablation study done on n, k?\n- Why wasn't wall clock time reported? Please provide.\n- Why weren't there comparison to standard (non-MR) SDE fast samplers? This comparison should include both reconstruction quality for inverse problems (since this is given as the main motivation for MR diffusion) and computational time.\n- How was guidance incorporated for image restoration? Does the particular method for incorporating the degraded observation make a difference for the performance of MR sampler?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Addresses a gap in fast sampling for mean-reverting diffusion\n- Provides two alternatives focusing on noise/data prediction\n- Relevant ablation studies for some parameter choices\n- Evaluates performance in image restoration tasks", "weaknesses": "- No ablation study on n, k\n- Wall clock time not reported, only NFE improvement is discussed\n- No comparison to standard SDE fast samplers \n- The whole description is for unconditional sampling, but all results are for image restoration. How is the guidance incorporated?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "62uMmVEukE", "reviewer": "Reviewer_wsmd", "summary": "The authors propose a novel algorithm, named MRS, designed as a sampler to reduce the number of sampling steps required for Mean Reverting Diffusion. To achieve this, they solve the reverse-time stochastic differential equation (SDE) alongside the probability flow ordinary differential equation, resulting in a reduction of required steps to between 5 and 15 while maintaining high-quality outcomes. The method is compatible with mainstream parameterizations, making it broadly applicable. The authors conduct an extensive evaluation across 16 different image perturbations, demonstrating that their proposed sampler generally outperforms others in both efficiency and speed.", "questions": "- How did you determine the optimal number of sampling steps? The choice of steps varies across different experiments, but the selection criteria are not immediately clear. Could you provide more explanation on how to select the number of steps to balance efficiency and performance?\n- What was the rationale behind choosing a low-light and motion-blurry dataset for the visualizations in Figure 2? \n- In reference to the related paper [1], it is mentioned that solving the SDE typically requires only 22 steps, although they still use 100 steps. Since you opted to use 100 steps aswell, I am wondering why the performance with 20 steps is still so substantially different.\n\n[1] Luo, Ziwei, et al. \"Image restoration with mean-reverting stochastic differential equations.\" arXiv preprint arXiv:2301.11699 (2023).", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The authors' proposed method demonstrates robust performance across a variety of experiments, consistently showing strong results in both quality and efficiency. Notably, it outperforms posterior-based and Euler-based samplers when using a high number of sampling steps. The performance advantage becomes even more pronounced as the number of steps is reduced, underscoring the method's effectiveness with low-step counts.\n\n- It is particularly encouraging to see that the proposed method performs best across almost all of the 16 tested image perturbations. This breadth of performance suggests that the approach is not only effective but also adaptable to different types of image perturbations.\n\n- The paper is generally well-organized and written. Key concepts and technical choices are clearly explained.\n\n- Overall, this is a well-rounded paper that offers convincing results and provides a method that should be straightforward to integrate with existing frameworks. I also appreciate the authors’ decision to release the code for easy reproducibility.", "weaknesses": "- The choice of the number of sampling steps is not entirely clear. While it seems that setting this value to 20 yields acceptable results, this may not fully leverage the acceleration benefits of the proposed method. Further guidance on selecting an optimal number of steps or a discussion on its trade-offs would make this aspect more transparent.\n- There is some ambiguity regarding the number of function evaluations (NFEs) required for MR Diffusion. The authors mention that MR Diffusion typically requires hundreds of NFEs, but in the current paper, this value is set to 100, consistent with the original paper. This raises questions about whether the maximum performance of the Posterior and Euler-based sampling methods could be higher if a larger number of steps were used. Given the performance gains depicted in Figure 2, it would be helpful to clarify if the efficiency improvements observed are partly due to these optimized sampling steps.\n- Table 15 in the appendix contains incorrect highlight", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zl1TupfGj3", "reviewer": "Reviewer_pnRJ", "summary": "The paper proposes a new perspective on improving the efficiency of sampling in Mean Reverting  Diffusion models. This approach introduces MR Sampler, which leverages both ordinary differential equations (ODE) and stochastic differential equations (SDE) to enhance the speed of the sampling process. While MR Diffusion provides a fresh perspective by modifying the structure of the SDE to make controllable image generation simpler and more natural, the main challenge it addresses is the inefficiency of sampling, which typically requires hundreds of function evaluations (NFEs) to generate high-quality outputs.\n\nThe authors build on prior work that primarily focused on denoising tasks using ODE solvers, and they extend this framework to a broader set of tasks without requiring additional training. The key technical contribution is the semi-analytical solution they derive, which reduces the number of NFEs significantly while maintaining competitive sampling quality. In fact, MR Sampler shows a speedup of 10 to 20 times across ten different image restoration tasks.", "questions": "1. How does reducing NFEs affect sampling quality? Could you compare this with speed-quality trade-off methods like \"Come-Closer-Diffuse-Faster\"?\n\n2. Do different image restoration tasks benefit differently from the MR Sampler in terms of speed or quality?\n\n3. Why is Mean-Reverting SDE better for posterior sampling compared to other methods like Optimal Transport or Schrödinger Bridge?\n\n4. Can you provide more details on the stability of data prediction models at low NFEs, and when would noise prediction be preferable?\n\n5. Can MR Sampler be extended to larger or multi-modal tasks, such as video generation or text-to-image models?\n\n6. Why are there no comparisons with SOTA models in terms of both speed and quality? Would this strengthen the paper?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper introduces a novel approach to sampling in Mean Reverting  Diffusion models, which itself is a relatively new paradigm compared to more conventional diffusion processes like Variance Preserving and Variance Exploding SDEs. The key originality lies in the way the authors combine semi-analytical solutions derived from ODE and SDE solvers with MR Diffusion, enabling faster sampling. This combination is innovative as it offers a fresh perspective on accelerating diffusion models by altering the underlying stochastic differential equation structure rather than focusing solely on the score function as in prior work.\n\nAdditionally, the Mean Reverting SDE framework offers a natural integration of image conditions into the generation process, making it more applicable to tasks that require controllable generation, such as image restoration and inpainting. This is a novel contribution compared to prior work that has typically used Diffusion Schrödinger Bridge or Optimal Transport methods for similar purposes. The MR Diffusion model’s ability to address multiple tasks beyond denoising is a creative and valuable extension of existing frameworks.\n\nThe technical depth of the paper is commendable. The authors offer a rigorous derivation of the semi-analytical solution and provide substantial theoretical grounding for their approach. The use of probability flow ODEs (PF-ODEs) and reverse-time SDEs in the MR Diffusion context is clearly explained and well-justified. The semi-analytical solution, which combines an analytical function and a neural network-based integral, reduces computational complexity without compromising sampling quality.\n\nThe experiments are comprehensive and cover ten different image restoration tasks, providing strong empirical support for the proposed method. The use of multiple performance metrics (e.g., FID, LPIPS, PSNR, SSIM) demonstrates that the authors took a thorough approach to evaluating both the quality and speed of the generated samples. The speedups of 10-20x, without a significant drop in sample quality, highlight the robustness and practical value of the proposed method.\n\nThe paper is well-structured, with clear sections on the methodology, theoretical contributions, and experimental validation. The technical content, while complex, is made accessible through the use of visual aids (e.g., Figure 1 comparing qualitative sampling results, and charts showing performance metrics) and clear explanations of the mathematical formulations. The distinction between noise prediction and data prediction models, and the impact of these choices on sampling quality and stability, is clearly delineated and contributes to a deeper understanding of the technique.\n\nThe authors also provide appendices with detailed proofs and further experimental results, ensuring that the methodology is reproducible and the claims are verifiable. Overall, the clarity in presenting a technically complex subject is a strong point of the paper.\n\nThe proposed MR Sampler addresses a significant challenge in the field of diffusion models: accelerating the sampling process without sacrificing quality. The speedups achieved in this work, which range from 10x to 20x across multiple tasks, are substantial and have clear practical implications, particularly in real-time applications such as image restoration. This makes the method highly relevant for use cases that demand controllable and fast generation, such as medical imaging, video processing, and computational photography.\n\nThe method’s plug-and-play nature is another strength. It is adaptable to a variety of existing diffusion models and does not require retraining, making it easy to integrate into different applications. The broad applicability to various image restoration tasks (e.g., dehazing, inpainting, motion-blur reduction) enhances the significance of the work.", "weaknesses": "A notable limitation is that the paper places more emphasis on the speed of sampling rather than the quality. While the authors claim comparable quality in terms of metrics such as FID and LPIPS, they do not provide an in-depth analysis of the trade-off between sampling speed and output quality. This is a  gap, as accelerating sampling without sacrificing quality is one of the central challenges in diffusion models. The authors could have benefited from comparing their approach with other speed-quality trade-off techniques, such as those mentioned in prior work (e.g., the 'Come-Closer-Diffuse-Faster' approach). Although these methods might be older, they are relevant in establishing a clear benchmark and providing a deeper understanding of the trade-offs at play. (like https://openaccess.thecvf.com/content/CVPR2022/papers/Chung_Come-Closer-Diffuse-Faster_Accelerating_Conditional_Diffusion_Models_for_Inverse_Problems_Through_Stochastic_CVPR_2022_paper.pdf  and https://arxiv.org/abs/2108.01073 )\n\nThe experiments demonstrate a clear focus on speedups, and while they are rigorous and cover a variety of tasks (e.g., image dehazing, inpainting), the lack of comparison with more state-of-the-art methods suggests that the method may not yet be positioned as a new state of the art but rather as a faster alternative with comparable performance under specific conditions.\n\nOverall, the introduction of the Mean Reverting SDE into the posterior sampling stage with ODE solvers is a fresh and promising perspective. However, the paper would benefit from addressing the broader implications of their method, particularly how the mean-reverting approach impacts the balance between speed and quality. Including more comprehensive comparisons with established trade-off techniques would further strengthen the contribution but overall this is a good paper, potentially worth discussing at ICLR.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)