R EVISITING A UXILIARY L ATENT V ARIABLES IN G ENERATIVE M ODELS
A BSTRACT
Extending models with auxiliary latent variables is a well-known technique to increase model expressivity.
Bachman & Precup (2015); Naesseth et al.
Cremer et al.
Domke & Sheldon (2018) show that Importance Weighted Autoencoders (IWAE) (Burda et al., 2015) can be viewed as extending the variational family with auxiliary latent variables.
Similarly, we show that this view encompasses many of the recent developments in variational bounds (Maddison et al., 2017; Naesseth et al., 2018; Le et al., 2017; Yin & Zhou, 2018; Molchanov et al., 2018; Sobolev & Vetrov, 2018).
The success of enriching the variational family with auxiliary latent variables motivates applying the same techniques to the generative model.
We develop a generative model analogous to the IWAE bound and empirically show that it outperforms the recently proposed Learned
Accept/Reject
Sampling algorithm (Bauer & Mnih, 2018) while being substantially easier to implement.
Furthermore, we show that this generative process provides new insights on ranking Noise Contrastive Estimation (Jozefowicz et al., 2016; Ma & Collins, 2018) and Contrastive Predictive Coding (Oord et al., 2018)
I NTRODUCTION
Deep generative models with latent variables have seen a resurgence due to the inﬂuential work by Kingma & Welling (2013); Rezende et al.
(2014) and their success at modeling data such as natural images (Rezende & Mohamed, 2015; Kingma et al., 2016; Chen et al., 2016; Gulrajani et al., 2016), speech and music time-series (Chung et al., 2015; Fraccaro et al., 2016; Krishnan et al., 2015), and video (Babaeizadeh et al., 2017; Ha & Schmidhuber, 2018; Denton & Fergus, 2018).
The power of these models lies in the use of auxiliary latent variables to construct complex marginal distributions from tractable conditional distributions.
While directly optimizing the marginal likelihood of latent variable models is intractable, we can instead maximize a variational lower bound on the likelihood such as the evidence lower bound (ELBO) (Jordan et al., 1999; Blei et al., 2017).
The tightness of the bound is determined by the expressiveness of the variational family (Zellner, 1988).
Recently, there have been many advances in constructing tighter variational lower bounds for latent variable models (e.g., Burda et al.
(2015); Maddison et al.
; Naesseth et al.
Le et al.
(2017); Yin & Zhou (2018); Molchanov et al.
; Sobolev & Vetrov (2018)).
Each bound requires a separate derivation and evaluation, however, and the relationship between bounds is unclear.
We show that these bounds can be viewed as speciﬁc instances of auxiliary variable variational inference (Agakov & Barber, 2004; Ranganath et al., 2016;
Maaløe et al., 2016).
In particular, many partition function estimators can be justiﬁed from an auxiliary latent variable or extended state space view (e.g., Sequential Monte Carlo (Doucet et al., 2001), Hamiltonian Monte Carlo (Neal et al., 2011), Annealed Importance Sampling (Neal, 2001)).
Viewed from this perspective, they can be embedded in the variational family as a choice of auxiliary latent variables.
Based on the general results for auxiliary latent variables, this immediately gives rise to a variational lower bound with a characterization of the tightness of the bound.
Furthermore, this view highlights the implicit (potentially suboptimal) choices made and exposes the reusable components that can be combined to form
The success of augmenting variational distributions with auxiliary latent variables motivates investigating a similar augmentation for generative models.
When augmenting the variational distribution, the natural target distribution is the intractable posterior over the latent variables in the model.
With the generative model, this introduces an extra degree of learnable ﬂexibility (i e., we can learn the unnormalized potential function).
To illustrate this, we develop a latent variable model based on self-normalized importance sampling (Algorithm 1) which can be sampled from exactly and has a tractable lower bound on its log-likelihood.
It interpolates between a tractable proposal distribution and an energy model.
We show that this model is closely related to ranking NCE (Jozefowicz et al., 2016; Ma & Collins, 2018) and suggests a principled objective for training the noise distribution in
NCE.
In summary, our contributions are:
We view recent tighter variational lower bounds through the lens of auxiliary variable variational inference, unifying their analysis and exposing sub-optimal design choices in algorithms such as IWAE.
We apply similar ideas to generative models, developing a new model based on selfnormalized importance sampling which can be ﬁt by maximizing a tractable lower bound on its log-likelihood.
We show that the new model generalizes ranking NCE (Ma & Collins, 2018) and provides a proof that the CPC objective (Oord et al., 2018) is a lower bound on mutual information.
We evaluate the proposed model and ﬁnd it outperforms the recently developed approach in (Bauer & Mnih, 2018) despite being more computationally efﬁcient and simpler to implement.
B ACKGROUND
In this work, we consider learned probabilistic models of data p ( x ) .
Latent variables z allow us to construct complex distributions by deﬁning the likelihood p ( x ) =
(cid:82) p ( x | z ) p ( z ) dz in terms of tractable components p ( z ) and p ( x | z ) .
While marginalizing z is generally intractable, we can instead optimize a tractable lower bound on log p ( x ) using the identity
[Equation]
where q ( z | x ) is a variational distribution and the positive D KL term can be omitted to form a lower bound commonly referred to as the evidence lower bound (ELBO)
(Jordan et al., 1999; Blei et al., 2017).
The tightness of the bound is controlled by how accurately q ( z | x ) models p ( z | x ) , so limited expressivity in the variational family can negatively impact the learned model.
2. A UXILIARY V ARIABLE V ARIATIONAL I NFERENCE (AVVI)
Latent variables can also be used to deﬁne complex variational distributions q .
As before, we deﬁne q ( z | x ) =
(cid:82) q ( z | λ, x ) q ( λ | x ) dλ in terms of tractable conditional distributions q ( z | λ, x ) and q ( λ | x ) .
Agakov & Barber (2004) show that
[Equation]
where r ( λ | z, x ) is a variational distribution meant to model q ( λ | z, x ) , and the identity follows from the fact that q ( z | x ) =
q ( z,λ | x ) q ( λ | z,x ) .
Similar to Eq.
(1), Eq.
(2) shows the gap introduced by using r ( λ | z, x ) to deal with the intractability of q ( z | x ) .
We can form a lower bound on the original
ELBO and thus a lower bound on the log marginal by omitting the positive D KL term.
2. M ONTE C ARLO O BJECTIVES
To tighten the variational bound without explicitly expanding the variational family, Burda et al.
(2015) introduced the importance weighted autoencoder (IWAE) bound,
[Equation]
The IWAE bound reduces to the ELBO when K = 1 , is non-decreasing as K increases, and converges to log p ( x ) as K → ∞ under mild conditions (Burda et al., 2015).
Mnih & Rezende (2016) developed Monte Carlo Objectives (MCOs), which extend this notion to any unbiased estimator
ˆ p ( x ) of p ( x )
by noting that
[Equation]
by Jensen’s inequality.
IWAE is the special case where the unbiased estimator is the K -sample importance sampling estimator.
Maddison et al.
; Naesseth et al.
Le et al.
(2017) investigate MCOs in sequential models based on the unbiased estimator produced by Sequential
Monte Carlo.
Many unbiased estimators can be justiﬁed as performing simple importance sampling on an extended state space (e.g., Hamiltonian Importance Sampling (Neal, 2005), Annealed Importance
Sampling (Neal, 2001), and Sequential Monte Carlo (Doucet et al., 2001; Maddison et al., 2017)).
In other words, we can deﬁne auxiliary variables λ and distributions q ( λ | x ) , q ( z | λ, x ) , r ( λ | z, x ) such that
[Equation]
with z, λ ∼ q ( z, λ | x ) .
It immediately follows that the estimator is unbiased and leads to a variational bound Eq.
A UXILIARY L ATENT V ARIABLES IN V ARIATIONAL F AMILIES
Viewing recent improvements in variational bounds as augmenting variational families with latent variables allows us to apply the tools of auxiliary variable variational inference to understand the tradeoffs and derivation of these algorithms.
This uniﬁed view suggests novel bounds and reveals implicit design choices that may be sub-optimal.
3. I MPORTANCE W EIGHTED A UTO - ENCODERS (IWAE)
First, we explicitly work through an example with the IWAE bound.
Bachman & Precup (2015) introduced the idea of viewing IWAE as auxiliary variable variational inference and Naesseth et al.
(2018); Cremer et al.
Domke & Sheldon (2018) formalized the notion.
Consider the variational family deﬁned by ﬁrst sampling a set of K candidate
z i s from a proposal distribution ˜ q ( z i | x ) , and then sampling z from the empirical distribution composed of atoms located at each z i and weighted proportionally to p ( x, z i ) /
˜ q
( z i | x ) .
In this case, the auxiliary latent variables λ are the locations of the proposal samples z 1: K and the index of the selected sample, i .
Explicitly, let w i = p ( x, z i ) /
˜ q
( z i | x ) .
Then choosing the generalized densities of q and r as
[Equation]
yields the IWAE bound Eq.
(3) when plugged into to Eq.
(2) (see Appendix A for details).
From Eq.
(2), it is clear that IWAE is a lower bound on the standard ELBO for q ( z | x ) and the gap is due to D KL ( q ( z 1: K , i | z, x ) ||
r
( z 1: K , i | z, x )) .
The choice of r ( z 1: K , i | z, x ) in Eq.
(5) was for convenience and is suboptimal.
The optimal choice of r is
[Equation]
Compared to the optimal choice, Eq.
(5) makes the approximation q ( z − i | i, z, x ) ≈ (cid:81) [Equation]
= i ˜
q
( z j | x ) which ignores the inﬂuence of z on z − i and the fact that z − i are not independent given z .
A simple extension could be to learn a factored variational distribution conditional on z
[Equation]
Learning such an r could improve the tightness of the bound, and we plan to explore this in future work.
3. E XTENDED S TATE S PACES
More generally, many of the recent improvements in variational bounds (e.g., (Maddison et al., 2017;
Naesseth et al., 2018; Le et al., 2017; Yin & Zhou, 2018; Molchanov et al., 2018; Sobolev & Vetrov, 2018)) can be viewed as importance sampling on an extended state space.
By making the choice of r explicit, the gap between the bound and the ELBO bound with the marginalized variational distribution is clear and this can reveal novel choices for r .
A UXILIARY L ATENT V ARIABLES IN G ENERATIVE M ODELS
In Section 3.1, we showed how IWAE uses self-normalized importance sampling to expand the family of q .
Analogously, we can develop a generative model based on self-normalized importance sampling.
This model draws samples from a proposal π ( x ) , weights them according to a potential function U ( x ) , and then draws a sample from the empirical distribution formed by the weighted samples.
We deﬁne the self-normalized importance sampling (SNIS) generative process in Algorithmand denote the density of the process by p SNIS ( x ) .
The marginal log-likelihood, log p SNIS ( x ) , can be lower bounded as
[Equation]
for details see Appendix B. To summarize, p SNIS ( x ) can be sampled from exactly and has a tractable lower bound on its log-likelihood.
As K → ∞ , p SNIS ( x ) becomes proportional to π ( x ) exp
( U ( x )) .
For ﬁnite K , p SNIS ( x ) interpolates between the tractable π ( x ) and the energy model π ( x ) exp( U ( x )) .
Interestingly, log π ( x ) only shows up once in the lower bound, and simply lower-bounding it still gives a lower bound on log p SNIS ( x ) .
This expands the class of allowable distributions for the proposal π to include
Variational Autoencoders (VAEs) (Kingma & Welling, 2013; Rezende et al., 2014).
To train the SNIS generative model, we can perform stochastic gradient ascent on Eq.
(6) with respect to the parameters of the proposal distribution π
and the potential function U .
When the data x are continuous, reparameterization gradients can be used to estimate the gradients to the proposal distribution (Rezende et al., 2014; Kingma & Welling, 2013).
When the data are discrete, score function gradient estimators such as REINFORCE (Williams, 1992) or relaxed gradient estimators such as the Gumbel-Softmax (Maddison et al., 2016; Jang et al., 2016) can be used.
Simple importance sampling scales poorly to high dimensions, so it is natural to consider augmenting the generative model with latent variables from Hamiltonian Monte Carlo or more complex samplers.
We are currently exploring this.
4. C ONNECTION WITH RANKING NCE AND CPC
Equation (6) is closely connected with the ranking NCE loss (Ma & Collins, 2018), a popular objective for training energy based models.
In fact, if we consider π ( x ) as our noise distribution p N ( x ) and set U ( x ) =
˜ U ( x ) − log p N ( x ) , then up to a constant, we recover the ranking NCE loss.
The ranking NCE loss is motivated by the fact that it is a consistent objective for any K > 1 when the true data distribution is in our model family.
As a result, it is straightforward to adapt the consistency proof from (Ma & Collins, 2018) to our setting.
Furthermore, our perspective gives a coherent objective for jointly learning the noise distribution and the potential function and shows that the ranking NCE loss can be viewed as a lower bound on the log likelihood of a well-speciﬁed model regardless of whether the true data distribution is in our model family.
Moreover, this distribution provides a novel perspective on Contrastive Predictive Coding (Oord et al., 2018), a recent approach to bounding mutual information for representation learning.
Starting from the well-known variational bound on mutual information due to Barber & Agakov (2003)
[Equation]
x
[Equation]
y [Equation] for a variational distribution q ( x | y ) , we can use the self-normalized importance sampling distribution and choose the proposal to be p ( x ) (i.e., p SNIS ( p,U ) ).
Applying the bound in Eq.
(6), we have
I
[Equation]
This recovers the CPC bound and proves that it is indeed a lower bound on mutual information whereas the justiﬁcation in the original paper relied on approximations.
E XPERIMENTS
We evaluated generative models based on self-normalized importance sampling (SNIS) on a small, synthetic dataset as well as the MNIST dataset.
To provide a competitive baseline, we use the recently developed Learned Accept/Reject Sampling (LARS) model (Bauer & Mnih, 2018).
LARS trains a proposal distribution and an acceptance function (analogous to our potential function), which are used to perform rejection sampling.
The output of the rejection sampling process is the generated sample.
Such a process is attractive because unbiased gradients of its log likelihood can be easily computed without knowing the normalizing constant.
To ensure a sample can be generated in ﬁnite time, LARS truncates the rejection sampling after a set number of steps.
Unfortunately, this change requires estimating a normalizing constant.
In practice, Bauer & Mnih (2018) estimate the normalizing constant using 1024 samples during training and 10 10 samples during evaluation.
Even so, LARS requires additional implementation tricks (e.g., evaluating the target density, using an exponential moving average to estimate the normalizing constant) to ensure successful training, which complicate the implementation and analysis of the algorithm.
On the other hand, SNIS is well-speciﬁed and has a tractable log likelihood lower bound for any K .
As a result, no implementation tricks are necessary to train SNIS models.
Moreover,
SNIS weights and uses all samples instead of choosing a single sample, which we expect to be advantageous.
Target density
LARS approx.
density
SNIS approx.
density  200 1k training steps −2.
−2.
−1. −1.
−0.
Log ­ li k .
l o w e r bound
Log­lik.
lower bound
SNIS
LARS
Method
Static MNIST
Dynamic
[Table]
VAE
[Table]
VAE
[Table]
5. S YNTHETIC DATA
As a preliminary experiment, we reproduce the synthetic data experiment from (Bauer & Mnih, 2018) which models a mixture of Gaussian densities.
The target distribution is a mixture ofequally-weighted Gaussian densities with variance 0 . 01 and means ( x, y ) ∈ {− 1 , 0 , 1 } 2 .
Both
LARS and SNIS used a ﬁxed 2-D N (0 , 1) proposal distribution and a learned acceptance/potential function U ( x ) parameterized by a neural network with 2 hidden layers of size 20 and tanh activations.
For both methods the number of proposal samples drawn, K , was set to 128.
We used batch sizes of 128 and ADAM (Kingma & Ba, 2014) with a learning rate of 3 × 10 − 4 to ﬁt the models.
We plot the resulting densities and log-likelihood lower bounds in Fig.
As expected, SNIS quickly converges to the solution, and the potential function learns to cut out the mass between the mixture modes.
5. MNIST
Next, we evaluated SNIS on modeling the MNIST handwritten digit dataset (LeCun, 1998).
MNIST digits can be either statically or dynamically binarized — for the statically binarized dataset we used the binarization from (Salakhutdinov & Murray, 2008), and for the dynamically binarized dataset we sampled images from Bernoulli distributions with probabilities equal to the continuous values of the images in the original MNIST dataset.
We tested two different model conﬁgurations: a VAE with an SNIS prior, and an SNIS model with a VAE proposal.
In the ﬁrst case, the SNIS prior had a Gaussian proposal distribution, and in the second case, the VAE proposal had a Gaussian prior.
We chose hyperparameters to match the MNIST experiments in Bauer & Mnih (2018).
Speciﬁcally, we parameterized the SNIS potential function by a neural network with two hidden layers of size 100 and tanh activations, and parameterized the VAE observation model by neural networks with two layers of 300 units and tanh activations.
The latent spaces of the VAEs were 50-dimensional, and SNIS’s K was set to 1024.
We also lin- early annealed the weight of the KL term in the ELBO from 0 to 1 over the ﬁrst 1 × 10 5 steps and dropped the learning rate from 3 × 10 − 4 to 1 × 10 − 4 on step 1 × 10 6 .
All models were trained with ADAM (Kingma & Ba, 2014).
In the SNIS model with VAE proposal, we originally used the Straight-Through Gumbel estimator (Jang et al., 2016) to estimate gradients through the discrete samples proposed by the VAE, but found that method performed worse than ignoring those gradients altogether.
We suspect that this may be due to bias in the gradients.
Thus, for the SNIS model with VAE proposal, we report numbers on training runs which ignore those gradients, and we plan to investigate unbiased gradient estimators in future work.
We summarize log-likelihood lower bounds on the test set in Table 1.
We found that SNIS outperformed LARS even though it used only 1024 samples for training and evaluation, whereas LARS used 1024 samples during training and 10 10 samples for evaluation.
D ISCUSSION
In this paper, we viewed recent work on improving variational bounds through the lens of auxiliary variable variational inference.
This perspective allowed us to expose suboptimal choices in existing algorithms such as IWAE, unify analysis of other methods such as ranking NCE and CPC, and derive new methods for generative modeling such as SNIS.
We plan to further develop this view by embedding methods such as Hamiltonian Importance Sampling and Annealed Importance Sampling in generative models which we expect to scale better with dimension of the data space.
A PPENDICES A IWAE BOUND AS AVVI PROOF SKETCH
We provide a sketch of a proof that the IWAE bound can be interpreted as auxiliary variable variational inference by choosing speciﬁc values for q and r .
Recall the auxiliary variable variational inference bound,
[Equation]
Let q and r be
[Equation]
Then, plugging Eqs.
(8) and (9) into Eq.
(7) with λ =
( z 1: K , i ) gives
[Equation]
which is the IWAE bound.
B D ERIVATION OF LOWER BOUND ON SNIS DENSITY
We denote the density of the SNIS sampling process Algorithm 1 as p SNIS ( x ) .
Starting from the deﬁnition of p SNIS ( x ) , we obtain Eq.
(6) as
[Equation]