Is the Discrete VAE’s Power Stuck in its Prior?
Abstract
We investigate why probabilistic neural models with discrete latent variables are effective at generating high-quality images.
We hypothesize that ﬁtting a more ﬂexible variational posterior distribution and performing joint training of the encoder, decoder and prior distribution should improve model ﬁt.
However, we ﬁnd that modifying the training procedure for the well-known vector quantized variational autoencoder (VQ-VAE) leads to models with lower marginal likelihood for held-out data and degraded sample quality.
These results indicate that current discrete VAEs use their encoder and decoder as a deterministic compression bottleneck.
The distribution-matching power of these models lies solely in the prior distribution, which is typically trained after clamping the encoder and decoder.
Introduction
While generative adversarial networks (GANs) excel at synthesizing realistic-looking natural images [Karras et al., 2020, Park et al., 2019], probabilistic models including variational autoencoders (VAEs) [Kingma and Welling, 2014, Higgins et al., 2017] and normalizing ﬂows
[Dinh et al., 2015, Kingma and Dhariwal, 2018, Ho et al., 2019] generally lag behind as measured by qualitative assessments of sample images.
However, GANs do not offer many attractive features that probabilistic models naturally provide, such as efﬁcient encoding of inputs into a latent space representation and likelihood estimation for held-out data.
Indeed, while probabilistic models must posit some probability mass over every possible image, there is no constraint in the GAN architecture ensuring that all images correspond to a representation in latent space; due to mode collapse, the typical samples from a GAN regularly under-represent the training data distribution.
The current state-of-the-art VAE, as judged by image sample quality, is the vector quantized variational autoencoder (VQ-VAE).
VQ-VAEs introduce discrete latent variables and train a powerful autoregressive prior to capture high-level correlations in this discrete space [van den Oord et al., 2017].
A recent extension, dubbed VQ-VAE-2, competes with state of the art GANs as measured by
Inception Score and Fréchet Inception Distance [Razavi et al., 2019b] and improves sample image diversity and resolution by pairing a hierarchical encoder/decoder with a hierarchical autoregressive prior.
However, VQ-VAEs are trained in two steps.
The encoder and decoder networks are ﬁrst trained to minimize reconstruction error, like a traditional bottleneck autoencoder.
The learned encoding from inputs to latent variables is deterministic.
Next, a prior distribution is trained to model the discrete latent variables.
In this paper, we investigate the extent to which the VQ-VAE’s model ﬁt and sample quality depend on this speciﬁc training procedure.
We hypothesized that jointly training the encoder/decoder with the prior would improve the model ﬁt.
Surprisingly, we ﬁnd that “smoothing” the latent space through sampling joint training leads to poorer model ﬁt and lower sample image quality.
Unlike continuous
VAEs, which often perform well on interpolation tasks, our results show that the distribution of latent codes in the VQ-VAE is highly structured, and most latent codes do not decode to natural images.
Background
Variational autoencoders (VAEs) are a ﬂexible class of probabilistic latent variable models that leverage neural networks to learn nonlinear mappings between latent variables and observed variables [Kingma and Welling, 2014, Rezende et al., 2014].
A VAE posits a joint distribution over observed variables x and latent variables z : p ( x , z ; θ, ω )
=
p ( x | z ; θ ) p ( z ; ω ) .
The ﬂexibility in VAEs comes from training a (typically deterministic) mapping dec θ ( z ) .
The VAE likelihood p ( x | z ; θ ) is then parameterized by this mapping; for continuous observations, a commonly-used modeling choice is a
Gaussian distribution centered on the decoder output N ( dec θ ( z ) , σ 2 I ) with ﬁxed variance σ 2 .
VAEs are ﬁt by maximizing the marginal likelihood, or evidence , p ( x ; θ ) .
Since this quantity is usually intractable, approximate inference is performed by positing a variational posterior distribution q ( z | x ; φ ) .
Rather than optimizing a variational distribution for each observation x , VAEs perform amortized inference by training an encoder enc φ ( x ) , which in turn parameterizes q ( z | x ; φ ) .
The model is trained using the evidence lower bound (ELBO):
[Equation]
The outer expectation is approximated via Monte Carlo samples from the variational posterior q and the key “reparameterization trick” allow gradients to ﬂow through the sampling operation [Kingma and Welling, 2014, Rezende et al., 2014].
2.1 Vector Quantized VAEs
Our work focuses on extending the VQ-VAE, a VAE with discrete latent variables introduced by van den Oord et al.
In this model, each latent variable z i ∈ { 1 , . . .
, K } indexes into a codebook C ∈ R K × D to produce a row vector C z i .
The notation C z will refer to a matrix of codebook elements indexed by latent variables in z , where the i th row in C z is C z i .
The full model is speciﬁed by setting a uniform categorical prior and an isotropic Gaussian likelihood:
[Equation]
The authors ﬁt this model using a maximum a posteriori (MAP) estimate for the variational posterior by quantizing the output of the encoder to the nearest codebook vector:
[Equation]
Because the prior distribution is uniform and the posterior is a categorical distribution parameterized by a one-hot probability vector, the KL-divergence term in ELBO is a constant ( log K ).
Moreover, the ELBO expectation term is over a deterministic distribution q , so the model is simply trained by maximizing the likelihood max θ,φ log p ( x | z = dec θ ( x ); θ ) .
The sampling and quantization operations in Eq 4 are not differentiable, so the authors apply a “straight-through” gradient estimator, which propagates the gradients with respect to the codebook element C z i through to the outputs of the decoder enc φ ( x ) i .
If the encoder output and codebook element are sufﬁciently close, this is a close approximation, but there is no bound on the bias of the gradients with this method.
The VQ-VAE introduces additional terms to the loss function to move codebook entries and outputs of the encoder toward one another during training.
After training the encoder/decoder end-to-end, a ﬂexible autoregressive prior p ( z ; ω ) is ﬁt to the latent codes, obtained by the deterministic mapping in Eq 4.
The original VQ-VAE paper applies a
PixelCNN model for images [van der Oord et al., 2016].
In this work, we will apply a PixelCNN++ [Salimans et al., 2017].
Both the PixelCNN and PixelCNN++ are autoregressive convolutional neural networks that are trained efﬁciently using masked convolutions.
Model Variants
We make two changes to the VQ-VAE’s inference procedure to address approximations made in the original paper.
First, instead of ﬁtting a MAP estimate during variational inference, we ﬁt a
[Table]
mean-ﬁeld categorical distribution.
Second, instead of ﬁtting the encoder/decoder with a uniform prior using a straight-through gradient approximation, we train the model jointly end-to-end.
The model variants we explore are shown in Table 1.
In addition to the variants discussed below, we train baseline models with a Uniform categorical prior; this prior is never ﬁne-tuned.
The ﬁrst model “Quantized, PixelCNN++, Marginal” is the original VQ-VAE.
3.1 Quantization vs. Categorical Posterior
The ﬁrst question we ask is whether the quantization step in the VQ-VAE is responsible for the model’s success, or whether an encoder can perform equivalently by directly learning to parameterize a categorical distribution.
Because the variational posterior is simply a tool for amortized inference, we expect this equally-expressive alternative model to perform comparably.
For simplicity, we choose a mean-ﬁeld approximation q ( z | x ; φ ) =
(cid:81)
Ni =
1 q i
( z i ; x , φ ) where each q i is a categorical distribution.
The Quantized model uses Euclidean distance between outputs of the encoder dec θ ( x ) and codebook entries, similar to the original VQ-VAE.
Our “Quantized, PixelCNN++, Marginal” model is equivalent to the VQ-VAE with regularization parameter β = 0 .
All other Quantized models use a categorical variational distributions parameterized by distance:
[Equation]
We use the notation enc φ ( x ) i to refer to the encoder’s prediction for the i th latent variable.
The Categorical models directly estimate a categorical distribution without reference to the codebook:
[Equation]
3.2 Marginal vs. Joint Training
The second question we ask is whether the joint VQ-VAE is beneﬁting from the VAE framework at all.
Since the model is initially trained without any sampling operations, it is possible that the
VQ-VAE’s capability to model complex distributions, such as natural images, is entirely conferred by the subsequently-trained prior distribution.
Similar to the VQ-VAE loss function, we train the encoder/decoder in the Marginal models to maximize the expected likelihood term from ELBO: E z ∼ q ( z | x ; φ ) log p ( x | z ; θ ) .
This model can equivalently be viewed as a bottleneck autoencoder with sampling at the bottleneck layer, or as an amortized expectation-maximization algorithm to perform maximum likelihood estimation max φ,θ E z ∼ q ( z | x ; φ ) log p ( x | z ; θ )
[Roy et al., 2018].
Once the encoder/decoder converge, we train a PixelCNN++ to minimize the KL divergence: min ω D KL ( q ( z | x ; φ ) ||
p
z ; ω )) 1 .
Our Joint models are straightforwardly trained to maximize ELBO (Eq 1).
The model is jointly trained by allowing gradients to ﬂow between the encoder, decoder, and prior.
We ﬁt the prior distribution by optimizing its hyperparameters ω (i.e. empirical Bayes).
We use the Gumbel softmax trick from Jang et al.
[2017] to reparameterize the categorical sampling operation.
For our mean-ﬁeld variational posterior, this sampling is computed by:
[Equation]
where enc φ ( x ) i represents the K -length vector of logits for the i th latent variable, and g is a K -length vector where g k ∼ Gumbel (0 , 1) .
As τ → 0 , z i approaches a one-hot vector representing a draw from the distribution q
( z i = k ; φ ) =
softmax ( enc φ ( x ) i ) .
Rather than indexing into the codebook, the joint model matrix-multiplies the Gumbel-softmax samples into the codebook matrix C .
We begin training with a relatively high value of τ , resulting in a weighted mixture of codebook elements being passed to the decoder.
As the temperature is annealed, the samples for z i approach one-hot vectors.
In our experiments we found that annealing τ to a value of 0 .
1 over the course of training brought samples sufﬁciently close to one-hot vectors.
Related Work
Sønderby et al.
[2017] develop the most closely-related model to ours in the literature.
Rather than jointly training a ﬂexible prior, however, this work assumes that the prior distribution is uniform.
Our “Categorical, Uniform” model is similar to the GS-Soft model described in the paper.
Additionally,
Roy et al.
[2018] develop an EM inference algorithm for the VQ-VAE model.
Kaiser et al.
[2018] demonstrates that VQ-VAE model ﬁt degrades as the length of quantization vectors grows, a phenomenon they call “index collapse.”
Rather than introducing sampling to avoid the problem of index collapse, the authors develop a model to quantize sub-vectors of the encoder’s output using multiple codebooks.
It does not appear that our Joint models suffer from index collapse (see Section 5.2).
Similar to—yet distinct from—our ﬁndings, many recent works have found that VAE models with a ﬂexible decoder experience “posterior collapse”, whereby a powerful autoregressive decoder completely ignores outputs of the encoder.
The lack of a gradient from the decoder leads the variational posterior distribution q to collapse to match the prior distribution [Lucas et al., 2019,
Razavi et al., 2019a].
In the models we investigate, we restrict our attention to decoders that are convolutional neural networks, so our ﬁndings cannot be explained by posterior collapse.
Several recent works have tried to improve the quality of sample images by increasing the expressiveness of VAE models.
In particular, Rezende and Mohamed [2015] replaces the variational distribution in the encoder with a normalizing ﬂow, and Gulrajani et al.
[2017] uses a powerful autoregressive model as the decoder.
Both approaches are complementary to our model’s joint training of an expressive autoregressive prior.
Experiments
We present a qualitative and quantitative analysis of the model variants described in Section 3.
While the ImageNet samples in the original VQ-VAE paper appear high quality, it becomes clear upon close inspection that the model is only matching general textures (for example, most of the animals are missing heads).
We study facial images from the Large-scale CelebFaces Attributes (CelebA) dataset
[Liu et al., 2015] because we believe the quality of samples is more immediately apparent to human evaluators [Seyama and Nagayama, 2007].
All images are downscaled to 64x64 pixels, and intensities are scaled to the range (cid:2) 0 , 255256 (cid:3) .
We dequantize all images in order to be able to calculate discrete bits per dimension by adding uniform noise u ∼ U (cid:0) 0 , 1256 (cid:1)
[Kingma and Dhariwal, 2018].
In all model variants, we take one Monte Carlo sample per data point from the variational posterior distribution during training.
The models use the same encoder and decoder structure, consisting of 4 residual layers with 128 feature maps in both the encoder and decoder.
We set the number of codebook elements K = 256 .
We use strided convolutions to down/up sample the image spatial resolution.
Our PixelCNN++ prior is slightly modiﬁed from the original implementation, using softmax output rather than a logistic mixture and 3 residual layers with 64 channels instead of the original 5 residual layers .
Models are trained using the Adam optimizer for 500 epochs with a batch size of 128 and a learning rate of 3e-4.
The Gumbel softmax temperature is initialized at 1 . 0 and is annealed exponentially to 0 .
1 over the course of training.
The model was trained 10 times on random train/test splits to characterize variance; 70% of the dataset was used for training, 15% for testing, and 15% for validation.
5.1 Marginal Likelihood of Held Out Data
We apply the IWAE estimator described in Burda et al.
[2016] and Cremer et al.
[2018] to obtain an asymptotically tight lower bound on the model evidence; note that IWAE with 1 sample is equivalent to ELBO.
We compute IWAE by drawing S Monte Carlo samples z ( s ) ∼ q ( z | x ; φ ) and computing:
[Equation]
We convert our results from bounds on model evidence to bits per dimension:
[Equation]
where ˜ x ∼ U ( x, x + a ) , W , H , and C , are the width, height, and number of color channels in the image, respectively.
a measures the discretization level of the data; for 8-bit images standardized to the range (cid:0) 0 , 255256 (cid:1) , a will be 1256 .
Figure 1 shows average bits per dimension, along with variance across 10 random train/test splits.
The original VQ-VAE model (“Quantized, PixelCNN++, Marginal“) outperforms all other variants with 5.35 bits per dimension.
Joint training is not too far behind, with the “Categorical, PixelCNN++,
Joint“ model achieving 5.39 bits per dimension.
The models with Uniform priors underperform, and the “Categorical, PixelCNN++, Marginal” model, which is not trained via a valid ELBO objective, performs worst.
5.2 Index Collapse
Figure 2 shows the usage patterns across all codebook elements, for each model.
A large spike atfrequency indicates “index collapse”, whereby a large set of codebook elements are rarely used; a spike at 1 / 256 = 0 .
0039 would indicate “posterior collapse”, whereby every posterior distribution q converged to a uniform prior.
We see from these results that the models trained marginally do suffer from index collapse, but joint training with either a PixelCNN++ or a Uniform prior mitigates the problem.
Surprisingly, the VQ-VAE has the second-largest spike at 0, but achieves the highest Bits per Dimension.
No models appear to suffer from posterior collapse, as we expect since the decoder is a straightforward convolutional network.
5.3 Qualitative Evaluation
It is commonly accepted that likelihood of held-out data is not a good metric for sample quality [Theis et al., 2016].
Thus, we present samples from all model variants in Figure 3.
It is clear that these samples are not comparable to the quality of state of the art GAN samples; however, the VQ-VAE and prior architecture in Razavi et al.
[2019b] that generates samples competing with GANs is multi-scale and hierarchical whereas ours is not.
We believe this to be the reason for the comparatively poor samples.
We see that the samples from the original VQ-VAE (”Quantized, PixelCNN++, Marginal“) are moderately superior to the other models.
Interestingly, the “Categorical, PixelCNN++, Marginal” model, which performed worst as measured by model evidence, produces far superior samples to the models with a Uniform prior.
We note that in preliminary experiments (samples not shown), all of our model variants consistently produce high quality samples on simple images from CIFAR-10 and
MNIST.
Conclusions
Our experiments demonstrate that modifying the inference procedure for the VQ-VAE—by introducing a more expressive posterior distribution, adding sampling, and jointly training—have a large negative impact on the model ﬁt and small negative impact on the quality of images sampled from the model.
This ﬁnding is surprising, because we would expect to achieve better model ﬁt.
After all, the original VQ-VAE’s MAP posterior, i.e. a categorical distribution with all probability mass assigned to one location, could be learned by our Categorical models.
Although joint training leads to a worse model ﬁt, it offers the possibility of learning an accurate posterior distribution.
Being able to infer an accurate posterior distribution over discrete latent codes is important capability for a wide variety of inference tasks.
Our ﬁnding that learning a full variational posterior does not outperform a simple, quantized point estimate leads us to conclude that existing discrete neural models don’t truly amortize posterior model inference for discrete latent variable models.
Rather the encoder/decoder serve as a deterministic compression function.
All of the generative modelling functionality is contained within the prior distribution.
Allowing the PixelCNN++ prior to provide gradients to the encoder/decoder does not lead to improved model ﬁt.
This indicates that the PixelCNN++ is ﬂexible enough to model a somewhat arbitrary discrete distribution produced by a pre-trained, deterministic encoder.
Acknowledgments and Disclosure of Funding
We would like to thank our anonymous reviewers for insightful suggestions and comments.
Research presented in this paper was supported by the Laboratory Directed Research and Development program of Los Alamos National Laboratory under project numbers 20200666DI and 20210043DR.