Score-based Generative Modeling in Latent Space
Year: 2021

Abstract:
Score-based generative models (SGMs) have recently demonstrated impressive results in terms of both sample quality and distribution coverage. However, they are usually applied directly in data space and often require thousands of network evaluations for sampling. Here, we propose the Latent Score-based Generative Model (LSGM), a novel approach that trains SGMs in a latent space, relying on the variational autoencoder framework. Moving from data to latent space allows us to train more expressive generative models, apply SGMs to non-continuous data, and learn smoother SGMs in a smaller space, resulting in fewer network evaluations and faster sampling. To enable training LSGMs end-to-end in a scalable and stable manner, we (i) introduce a new score-matching objective suitable to the LSGM setting, (ii) propose a novel parameterization of the score function that allows SGM to focus on the mismatch of the target distribution with respect to a simple Normal one, and (iii) analytically derive multiple techniques for variance reduction of the training objective. LSGM obtains a state-of-the-art FID score of 2.10 on CIFAR-10, outperforming all existing generative results on this dataset. On CelebA-HQ-256, LSGM is on a par with previous SGMs in sample quality while outperforming them in sampling time by two orders of magnitude. In modeling binary images, LSGM achieves state-of-the-art likelihood on the binarized OMNIGLOT dataset. Our implementation is available at https://github.com/NVlabs/LSGM.

heading: 1 Introduction
text: The long-standing goal of likelihood-based generative learning is to faithfully learn a data distribution, while also generating high-quality samples. Achieving these two goals simultaneously is a tremendous challenge, which has led to the development of a plethora of different generative models. Recently, score-based generative models (SGMs) demonstrated astonishing results in terms of both high sample quality and likelihood [1, 2]. These models define a forward diffusion process that maps data to noise by gradually perturbing the input data. Generation corresponds to a reverse process that synthesizes novel data via iterative denoising, starting from random noise. The problem then reduces to learning the score function—the gradient of the log-density—of the perturbed data [3]. In a seminal work, Song et al. [2] show how this modeling approach is described with a stochastic differential equation (SDE) framework which can be converted to maximum likelihood training [4]. Variants of SGMs have been applied to images [1, 2, 5, 6], audio [7, 8, 9, 10], graphs [11] and point clouds [12, 13].
Albeit high quality, sampling from SGMs is computationally expensive. This is because generation amounts to solving a complex SDE, or equivalently ordinary differential equation (ODE) (denoted as the probability flow ODE in [2]), that maps a simple base distribution to the complex data distribution. The resulting differential equations are typically complex and solving them accurately requires numerical integration with very small step sizes, which results in thousands of neural network evaluations [1, 2, 6]. Furthermore, generation complexity is uniquely defined by the underlying data distribution and the forward SDE for data perturbation, implying that synthesis speed cannot be ∗Equal contribution.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
increased easily without sacrifices. Moreover, SDE-based generative models are currently defined for continuous data and cannot be applied effortlessly to binary, categorical, or graph-structured data.
Here, we propose the Latent Score-based Generative Model (LSGM), a new approach for learning SGMs in latent space, leveraging a variational autoencoder (VAE) framework [14, 15]. We map the input data to latent space and apply the score-based generative model there. The score-based model is then tasked with modeling the distribution over the embeddings of the data set. Novel data synthesis is achieved by first generating embeddings via drawing from a simple base distribution followed by iterative denoising, and then transforming this embedding via a decoder to data space (see Fig. 1). We can consider this model a VAE with an SGM prior. Our approach has several key advantages: Synthesis Speed: By pretraining the VAE with a Normal prior first, we can bring the marginal distribution over encodings (the aggregate posterior) close to the Normal prior, which is also the SGM’s base distribution. Consequently, the SGM only needs to model the remaining mismatch, resulting in a less complex model from which sampling becomes easier. Furthermore, we can tailor the latent space according to our needs. For example, we can use hierarchical latent variables and apply the diffusion model only over a subset of them, further improving synthesis speed. Expressivity: Training a regular SGM can be considered as training a neural ODE directly on the data [2]. However, previous works found that augmenting neural ODEs [16, 17] and more generally generative models [18, 19, 20, 21] with latent variables improves their expressivity. Consequently, we expect similar performance gains from combining SGMs with a latent variable framework. Tailored Encoders and Decoders: Since we use the SGM in latent space, we can utilize carefully designed encoders and decoders mapping between latent and data space, further improving expressivity. Additionally, the LSGM method can therefore be naturally applied to non-continuous data.
LSGMs can be trained end-to-end by maximizing the variational lower bound on the data likelihood. Compared to regular score matching, our approach comes with additional challenges, since both the score-based denoising model and its target distribution, formed by the latent space encodings, are learnt simultaneously. To this end, we make the following technical contributions: (i) We derive a new denoising score matching objective that allows us to efficiently learn the VAE model and the latent SGM prior at the same time. (ii) We introduce a new parameterization of the latent space score function, which mixes a Normal distribution with a learnable SGM, allowing the SGM to model only the mismatch between the distribution of latent variables and the Normal prior. (iii) We propose techniques for variance reduction of the training objective by designing a new SDE and by analytically deriving importance sampling schemes, allowing us to stably train deep LSGMs. Experimentally, we achieve state-of-the-art 2.10 FID on CIFAR-10 and 7.22 FID on CelebA-HQ-256, and significantly improve upon likelihoods of previous SGMs. On CelebA-HQ-256, we outperform previous SGMs in synthesis speed by two orders of magnitude. We also model binarized images, MNIST and OMNIGLOT, achieving state-of-the-art likelihood on the latter.

heading: 2 Background
text: Here, we review continuous-time score-based generative models (see [2] for an in-depth discussion). Consider a forward diffusion process {zt}t=1t=0 for continuous time variable t ∈ [0, 1], where z0 is the starting variable and zt its perturbation at time t. The diffusion process is defined by an Itô SDE:
dz = f(t)z dt+ g(t) dw (1) where f : R→ R and g : R→ R are scalar drift and diffusion coefficients, respectively, and w is the standard Wiener process. f(t) and g(t) can be designed such that z1 ∼ N (z1;0, I) follows a
Normal distribution at the end of the diffusion process.2 Song et al. [2] show that the SDE in Eq. 1 can be converted to a generative model by first sampling from z1 ∼ N (z1;0, I) and then running the reverse-time SDE dz = [f(t)z−g(t)2∇z log qt(z)] dt+g(t) dw̄, where w̄ is a reverse-time standard Wiener process and dt is an infinitesimal negative time step. The reverse SDE requires knowledge of ∇zt log qt(zt), the score function of the marginal distribution under the forward diffusion at time t. One approach for estimating it is via the score matching objective3:
min θ
Et∼U [0,1] [ λ(t)Eq(z0)Eq(zt|z0)[||∇zt log q(zt)−∇zt log pθ(zt)|| 2 2] ]
(2)
that trains the parameteric score function ∇zt log pθ(zt) at time t ∼ U [0, 1] for a given weighting coefficient λ(t). q(z0) is the z0-generating distribution and q(zt|z0) is the diffusion kernel, which is available in closed form for certain f(t) and g(t). Since ∇zt log q(zt) is not analytically available, Song et al. [2] rely on denoising score matching [22] that converts the objective in Eq. 2 to:
min θ
Et∼U [0,1] [ λ(t)Eq(z0)Eq(zt|z0)[||∇zt log q(zt|z0)−∇zt log pθ(zt)|| 2 2] ] + C (3)
Vincent [22] shows C = Et∼U [0,1][λ(t)Eq(z0)Eq(zt|z0)[||∇zt log q(zt)||22 − ||∇zt log q(zt|z0)||22]] is independent of θ, making the minimizations in Eq. 3 and Eq. 2 equivalent. Song et al. [4] show that for λ(t) = g(t)2/2, the minimizations correspond to approximate maximum likelihood training based on an upper on the Kullback-Leibler (KL) divergence between the target distribution and the distribution defined by the reverse-time generative SDE with the learnt score function. In particular, the objective of Eq. 2 can then be written:
KL ( q(z0)||pθ(z0) ) ≤ Et∼U[0,1]
[ g(t)2
2 Eq(z0)Eq(zt|z0)
[ ||∇zt log q(zt)−∇zt log pθ(zt)|| 2 2 ]] (4)
which can again be transformed into denoising score matching (Eq. 3) following Vincent [22].

heading: 3 Score-based Generative Modeling in Latent Space
text: The LSGM framework in Fig. 1 consists of the encoder qφ(z0|x), SGM prior pθ(z0), and decoder pψ(x|z0). The SGM prior leverages a diffusion process as defined in Eq. 1 and diffuses z0 ∼ qφ(z0|x) samples in latent space to the standard Normal distribution p(z1) = N (z1;0, I). Generation uses the reverse SDE to sample from pθ(z0) with time-dependent score function∇zt log pθ(zt), and the decoder pψ(x|z0) to map the synthesized encodings z0 to data space. Formally, the generative process is written as p(z0,x) = pθ(z0)pψ(x|z0). The goal of training is to learn {φ,θ,ψ}, the parameters of the encoder qφ(z0|x), score function∇zt log pθ(zt), and decoder pψ(x|z0), respectively. We train LSGM by minimizing the variational upper bound on negative data log-likelihood log p(x):
L(x,φ,θ,ψ) = Eqφ(z0|x) [ − log pψ(x|z0) ] +KL ( qφ(z0|x)||pθ(z0) ) (5)
= Eqφ(z0|x) [ − log pψ(x|z0) ]︸ ︷︷ ︸ reconstruction term +Eqφ(z0|x) [ log qφ(z0|x) ]︸ ︷︷ ︸ negative encoder entropy +Eqφ(z0|x) [ − log pθ(z0) ]︸ ︷︷ ︸ cross entropy
(6)
following a VAE approach [14, 15], where qφ(z0|x) approximates the true posterior p(z0|x). In this paper, we use Eq. 6 with decomposed KL divergence into its entropy and cross entropy terms. The reconstruction and entropy terms are estimated easily for any explicit encoder as long as the reparameterization trick is available [14]. The challenging part in training LSGM is to train the cross entropy term that involves the SGM prior. We motivate and present our expression for the cross-entropy term in Sec. 3.1, the parameterization of the SGM prior in Sec. 3.2, different weighting mechanisms for the training objective in Sec. 3.3, and variance reduction techniques in Sec. 3.4.

heading: 3.1 The Cross Entropy Term
text: One may ask, why not train LSGM with Eq. 5 and rely on the KL in Eq. 4. Directly using the KL expression in Eq. 4 is not possible, as it involves the marginal score ∇zt log q(zt), which is unavailable analytically for common non-Normal distributions q(z0) such as Normalizing flows.
2Other distributions at t = 1 are possible; for instance, see the “variance-exploding” SDE in [2]. In this paper, however, we use only SDEs converging towardsN (z1;0, I) at t = 1.
3We omit the t-subscript of the diffused distributions qt in all score functions of the form∇zt log qt(zt).
Transforming into denoising score matching does not help either, since in that case the problematic ∇zt log q(zt) term appears in the C term (see Eq. 3). In contrast to previous works [2, 22], we cannot simply drop C, since it is, in fact, not constant but depends on q(zt), which is trainable in our setup.
To circumvent this problem, we instead decompose the KL in Eq. 5 and rather work directly with the cross entropy between the encoder distribution q(z0|x) and the SGM prior p(z0). We show: Theorem 1. Given two distributions q(z0|x) and p(z0), defined in the continuous space RD, denote the marginal distributions of diffused samples under the SDE in Eq. 1 at time t with q(zt|x) and p(zt). Assuming mild smoothness conditions on log q(zt|x) and log p(zt), the cross entropy is:
CE(q(z0|x)||p(z0)) = Et∼U[0,1]
[ g(t)2
2 Eq(zt,z0|x)
[ ||∇zt log q(zt|z0)−∇zt log p(zt)|| 2 2 ]] + D 2 log ( 2πeσ20 ) ,
with q(zt, z0|x) = q(zt|z0)q(z0|x) and a Normal transition kernel q(zt|z0) = N (zt;µt(z0), σ2t I), where µt and σ 2 t are obtained from f(t) and g(t) for a fixed initial variance σ 2 0 at t = 0.
A proof with generic expressions for µt and σ 2 t as well as an intuitive interpretation are in App. A.
Importantly, unlike for the KL objective of Eq. 4, no problematic terms depending on the marginal score ∇zt log q(zt|x) arise. This allows us to use this denoising score matching objective for the cross entropy term in Theorem 1 not only for optimizing p(z0) (which is commonly done in the score matching literature), but also for the q(z0|x) encoding distribution. It can be used even with complex q(z0|x) distributions, defined, for example, in a hierarchical fashion [20, 21] or via Normalizing flows [23, 24]. Our novel analysis shows that, for diffusion SDEs following Eq. 1, only the cross entropy can be expressed purely with ∇zt log q(zt|z0). Neither KL nor entropy in [4] can be expressed without the problematic term∇zt log q(zt|x) (details in the Appendix). Note that in Theorem 1, the term∇zt log p(zt) in the score matching expression corresponds to the score that originates from diffusing an initial p(z0) distribution. In practice, we use the expression to learn an SGM prior pθ(z0), which models∇zt log p(zt) by a neural network. With the learnt score ∇zt log pθ(zt) (here we explicitly indicate the parameters θ to clarify that this is the learnt model), the actual SGM prior is defined via the generative reverse-time SDE (or, alternatively, a closely-connected ODE, see Sec. 2 and App. D), which generally defines its own, separate marginal distribution pθ(z0) at t = 0. Importantly, the learnt, approximate score∇zt log pθ(zt) is not necessarily the same as one would obtain when diffusing pθ(z0). Hence, when considering the learnt score∇zt log pθ(zt), the score matching expression in our Theorem only corresponds to an upper bound on the cross entropy between q(z0|x) and pθ(z0) defined by the generative reverse-time SDE. This is discussed in detail in concurrent works [4, 25]. Hence, from the perspective of the learnt SGM prior, we are training with an upper bound on the cross entropy (similar to the bound on the KL in Eq. 4), which can also be considered as the continuous version of the discretized variational objective derived by Ho et al. [1].

heading: 3.2 Mixing Normal and Neural Score Functions
text: In VAEs [14], p(z0) is often chosen as a standard Normal N (z0;0, I). For recent hierarchical VAEs [20, 21], using the reparameterization trick, the prior can be converted to N (z0;0, I) (App. E). Considering a single dimensional latent space, we can assume that the prior at time t is in the form of a geometric mixture p(zt) ∝ N (zt; 0, 1)1−αp′θ(zt)α where p′θ(zt) is a trainable SGM prior and α ∈ [0, 1] is a learnable scalar mixing coefficient. Formulating the prior this way has crucial advantages: (i) We can pretrain LSGM’s autoencoder networks assuming α=0, which corresponds to training the VAE with a standard Normal prior. This pretraining step will bring the distribution of latent variable close to N (z0; 0, 1), allowing the SGM prior to learn a much simpler distribution in the following end-to-end training stage. (ii) The score function for this mixture is of the form ∇zt log p(zt) = −(1− α)zt + α∇zt log p′θ(zt). When the score function is dominated by the linear term, we expect that the reverse SDE can be solved faster, as its drift is dominated by this linear term.
For our multivariate latent space, we obtain diffused samples at time t by sampling zt ∼ q(zt|z0) with zt = µt(z0) + σt , where ∼ N ( ;0, I). Since we have ∇zt log q(zt|z0) = − /σt, similar to [1], we parameterize the score function by ∇zt log p(zt) := − θ(zt, t)/σt, where θ(zt, t) := σt(1 − α) zt + α ′θ(zt, t) is defined by our mixed score parameterization that is applied elementwise to the components of the score. With this, we simplify the cross entropy expression to:
CE(qφ(z0|x)||pθ(z0)) = Et∼U[0,1] [ w(t)
2 Eqφ(zt,z0|x),
[ || − θ(zt, t)||22 ]] + D 2 log ( 2πeσ20 ) , (7)
where w(t) = g(t)2/σ2t is a time-dependent weighting scalar.
3.3 Training with Different Weighting Mechanisms Table 1: Weighting mechanisms
Mechanism Weights
Weighted wll(t) = g(t)2/σ2t Unweighted wun(t) = 1 Reweighted wre(t) = g(t)2 The weighting term w(t) in Eq. 7 trains the prior with maximum likelihood. Similar to [1, 2], we observe that when w(t) is dropped while training the SGM prior (i.e., w(t) = 1), LSGM often yields higher quality samples at a small cost in likelihood. However, in our case, we can only drop the weighting when training the prior. When updating the encoder parameters, we still need to use the maximum likelihood weighting to ensure that the encoder q(z0|x) is brought closer to the true posterior p(z0|x)4. Tab. 1 summarizes three weighting mechanisms we consider in this paper: wll(t) corresponds to maximum likelihood, wun(t) is the unweighted objective used by [1, 2], and wre(t) is a variant obtained by dropping only 1/σ2t . This weighting mechanism has a similar affect on the sample quality as wun(t) = 1; however, in Sec. 3.4, we show that it is easier to define a variance reduction scheme for this weighting mechanism.
The following summarizes our training objectives (with t ∼ U [0, 1] and ∼ N ( ;0, I)):
min φ,ψ
Eqφ(z0|x) [ −log pψ(x|z0) ] +Eqφ(z0|x) [ log qφ(z0|x) ] +Et, ,q(zt|z0),qφ(z0|x)
[ wll(t)
2 || − θ(zt, t)||22
] (8)
min θ Et, ,q(zt|z0),qφ(z0|x)
[ wll/un/re(t)
2 || − θ(zt, t)||22
] with q(zt|z0) = N (zt;µt(z0), σ 2 t I), (9)
where Eq. 8 trains the VAE encoder and decoder parameters {φ,ψ} using the variational bound L(x,φ,θ,ψ) from Eq. 6. Eq. 9 trains the prior with one of the three weighting mechanisms. Since the SGM prior participates in the objective only in the cross entropy term, we only consider this term when training the prior. Efficient algorithms for training with the objectives are presented in App. G.

heading: 3.4 Variance Reduction
text: The objectives in Eqs. 8 and 9 involve sampling of the time variable t, which has high variance [26]. We introduce several techniques for reducing this variance for all three objective weightings. We focus on the “variance preserving” SDEs (VPSDEs) [2, 1, 27], defined by dz = − 12β(t)z dt+ √ β(t) dw where β(t) = β0 + (β1 − β0)t linearly interpolates in [β0, β1] (other SDEs discussed in App. B). We denote the marginal distribution of latent variables by q(z0) := Epdata(x)[q(z0|x)]. Here, we derive variance reduction techniques for CE(q(z0)||p(z0)), assuming that both q(z0) = p(z0) = N (z0;0, I). This is a reasonable simplification for our analysis because pretraining our LSGM model with a N (z0;0, I) prior brings q(z0) close to N (z0;0, I) and our SGM prior is often dominated by the fixed Normal mixture component. We empirically observe that the variance reduction techniques developed with this assumption still work well when q(z0) and p(z0) are not exactly N (z0;0, I). Variance reduction for likelihood weighting: In App. B, for q(z0) = p(z0) = N (z0;0, I), we show CE(q(z0)||p(z0)) is given by D2 Et∼U [0,1][d log σ 2 t /dt] + const. We consider two approaches:
(1) Geometric VPSDE: To reduce the variance sampling uniformly from t, we can design the SDE such that d log σ2t /dt is constant for t ∈ [0, 1]. We show in App. B that a β(t) = log(σ2max/σ2min) σ2t (1−σ2t ) with geometric variance σ2t = σ 2 min(σ 2 max/σ 2 min)
t satisfies this condition. We call a VPSDE with this β(t) a geometric VPSDE. σ2min and σ 2 max are the hyperparameters of the SDE, with 0<σ 2 min<σ 2 max<1. Although our geometric VPSDE has a geometric variance progression similar to the “variance exploding” SDE (VESDE) [2], it still enjoys the “variance preserving” property of the VPSDE. In App. B, we show that the VESDE does not come with a reduced variance for t-sampling by default.
(2) Importance sampling (IS): We can keep β(t) and σ2t unchanged for the original linear VPSDE, and instead use IS to minimize variance. The theory of IS shows that the proposal r(t) ∝ d log σ2t /dt has minimum variance [28]. In App. B, we show that we can sample from r(t) using inverse transform sampling t = var−1((σ21) ρ(σ20) 1−ρ) where var−1 is the inverse of σ2t and ρ ∼ U [0, 1]. This variance reduction technique is available for any VPSDE with arbitrary β(t).
In Fig. 2, we train a small LSGM on CIFAR-10 with wll weighting using (i) the original VPSDE with uniform t sampling, (ii) the same SDE but with our IS from t, and (iii) the proposed geometric
4Minimizing L(x,φ,θ,ψ) w.r.t φ is equivalent to minimizing KL ( q(z0|x)||p(z0|x) ) w.r.t q(z0|x).
VPSDE. Note how both (ii) and (iii) significantly reduce the variance and allow us to monitor the progress of the training objective. In this case, (i) has difficulty minimizing the objective due to the high variance. In App. B, we show how IS proposals can be formed for other SDEs, including the VESDE and Sub-VPSDE from [2].
Variance reduction for unweighted and reweighted objectives: When training with wun, analytically deriving IS proposal distributions for arbitrary β(t) is challenging. For linear VPSDEs, we provide a derivation in App. B to obtain the optimal IS distribution. In contrast, defining IS proposal distributions is easier when training with wre. In App. B, we show that the optimal distribution is in the form r(t) ∝ dσ2t /dtwhich is sampled by t=var−1((1−ρ)σ20 +ρσ21) with ρ ∼ U [0, 1]. In Fig. 3, we visualize the IS distributions for the three weighting mechanisms for the linear VPSDE with the original [β0, β1] parameters from [2]. r(t) for the likelihood weighting is more tilted towards t = 0 due to the 1/σ2t term in wll.
When using differently weighted objectives for training, we can either sample separate t with different IS distributions for each objective, or use IS for the SGM objective (Eq. 9) and reweight the samples according to the likelihood objective for encoder training (Eq. 8). See App. G for details.

heading: 4 Related Work
text: Our work builds on score-matching [29, 30, 31, 32, 33, 34, 35, 36, 37], specifically denoising score matching [22], which makes our work related to recent generative models using denoising score matching- and denoising diffusion-based objectives [3, 38, 1, 2, 6]. Among those, [1, 6] use a discretized diffusion process with many noise scales, building on [27], while Song et al. [2] introduce the continuous time framework using SDEs. Experimentally, these works focus on image modeling and, contrary to us, work directly in pixel space. Various works recently tried to address the slow sampling of these types of models and further improve output quality. [39] add an adversarial objective, [5] introduce non-Markovian diffusion processes that allow to trade off synthesis speed, quality, and sample diversity, [40] learn a sequence of conditional energy-based models for denoising, [41] distill the iterative sampling process into single shot synthesis, and [42] learn an adaptive noise schedule, which is adjusted during synthesis to accelerate sampling. Further, [26] propose empirical variance reduction techniques for discretized diffusions and introduce a new, heuristically motivated, noise schedule. In contrast, our proposed noise schedule and our variance reduction techniques are analytically derived and directly tailored to our learning setting in the continuous time setup.
Recently, [11] presented a method to generate graphs using score-based models, relaxing the entries of adjacency matrices to continuous values. LSGM would allow to model graph data more naturally using encoders and decoders tailored to graphs [43, 44, 45, 46].
Since our model can be considered a VAE [14, 15] with score-based prior, it is related to approaches that improve VAE priors. For example, Normalizing flows and hierarchical distributions [23, 24, 47, 48, 20, 21], as well as energy-based models [49, 50, 51, 52, 53] have been proposed as VAE priors. Furthermore, classifiers [54, 55, 56], adversarial methods [57], and other techniques [58, 59] have been used to define prior distributions implicitly. In two-stage training, a separate generative model is trained in latent space as a new prior after training the VAE itself [60, 61, 62, 63, 64, 10]. Our work also bears a resemblance to recent methods on improving the sampling quality in generative adversarial networks using gradient flows in the latent space [65, 66, 67, 68], with the main difference that these prior works use a discriminator to update the latent variables, whereas we train an SGM.
Concurrent works: [10] proposed to learn a denoising diffusion model in the latent space of a VAE for symbolic music generation. This work does not introduce an end-to-end training framework of the combined VAE and denoising diffusion model and instead trains them in two separate stages. In contrast, concurrently with us [69] proposed an end-to-end training approach, and [70] combines contrastive learning with diffusion models in the latent space of VAEs for controllable generation. However, [10, 69, 70] consider the discretized diffusion objective [1], while we build on the continuous time framework. Also, these models are not equipped with the mixed score parameterization and variance reduction techniques, which we found crucial for the successful training of SGM priors.
Additionally, [71, 4, 25] concurrently with us proposed likelihood-based training of SGMs in data space5. [4] developed a bound for the data likelihood in their Theorem 3 of their second version, using a denoising score matching objective, closely related to our cross entropy expression. However, our cross entropy expression is much simpler as we show how several terms can be marginalized out analytically for the diffusion SDEs employed by us (see our proof in App. A). The same marginalization can be applied to Theorem 3 in [4] when the drift coefficient takes a special affine form (i.e., f(z, t) = f(t)z). Moreover, [25] discusses the likelihood-based training of SGMs from a fundamental perspective and shows how several score matching objectives become a variational bound on the data likelihood. [71] introduced a notion of signal-to-noise ratio (SNR) that results in a noise-invariant parameterization of time that depends only on the initial and final noise. Interestingly, our importance sampling distribution in Sec. 3.4 has a similar noise-invariant parameterization of time via t = var−1((σ21) ρ(σ20) 1−ρ), which also depends only on the initial and final diffusion process variances. We additionally show that this time parameterization results in the optimal minimumvariance objective, if the distribution of latent variables follows a standard Normal distribution. Finally, [72] proposed a modified time parameterization that allows modeling unbounded data scores.

heading: 5 Experiments
text: Here, we examine the efficacy of LSGM in learning generative models for images.
Implementation details: We implement LSGM using the NVAE [20] architecture as VAE backbone and NCSN++ [2] as SGM backbone. NVAE has a hierarchical latent structure. The diffusion process input z0 is constructed by concatenating the latent variables from all groups in the channel dimension. For NVAEs with multiple spatial resolutions in latent groups, we only feed the smallest resolution groups to the SGM prior and assume that the remaining groups have a standard Normal distribution.
Sampling: To generate samples from LSGM at test time, we use a black-box ODE solver [73] to sample from the prior. Prior samples are then passed to the decoder to generate samples in data space.
Evaluation: We measure NELBO, an upper bound on negative log-likelihood (NLL), using Eq. 6. For estimating log p(z0), we rely on the probability flow ODE [2], which provides an unbiased but stochastic estimation of log p(z0). This stochasticity prevents us from performing an importance weighted estimation of NLL [74] (see App. F for details). For measuring sample quality, Fréchet inception distance (FID) [75] is evaluated with 50K samples. Implementation details in App. G.

heading: 5.1 Main Results
text: Unconditional color image generation: Here, we present our main results for unconditional image generation on CIFAR-10 [89] (Tab. 2) and CelebA-HQ-256 (5-bit quantized) [88] (Tab. 3). For CIFAR-10, we train 3 different models: LSGM (FID) and LSGM (balanced) both use the VPSDE with linear β(t) and wun-weighting for the SGM prior in Eq. 9, while performing IS as derived in Sec. 3.4. They only differ in how the backbone VAE is trained. LSGM (NLL) is a model that is trained with our novel geometric VPSDE, using wll-weighting in the prior objective (further details in App. G). When set up for high image quality, LSGM achieves a new state-of-the-art FID of 2.10. When tuned towards NLL, we achieve a NELBO of 2.87, which is significantly better than previous score-based models. Only autoregressive models, which come with very slow synthesis, and VDVAE [21] reach similar or higher likelihoods, but they usually have much poorer image quality.
For CelebA-HQ-256, we observe that when LSGM is trained with different SDE types and weighting mechanisms, it often obtains similar NELBO potentially due to applying the SGM prior only to small latent variable groups and using Normal priors at the larger groups. With wre-weighting and linear VPSDE, LSGM obtains the state-of-the-art FID score of 7.22 on a par with the original SGM [2].
For both datasets, we also report results for the VAE backbone used in our LSGM. Although this baseline achieves competitive NLL, its sample quality is behind our LSGM and the original SGM.
Modeling binarized images: Next, we examine LSGM on dynamically binarized MNIST [93] and OMNIGLOT [74]. We apply LSGM to binary images using a decoder with pixel-wise independent Bernoulli distributions. For these datasets, we report both NELBO and NLL in nats in Tab. 4 and Tab. 5. On OMNIGLOT, LSGM achieves state-of-the-art likelihood of ≤87.79 nat, outperforming previous models including VAEs with autoregressive decoders, and even when comparing its NELBO
5We build on the V1 version of [4], which was substantially updated after the NeurIPS submission deadline.
Table 5: Dynamically binarized MNIST results.
Method NELBO↓ NLL↓ Ours LSGM 78.47 ≤78.47
VAEs
NVAE [20] 79.56 78.01 BIVA [48] 80.06 78.41 IAF-VAE [24] 80.80 79.10 DVAE++ [51] - 78.49
Aut. Reg. PixelVAE++ [91] - 78.00 VampPrior [59] - 78.45 MAE [92] - 77.98
against importance weighted estimation of NLL for other methods. On MNIST, LSGM outperforms previous VAEs in NELBO, reaching a NELBO 1.09 nat lower than the state-of-the-art NVAE.
Qualitative results: We visualize qualitative results for all datasets in Fig. 5. On the complex multimodal CIFAR-10 dataset, LSGM generates sharp and high-quality images. On CelebA-HQ-256, LSGM generates diverse samples from different ethnicity and age groups with varying head poses and facial expressions. On MNIST and OMNIGLOT, the generated characters are sharp and high-contrast.
Sampling time: We compare LSGM against the original SGM [2] trained on the CelebA-HQ-256 dataset in terms of sampling time and number of function evaluations (NFEs) of the ODE solver. Song et al. [2] propose two main sampling techniques including predictor-corrector (PC) and probability flow ODE. PC sampling involves 4000 NFEs and takes 44.6 min. on a Titan V for a batch of 16 images. It yields 7.23 FID score (see Tab. 3). ODE-based sampling from SGM takes 3.91 min. with 335 NFEs, but it obtains a poor FID score of 128.13 with 10−5 as ODE solver error tolerance6.
In a stark contrast, ODE-based sampling from our LSGM takes 0.07 min. with average of 23 NFEs, yielding 7.22 FID score. LSGM is 637× and 56× faster than original SGM’s [2] PC and ODE
6We use the VESDE checkpoint at https://github.com/yang-song/score_sde_pytorch. Song et al. [2] report that ODE-based sampling yields worse FID scores for their models (see D.4 in [2]). The problem is more severe for VESDEs. Unfortunately, at submission time only a VESDE model was released.
sampling, respectively. In Fig. 4, we visualize FID scores and NFEs for different ODE solver error tolerances. Our LSGM achieves low FID scores for relatively large error tolerances.
We identify three main reasons for this significantly faster sampling from LSGM: (i) The SGM prior in our LSGM models latent variables with 32×32 spatial dim., whereas the original SGM [2] directly models 256×256 images. The larger spatial dimensions require a deeper network to achieve a large receptive field. (ii) Inspecting the SGM prior in our model suggests that the score function is heavily dominated by the linear term at the end of training, as the mixing coefficients α are all < 0.02. This makes our SGM prior smooth and numerically faster to solve. (iii) Since SGM is formed in the latent space in our model, errors from solving the ODE can be corrected to some degree using the VAE decoder, while in the original SGM [2] errors directly translate to artifacts in pixel space.

heading: 5.2 Ablation Studies
text: SDEs, objective weighting mechanisms and variance reduction. In Tab. 6, we analyze the different weighting mechanisms and variance reduction techniques and compare the geometric VPSDE with the regular VPSDE with linear β(t) [1, 2]. In the table, SGM-obj.-weighting denotes the weighting mechanism used when training the SGM prior (via Eq. 9). t-sampling (SGM-obj.) indicates the sampling approach for t, where rll(t), run(t) and rre(t) denote the IS distributions for the weighted (likelihood), the unweighted, and the reweighted objective, respectively. For training the VAE encoder qφ(z0|x) (last term in Eq. 8), we either sample a separate batch t with importance sampling following rll(t) (only necessary when the SGM prior is not trained with wll itself), or we reweight the samples drawn for training the prior according to the likelihood objective (denoted by rew.). n/a indicates fields that do not apply: The geometric VPSDE has optimal variance for the weighted (likelihood) objective already with uniform sampling; there is no additional IS distribution. Also, we did not derive IS distributions for the geometric VPSDE for wun. NaN indicates experiments that failed due to training instabilities. Previous work [20, 21] have reported instability in training large VAEs. We find that our method inherits similar instabilities from VAEs; however, importance sampling often stabilizes training our LSGM. As expected, we obtain the best NELBOs (red) when training with the weighted, maximum likelihood objective (wll). Importantly, our new geometric VPSDE achieves the best NELBO. Furthermore, the best FIDs (blue) are obtained either by unweighted (wun) or reweighted (wre) SGM prior training, with only slightly worse NELBOs. These experiments were run on the CIFAR10 dataset, using a smaller model than for our main results above (details in App. G).
End-to-end training. We proposed to train LSGM end-to-end, in contrast to [10]. Using a similar setup as above we compare end-to-end training of LSGM during the second stage with freezing the VAE encoder and decoder and only training the SGM prior in latent space during the second stage. When training the model end-to-end, we achieve an FID of 5.19 and NELBO of 2.98; when freezing the VAE networks during the second stage, we only get an FID of 9.00 and NELBO of 3.03. These results clearly motivate our end-to-end training strategy.
Mixing Normal and neural score functions. We generally found training LSGM without our proposed “mixed score” formulation (Sec. 3.2) to be unstable during end-to-end training, highlighting its importance. To quantify the contribution of the mixed score parametrization for a stable model, we train a small LSGM with only one latent variable group. In this case, without the mixed score, we reached an FID of 34.71 and NELBO of 3.39; with it, we got an FID of 7.60 and NELBO of 3.29. Without the inductive bias provided by the mixed score, learning that the marginal distribution is close to a Normal one for large t purely from samples can be very hard in the high-dimensional latent space, where our diffusion is run. Furthermore, due to our importance sampling schemes, we tend to oversample small, rather than large t. However, synthesizing high-quality images requires an accurate score function estimate for all t. On the other hand, the log-likelihood of samples is highly sensitive to local image statistics and primarily determined at small t. It is plausible that we are still able to learn a reasonable estimate of the score function for these small t even without the mixed score formulation. That may explain why log-likelihood suffers much less than sample quality, as estimated by FID, when we remove the mixed score parameterization.
Additional experiments and model samples are presented in App. H.

heading: 6 Conclusions
text: We proposed the Latent Score-based Generative Model, a novel framework for end-to-end training of score-based generative models in the latent space of a variational autoencoder. Moving from data to latent space allows us to form more expressive generative models, model non-continuous data, and reduce sampling time using smoother SGMs. To enable training latent SGMs, we made three core contributions: (i) we derived a simple expression for the cross entropy term in the variational objective, (ii) we parameterized the SGM prior by mixing Normal and neural score functions, and (iii) we proposed several techniques for variance reduction in the estimation of the training objective. Experimental results show that latent SGMs outperform recent pixel-space SGMs in terms of both data likelihood and sample quality, and they can also be applied to binary datasets. In large image generation, LSGM generates data several orders of magnitude faster than recent SGMs. Nevertheless, LSGM’s synthesis speed does not yet permit sampling at interactive rates, and our implementation of LSGM is currently limited to image generation. Therefore, future work includes further accelerating sampling, applying LSGMs to other data types, and designing efficient networks for LSGMs.

heading: 7 Broader Impact
text: Generating high-quality samples while fully covering the data distribution has been a long-standing challenge in generative learning. A solution to this problem will likely help reduce biases in generative models and lead to improving overall representation of minorities in the data distribution. SGMs are perhaps one of the first deep models that excel at both sample quality and distribution coverage. However, the high computational cost of sampling limits their widespread use. Our proposed LSGM reduces the sampling complexity of SGMs by a large margin and improves their expressivity further. Thus, in the long term, it can enable the usage of SGMs in practical applications.
Here, LSGM is examined on the image generation task which has potential benefits and risks discussed in [94, 95]. However, LSGM can be considered a generic framework that extends SGMs to non-continuous data types. In principle LSGM could be used to model, for example, language [96, 97], music [98, 10], or molecules [99, 100]. Furthermore, like other deep generative models, it can potentially be used also for non-generative tasks such as semi-supervised and representation learning [101, 102, 103]. This makes the long-term social impacts of LSGM dependent on the downstream applications.
Funding Statement All authors were funded by NVIDIA through full-time employment.

Review summaries:
The paper at hand proposes end-to-end training of score-based generative models in the latent space of a variational autoencoder. The VAE is pre-trained using a normal prior which is then replaced by the score-based generative model that is then jointly trained with the VAE. The paper also introduces a novel training objective based on the cross entropy between the encoder distribution and the SGM prior. Two techniques for variance reduction of the loss function are discussed. 

The paper proposes Latent Score-based Generative Model (LSGM) which introduces a score-based prior in the Variational Autoencoder (VAE) framework. In contrast to previous score-based models that operate in the data space, LSGM uses a score-based model in the latent space. A sample from a base distribution (standard Gaussian) is denoised using the score-based prior in the latent space and is then mapped to the data space using a decoder. The authors further discuss how the various terms in the ELBO can be computed to train the model in an end-to-end fashion without requiring the time-dependent marginal score function. Multiple variance reduction techniques and training tricks have also been proposed for the resulting objective function. Both quantitative and qualitative results demonstrate the ability of the model to generate high fidelity images.

The paper proposes a generative model based on the VAE framework with a sophisticated prior / inference scheme based on denoising score-matching. Authors develop extensive machinery to adapt related ideas from score-based generation in the observed space to latent space. This greatly improves sampling time while preserving and sometimes improving sample quality and/or likelihood of data. I believe the paper is an interesting addition to the existing collection of generative modelling techniques.

The paper proposes to learn a flexible VAE prior using score-based generative model. Contrary to the existing score-based model that builds on high-dim pixel space, the paper applies the score-based model directly on the latent space which is typically low-dimensional.  The effectiveness of the model has been verified using generative FID score, negative log-likelihood as well as various ablation studies.