* [11] Brown R.B., Gray A., Vector cross products, Comment. Math. Helv. 42 (1967), 222–236.
  * [12] Bryant R.L., Some remarks on ${\rm G}_{2}$-structures, in Proceedings of Gökova Geometry-Topology Conference 2005, Editors S. Akbulut, T. Onder, R.J. Stern, Gökova Geometry/Topology Conference (GGT), Gökova, 2006, 75–109, math.DG/0305124.
  * [13] Bryant R.L., Hsu L., Rigidity of integral curves of rank $2$ distributions, Invent. Math. 114 (1993), 435–461.
  * [14] Calderbank D.M.J., Diemer T., Differential invariants and curved Bernstein–Gelfand–Gelfand sequences, J. Reine Angew. Math. 537 (2001), 67–103, math.DG/0001158.
  * [15] Čap A., Infinitesimal automorphisms and deformations of parabolic geometries, J. Eur. Math. Soc. 10 (2008), 415–437, math.DG/0508535.
  * [16] Čap A., Gover A.R., A holonomy characterisation of Fefferman spaces, Ann. Global Anal. Geom. 38 (2010), 399–412, math.DG/0611939.
  * [17] Čap A., Gover A.R., Hammerl M., Holonomy reductions of Cartan geometries and curved orbit decompositions, Duke Math. J. 163 (2014), 1035–1070, arXiv:1103.4497.
  * [18] Čap A., Slovák J., Parabolic geometries. I. Background and general theory, Mathematical Surveys and Monographs, Vol. 154, Amer. Math. Soc., Providence, RI, 2009.
  * [19] Čap A., Slovák J., Souček V., Bernstein–Gelfand–Gelfand sequences, Ann. of Math. 154 (2001), 97–113, math.DG/0001164.
  * [20] Cartan É., Sur la structure des groupes simples finis et continus, C. R. Acad. Sci. Paris 113 (1893), 784–786.
  * [21] Cartan É., Les systèmes de Pfaff, à cinq variables et les équations aux dérivées partielles du second ordre, Ann. Sci. École Norm. Sup. (3) 27 (1910), 109–192.
  * [22] Chudecki A., On some examples of para-Hermite and para-Kähler Einstein spaces with $\Lambda\neq 0$, J. Geom. Phys. 112 (2017), 175–196, arXiv:1602.02913.
  * [23] Cortés V., Leistner T., Schäfer L., Schulte-Hengesbach F., Half-flat structures and special holonomy, Proc. Lond. Math. Soc. 102 (2011), 113–158, arXiv:0907.1222.
  * [24] Curry S., Gover A.R., An introduction to conformal geometry and tractor calculus, with a view to applications in general relativity, arXiv:1412.7559.
  * [25] Doubrov B., Komrakov B., The geometry of second-order ordinary differential equations, arXiv:1602.00913.
  * [26] Doubrov B., Kruglikov B., On the models of submaximal symmetric rank 2 distributions in 5D, Differential Geom. Appl. 35 (2014), suppl., 314–322, arXiv:1311.7057.
  * [27] Dunajski M., Przanowski M., Null Kähler structures, symmetries and integrability, in Topics in Mathematical Physics, General Relativity and Cosmology in Honor of Jerzy Plebański, Editors H. García-Compeán, B. Mielnik, M. Montesinos, M. Przanowski, World Sci. Publ., Hackensack, NJ, 2006, 147–155, gr-qc/0310005.
  * [28] Engel F., Sur un groupe simple à quatorze paramètres, C. R. Acad. Sci. Paris 116 (1893), 786–788.
  * [29] Fefferman C., Graham C.R., The ambient metric, Annals of Mathematics Studies, Vol. 178, Princeton University Press, Princeton, NJ, 2012, arXiv:0710.0919.
  * [30] Fefferman C.L., Monge–Ampère equations, the Bergman kernel, and geometry of pseudoconvex domains, Ann. of Math. 103 (1976), 395–416, correction, Ann. of Math. 104 (1976), 393–394.
  * [31] Fox D.J.F., Contact projective structures, Indiana Univ. Math. J. 54 (2005), 1547–1598, math.DG/0402332.
  * [32] Gover A.R., Almost conformally Einstein manifolds and obstructions, in Differential Geometry and its Applications, Matfyzpress, Prague, 2005, 247–260, math.DG/0412393.
  * [33] Gover A.R., Almost Einstein and Poincaré–Einstein manifolds in Riemannian signature, J. Geom. Phys. 60 (2010), 182–204, arXiv:0803.3510.
  * [34] Gover A.R., Macbeth H.R., Detecting Einstein geodesics: Einstein metrics in projective and conformal geometry, Differential Geom. Appl. 33 (2014), suppl., 44–69, arXiv:1212.6286.
  * [35] Gover A.R., Neusser K., Willse T., Sasaki and Kähler structures and their compactifications via projective geometry, in preparation.
  * [36] Gover A.R., Nurowski P., Obstructions to conformally Einstein metrics in $n$ dimensions, J. Geom. Phys. 56 (2006), 450–484, math.DG/0405304.
  * [37] Gover A.R., Panai R., Willse T., Nearly Kähler geometry and $(2,3,5)$-distributions via projective holonomy, Indiana Univ. Math. J., to appear, arXiv:1403.1959.
  * [38] Graham C.R., Willse T., Parallel tractor extension and ambient metrics of holonomy split ${\rm G}_{2}$, J. Differential Geom. 92 (2012), 463–505, arXiv:1109.3504.
  * [39] Haantjes J., Wrona W., Über konformeuklidische und Einsteinsche Räume gerader Dimension, Proc. Nederl. Akad. Wetensch. 42 (1939), 626–636.
  * [40] Hammerl M., Natural prolongations of BGG-operators, Ph.D. Thesis, Universität Wien, 2009.
  * [41] Hammerl M., Sagerschnig K., Conformal structures associated to generic rank 2 distributions on 5-manifolds – characterization and Killing-field decomposition, SIGMA 5 (2009), 081, 29 pages, arXiv:0908.0483.
  * [42] Hammerl M., Sagerschnig K., The twistor spinors of generic 2- and 3-distributions, Ann. Global Anal. Geom. 39 (2011), 403–425, arXiv:1004.3632.
  * [43] Hammerl M., Sagerschnig K., Šilhan J., Taghavi-Chabert A., Žádník V., A projective-to-conformal Fefferman-type construction, arXiv:1510.03337.
  * [44] Kath I., Killing spinors on pseudo-Riemannian manifolds, Humboldt-Universität Berlin, Habilitationsschrift, 1999.
  * [45] Kath I., $G_{2(2)}^{*}$-structures on pseudo-Riemannian manifolds, J. Geom. Phys. 27 (1998), 155–177.
  * [46] Kolář I., Michor P.W., Slovák J., Natural operations in differential geometry, Springer-Verlag, Berlin, 1993.
  * [47] Kozameh C.N., Newman E.T., Tod K.P., Conformal Einstein spaces, Gen. Relativity Gravitation 17 (1985), 343–352.
  * [48] Kruglikov B., The D., The gap phenomenon in parabolic geometries, J. Reine Angew. Math., to appear, arXiv:1303.1307.
  * [49] Leistner T., Nurowski P., Ambient metrics with exceptional holonomy, Ann. Sc. Norm. Super. Pisa Cl. Sci. (5) 11 (2012), 407–436, arXiv:0904.0186.
  * [50] Leitner F., Conformal Killing forms with normalisation condition, Rend. Circ. Mat. Palermo (2) Suppl. (2005), 279–292, math.DG/0406316.
  * [51] Leitner F., A remark on unitary conformal holonomy, in Symmetries and Overdetermined Systems of Partial Differential Equations, IMA Vol. Math. Appl., Vol. 144, Editors M.G. Eastwood, W. Miller, Springer, New York, 2008, 445–460, math.DG/0604393.
  * [52] Nurowski P., Differential equations and conformal structures, J. Geom. Phys. 55 (2005), 19–49, math.DG/0406400.
  * [53] Nurowski P., Sparling G.A., Three-dimensional Cauchy–Riemann structures and second-order ordinary differential equations, Classical Quantum Gravity 20 (2003), 4995–5016, math.DG/0306331.
  * [54] Sagerschnig K., Split octonions and generic rank two distributions in dimension five, Arch. Math. (Brno) 42 (2006), suppl., 329–339.
  * [55] Sagerschnig K., Willse T., The almost Einstein operator for $(2,3,5)$ distributions, in preparation.
  * [56] Sagerschnig K., Willse T., Fefferman constructions for Kähler surfaces, in preparation.
  * [57] Sasaki S., On a relation between a Riemannian space which is conformal with Einstein spaces and normal conformally connected spaces whose groups of holonomy fix a point or a hypersphere, Tensor 5 (1942), 66–72.
  * [58] Schouten J.A., Ricci-calculus. An introduction to tensor analysis and its geometrical applications, Die Grundlehren der Mathematischen Wissenschaften, Vol. 10, 2nd ed., Springer-Verlag, Berlin – Göttingen – Heidelberg, 1954.
  * [59] Schulte-Hengesbach F., Half-flat structure on Lie groups, Ph.D. Thesis, Universität Hamburg, 2010.
  * [60] Semmelmann U., Conformal Killing forms on Riemannian manifolds, Math. Z. 245 (2003), 503–527, math.DG/0206117.
  * [61] Sharpe R.W., Differential geometry. Cartan’s generalization of Klein’s Erlangen program, Graduate Texts in Mathematics, Vol. 166, Springer-Verlag, New York, 1997.
  * [62] Susskind L., The world as a hologram, J. Math. Phys. 36 (1995), 6377–6396, hep-th/9409089.
  * [63] Szekeres P., Spaces conformal to a class of spaces in general relativity, Proc. Roy. Soc. Ser. A 274 (1963), 206–212.
  * [64] Willse T., Highly symmetric 2-plane fields on 5-manifolds and 5-dimensional Heisenberg group holonomy, Differential Geom. Appl. 33 (2014), suppl., 81–111, arXiv:1302.7163.
  * [65] Wolf J.A., Isotropic manifolds of indefinite metric, Comment. Math. Helv. 39 (1964), 21–64.
  * [66] Wong Y.-C., Some Einstein spaces with conformally separable fundamental tensors, Trans. Amer. Math. Soc. 53 (1943), 157–194.
  * [67] Yano K., Conformal and concircular geometries in Einstein spaces, Proc. Imp. Acad. Tokyo 19 (1943), 444–453.


# High Fidelity Image Synthesis With Deep VAEs In Latent Space

Troy Luhman  Eric Luhman

###### Abstract

We present fast, realistic image generation on high-resolution, multimodal
datasets using hierarchical variational autoencoders (VAEs) trained on a
deterministic autoencoder’s latent space. In this two-stage setup, the
autoencoder compresses the image into its semantic features, which are then
modeled with a deep VAE. With this method, the VAE avoids modeling the fine-
grained details that constitute the majority of the image’s code length,
allowing it to focus on learning its structural components. We demonstrate the
effectiveness of our two-stage approach, achieving a FID of 9.34 on the
ImageNet-256 dataset which is comparable to BigGAN. We make our implementation
available online111 Torch code: https://github.com/ericl122333/latent-vae, Jax
code: https://github.com/ericl122333/latent-vae-jax.

##  1 Introduction

There has recently been a surge of interest in deep generative modeling for
image generation, which has made great progress in the last few years. In
particular, advances in diffusion and autoregressive models have enabled them
to create photo-realistic images for even the most complex object compositions
(Ramesh et al., 2022; Rombach et al., 2022; Yu et al., 2022). Generative
models like these will likely have significant impact on many creative
applications as they become continually adopted and improved.

Despite their impressive capabilities, both these model classes have a serious
drawback which is their slow sampling speed. Producing an image generally
requires evaluating a large neural network 50-1000 times since each transition
in their generative process reuses the function. This inefficiency is an
important consideration in practical settings, and can pose a challenge to
widespread deployment of the model.

Variational Autoencoders (VAEs) are a class of likelihood-based models that
use learned latent variables to generate data (Kingma & Welling, 2013; Rezende
et al., 2014). Hierarchical VAEs separate latents into multiple conditionally
dependent groups for greater expressitivity (Sønderby et al., 2016; Kingma et
al., 2016). Notable examples of hierarchical VAEs for images include NVAE
(Vahdat & Kautz, 2020) and VDVAE (Child, 2020). Unlike diffusion and
autoregressive models which use a fixed analytic posterior, deep VAE
posteriors are trainable and parameterized by neural networks. Their learned
ordering of latent variables results in an efficient generation process that
only uses one function evaluation. Beyond their fast sampling speed,
hierarchical VAEs are also desirable due to their easy inference, stable
training, and interpretable latent representations.

However, to date hierarchical VAEs have not suceeded in generating convincing
images on large, multimodal datasets like ImageNet (Deng et al., 2009). This
is especially surprising given that their hierarchical generation process
seems naturally suited for image generation. Indeed, autoregressive models
have proven far more successful despite a counter-intuitive inductive bias
that simply generates images in raster-scan order. Thus, we start by observing
ingredients that make autoregressive image models so successful and ask
whether similar techniques can be applied to VAEs.

Figure 1: Selected samples from our guided ImageNet-256 model (FID=9.34).

A crucial component to the success of autoregressive models is training on a
sequence of compressed image tokens instead of direct pixel values (Razavi et
al., 2019a; Esser et al., 2021). Doing so enables them to focus on learning
the relationships between image semantics while ignoring imperceptible image
details, which is especially important because imperceptible details make up a
great majority of an image’s code length (Pope et al., 2021; Ho et al., 2020;
Rombach et al., 2022). We hypothesize that, like pixel-space autoregressive
models, existing pixel-space hierarchical VAEs focus primarily on learning
fine-grained features, weakening their ability to learn the underlying
composition of image concepts.

Motivated by this theory, we propose to improve deep VAEs by applying them to
the latent space of a deterministic autoencoder (DAE). Our pipeline consists
of two parts trained in succession: a DAE that reconstruct images from low-
dimensional latents, and a VAE that forms a generative model of these latents.
Compared to acting directly on pixel-space, training the VAE on low-
dimensional latents has two major advantages. Firstly, the code length of the
compressed latent is much smaller than its RGB representation, but it
preserves nearly all the image’s perceptual information. Having a smaller
total code length is helpful because it makes global features, which make up
only a few bits, more strongly emphasized. Imperceptible details which would
otherwise be modeled by the VAE are now discarded, allowing the VAE to focus
solely on image structure. Secondly, the latent variable has greatly reduced
dimensionality, which reduces computational costs and allows for larger models
to be trained with the same resources.

We evaluate our method on several image benchmarks, and obtain strong
performance that greatly exceeds previous hierarchical VAE capabilities. When
combining the proposed two-stage method with classifier-free guidance on
ImageNet-256, we obtain a FID score (Heusel et al., 2017) score of 9.34 that
is competitive landmark diffusion and GAN baselines (Brock et al., 2018;
Dhariwal & Nichol, 2021). Additionally, we experiment with VAE’s latent
representation to showcase their good interpretability and flexibility.

##  2 Related Work

Hierarchical VAEs. Our work builds off previous literature in VAEs (Kingma &
Welling, 2013; Rezende et al., 2014), and their hierarchical extensions
(Sønderby et al., 2016; Kingma et al., 2016; Maaløe et al., 2019; Vahdat &
Kautz, 2020). In these hierarchical VAEs, early latent groups start at low
dimensionality and increase in dimensionality over the course of generation.
For image data, this progression from low to high resolution naturally leads
to a coarse to fine generation process. These hierarchical VAEs benefit from
having many latent groups, where (Child, 2020) showed increasing statistical
depth substantially improves likelihood. A downside of higher depth is the
increased compute and memory costs, which can impede scaling to high-
resolution data. Our method addresses this by using a shallow DAE for high-
frequency relationships, while learning complex semantic relationships with a
deep VAE.

Compared to other likelihood-based models, VAEs are distinguished by their
powerful posteriors parameterized by neural networks. This enables them to
learn an efficient ordering of latent variables and results in fast sampling
speed. In contrast, diffusion and autoregressive models train a powerful prior
to match a simple, hand-crafted posterior. The neural network parameterizing
the prior shares parameters across latent layers, and its repeated evaluation
at sampling time leads to slow generation. Normalizing flows (Rezende &
Mohamed, 2015; Dinh et al., 2016) are a likelihood-based model with fast
sampling speed, but the need for explicit invertibility weakens their modeling
capabilities in practice.

Two-stage Image Generation. Instead of directly generating high-resolution
images, many works have found dividing the generation process into two stages
is beneficial for both efficiency and performance (Razavi et al., 2019b; Esser
et al., 2021; Ho et al., 2022; Rombach et al., 2022). One leading method is to
train an autoregressive model on discrete image tokens produced by an encoder
neural network (Razavi et al., 2019a; Esser et al., 2021; Yu et al., 2021);
this sequence-like tokenization is a helpful inductive bias for the
transformer architecture these models use. VQ-VAEs train a discrete VAE for
the first-stage encoder, and VQ-GANs achieve sharper images by adding
auxiliary perceptual and adversarial losses to the reconstruction objective
(Esser et al., 2021). (Rombach et al., 2022) also trains an encoder with the
same loss objectives, but it outputs continuous values and trains a diffusion
model over these latents. Our work adapts the pretrained autoencoders from
(Rombach et al., 2022) since our VAEs act in continuous space.

An alternative two-stage approach uses an upsampling stack which first creates
low-resolution RGB images before upsampling them to a higher resolution
(Saharia et al., 2022c; Ho et al., 2022). This method works very well for
diffusion models (Dhariwal & Nichol, 2021; Ramesh et al., 2022) as diffusion
models have good super-resolution capabilities (Li et al., 2022; Saharia et
al., 2022a). In our hierarchical VAE experiments, we found that upsampling
stack is outperformed by a latent-space approach (Section 4.3).

Our work is perhaps most similar to (Dai & Wipf, 2019), who propose a VAE on
the latent space of another VAE learned directly on pixels. This work is
similarly motivated to ours, accounting for the fact that the ambient
dimension is much greater than the manifold dimension. However, samples from
their model are blurry since their first stage is a regular VAE trained with a
pixel-space distortion term. In contrast, we learn the underlying manifold
using a DAE augmented with perceptual and aversarial losses, and achieve much
higher compression rates with little to no blurriness. Furthermore, our second
stage is based off hierarchical VAEs instead of single-group ones to reap the
benefits of high stochastic depth observed in (Child, 2020).

Perceptual and Adversarial Losses in VAEs. Previous works (Larsen et al.,
2016; Hou et al., 2017) have found that the pixel-wise reconstruction
objective in VAEs encourages them to produce blurry, unrealistic samples. As a
result, they modify the reconstruction objective to include a perceptual or
adversarial loss, which encourages the reconstructions to be visually similar
to real images in the feature space of a deep CNN. While both our method and
theirs include perceptual or adversarial losses in the reconstruction
objective, our adversarial component is trained separately from the MLE-based
VAE component, whereas these works combine the two objectives in a joint
training scheme.

##  3 Method

###  3.1 Perceptual Compression Stage

In a two-stage image generation approach, the first step often involves
training an autoencoder that compresses high-resolution images onto a lower-
dimensional latent space. Specifically, a neural network encoder takes in an
image $\mathbb{x}\in\mathbb{R}^{H\times W\times 3}$ and outputs a latent code
$\textbf{h}\in\mathbb{R}^{H/f\times W/f\times C}$ for some downsampling factor
$f$ and output channel size $C$. From this, a decoder tries to reconstruct the
original image $\mathbb{x}$ as best as possible. The autoencoder’s objective
combines an L1 loss with perceptual and patch-based adversarial losses (Esser
et al., 2021; Rombach et al., 2022). Perceptual and adversarial losses are
crucial to attaining sharp reconstructions, since the decoder would otherwise
produce a blurry average of prediction to minimize pixel-space distortion.
Meanwhile, instabilities caused by adversarial training are tempered by the L1
and perceptual losses.

Our motivation for the autoencoder stems from the manifold hypothesis
(Goodfellow et al., 2016), which states that real-world data modalities lie on
a low dimensional manifold in a high dimensional space. This implies that the
autoencoder can learn representations with much lower dimensionlity but
contain all factors of variation in the data. Such representations might be
beneficial to a generative model for several reasons. One reason could be the
excessive amounts of trivial relationships present in high-dimensional data.
These perceptually meaningless relationships make up many of the dependencies
in the ambient space, so removing them allows focusing on the most important
concepts. Additionally, data distributions in high dimensional spaces tend to
be extremely sharp, with data modes separated by wide boundaries with no
density at all. Reducing this sharpness could be especially beneficial for
VAEs which must fit Guassian posteriors to Gaussian priors.

###  3.2 Hierarchical VAEs

Hierarchical VAEs are latent variable models that train a generator
$p(\mathbb{x}|\mathbb{z})$ to reconstruct data $\mathbb{x}$ from a sequence of
latent variables
$\mathbb{z}\coloneqq\\{\mathbb{z}_{1},\ldots,\mathbb{z}_{N}\\}$ sampled from a
learned posterior. The posterior encodes information about the data necessary
for reconstruction, but is also incentivized to be close to a prior
distribution. The posterior and prior are generally factorized normal
distributions of the form
$q(\mathbb{z}|\mathbb{x})\coloneqq\prod_{i=1}^{N}q(\mathbb{z}_{i}|\mathbb{z}_{<i},\mathbb{x})$
and $p(\mathbb{z})\coloneqq\prod_{i=1}^{N}p(\mathbb{z}_{i}|\mathbb{z}_{<i})$
respectively. The generator, posterior, and prior are trained jointly to
maximize a lower bound on the data log-likelihood:

$\log p(\mathbb{x})\geq\mathbb{E}_{q(\mathbb{z}|\mathbb{x})}\left[\log
p(\mathbb{x}|\mathbb{z})\right]-\text{KL}({q(\mathbb{z}_{1}|\mathbb{x})}\|{p(\mathbb{z}_{1})})\\\
-\sum_{i=2}^{N}\mathbb{E}_{q(\mathbb{z}_{<i}|\mathbb{x})}\left[\text{KL}({q(\mathbb{z}_{i}|\mathbb{z}_{<i},\mathbb{x})}\|{p(\mathbb{z}_{i}|\mathbb{z}_{<i})})\right]$
(1)

The objective in Equation 1 is analogous to minimizing a lossless codelength
of the data, with the reconstruction loss describing the distortion measured
in nats, and KL terms representing the rate (Chen et al., 2016). However, if
most bits in a natural image encode barely perceptible details, then this
lossless compression task is dominated by the perceptual compression
component. In other words, a VAE trying to minimize its loss would do so by
mastering high-level features while ignoring the image’s semantics, since the
former make up the majority of the image’s codelength.

This problem disappears when training in the latent space since high-frequency
details have already been removed by the autoencoder. Thus to minimize its
loss, a latent-space VAE must learn the complex dependencies between image
concepts. Fortunately, it can dedicate all its capacity to this single task
without spending any on learning details. This is especially important for
deep VAEs, since it redistributes not only parameters but also latent groups.

###  3.3 Classifier-Free Guidance

In addition to learning on lower dimensional representations, large-scale
diffusion and autoregressive models employ classifier-free guidance (Ho &
Salimans, 2022) to boost image fidelity (Nichol et al., 2021; Saharia et al.,
2022b; Gafni et al., 2022; Yu et al., 2022). This technique is designed to
trade off diversity for sample quality, since poor likelihood-based models
tend to cover regions not present in the data distribution. Guidance pushes
samples towards regions that more closely resemble a desired label by
comparing conditional and unconditional likelihood functions.

(Luhman & Luhman, 2022) extends this technique to hierarchical VAEs where
means and log-variances are extrapolated towards class-conditional values and
away from unconditional ones. This VAE guidance strategy considers diagonal
Gaussian priors $p(\mathbb{z}_{i}|\mathbb{z}_{<i})$ conditioned only on
previous latents, and priors $p(\mathbb{z}_{i}|\mathbb{z}_{<i},\mathbb{c})$
conditioned on both previous latents and an auxiliary label $\mathbb{c}$. The
guided latent variables are drawn from
$p_{\text{guided},w_{\mu},w_{\sigma}}(\mathbb{z}_{i}|\mathbb{z}_{<i},\mathbb{c})$
defined as:

$p_{\text{guided},w_{\mu},w_{\sigma}}(\mathbb{z}_{i}|\mathbb{z}_{<i},\mathbb{c})\coloneqq\\\
\mathcal{N}\left(\mathbb{z}_{i};\boldsymbol{\mu}_{c}+w_{\mu}(\boldsymbol{\mu}_{c}-\boldsymbol{\mu}_{u}),\text{diag}\left(\boldsymbol{\sigma}_{c}^{2}\left(\frac{\boldsymbol{\sigma}_{c}}{\boldsymbol{\sigma}_{u}}\right)^{2w_{\sigma}}\right)\right)$
(2)

for conditional means and covariances
$\boldsymbol{\mu}_{c},\text{diag}\left(\boldsymbol{\sigma}_{c}^{2}\right)$ and
unconditional ones
$\boldsymbol{\mu}_{u},\text{diag}\left(\boldsymbol{\sigma}_{u}^{2}\right)$.
The strength of the mean and variance guidance are controlled by $w_{\mu}$ and
$w_{\sigma}$ respectively, which may be chosen independently.

Intuitively, mean guidance accentuates components that are present in the
class-conditional value but not in the unconditional one. This helps give
generated images a distinct visual structure that is characteristic of other
images in the class. Variance guidance acts similarly to low-temperature
sampling, adaptively adjusting the temperature based on how the ratio of
class-conditional variance compares to the unconditional one.

To implement guidance in VAEs, we train a model where all priors are
conditioned on a label that is replaced with a dummy label 10% of the time. At
sampling time, we store running hidden states for both the conditional and
unconditional generation paths. Each path outputs its own prior, and we
combine the two into a single guided distribution according to Equation 2. A
single sample $\mathbb{z}\sim p_{\text{guided}}$ is then projected back into
both hidden states (see Appendix C for a PyTorch implementation).

##  4 Experiments

In this section, we ablate the effect of training in the latent space for
various image benchmarks, and evaluate our method’s image generation
capabilities and its efficiency. Full details of our setup is listed in
Appendix B. For each dataset considered in this section, we use an off-the-
shelf autoencoder from (Rombach et al., 2022) as our first stage model.

###  4.1 Ablation Studies

We first test whether training in the latent space improves over a pixel space
VAE on the LSUN Church-256 dataset. We also compare to an upsampling stack,
where a deep VAE is trained to generate low-resolution images, which are then
upsampled with a second VAE. Both latent and upsampling methods produce
$32^{2}$ outputs in the first stage, and have 119 million parameters. Our VAE
upsampler has 63 million, while our DAE upsampler has slightly more at 80
million. All models are trained for 250 thousand steps and roughly the same
wall clock time ($\sim 350$ TPUv2-8 hrs).

Table 1: Comparison between methods on LSUN Church-256. Method | FID $\downarrow$ | Prec $\uparrow$ | Rec $\uparrow$
---|---|---|---
Latent-space VAE | 7.89 | 0.69 | 0.35
Upsampling Stack | 33.53 | 0.58 | 0.10
Pixel-space VAE | 44.36 | 0.48 | 0.13

Table 1 shows the ablation results. The latent-space VAE obtains much better
FID than the upsampling stack and pixel-space VAE, performing comparably to a
diffusion model from (Ho et al., 2020). The good FID and precision verify the
latent VAE’s more realistic samples, but its higher recall indicates better
mode coverage of the data manifold. The pixel-space VAE covers modes in the
ambient space better because of its data log-likelihood objective, but its
samples’ lack of realism make them fall outside the data manifold. The
compressed DAE representation more closely resembles this manifold, allowing
useful learning signal to flow into the generative model.

Next, we ablate the effect of the downsampling factor on performance. Three
models with $f=4$, $f=8$ and $f=16$ downsampling are trained on the LSUN-
Bedroom dataset. All models have $\sim 100$ million parameters and trained for
250 thousand steps.

Table 2: Comparison between downsampling factors on LSUN Bedroom-256 Downsampling Factor | FID $\downarrow$ | Prec $\uparrow$ | Rec $\uparrow$
---|---|---|---
$f=4$ | 9.35 | 0.56 | 0.36
$f=8$ | 11.16 | 0.47 | 0.36
$f=16$ | 17.46 | 0.38 | 0.28

As shown in Table 2, a compression rate of $f=4$ performs only slightly better
than $f=8$, but $f=16$ performs notably worse. This indicates that the benefit
of dimensionality reduction has diminishing returns beyond a certain point,
and excessive compression rates such as $f=16$ can hurt overall performance.
Our observations are consistent with (Rombach et al., 2022), who find that
compression rates of $f=4$ and $f=8$ work best with diffusion models.

###  4.2 Quantitative results

Table 3 evaluates our method on the class-conditional ImageNet $256^{2}$
dataset using the Frechet Inception Distance (FID) and Inception Score (IS)
metrics (Heusel et al., 2017; Salimans et al., 2016), as well as the Precision
and Recall metrics to assess fidelity and diversity respectively [precrecall].
We train a 386M parameter model for 475k iterations on a batch size of 1024,
with an $f=8$ autoencoder from (Rombach et al., 2022) as our first stage.

We obtain a FID of 9.34 when performing guided sampling with $w_{\mu}=1.5$ and
$w_{\sigma}=3$222Mean guidance not used for $1\times 1$ latents to reduce
saturation.. Surprisingly, this result is only slightly worse than LDM-8,
which was with the same methodology and downsampling factor but uses a
diffusion backbone. This provides evidence that the gap between diffusion
models and VAEs greatly shrinks when moving to a compressed latent space.

When not applying guidance, our VAE’s sample quality degrades significantly as
measured by FID and Recall. Our strong performance in the guided setting
indicates useful features have been learned, but the noisiness of unguided
sampling prevents this from being translated into realistic images. Guidance
directs the model to regions where it can effectively utilize these learned
features. Nevertheless, our unguided samples exhibit high recall, implying
diverse mode coverage. Even if these samples are not as impressive visually,
reliably covering modes is a desirable feature and an advantage of VAEs over
GANs.

Table 3: Performance on ImageNet $256^{2}$. $\dagger$ indicates methods using classifier-free guidance. $\ddagger$ indicates two-stage approaches. All numbers taken from (Dhariwal & Nichol, 2021) except for LDM and StyleGAN-XL, which we take from their corresponding papers. Method | FID$\downarrow$ | Prec$\uparrow$ | Rec$\uparrow$
---|---|---|---
Latent VAE†‡ (ours, guided) | 9.34 | 0.85 | 0.32
Latent VAE‡ (ours, unguided) | 32.7 | 0.56 | 0.59
VQVAE-2‡ (Razavi et al., 2019b) | 31.11 | 0.36 | 0.57
LDM-8 †‡ (Rombach et al., 2022) | 8.11 | 0.83 | 0.36
BigGAN-deep (Brock et al., 2018) | 6.95 | 0.87 | 0.28
LDM-4†‡ (Rombach et al., 2022) | 3.60 | 0.87 | 0.48
ADM (Dhariwal & Nichol, 2021) | 10.94 | 0.69 | 0.63
ADM-G† (Dhariwal & Nichol, 2021) | 4.59 | 0.82 | 0.52
StyleGAN-XL (Sauer et al., 2022) | 2.30 | 0.78 | 0.53

###  4.3 Qualitative results

Figure 2 shows random samples from our ImageNet model. The guided samples
demonstrate VAEs’ ability to create realistic pictures with both coherent
global structures and realistic textures. Although the unguided samples are
less realistic, they have greater variety in their composition and
backgrounds, demonstrating VAEs ability to trade sample diversity for
perceptual quality. The effect of guidance also differs from image to image:
the landscape in the last column of the first row changes little when applying
guidance, while the photos of butterflies are changed drastically.

Figure 2: Class-conditional samples from our ImageNet model without guidance
(left, FID 32.7), and with guidance (right, FID 9.34). Guidance significantly
improves the realism of the generated samples.

In Figure 3, we compare samples from our model, BigGAN, and ADM for the tinca
class. Our model is capable of generating diverse settings, such as tincas not
held by a person. Furthermore, our model renders human faces more
realistically than BigGAN. The diversity of model samples and realistic
generation of difficult concepts show how VAEs emulate the most desirable
aspects of diffusion models. Finally, we visualize nearest neighbors in
Appendix D to demonstrate that our model does not simply memorize the training
set.

Figure 3: Examples for the tinca class from BigGAN-deep (left), ADM (center),
and ours (right). Like diffusion models, VAEs show significant improvement
over GANs at generating faces.

###  4.4 Efficiency

Similar to GANs and unlike diffusion models, VAEs perform their entire
generative process in a single forward pass. As such, we would expect VAEs to
vastly outperform diffusion models in sampling speed. Table 4 supports this
claim empirically by comparing our VAE’s efficiency to common diffusion and
GAN methods. Our method is similarly fast as StyleGAN-XL (Sauer et al., 2022),
while boasting a $20\times$ speedup over a 50-step LDM and a $300\times$
speedup over ADM (Rombach et al., 2022; Dhariwal & Nichol, 2021).

Table 4: Sampling time (s) on ImageNet-256, measured for 1000 images on a Tesla T4 GPU. Guidance is used when applicable. Method | Sampling Time | Speedup Factor
---|---|---
Latent VAE (ours) | 213.8 | 1$\times$
StyleGAN-XL | 272.9 | 0.781$\times$
LDM-4 (50 steps) | 4560 | 0.047$\times$
ADM (250 steps) | 72875 | 0.003$\times$

Although VAEs have an inherently fast sampling process, existing VAEs are
notoriously expensive to train (Harvey et al., 2021). This inefficiency stems
from the need to back-propagate through a deep hierarchy of latent variables.
As an example, a VD-VAE on FFHQ-256 uses 62 latent layers and roughly 1000
convolutional layers, and takes over 300 V-100 days of compute to train to
completion (Child, 2020). To present a more complete picture of our method’s
efficiency, we report the training speed of our method.

Our base ImageNet generator was trained for 26 days on TPUv3-8 hardware, which
equates to only 104 days on a single V-100. Even adding $\sim 70$ V-100 days
for training the autoencoder, this training time is significantly better than
VD-VAE’s, despite our generator having three times more parameters. We achieve
this fast training speed by only applying deep VAE layers only to low-
resolution components, using a wide but shallow autoencoder at high-
resolutions.

##  5 Image Manipulations

A distinguishing feature of VAEs is their powerful inference network where
data points are projected onto semantically meaningful latent space. This
allows for both good interpretability as well as flexible manipulation.
Section 5.1 examines learned representations using visualizing two latent
dimensions in a grid, and Section 5.2 creates interpolation paths between two
dataset images. Section 5.3 showcases how latents can be manipulated to
perform image outpainting without any task-specific training or auxiliary
networks. For simplicity, this section considers only our unconditional
models, which were trained with a standard normal prior.

###  5.1 Visualizing a Manifold

One desirable aspect of a latent variable model is good interpretability. This
lets researchers and practitioners understand what representations have been
learned and how they affect the model’s decisions. The seminal work of (Kingma
& Welling, 2013) used a VAE with a two-dimensional latent code to visualize
the entire learned manifold of a small-scale dataset. While our models and
datasets are far more complex and high-dimensional, we are interested in
understanding the learned latent representations.

Figure 4: Traversing two dimensions of the learned latent space. We vary the
noise seed for the most informative dimension of the first two layers, with
linearly spaced intervals from -2.25 to 2.25.

Our deep VAEs have thousands of latent dimensions starting from $1\times 1$
resolution all the way to $32\times 32$. Despite this high dimensionality,
most global features are encoded in $1\times 1$ and $4\times 4$ latents
(Child, 2020). To retrieve a two-dimensional representation, we consider the
first two latent groups at the $1\times 1$ resolution. Each of these groups
has 16 channels each, but similar to (Hazami et al., 2022), we find that most
the information is stored in only one of these 16 channels. We create a grid
by sweeping over the most active channel in the first two latent groups,
sampling all other dimensions with a fixed noise seed.

Figure 4 shows how these two dimensions affect samples of our LSUN Church
model. The churches become progressively more intricate when moving left to
right, while moving top to down changes the image’s vibrance and color
palette. The structure of the church stays relatively constant as we move
across these two dimensions, indicating this information is encoded in later
latent variables.

###  5.2 Interpolation

Another way of examining learned representations is by interpolating between
different samples. We are interested in generating between real images from
the dataset instead of synthetic model samples. For GANs, this is a more
complicated problem which requires either an auxiliary encoding network or
multiple optimization steps. However, in VAEs the inference network is trained
directly in the objective and is directly applicable to downstream tasks.

A simple way of interpolating between images is to simply interpolate between
the latent codes that created them. In other words, one could sample
$\mathbb{z}^{(1)}\sim q(\mathbb{z}|\mathbb{x}^{(1)})$ and
$\mathbb{z}^{(2)}\sim q(\mathbb{z}|\mathbb{x}^{(2)})$ to produce an
interpolated latent
$\mathbb{z}_{\text{interp}}\coloneqq\lambda\mathbb{z}^{(1)}+(1-\lambda)\mathbb{z}^{(2)}$.
The interpolated latent is converted to an image using the decoder
$p(\mathbb{x}|\mathbb{z}_{\text{interp}})$.

Figure 5 shows interpolations between samples from the LSUN Church-256
dataset. Despite its simplicity, this interpolation method works for models
with a standard normal prior. This makes sense because each posterior is
trained to fall within the same prior, so interpolating between points sampled
from these posteriors should produce another point within the prior. In models
with hierarchical priors, this direct interpolation fails because the
interpolated latent falls outside the priors learned during training. It may
be possible to reparameterize the hierarchical prior into a standard normal
one, as done in Appendix E of (Vahdat et al., 2021), but we leave that to
future work.

Figure 5: Interpolation results. Leftmost and rightmost columns are samples
from the dataset, while the middle three columns are generated interpolations.

###  5.3 Outpainting

We now explore how to manipulate latent variables for an image outpainting
task. Outpainting involves generating a larger image while its context is
provided by a reference image. Although our models are not explicitly
conditioned with the context provided, we can store this information in latent
variables that are used alongside newly sampled latents. Latents corresponding
to outpainted regions are sampled from a standard normal prior, while regions
corresponding to the original image use latents from the posterior.

To outpaint a $256^{2}$ image into $256\times 512$ image, we perform the
following steps. First, we encode the $256^{2}$ image which outputs a latent
code with $1\times 1$, $4\times 4$, $8\times 8$, $16\times 16$, and $32\times
32$ resolutions. We then double the width of each latent with resolution 4 or
greater by sampling random Gaussian noise for the new dimensions; this
gaussian noise can be considered a sample from our standard normal prior. We
leave the $1\times 1$ latent intact to ensure the new regions share global
context of the original image. These new latents are then decoded to produce
an image.

Figure 6 visualizes our outpainting results. The model can generate convincing
outpainted regions that are smoothly integrated into the original image. This
indicates impressive generalization abilities considering our model was not
trained explicitly for this task, and has not seen any images of this
resolution during training. While the samples are not perfect and modify the
original image in some places, it nevertheless highlights the usefulness of
VAE representations in downstream tasks. Such useful representations, combined
with easy access provided by the inference model, extend VAEs’ practical
applications for endeavors beyond simply generating samples.

Figure 6: The left column shows images from the dataset, and the right shows
outpainted samples from our model.

##  6 Conclusions

We have demonstrated how hierarchical VAEs trained in latent space can
generate highly realistic images on complex multimodal datasets. This
performance, combined with their fast sampling, stable training, and flexibile
manipulation make them an exciting class of generative models. Furthermore,
scaling to large text-to-image datasets might have great benefits for VAEs,
the way it has for other likelihood-based models.

One drawback of our method compared to pixel-space VAEs is the inability to
compute the data likelihood, which is useful for tasks such as density
estimation and anomaly detection. Additionally, our VAEs are still
outperformed by best-in-class diffusion models and GANs, particularly in the
unguided setting. We believe that further advancements in VAEs might help
close this gap.

## References

  * Bradbury et al. (2018) Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., et al.  Jax: composable transformations of python+ numpy programs.  _Version 0.2_ , 5:14–24, 2018.
  * Brock et al. (2018) Brock, A., Donahue, J., and Simonyan, K.  Large scale gan training for high fidelity natural image synthesis.  _ArXiv_ , abs/1809.11096, 2018.
  * Chen et al. (2016) Chen, X., Kingma, D. P., Salimans, T., Duan, Y., Dhariwal, P., Schulman, J., Sutskever, I., and Abbeel, P.  Variational lossy autoencoder.  _arXiv preprint arXiv:1611.02731_ , 2016.
  * Child (2020) Child, R.  Very deep vaes generalize autoregressive models and can outperform them on images.  _arXiv preprint arXiv:2011.10650_ , 2020.
  * Dai & Wipf (2019) Dai, B. and Wipf, D.  Diagnosing and enhancing vae models.  _arXiv preprint arXiv:1903.05789_ , 2019.
  * Deng et al. (2009) Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L.  Imagenet: A large-scale hierarchical image database.  In _2009 IEEE conference on computer vision and pattern recognition_ , pp. 248–255. Ieee, 2009.
  * Dhariwal & Nichol (2021) Dhariwal, P. and Nichol, A.  Diffusion models beat gans on image synthesis.  _Advances in Neural Information Processing Systems_ , 34:8780–8794, 2021.
  * Dinh et al. (2016) Dinh, L., Sohl-Dickstein, J., and Bengio, S.  Density estimation using real nvp.  _arXiv preprint arXiv:1605.08803_ , 2016.
  * Esser et al. (2021) Esser, P., Rombach, R., and Ommer, B.  Taming transformers for high-resolution image synthesis.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pp. 12873–12883, 2021.
  * Gafni et al. (2022) Gafni, O., Polyak, A., Ashual, O., Sheynin, S., Parikh, D., and Taigman, Y.  Make-a-scene: Scene-based text-to-image generation with human priors.  _ArXiv_ , abs/2203.13131, 2022.
  * Goodfellow et al. (2016) Goodfellow, I., Bengio, Y., and Courville, A.  _Deep Learning_.  MIT Press, 2016.  http://www.deeplearningbook.org.
  * Harvey et al. (2021) Harvey, W., Naderiparizi, S., and Wood, F.  Conditional image generation by conditioning variational auto-encoders.  In _International Conference on Learning Representations_ , 2021.
  * Hazami et al. (2022) Hazami, L., Mama, R., and Thurairatnam, R.  Efficient-vdvae: Less is more.  _arXiv preprint arXiv:2203.13751_ , 2022.
  * Heusel et al. (2017) Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S.  Gans trained by a two time-scale update rule converge to a local nash equilibrium.  _Advances in neural information processing systems_ , 30, 2017.
  * Higgins et al. (2017) Higgins, I., Matthey, L., Pal, A., Burgess, C., Glorot, X., Botvinick, M., Mohamed, S., and Lerchner, A.  beta-VAE: Learning basic visual concepts with a constrained variational framework.  In _International Conference on Learning Representations_ , 2017.  URL https://openreview.net/forum?id=Sy2fzU9gl.
  * Ho & Salimans (2022) Ho, J. and Salimans, T.  Classifier-free diffusion guidance.  _arXiv preprint arXiv:2207.12598_ , 2022.
  * Ho et al. (2020) Ho, J., Jain, A., and Abbeel, P.  Denoising diffusion probabilistic models.  _Advances in Neural Information Processing Systems_ , 33:6840–6851, 2020.
  * Ho et al. (2022) Ho, J., Saharia, C., Chan, W., Fleet, D. J., Norouzi, M., and Salimans, T.  Cascaded diffusion models for high fidelity image generation.  _J. Mach. Learn. Res._ , 23:47–1, 2022.
  * Hou et al. (2017) Hou, X., Shen, L., Sun, K., and Qiu, G.  Deep feature consistent variational autoencoder.  In _2017 IEEE winter conference on applications of computer vision (WACV)_ , pp. 1133–1141. IEEE, 2017.
  * Kingma & Ba (2014) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kingma & Welling (2013) Kingma, D. P. and Welling, M.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_ , 2013.
  * Kingma et al. (2016) Kingma, D. P., Salimans, T., Jozefowicz, R., Chen, X., Sutskever, I., and Welling, M.  Improved variational inference with inverse autoregressive flow.  _Advances in neural information processing systems_ , 29, 2016.
  * Larsen et al. (2016) Larsen, A. B. L., Sønderby, S. K., Larochelle, H., and Winther, O.  Autoencoding beyond pixels using a learned similarity metric.  In _International conference on machine learning_ , pp. 1558–1566. PMLR, 2016.
  * Li et al. (2022) Li, H., Yang, Y., Chang, M., Chen, S., Feng, H., Xu, Z., Li, Q., and Chen, Y.  Srdiff: Single image super-resolution with diffusion probabilistic models.  _Neurocomputing_ , 479:47–59, 2022.
  * Luhman & Luhman (2022) Luhman, E. and Luhman, T.  Optimizing hierarchical image vaes for sample quality.  _arXiv preprint arXiv:2210.10205_ , 2022.
  * Maaløe et al. (2019) Maaløe, L., Fraccaro, M., Liévin, V., and Winther, O.  Biva: A very deep hierarchy of latent variables for generative modeling.  _Advances in neural information processing systems_ , 32, 2019.
  * Nichol et al. (2021) Nichol, A., Dhariwal, P., Ramesh, A., Shyam, P., Mishkin, P., McGrew, B., Sutskever, I., and Chen, M.  Glide: Towards photorealistic image generation and editing with text-guided diffusion models.  _arXiv preprint arXiv:2112.10741_ , 2021.
  * Pope et al. (2021) Pope, P., Zhu, C., Abdelkader, A., Goldblum, M., and Goldstein, T.  The intrinsic dimension of images and its impact on learning.  _arXiv preprint arXiv:2104.08894_ , 2021.
  * Ramesh et al. (2022) Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., and Chen, M.  Hierarchical text-conditional image generation with clip latents.  _arXiv preprint arXiv:2204.06125_ , 2022.
  * Razavi et al. (2019a) Razavi, A., Van den Oord, A., and Vinyals, O.  Generating diverse high-fidelity images with vq-vae-2.  _Advances in neural information processing systems_ , 32, 2019a.
  * Razavi et al. (2019b) Razavi, A., Van den Oord, A., and Vinyals, O.  Generating diverse high-fidelity images with vq-vae-2.  _Advances in neural information processing systems_ , 32, 2019b.
  * Rezende & Mohamed (2015) Rezende, D. and Mohamed, S.  Variational inference with normalizing flows.  In _International conference on machine learning_ , pp. 1530–1538. PMLR, 2015.
  * Rezende et al. (2014) Rezende, D. J., Mohamed, S., and Wierstra, D.  Stochastic backpropagation and approximate inference in deep generative models.  In _International conference on machine learning_ , pp. 1278–1286. PMLR, 2014.
  * Rombach et al. (2022) Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B.  High-resolution image synthesis with latent diffusion models.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 10684–10695, 2022.
  * Saharia et al. (2022a) Saharia, C., Chan, W., Chang, H., Lee, C., Ho, J., Salimans, T., Fleet, D., and Norouzi, M.  Palette: Image-to-image diffusion models.  In _ACM SIGGRAPH 2022 Conference Proceedings_ , pp. 1–10, 2022a.
  * Saharia et al. (2022b) Saharia, C., Chan, W., Saxena, S., Li, L., Whang, J., Denton, E., Ghasemipour, S. K. S., Ayan, B. K., Mahdavi, S. S., Lopes, R. G., et al.  Photorealistic text-to-image diffusion models with deep language understanding.  _arXiv preprint arXiv:2205.11487_ , 2022b.
  * Saharia et al. (2022c) Saharia, C., Ho, J., Chan, W., Salimans, T., Fleet, D. J., and Norouzi, M.  Image super-resolution via iterative refinement.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 2022c.
  * Salimans et al. (2016) Salimans, T., Goodfellow, I. J., Zaremba, W., Cheung, V., Radford, A., and Chen, X.  Improved techniques for training gans.  _ArXiv_ , abs/1606.03498, 2016.
  * Sauer et al. (2022) Sauer, A., Schwarz, K., and Geiger, A.  Stylegan-xl: Scaling stylegan to large diverse datasets.  In _ACM SIGGRAPH 2022 conference proceedings_ , pp. 1–10, 2022\.
  * Sønderby et al. (2016) Sønderby, C. K., Raiko, T., Maaløe, L., Sønderby, S. K., and Winther, O.  Ladder variational autoencoders.  _Advances in neural information processing systems_ , 29, 2016.
  * Vahdat & Kautz (2020) Vahdat, A. and Kautz, J.  Nvae: A deep hierarchical variational autoencoder.  _Advances in Neural Information Processing Systems_ , 33:19667–19679, 2020.
  * Vahdat et al. (2021) Vahdat, A., Kreis, K., and Kautz, J.  Score-based generative modeling in latent space.  _Advances in Neural Information Processing Systems_ , 34:11287–11302, 2021.
  * Yoshida & Miyato (2017) Yoshida, Y. and Miyato, T.  Spectral norm regularization for improving the generalizability of deep learning.  _arXiv preprint arXiv:1705.10941_ , 2017.
  * Yu et al. (2021) Yu, J., Li, X., Koh, J. Y., Zhang, H., Pang, R., Qin, J., Ku, A., Xu, Y., Baldridge, J., and Wu, Y.  Vector-quantized image modeling with improved vqgan.  _ArXiv_ , abs/2110.04627, 2021.
  * Yu et al. (2022) Yu, J., Xu, Y., Koh, J. Y., Luong, T., Baid, G., Wang, Z., Vasudevan, V., Ku, A., Yang, Y., Ayan, B. K., et al.  Scaling autoregressive models for content-rich text-to-image generation.  _arXiv preprint arXiv:2206.10789_ , 2022.

##  Appendix A Introduction to Variational Autoencoders

###  A.1 Single group VAEs

For completeness, we provide an introduction to single-group and hierarchical
VAEs in this appendix. Variational autoencoders are generative models of the
form $p(\mathbb{x},\mathbb{z})=p(\mathbb{x}|\mathbb{z})p(\mathbb{z})$ where
latent variables $\mathbb{z}$ are used to estimate a conditional probability
of the data $\mathbb{x}$. While the true posterior $p(\mathbb{z}|\mathbb{x})$
is generally intractable, introducing an approximate posterior
$q(\mathbb{z}|\mathbb{x})$ enables us to obtain an evidence lower bound (ELBO)
on the data log-likelihood:

$\displaystyle\log p(\mathbb{x})$ $\displaystyle=\log\int
p(\mathbb{x}|\mathbb{z})p(\mathbb{z})\text{d}\mathbb{z}$ (3)
$\displaystyle=\log\left(\int
q(\mathbb{z}|\mathbb{x})\frac{p(\mathbb{x}|\mathbb{z})p(\mathbb{z})}{q(\mathbb{z}|\mathbb{x})}\text{d}\mathbb{z}\right)$
(4)
$\displaystyle=\log\left(\mathbb{E}_{q(\mathbb{z}|\mathbb{x})}\left[\frac{p(\mathbb{x}|\mathbb{z})p(\mathbb{z})}{q(\mathbb{z}|\mathbb{x})}\right]\right)$
(5)
$\displaystyle\geq\mathbb{E}_{q(\mathbb{z}|\mathbb{x})}\left[\log\frac{p(\mathbb{x}|\mathbb{z})p(\mathbb{z})}{q(\mathbb{z}|\mathbb{x})}\right]$
(6) $\displaystyle\geq\mathbb{E}_{q(\mathbb{z}|\mathbb{x})}\left[\log
p(\mathbb{x}|\mathbb{z})\right]-\mathbb{E}_{q(\mathbb{z}|\mathbb{x})}\left[\log
q(\mathbb{z}|\mathbb{x})-\log p(\mathbb{z})\right]$ (7)
$\displaystyle\geq\mathbb{E}_{q(\mathbb{z}|\mathbb{x})}\left[\log
p(\mathbb{x}|\mathbb{z})\right]-\text{KL}({q(\mathbb{z}|\mathbb{x})}\|{p(\mathbb{z})})$
(8)

where Equation 4 is obtained using Jensen’s inequality. In a VAE, the
posterior $q(\mathbb{z}|\mathbb{x})$ and prior $p(\mathbb{z})$ are usually
diagonal Gaussian distributions, which enables us to learn the posterior using
stochastic backpropogation. By reparameterizing
$\mathbb{z}=\boldsymbol{\mu}_{\theta}(\mathbb{x})+\boldsymbol{\epsilon}\odot\boldsymbol{\sigma}_{\theta}(\mathbb{x})$
for $\boldsymbol{\epsilon}\sim\mathcal{N}(\textbf{0},\mathbb{I})$, we can
differentiate Equation 6 with respect to the mean $\boldsymbol{\mu}_{\theta}$
and variance $\boldsymbol{\sigma}_{\theta}$ of the posterior.

###  A.2 Hierarchical VAEs

For complex data distributions such as natural images, a single Gaussian
latent layer is often not expressive enough to capture important dependencies.
One way of improving expressitivity is to factorize latent distributions over
multiple groups
$\mathbb{z}\coloneqq\\{\mathbb{z}_{1},\ldots,\mathbb{z}_{N}\\}$, where the
prior is defined as
$p(\mathbb{z})\coloneqq\prod_{i=1}^{N}p(\mathbb{z}_{i}|\mathbb{z}_{<i})$ and
the posterior as
$q(\mathbb{z}|\mathbb{x})\coloneqq\prod_{i=1}^{N}q(\mathbb{z}_{i}|\mathbb{z}_{<i},\mathbb{x})$.
This conditional dependence creates a hierarchy among latent groups where
early latents encode global information that is then used by later groups for
modeling finer details.

A hierarchical VAE’s main components include a generator network that
assembles latents to reconstruct the data and parameterize the prior, an
encoder to extract useful representations from the input, and a sequence of
posterior blocks conditioned on the encoded input and generator hidden state.
Both the encoder and generator are implemented with residual neural networks;
this is especially important for the generator as it allows non-Markovian
dependencies where future latents depend on a combination of all previous
ones. When considering image data, it is helpful to have the generator’s
hidden state start at low resolution and progressively increase it to reflect
a coarse-to-fine generation process (Child, 2020).

Figure 7: A visual depiction of a hierarchal VAE.

##  Appendix B Implentation Details

Our models are based on the deep VAEs from (Child, 2020), with several
modifications that we detail here. While we did not do official ablations for
these changes, we do not believe they have a major positive or negative impact
on the model’s performance (with the exception of class-conditioning and
guidance).

  * •

We replace all bottleneck blocks with inverted bottleneck blocks for better
memory efficiency. Our inverted bottleneck blocks consist of two $3\times 3$
convolutions followed by two $1\times 1$ convolutions in a feed-forward
network.

  * •

We introduce multi-headed attention in the encoder and generator for modeling
non-local relationships.

  * •

In layers that output prior or posterior distributions, we use spectral
regularization (Yoshida & Miyato, 2017). This is similar to (Vahdat & Kautz,
2020), with the difference being they apply it to all layers in the network.

  * •

In our unconditional experiments, we opt to use a (non-hierarchical) standard
normal prior instead of a hierarchical one due to its simplicity. We continue
to use a hierarchical prior in class-conditional experiments so that the prior
is conditioned on labels, and to facilitate guidance in the prior.

  * •

We shift and scale the data by constants so it has means and variances close
to zero and one respectively. This normalized data is used as both the input
and reconstruction target.

  * •

Similar to (Luhman & Luhman, 2022), our reconstruction loss is a KL divergence
between Gaussians instead of a negative log-likelihood loss. The target
distribution is centered at the input and has fixed variance $\sigma^{2}_{q}$,
while the prediction has learned means and variances. This makes it similar to
the negative log-likelihood loss, except that it discourages the prior
variance from being less than $\sigma^{2}_{q}$. The parameter $\sigma^{2}_{q}$
controls the strength of the reconstruction term, similar to the $\beta$
parameter in a $\beta$-VAE (Higgins et al., 2017).

  * •

To prevent posterior collapse, we use the KL-scheduling technique from (Luhman
& Luhman, 2022) which encourages a designated amount of KL to be used in each
latent layer.

Table 5: Training details for ablation and image generation experiments.

Dataset | Church Latent-VAE | Church $32^{2}$ RGB | Church Upsampler | Church Pixel-space
---|---|---|---|---
Parameters | 119.1M | 119.1M | 62.8M | 143.3M
Base channel size | 192 | 192 | 64 | 64
Resolutions | 32,16,8,4,1 | 32,16,8,4,1 | 256,128,64,32 | 256,128,64,32,16,8,4,1
Channel multiplier | 1,1,1,1,1 | 1,1,1,1,1 | 2,2,3,3 | 2,2,3,3,3,3,3,3
FFN expansion ratio | 4,4,4,4,4 | 4,4,4,4,4 | 0,0,4,4 | 0,0,0,4,4,4,4,4
Num latent groups | 8,8,8,8,8 | 8,8,8,8,8 | 3,6,8,8 | 2,5,8,8,8,8,8,8
Num encoder layers | 4,4,4,4,4 | 4,4,4,4,4 | 2,3,4,4 | 2,3,4,4,4,4,4,4
Num attention layers | 4,4,4,0,0 | 4,4,4,0,0 | 0,0,0,4 | 0,0,0,4,4,4,0,0
Microbatch size $\times$ # accums | $128\times 1$ | $128\times 1$ | $32\times 4$ | $32\times 4$
Learning rate | $2\times 10^{-4}$ | $2\times 10^{-4}$ | $1\times 10^{-4}$ | $1\times 10^{-4}$
$\sigma_{q}$ | 0.07 | 0.035 | 0.05 | 0.05
SR penalty | 0.1 | 0.1 | 0.25 | 0.25
EMA decay | 0.9999 | 0.9999 | 0.9998 | 0.9999
Hardware | TPUv2-8 | TPUv2-8 | TPUv2-8 | TPUv3-8
Wall clock time (hr) | 36 | 312 | 345 | 175

Dataset | Bedroom f-4 | Bedroom f-8 | Bedroom f-16 | ImageNet-256
---|---|---|---|---
Parameters | 96.2M | 107.2M | 99.8M | 386.4M
Base channel size | 64 | 192 | 192 | 256
Resolutions | 64,32,16,8,4,1 | 32,16,8,4,1 | 16,8,4,1 | 32,16,8,4,1
Channel multiplier | 1,1,1,1,1 | 1,1,1,1,1 | 1,1,1,1 | 1,1,1,1,1
FFN expansion ratio | 0,4,4,4,4 | 4,4,4,4,4 | 4,4,4,4 | 4,6,8,10,12
Num latent groups | 8,12,10,8,6,6 | 8,8,7,6,6 | 10,9,8,6 | 10,12,10,8,8
Num encoder layers | 4,6,5,4,3,3 | 4,4,4,3,3 | 5,5,4,3 | 4,5,4,3,3
Num attention layers | 0,4,4,4,0,0 | 4,4,4,0,0 | 4,4,0,0 | 4,4,4,0,0
Microbatch size $\times$ # accums | $32\times 4$ | $128\times 1$ | $128\times 1$ | $128\times 8$
Learning Rate | $2\times 10^{-4}$ | $2\times 10^{-4}$ | $2\times 10^{-4}$ | $3\times 10^{-4}$
$\sigma_{q}$ | 0.1 | 0.07 | 0.07 | 0.07
SR Penalty | 0.1 | 0.1 | 0.1 | 0.1
EMA Decay | 0.9999 | 0.9999 | 0.9999 | 0.9998
Hardware | TPUv2-8 | TPUv2-8 | TPUv2-8 | TPUv3-8
Wall clock time (hr) | 104 | 28 | 17 | 630

Table 5 shows hyperparameters for our experiments. All models are trained with
the Adam (Kingma & Ba, 2014) optimizer with $\beta_{1}=0.9,\beta_{2}=0.9$ with
a linear learning rate warmup for the first 100 steps. We implement our
repository in JAX (Bradbury et al., 2018) and train on either TPUv2-8 or
TPUv3-8 hardware. Gradient accumulation is used when a batch cannot fit into
memory. When interacting with pretrained autoencoder models, we use LDM
PyTorch repository at https://github.com/CompVis/latent-diffusion. All metrics
are computed with the evaluation script from (Dhariwal & Nichol, 2021) at
https://github.com/openai/guided-diffusion/tree/main/evaluations.

##  Appendix C Implementation of Classifier-Free Guidance

This appendix provides a code implementation template for performing guided
sampling from a hierarchical VAE stochastic layer. While our VAEs are
conditioned only simple class labels, this guidance strategy can be easily
extended to other conditioning variables such as text.

⬇

# Hierarchical VAE stochastic layer that supports guidance

class StochasticLayer(nn.Module):

...

def sample_guided(self, h_c, h_u, label, w_mu, w_sigma):

uncond_label = torch.ones_like(label) * NUM_CLASSES

cond_output = self.block1(h_c, label)

uncond_output = self.block1(h_u, uncond_label)

pmean_c, plogvar_c, residual_c = torch.split(cond_output, [self.z_dim,
self.z_dim, self.hidden_dim], dim=1)

pmean_u, plogvar_u, residual_u = torch.split(uncond_output, [self.z_dim,
self.z_dim, self.hidden_dim], dim=1)

guided_mean = pmean_c + w_mu * (pmean_c - pmean_u)

guided_logvar = plogvar_c + w_sigma * (plogvar_c - plogvar_u)

z = guided_mean + torch.randn(pmean_c.shape) * torch.exp(0.5 * guided_logvar)

z = self.z_projection(z)

h_c = h_c + residual_c + z

h_u = h_u + residual_u + z

h_c = h_c + self.block2(h_c, cond_label)

h_u = h_u + self.block2(h_u, uncond_label)

return h_c, h_u

##  Appendix D Nearest Neighbor Visualizations

Figure 8: Nearest neighbors computed in the feature space of Inception-v3. The
leftmost column shows samples from our model, and columns to the right are
neighbors from the training set.

##  Appendix E ImageNet Samples

Figure 9: Random, guided samples from our ImageNet-256 model (FID=9.34) Figure
10: Random, un-guided samples from our ImageNet-256 model (FID=32.7)

##  Appendix F LSUN Samples

Figure 11: Random samples from our LSUN Church-256 model (FID=7.89) Figure 12:
Random samples from a church upsampling stack in Section 4.1’s ablation
(FID=33.53).
Samples exhibit decent large-scale structure but have unrealistic details. It
is possible that a
more powerful upsampler could improve performance. Figure 13: Random samples
from a pixel-space model in Section 4.1’s ablation (FID=44.36). Figure 14:
Random samples from our LSUN Bedroom-256 model on a $4\times$ downsampled
latent space (FID=9.35). Figure 15: Random samples from our LSUN Bedroom-256
model on a $8\times$ downsampled latent space (FID=11.16). Figure 16: Random
samples from our LSUN Bedroom-256 model on a $16\times$ downsampled latent
space (FID=17.46).