S CORE -B ASED G ENERATIVE M ODELING THROUGH S TOCHASTIC D IFFERENTIAL E QUATIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
Creating noise from data is easy; creating data from noise is generative modeling.
We present a stochastic differential equation (SDE) that smoothly transforms a complex data distribution to a known prior distribution by slowly injecting noise, and a corresponding reverse-time SDE that transforms the prior distribution back into the data distribution by slowly removing the noise.
Crucially, the reversetime SDE depends only on the time-dependent gradient ﬁeld (a.k a., score) of the perturbed data distribution.
By leveraging advances in score-based generative modeling we can accurately estimate these scores with neural networks, and use numerical SDE solvers to generate samples.
We show that this framework encapsulates previous approaches in diffusion probabilistic modeling and scorebased generative modeling, and allows for new sampling procedures.
In particular we introduce a predictor-corrector framework to correct errors in the evolution of the discretized reverse-time SDE
We also derive an equivalent neural ODE that samples from the same distribution as the SDE which enables exact likelihood computation, and improved sampling efﬁciency.
In addition our framework enables conditional generation with an unconditional model as we demonstrate with experiments on class-conditional generation image inpainting, and colorization.
Combined with multiple architectural improvements we achieve record-breaking performance for unconditional image generation on CIFAR-10 with an Inception score of 9.89 and FID of 2.
a competitive likelihood of 3.10 bits/dim, and demonstrate high ﬁdelity generation of 1024 ˆ 1024 images for the ﬁrst time from a score-based generative model
I NTRODUCTION
Two successful classes of probabilistic generative models involve sequentially corrupting training data with slowly increasing noise, and then learning to reverse this corruption in order to form a generative model of the data.
Score matching with Langevin dynamics (SMLD)
(Song & Ermon, 2019) estimates the score ( i.e ., the gradient of the log probability density) at each noise scale, and then uses Langevin dynamics to sample from a sequence of decreasing noise scales during generation.
Denoising diffusion probabilistic modeling (DDPM)
(Sohl-Dickstein et al., 2015; Ho et al., 2020) trains a sequence of probabilistic models to reverse each step of the noise corruption, using knowledge of the functional form of the reverse distributions to make training tractable.
For continuous state spaces, the DDPM training objective implicitly computes scores at each noise scale.
We therefore refer to these two model classes together as score-based generative models .
Score-based generative models, and related techniques (Bordes et al., 2017;
Goyal et al., 2017), have proven effective at generation of images (Song & Ermon, 2019; 2020; Ho et al., 2020), audio (Chen et al., 2020; Kong et al., 2020), graphs (Niu et al., 2020), and shapes (Cai et al., 2020).
However, the practical performance of these two model classes is often quite different for reasons that are not fully understood.
More generally, despite obvious surface similarities, the relationship between SMLD and DDPM approaches is largely unexplored.
We address these open questions by unifying and generalizing both approaches through the lens of stochastic differential equations (SDEs).
Instead of perturbing data with a ﬁnite number of noise distributions, we consider a continuum of distributions that evolve over time according to a diffusion process.
This process progressively diffuses a data point into random noise, and is given by a
[Figure]
prescribed SDE that does not depend on the data and has no trainable parameters.
By reversing this process, we can smoothly mold random noise into data for sample generation.
Crucially, this reverse process satisﬁes a reverse-time SDE (Anderson, 1982), which can be derived from the forward
SDE given the score of the marginal probability densities as a function of time.
We can therefore approximate the reverse-time SDE by training a time-dependent neural network to estimate the scores, and then produce samples using numerical SDE solvers.
Our key idea is summarized in Fig.
Our proposed framework has several theoretical and practical contributions:
Flexible sampling: We can employ any general-purpose SDE solver to integrate the reverse-time
SDE for sampling.
In addition, we propose two special methods not viable for general SDEs: (i)
Predictor-Corrector (PC) samplers that combine numerical SDE solvers with score-based MCMC approaches, such as Langevin MCMC (Parisi, 1981; Grenander & Miller, 1994); and (ii) deterministic samplers based on the probability ﬂow ordinary differential equation (ODE).
The former uniﬁes and improves over existing sampling methods for score-based models.
The latter allows for exact likelihood computation , efﬁcient and adaptive sampling via black-box ODE solvers, ﬂexible data manipulation via latent codes, and a unique encoding .
Controllable generation: We can modulate the generation process by conditioning on information not available during training, because the conditional reverse-time SDE can be efﬁciently computed from unconditional scores.
This enables applications such as class-conditional generation, image inpainting, and colorization using a single unconditional score-based model without re-training.
Uniﬁed picture: The approaches of SMLD and DDPM can be uniﬁed into our framework as discretizations of different SDEs.
Although Ho et al.
(2020) has reported higher sample quality than Song & Ermon (2019; 2020), we show that with better architectures and our new sampling algorithms, the latter can achieve new state-of-the-art Inception and FID scores on CIFAR-10, as well as high-ﬁdelity generation of 1024 ˆ 1024 samples.
This indicates that either SDE could be advantageous, and should be tuned jointly with other design choices for score-based generative modeling.
B ACKGROUND 2. D ENOISING SCORE MATCHING WITH L ANGEVIN DYNAMICS (SMLD)
Let p σ p ˜ x
| x q ﬁ N p ˜ x ; x , σ 2 I q be a perturbation kernel, and p σ p ˜ x
q ﬁ ş p data p x q p σ p ˜ x
| x q d x , where p data p x q denotes the data distribution.
Consider a sequence of positive noise scales σ min “ σ 1 ă σ 2 ă ¨ ¨ ¨ ă σ N “ σ max .
Typically, σ min is small enough such that p σ min p x q « p data p x q , and σ max is large enough such that p σ max p x q « N p x ; 0 , σ 2 max I q .
Song & Ermon (2019) propose to train a Noise
Conditional Score Network (NCSN), denoted by s θ p x , σ q , with a weighted sum of denoising score matching (Vincent, 2011) objectives:
[Equation]
Given sufﬁcient data and model capacity, the optimal score-based model s θ ˚ p x , σ q matches ∇ x log p σ p x q almost everywhere for σ P t σ i u N i “ 1 .
For sampling, Song & Ermon (2019) run M steps of
[Equation]
where (cid:15) i ą 0 is the step size, and z mi is standard normal.
The above is repeated for i “ N, N ´ 1 , ¨ ¨ ¨ , 1 in turn with x 0 N „ N p x | 0 , σ 2 max I q and x 0 i “ x Mi `
when i ă N .
As M Ñ 8 and (cid:15) i Ñfor all i , x N 1 becomes an exact sample from p σ min p x q « p data p x q under some regularity conditions.
2. D ENOISING DIFFUSION PROBABILISTIC MODELS (DDPM)
Sohl-Dickstein et al.
; Ho et al.
(2020) consider a sequence of positive noise scales
ă β 1 , β 2 , ¨ ¨ ¨ , β N ă 1 .
For each training data point x 0 „ p data p x q , a discrete Markov chain t x 0 , x 1 , ¨ ¨ ¨ , x N u is constructed such that p p x i | x i ´ 1 q “ N p x
i ; ?
1 ´ β i x i ´ 1 , β
i
I q , and therefore p α i p x i | x 0 q “ N p x i ; ?
α i x 0 , p 1 ´ α i q I q , where α i ﬁ ś ij “ 1 p 1 ´ β j q .
Similar to SMLD, we can denote the perturbed data distribution as p α i p ˜ x
q ﬁ ş p data p x q p
α
i p ˜ x
| x q d x .
The noise scales are prescribed such that x N is an approximate sample from N p 0 , I q .
A variational Markov chain in the reverse direction is parameterized with p θ p x i ´ 1 | x i q “ N p x
i ´ 1 ;
[Equation]
β
[Equation]
θ p x
i , i qq , β i
I q , and trained with a re-weighted variant of the evidence lower bound (ELBO):
[Equation]
After solving Eq.
(3) to get the optimal model s θ ˚ p x , i q , samples can be generated by starting from x N „ N p 0 , I q and following the estimated reverse Markov chain as below
[Equation]
We call this method ancestral sampling , since it amounts to performing ancestral sampling from the graphical model ś Ni “ 1 p θ p x i ´ 1 | x i q .
The objective Eq.
(3) described here is equivalent to
L simple in Ho et al.
but we re-write it in a slightly different form to expose more similarity to Eq.
Like Eq.
(1), Eq.
(3) is also a weighted sum of denoising score matching objectives, which implies that the optimal model, s θ ˚ p ˜ x , i q , matches the score of the perturbed data distribution, ∇ x log p
α i p x q .
Moreover, we observe that the weights of the i -th summand in Eq.
(1) and Eq.
(3), namely σ 2 i and p 1 ´ α i q , are related to corresponding perturbation kernels in the same functional form:
[Equation]
{ E
[Equation]
(cid:107)∇
[Equation]
S CORE - BASED GENERATIVE MODELING WITH SDE S Perturbing
data with multiple noise scales is key to the success of previous methods.
We propose to generalize this idea further to an inﬁnite number of noise scales, such that perturbed data distributions evolve according to an SDE as the noise intensiﬁes.
An overview of our framework is given in Fig.
3. P ERTURBING DATA WITH SDE S
Our goal is to construct a diffusion process t x p t qu Tt “ 0 indexed by a continuous time variable t P r 0 , T s , such that x p 0 q „ p 0 , for which we have a dataset of i.i.d.
samples, and x p T q „ p T , for which we have a tractable form to generate samples efﬁciently.
In other words, p 0 is the data distribution and p T is the prior distribution.
This diffusion process can be modeled as the solution to an Itˆo SDE:
[Equation]
where w is the standard Wiener process (a.k.a., Brownian motion), f p¨ , t q : R d Ñ R d is a vectorvalued function called the drift coefﬁcient of x p t q , and g p¨q : R Ñ R is a scalar function known as the diffusion coefﬁcient of x p t q .
For ease of presentation we assume the diffusion coefﬁcient is a scalar (instead of
a d ˆ d matrix) and does not depend on x , but our theory can be generalized to hold in those cases (see Appendix A).
The SDE has a unique strong solution when the coefﬁcients are globally Lipschitz in both state and time (Øksendal, 2003).
We hereafter denote by p t p x q the
[Figure]
probability density of x p t q , and use p st p x p t q | x p s qq to denote the transition kernel from x p s q to x p t q , where 0 ď s ă t ď T .
Typically, p T is an unstructured prior distribution that contains no information of p 0 , such as a
Gaussian distribution with simple mean and variance.
There are various ways of designing the SDE in Eq.
(5) such that it diffuses the data distribution into a ﬁxed prior distribution.
We provide two examples later in Section 3.4.
3. G ENERATING SAMPLES BY REVERSING THE SDE
By starting from samples of x p T q „ p T and reversing the process, we can obtain samples x p 0 q „ p 0 .
A remarkable result from Anderson (1982) states that the reverse of a diffusion process is also a diffusion process, running backwards in time and given by the reverse-time SDE:
[Equation]
(6) where ¯ w is a standard Wiener process when time ﬂows backwards from T to 0 , and d t is an inﬁnitesimal negative timestep.
Once the score of each marginal distribution, ∇ x log p t p x q , is known for all t , we can derive the reverse diffusion process from Eq.
(6) and simulate it to sample from p 0 .
3. E STIMATING SCORES FOR THE SDE
The score of a distribution can be estimated by training a score-based model on samples with score matching (Hyv¨arinen, 2005; Song et al., 2019a).
To estimate ∇ x log p t p x q , we can train a time-dependent score-based model s θ p x , t q via a continuous generalization to Eqs.
(1) and (3):
[Equation]
Here λ : r 0 , T s Ñ R ` is a positive weighting function, t is sampled from a uniform distribution over r 0 , T s , x p 0 q „ p 0 p x q and x p t q „ p 0 t p x p t q | x p 0 qq .
With sufﬁcient data and model capacity, score matching ensures that the optimal solution to Eq.
(7), denoted by s θ ˚ p x , t q , equals ∇ x log p t p x q for almost all x and t .
As in SMLD and DDPM, we can typically choose λ
[Equation]
Note that Eq.
(7) uses denoising score matching, but other score matching objectives, such as sliced score matching (Song et al., 2019a) and ﬁnite-difference score matching (Pang et al., 2020) are also applicable here.
We need to know the transition kernel p 0 t p x p t q | x p 0 qq to efﬁciently solve Eq.
When f p¨ , t q is afﬁne, the transition kernel is always a Gaussian distribution, where the mean and variance are often known in closed-forms and can be obtained with standard techniques (see Section 5.5 in S¨arkk¨a &
Solin (2019)).
For more general SDEs, we may solve Kolmogorov’s forward equation (Øksendal,
2003) to obtain p 0 t
p x p t q | x p 0 qq .
Alternatively, we can simulate the SDE to sample from p 0 t p x p t q | x p 0 qq and replace denoising score matching in Eq.
(7) with sliced score matching for model training, which does not require computing ∇ x p t q log p 0 t p x p t q | x p 0 qq (see Appendix A).
3. S PECIAL CASES : VE AND VP SDE S
The noise perturbations used in SMLD and DDPM can be regarded as discretizations of two different
SDEs, which we name as Variance Exploding (VE) and Variance Preserving (VP) SDEs respectively.
Below we provide a brief discussion and relegate more details to Appendix B.
When using a total of N noise scales, each perturbation kernel p σ i p x | x 0 q of SMLD corresponds to the distribution of x i in the following Markov chain:
[Equation]
where z i ´ 1 „ N p 0 , I q , and we have introduced σ 0 “ 0 to simplify the notation.
In the limit of
N Ñ 8 , t
σ i u Ni “ 1 becomes a function σ p t q , z i becomes z p t q , and the Markov chain t x i u Ni “ 1 becomes a continuous stochastic process t x p t qu 1 t “ 0 , where we have used a continuous time variable t P r 0 , 1 s for indexing, rather than an integer i .
The process t x p t qu 1 t “ 0 is given by the following SDE
[Equation]
Likewise for the perturbation kernels t p α i p x | x 0 qu Ni “ 1 of DDPM, the discrete Markov chain is
[Equation]
As N Ñ 8 , Eq.
(10) converges to the following SDE,
[Equation]
Therefore, the noise perturbations used in SMLD and DDPM correspond to discretizations of SDEs
Eq.
(9) and Eq.
(11) respectively.
Interestingly, the SDE of Eq.
(9) always gives a process with exploding variance when t Ñ 8 , whilst the SDE of Eq.
(11) yields a process with bounded variance (proof in Appendix B).
Due to this difference, we hereafter refer to Eq.
(9) as the Variance Exploding (VE) SDE, and Eq.
(11) the Variance Preserving (VP) SDE.
Since both VE and VP SDEs have afﬁne drift coefﬁcients, their perturbation kernels p 0 t p x p t q | x p 0 qq are both Gaussian and can be computed in closed-form, as discussed in Section 3.3.
This makes training with Eq.
(7) particularly efﬁcient.
S OLVING THE REVERSE SDE
After training a time-dependent score-based model, we can use it to construct the reverse-time SDE and then solve it with numerical approaches to generate samples from p 0 .
4. G ENERAL - PURPOSE NUMERICAL SDE SOLVERS
Numerical solvers provide approximate trajectories from SDEs.
Many general-purpose numerical methods exist for solving SDEs, such as Euler-Maruyama and stochastic Runge-Kutta methods (Kloeden & Platen, 2013), which correspond to different discretizations of the stochastic dynamics.
We can apply any of them to the reverse-time SDE for sample generation.
Ancestral sampling—the sampling method of DDPM (Eq.
(4))—actually corresponds to one special discretization of the reverse-time VP SDE (Eq.
(see Appendix E).
In theory, other discretizations should be able to perform comparably or better.
To verify this, we propose to apply the same discretization as the forward SDE to the reverse-time SDE (details in Appendix E).
The resulting samplers for SMLD and DDPM are given as the “predictor” part in Algs.
1 and 2, and we call this family of methods the reverse diffusion sampler .
As shown in Tab. 1, these sampling methods perform slightly better than ancestral sampling for both SMLD and DDPM models on CIFAR-
(DDPM-type ancestral sampling can be applied to SMLD models, see Appendix F).
1: x N „ N p 0 , σ 2 max I q 2: for i “ N ´ 1 to 0 do 3: x 1 i Ð x i ` 1 ` p σ 2 i ` 1 ´ σ 2 i q s θ ˚ p x i ` 1 , σ i ` 1 q 4: z „ N p 0 , I q 5: x i Ð x 1 i ` b σ 2 i ` 1 ´ σ 2 i z 6: for j “ 1 to M do 7: z „ N p 0 , I q 8: x i Ð x i ` (cid:15) i s θ ˚ p x i , σ i q ` ? 2 (cid:15)
i z 9: return x
Algorithm 2 PC sampling (VP SDE) 1: x N „ N p 0 , I q 2: for i “ N ´ 1 to 0 do 3: x 1 i Ð p 2 ´ ?
1 ´ β i ` 1 q x i ` 1 ` β i ` 1 s θ ˚ p x
i ` 1 ,i ` 1 q 4: z „ N p 0 , I q 5: x i Ð x 1 i ` ?
β
i ` 1 z 6: for j “ 1 to M do 7: z „ N p 0 , I q 8: x i Ð x i ` (cid:15)
i s θ ˚ p
x
i , i q ` ? 2 (cid:15)
i z 9: return x
4. P REDICTOR - CORRECTOR SAMPLERS
Unlike generic SDEs, we have additional information that can be used to improve solutions.
Since we have a score-based model s θ ˚ p x , t q « ∇ x log p t p x q , we can employ score-based MCMC approaches, such as Langevin MCMC (Parisi, 1981; Grenander & Miller, 1994) or HMC (Neal et al., 2011) to sample from p t directly, and correct the solution of a numerical SDE solver.
At each time step, the numerical SDE solver gives an estimate of the sample at the next time step, playing the role of a “predictor”.
The score-based MCMC approach corrects the marginal distribution of the estimated sample, playing the role of a “corrector”.
The idea is analogous to Predictor-Corrector methods for solving systems of equations (Allgower & Georg, 2012), and we similarly name our hybrid sampling algorithms Predictor-Corrector (PC) samplers.
When using the reverse diffusion SDE solver (Appendix E) as the predictor, and annealed Langevin dynamics (Song & Ermon, 2019) as the corrector, we have Algs.
1 and 2 for VE and VP SDEs respectively, where t (cid:15) i u N ´
i “ 0 are step sizes for Langevin dynamics (speciﬁed in Appendix G).
PC samplers generalize the original sampling methods of SMLD and DDPM:
We can recover annealed
Langevin dynamics for SMLD by removing the predictor part, or regain the ancestral sampling method for DDPM when using it as the predictor and removing the corrector part.
We test PC samplers on SMLD/DDPM models trained with Eqs.
(1) and (3), instead of our new continuous objective Eq.
This exhibits the compatibility of PC samplers to score-based models trained with a ﬁxed number of noise scales.
We compare different samplers in Tab. 1,
where probability ﬂow is a predictor to be discussed in Section 4.3.
Experimental details are given in
Appendix G. Due to additional corrector steps, PC samplers require twice as much computation as predictor-only samplers.
To match the computation, one has to use more noise scales for predictor-only samplers, which is difﬁcult for models trained with a ﬁxed number of noise scales as in SMLD/DDPM.
As a workaround, we use an ad-hoc method to interpolate between noise scales (Appendix G).
Observations from Tab. 1 include: (i) our reverse diffusion method uniformly outperforms ancestral sampling; (ii) PC samplers uniformly outperform predictor-only methods
with double computation; (iii) PC samplers can outperform predictor-only and corrector-only methods in most cases using
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
matched computation.
The with different noise interpolation strategies (Tab. 4 in Appendix G).
In Fig.
3, we additionally provide qualitative comparison for models trained with the continuous objective Eq.
(7), where PC samplers clearly surpass predictoronly samplers under comparable computation, when using a proper number of corrector steps.
4. P ROBABILITY FLOW AND EQUIVALENCE TO NEURAL ODE S
Score-based models enable another numerical method for solving the reverse-time SDE.
For all diffusion processes, there exists a corresponding deterministic process whose trajectories share the same marginal probability densities t p t p x qu Tt “ 0 .
This deterministic process satisﬁes an ODE (more details in Appendix D.1):
[Equation]
which can be determined from the SDE once scores are known.
We name the ODE in Eq.
(12) the probability ﬂow ODE .
As the score functions are typically parameterized by a neural network, this is an example of a neural ODE (Chen et al., 2018).
Efﬁcient sampling As with neural ODEs, we can sample x p 0 q „ p 0 by solving Eq.
(12) from different ﬁnal conditions x p T q „ p T .
Using a ﬁxed discretization strategy we can generate competitive samples, especially when used in conjuction with correctors (Tab. 1, “probability ﬂow sampler”, details in Appendix D.3).
Using a black-box ODE solver (Dormand & Prince, 1980) allows us to explicitly trade-off accuracy for efﬁciency.
With a larger error tolerance, the number of function evaluations can be reduced by over 90% without affecting the visual quality of samples (Fig. 4).
Exact likelihood computation Leveraging the connection to neural ODEs, we can compute the density deﬁned by Eq.
(12) via the instantaneous change of variables formula (Chen et al., 2018).
This allows us to compute the exact likelihood on any input data (details in Appendix D.2).
As an example, we report negative log-likelihoods (NLLs) measured in bits/dim on the CIFAR-10 dataset in Tab. 2.
We compute log-likelihoods on uniformly dequantized data, and only compare to models evaluated in the same way (excluding models evaluated with variational dequantization (Ho et al., 2019) or
[Table] [Table]
discrete data).
Main results: (i) For the same DDPM model in Ho et al.
(2020), we obtain better bits/dim compared to the upper bound given by ELBO, since our likelihoods are exact; (ii) Using the same architecture, we trained another DDPM model with our continuous denoising score matching objective in Eq.
(7) on uniformly dequantized data.
This continuously-trained model (DDPM cont.
in Tab. 2) has much lower bits/dim compared to the original one, setting a new record bits/dim on uniformly dequantized CIFAR-10 even without maximum likelihood training .
Manipulating latent representations By integrating Eq.
(12), we can encode any datapoint x p 0 q into a latent space x p T q .
Decoding can be achieved by integrating a corresponding ODE for the reverse-time SDE.
As is done with other invertible models such as neural ODEs and normalizing ﬂows (Dinh et al., 2016; Song et al., 2019b), we can manipulate this latent representation for image editing, such as interpolation, and temperature scaling (see Fig.
4 and Appendix D.4).
Uniquely identiﬁable encoding Unlike all current invertible models, our encoding is uniquely identiﬁable , meaning that with sufﬁcient training data, model capacity, and optimization accuracy, the encoding for an input is uniquely determined by the data distribution (Roeder et al., 2020).
This is because our forward SDE, Eq.
(5), has no trainable parameters, and its associated probability ﬂow
ODE, Eq.
(12), provides the same trajectories given perfectly estimated scores.
We provide additional empirical veriﬁcation on this property in Appendix D.5.
4. A RCHITECTURE CHOICES
Although SMLD and DDPM can be viewed as different instantiations of a uniﬁed framework, there is a performance gap reported in previous papers.
The best FID values of SMLD models on CIFARis 10.23 (Song & Ermon, 2020), whereas for DDPM it is 3.17 (Ho et al., 2020).
With PC samplers and the same model architecture in Ho et al.
(2020), the gaps can be reduced, but score-based models trained with the VE SDE still perform slightly worse (see Tab. 1).
This raises the question of whether the variance preserving property of the VP SDE is always preferable or whether there is an interaction with model architecture.
To answer this question, we performed an extensive architecture search (see Appendix H).
Our optimal architecture for the VE SDE, dubbed NCSN++, achieves an FID of 2.45 on CIFAR-10 with
PC samplers, whereas the optimal architecture for the VP SDE over the same architecture sweep obtains an FID of 2.88.
This indicates that the VE SDE can be advantageous, and practitioners likely need to experiment with different SDEs for new domains and architectures.
Based on these improved architectures, we were able to further improve FID with our continuous training objective in Eq.
(7) (FID 2.38), and increasing the number of residual blocks (FID 2.2, see Tab. 3 and Appendix H.2).
This model sets new records for both inception score and FID on unconditional generation for CIFAR-10.
Surprisingly, we can achieve better FID than the previous best conditional generative model without requiring labeled data.
Finally, with all improvements together, we are able to obtain the ﬁrst set of high-ﬁdelity samples on CelebA-HQ 1024 ˆ 1024 from score-based models (see Appendix H.3).
C ONTROLLABLE GENERATION
The continuous structure of our framework allows us to not only produce data samples from p 0 , but also from p
p x p 0 q | y q if p t p y
| x p t qq is known.
Given a forward SDE as in Eq.
(5), we can sample from p t p x p t q | y q by starting from p T p x p T q | y q and solving a conditional reverse-time SDE:
[Equation]
When y represents class labels, we can train a time-dependent classiﬁer p t p y | x p t qq for classconditional sampling.
Since the forward SDE is tractable, we can easily create training data p x p t q , y q for the time-dependent classiﬁer by ﬁrst sampling p x p 0 q , y q from a dataset, and then sampling x p t q „ p 0 t p x p t q | x p 0 qq .
Afterwards, we may employ a mixture of cross-entropy losses over different time steps, like Eq.
(7), to train the time-dependent classiﬁer
p t
p
y
| x p t qq .
Class-conditional
CIFAR-10 samples can be found in Fig.
5 (left).
See more details and results are in Appendix I.
Imputation is a special case of conditional sampling.
Suppose we have an incomplete data point y where only some subset, Ω p y q is known.
Imputation amounts to sampling from p p x p 0 q | Ω p y qq , which we can accomplish using an unconditional model (see Appendix I.2).
Colorization is a special case of imputation, except that the known data dimensions are coupled.
We can decouple these data dimensions with an orthogonal linear transformation, and perform imputation in the transformed space (details in Appendix I.3).
Fig.
5 (right) shows results for inpainting and colorization achieved with unconditional time-dependent score-based models.
D ISCUSSION
We presented a framework for score-based generative modeling based on SDEs.
Our work enables a better understanding of existing approaches, new sampling algorithms, exact likelihood computation, uniquely identiﬁable encoding, latent code manipulation, and brings new conditional generation abilities to the family of score-based generative models.
While our proposed sampling approaches improve results and enable more efﬁcient sampling, they remain slower at sampling than GANs (Goodfellow et al., 2014) on the same datasets.
Identifying ways of combining the stable learning of score-based generative models with the fast sampling of implicit models like GANs remains an important research direction.
Additionally, the breadth of samplers one can use when given access to score functions introduces a number of hyper-parameters.
Future work would beneﬁt from improved methods to automatically select and tune these hyperparameters, as well as more extensive investigation on the merits and limitations of various samplers.
A PPENDIX
We include several appendices with additional details, derivations, and results.
Our framework allows general SDEs with matrix-valued diffusion coefﬁcients that depend on the state, for which we provide a detailed discussion in Appendix A.
We give a full derivation of VE and VP SDEs in Appendix B, and discuss how to use them from a practitioner’s perspective in Appendix C. We elaborate on the probability ﬂow formulation of our framework in Appendix D, including a derivation of the probability ﬂow ODE (Appendix D.1), exact likelihood computation (Appendix D.2), probability ﬂow sampling (Appendices D.3 and D.4) and experimental veriﬁcation on uniquely identiﬁable encoding (Appendix D.5).
We give a full description of the reverse diffusion sampler in Appendix E, and the DDPM-type ancestral sampler for SMLD models in Appendix F.
We provide full experimental details on Predictor-Corrector samplers in Appendix G.
We explain our model architectures and detailed experimental settings in Appendix H, with 1024 ˆ 1024 CelebA-HQ samples therein.
Finally, we detail on the algorithms for controllable generation in Appendix I, and include extended results for class-conditional generation (Appendix I.1), image inpainting (Appendix I.2) and colorization (Appendix I.3).
A T HE FRAMEWORK FOR MORE GENERAL SDE S
In the main text, we introduced our framework based on a simpliﬁed SDE Eq.
(5) where the diffusion coefﬁcient is independent of x p t q .
It turns out that our framework can be extended to hold for more general diffusion coefﬁcients.
We can consider SDEs in the following form:
[Equation]
(14) where f p¨ , t q : R d Ñ R d and G p¨ , t q : R d Ñ R d ˆ d .
We follow the It ˆ o interpretation of SDEs throughout this paper.
According to (Anderson, 1982), the reverse-time SDE is given by ( cf ., Eq.
d x “ t
[Equation]
where we deﬁne ∇ ¨ F p x q : “ p ∇ ¨ f 1 p x q , ∇ ¨ f 2 p x q , ¨ ¨ ¨ , ∇ ¨ f d p x qq T for a matrix-valued function
F p x q : “ p f 1 p x q , f 2 p x q , ¨ ¨ ¨ , f d p x qq T throughout the paper.
The probability ﬂow ODE corresponding to Eq.
(14) has the following form ( cf ., Eq.
(12), see a detailed derivation in Appendix D.1):
[Equation]
Finally for conditional generation with the general SDE Eq.
(14), we can solve the conditional reverse-time SDE below ( cf ., Eq.
(13), details in Appendix I):
[Equation]
When the drift and diffusion coefﬁcient of an SDE are not afﬁne, it can be difﬁcult to compute the transition kernel p 0 t p x p t q | x p 0 qq in closed form.
This hinders the training of score-based models, because Eq.
(7) requires knowing ∇ x p t q log p 0 t
p x p t q | x p 0 qq .
To overcome this difﬁculty, we can replace denoising score matching in Eq.
(7) with other efﬁcient variants of score matching that do not require computing ∇ x p t q log p 0 t p x p t q | x p 0 qq .
For example, when using sliced score matching (Song et al., 2019a), our training objective Eq.
(7) becomes
[Equation]
where λ : r 0 , T s Ñ R ` is a positive weighting function, t „ U p 0 , T q , E r v s “ 0 , and Cov r v s “ I .
We can always simulate the SDE to sample from p 0 t p x p t q | x p 0 qq , and solve Eq.
(18) to train the time-dependent score-based model s θ p x , t q .
B D ERIVATIONS OF VE AND VP SDE S
Below we provide detailed derivations to show that the noise perturbations of SMLD and DDPM are discretizations of the Variance Exploding (VE) and Variance Preserving (VP) SDEs respectively.
First, when using a total of N noise scales, each perturbation kernel p σ i p x | x 0 q of SMLD can be derived from the following Markov chain:
[Equation]
where z i ´ 1 „ N p 0 , I q , x 0 „ p data , and we have introduced σ 0 “ 0 to simplify the notation.
In the limit of N Ñ 8 , the Markov chain t x i u Ni “ 1 becomes a continuous stochastic process t x p t qu 1 t “ 0 , t σ i u Ni “ 1 becomes a function σ p t q , and z i becomes z p t q , where we have used a continuous time variable t P r 0 , 1 s for indexing, rather than an integer i P t 1 , 2 , ¨ ¨ ¨ , N u .
Let ∆ t “ 1 N ´ 1 , x ´ i ´ 1 N ´¯
[Equation]
, and z
[Equation]
z i .
Eq.
(19) becomes the following at
[Equation]
where the approximate equality holds when ∆ t ! 1 .
In the limit of ∆ t Ñ 0 , this converges to
[Equation]
which is the VE SDE.
Note that x p 0 q “ x 1 “ x 0 ` σ 1 z 0 “ x 0 ` σ 2 p
q
z 0 and x 0 „ p data , so we have p
p x p 0 qq “ ş p data p x q N p x p 0 q ; x , σ 21 I q d x .
In other words, p 0 p x p 0 qq for the VE SDE is a perturbed data distribution with a small Gaussian noise.
For the perturbation kernels t p α i p x | x 0 qu Ni “ 1 used in DDPM, the discrete Markov chain is
[Equation]
where z i ´ 1 „ N p 0 , I q .
To obtain the limit of this Markov chain when N Ñ 8 , we deﬁne an auxiliary set of noise scales t ¯ β i “ p N ´ 1 q β i u Ni “ 1 , and re-write Eq.
(21) as below
[Equation]
In the limit of N Ñ 8 , t ¯ β i u Ni “ 1 becomes a function β p t q indexed by t P r 0 , 1 s .
Let β
´ i ´ 1 N ´¯ “ ¯ β i , x p
i ´
1 N ´ 1 q “ x i , z p i ´
1 N ´ 1 q
[Equation]
z i and ∆ t “ 1 N ´ 1 .
The Markov chain Eq.
(22) becomes the following at t P
[Equation]
0 , 1 , ¨ ¨ ¨ , N ´ 2 N ´
[Equation]
where the approximate equality holds when ∆ t ! 1 .
Therefore, in the limit of ∆ t Ñ 0 , Eq.
(23) converges to the following VP SDE:
[Equation] [Equation]
Because x p 0 q “ lim N Ñ8 ?
1 ´ β 1 x 0 ` ?
β 1 z 0 “ lim N Ñ b 1 ´ β p
q N ´ 1 x 0 ` b β p 0 q N ´
1 z 0 “ x 0 and x 0 „ p data , we have p
p x p 0 qq “ p data p x q .
So far, we have demonstrated that the noise perturbations used in SMLD and DDPM correspond to discretizations of VE and VP SDEs respectively.
The VE SDE always yields a process with exploding variance when t Ñ 8 .
In contrast, the VP SDE yields a process with bounded variance.
In addition, the process has a constant unit variance for all t P r 0 , 8q when p p x p 0 qq has a unit variance.
Since the
VP SDE has afﬁne drift and diffusion coefﬁcients, we can use Eq.
(5.51) in S¨arkk¨a & Solin (2019) to obtain an ODE that governs the evolution of variance
[Equation]
where Σ p t q “ Cov r x p t qs .
Solving this ODE, we obtain
[Equation]
from which it is clear that the variance Σ p t q is always bounded given Σ p 0 q .
Moreover, Σ p t q ” I if Σ p 0 q “ I .
Due to this difference, we name Eq.
(9) as the Variance Exploding (VE) SDE , and Eq.
(11) the Variance Preserving (VP) SDE .
Since both VE and VP SDEs have afﬁne drift coefﬁcients, their perturbation kernels p 0 t p x p t q | x p 0 qq are both Gaussian and can be computed with Eqs.
(5.50) and (5.51) in S¨arkk¨a & Solin (2019):
[Equation]
Therefore, both SDEs allow efﬁcient training with the objective in Eq.
C SDE S IN THE WILD
Below we discuss SDEs that directly correspond to implementations for practical SMLD and DDPM models.
In SMLD, the noise scales t σ i u Ni “ 1 is typically a geometric sequence where σ min is ﬁxed to 0 .
01 and σ max is chosen according to Technique 1 in Song & Ermon (2020).
Usually, SMLD models normalize image inputs to the range r 0 , 1 s .
Since t σ i u Ni “ 1 is a geometric sequence, we have σ i “ σ min
[Equation]
and therefore σ p t q “ σ
[Equation]
¯ t .
The VE SDE in this case is
[Equation]
[Equation]
where x p 0 q „ ş p data p x q
[Equation]
p 0 q ;
[Equation] [Equation]
d x .
The perturbation kernel can be computed via Eq.
[Equation]
For DDPM models, t β i u Ni “ 1 is typically an arithmetic sequence where β
i “ ¯ β min N ´ 1 `
[Equation]
¯ β min q .
Therefore, β p t q “ ¯ β min ` t p ¯ β max ´ ¯ β min q .
This corresponds to the following instantiation of the VP SDE:
[Equation] [Equation]
where x p 0 q „ p data p x q .
In our experiments, we let ¯ β min “ 0 . 1 and ¯ β max “ 20 , which correspond to the settings in Ho et al.
The perturbation kernel is given by
[Equation]
As a sanity check for our SDE generalizations to SMLD and DDPM, we compare the perturbation kernels of SDEs and original discrete Markov chains in Fig.
The SMLD and DDPM models both use N “ 1000 noise scales.
For SMLD, we only need to compare the variances of perturbation kernels.
For DDPM, we compare the scaling factors of means and the variances.
As demonstrated in Fig.
6, the discrete perturbation kernels of original SMLD and DDPM models align well with perturbation kernels derived from VE and VP SDEs.
D P ROBABILITY FLOW ODE D. D ERIVATION
The idea of probability ﬂow ODE is inspired by Maoutsa et al.
(2020), and one can ﬁnd the derivation of a simpliﬁed case therein.
Below we provide a derivation for the fully general ODE in Eq.
We consider the SDE in Eq.
(14), which possesses the following form:
[Equation]
[Figure]
where f p¨ , t q : R d Ñ R d and G p¨ , t q : R d Ñ R d ˆ d .
The marginal probability density p t p x p t qq evolves according to the forward Kolmogorov equation (Fokker-Planck equation) (Øksendal, 2003)
[Equation]
ı .
We can easily rewrite Eq.
(30) to obtain
B p t p x q
[Equation]
t
[Equation]
Note that d
[Equation]
based on which we can continue the rewriting of Eq.
(31) to obtain
[Equation]
“ ´
[Equation] [Equation]
“ ´
[Equation] [Equation]
ı
[Equation]
where we deﬁne
[Equation]
Inspecting Eq.
(32), we observe that it equates the forward Kolmogorov equation of the following
SDE with ˜ G p x , t q ﬁ 0 (the forward Kolmogorov equation holds in this case, as can be justiﬁed by the instantaneous change of variables formula in Chen et al.
[Equation]
which is essentially an ODE:
[Equation]
same as the probability ﬂow ODE given by Eq.
Therefore, we have shown that the probability ﬂow ODE Eq.
(16) induces the same marginal probability density p t p x q as the SDE in Eq.
D. L IKELIHOOD COMPUTATION
The probability ﬂow ODE in Eq.
(16) has the following form when we replace the score ∇ x log p t p x q with the time-dependent score-based model s θ p x , t q :
[Equation]
With the instantaneous change of variables formula (Chen et al., 2018), we can compute the loglikelihood of p
p x q using
[Equation]
(34) where the random variable x p t q as a function of t can be obtained by solving the probability ﬂow
ODE in Eq.
In many cases computing ∇ ¨ ˜ f θ p x , t q is expensive, so we follow Grathwohl et al.
(2018) to estimate it with the Skilling-Hutchinson trace estimator (Skilling, 1989; Hutchinson, 1990).
In particular, we have
[Equation]
(35) where ∇ ˜ f θ denotes the Jacobian of ˜ f θ p¨ , t q , and the random variable (cid:15) satisﬁes E p p (cid:15) q r (cid:15) s “ 0 and
Cov p p (cid:15) q r (cid:15) s “ I .
The vector-Jacobian product (cid:15) T ∇ ˜ f
θ p x , t q can be efﬁciently computed using reversemode automatic differentiation, at approximately the same cost as evaluating ˜ f
θ p x , t q .
As a result, we can sample (cid:15) „ p p (cid:15) q and then compute an efﬁcient unbiased estimate to ∇ ¨ ˜ f θ p x , t q using (cid:15) T ∇ ˜ f
θ p x , t q (cid:15) .
Since this estimator is unbiased, we can attain an arbitrarily small error by averaging over a sufﬁcient number of runs.
Therefore, by applying the Skilling-Hutchinson estimator Eq.
(35) to Eq.
(34), we can compute the log-likelihood to any accuracy.
D. P ROBABILITY FLOW SAMPLING
Suppose we have a forward SDE
[Equation]
, and one of its discretization
[Equation] [Equation]
where z i „ N p 0 , I q .
We assume the discretization schedule of time is ﬁxed beforehand, and thus we absorb the dependency on ∆ t into the notations of f i and G i .
Using Eq.
(16), we can obtain the following probability ﬂow ODE:
[Equation]
We may employ any numerical method to integrate the probability ﬂow ODE backwards in time for sample generation.
In particular, we propose a discretization in a similar functional form to Eq.
[Equation]
where the score-based model s θ ˚ p x i , i q is conditioned on the iteration number i .
This is a deterministic iteration rule.
Unlike reverse diffusion samplers or ancestral sampling, there is no additional randomness once the initial sample x N is obtained from the prior distribution.
When applied to
SMLD models, we can get the following iteration rule for probability ﬂow sampling:
[Equation]
Similarly, for DDPM models, we have
[Equation]
D. P ROBABILITY FLOW SAMPLING WITH BLACK - BOX ODE SOLVERS
For producing ﬁgures in Fig.
4, we use a DDPM model trained on 256 ˆ 256 CelebA-HQ with the same settings in Ho et al.
We use the RK45 ODE solver (Dormand & Prince, 1980) provided by scipy.integrate.solve_ivp in all cases.
The bits/dim values in Tab. 2 are computed with atol=1e-5 and rtol=1e-5 , same as Grathwohl et al.
To give the results for DDPM (probability ﬂow) and DDPM cont.
(probability ﬂow) in Tab. 2, we average the bits/dim obtained on the test dataset over ﬁve different runs.
The standard deviation is below 0 .
005 for both of them.
FID scores are computed on samples from the probability ﬂow sampler.
Aside from the interpolation results in Fig.
4, we demonstrate more examples of latent space manipulation in Fig.
7, including interpolation and temperature scaling.
The model tested here is a
DDPM model trained with the same settings in Ho et al.
D. U NIQUELY IDENTIFIABLE ENCODING
As a sanity check, we train two models (denoted as “Model A” and “Model B”) with different architectures using the VE SDE on CIFAR-10.
Here Model A is an NCSN++ model with 4 layers per resolution trained using the continuous objective in Eq.
(7), and Model B is all the same except that it uses 8 layers per resolution.
Model deﬁnitions are in Appendix H.
We report the latent codes obtained by Model A and Model B for a random CIFAR-10 image in
Fig. 8.
In Fig.
9, we show the dimension-wise differences and correlation coefﬁcients between latent encodings on a total of 16 CIFAR-10 images.
Our results demonstrate that for the same inputs, Model
A and Model B provide encodings that are close in every dimension, despite having different model architectures and training runs.
E R EVERSE DIFFUSION SAMPLING
Given a forward SDE d x “ f p x , t q d t ` G p t q d w , and suppose the following iteration rule is a discretization of it:
[Equation] [Equation]
where z i „ N p 0 , I q .
Here we assume the discretization schedule of time is ﬁxed beforehand, and thus we can absorb it into the notations of f i and G i .
Based on Eq.
(40), we propose to discretize the reverse-time SDE
[Equation]
with a similar functional form, which gives the following iteration rule for i P t 0 ,
[Equation] [Equation]
[Figure] [Figure]
where our trained score-based model s θ ˚ p x i , i q is conditioned on iteration number i .
When applying Eq.
(41) to Eqs.
(10) and (19), we obtain a new set of numerical solvers for the reverse-time VE and VP SDEs, resulting in sampling algorithms as shown in the “predictor” part of Algs.
1 and 2.
We name these sampling methods (that are based on the discretization strategy in
Eq.
(41)) reverse diffusion samplers .
Note that the ancestral sampling of DDPM (Ho et al., 2020)
(Eq.
(4)) matches its reverse diffusion counterpart when β i Ñ
for all i (which happens when ∆ t Ñ 0 since β i “ ¯ β i ∆ t ), because
[Equation]
[Table] [Equation] [Table] [Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
Therefore, the original ancestral sampler of Eq.
(4) is essentially a different discretization to the same reverse-time SDE.
This uniﬁes the sampling method in Ho et al.
(2020) as a numerical solver to the reverse-time VP SDE in our continuous framework.
F A NCESTRAL SAMPLING FOR SMLD MODELS
The ancestral sampling method for DDPM models can also be adapted to SMLD models.
Consider a sequence of noise scales σ 1 ă σ 2 ă ¨ ¨ ¨ ă σ N as in SMLD.
By perturbing a data point x 0 with these noise scales sequentially, we obtain a Markov chain x
Ñ x 1 Ñ ¨ ¨ ¨ Ñ x N , where
[Equation]
Here we assume σ 0 “ 0 to simplify notations.
Following Ho et al.
[Equation]
we can compute
[Equation]
If we parameterize the reverse transition kernel as p θ p x i ´
| x i q “ N p x
i ´ 1 ; µ θ p x i , i q , τ 2 i I q
[Equation]
then
[Equation]
where L t ´ 1 is one representative term in the ELBO objective (see Eq.
(8) in Ho et al.
, C is a constant that does not depend on θ , z „ N p 0 , I q , and x i p x 0 , z q “ x 0 ` σ i z .
We can therefore parameterize µ θ p x
i , i q via
[Equation]
i
[Equation]
where s θ p x i , i q is to estimate z { σ i .
As in Ho et al.
, we let τ i “ c σ
i ´ 1 p σ 2 i ´
σ 2 i ´ 1 q σi .
Through ancestral sampling on ś Ni “ 1 p θ p x i ´ 1 | x i q , we obtain the following iteration rule
[Equation]
where x N „ N p 0 , σ 2 N I q , θ ˚ denotes the optimal parameter of s θ , and z i „ N p 0 , I q .
We call
Eq.
(42) the ancestral sampling method for SMLD models.
G A DDITIONAL DETAILS ON P REDICTOR -C ORRECTOR SAMPLERS
Training We use the same architecture in Ho et al.
(2020) for our score-based models.
For the
VE SDE, we train a model with the original SMLD objective in Eq.
(1); similarly for the VP SDE, we use the original DDPM objective in Eq.
We apply a total number of 1000 noise scales for training both models.
For results in Fig.
3, we train an NCSN++ model (deﬁnition in Appendix H) on 256 ˆ 256 LSUN bedroom and church outdoor (Yu et al., 2015) datasets with the VE SDE and our continuous objective Eq.
The corrector algorithms We take the schedule of annealed Langevin dynamics in Song & Ermon (2019), but re-frame it with slight modiﬁcations in order to get better interpretability and empirical performance.
We provide the corrector algorithms in Algs.
3 and 4 respectively, where we call r the “signal-to-noise” ratio.
We determine the step size (cid:15) using the norm of the Gaussian noise (cid:107) z (cid:107) 2 , norm of the score-based model (cid:107) s θ ˚ (cid:107) 2 and the signal-to-noise ratio r .
When sampling a large batch of samples together, we replace the norm (cid:107) ¨ (cid:107) 2 with the average norm across the mini-batch.
When the batch size is small, we suggest replacing (cid:107) z (cid:107) 2 with ?
d , where d represents the dimensionality of z .
Require: t
σ i u Ni “ 1 , r, N, M .
1: x 0 N „ N p 0 , σ 2 max I q 2: for i Ð N to 1 do 3: for j Ð 1 to M do 4:
z „ N p 0 , I q 5: g Ð s θ ˚ p x j ´ 1 i , σ i q 6: (cid:15)
Ð 2 p r (cid:107) z (cid:107) 2 { (cid:107) g
(cid:107) 2 q7: x ji Ð x j ´
i ` (cid:15) g ` ? 2 (cid:15) z 8: x 0 i ´ 1 Ð x Mi return x
Require:
t β i u Ni “ 1 , t α i u Ni “ 1 , r, N, M .
1: x 0 N „ N p 0 , I q 2: for i Ð N to 1 do 3: for j Ð 1 to M do 4:
z „ N p 0 , I q 5: g Ð s θ ˚ p x j ´ 1 i , i q 6: (cid:15)
Ð
α i p r (cid:107) z (cid:107) 2 { (cid:107) g (cid:107) 2 q7: x ji Ð x j ´
i ` (cid:15) g ` ? 2 (cid:15) z 8: x 0 i ´ 1 Ð x Mi return x
Denoising For both SMLD and DDPM models, the generated samples typically contain small noise that is hard to detect by humans.
As noted by Jolicoeur-Martineau et al.
(2020), FIDs can be signiﬁcantly worse without removing this noise.
This unfortunate sensitivity to noise is also part of the reason why NCSN models trained with SMLD has been performing worse than DDPM models in terms of FID, because the former does not use a denoising step at the end of sampling, while the latter does.
In all experiments of this paper we ensure there is a single denoising step at the end of sampling, using Tweedie’s formula (Efron, 2011).
Ad-hoc interpolation methods for noise scales Models in this experiment are all trained with 1000 noise scales.
To get results for P2000 (predictor-only sampler using 2000 steps) which requires 2000 noise scales, we need to interpolate between 1000 noise scales at test time.
The speciﬁc architecture of the noise-conditional score-based model in Ho et al.
(2020) uses sinusoidal positional embeddings for conditioning on integer time steps.
This allows us to interpolate between noise scales at test time in an ad-hoc way (while it is hard to do so for other architectures like the one in Song &
Ermon (2019)).
Speciﬁcally, for SMLD models, we keep σ min and σ max ﬁxed and double the number of time steps.
For DDPM models, we halve ¯ β min and ¯ β max before doubling the number of time steps.
Suppose t s θ p x , i qu N ´ 1 i “ 0 is a score-based model trained on N time steps, and let t s 1 θ p x , i qu 2 N ´
i “denote the corresponding interpolated score-based model at 2 N time steps.
We test two different interpolation strategies for time steps: linear interpolation where s 1 θ p x , i q “ s θ p x , i { 2 q and rounding interpolation where s 1 θ p x , i q “ s θ p x , t i { 2 u q .
We provide results with rounding interpolation in Tab. 1, and give results of linear interpolation in Tab. 4.
We observe that different interpolation methods result in performance differences but maintain the general trend of predictor-corrector methods performing on par or better than predictor-only or corrector-only samplers.
Hyper-parameters of the samplers For Predictor-Corrector and corrector-only samplers on
CIFAR-10, we search for the best signal-to-noise ratio ( r ) over a grid that increments at 0.01.
We report the best r in Tab. 5.
For LSUN bedroom/church outdoor, we ﬁx r to 0.075.
Unless otherwise noted, we use one corrector step per noise scale for all PC samplers.
We use two corrector steps per noise scale for corrector-only samplers on CIFAR-10.
We always use a batch size of 128 for training.
For sample generation, the batch size is 1024 on CIFAR-10 and 8 on LSUN bedroom/church outdoor.
H A RCHITECTURAL IMPROVEMENTS FOR MODELS WITH VE PERTURBATIONS
To study whether the VE perturbations are inherently weaker than VP perturbations, we conducted a comprehensive architecture search to compare their best performance.
Our results demonstrate that the VE perturbation can be advantageous in certain cases.
Further, our endeavor gives rise to new state-of-the-art models on CIFAR-10, and enables the ﬁrst high-ﬁdelity image samples of resolution 1024 ˆ 1024 from score-based generative models.
Code and checkpoints will be open sourced.
[Table] [Table]
H. S ETTINGS
We consider two datasets: 32 ˆ 32 CIFAR-10 (Krizhevsky et al., 2009) and 64 ˆ 64 CelebA
(Liu et al., 2015), which is pre-processed following Song & Ermon (2020).
All models are trained with 1.3M iterations, and we save one checkpoint per 50k iterations.
We compare different conﬁgurations based on their FID scores averaged over checkpoints after 0.5M iterations.
The FIDs are computed on 50k samples with TF-GAN.
For sampling, we use the PC sampler discretized at 1000 noise scales.
We follow Ho et al.
(2020) for optimization, including the learning rate, gradient clipping, and learning rate warm-up schedules.
Unless otherwise noted, all models are trained with the original SMLD objective Eq.
(1) and use a batch size of 128.
[Figure]
Our architecture is mostly based on Ho et al.
We additionally search over the following components to explore the potential of score-based models trained with VE perturbations.
• Upsampling and downsampling images with anti-aliasing based on Finite Impulse Response (FIR) (Zhang, 2019).
We follow the same implementation and hyper-parameters in StyleGAN-2 (Karras et al., 2020b).
• Rescaling all skip connections by 1 { ? 2 .
This has been demonstrated effective in several works, including ProgressiveGAN (Karras et al., 2018), StyleGAN (Karras et al., 2019) and StyleGAN-2 (Karras et al., 2020b).
• Replacing the original residual blocks in DDPM with residual blocks from BigGAN (Brock et al., 2018).
• Increasing the number of residual blocks per resolution from 2 to 4 .
Incorporating progressive growing architectures.
We consider two progressive architectures for input: “input skip” and “residual”, and two progressive architectures for output: “output skip” and “residual”.
These progressive architectures are deﬁned and implemented according to StyleGAN-2.
We also tested equalized learning rates, a trick used in very successful models like ProgressiveGAN (Karras et al., 2018) and StyleGAN (Karras et al., 2019).
However, we found it harmful at an early stage of our experiments, and therefore decided not to include it in the architecture search.
The exponential moving average (EMA) rate has a signiﬁcant impact on performance.
For models trained with VE perturbations, we notice that 0.999 works better than 0.9999, whereas for models trained with VP perturbations it is the opposite.
We therefore use an EMA rate of 0.999 and 0.
for VE and VP models respectively.
For sampling from models with VE perturbations, we use the PC sampler discretized at 1000 noise scales with one Langevin step per scale and a signal-of-noise ratio of 0.16.
For VP perturbations, we use the reverse diffusion sampler discretized at
H. R ESULTS ON CIFAR-
We provide the architecture search results in Fig.
The box plots demonstrate the importance of each component when other components can vary freely.
On both CIFAR-10 and CelebA, the additional components that we searched over can always improve the performance on average.
For progressive growing, it is not clear which combination of conﬁgurations consistently performs the best, but the results are typically better than when no progressive growing architecture is used.
Our best score-based model uses FIR upsampling/downsampling, rescales skip connections, employs
BigGAN-type residual blocks, uses 4 residual blocks per resolution instead of 2, and uses “residual” for input and no progressive growing architecture for output.
We name this model “NCSN++”, following the naming convention of previous SMLD models (Song & Ermon, 2019; 2020).
The basic NCSN++ model with 4 residual blocks achieves an FID of 2.45 on CIFAR-10.
Here in order to match the convention used in Karras et al.
(2018); Song & Ermon (2019);
Ho et al.
, the FID value reported here is the lowest over the course of training, rather than the average over checkpoints after 0.5M iterations (as done in our architecture search).
We conducted the same architecture search for a score-based model trained with the DDPM objective Eq.
The optimal model obtains an FID of 2.88 on CIFAR-10, signiﬁcantly worse than the model trained with VE perturbations.
Interestingly, the optimal score-based model for VP perturbations has signiﬁcantly different conﬁgurations than those trained with VE perturbation.
Some of the components like rescaling skip connections and progressive growing architectures seem to be harmful for them.
This further conﬁrms that the relative advantage of VE/VP SDEs is strongly correlated with model architecture, and should be jointly tuned as part of the design choices.
To further improve the NCSN++ model upon conditioning on continuous time variables, we change positional embeddings, the layers in Ho et al.
(2020) for conditioning on discrete time steps, to random Fourier feature embeddings, as advocated in Tancik et al.
The scale parameter of these random Fourier feature embeddings is set to 16.
We additionally train the model with our continuous objective Eq.
These changes improve the FID on CIFAR-10 from 2.45 to 2.38, resulting in a model dubbed NCSN++ cont.
(4 blocks/resolution).
Finally, we can further improve the FID from 2.38 to 2.2 by doubling the number of residual blocks per resolution, resulting in the model denoted as NCSN++ cont.
(8 blocks/resolution).
All results are summarized in Tab. 3.
H. H IGH RESOLUTION IMAGES
Encouraged by the success of NCSN++ on CIFAR-10, we proceed to test it on 1024 ˆ 1024 CelebAHQ
(Karras et al., 2018), a task that was previously only achievable by some GAN models and
VQ-VAE-2 (Razavi et al., 2019).
We used a batch size of 8, increased the EMA rate to 0.9999, and trained the NCSN++ model with the continuous objective (Eq.
for around 700k iterations.
We use the PC sampler discretized at 2000 steps with one Langevin step per noise scale and a signal-to-noise ratio of 0.15.
The scale parameter for the random Fourier feature embeddings is ﬁxed to 4.
We use the “residual” progressive architecture for the input, and “output skip” progressive architecture for the output.
We provide samples in Fig.
Although these samples are not perfect ( e.g ., there are visible ﬂaws on facial symmetry), we believe these results are encouraging and can demonstrate the scalability of our approach.
Future work on more effective architectures are likely to signiﬁcantly advance the performance of score-based generative models on this task.
I C ONTROLLABLE GENERATION
Consider a forward SDE with the following general form
[Equation]
and suppose the initial state distribution is p
p x p 0 q | y q .
The density at time t is p t p x
p t q | y q when conditioned on y .
Therefore, using Anderson (1982), the reverse-time SDE is given by
[Equation]
d
[Equation] [Equation]
Since p t p x p t q | y q9 p t p x p t qq p p y | x p t qq , the score ∇ x log p t p x p t q | y q can be computed easily by
[Equation]
This subsumes the conditional reverse-time SDE in Eq.
(13) as a special case.
All sampling methods we have discussed so far can be applied to the conditional reverse-time SDE for sample generation.
I. C LASS - CONDITIONAL SAMPLING
When y represents class labels, we can train a time-dependent classiﬁer p t p y | x p t qq for classconditional sampling.
Since the forward SDE is tractable, we can easily create a pair of training data p x p t q , y q by ﬁrst sampling p x p 0 q , y q from a dataset and then obtaining x p t q „ p 0 t p x p t q | x p 0 qq .
Afterwards, we may employ a mixture of cross-entropy losses over different time steps, like Eq.
(7), to train the time-dependent classiﬁer
p t
p
y
| x p t qq .
To test this idea, we trained a Wide ResNet (Zagoruyko & Komodakis, 2016)
( Wide-ResNet-28-10 ) on CIFAR-10 with VE perturbations.
The classiﬁer is conditioned on log σ i using random Fourier features (Tancik et al., 2020), and the training objective is a simple sum of cross-entropy losses sampled at different scales.
We provide a plot to show the accuracy of this classiﬁer over noise scales in Fig.
The score-based model is an unconditional
NCSN++ (4 blocks/resolution) in Tab. 3, and we generate samples using the PC algorithm withdiscretization steps.
The class-conditional samples are provided in Fig.
5, and an extended set of conditional samples is given in Fig.
I. I MPUTATION
Imputation is a special case of conditional sampling.
Denote by Ω p x q and ¯Ω p x q the known and unknown dimensions of x respectively, and let f ¯Ω p¨ , t q and G ¯Ω p¨ , t q denote f p¨ , t q and G p¨ , t q restricted to the unknown dimensions.
For VE/VP SDEs, the drift coefﬁcient f p¨ , t q is element-wise, and the diffusion coefﬁcient G p¨ , t q is diagonal.
When f p¨ , t q is element-wise, f ¯Ω p¨ , t q denotes the same element-wise function applied only to the unknown dimensions.
When G p¨ , t q is diagonal, G ¯Ω p¨ , t q denotes the sub-matrix restricted to unknown dimensions.
[Figure]
For imputation, our goal is to sample from p p ¯Ω p x p 0 qq | Ω p x p 0 qq “ y q .
Deﬁne a new diffusion process z p t q “ ¯Ω p x p
t
qq , and note that the SDE for z p t q can be written as
[Equation]
The reverse-time SDE, conditioned on Ω p x p 0 qq “ y , is given by
[Equation]
Although p t p z p t q | Ω p x p 0 qq
“ y q is in general intractable, it can be approximated.
Let A ﬁt Ω p x p 0 qq
“ y u , where 1 t¨u denotes the indicator function.
We have
[Equation]
where ˆΩ p x p t qq is a random sample from p t
p Ω p x p t qq | A q , which is typically a tractable distribution.
Therefore, ∇
[Equation]
“ ∇ z log p t pr z p t q ; ˆΩ p x p t qqsq , where r z p t q ; ˆΩ p x p t qqs denotes a vector u p t q such that Ω p u p t qq “ ˆΩ p x
p t qq and ¯Ω p u p t qq “ z p t q , and the identity holds because ∇ z log p t pr z p t q ; ˆΩ p x p t qqsq “ ∇ z log p t p
z p t q | ˆΩ p x p t qqq ` ∇ z log p p
ˆΩ p x p t qqq “ ∇ z log p t
p
z
p
t
q | ˆΩ p x p t qqq .
We provided an extended set of inpainting results in Figs.
13 and 14.
I. C OLORIZATION
Colorization is a special case of imputation, except that the known data dimensions are coupled.
We can decouple these data dimensions by using an orthogonal linear transformation to map the gray-scale image to a separate channel in a different space, and then perform imputation to complete the other channels before transforming everything back to the original image space.
The orthogonal matrix we used to decouple color channels is ˜ 0 .
577 ´ 0 . 0 .0 .0 .
´
.¸ .
Because the transformations are all orthogonal matrices, the standard Wiener process w p t q will still be a standard Wiener process in the transformed space, allowing us to build an SDE and use the same imputation method in Appendix I.2.
We provide an extended set of colorization results in Figs.
15 and 16.