A N
INFORMATION - THEORETIC ANALYSIS OF DEEP LATENT - VARIABLE MODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
We present an information-theoretic framework for understanding trade-offs in unsupervised learning of deep latent-variables models using variational inference.
This framework emphasizes the need to consider latent-variable models along two dimensions: the ability to reconstruct inputs (distortion) and the communication cost (rate).
We derive the optimal frontier of generative models in the two-dimensional rate-distortion plane, and show how the standard evidence lower bound objective is insufﬁcient to select between points along this frontier.
However, by performing targeted optimization to learn generative models with different rates, we are able to learn many models that can achieve similar generative performance but make vastly different trade-offs in terms of the usage of the latent variable.
Through experiments on MNIST and Omniglot with a variety of architectures, we show how our framework sheds light on many recent proposed extensions to the variational autoencoder family.
I NTRODUCTION
Deep learning has led to tremendous advances in supervised learning (Szegedy et al., 2016; Huang et al., 2017; Vaswani et al., 2017); however, unsupervised learning remains a challenging area.
Recent advances in variational inference (VI) (Kingma & Welling, 2014; Rezende et al., 2014), have led to an explosion of research in the area of deep latent-variable models and breakthroughs in our ability to model natural high-dimensional data.
This class of models typically optimize a lower bound on the log-likelihood of the data known as the evidence lower bound (ELBO), and leverage the “reparameterization trick” to make large-scale training feasible.
Improvements in building deep latent-variable generative models have primarily focused on deriving better algorithms and objectives for training that go beyond the ELBO, and identifying more powerful and ﬂexible architectures.
On the algorithmic side, new objectives (Burda et al., 2015;
Li & Turner, 2016; Ranganath et al., 2016a) have been presented to speed up training, learn better features, and combat some of the deﬁciencies of the ELBO.
On the architectural side, there are three directions of focus: (1) more complex encoders that extend the variational family (Rezende &
Mohamed, 2015; Kingma et al., 2016; Ranganath et al., 2016b) (2) more complex priors (Tomczak & Welling, 2017; Larochelle & Murray, 2011b; Nalisnick & Smyth, 2016), and (3) more complex decoders (Chen et al., 2016; Gulrajani et al., 2016; Makhzani & Frey, 2017).
These extensions have allowed us to scale VI to many exciting applications, but it remains unclear how these different pieces of complexity ﬁt together and interact.
Here we are interested in understanding the behavior of large families of latent-variable models in terms of the amount of information that the latent variable contains about the input.
We accomplish this by leveraging tools from information theory, which provides a reparameterization-independent measure of dependence.
Computing mutual information in high dimensions is problematic (Paninski, 2003; Zheng & Benjamini, 2016), thus we turn to recently developed tools in variational inference to approximate it.
We ﬁnd that a natural lower and upper bound on the mutual information between the input and latent variable can be simply related to the ELBO, and understood in terms of two terms: (1) a lower bound that depends on the distortion, or how well an input can be reconstructed through the encoder and decoder, and (2) an upper bound that measures the rate, or how costly it is to transmit information about the latent variable.
Together these terms provide a unifying perspective on the set of optimal models given a dataset, and show that there exists a continuum of models that make very different trade-offs in terms of rate and distortion.
By leveraging additional information about the amount of information contained in the latent variable, we show that we can recover the ground-truth generative model used to create the data in a toy model.
We perform extensive experiments on MNIST and Omniglot using a variety of encoder, decoder, and prior architectures and demonstrate how our framework provides a simple and intuitive mechanism for understanding the trade-offs made by these models.
By optimizing for models along this frontier using the method of Lagrange multipliers we recover the β -VAE objective (Higgins et al., 2017).
By varying β , we can learn many models with the same architecture and comparable generative performance (in terms of marginal data log likelihood), but that exhibit drastically different behavior in terms of the usage of the latent variable and variability of the decoder.
F RAMEWORK
Unsupervised Representation Learning
Depending on the task, there are many desiderata for a good representation.
Here we focus on one aspect of a learned representation: the amount of information that the latent variable contains about the input.
In the absence of additional knowledge of a “downstream” task, we focus on the ability to recover or reconstruct the input from the representation.
Given a set of samples from a true data distribution p ∗ ( x ) , our goal is to learn a representation that contains a particular amount of information and from which the input can be reconstructed as well as possible.
We will convert each observed data vector x into a latent representation z using any stochastic encoder e ( z | x ) of our choosing.
This then induces the joint distribution p e ( x, z ) =
p ∗ ( x ) e ( z | x ) and the corresponding marginal p e ( z ) =
(cid:82) dx p ∗ ( x ) e ( z | x ) (the “aggregated posterior” in (Makhzani et al., 2015; Tomczak & Welling, 2017)) and conditional p e ( x | z )
[Equation]
A good representation Z must contain information about the input X which we deﬁne as follows:
[Equation]
Unfortunately, Equation 1 is hard to compute, since we do not have access to the true data density p ∗ ( x ) , and computing the marginal p e ( z )
= (cid:82) dx p e ( x, z ) can be challenging.
However, in Appendices A.1 and A.2, we derive the following tractable lower and upper bounds on this quantity:
[Equation]
where d ( x | z ) (the “ decoder ”) is a variational approximation to p e ( x | z ) , m ( z ) (the “ marginal ”) is a variational approximation to p e ( z ) , and all the integrals can be approximated using Monte Carlo given a ﬁnite sample of data from p ∗ ( x ) , as we discuss in Section 2.
In connection with rate-distortion theory, we can interpret the upper bound as the rate R of our representation (Tishby & Zaslavsky, 2015).
This rate term measures the average number of additional nats necessary to encode samples from the encoder using an entropic code constructed from the marginal, being an average KL divergence.
Unlike most rate-distortion work (Cover & Thomas, 2012), where the marginal is assumed a ﬁxed property of the channel, here the marginal is a completely general distribution, which we assume is learnable.
Similarly, we can interpret the lower bound as the data entropy H , which measures the complexity of the dataset (a ﬁxed but unknown constant), minus the distortion D , which measures our ability to accurately reconstruct samples:
[Equation]
This distortion term is deﬁned in terms of an arbitrary decoding distribution d ( x | z ) , which we consider a learnable distribution.
This contrasts with most of the compression literature where distortion is typically measured using a ﬁxed perceptual loss (Ball´e et al., 2017).
Combining these equations, we get the “sandwich equation” H − D ≤ I ≤ R. Notice also that the usual ELBO takes the form
ELBO = − D − R.
Phase Diagram From the sandwich equation, we see that H − D − R ≤ 0 .
This is a bound that must hold for any set of four distributions p ∗ ( x ) , e ( z | x ) , d ( x | z ) , m ( z ) .
The inequality places strict limits on which values of rate and distortion are achievable, and allows us to reason about all possible solutions in a two dimensional RD -plane.
A sketch of this phase diagram is shown in
Figure 1.
First, we consider the data entropy term.
For discrete data 1 , all probabilities in X are bounded above by one and both the data entropy and distortion are non-negative ( H ≥ 0 , D ≥ 0 ).
The rate is also non-negative ( R ≥ 0 ), because it is an average KL divergence, for either continuous or discrete
Z .
The positivity constraints and the sandwich equation separate the RD -plane into feasible and infeasible regions, visualized in Figure 1.
The boundary between these regions is a convex curve (thick black line).
Even given complete freedom in specifying the encoder e ( z | x ) , decoder d ( x | z ) and marginal approximation m ( z ) , and inﬁnite data, we can never cross this bounding line.
We now explain qualitatively what the different areas of this diagram correspond to.
For simplicity, we will consider the inﬁnite model family limit, where we have complete freedom in specifying e ( z | x ) , d ( x | z ) and m ( z ) but consider the data distribution p ∗ ( x ) ﬁxed.
The bottom horizontal line corresponds to the zero distortion setting, which implies that we can perfectly encode and decode our data; we call this the auto-encoding limit .
The lowest possible rate is given by H , the entropy of the data.
This corresponds to the point ( R = H, D = 0) .
(In this case, our lower bound is tight, and hence d ( x | z ) =
p e ( x | z ) .)
We can obtain higher rates at ﬁxed distortion by making the marginal approximation m ( z ) a weaker approximation to p e ( z ) , since only the rate and not the distortion depends on m ( z ) .
The left vertical line corresponds to the zero rate setting.
Since R =
=
⇒
e
( z | x ) = m
( z ) , we see that our encoding distribution e ( z | x ) must itself be independent of x .
Thus the latent representation is not encoding any information about the input and we have failed to create a useful learned representation.
However, by using a suitably powerful decoder, d ( x | z ) , that is able to capture correlations between the components of x (e.g., an autoregressive model, such as pixelCNN (Salimans et al., 2017), or an acausal MRF model, such as (Dai et al., 2015)), we can still reduce the distortion to the lower bound of H , thus achieving the point ( R = 0 , D = H ) ; we call this the auto-decoding limit .
Hence we see that we can do density estimation without learning a good representation, as we will verify empirically in Section 4.
(Note that since R is an upper bound on the mutual information, in the limit that R = 0 , the bound must be tight, which guarantees that m ( z ) =
p e ( z ) .)
We can achieve solutions further up on the D -axis, while keeping the rate ﬁxed, simply by making the decoder worse, since only the distortion and not the rate depends on d ( x | z ) .
Finally, we discuss solutions along the diagonal line.
Such points satisfy D
= H − R , and hence both of our bounds are tight, so m ( z ) =
p e ( z ) and d ( x | z ) =
p e ( x | z ) .
(Proofs of these claims are given in Sections A.3 and A.4 respectively.)
So far, we have considered the inﬁnite model family limit.
If we have only ﬁnite parametric families for each of d ( x | z ) , m ( z ) , e ( z | x ) , we expect in general that our bounds will not be tight.
Any failure of the approximate marginal m ( z ) to model the true marginal p e ( z ) , or the decoder d ( x | z ) to model the true likelihood p e ( x | z ) , will lead to a gap with respect to the optimal black surface.
However, it will still be the case that H − D − R ≤ 0 .
This suggests that there will still be a one dimensional optimal surface, D ( R ) , or R ( D ) where optimality is deﬁned to be the tightest achievable sandwiched bound within the parametric family.
We will use the term RD curve to refer to this optimal surface in the rate-distortion ( RD ) plane.
Since the data entropy H is outside our control, this surface can be found by means of constrained optimization, either minimizing the distortion at some ﬁxed rate, or minimizing the rate at some ﬁxed distortion, as we show in Section 2.
Furthermore, by the same arguments as above, this surface should be monotonic in both R and D , since for any solution, with only very mild assumptions on the form of the parametric families, we should always be able to make m ( z ) less accurate in order to increase the rate at ﬁxed distortion (see shift from red curve to blue curve in ﬁg. 1), or make the decoder d ( x | z ) less accurate to increase the distortion at ﬁxed rate (see shift from red curve to green curve in ﬁg. 1).
Optimization In this section, we discuss how we can ﬁnd models that target a given point on the
RD curve.
Recall that the rate R and distortion D are given by
[Equation]
[Equation]
[Equation]
These can both be approximated using a Monte Carlo sample from our training set.
We also require that the terms log d ( x | z ) , log m ( z ) and log e ( z | x ) be efﬁcient to compute, and that e ( z | x ) be efﬁcient to sample from.
In Section 4, we will describe the modeling choices we made for our experiments.
In order to explore the qualitatively different optimal solutions along the frontier, we need to explore different rate-distortion trade-offs.
One way to do this would be to perform some form of constrained optimization at ﬁxed rate.
Alternatively, instead of considering the rate as ﬁxed, and tracing out the optimal distortion as a function of the rate D ( R ) , we can perform the Legendre transformation and can ﬁnd the optimal rate and distortion for a ﬁxed β = ∂D∂R , by minimizing min e ( z | x ) ,m ( z ) ,d ( x | z ) D + βR .
Writing this objective out in full, we get
[Equation]
If we set β = 1 , this matches the ELBO (evidence lower bound) objective used when training a
VAE (Kingma & Welling, 2013), with the distortion term matching the reconstruction loss, and the rate term matching the “KL term.”
Note, however, that this objective does not distinguish between any of the points along the diagonal of the optimal RD curve, all of which have β = 1 and the same
ELBO.
Thus the ELBO objective alone (and the marginal likelihood) cannot distinguish between models that make no use of the latent variable (autodecoders) versus models that make large use of the latent variable and learn useful representations for reconstruction (autoencoders).
This is demonstrated experimentally in Section 4.
If we allow a general β ≥ 0 , we get the β -VAE objective used in (Higgins et al., 2017; Alemi et al., 2017).
This allows us to smoothly interpolate between auto-encoding behavior ( β = 0 ), where the distortion is low but the rate is high, to auto-decoding behavior ( β = ∞ ), where the distortion is high but the rate is low, all without having to change the model architecture.
However, unlike Higgins et al.
; Alemi et al.
(2017), we additionally optimize over the marginal m ( z ) and compare across a variety of architectures, thus exploring a much larger solution space, which we illustrate empirically in Section 4.
R ELATED W ORK
Here we present an overview of the most closely related work.
A more detailed treatment can be found in Appendix B.
Model families for unsupervised learning with neural networks.
There are two broad areas of active research in deep latent-variable models with neural networks: methods based on the variational autoencoder (VAE), introduced by Kingma & Welling (2013);
Rezende et al.
(2014), and methods based on generative adversarial networks (GANs), introduced by Goodfellow et al.
In this paper, we focus on the VAE family of models.
In particular, we consider recent variants using inverse autoregressive ﬂow (IAF) (Kingma et al., 2016), masked autoregressive ﬂow (MAF)
(Papamakarios et al., 2017), PixelCNN++ (Salimans et al., 2017), and the VampPrior (Tomczak &
Welling, 2017), as well as common Conv/Deconv encoders and decoders.
Information Theory and machine learning.
Agakov (2006) was the ﬁrst to introduce tractable variational bounds on mutual information, and made close analogies and comparisons to maximum likelihood learning and variational autoencoders.
The information bottleneck framework (Tishby et al., 1999; Shamir et al., 2010; Tishby & Zaslavsky, 2015; Alemi et al., 2017; Achille & Soatto, 2016; 2017) allows a model to smoothly trade off the minimality of the learned representation ( Z ) from data ( X ) by minimizing their mutual information, I ( X ; Z ) , against the informativeness of the representation for the task at hand ( Y ) by maximizing their mutual information, I ( Z ; Y ) .
This constrained optimization problem is rephrased with the Lagrange multiplier, β , to the unconstrained optimization of I ( X ; Z ) − βI ( Z ; Y ) .
Tishby & Zaslavsky (2015) plot an RD curve similar to the one in this paper, but they only consider the supervised setting, and they do not consider the information content that is implicit in powerful stochastic decoders.
Higgins et al.
(2017) proposed the β -VAE for unsupervised learning, which is a generalization of the original VAE in which the KL term is scaled by β , similar to this paper.
However, they only considered β > 1 .
In this paper, we show that when using powerful autoregressive decoders, using β ≥ 1 results in the model ignoring the latent code, so it is necessary to use β < 1 .
Generative Models and Compression.
Much recent work has explored the use of latent-variable generative models for image compression.
Ball ´ e et al.
(2017) studies the problem explicitly in terms of the rate/distortion plane, adjusting a Lagrange multiplier on the distortion term to explore the convex hull of a model’s optimal performance.
Johnston et al.
(2017) uses a recurrent VAE architecture to achieve state-of-the-art image compression rates, posing the loss as minimizing distortion at a ﬁxed rate.
Theis et al.
(2017) writes the VAE loss as R + βD .
Rippel & Bourdev (2017) shows that a GAN optimization procedure can also be applied to the problem of compression.
All of these efforts focus on rate/distortion tradeoffs for individual models, but don’t explore how the selection of the model itself affects the rate/distortion curve.
Because we explore many combinations of modeling choices, we are able to more deeply understand how model selection impacts the rate/distortion curve, and to point out the area where all current models are lacking – the autoencoding point.
Generative compression models also have to work with both quantized latent spaces and approximately ﬁxed decoder model families trained with perceptual losses such as MS-SSIM (Wang et al., 2003), which constrain the form of the learned distribution.
Our work does not assume either of these constraints are present for the tasks of interest.
E XPERIMENTS
Toy Model
In this section, we empirically show a case where the usual ELBO objective can learn a model which perfectly captures the true data distribution, p ∗ ( x ) , but which fails to learn a useful latent representation.
However, by training the same model such that we minimize the distortion, subject to achieving a desired target rate R ∗ , we can recover a latent representation that closely matches the true generative process (up to a reparameterization), while also perfectly capturing the true data distribution.
We create a simple data generating process that consists of a true latent variable Z ∗
=
{ z 0 , z 1 } ∼
Ber(0 .
7) with added Gaussian noise and discretization.
The magnitude of the noise was chosen so that the true generative model had I( x ; z ∗ )
= 0 .
5 nats of mutual information between the observations and the latent.
We additionally choose a model family with sufﬁcient power to perfectly autoencode or autodecode.
See Appendix C for more detail on both the data generation and the model.
Figure 2 shows various distributions computed using three models.
For the left column, we use a hand-engineered encoder e ( z | x ) , decoder d ( x | z ) , and marginal m ( z ) constructed with knowledge of the true data generating mechanism to illustrate an optimal model.
For the middle and right columns, we learn e ( z | x ) , d ( x | z ) , and m ( z ) using effectively inﬁnite data sampled from p ∗ ( x ) directly.
The middle column is trained with ELBO.
The right column is trained by targeting R = 0 .
5 while minimizing D .
In both cases, we see that p ∗ ( x ) ≈ g ( x ) ≈ d ( x ) for both trained models, indicating that optimization found the global optimum of the respective objectives.
However, the VAE fails to learn a useful representation, only yielding a rate of R = 0 .
0002 nats, 3 while the target rate model achieves R = 0 .
4999 nats and nearly perfectly reproduces the true generative process.
[Figure]
MNIST.
In this section, we show how our view of unsupervised learning can shed light on various kinds of VAE-like models that have been proposed in the literature.
We use the static binary MNIST dataset originally produced for (Larochelle & Murray, 2011a) 4 .
In appendix E, we show analogous results for the Omniglot dataset (Lake et al., 2015).
We will consider simple and complex variants for the encoder and decoder, and three different types of marginal.
The simple encoder is a CNN with a fully factored 64 dimensional Gaussian for e ( z | x ) ; the more complex encoder is similar, but followed by 4 steps of mean-only Gaussian inverse autoregressive ﬂow (Kingma et al., 2016), with each step implemented as a 3 hidden layer
MADE (Germain et al., 2015) with 640 units in each hidden layer.
The simple decoder is a multilayer deconvolutional network; the more powerful decoder is a PixelCNN++ (Salimans et al., 2017) model.
The simple marginal is a ﬁxed isotropic Gaussian, as is commonly used with VAEs; the more complicated version has a 4 step 3 layer MADE
(Germain et al., 2015) mean-only Gaussian autoregressive ﬂow (Papamakarios et al., 2017).
We also consider the setting in which the marginal uses the VampPrior from (Tomczak & Welling, 2017).
We will denote the particular model combination by the tuple (+ / − , + / − , + / − /v ) , depending on whether we use a simple ( − ) or complex (+) (or ( v ) VampPrior) version for the (encoder, decoder, marginal) respectively.
In total we consider 2 × 2 × 3 = 12 models.
We train them all to minimize the objective in Equation 4.
Full details can be found in Appendix D. Runs were performed at various values of β ranging from 0.1 to 10.0, both with and without KL annealing (Bowman et al., 2015).
RD curve.
Figure 3a show the RD plot for 12 models on the MNIST dataset.
Dashed lines represent the best achieved test ELBO of 80.2 nats, which then sets an upper bound on the true data [Figure] entropy H for the static MNIST dataset.
This implies that any RD value above the dashed line is in principle achievable in a powerful enough model.
The stepwise black curves show the monotonic pareto frontier of achieved RD points across all model families.
Points participating in this curve are denoted with a × .
The grey solid line shows the corresponding convex hull, which we approach closely across all rates.
Strong decoder model families dominate at the lowest and highest rates.
Weak decoder models dominate at intermediate rates.
Strong marginal models dominate strong encoder models at most rates.
Across our model families we appear to be pushing up against an approximately smooth RD curve on both datasets.
Generically, the 12 model families we considered here, arguably a representation of the classes of models considered in the VAE literature in general perform much worse in the auto-encoding limit (bottom right corner) of the RD plane.
Figure 3b shows the same raw data, but where we plot ELBO= R + D versus R .
Here some of the differences between individual model families performances are more easily resolved.
Broadly, models with a deconvolutional decoder perform well at intermediate ~22 nat rates, but quickly suffer large distortion penalties as they move away from that point.
This is perhaps unsurprising considering we trained on the binary MNIST dataset, for which the measured pixel level sampling entropy on the test set is approximately 22 nats.
Models with a powerful autoregressive decoder perform well at low rates, but for values of β ≥ 1 tend to collapse to pure autodecoding models.
With the use of the VampPrior and KL annealing however, β = 1 models can exist at ﬁnite rates of aroundnats.
Comparing model performances in terms of their achieved RD curves gives a much more insightful comparison of relative model performance than simply comparing their marginal data log likelihoods.
In particular, we managed to achieve models with ﬁve-sample IWAE (Burda et al., 2015) estimates below 82 nats (a competitive rate for single layer latent variable models Tomczak & Welling (2017)) for rates spanning from 10 − 4 to 30 nats.
While all of those models have competitive ELBOs or marginal log likelihood, they differ substantially in the tradeoffs they make between rate and distortion.
With the most powerful current models, which include autoregressive decoders and learned marginals, particularly the VampPrior, training with an appropriately selected β is a simple way to achieve any desired rate, which does not require the architectural adjustments needed in other proposals like the variational lossy autoencoder (Chen et al., 2016).
[Figure]
The interaction between latent variables and powerful decoders.
Within any particular model family, we can smoothly move between and explore its performance at varying rates. An illustrative example is shown in Fig. 4, where we study the effect of changing β (using KL annealing from low to high) on the same -+v model, corresponding to a VAE with a simple encoder, a powerful
PixelCNN++ decoder, and a powerful VampPrior marginal.
In Fig. 4a we assess how well the models do at reconstructing their inputs.
We pick an image x at random encode it using z ∼ e ( z | x ) , and then reconstruct it using ˆ x ∼ d ( x | z ) .
When β = 1 . 10 (left column), the model obtains R = 0 .
0004 , D = 80 . 6 , ELBO = 80 .
6 nats.
The tiny rate indicates that the decoder ignores its latent code, and hence the reconstructions are independent of the input x .
For example, when the input is x = 8 , the reconstruction is ˆ x = 3 .
However, the generated images sampled from the decoder look good (this is an example of an auto-decoder).
At the other extreme, when β = 0 .
05 (right column), the model obtains R = 156 , D = 4 . 8 , ELBO=161 nats.
Here the model does an excellent job of auto-encoding, generating nearly pixel perfect reconstructions.
However, samples from this model’s prior, as show on the right, are of very poor quality, reﬂected in the worse ELBO and IWAE values.
At intermediate values, such as β = 1 . 0 , ( R = 6 . 2 , D = 74 . 1 , ELBO=80.3) the model seems to retain semantically meaningful information about the input, such as its class and width of the strokes, but maintains variation in the individual reconstructions.
In particular, notice that the individual “2” sent in is reconstructed as a similar “2” but with a visible loop at the bottom.
This model also has very good generated samples.
This intermediate rate encoding arguably typiﬁes what we want to achieve in unsupervised learning: we have learned a highly compressed representation that retains salient features of the data.
In the third column, the β = 0 .
15 model ( R = 120 .
3 , D = 8 . 1 , ELBO=128)
we have very good reconstructions while still obtaining a good degree of compression.
This model arguably typiﬁes the domain most compression work is interested in, where most perceivable variations in the digit are retained in the compression.
However, at these higher rates the failures of our current architectures to approach their theoretical performance becomes more apparent, as the corresponding ELBO of 128 nats is much higher than the 81 nats we obtain at low rates.
This is also evident in the visual degradation in the generated samples from the model on the right.
D ISCUSSION AND FURTHER WORK
We have shown how the rate-distortion tradeoff can shed light on many different models and objectives that have been proposed in the literature.
It also provides new methods for training VAE-type models which can hopefully advance the state of the art in unsupervised representation learning.
A P ROOFS A. L OWER B OUND ON M UTUAL I NFORMATION
Our lower bound is established by the fact that Kullback-Leibler (KL) divergences are positive semideﬁnite
[Equation]
which implies for any distribution p ( x |
[Equation] [Equation]
A. U PPER B OUND ON M UTUAL I NFORMATION
The upper bound is establish again by the positive semideﬁnite quality of KL divergence
[Equation] [Equation]
A. O PTIMAL M ARGINAL FOR F IXED E NCODER
Here we establish that the optimal marginal approximation p ( z ) , is precisely the marginal distribution of the encoder
[Equation]
Consider the variational derivative of the rate with respect to the marginal approximation:
[Equation]
Where in the last line we have taken the ﬁrst order variation, which must vanish if the total variation is to vanish.
In particular, in order for this variation to vanish, since we are considering an arbitrary δm ( z ) , except for the fact that the integral of this variation must vanish, in order for the ﬁrst order variation in the rate to vanish it must be true that for every value of x, z we have that:
[Equation]
which when normalized gives:
[Equation]
or that the marginal approximation is the true encoder marginal.
A. O PTIMAL D ECODER FOR F IXED E NCODER
Next consider the variation in the distortion in terms of the decoding distribution with a ﬁxed encoding distribution.
[Equation]
[Equation]
Similar to the section above, we took only the leading variation into account, which itself must vanish for the full variation to vanish.
Since our variation in the decoder must integrate to 0, this term will vanish if for ever x, z we have that:
[Equation]
when normalized this gives:
[Equation]
which ensures that our decoding distribution is the correct posterior induced by our data and encoder.
B D ETAILED R ELATED W ORK
Here we expand on the brief related work in Section 3.
B. I NFORMATION T HEORY AND MACHINE LEARNING
Much recent work has leveraged information theory to improve our understanding of machine learning in general, and unsupervised learning speciﬁcally.
In Tishby & Zaslavsky (2015), the authors present theory for the success of supervised deep learning as approximately optimizing the information bottleneck objective, and also theoretically predict a supervised variant of the rate/distortion plane we describe here.
Shwartz-Ziv & Tishby (2017) further proposes that training of deep learning models proceeds in two phases: error minimization followed by compression.
They suggest that the compression phase diffuses the conditional entropy of the individual layers of the model, and when the model has converged, it lies near the information bottleneck optimal frontier on the proposed rate/distortion plane.
In Higgins et al.
the authors motivate the β -VAE objective from a combined neuroscience and information theoretic perspective.
The Higgins et al.
(2017) propose that β should be greater than 1 to properly learn disentangled representations in an unsupervised manner.
Chen et al.
(2016) described the issue of the too-powerful decoder when training standard VAEs, where β
= 1 .
They proposed a bits-back (Hinton & Van Camp, 1993) model to understand this phenomenon, as well as a noise-injection technique to combat it.
Our approach removes the need for an additional noise source in the decoder, and concisely rephrases the problem as ﬁnding the optimal β for the chosen model family, which can now be as powerful as we like, without risk of ignoring the latent space and collapsing to an autodecoding model.
Bowman et al.
(2015) suggested annealing the weight of the KL term of the ELBO ( KL[ q ( z | x ) || p ( z )] ) from 0 to 1 to make it possible to train an RNN decoder without ignoring the latent space.
Sønderby et al.
(2016) applies the same idea to ladder network decoders.
We relate this idea of KL annealing to our optimal rate/distortion curve, and show empirically that KL annealing does not in general attain the performance possible when setting a ﬁxed β or a ﬁxed target rate.
In Achille & Soatto (2016), the authors proposed an information bottleneck approach to the activations of a network, termed Information Dropout, as a form of regularization that explains and generalizes Dropout (Srivastava et al., 2014).
They suggest that, without such a form of regularization, standard SGD training only provides a sufﬁcient statistic, but does not in general provide two other desiderata: minimality and invariance to nuisance factors.
Both of these would be provided by a procedure that directly optimized the information bottleneck.
They propose that simply injecting noise adaptively into the otherwise deterministic function computed by the deep network is sufﬁcient to cause SGD to optimize
(2017) expands on this exploration of sufﬁciency, minimality, and invariance in deep networks.
In particular they propose that architectural bottlenecks and depth both promote invariance directly, and they decompose the standard cross entropy loss used in supervised learning into four terms, including one which they name ‘overﬁtting’, and which, without other regularization, an optimization procedure can easily increase in order to reduce the total loss.
Other recent work explores related theoretical frameworks for unsupervised learning, including Pu et al.
; Hu et al.
; Zhang et al.
B. M ODEL FAMILIES FOR UNSUPERVISED LEARNING WITH NEURAL NETWORKS
Burda et al.
(2015) presented an importance-weighted variant of the VAE objective.
By increasing the number of samples taken from the encoder during training, they are able to tighten the variational lower bound and improve the test log likelihood.
Rezende & Mohamed (2015) proposed to use normalizing ﬂows to approximate the true posterior during inference, in order to overcome the problem of the standard mean-ﬁeld posterior approximation used in VAEs lacking sufﬁcient representational power to model complex posterior distributions.
Normalizing ﬂow permits the use of a deep network to compute a differentiable function with a computable determinant of a random variable and have the resulting function be a valid normalized distribution.
Kingma et al.
(2016) expanded on this idea by introducing inverse autoregressive ﬂow (IAF).
IAF takes advantage of properties of current autoregressive models, including their expressive power and particulars of their Jacobians when inverted, and used them to learn expressive, parallelizeable normalizing ﬂows that are efﬁcient to compute when using high dimensional latent spaces for the posterior.
Autoregressive models have also been applied successfully to the density estimation problem, as well as high quality sample generation.
MADE (Germain et al., 2015)
proposed directly masking the parameters of an autoencoder during generation such that a given unit makes its predictions based solely on the ﬁrst d activations of the layer below.
This enforces that the autoencoder maintains the “autoregressive” property.
In van den Oord et al.
(2016b), the authors presented a recurrent neural network that can autoregressively predict the pixels of an image, as well as provide tractable density estimation.
This work was expanded to a convolutional model called PixelCNN (van den
Oord et al., 2016a), which enforced the autoregressive property by masking the convolutional ﬁlters.
In Salimans et al.
, the authors further improved the performance with PixelCNN++ with a collection of architecture changes that allow for much faster training.
Finally, Papamakarios et al.
(2017) proposed another uniﬁcation of normalizing ﬂow models with autoregressive models for density estimation.
The authors observe that the conditional ordering constraints required for valid autoregressive modeling enforces a choice which may be arbitrarily incorrect for any particular problem.
In their proposal, Masked Autoregressive Flow (MAF), they explicitly model the random number generation process with stacked MADE layers.
This particular choice means that MAF is fast at density estimation, whereas the nearly identical IAF architecture is fast at sampling.
Tomczak & Welling (2017) proposed a novel method for learning the marginal posterior, m ( z ) (written q ( z ) in that work): learn k pseudo-inputs that can be mixed to approximate any of the true samples x ∼ p ∗ ( x ) .
C T OY M ODEL D ETAILS
Data generation.
The true data generating distribution is as follows.
We ﬁrst sample a latent binary variable, z ∼ Ber(0 .
7) , then sample a latent 1d continuous value from that variable, h | z ∼
N ( h | µ z , σ z ) , and ﬁnally we observe a discretized value, x = discretize( h ; B ) , where B is a set of 30 equally spaced bins.
We set µ z and σ z such that R ∗ ≡
I( x ; z ) = 0 .
5 nats, in the true generative process, representing the ideal rate target for a latent variable model.
Model details.
We choose to use a discrete latent representation with K = 30 values, with an encoder of the form e ( z i | x j ) ∝ − exp[( w ei x j − b ei ) 2 ] , where z is the one-hot encoding of the latent categorical variable, and x is the one-hot encoding of the observed categorical variable.
Thus the encoder has 2 K = 60 parameters.
We use a decoder of the same form, but with different parameters:
d ( x j | z i ) ∝ − exp[( w d i x j − b d i ) 2 ] .
Finally, we use a variational marginal, m ( z i )
=
π i .
Given this, the true joint distribution has the form p e ( x, z ) =
p ∗ ( x ) e ( z | x ) , with marginal m ( z ) =
(cid:80) x
p e ( x, z ) and conditional p e ( x | z ) =
p e ( x, z ) /p
e ( z ) .
D D ETAILS FOR MNIST AND O MNIGLOT E XPERIMENTS
We used the static binary MNIST dataset originally produced for (Larochelle & Murray, 2011a) 5 , and the Omniglot dataset from Lake et al.
; Burda et al.
As stated in the main text, for our experiments we considered twelve different model families corresponding to a simple and complex choice for the encoder and decoder and three different choices for the marginal.
Unless otherwise speciﬁed, all layers used a linearly gated activation function activation function (Dauphin et al., 2016), h ( x )
=
( W 1 x + b 2 ) σ ( W 2 x + b 2 ) .
D. E NCODER ARCHITECTURES
For the encoder, the simple encoder was a convolutional encoder outputting parameters to a diagonal Gaussian distribution.
The inputs were ﬁrst transformed to be between -1 and 1.
The architecture contained 5 convolutional layers, summarized in the format Conv (depth, kernel size, stride, padding), followed by a linear layer to read out the mean and a linear layer with softplus nonlinearity to read out the variance of the diagonal Gaussiann distribution.
• Input (28, 28, 1) • Conv (32, 5, 1, same) • Conv (32, 5, 2, same) • Conv (64, 5, 1, same) • Conv (64, 5, 2, same) • Conv (256, 7, 1,
valid) • Gauss (Linear (64), Softplus (Linear (64)))
For the more complicated encoder, the same 5 convolutional layer architecture was used, followed by 4 steps of
D. D ECODER ARCHITECTURES
The simple decoder was a transposed convolutional network, with 6 layers of transposed convolution, denoted as Deconv (depth, kernel size, stride, padding) followed by a linear convolutional layer parameterizing an independent Bernoulli distribution over all of the pixels:
• Input (1, 1, 64) •
Deconv (64, 7, 1, valid) •
Deconv (64, 5, 1, same) • Deconv (64, 5, 2,
same) • Deconv (32, 5, 1, same) • Deconv (32, 5, 2, same) • Deconv (32, 4, 1, same) • Bernoulli (Linear Conv (1, 5, 1, same))
The complicated decoder was a slightly modiﬁed PixelCNN++ style network (Salimans et al., 2017) 6 .
However in place of the original RELU activation functions we used linearly gated activation functions and used six blocks (with sizes (28 × 28) – (14 × 14) – (7 × 7) – (7 × 7) – (14 × 14) – (28 × 28) ) of two resnet layers in each block.
All internal layers had a feature depth of 64.
Shortcut connections were used throughout between matching sized featured maps.
The 64-dimensional latent representation was sent through a dense lineary gated layer to produce a 784-dimensional representation that was reshaped to (28 × 28 × 1) and concatenated with the target image to produce a (28 × 28 × 2) dimensional input.
The ﬁnal output (of size (28 × 28 × 64) ) was sent through a (1 × 1) convolution down to depth 1.
These were interpreted as the logits for a Bernoulli distribution deﬁned on each pixel.
D. M ARGINAL ARCHITECTURES
We used three different types of marginals.
The simplest architecture (denoted (-)), was just a ﬁxed isotropic gaussian distribution in 64 dimensions with means ﬁxed at 0 and variance ﬁxed at 1.
The complicated marginal (+) was created by transforming the isotropic Gaussian base distribution with 4 layers of mean-only Gaussian autoregressive ﬂow, with each steps location parameters computed using a 3 layer MADE style masked network with 640 units in the hidden layers and relu activations.
This network resembles the architecture used in Papamakarios et al.
The last choice of marginal was based on VampPrior and denoted with (v), which uses a mixture of the encoder distributions computed on a set of pseudo-inputs to parameterize the prior (Tomczak & Welling, 2017).
We add an additional learned set of weights on the mixture distributions that are constrained to sum to one using a softmax function: m ( z ) = (cid:80) Ni =
w i e ( z | φ i ) where N are the number of pseudo-inputs, w are the weights, e is the encoder, and φ are the pseudo-inputs that have the same dimensionality as the inputs.
D. O PTIMIZATION
The models were all trained using the β -VAE objective (Higgins et al., 2017) at various values of β .
No form of explicit regularization was used.
The models were trained with Adam (Kingma & Ba, 2014) with normalized gradients (Yu et al., 2017) for 200 epochs to get good convergence on the training set, with a ﬁxed learning rate of 3 × 10 − 4 for the ﬁrst 100 epochs and a linearly decreasing learning rate towards 0 at the 200th epoch.
E R ESULTS ON OMNIGLOT
Figure 5 plots the RD curve for various models ﬁt to the Omniglot dataset (Lake et al., 2015), in the same form as the MNIST results in Figure 3.
Due to lack of time, these experiments were not as extensive and we only explored β values of 0.9, 1.0, 1.1 to get a sense of the regions in the RD -phase diagram each model architecture performed best at. Qualitatively the story is the same as in the MNIST case. The powerful decoder models with their autoregressive form most naturally sit at very low rates. We were able to obtain ﬁnite rates by means of KL annealing.
Further experiments will help to ﬁll in the details especially as we explore differing β values for these architectures on the Omniglot dataset.
Our best achieved ELBO was at 90.37 nats, set by the ++- model with β = 1 . 0 and KL annealing.
This model obtains R = 0 .
77 , D = 89 .
60 , ELBO = 90 .
and is nearly auto-decoding.
We found 13 models with ELBOs below 91.2 nats ranging in rates from 0.0074 nats to 10.92 nats.
[Figure]