Score-based Generative Modeling in Latent Space

1 Introduction . 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 can not 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 can not 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 . 2 Background . 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 ] . 3 Score-based Generative Modeling in Latent Space . 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 .

Summary: 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.