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