Score-Based Generative Modeling through Stochastic Differential Equations（附录_英文原文_2: E-I）

E REVERSE DIFFUSION SAMPLING

Given a forward SDE

$$\mathrm{dx}=\mathbf{f}(\mathbf{x},t)\mathrm{d}t+\mathbf{G}(t)\mathrm{d}\mathbf{w},$$

and suppose the following iteration rule is a discretization of it:

$${\mathbf{x}}_{i+1}={\mathbf{x}}_i+{\mathbf{f}}_i({\mathbf{x}}_i)+{\mathbf{G}}_i{\mathbf{z}}_i,i=0,1,\cdots ,N-1,(45)$$

where ${\mathbf{z}}_i\sim \mathcal{N}(0,\mathbf{I})$. Here we assume the discretization schedule of time is fixed beforehand, and thus we can absorb it into the notations of ${\mathbf{f}}_i$ and ${\mathbf{G}}_i$.

Based on Eq. (45), we propose to discretize the reverse-time SDE

$$\mathrm{dx}=\left[\mathbf{f}(\mathbf{x},t)-\mathbf{G}(t)\mathbf{G}(t{)}^{\top }{\nabla }_{\mathbf{x}}\log p_t(\mathbf{x})\right]\mathrm{d}t+\mathbf{G}(t)\mathrm{d}\mathbf{w},$$

with a similar functional form, which gives the following iteration rule for i∈{0,1,⋯ ,N−1}i \in \{0,1, \cdots, N-1\}i∈{0,1,⋯,N−1}:

$${\mathbf{x}}_i={\mathbf{x}}_{i+1}-{\mathbf{f}}_{i+1}({\mathbf{x}}_{i+1})+{\mathbf{G}}_{i+1}{\mathbf{G}}_{i+1}^{\top }{\mathbf{s}}_{{\theta }^{\ast }}({\mathbf{x}}_{i+1},i+1)+{\mathbf{G}}_{i+1}{\mathbf{z}}_{i+1},(46)$$

where our trained score-based model ${\mathbf{s}}_{{\theta }^{\ast }}({\mathbf{x}}_i,i)$ is conditioned on iteration number $i$.

When applying Eq. (46) to Eqs. (10) and (20), 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 Algorithms 2 and 3. We name these sampling methods (that are based on the discretization strategy in Eq. (46)) reverse diffusion samplers.

As expected, the ancestral sampling of DDPM (Ho et al., 2020) (Eq. (4)) matches its reverse diffusion counterpart when ${\beta }_i\to 0$ for all $i$ (which happens when $\Delta t\to 0$ since ${\beta }_i={\tilde{\beta }}_i\Delta t$, see Appendix B),

because

$x_i=\frac{1}{\sqrt{1-{\beta }_{i+1}}}(x_{i+1}+{\beta }_{i+1}{S}_{\theta \ast }(x_{i+1},i+1))+\sqrt{{\beta }_{i+1}}{z}_{i+1}$
$=\left(1+\frac{1}{2}{\beta }_{i+1}+o({\beta }_{i+1})\right)(x_{i+1}+{\beta }_{i+1}{S}_{\theta \ast }(x_{i+1},i+1))+\sqrt{{\beta }_{i+1}}{z}_{i+1}$
$\approx \left(1+\frac{1}{2}{\beta }_{i+1}\right)(x_{i+1}+{\beta }_{i+1}{S}_{\theta \ast }(x_{i+1},i+1))+\sqrt{{\beta }_{i+1}}{z}_{i+1}$
$=\left(1+\frac{1}{2}{\beta }_{i+1}\right)x_{i+1}+{\beta }_{i+1}{S}_{\theta \ast }(x_{i+1},i+1)+\frac{1}{2}{\beta }_{i+1}^2{S}_{\theta \ast }(x_{i+1},i+1)+\sqrt{{\beta }_{i+1}}{z}_{i+1}$
$\approx \left(1+\frac{1}{2}{\beta }_{i+1}\right)x_{i+1}+{\beta }_{i+1}{S}_{\theta \ast }(x_{i+1},i+1)+\sqrt{{\beta }_{i+1}}{z}_{i+1}$
$=\left[2-\left(1-\frac{1}{2}{\beta }_{i+1}\right)\right]x_{i+1}+{\beta }_{i+1}{S}_{\theta \ast }(x_{i+1},i+1)+\sqrt{{\beta }_{i+1}}{z}_{i+1}$
$\approx \left[2-\left(1-\frac{1}{2}{\beta }_{i+1}\right)+o({\beta }_{i+1})\right]x_{i+1}+{\beta }_{i+1}{S}_{\theta \ast }(x_{i+1},i+1)+\sqrt{{\beta }_{i+1}}{z}_{i+1}$
$=(2-\sqrt{1-{\beta }_{i+1}})x_{i+1}+{\beta }_{i+1}{S}_{\theta \ast }(x_{i+1},i+1)+\sqrt{{\beta }_{i+1}}{z}_{i+1}$.

Therefore, the original ancestral sampler of Eq. (4) is essentially a different discretization to the same reverse-time SDE. This unifies the sampling method in Ho et al. (2020) as a numerical solver to the reverse-time VP SDE in our continuous framework.

Therefore, the original ancestral sampler of Eq. (4) is essentially a different discretization to the same reverse-time SDE. This unifies the sampling method in Ho et al. (2020) as a numerical solver to the reverse-time VP SDE in our continuous framework.

F ANCESTRAL SAMPLING FOR SMLD MODELS

The ancestral sampling method for DDPM models can also be adapted to SMLD models. Consider a sequence of noise scales ${\sigma }_1<{\sigma }_2<\cdots <{\sigma }_N$ as in SMLD. By perturbing a data point ${\mathbf{x}}_0$ with these noise scales sequentially, we obtain a Markov chain ${\mathbf{x}}_0\to {\mathbf{x}}_1\to \cdots \to {\mathbf{x}}_N$, where

$$p({\mathbf{x}}_i\mid {\mathbf{x}}_{i-1})=\mathcal{N}\left({\mathbf{x}}_i;{\mathbf{x}}_{i-1},\left({\sigma }_i^2-{\sigma }_{i-1}^2\right)\mathbf{I}\right),i=1,2,\cdots ,N.$$

Here we assume ${\sigma }_0=0$ to simplify notations. Following Ho et al. (2020), we can compute

$$q({\mathbf{x}}_{i-1}\mid {\mathbf{x}}_i,{\mathbf{x}}_0)=\mathcal{N}\left({\mathbf{x}}_{i-1};\frac{{\sigma }_{i-1}^2}{{\sigma }_i^2}{\mathbf{x}}_i+\left(1-\frac{{\sigma }_{i-1}^2}{{\sigma }_i^2}\right){\mathbf{x}}_0,\frac{{\sigma }_{i-1}^2\left({\sigma }_i^2-{\sigma }_{i-1}^2\right)}{{\sigma }_i^2}\mathbf{I}\right).$$

If we parameterize the reverse transition kernel as $p_{\theta }({\mathbf{x}}_{i-1}\mid {\mathbf{x}}_i)=\mathcal{N}\left({\mathbf{x}}_{i-1};{\mu }_{\theta }({\mathbf{x}}_i,i),{\tau }_i^2\mathbf{I}\right)$, then

$$\begin{array}{rl}{L}_{t-1}& ={\mathbb{E}}_q\left[D_{\text{KL}}\left(q({\mathbf{x}}_{i-1}\mid {\mathbf{x}}_i,{\mathbf{x}}_0)\Vert p_{\theta }({\mathbf{x}}_{i-1}\mid {\mathbf{x}}_i)\right)\right]\\ & ={\mathbb{E}}_q\left[\frac{1}{2{\tau }_i^2}{\Vert \frac{{\sigma }_{i-1}^2}{{\sigma }_i^2}{\mathbf{x}}_i+\left(1-\frac{{\sigma }_{i-1}^2}{{\sigma }_i^2}\right){\mathbf{x}}_0-{\mu }_{\theta }({\mathbf{x}}_i,i)\Vert }_2^2\right]+C\\ & ={\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{z}}\left[\frac{1}{2{\tau }_i^2}{\Vert {\mathbf{x}}_i({\mathbf{x}}_0,\mathbf{z})-\frac{{\sigma }_i^2-{\sigma }_{i-1}^2}{{\sigma }_i}\mathbf{z}-{\mu }_{\theta }({\mathbf{x}}_i({\mathbf{x}}_0,\mathbf{z}),i)\Vert }_2^2\right]+C,\end{array}$$

where $L_{t-1}$ is one representative term in the ELBO objective (see Eq. (8) in Ho et al. (2020)), $C$ is a constant that does not depend on $\mathbf{\theta }$, $\mathbf{z}\sim \mathcal{N}(0,\mathbf{I})$, and ${\mathbf{x}}_i({\mathbf{x}}_0,\mathbf{z})={\mathbf{x}}_0+{\sigma }_i\mathbf{z}$. We can therefore parameterize ${\mu }_{\theta }({\mathbf{x}}_i,i)$ via

$${\mu }_{\theta }({\mathbf{x}}_i,i)={\mathbf{x}}_i+\left({\sigma }_i^2-{\sigma }_{i-1}^2\right){\mathbf{s}}_{\theta }({\mathbf{x}}_i,i),$$

where ${\mathbf{s}}_{\theta }({\mathbf{x}}_i,i)$ is to estimate $\mathbf{z}/{\sigma }_i$. As in Ho et al. (2020), we let ${\tau }_i=\sqrt{\frac{{\sigma }_{i-1}^2\left({\sigma }_i^2-{\sigma }_{i-1}^2\right)}{{\sigma }_i^2}}$. Through ancestral sampling on ${\prod }_{i=1}^N{p}_{\theta }({\mathbf{x}}_{i-1}\mid {\mathbf{x}}_i)$, we obtain the following iteration rule

$${\mathbf{x}}_{i-1}={\mathbf{x}}_i+\left({\sigma }_i^2-{\sigma }_{i-1}^2\right){\mathbf{s}}_{{\theta }^{\ast }}({\mathbf{x}}_i,i)+\sqrt{\frac{{\sigma }_{i-1}^2\left({\sigma }_i^2-{\sigma }_{i-1}^2\right)}{{\sigma }_i^2}}{\mathbf{z}}_i,i=1,2,\cdots ,N,$$

where ${\mathbf{x}}_N\sim \mathcal{N}(0,{\sigma }_N^2\mathbf{I})$, ${\mathbf{\theta }}^{\ast }$ denotes the optimal parameter of ${\mathbf{s}}_{\theta }$, and ${\mathbf{z}}_i\sim \mathcal{N}(0,\mathbf{I})$. We call Eq. (47) the ancestral sampling method for SMLD models.

Algorithm 1 Predictor-Corrector (PC) sampling

Require:

• $N$: Number of discretization steps for the reverse-time SDE
• $M$: Number of corrector steps

Initialize ${\mathbf{x}}_N\sim p_T(\mathbf{x})$
for $i=N-1$ to 0 do
${\mathbf{x}}_i\leftarrow \text{Predictor}\left({\mathbf{x}}_{i+1}\right)$
for $j=1$ to $M$ do
${\mathbf{x}}_i\leftarrow \text{Corrector}\left({\mathbf{x}}_i\right)$
return ${\mathbf{x}}_0$

G PREDICTOR-CORRECTOR SAMPLERS

Predictor-Corrector (PC) sampling

The predictor can be any numerical solver for the reverse-time SDE with a fixed discretization strategy. The corrector can be any score-based MCMC approach. In PC sampling, we alternate between the predictor and corrector, as described in Algorithm 1. For example, when using the reverse diffusion SDE solver (Appendix E) as the predictor, and annealed Langevin dynamics (Song & Ermon, 2019) as the corrector, we have Algorithms 2 and 3 for VE and VP SDEs respectively, where {ϵi}i=0N−1\left\{\epsilon_i\right\}_{i=0}^{N-1}{ϵi​}i=0N−1​ are step sizes for Langevin dynamics as specified below.

The corrector algorithms

We take the schedule of annealed Langevin dynamics in Song & Ermon (2019), but re-frame it with slight modifications in order to get better interpretability and empirical performance. We provide the corrector algorithms in Algorithms 4 and 5 respectively, where we call $r$ the “signal-to-noise” ratio. We determine the step size $ϵ$ using the norm of the Gaussian noise $\Vert \mathbf{z}{\Vert }_2$, norm of the score-based model ${\Vert {\mathbf{s}}_{{\theta }^{\ast }}\Vert }_2$ and the signal-to-noise ratio $r$. When sampling a large batch of samples together, we replace the norm $\Vert \cdot {\Vert }_2$ with the average norm across the mini-batch. When the batch size is small, we suggest replacing $\Vert \mathbf{z}{\Vert }_2$ with $\sqrt{d}$, where $d$ is the dimensionality of $\mathbf{z}$.

Algorithm 4 Corrector algorithm (VE SDE)
Require: {σi}i=1N,r,N,M\left\{\sigma_i\right\}_{i=1}^N, r, N, M{σi​}i=1N​,r,N,M.
${\mathbf{x}}_0^0\sim \mathcal{N}\left(0,{\sigma }_{\max }^2\mathbf{I}\right)$
for $i\leftarrow N$ to 1 do
for $j\leftarrow 1$ to $M$ do
$\mathbf{z}\sim \mathcal{N}(0,\mathbf{I})$
$\mathbf{g}\leftarrow {\mathbf{s}}_{0\ast }\left({\mathbf{x}}_i^{j-1},{\sigma }_i\right)$
$ϵ\leftarrow 2{\left(r\Vert \mathbf{z}{\Vert }_2/\Vert \mathbf{g}{\Vert }_2\right)}^2$
${\mathbf{x}}_i^j\leftarrow {\mathbf{x}}_i^{j-1}+ϵ\mathbf{g}+\sqrt{2ϵ}\mathbf{z}$
${\mathbf{x}}_i^i\leftarrow {\mathbf{x}}_i^{i+1}$
return ${\mathbf{x}}_0^0$

Algorithm 5 Corrector algorithm (VP SDE)
Require: {βi}i=1N,{αi}i=1N,r,N,M\left\{\beta_i\right\}_{i=1}^N,\left\{\alpha_i\right\}_{i=1}^N, r, N, M{βi​}i=1N​,{αi​}i=1N​,r,N,M.
${\mathbf{x}}_i^0\sim \mathcal{N}(0,\mathbf{I})$
for $i\leftarrow N$ to 1 do
for $j\leftarrow 1$ to $M$ do
$\mathbf{z}\sim \mathcal{N}(0,\mathbf{I})$
$\mathbf{g}\leftarrow {\mathbf{s}}_{0\ast }\left({\mathbf{x}}_i^{j-1},i\right)$
$ϵ\leftarrow 2{\left(r\Vert \mathbf{z}{\Vert }_2/\Vert \mathbf{g}{\Vert }_2\right)}^2$
${\mathbf{x}}_i^j\leftarrow {\mathbf{x}}_i^{j-1}+ϵ\mathbf{g}+\sqrt{2ϵ}\mathbf{z}$
${\mathbf{x}}_i^i\leftarrow {\mathbf{x}}_i^{i+1}$
return ${\mathbf{x}}_0^0$

Denoising

For both SMLD and DDPM models, the generated samples typically contain small noise that is hard to detect by humans. As noted by Jolicocur-Martineau et al. (2020), FIDs can be significantly 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).

Figure 9: PC sampling for LSUN bedroom and church. The vertical axis corresponds to the total computation, and the horizontal axis represents the amount of computation allocated to the corrector. Samples are the best when computation is split between the predictor and corrector.

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. (3). We apply a total number of 1000 noise scales for training both models. For results in Fig. 9, we train an NCSN++ model (definition in Appendix H) on $256\times 256$ LSUN bedroom and church_outdoor (Yu et al., 2015) datasets with the VE SDE and our continuous objective Eq. (7). The batch size is fixed to 128 on CIFAR-10 and 64 on LSUN.

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 specific 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 hard to do so for other architectures like the one in Song & Ermon (2019)). Specifically, for SMLD models, we keep ${\sigma }_{\min }$ and ${\sigma }_{\max }$ fixed and double the number of time steps. For DDPM models, we halve ${\beta }_{\min }$ and ${\beta }_{\max }$ before doubling the number of time steps. Suppose {sθ(x,i)}i=0N−1\left\{s_\theta(\mathbf{x}, i)\right\}_{i=0}^{N-1}{sθ​(x,i)}i=0N−1​ is a score-based model trained on $N$ time steps, and let {sθ′(x,i)}i=02N−1\left\{s_\theta^{\prime}(\mathbf{x}, i)\right\}_{i=0}^{2 N-1}{sθ′​(x,i)}i=02N−1​ denote the corresponding interpolated score-based model at $2N$ time steps. We test two different interpolation strategies for time steps: linear interpolation where ${\mathrm{s}}_{\theta }^{\prime }(\mathrm{x},i)={\mathrm{s}}_{\theta }(\mathrm{x},i/2)$ and rounding interpolation where ${\mathrm{s}}_{\theta }^{\prime }(\mathrm{x},i)={\mathrm{s}}_{\theta }(\mathrm{x},\lfloor i/2\rfloor )$. We provide results with linear interpolation in Table 1, and give results of rounding interpolation in Table 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 Table 5. For LSUN bedroom/church_outdoor, we fix $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. For sample generation, the batch size is 1024 on CIFAR-10 and 8 on LSUN bedroom/church_outdoor.

Table 4: Comparing different samplers on CIFAR-10, where “P2000” uses the rounding interpolation between noise scales. Shaded regions are obtained with the same computation (number of score function evaluations). Mean and standard deviation are reported over five sampling runs.

Table 5: Optimal signal-to-noise ratios of different samplers. “P1000” or “P2000”: predictor-only samplers using 1000 or 2000 steps. “C2000”: corrector-only samplers using 2000 steps. “PC1000”: PC samplers using 1000 predictor and 1000 corrector steps.

H ARCHITECTURE IMPROVEMENTS

We explored several architecture designs to improve score-based models for both VE and VP SDEs. Our endeavor gives rise to new state-of-the-art sample quality on CIFAR-10, new state-of-the-art likelihood on uniformly dequantized CIFAR-10, and enables the first high-fidelity image samples of resolution $1024\times 1024$ from score-based generative models. Code and checkpoints are open-sourced at https://github.com/yang-song/score_sde.

H.1 SETTINGS FOR ARCHITECTURE EXPLORATION

Unless otherwise noted, all models are trained for $1.3\mathrm{M}$ iterations, and we save one checkpoint per $50\mathrm{k}$ iterations. For VE SDEs, we consider two datasets: $32\times 32$ CIFAR-10 (Krizhevsky et al., 2009) and $64\times 64$ CelebA (Liu et al., 2015), pre-processed following Song & Ermon (2020). We compare different configurations based on their FID scores averaged over checkpoints after $0.5\mathrm{M}$ iterations. For VP SDEs, we only consider the CIFAR-10 dataset to save computation, and compare models based on the average FID scores over checkpoints obtained between $0.25\mathrm{M}$ and $0.5\mathrm{M}$ iterations, because FIDs turn to increase after $0.5\mathrm{M}$ iterations for VP SDEs.
All FIDs are computed on $50\mathrm{k}$ samples with tensorflow_gan. For sampling, we use the PC sampler discretized at 1000 time steps. We choose reverse diffusion (see Appendix E) as the predictor. We use one corrector step per update of the predictor for VE SDEs with a signal-to-noise ratio of 0.16, but save the corrector step for VP SDEs since correctors there only give slightly better results but require double computation. We follow Ho et al. (2020) for optimization, including the learning rate, gradient clipping, and learning rate warm-up schedules. Unless otherwise noted, models are trained with the original discrete SMLD and DDPM objectives in Eqs. (1) and (3) and use a batch size of 128. The optimal architectures found under these settings are subsequently transferred to continuous objectives and deeper models. We also directly transfer the best architecture for VP SDEs to sub-VP SDEs, given the similarity of these two SDEs.

Figure 10: The effects of different architecture components for score-based models trained with VE perturbations.

Our architecture is mostly based on Ho et al. (2020). We additionally introduce the following components to maximize the potential improvement of score-based models.

1. 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).
2. Rescaling all skip connections by $1/\sqrt{2}$. This has been demonstrated effective in several best-in-class GAN models, including ProgressiveGAN (Karras et al., 2018), StyleGAN (Karras et al., 2019) and StyleGAN-2 (Karras et al., 2020b).
3. Replacing the original residual blocks in DDPM with residual blocks from BigGAN (Brock et al., 2018).
4. Increasing the number of residual blocks per resolution from 2 to 4.
5. 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 defined 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 explore more on it.
The exponential moving average (EMA) rate has a significant 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.9999 for VE and VP models respectively.

H.2 RESULTS ON CIFAR-10

All architecture components introduced above can improve the performance of score-based models trained with VE SDEs, as shown in Fig. 10. 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 explored always improve the performance on average for VE SDEs. For progressive growing, it is not clear which combination of configurations consistently performs the best, but the results are typically better than when no progressive growing architecture is used. Our best score-based model for VE SDEs 1) uses FIR upsampling/downsampling, 2) rescales skip connections, 3) employs BigGAN-type residual blocks, 4) uses 4 residual blocks per resolution instead of 2, and 5) 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).

We followed a similar procedure to examine these architecture components for VP SDEs, except that we skipped experiments on CelebA due to limited computing resources. The NCSN++ architecture worked decently well for VP SDEs, ranked 4th place over all 144 possible configurations. The top configuration, however, has a slightly different structure, which uses no FIR upsampling/downsampling and no progressive growing architecture compared to NCSN++. We name this model “DDPM++”, following the naming convention of Ho et al. (2020).

The basic NCSN++ model with 4 residual blocks per resolution achieves an FID of 2.45 on CIFAR-10, whereas the basic DDPM++ model achieves an FID of 2.78. Here in order to match the convention used in Karras et al. (2018); Song & Ermon (2019) and Ho et al. (2020), we report the lowest FID value over the course of training, rather than the average FID value over checkpoints after $0.5\mathrm{M}$ iterations (used for comparing different models of VE SDEs) or between $0.25\mathrm{M}$ and $0.5\mathrm{M}$ iterations (used for comparing VP SDE models) in our architecture exploration.

Switching from discrete training objectives to continuous ones in Eq. (7) further improves the FID values for all SDEs. To condition the NCSN++ model 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 (Tancik et al., 2020). The scale parameter of these random Fourier feature embeddings is fixed to 16. We also reduce the number of training iterations to $0.95\mathrm{M}$ to suppress overfitting. These changes improve the FID on CIFAR-10 from 2.45 to 2.38 for NCSN++ trained with the VE SDE, resulting in a model called “NCSN++ cont.”. In addition, we can further improve the FID from 2.38 to 2.20 by doubling the number of residual blocks per resolution for NCSN++ cont., resulting in the model denoted as “NCSN++ cont. (deep)”. All quantitative results are summarized in Table 3, and we provide random samples from our best model in Fig. 11.

Similarly, we can also condition the DDPM++ model on continuous time steps, resulting in a model “DDPM++ cont.”. When trained with the VP SDE, it improves the FID of 2.78 from DDPM++ to 2.55. When trained with the sub-VP SDE, it achieves an FID of 2.61. To get better performance, we used the Euler-Maruyama solver as the predictor for continuously-trained models, instead of the ancestral sampling predictor or the reverse diffusion predictor. This is because the discretization strategy of the original DDPM method does not match the variance of the continuous process well when $t\to 0$, which significantly hurts FID scores. As shown in Table 2, the likelihood values are 3.21 and 3.05 bits/dim for VP and sub-VP SDEs respectively. Doubling the depth, and training with $0.95\mathrm{M}$ iterations, we can improve both FID and bits/dim for both VP and sub-VP SDEs, leading to a model “DDPM++ cont. (deep)”. Its FID score is 2.41, same for both VP and sub-VP SDEs. When trained with the sub-VP SDE, it can achieve a likelihood of 2.99 bits/dim. Here all likelihood values are reported for the last checkpoint during training.

Figure 11: Unconditional CIFAR-10 samples from NCSN++ cont. (deep, VE).

H.3 HIGH RESOLUTION IMAGES

Encouraged by the success of NCSN++ on CIFAR-10, we proceed to test it on $1024\times 1024$ CelebA-HQ (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 a model similar to NCSN++ with the continuous objective (Eq. (7)) for around 2.4M iterations (please find the detailed architecture in our code release.) We use the PC sampler discretized at 2000 steps with the reverse diffusion predictor, one Langevin step per predictor update and a signal-to-noise ratio of 0.15. The scale parameter for the random Fourier feature embeddings is fixed to 16. We use the “input skip” progressive architecture for the input, and “output skip” progressive architecture for the output. We provide samples in Fig. 12. Although these samples are not perfect (e.g., there are visible flaws 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 significantly advance the performance of score-based generative models on this task.

Figure 12: Samples on $1024\times 1024$ CelebA-HQ from a modified NCSN++ model trained with the VE SDE.

I CONTROLLABLE GENERATION

Consider a forward SDE with the following general form

$$\mathrm{dx}=\mathbf{f}(\mathbf{x},t)\mathrm{d}t+\mathbf{G}(\mathbf{x},t)\mathrm{d}\mathbf{w},$$

and suppose the initial state distribution is $p_0(\mathbf{x}(0)\mid \mathbf{y})$. The density at time $t$ is $p_t(\mathbf{x}(t)\mid \mathbf{y})$ when conditioned on $\mathbf{y}$. Therefore, using Anderson (1982), the reverse-time SDE is given by

$$\mathrm{dx}=\left\{\mathbf{f}(\mathbf{x},t)-\nabla \cdot \left[\mathbf{G}(\mathbf{x},t)\mathbf{G}(\mathbf{x},t{)}^{\top }\right]-\mathbf{G}(\mathbf{x},t)\mathbf{G}(\mathbf{x},t{)}^{\top }{\nabla }_{\mathbf{x}}\log p_t(\mathbf{x}\mid \mathbf{y})\right\}\mathrm{d}t+\mathbf{G}(\mathbf{x},t)\mathrm{d}\mathbf{w}.(48)$$

Since $p_t(\mathbf{x}(t)\mid \mathbf{y})\propto p_t(\mathbf{x}(t))p(\mathbf{y}\mid \mathbf{x}(t))$, the score ${\nabla }_{\mathbf{x}}\log p_t(\mathbf{x}(t)\mid \mathbf{y})$ can be computed easily by

$${\nabla }_{\mathbf{x}}\log p_t(\mathbf{x}(t)\mid \mathbf{y})={\nabla }_{\mathbf{x}}\log p_t(\mathbf{x}(t))+{\nabla }_{\mathbf{x}}\log p(\mathbf{y}\mid \mathbf{x}(t)).(49)$$

This subsumes the conditional reverse-time SDE in Eq. (14) 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.1 Class-conditional sampling

When $\mathbf{y}$ represents class labels, we can train a time-dependent classifier $p_t(\mathbf{y}\mid \mathbf{x}(t))$ for class-conditional sampling. Since the forward SDE is tractable, we can easily create a pair of training data $(\mathbf{x}(t),\mathbf{y})$ by first sampling $(\mathbf{x}(0),\mathbf{y})$ from a dataset and then obtaining $\mathbf{x}(t)\sim p_{0t}(\mathbf{x}(t)\mid \mathbf{x}(0))$. Afterwards, we may employ a mixture of cross-entropy losses over different time steps, like Eq. (7), to train the time-dependent classifier $p_t(\mathbf{y}\mid \mathbf{x}(t))$.

To test this idea, we trained a Wide ResNet (Zagoruyko & Komodakis, 2016) (Wide-ResNet-28-10) on CIFAR-10 with VE perturbations. The classifier is conditioned on $\log {\sigma }_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 classifier over noise scales in Fig. 13. The score-based model is an unconditional NCSN++ (4 blocks/resolution) in Table 3, and we generate samples using the PC algorithm with 2000 discretization steps. The class-conditional samples are provided in Fig. 4, and an extended set of conditional samples is given in Fig. 13.

Figure 13: Class-conditional image generation by solving the conditional reverse-time SDE with PC. The curve shows the accuracy of our noise-conditional classifier over different noise scales.

I.2 IMPUTATION

Imputation is a special case of conditional sampling. Denote by $\Omega (\mathbf{x})$ and $\bar{\Omega }(\mathbf{x})$ the known and unknown dimensions of $\mathbf{x}$ respectively, and let ${\mathbf{f}}_{\Omega }(\cdot ,t)$ and ${\mathbf{G}}_{\bar{\Omega }}(\cdot ,t)$ denote $\mathbf{f}(\cdot ,t)$ and $\mathbf{G}(\cdot ,t)$ restricted to the unknown dimensions. For VE/VP SDEs, the drift coefficient $\mathbf{f}(\cdot ,t)$ is element-wise, and the diffusion coefficient $\mathbf{G}(\cdot ,t)$ is diagonal. When $\mathbf{f}(\cdot ,t)$ is element-wise, ${\mathbf{f}}_{\Omega }(\cdot ,t)$ denotes the same element-wise function applied only to the unknown dimensions. When $\mathbf{G}(\cdot ,t)$ is diagonal, ${\mathbf{G}}_{\Omega }(\cdot ,t)$ denotes the sub-matrix restricted to unknown dimensions.

For imputation, our goal is to sample from $p(\Omega (\mathbf{x}(0))\mid \Omega (\mathbf{x}(0))=\mathbf{y})$. Define a new diffusion process $\mathbf{z}(t)=\Omega (\mathbf{x}(t))$, and note that the SDE for $\mathbf{z}(t)$ can be written as

$$\mathrm{d}\mathbf{z}={\mathbf{f}}_{\Omega }(\mathbf{z},t)\mathrm{d}t+{\mathbf{G}}_{\Omega }(\mathbf{z},t)\mathrm{d}\mathbf{w}.$$

The reverse-time SDE, conditioned on $\Omega (\mathbf{x}(0))=\mathbf{y}$, is given by

$$\mathrm{d}\mathbf{z}=\left\{{\mathbf{f}}_{\Omega }(\mathbf{z},t)-\nabla \cdot \left[{\mathbf{G}}_{\Omega }(\mathbf{z},t){\mathbf{G}}_{\Omega }(\mathbf{z},t{)}^{\top }\right]-{\mathbf{G}}_{\Omega }(\mathbf{z},t){\mathbf{G}}_{\Omega }(\mathbf{z},t{)}^{\top }{\nabla }_{\mathbf{z}}\log p_t(\mathbf{z}\mid \Omega (\mathbf{z}(0))=\mathbf{y})\right\}\mathrm{d}t+{\mathbf{G}}_{\Omega }(\mathbf{z},t)\mathrm{d}\mathbf{w}.$$

Although $p_t(\mathbf{z}(t)\mid \Omega (\mathbf{x}(0))=\mathbf{y})$ is in general intractable, it can be approximated. Let $A$ denote the event $\Omega (\mathbf{x}(0))=\mathbf{y}$. We have

$$\begin{array}{rl}{p}_t(\mathbf{z}(t)\mid \Omega (\mathbf{x}(0))=\mathbf{y})=p_t(\mathbf{z}(t)\mid A)& =\int p_t(\mathbf{z}(t)\mid \Omega (\mathbf{x}(t)),A)p_t(\Omega (\mathbf{x}(t))\mid A)\mathrm{d}\Omega (\mathbf{x}(t))\\ & ={\mathbb{E}}_{p_t(\Omega (\mathbf{x}(t))\mid A)}\left[p_t(\mathbf{z}(t)\mid \Omega (\mathbf{x}(t)),A)\right]\\ & \approx {\mathbb{E}}_{p_t(\Omega (\mathbf{x}(t))\mid A)}\left[p_t(\mathbf{z}(t)\mid \Omega (\mathbf{x}(t)))\right]\\ & \approx p_t(\mathbf{z}(t)\mid \Omega (\mathbf{x}(t))),\end{array}$$

where $\Omega (\mathbf{x}(t))$ is a random sample from $p_t(\Omega (\mathbf{x}(t))\mid A)$, which is typically a tractable distribution. Therefore,

$$\begin{array}{rl}{\nabla }_{\mathbf{z}}\log p_t(\mathbf{z}(t)\mid \Omega (\mathbf{x}(0))& =\mathbf{y})\approx {\nabla }_{\mathbf{z}}\log p_t(\mathbf{z}(t)\mid \Omega (\mathbf{x}(t)))\\ & ={\nabla }_{\mathbf{z}}\log p_t([\mathbf{z}(t);\Omega (\mathbf{x}(t))]),\end{array}$$

where $[\mathbf{z}(t);\Omega (\mathbf{x}(t))]$ denotes a vector $\mathbf{u}(t)$ such that $\Omega (\mathbf{u}(t))=\Omega (\mathbf{x}(t))$ and $\bar{\Omega }(\mathbf{u}(t))=\mathbf{z}(t)$, and the identity holds because ${\nabla }_{\mathbf{z}}\log p_t([\mathbf{z}(t);\Omega (\mathbf{x}(t))])={\nabla }_{\mathbf{z}}\log p_t(\mathbf{z}(t)\mid \Omega (\mathbf{x}(t)))+{\nabla }_{\mathbf{z}}\log p(\Omega (\mathbf{x}(t)))={\nabla }_{\mathbf{z}}\log p_t(\mathbf{z}(t)\mid \Omega (\mathbf{x}(t)))$.

We provided an extended set of inpainting results in Figs. 14 and 15.

I.3 COLORIZATION

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

$$\left(\begin{array}{ccc}0.577& -0.816& 0\\ 0.577& 0.408& 0.707\\ 0.577& 0.408& -0.707\end{array}\right).$$

Because the transformations are all orthogonal matrices, the standard Wiener process $\mathbf{w}(t)$ 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. 16 and 17.

I.4 SOLVING GENERAL INVERSE PROBLEMS

Suppose we have two random variables $\mathbf{x}$ and $\mathbf{y}$, and we know the forward process of generating $\mathbf{y}$ from $\mathbf{x}$, given by $p(\mathbf{y}\mid \mathbf{x})$. The inverse problem is to obtain $\mathbf{x}$ from $\mathbf{y}$, that is, generating samples from $p(\mathbf{x}\mid \mathbf{y})$. In principle, we can estimate the prior distribution $p(\mathbf{x})$ and obtain $p(\mathbf{x}\mid \mathbf{y})$ using Bayes’ rule: $p(\mathbf{x}\mid \mathbf{y})=p(\mathbf{x})p(\mathbf{y}\mid \mathbf{x})/p(\mathbf{y})$. In practice, however, both estimating the prior and performing Bayesian inference are non-trivial.

Leveraging Eq. (48), score-based generative models provide one way to solve the inverse problem. Suppose we have a diffusion process {x(t)}t=0t\{\mathbf{x}(t)\}_{t=0}^t{x(t)}t=0t​ generated by perturbing $\mathbf{x}$ with an SDE, and a time-dependent score-based model ${\mathrm{s}}_{\theta \ast }(\mathbf{x}(t),t)$ trained to approximate ${\nabla }_{\mathbf{x}}\log p_t(\mathbf{x}(t))$. Once we have an estimate of ${\nabla }_{\mathbf{x}}\log p_t(\mathbf{x}(t)\mid \mathbf{y})$, we can simulate the reverse-time SDE in Eq. (48) to sample from $p_0(\mathbf{x}(0)\mid \mathbf{y})=p(\mathbf{x}\mid \mathbf{y})$. To obtain this estimate, we first observe that

$${\nabla }_{\mathbf{x}}\log p_t(\mathbf{x}(t)\mid \mathbf{y})={\nabla }_{\mathbf{x}}\log \int p_t(\mathbf{x}(t)\mid \mathbf{y}(t),\mathbf{y})p(\mathbf{y}(t)\mid \mathbf{y})\mathrm{d}\mathbf{y}(t),$$

where $\mathbf{y}(t)$ is defined via $\mathbf{x}(t)$ and the forward process $p(\mathbf{y}(t)\mid \mathbf{x}(t))$. Now assume two conditions:

• $p(\mathbf{y}(t)\mid \mathbf{y})$ is tractable. We can often derive this distribution from the interaction between the forward process and the SDE, like in the case of image imputation and colorization.
• $p_t(\mathbf{x}(t)\mid \mathbf{y}(t),\mathbf{y})\approx p_t(\mathbf{x}(t)\mid \mathbf{y}(t))$. For small $t,\mathbf{y}(t)$ is almost the same as $\mathbf{y}$ so the approximation holds. For large $t,\mathbf{y}$ becomes further away from $\mathbf{x}(t)$ in the Markov chain, and thus have smaller impact on $\mathbf{x}(t)$. Moreover, the approximation error for large $t$ matters less for the final sample, since it is used early in the sampling process.

Given these two assumptions, we have

$$\begin{array}{rl}{\nabla }_{\mathbf{x}}\log p_t(\mathbf{x}(t)\mid \mathbf{y})& \approx {\nabla }_{\mathbf{x}}\log \int p_t(\mathbf{x}(t)\mid \mathbf{y}(t))p(\mathbf{y}(t)\mid \mathbf{y})\mathrm{d}\mathbf{y}\\ & \approx {\nabla }_{\mathbf{x}}\log p_t(\mathbf{x}(t)\mid \hat{\mathbf{y}}(t))\\ & ={\nabla }_{\mathbf{x}}\log p_t(\mathbf{x}(t))+{\nabla }_{\mathbf{x}}\log p_t(\hat{\mathbf{y}}(t)\mid \mathbf{x}(t))\\ & \approx {\mathbf{s}}_{\theta \ast }(\mathbf{x}(t),t)+{\nabla }_{\mathbf{x}}\log p_t(\hat{\mathbf{y}}(t)\mid \mathbf{x}(t)),\end{array}(50)$$

where $\hat{\mathbf{y}}(t)$ is a sample from $p(\mathbf{y}(t)\mid \mathbf{y})$. Now we can plug Eq. (50) into Eq. (48) and solve the resulting reverse-time SDE to generate samples from $p(\mathbf{x}\mid \mathbf{y})$.

Figure 14: Extended inpainting results for $256\times 256$ bedroom images.

Figure 15: Extended inpainting results for $256\times 256$ church images.

Figure 16: Extended colorization results for $256\times 256$ bedroom images.

Figure 17: Extended colorization results for $256\times 256$ church images.