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.

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

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

summaries: 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:
Yang Song
Jascha Sohl-Dickstein
Diederik P Kingma
Abhishek Kumar
Stefano Ermon
Ben Poole

title: Score-based generative modeling through stochastic differential equations

authors:
Yang Song
Stefano Ermon

title: Generative Modeling by Estimating Gradients of the Data Distribution

venue: In Advances in Neural Information Processing Systems

authors:
Yang Song
Conor Durkan
Iain Murray
Stefano Ermon

title: Maximum likelihood training of score-based diffusion models

venue: arXiv e-prints,

authors:
Jiaming Song
Chenlin Meng
Stefano Ermon

title: Denoising diffusion implicit models

authors:
Prafulla Dhariwal
Alex Nichol

title: Diffusion models beat gans on image synthesis

venue: arXiv preprint arXiv:2105.05233,

authors:
Nanxin Chen
Yu Zhang
Heiga Zen
Ron J. Weiss
Mohammad Norouzi
William Chan

title: WaveGrad: Estimating Gradients for Waveform Generation

authors:
Zhifeng Kong
Wei Ping
Jiaji Huang
Kexin Zhao
Bryan Catanzaro

title: DiffWave: A Versatile Diffusion Model for Audio Synthesis

authors:
Myeonghun Jeong
Hyeongju Kim
Sung Jun Cheon
Byoung Jin Choi
Nam Soo Kim

title: Diff-tts: A denoising diffusion model for text-to-speech

authors:
Gautam Mittal
Jesse Engel
Curtis Hawthorne
Ian Simon

title: Symbolic music generation with diffusion models

venue: arXiv preprint arXiv:2103.16091,

authors:
Chenhao Niu
Yang Song
Jiaming Song
Shengjia Zhao
Aditya Grover
Stefano Ermon

title: Permutation invariant graph generation via score-based generative modeling

venue: In The 23rd International Conference on Artificial Intelligence and Statistics,

authors:
Ruojin Cai
Guandao Yang
Hadar Averbuch-Elor
Zekun Hao
Serge J. Belongie
Noah Snavely
Bharath Hariharan

title: Learning gradient fields for shape generation

venue: In Computer Vision - ECCV 2020 - 16th European Conference,

authors:
Shitong Luo
Wei Hu

title: Diffusion probabilistic models for 3d point cloud generation

authors:
Diederik P Kingma
Max Welling

title: Auto-encoding variational bayes

venue: In The International Conference on Learning Representations (ICLR),

authors:
Danilo Jimenez Rezende
Shakir Mohamed
Daan Wierstra

title: Stochastic backpropagation and approximate inference in deep generative models

authors:
Emilien Dupont
Arnaud Doucet
Yee Whye Teh

title: Augmented neural ODEs

authors:
Zhifeng Kong
Kamalika Chaudhuri

title: The expressive power of a class of normalizing flow models

venue: arXiv preprint arXiv:2006.00392,

authors:
Chin-Wei Huang
Laurent Dinh
Aaron Courville

title: Augmented normalizing flows: Bridging the gap between generative flows and latent variable models

venue: arXiv preprint arXiv:2002.07101,

authors:
Jianfei Chen
Cheng Lu
Biqi Chenli
Jun Zhu
Tian Tian

title: Vflow: More expressive generative flows with variational data augmentation

venue: arXiv preprint arXiv:2002.09741,

authors:
Arash Vahdat
Jan Kautz

title: NVAE: A Deep Hierarchical Variational Autoencoder

title: Very deep VAEs generalize autoregressive models and can outperform them on images

title: A Connection between Score Matching and Denoising Autoencoders

venue: Neural Computation,

authors:
Danilo Jimenez Rezende
Shakir Mohamed

title: Variational inference with normalizing flows

venue: arXiv preprint arXiv:1505.05770,

authors:
Diederik P Kingma
Tim Salimans
Rafal Jozefowicz
Xi Chen
Ilya Sutskever
Max Welling

title: Improved variational inference with inverse autoregressive flow

authors:
Chin-Wei Huang
Jae Hyun Lim
Aaron Courville

title: A variational perspective on diffusion-based generative models and score matching

venue: arXiv preprint arXiv:2106.02808,

authors:
Alex Nichol
Prafulla Dhariwal

title: Improved denoising diffusion probabilistic models

venue: arXiv preprint arXiv:2102.09672,

authors:
Jascha Sohl-Dickstein
Eric A. Weiss
Niru Maheswaranathan
Surya Ganguli

title: Deep Unsupervised Learning Using Nonequilibrium Thermodynamics

venue: In Proceedings of the 32nd International Conference on International Conference on Machine Learning - Volume 37,

title: Estimation of Non-Normalized Statistical Models by Score Matching

venue: J. Mach. Learn. Res.,

title: Interpretation and Generalization of Score Matching

venue: In Proceedings of the Twenty-Fifth Conference on Uncertainty in Artificial Intelligence,

authors:
Durk P Kingma
Yann L. Cun

title: Regularized estimation of image statistics by Score Matching

authors:
Yoshua Bengio
Li Yao
Guillaume Alain
Pascal Vincent

title: Generalized Denoising Auto-Encoders as Generative Models

venue: In Proceedings of the 26th International Conference on Neural Information Processing Systems,

authors:
Krzysztof J. Geras
Charles A. Sutton

title: Scheduled denoising autoencoders

venue: In 3rd International Conference on Learning Representations,

authors:
Yang Song
Sahaj Garg
Jiaxin Shi
Stefano Ermon

title: Sliced Score Matching: A Scalable Approach to Density and Score Estimation

venue: In Proceedings of the Thirty-Fifth Conference on Uncertainty in Artificial Intelligence,

authors:
Zengyi Li
Yubei Chen
Friedrich T. Sommer

title: Learning Energy-Based Models in High-Dimensional Spaces with Multi-scale Denoising

venue: Score Matching

authors:
Tianyu Pang
Kun Xu
Chongxuan Li
Yang Song
Stefano Ermon
Jun Zhu

title: Efficient Learning of Generative Models via Finite-Difference Score Matching

title: Improved Techniques for Training Score-Based

venue: Generative Models

authors:
Alexia Jolicoeur-Martineau
Rémi Piché-Taillefer
Ioannis Mitliagkas
Remi Tachet des Combes

title: Adversarial score matching and improved sampling for image generation

authors:
Ruiqi Gao
Yang Song
Ben Poole
Ying Nian Wu
Diederik P Kingma

title: Learning energy-based models by diffusion recovery likelihood

authors:
Eric Luhman
Troy Luhman

title: Knowledge distillation in iterative generative models for improved sampling speed

venue: arXiv preprint arXiv:2101.02388,

authors:
Robin San-Roman
Eliya Nachmani
Lior Wolf

title: Noise estimation for generative diffusion models

venue: arXiv preprint arXiv:2104.02600,

authors:
Martin Simonovsky
Nikos Komodakis

title: Graphvae: Towards generation of small graphs using variational autoencoders

authors:
Wengong Jin
Regina Barzilay
Tommi Jaakkola

title: Junction tree variational autoencoder for molecular graph generation

venue: In Proceedings of the 35th International Conference on Machine Learning,

authors:
Aditya Grover
Aaron Zweig
Stefano Ermon

title: Graphite: Iterative generative modeling of graphs

venue: In Proceedings of the 36th International Conference on Machine Learning,

authors:
Renjie Liao
Yujia Li
Yang Song
Shenlong Wang
Will Hamilton
David K Duvenaud
Raquel Urtasun
Richard Zemel

title: Efficient graph generation with graph recurrent attention networks

authors:
Xi Chen
Diederik P Kingma
Tim Salimans
Yan Duan
Prafulla Dhariwal
John Schulman
Ilya Sutskever
Pieter Abbeel

title: Variational lossy autoencoder

venue: arXiv preprint arXiv:1611.02731,

authors:
Lars Maaløe
Marco Fraccaro
Valentin Liévin
Ole Winther

title: BIVA: A very deep hierarchy of latent variables for generative modeling

venue: In Advances in neural information processing systems,

title: Discrete variational autoencoders

venue: arXiv preprint arXiv:1609.02200,

authors:
Arash Vahdat
Evgeny Andriyash
William G Macready

title: DVAE#: Discrete variational autoencoders with relaxed Boltzmann priors

venue: In Neural Information Processing Systems,

authors:
Arash Vahdat
William G. Macready
Zhengbing Bian
Amir Khoshaman
Evgeny Andriyash

title: DVAE++: Discrete variational autoencoders with overlapping transformations

venue: In International Conference on Machine Learning (ICML),

title: Undirected graphical models as approximate posteriors

authors:
Bo Pang
Tian Han
Erik Nijkamp
Song-Chun Zhu
Ying Nian Wu

title: Learning latent space energybased prior model

venue: Advances in Neural Information Processing Systems,

authors:
Jesse Engel
Matthew Hoffman
Adam Roberts

title: Latent constraints: Learning to generate conditionally from unconditional generative models

authors:
Matthias Bauer
Andriy Mnih

title: Resampled priors for variational autoencoders

venue: Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics,

authors:
Jyoti Aneja
Alexander Schwing
Jan Kautz
Arash Vahdat

title: NCP-VAE: Variational autoencoders with noise contrastive priors

venue: arXiv preprint arXiv:2010.02917,

authors:
Hiroshi Takahashi
Tomoharu Iwata
Yuki Yamanaka
Masanori Yamada
Satoshi Yagi

title: Variational autoencoder with implicit optimal priors

venue: Proceedings of the AAAI Conference on Artificial Intelligence,

authors:
Jakub Tomczak
Max Welling

title: Vae with a vampprior

venue: In International Conference on Artificial Intelligence and Statistics,

authors:
Bin Dai
David Wipf

title: Diagnosing and enhancing VAE models

authors:
Aaron van den Oord
Oriol Vinyals
Koray Kavukcuoglu

title: Neural discrete representation learning

venue: arXiv preprint arXiv:1711.00937,

authors:
Ali Razavi
Aaron van den Oord
Oriol Vinyals

title: Generating diverse high-fidelity images with vq-vae-2

authors:
Partha Ghosh
Mehdi S.M. Sajjadi
Antonio Vergari
Michael Black
Bernhard

title: Scholkopf. From variational to deterministic autoencoders

authors:
Patrick Esser
Robin Rombach
Björn Ommer

title: Taming transformers for high-resolution image synthesis

venue: arXiv preprint arXiv:2012.09841,

authors:
Abdul Fatir Ansari
Ming Liang Ang
Harold Soh

title: Refining deep generative models via discriminator gradient flow

authors:
Tong Che
Ruixiang ZHANG
Jascha Sohl-Dickstein
Hugo Larochelle
Liam Paull
Yuan Cao
Yoshua Bengio

title: Your gan is secretly an energy-based model and you should use discriminator driven latent sampling

title: Discriminator optimal transport

authors:
Weili Nie
Arash Vahdat
Anima Anandkumar

title: Controllable and compositional generation with latent-space energy-based models

venue: In Neural Information Processing Systems (NeurIPS),

authors:
Antoine Wehenkel
Gilles Louppe

title: Diffusion priors in variational autoencoders

venue: In ICML Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models,

authors:
Abhishek Sinha
Jiaming Song
Chenlin Meng
Stefano Ermon

title: D2c: Diffusion-denoising models for few-shot conditional generation

venue: arXiv preprint arXiv:2106.06819,

authors:
Diederik P Kingma
Tim Salimans
Ben Poole
Jonathan Ho

title: Variational diffusion models

venue: arXiv preprint arXiv:2107.00630,

authors:
Dongjun Kim
Seungjae Shin
Kyungwoo Song
Wanmo Kang
Il-Chul Moon

title: Score matching model for unbounded data score

venue: arXiv preprint arXiv:2106.05527,

authors:
Ricky T.Q. Chen
Yulia Rubanova
Jesse Bettencourt
David Duvenaud

title: Neural ordinary differential equations

authors:
Yuri Burda
Roger Grosse
Ruslan Salakhutdinov

title: Importance weighted autoencoders

venue: arXiv preprint arXiv:1509.00519,

authors:
Martin Heusel
Hubert Ramsauer
Thomas Unterthiner
Bernhard Nessler
Sepp Hochreiter

title: Gans trained by a two time-scale update rule converge to a local nash equilibrium

authors:
Zhisheng Xiao
Karsten Kreis
Jan Kautz
Arash Vahdat

title: VAEBM: A symbiosis between variational autoencoders and energy-based models

authors:
Gaurav Parmar
Dacheng Li
Kwonjoon Lee
Zhuowen Tu

title: Dual contradistinctive generative autoencoder

venue: arXiv preprint arXiv:2011.10063,

authors:
Heewoo Jun
Rewon Child
Mark Chen
John Schulman
Aditya Ramesh
Alec Radford
Ilya Sutskever

title: Distribution augmentation for generative modeling

authors:
Rewon Child
Scott Gray
Alec Radford
Ilya Sutskever

title: Generating long sequences with sparse transformers

venue: arXiv preprint arXiv:1904.10509,

authors:
Ali Razavi
Aäron van den Oord
Ben Poole
Oriol Vinyals

title: Preventing posterior collapse with delta-vaes

authors:
XI Chen
Nikhil Mishra
Mostafa Rohaninejad
Pieter Abbeel

title: PixelSNAIL: An improved autoregressive generative model

authors:
Tim Salimans
Andrej Karpathy
Xi Chen
Diederik P Kingma

title: PixelCNN++: Improving the pixelCNN with discretized logistic mixture likelihood and other modifications

venue: arXiv preprint arXiv:1701.05517,

authors:
Xinyu Gong
Shiyu Chang
Yifan Jiang
Zhangyang Wang

title: Autogan: Neural architecture search for generative adversarial networks

venue: In The IEEE International Conference on Computer Vision (ICCV),

authors:
Tero Karras
Miika Aittala
Janne Hellsten
Samuli Laine
Jaakko Lehtinen
Timo Aila

title: Training generative adversarial networks with limited data

venue: In NeurIPS,

authors:
Durk P Kingma
Prafulla Dhariwal

title: Glow: Generative flow with invertible 1x1 convolutions

venue: Advances in Neural Information Processing Systems

authors:
Jacob Menick
Nal Kalchbrenner

title: Generating high fidelity images with subscale pixel networks and multidimensional upscaling

authors:
Stanislav Pidhorskyi
Donald A Adjeroh
Gianfranco Doretto

title: Adversarial latent autoencoders

venue: In Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR),

authors:
Tero Karras
Timo Aila
Samuli Laine
Jaakko Lehtinen

title: Progressive growing of GANs for improved quality, stability, and variation

title: Learning multiple layers of features from tiny images

authors:
Casper Kaae Sønderby
Tapani Raiko
Lars Maaløe
Søren Kaae Sønderby
Ole Winther

title: Ladder variational autoencoders

authors:
Hossein Sadeghi
Evgeny Andriyash
Walter Vinci
Lorenzo Buffoni
Mohammad H Amin

title: Pixelvae++: Improved pixelvae with discrete prior

authors:
Xuezhe Ma
Chunting Zhou
Eduard Hovy

title: MAE: Mutual posterior-divergence regularization for variational autoencoders

title: The mnist database of handwritten digits. http://yann

venue: lecun. com/exdb/mnist/,

title: The tools of generative art, from flash to neural networks

venue: Art in America,

authors:
Cristian Vaccari
Andrew Chadwick

title: Deepfakes and disinformation: Exploring the impact of synthetic political video on deception, uncertainty, and trust in news

venue: Social Media+

authors:
Samuel R Bowman
Luke Vilnis
Oriol Vinyals
Andrew Dai
Rafal Jozefowicz
Samy Bengio

title: Generating sentences from a continuous space

venue: In Proceedings of The 20th SIGNLL Conference on Computational Natural Language Learning,

authors:
Chunyuan Li
Xiang Gao
Yuan Li
Xiujun Li
Baolin Peng
Yizhe Zhang
Jianfeng Gao

title: Optimus: Organizing sentences via pre-trained modeling of a latent space

venue: In EMNLP,

authors:
Prafulla Dhariwal
Heewoo Jun
Christine Payne
Jong Wook Kim
Alec Radford
Ilya Sutskever

title: Jukebox: A generative model for music

venue: arXiv preprint arXiv:2005.00341,

authors:
Benjamin Sanchez-Lengeling
Alán Aspuru-Guzik

title: Inverse molecular design using machine learning: Generative models for matter

venue: engineering. Science,

authors:
Zaccary Alperstein
Artem Cherkasov
Jason Tyler Rolfe

title: All smiles variational autoencoder

venue: arXiv preprint arXiv:1905.13343,

authors:
Diederik P Kingma
Shakir Mohamed
Danilo Jimenez Rezende
Max Welling

title: Semi-supervised learning with deep generative models

title: Semi-supervised learning with generative adversarial networks

venue: arXiv preprint arXiv:1606.01583,

authors:
Daiqing Li
Junlin Yang
Karsten Kreis
Antonio Torralba
Sanja Fidler

title: Semantic segmentation with generative models: Semi-supervised learning and strong out-of-domain generalization

venue: In Conference on Computer Vision and Pattern Recognition (CVPR),

authors:
Simo Särkkä
Arno Solin

title: Applied Stochastic Differential Equations. Institute of Mathematical Statistics Textbooks

title: Reverse-time diffusion equation models

venue: Stochastic Process. Appl.,

authors:
Will Grathwohl
Ricky T.Q. Chen
Jesse Bettencourt
Ilya Sutskever
David Duvenaud

title: Ffjord: Free-form continuous dynamics for scalable reversible generative models

venue: International Conference on Learning Representations,

title: Making convolutional networks shift-invariant again

venue: In ICML,

authors:
Diederik Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

venue: arXiv preprint arXiv:1412.6980,

authors:
J.R. Dormand
P.J. Prince

title: A family of embedded Runge–Kutta formulae

venue: Journal of Computational and Applied Mathematics,

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.

title: Score-based Generative Modeling in Latent Space