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. 
abstractText: 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.
authors.name: Arash Vahdat
authors.name: Karsten Kreis
references.authors: Yang Song Jascha Sohl-Dickstein Diederik P Kingma Abhishek Kumar Stefano Ermon Ben Poole
references.title: Score-based generative modeling through stochastic differential equations
references.venue: In International Conference on Learning Representations,
references.authors: Yang Song Stefano Ermon
references.title: Generative Modeling by Estimating Gradients of the Data Distribution
references.venue: In Advances in Neural Information Processing Systems
references.authors: Yang Song Conor Durkan Iain Murray Stefano Ermon
references.title: Maximum likelihood training of score-based diffusion models
references.venue: arXiv e-prints,
references.authors: Jiaming Song Chenlin Meng Stefano Ermon
references.title: Denoising diffusion implicit models
references.venue: In International Conference on Learning Representations,
references.authors: Prafulla Dhariwal Alex Nichol
references.title: Diffusion models beat gans on image synthesis
references.venue: arXiv preprint arXiv:2105.05233,
references.authors: Nanxin Chen Yu Zhang Heiga Zen Ron J. Weiss Mohammad Norouzi William Chan
references.title: WaveGrad: Estimating Gradients for Waveform Generation
references.authors: Zhifeng Kong Wei Ping Jiaji Huang Kexin Zhao Bryan Catanzaro
references.title: DiffWave: A Versatile Diffusion Model for Audio Synthesis
references.authors: Myeonghun Jeong Hyeongju Kim Sung Jun Cheon Byoung Jin Choi Nam Soo Kim
references.title: Diff-tts: A denoising diffusion model for text-to-speech
references.authors: Gautam Mittal Jesse Engel Curtis Hawthorne Ian Simon
references.title: Symbolic music generation with diffusion models
references.venue: arXiv preprint arXiv:2103.16091,
references.authors: Chenhao Niu Yang Song Jiaming Song Shengjia Zhao Aditya Grover Stefano Ermon
references.title: Permutation invariant graph generation via score-based generative modeling
references.venue: In The 23rd International Conference on Artificial Intelligence and Statistics,
references.authors: Ruojin Cai Guandao Yang Hadar Averbuch-Elor Zekun Hao Serge J. Belongie Noah Snavely Bharath Hariharan
references.title: Learning gradient fields for shape generation
references.venue: In Computer Vision - ECCV 2020 - 16th European Conference,
references.authors: Shitong Luo Wei Hu
references.title: Diffusion probabilistic models for 3d point cloud generation
references.authors: Diederik P Kingma Max Welling
references.title: Auto-encoding variational bayes
references.venue: In The International Conference on Learning Representations (ICLR),
references.authors: Danilo Jimenez Rezende Shakir Mohamed Daan Wierstra
references.title: Stochastic backpropagation and approximate inference in deep generative models
references.venue: In International Conference on Machine Learning,
references.authors: Emilien Dupont Arnaud Doucet Yee Whye Teh
references.title: Augmented neural ODEs
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Zhifeng Kong Kamalika Chaudhuri
references.title: The expressive power of a class of normalizing flow models
references.venue: arXiv preprint arXiv:2006.00392,
references.authors: Chin-Wei Huang Laurent Dinh Aaron Courville
references.title: Augmented normalizing flows: Bridging the gap between generative flows and latent variable models
references.venue: arXiv preprint arXiv:2002.07101,
references.authors: Jianfei Chen Cheng Lu Biqi Chenli Jun Zhu Tian Tian
references.title: Vflow: More expressive generative flows with variational data augmentation
references.venue: arXiv preprint arXiv:2002.09741,
references.authors: Arash Vahdat Jan Kautz
references.title: NVAE: A Deep Hierarchical Variational Autoencoder
references.authors: Rewon Child
references.title: Very deep VAEs generalize autoregressive models and can outperform them on images
references.venue: In International Conference on Learning Representations,
references.authors: Pascal Vincent
references.title: A Connection between Score Matching and Denoising Autoencoders
references.venue: Neural Computation,
references.authors: Danilo Jimenez Rezende Shakir Mohamed
references.title: Variational inference with normalizing flows
references.venue: arXiv preprint arXiv:1505.05770,
references.authors: Diederik P Kingma Tim Salimans Rafal Jozefowicz Xi Chen Ilya Sutskever Max Welling
references.title: Improved variational inference with inverse autoregressive flow
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Chin-Wei Huang Jae Hyun Lim Aaron Courville
references.title: A variational perspective on diffusion-based generative models and score matching
references.venue: arXiv preprint arXiv:2106.02808,
references.authors: Alex Nichol Prafulla Dhariwal
references.title: Improved denoising diffusion probabilistic models
references.venue: arXiv preprint arXiv:2102.09672,
references.authors: Jascha Sohl-Dickstein Eric A. Weiss Niru Maheswaranathan Surya Ganguli
references.title: Deep Unsupervised Learning Using Nonequilibrium Thermodynamics
references.venue: In Proceedings of the 32nd International Conference on International Conference on Machine Learning - Volume 37,
references.authors: Aapo Hyvärinen
references.title: Estimation of Non-Normalized Statistical Models by Score Matching
references.venue: J. Mach. Learn. Res.,
references.authors: Siwei Lyu
references.title: Interpretation and Generalization of Score Matching
references.venue: In Proceedings of the Twenty-Fifth Conference on Uncertainty in Artificial Intelligence,
references.authors: Durk P Kingma Yann L. Cun
references.title: Regularized estimation of image statistics by Score Matching
references.venue: In Advances in Neural Information Processing Systems
references.authors: Yoshua Bengio Li Yao Guillaume Alain Pascal Vincent
references.title: Generalized Denoising Auto-Encoders as Generative Models
references.venue: In Proceedings of the 26th International Conference on Neural Information Processing Systems,
references.authors: Krzysztof J. Geras Charles A. Sutton
references.title: Scheduled denoising autoencoders
references.venue: In 3rd International Conference on Learning Representations,
references.authors: Yang Song Sahaj Garg Jiaxin Shi Stefano Ermon
references.title: Sliced Score Matching: A Scalable Approach to Density and Score Estimation
references.venue: In Proceedings of the Thirty-Fifth Conference on Uncertainty in Artificial Intelligence,
references.authors: Zengyi Li Yubei Chen Friedrich T. Sommer
references.title: Learning Energy-Based Models in High-Dimensional Spaces with Multi-scale Denoising
references.venue: Score Matching
references.authors: Tianyu Pang Kun Xu Chongxuan Li Yang Song Stefano Ermon Jun Zhu
references.title: Efficient Learning of Generative Models via Finite-Difference Score Matching
references.authors: Yang Song Stefano Ermon
references.title: Improved Techniques for Training Score-Based
references.venue: Generative Models
references.authors: Alexia Jolicoeur-Martineau Rémi Piché-Taillefer Ioannis Mitliagkas Remi Tachet des Combes
references.title: Adversarial score matching and improved sampling for image generation
references.venue: In International Conference on Learning Representations,
references.authors: Ruiqi Gao Yang Song Ben Poole Ying Nian Wu Diederik P Kingma
references.title: Learning energy-based models by diffusion recovery likelihood
references.venue: In International Conference on Learning Representations,
references.authors: Eric Luhman Troy Luhman
references.title: Knowledge distillation in iterative generative models for improved sampling speed
references.venue: arXiv preprint arXiv:2101.02388,
references.authors: Robin San-Roman Eliya Nachmani Lior Wolf
references.title: Noise estimation for generative diffusion models
references.venue: arXiv preprint arXiv:2104.02600,
references.authors: Martin Simonovsky Nikos Komodakis
references.title: Graphvae: Towards generation of small graphs using variational autoencoders
references.authors: Wengong Jin Regina Barzilay Tommi Jaakkola
references.title: Junction tree variational autoencoder for molecular graph generation
references.venue: In Proceedings of the 35th International Conference on Machine Learning,
references.authors: Aditya Grover Aaron Zweig Stefano Ermon
references.title: Graphite: Iterative generative modeling of graphs
references.venue: In Proceedings of the 36th International Conference on Machine Learning,
references.authors: Renjie Liao Yujia Li Yang Song Shenlong Wang Will Hamilton David K Duvenaud Raquel Urtasun Richard Zemel
references.title: Efficient graph generation with graph recurrent attention networks
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Xi Chen Diederik P Kingma Tim Salimans Yan Duan Prafulla Dhariwal John Schulman Ilya Sutskever Pieter Abbeel
references.title: Variational lossy autoencoder
references.venue: arXiv preprint arXiv:1611.02731,
references.authors: Lars Maaløe Marco Fraccaro Valentin Liévin Ole Winther
references.title: BIVA: A very deep hierarchy of latent variables for generative modeling
references.venue: In Advances in neural information processing systems,
references.authors: Jason Tyler Rolfe
references.title: Discrete variational autoencoders
references.venue: arXiv preprint arXiv:1609.02200,
references.authors: Arash Vahdat Evgeny Andriyash William G Macready
references.title: DVAE#: Discrete variational autoencoders with relaxed Boltzmann priors
references.venue: In Neural Information Processing Systems,
references.authors: Arash Vahdat William G. Macready Zhengbing Bian Amir Khoshaman Evgeny Andriyash
references.title: DVAE++: Discrete variational autoencoders with overlapping transformations
references.venue: In International Conference on Machine Learning (ICML),
references.authors: Arash Vahdat Evgeny Andriyash William G Macready
references.title: Undirected graphical models as approximate posteriors
references.venue: In International Conference on Machine Learning (ICML),
references.authors: Bo Pang Tian Han Erik Nijkamp Song-Chun Zhu Ying Nian Wu
references.title: Learning latent space energybased prior model
references.venue: Advances in Neural Information Processing Systems,
references.authors: Jesse Engel Matthew Hoffman Adam Roberts
references.title: Latent constraints: Learning to generate conditionally from unconditional generative models
references.venue: In International Conference on Learning Representations,
references.authors: Matthias Bauer Andriy Mnih
references.title: Resampled priors for variational autoencoders
references.venue: Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics,
references.authors: Jyoti Aneja Alexander Schwing Jan Kautz Arash Vahdat
references.title: NCP-VAE: Variational autoencoders with noise contrastive priors
references.venue: arXiv preprint arXiv:2010.02917,
references.authors: Hiroshi Takahashi Tomoharu Iwata Yuki Yamanaka Masanori Yamada Satoshi Yagi
references.title: Variational autoencoder with implicit optimal priors
references.venue: Proceedings of the AAAI Conference on Artificial Intelligence,
references.authors: Jakub Tomczak Max Welling
references.title: Vae with a vampprior
references.venue: In International Conference on Artificial Intelligence and Statistics,
references.authors: Bin Dai David Wipf
references.title: Diagnosing and enhancing VAE models
references.venue: In International Conference on Learning Representations,
references.authors: Aaron van den Oord Oriol Vinyals Koray Kavukcuoglu
references.title: Neural discrete representation learning
references.venue: arXiv preprint arXiv:1711.00937,
references.authors: Ali Razavi Aaron van den Oord Oriol Vinyals
references.title: Generating diverse high-fidelity images with vq-vae-2
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Partha Ghosh Mehdi S.M. Sajjadi Antonio Vergari Michael Black Bernhard
references.title: Scholkopf. From variational to deterministic autoencoders
references.venue: In International Conference on Learning Representations,
references.authors: Patrick Esser Robin Rombach Björn Ommer
references.title: Taming transformers for high-resolution image synthesis
references.venue: arXiv preprint arXiv:2012.09841,
references.authors: Abdul Fatir Ansari Ming Liang Ang Harold Soh
references.title: Refining deep generative models via discriminator gradient flow
references.venue: In International Conference on Learning Representations,
references.authors: Tong Che Ruixiang ZHANG Jascha Sohl-Dickstein Hugo Larochelle Liam Paull Yuan Cao Yoshua Bengio
references.title: Your gan is secretly an energy-based model and you should use discriminator driven latent sampling
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Akinori Tanaka
references.title: Discriminator optimal transport
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Weili Nie Arash Vahdat Anima Anandkumar
references.title: Controllable and compositional generation with latent-space energy-based models
references.venue: In Neural Information Processing Systems (NeurIPS),
references.authors: Antoine Wehenkel Gilles Louppe
references.title: Diffusion priors in variational autoencoders
references.venue: In ICML Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models,
references.authors: Abhishek Sinha Jiaming Song Chenlin Meng Stefano Ermon
references.title: D2c: Diffusion-denoising models for few-shot conditional generation
references.venue: arXiv preprint arXiv:2106.06819,
references.authors: Diederik P Kingma Tim Salimans Ben Poole Jonathan Ho
references.title: Variational diffusion models
references.venue: arXiv preprint arXiv:2107.00630,
references.authors: Dongjun Kim Seungjae Shin Kyungwoo Song Wanmo Kang Il-Chul Moon
references.title: Score matching model for unbounded data score
references.venue: arXiv preprint arXiv:2106.05527,
references.authors: Ricky T.Q. Chen Yulia Rubanova Jesse Bettencourt David Duvenaud
references.title: Neural ordinary differential equations
references.venue: Advances in Neural Information Processing Systems,
references.authors: Yuri Burda Roger Grosse Ruslan Salakhutdinov
references.title: Importance weighted autoencoders
references.venue: arXiv preprint arXiv:1509.00519,
references.authors: Martin Heusel Hubert Ramsauer Thomas Unterthiner Bernhard Nessler Sepp Hochreiter
references.title: Gans trained by a two time-scale update rule converge to a local nash equilibrium
references.venue: In Advances in neural information processing systems,
references.authors: Zhisheng Xiao Karsten Kreis Jan Kautz Arash Vahdat
references.title: VAEBM: A symbiosis between variational autoencoders and energy-based models
references.venue: In International Conference on Learning Representations,
references.authors: Gaurav Parmar Dacheng Li Kwonjoon Lee Zhuowen Tu
references.title: Dual contradistinctive generative autoencoder
references.venue: arXiv preprint arXiv:2011.10063,
references.authors: Heewoo Jun Rewon Child Mark Chen John Schulman Aditya Ramesh Alec Radford Ilya Sutskever
references.title: Distribution augmentation for generative modeling
references.venue: In International Conference on Machine Learning,
references.authors: Rewon Child Scott Gray Alec Radford Ilya Sutskever
references.title: Generating long sequences with sparse transformers
references.venue: arXiv preprint arXiv:1904.10509,
references.authors: Ali Razavi Aäron van den Oord Ben Poole Oriol Vinyals
references.title: Preventing posterior collapse with delta-vaes
references.venue: In The International Conference on Learning Representations (ICLR),
references.authors: XI Chen Nikhil Mishra Mostafa Rohaninejad Pieter Abbeel
references.title: PixelSNAIL: An improved autoregressive generative model
references.venue: In International Conference on Machine Learning,
references.authors: Tim Salimans Andrej Karpathy Xi Chen Diederik P Kingma
references.title: PixelCNN++: Improving the pixelCNN with discretized logistic mixture likelihood and other modifications
references.venue: arXiv preprint arXiv:1701.05517,
references.authors: Xinyu Gong Shiyu Chang Yifan Jiang Zhangyang Wang
references.title: Autogan: Neural architecture search for generative adversarial networks
references.venue: In The IEEE International Conference on Computer Vision (ICCV),
references.authors: Tero Karras Miika Aittala Janne Hellsten Samuli Laine Jaakko Lehtinen Timo Aila
references.title: Training generative adversarial networks with limited data
references.venue: In NeurIPS,
references.authors: Durk P Kingma Prafulla Dhariwal
references.title: Glow: Generative flow with invertible 1x1 convolutions
references.venue: Advances in Neural Information Processing Systems
references.authors: Jacob Menick Nal Kalchbrenner
references.title: Generating high fidelity images with subscale pixel networks and multidimensional upscaling
references.venue: In International Conference on Learning Representations,
references.authors: Stanislav Pidhorskyi Donald A Adjeroh Gianfranco Doretto
references.title: Adversarial latent autoencoders
references.venue: In Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR),
references.authors: Tero Karras Timo Aila Samuli Laine Jaakko Lehtinen
references.title: Progressive growing of GANs for improved quality, stability, and variation
references.venue: In International Conference on Learning Representations,
references.authors: Alex Krizhevsky
references.title: Learning multiple layers of features from tiny images
references.authors: Casper Kaae Sønderby Tapani Raiko Lars Maaløe Søren Kaae Sønderby Ole Winther
references.title: Ladder variational autoencoders
references.venue: In Advances in neural information processing systems,
references.authors: Hossein Sadeghi Evgeny Andriyash Walter Vinci Lorenzo Buffoni Mohammad H Amin
references.title: Pixelvae++: Improved pixelvae with discrete prior
references.authors: Xuezhe Ma Chunting Zhou Eduard Hovy
references.title: MAE: Mutual posterior-divergence regularization for variational autoencoders
references.venue: In The International Conference on Learning Representations (ICLR),
references.authors: Yann LeCun
references.title: The mnist database of handwritten digits. http://yann
references.venue: lecun. com/exdb/mnist/,
references.authors: J. Bailey
references.title: The tools of generative art, from flash to neural networks
references.venue: Art in America,
references.authors: Cristian Vaccari Andrew Chadwick
references.title: Deepfakes and disinformation: Exploring the impact of synthetic political video on deception, uncertainty, and trust in news
references.venue: Social Media+
references.authors: Samuel R Bowman Luke Vilnis Oriol Vinyals Andrew Dai Rafal Jozefowicz Samy Bengio
references.title: Generating sentences from a continuous space
references.venue: In Proceedings of The 20th SIGNLL Conference on Computational Natural Language Learning,
references.authors: Chunyuan Li Xiang Gao Yuan Li Xiujun Li Baolin Peng Yizhe Zhang Jianfeng Gao
references.title: Optimus: Organizing sentences via pre-trained modeling of a latent space
references.venue: In EMNLP,
references.authors: Prafulla Dhariwal Heewoo Jun Christine Payne Jong Wook Kim Alec Radford Ilya Sutskever
references.title: Jukebox: A generative model for music
references.venue: arXiv preprint arXiv:2005.00341,
references.authors: Benjamin Sanchez-Lengeling Alán Aspuru-Guzik
references.title: Inverse molecular design using machine learning: Generative models for matter
references.venue: engineering. Science,
references.authors: Zaccary Alperstein Artem Cherkasov Jason Tyler Rolfe
references.title: All smiles variational autoencoder
references.venue: arXiv preprint arXiv:1905.13343,
references.authors: Diederik P Kingma Shakir Mohamed Danilo Jimenez Rezende Max Welling
references.title: Semi-supervised learning with deep generative models
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Augustus Odena
references.title: Semi-supervised learning with generative adversarial networks
references.venue: arXiv preprint arXiv:1606.01583,
references.authors: Daiqing Li Junlin Yang Karsten Kreis Antonio Torralba Sanja Fidler
references.title: Semantic segmentation with generative models: Semi-supervised learning and strong out-of-domain generalization
references.venue: In Conference on Computer Vision and Pattern Recognition (CVPR),
references.authors: Simo Särkkä Arno Solin
references.title: Applied Stochastic Differential Equations. Institute of Mathematical Statistics Textbooks
references.authors: Brian D.O. Anderson
references.title: Reverse-time diffusion equation models
references.venue: Stochastic Process. Appl.,
references.authors: Will Grathwohl Ricky T.Q. Chen Jesse Bettencourt Ilya Sutskever David Duvenaud
references.title: Ffjord: Free-form continuous dynamics for scalable reversible generative models
references.venue: International Conference on Learning Representations,
references.authors: Richard Zhang
references.title: Making convolutional networks shift-invariant again
references.venue: In ICML,
references.authors: Diederik Kingma Jimmy Ba
references.title: Adam: A method for stochastic optimization
references.venue: arXiv preprint arXiv:1412.6980,
references.authors: J.R. Dormand P.J. Prince
references.title: A family of embedded Runge–Kutta formulae
references.venue: Journal of Computational and Applied Mathematics,
sections.heading: 1 Introduction
sections.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.
sections.heading: 2 Background
sections.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].
sections.heading: 3 Score-based Generative Modeling in Latent Space
sections.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.
sections.heading: 3.1 The Cross Entropy Term
sections.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].
sections.heading: 3.2 Mixing Normal and Neural Score Functions
sections.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.
sections.heading: 3.4 Variance Reduction
sections.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.
sections.heading: 4 Related Work
sections.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.
sections.heading: 5 Experiments
sections.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.
sections.heading: 5.1 Main Results
sections.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.
sections.heading: 5.2 Ablation Studies
sections.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.
sections.heading: 6 Conclusions
sections.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.
sections.heading: 7 Broader Impact
sections.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.
title: Score-based Generative Modeling in Latent Space