Gotta Go Fast with Score-Based Generative Models
Abstract
Score-based (denoising diffusion) generative models have recently gained a lot of success in generating realistic and diverse data.
These approaches deﬁne a forward diffusion process for transforming data to noise and generate data by reversing it.
Unfortunately, current score-based models generate data very slowly due to the sheer number of score network evaluations required by numerical SDE solvers.
In this work, we aim to accelerate this process by devising a more efﬁcient SDE solver.
Our solver requires only two score function evaluations per step, rarely rejects samples, and leads to high-quality samples.
Our approach generates data 2 to 10 times faster than EM while achieving better or equal sample quality.
For high-resolution images, our method leads to signiﬁcantly higher quality samples than all other methods tested.
Our SDE solver has the beneﬁt of requiring no step size tuning.
Introduction
Score-based generative models [Song and Ermon, 2019, Ho et al., 2020, Jolicoeur-Martineau et al., 2020, Piché-Taillefer, 2021] have been very successful at generating data from various modalities[Song et al., 2020a, Chen et al., 2020, Kong et al., 2020, Niu et al., 2020].
These models generally achieve superior performances in terms of quality and diversity than the historically dominant Generative Adversarial Networks (GANs)
[Goodfellow et al., 2014].
Although very powerful, score-based models generate data through an undesirably long iterative process; meanwhile, other state-of-the-art methods such as GANs generate data from a single forward pass of a neural network.
Increasing the speed of the generative process is thus an active area of research.
Existing methods for acceleration Chen et al.
[2020], San-Roman et al.
[2021], Song et al.
[2020a,b] often require considerable step size/schedule tuning and do not always work for both Variance
Exploding (VE) and Variance Preserving (VP) processes (the two most popular diffusion processes for score-based models).
To improve speed and remove the need for step size/schedule tuning, we propose to solve the reverse diffusion process using SDE solvers with adaptive step sizes.
Maruyama for an equal computational budget.
We measure speed through the Number of score
Function Evaluations (NFE) and the quality of the generated images through the Fréchet Inception Distance (FID; lower is better).
See Table 1-2 for more details.
It turns out that off-the-shelf SDE solvers are ill-suited for generative modeling and exhibit either (1) divergence, (2) slower data generation than the baseline, or (3) signiﬁcantly worse quality than the baseline.
This can be attributed to: (1) the extremely high-dimensionality; (2) the high compute of evaluating the score function; (3) the low required precision of the solution because we are satisﬁed as long as the error is not perceptible (e.g., one RGB increment on an
We devise our own SDE solver with these features in mind, resulting in an algorithm that can get around the problems encountered by off-the-shelf solvers.
Background 2.1 Score-based modeling with SDEs
Let x (
) ∈
R d be a sample from the data distribution p data .
The sample is gradually corrupted over time through a Forward Diffusion Process (FDP), a common type of Stochastic Differential Equation (SDE):
[Equation]
where f ( x , t ) : R d × R → R d is the drift, g ( t ) :
R → R is the diffusion coefﬁcient and w ( t ) is the Wiener process indexed by t ∈ [0 , 1] .
Data points and their probability distribution evolve along the trajectories { x ( t ) } 1 t =0 and { p t ( x ) } 1 t =0 respectively, with p 0 ≡ p data .
The functions f and g are chosen such that x (1) be approximately Gaussian and independent from x (0) .
Inference is achieved by reversing this diffusion, drawing x (1) from its Gaussian distribution and solving the Reverse
Diffusion Process (RDP) equal to:
[Equation]
where ∇ x log p t ( x ) is referred to as the score of the distribution at time t [Hyvärinen, 2005] and ¯w ( t ) is the Wiener process in which time ﬂows backward [Anderson, 1982].
The RDP requires the score (or p t ), which can be estimated by a neural network (referred to as the score network) by optimizing the following objective:
[Equation]
where λ ( t ) : R → R is chosen to be inversely proportional to E (cid:104)(cid:13)(cid:13) ∇ x ( t ) log p t ( x ( t )
| x (0))
(cid:13)(cid:13) 2 2 (cid:105) .
There are two primary choices for the FDP in the literature, which we discuss below.
2.2 Variance Exploding (VE) process
The Variance Exploding (VE) process consists in the following FDP:
[Equation]
Its associated transition kernel is:
[Equation]
Given a large enough σ 2 (1) , x (1) is approximately distributed as N ( 0 , σ 2 (1) I ) .
2.3 Variance Preserving (VP) process
The Variance Preserving (VP) process consists in the following FDP:
[Equation]
Its associated transition kernel is:
[Equation]
x (1) is approximately distributed as N ( 0 , I ) and does not depend on x (0) .
2.4 Solving the Reverse Diffusion Process (RDP)
There are many ways to solve the RDP; the most basic one being Euler-Maruyama (EM) [Kloeden and Platen, 1992], the SDE analog to Euler’s method for solving ODEs.
Song et al.
[2020a] also proposed Reverse-Diffusion , which consists in ancestral sampling
[Ho et al., 2020] with the same discretization used in the FDP.
With the Reverse-Diffusion, [Song et al., 2020a] obtained poor results unless applying an additional Langevin dynamics step after each Reversion-Diffusion step; this approach is only heuristically motivated.
Efﬁcient Method for Solving Reverse Diffusion Processes 3.1 Integration method
To dynamically adjust the step size over time, thereby gaining speed over a ﬁxed-step size algorithm, two integration methods are employed.
A lower-order ( x (cid:48) ) method is used conjointly with a higherorder ( x (cid:48)(cid:48) ) one.
The local error E ( x (cid:48) , x (cid:48)(cid:48) )
=
x (cid:48) − x (cid:48)(cid:48) is used to determine how stable the lower-order method is at the current step size; the closer to zero, the more appropriate the step size is.
From this information, the step size can be dynamically adjusted and x (cid:48) can be accept or rejected.
While higher-order solvers may achieve lower discretization errors, they require more function evaluations, and the improved precision might not be worth the increased computation cost [Lehn et al., 2002, Lamba, 2003].
For this reason, we use EM (order 1) and stochastic Improved Euler [Roberts, 2012] (order 2) for the lower and higher integration methods .
This approach only requires two score function evaluations; however, it leads to images of poor quality.
Thankfully, by using extrapolation (accepting x (cid:48)(cid:48) instead of x (cid:48) as our proposal), we improve over the baseline approach.
3.2 Tolerance
In ODE/SDE solvers, the local error is divided by a tolerance term.
We calculate the mixed tolerance through the maximum of the current and previous sample:
[Equation]
For image generation, we can set (cid:15) abs a priori.
During training, images are represented as ﬂoatingpoint tensors with range [ y min , y max ] , but when evaluated, they are converted to 8-bit color images: scaled to [0 , 255] and converted to the nearest integer.
This means that an absolute tolerance (cid:15) abs
=
y max − y min corresponds to tolerating local errors of at most one color (e.g., x (cid:48) ij with Red=5 and x (cid:48)(cid:48) ij with Red=6 is accepted, while Red=7 is not) channel-wise.
One-color differences are not perceptible and should not inﬂuence the metrics used for evaluating the generated images.
To control speed/quality, we vary (cid:15) rel .
3.3 Norm of the scaled error
The scaled error (the error scaled by the mixed tolerance) is calculated as
[Equation]
Many algorithms use q = ∞
[Lamba, 2003, Rackauckas and Nie, 2017a], where || x || ∞ =
max( x 1 , ..., x k ) over all k elements of x .
This is highly problematic for high-dimensional SDEs such as those in image-space because a single channel of a single pixel (out of 65536 pixels for a 256 × 256 color image) with a large local error will cause the step size to be reduced for all pixels.
To that effect, we instead use a scaled (cid:96) 2 norm (the (cid:96) 2 norm multiplied by √ k ).
3.4 Hyperparameters of the dynamic step size algorithm
Upon calculating the scaled error, we accept the proposal x (cid:48)(cid:48) if E q ≤ 1 and increment the time by h .
Whether or not it is accepted, we update the next step size h in the usual way:
[Equation]
where h max is the maximum step size, θ is the safety parameter, and r is an exponent-scaling term.
Although ODE theory tells which r is optimal, there is no such theory for SDEs [Rackauckas and
Nie, 2017a].
Following Rackauckas and Nie [2017a], we empirically testing values and found r ∈ [0 . 8 , 0 . 9] to work well (see Appendix E).
We arbitrarily chose r = 0 .
9 , θ = 0 .
9 , and h max as the largest step size possible, namely the remaining time t .
We have now deﬁned every aspect of the algorithm needed to numerically solve the Equation (2) for images.
The resulting algorithm is described in Algorithm 1.
Experiments
We apply our method to pre-trained VP or VE models from Song et al. [2020a] on CIFAR- [Krizhevsky et al., 2009] (32x32) and higher-res (256x256) LSUN-Church [Yu et al., 2015], and
Flickr-Faces-HQ (FFHQ) [Karras et al., 2019].
We measure performance through the Fréchet Inception Distance (FID) [Heusel et al., 2017] and the Inception Score (IS) [Salimans et al., 2016], whereas low FID and high IS correspond to higher quality/diversity.
We used ≤ 4 V100 GPUs to run the experiments.
Results are presented in Figure 1 (and Tables 1 and 2).
{+Probability Flow [Song et al., 2020a] solves an ODE instead of an SDE using Runge-Kutta 
+}
[Dormand and Prince, 1980] and denoising diffusion implicit models (DDIM)
[Song et al., 2020b] is only deﬁned for VP models.
Conclusion
We built an SDE solver to generate images of comparable (or better) quality to Euler-Maruyama at a much faster speed.
Our approach makes image generation with score-based models more accessible by shrinking the required computational budgets by a factor of 2 to 5 × , and presenting a sensible way of compromising quality for additional speed.
Nevertheless, data generation remains slow (a few minutes) compared to other generative models, which can generate data in a single forward pass of a neural network.
In practice, we let σ ( t ) = σ min (cid:16) σ max σ min (cid:17) t , where σ min = 0 . 01 and σ max ≈ max i (cid:80) Nj =
{+1 || x ( i ) − x ( j ) || is the maximum Euclidean distance between two samples from the dataset { x ( i ) } Ni =
 [Song and Ermon, 2020]. Using the maximum Euclidean distance ensures that x (1) does not depend on x (0) ; thus, x (1) is approximately distributed as N ( 0 , σ 2 (1) I ) . A.2 Variance Preserving (VP) process+}
In practice, we let β ( t ) = β min + t ( β max − β min ) , where β min = 0 . 1 and β max = 20 . Thus, x (1) is approximately distributed as N ( 0 , I ) and does not depend on x (0) .
B Algorithm
Algorithm
Require:
s θ , (cid:15) rel , (cid:15) abs , h init = 0 .
01 , r = 0 .
9 , θ = 0 .(cid:46)
for images: (cid:15) abs
=
y max − y mint ←h ← h init
Initialize x x (cid:48)
prev ← x while t > 0 do
Draw z ∼ N ( 0 , I ) x (cid:48) ← x − hf ( x , t ) + hg ( t ) 2 s θ ( x , t ) + √ hg ( t )
z (cid:46) Euler-Maruyama
˜x ← x − hf
( x (cid:48) , t − h )
+ hg ( t − h )
2 s θ ( x (cid:48) , t − h )
+
√ hg ( t − h )
z x (cid:48)(cid:48)
← 12 ( x (cid:48)
+ ˜x ) (cid:46)
Improved Euler (SDE version) δ ← max( (cid:15) abs , (cid:15) rel max( | x (cid:48)
| , | x (cid:48) prev | ))
(cid:46) Element-wise operations
E 2 ← 1 √ n (cid:107) ( x (cid:48) − x (cid:48)(cid:48) ) /
δ (cid:107)if E 2 ≤ 1 then (cid:46)
Accept x ← x (cid:48)(cid:48) (cid:46) Extrapolation t ← t − h x (cid:48)
prev ← x (cid:48)
h ← min( t, θhE − r 2 ) (cid:46) Dynamic step size update return x
C Tables of the results
[Table] [Table]
D DifferentialEquations.jl
Here, we report the preliminary experiments we ran with the DifferentialEquations.jl Julia package [Rackauckas and Nie, 2017b] before devising our own SDE solver.
As can be seen, most methods either did not converge (with warnings of "instability detected") or converged, but were much slower than Euler-Maruyama.
The only promising method was Lamba’s method [Lamba, 2003].
Note that an algorithm has strong-order p when the local error from t to t
+ h is O
( h p +1 ) ).
[Table]
E Effects of modifying Algorithm
Change(s) in Algorithm
[Table]
δ
[Table]
No Extrapolation (thus, using Euler–Maruyama) 9.
q = ∞ 9.
r = .9. r = .9.
r =9.
Variations of Lamba [2003] Algorithm r = 0 .
5 , Lamba integration 7.
r = 0 .
5 , Lamba integration, Extrapolation 7.
r = 0 .
5 , Lamba integration, q = ∞ 9.
r = 0 .
5 , Lamba integration, q = ∞ , θ = 0 .9.
As can be seen, most chosen settings lead to better results.
However, r seems to have little impact on the FID.
Still, using r ∈ [0 . 8 , 0 . 9] lead to a little bit less score function evaluations and sometimes lead to lower FID.
Change(s) in Algorithm
IS
[Table]
No Extrapolation (thus, using Euler–Maruyama)
[Table]
∞
[Table]
Variations
of Lamba [2003] Algorithm
r = 0 .
5 , Lamba integration
[Table]
r = 0 .
5 , Lamba integration, Extrapolation
[Table]
r = 0 .
5 , Lamba integration, q = ∞
[Table]
r = 0 .
5 , Lamba integration, q = ∞ , θ = 0 .
[Table]
F Implementation Details
We started from the original code by Song et al.
[2020a]
but changed a few settings concerning the
SDE solving.
This creates some very minor difference between their reported results and ours.
For the VP and VP-deep models, we obtained 2.55 and 2.49 instead of the original 2.55 and 2.41 for the baseline method (EM).
For the VE and VE-deep models, we obtained 2.40 and 2.21 instead of the original 2.38 and 2.20 for the baseline method (Reverse-Diffusion with Langevin).
When solving the SDE, time followed the sequence t
= 1 , t i
=
t i − 1 − 1 − (cid:15)N , where N = 1000 for
CIFAR-10, N = 2000 for LSUN, (cid:15) = 1 e − 3 for VP models, and (cid:15) = 1 e − 5 for VE models.
Meanwhile, the actual step size h used in the code for Euler-Maruyama (EM) was equal to 1 N .
Thus, there was a negligible difference between the step size used in the algorithm ( h = 1 N ) and the actual step size implied by t ( h = 1 − (cid:15)N ).
Note that this has little to no impact.
The bigger issue is at the last predictor step was going from t = (cid:15) to t = (cid:15) − 1 N < 0 .
Thus, t was made negative.
Furthermore the sample was denoised at t < 0 while assuming t = (cid:15) .
There are two ways to ﬁx this issue: 1) take only a step from t = (cid:15) to t = 0 and do not denoise (since you cannot denoise with the incorrect t or with t = 0 ), or 2) stop at t = (cid:15) and then denoise.
Since denoising is very helpful, we took approach 2; however, both approaches are sensible.
Finally, denoising was not implemented correctly before.
Denoising was implemented as one predictor step (Reverse-Diffusion or EM) without adding noise.
This corresponds to:
[Equation]
x )
[Equation]
At the last iteration, this incorrect denoising would be:
[Equation]
for VE and
[Equation]
for VP.
Meanwhile, the correct way to denoise based on Tweedie formula [Efron, 2011] is:
[Equation]
where Var [ x ( t ) | x (0)] is the variance of the transition kernel:
Var [ x ( t ) | x (0)]
= σ min = 0 .
01 for VE and Var [ x ( t ) | x (0)]
= 1 .
This means that the correct Tweedie formula corresponds to
[Equation]
for VE and
[Equation]
for VP.
As can be seen, denoising has a very small impact on VE so the difference between the correct and incorrect denoising is minor.
Meanwhile, for VP the incorrect denoising lead to a tiny change, while the correct denoising lead to a large change.
In practice, we observe that changing the denoising method to the correct one does not signiﬁcantly affect the FID with VE, but lowers down the FID signiﬁcantly with VP. G Inception Score on CIFAR-
[Table]
H Stability and Bias of the Numerical Scheme
The following constructions rely on the underlying assumption of the stochastic dynamics being driven by a wiener process.
More so, we also assume that the Brownian motion is time symmetrical.
Both assumptions are consistent and widely used in the literature; for example, see [Gardiner, 2009] [Arnold, 1974].
The method described in Algorithm 1 gives us a signiﬁcant speedup in terms of computing time and actions.
Albeit the speed up comes from a piece-wise step in the algorithm combining the traditional
Euler Maruyama (EM) with a form of adaptive step size predictor-corrector.
Here we show that both the stability and the convergence of the EM scheme are conserved by introducing the extra adaptive stepsize of our new scheme.
As a ﬁrst step, we deﬁne the stability and bias in a Stochastic Differential
Equation (SDE) numerical solution.
We denote (cid:60) ( λ ) as the real value of a complex-valued λ .
The linear test SDE is deﬁned in the following way:
[Equation]
with its numerical counterpart
[Equation]
where the z n are random variables that do not depend on y 0 , y
...... y n or λ and the EM scheme is y n {++
+}
=
(1 + hλ ) y n + z n .
A numerical scheme is asymptotically unbiased with step size h > [-
-]0 if, for a given linear SDE (5) driven by a two-sided Wiener process, the distribution of the numerical solution y n converges as n → ∞ to the normal distribution with zero mean and variance σ 2 2 | λ | [Artemiev and Averina, 2011].
This stems from the fact that a solution of a linear SDE (5) is a Gaussian process whenever the initial condition is Gaussian (or deterministic); thus, there are only two moments that control the bias in the algorithm:
[Equation]
A numerical scheme with step size h is numerically stable in mean if the numerical solution y ( h ) n applied to a linear SDE satisﬁes
[Equation]
and is stable in mean square
[Saito and Mitsui, 1996] if we have that
[Equation]
In what follows, we will trace the criteria for bias through our algorithm and show that it remains unbiased.
By construction, the ﬁrst EM step remains unbiased, while for the RDP, we write down the time reverse Wiener process as
[Equation]
in the reverse time steps h i.e., t − nh , t − 2 nh ,
[Equation]
Thus
[Equation]
if
[Equation]
then
[Equation]
In Algorithm [Equation] we are performing consecutive steps forward and backwards in time so t = 2 h such that
[Equation]
Thus, the scheme is both numerically stable and unbiased with respect to the mean.
Next, we focus on the numerical solution in mean square:
[Equation]
Under the same assumption of consecutive steps
[Equation]
we have that
[Equation]
Assuming the imaginary part of λ is null [Equation] we have that
[Equation]
Thus, the numerical scheme is stable and unbiased in the mean square.
Following the two steps for computation of x (cid:48) and ˜x , the step size decreases and does not change size; thus, all the above statements hold [Equation] and the entire algorithm is stable and unbiased with respect to both the mean and square mean.
I
Samples
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]