DISCRETE VARIATIONAL AUTOENCODERS
Probabilistic models with discrete latent variables naturally capture datasets composed of discrete classes. However, they are difficult to train efficiently, since backpropagation through discrete variables is generally not possible. We present a novel method to train a class of probabilistic models with discrete latent variables using the variational autoencoder framework, including backpropagation through the discrete latent variables. The associated class of probabilistic models comprises an undirected discrete component and a directed hierarchical continuous component. The discrete component captures the distribution over the disconnected smooth manifolds induced by the continuous component. As a result, this class of models efficiently learns both the class of objects in an image, and their specific realization in pixels, from unsupervised data; and outperforms state-ofthe-art methods on the permutation-invariant MNIST, Omniglot, and Caltech-101 Silhouettes datasets.
1 INTRODUCTION
Unsupervised learning of probabilistic models is a powerful technique, facilitating tasks such as denoising and inpainting, and regularizing supervised tasks such as classification (Hinton et al., 2006; Salakhutdinov & Hinton, 2009; Rasmus et al., 2015). Many datasets of practical interest are projections of underlying distributions over real-world objects into an observation space; the pixels of an image, for example. When the real-world objects are of discrete types subject to continuous transformations, these datasets comprise multiple disconnected smooth manifolds. For instance, natural images change smoothly with respect to the position and pose of objects, as well as scene lighting. At the same time, it is extremely difficult to directly transform the image of a person to one of a car while remaining on the manifold of natural images.
It would be natural to represent the space within each disconnected component with continuous variables, and the selection amongst these components with discrete variables. In contrast, most stateof-the-art probabilistic models use exclusively discrete variables — as do DBMs (Salakhutdinov & Hinton, 2009), NADEs (Larochelle & Murray, 2011), sigmoid belief networks (Spiegelhalter & Lauritzen, 1990; Bornschein et al., 2016), and DARNs (Gregor et al., 2014) — or exclusively continuous variables — as do VAEs (Kingma & Welling, 2014; Rezende et al., 2014) and GANs (Goodfellow et al., 2014).1 Moreover, it would be desirable to apply the efficient variational autoencoder framework to models with discrete values, but this has proven difficult, since backpropagation through discrete variables is generally not possible (Bengio et al., 2013; Raiko et al., 2015).
We introduce a novel class of probabilistic models, comprising an undirected graphical model defined over binary latent variables, followed by multiple directed layers of continuous latent variables. This class of models captures both the discrete class of the object in an image, and its specific continuously deformable realization. Moreover, we show how these models can be trained efficiently using the variational autoencoder framework, including backpropagation through the binary latent variables. We ensure that the evidence lower bound remains tight by incorporating a hierarchical approximation to the posterior distribution of the latent variables, which can model strong correlations. Since these models efficiently marry the variational autoencoder framework with discrete latent variables, we call them discrete variational autoencoders (discrete VAEs).
1Spike-and-slab RBMs (Courville et al., 2011) use both discrete and continuous latent variables.
1.1 VARIATIONAL AUTOENCODERS ARE INCOMPATIBLE WITH DISCRETE DISTRIBUTIONS
Conventionally, unsupervised learning algorithms maximize the log-likelihood of an observed dataset under a probabilistic model. Even stochastic approximations to the gradient of the loglikelihood generally require samples from the posterior and prior of the model. However, sampling from undirected graphical models is generally intractable (Long & Servedio, 2010), as is sampling from the posterior of a directed graphical model conditioned on its leaf variables (Dagum & Luby, 1993).
In contrast to the exact log-likelihood, it can be computationally efficient to optimize a lower bound on the log-likelihood (Jordan et al., 1999), such as the evidence lower bound (ELBO, L(x, θ, φ); Hinton & Zemel, 1994):
L(x, θ, φ) = log p(x|θ)− KL[q(z|x, φ)||p(z|x, θ)], (1) where q(z|x, φ) is a computationally tractable approximation to the posterior distribution p(z|x, θ). We denote the observed random variables by x, the latent random variables by z, the parameters of the generative model by θ, and the parameters of the approximating posterior by φ. The variational autoencoder (VAE; Kingma & Welling, 2014; Rezende et al., 2014; Kingma et al., 2014) regroups the evidence lower bound of Equation 1 as:
L(x, θ, φ) = −KL [q(z|x, φ)||p(z|θ)]︸ ︷︷ ︸ KL term +Eq [log p(x|z, θ)]︸ ︷︷ ︸ autoencoding term . (2)
In many cases of practical interest, such as Gaussian q(z|x) and p(z), the KL term of Equation 2 can be computed analytically. Moreover, a low-variance stochastic approximation to the gradient of the autoencoding term can be obtained using backpropagation and the reparameterization trick, so long as samples from the approximating posterior q(z|x) can be drawn using a differentiable, deterministic function f(x, φ, ρ) of the combination of the inputs, the parameters, and a set of inputand parameter-independent random variables ρ ∼ D. For instance, samples can be drawn from a Gaussian distribution with mean and variance determined by the input, N (m(x, φ), v(x, φ)), using f(x, φ, ρ) = m(x, φ) + √ v(x, φ) · ρ, where ρ ∼ N (0, 1). When such an f(x, φ, ρ) exists,
∂
∂φ Eq(z|x,φ) [log p(x|z, θ)] ≈
1
N ∑ ρ∼D ∂ ∂φ log p(x|f(x, ρ, φ), θ). (3)
The reparameterization trick can be generalized to a large set of distributions, including nonfactorial approximating posteriors. We address this issue carefully in Appendix A, where we find that an analog of Equation 3 holds. Specifically, Di is the uniform distribution between 0 and 1, and
f(x) = F−1(x), (4)
where F is the conditional-marginal cumulative distribution function (CDF) defined by:
Fi(x) = ∫ x x′i=−∞ p (x′i|x1, . . . , xi−1) . (5)
However, this generalization is only possible if the inverse of the conditional-marginal CDF exists and is differentiable.
A formulation comparable to Equation 3 is not possible for discrete distributions, such as restricted Boltzmann machines (RBMs) (Smolensky, 1986):
p(z) = 1 Zp e−Ep(z) = 1 Zp · e(z>Wz+b>z), (6)
where z ∈ {0, 1}n, Zp is the partition function of p(z), and the lateral connection matrix W is triangular. Any approximating posterior that only assigns nonzero probability to a discrete domain corresponds to a CDF that is piecewise-contant. That is, the range of the CDF is a proper subset of the interval [0, 1]. The domain of the inverse CDF is thus also a proper subset of [0, 1], and its derivative is not defined, as required in Equations 3 and 4.2
2This problem remains even if we use the quantile function, F−1p (ρ) = inf { z ∈ R : ∫ z z′=−∞ p(z ′) ≥ ρ } ,
the derivative of which is either zero or infinite if p is a discrete distribution.
In the following sections, we present the discrete variational autoencoder (discrete VAE), a hierarchical probabilistic model consising of an RBM,3 followed by multiple directed layers of continuous latent variables. This model is efficiently trainable using the variational autoencoder formalism, as in Equation 3, including backpropagation through its discrete latent variables.
1.2 RELATED WORK
Recently, there have been many efforts to develop effective unsupervised learning techniques by building upon variational autoencoders. Importance weighted autoencoders (Burda et al., 2016), Hamiltonian variational inference (Salimans et al., 2015), normalizing flows (Rezende & Mohamed, 2015), and variational Gaussian processes (Tran et al., 2016) improve the approximation to the posterior distribution. Ladder variational autoencoders (Sønderby et al., 2016) increase the power of the architecture of both approximating posterior and prior. Neural adaptive importance sampling (Du et al., 2015) and reweighted wake-sleep (Bornschein & Bengio, 2015) use sophisticated approximations to the gradient of the log-likelihood that do not admit direct backpropagation. Structured variational autoencoders use conjugate priors to construct powerful approximating posterior distributions (Johnson et al., 2016).
It is easy to construct a stochastic approximation to the gradient of the ELBO that admits both discrete and continuous latent variables, and only requires computationally tractable samples. Unfortunately, this naive estimate is impractically high-variance, leading to slow training and poor performance (Paisley et al., 2012). The variance of the gradient can be reduced somewhat using the baseline technique, originally called REINFORCE in the reinforcement learning literature (Mnih & Gregor, 2014; Williams, 1992; Mnih & Rezende, 2016), which we discuss in greater detail in Appendix B.
Prior efforts by Makhzani et al. (2015) to use multimodal priors with implicit discrete variables governing the modes did not successfully align the modes of the prior with the intrinsic clusters of the dataset. Rectified Gaussian units allow spike-and-slab sparsity in a VAE, but the discrete variables are also implicit, and their prior factorial and thus unimodal (Salimans, 2016). Graves (2016) computes VAE-like gradient approximations for mixture models, but the component models are assumed to be simple factorial distributions. In contrast, discrete VAEs generalize to powerful multimodal priors on the discrete variables, and a wider set of mappings to the continuous units.
The generative model underlying the discrete variational autoencoder resembles a deep belief network (DBN; Hinton et al., 2006). A DBN comprises a sigmoid belief network, the top layer of which is conditioned on the visible units of an RBM. In contrast to a DBN, we use a bipartite Boltzmann machine, with both sides of the bipartite split connected to the rest of the model. Moreover, all hidden layers below the bipartite Boltzmann machine are composed of continuous latent variables with a fully autoregressive layer-wise connection architecture. Each layer j receives connections from all previous layers i < j, with connections from the bipartite Boltzmann machine mediated by a set of smoothing variables. However, these architectural differences are secondary to those in the gradient estimation technique. Whereas DBNs are traditionally trained by unrolling a succession of RBMs, discrete variational autoencoders use the reparameterization trick to backpropagate through the evidence lower bound.
2 BACKPROPAGATING THROUGH DISCRETE LATENT VARIABLES BY ADDING CONTINUOUS LATENT VARIABLES
When working with an approximating posterior over discrete latent variables, we can effectively smooth the conditional-marginal CDF (defined by Equation 5 and Appendix A) by augmenting the latent representation with a set of continous random variables. The conditional-marginal CDF over the new continuous variables is invertible and its inverse is differentiable, as required in Equations 3 and 4. We redefine the generative model so that the conditional distribution of the observed variables given the latent variables only depends on the new continuous latent space. This does not alter
3Strictly speaking, the prior contains a bipartite Boltzmann machine, all the units of which are connected to the rest of the model. In contrast to a traditional RBM, there is no distinction between the “visible” units and the “hidden” units. Nevertheless, we use the familiar term RBM in the sequel, rather than the more cumbersome “fully hidden bipartite Boltzmann machine.”
the fundamental form of the model, or the KL term of Equation 2; rather, it can be interpreted as adding a noisy nonlinearity, like dropout (Srivastava et al., 2014) or batch normalization with a small minibatch (Ioffe & Szegedy, 2015), to each latent variable in the approximating posterior and the prior. The conceptual motivation for this approach is discussed in Appendix C.
Specifically, as shown in Figure 1a, we augment the latent representation in the approximating posterior with continuous random variables ζ,4 conditioned on the discrete latent variables z of the RBM:
q(ζ, z|x, φ) = r(ζ|z) · q(z|x, φ), where r(ζ|z) = ∏ i r(ζi|zi).
The support of r(ζ|z) for all values of z must be connected, so the marginal distribution q(ζ|x, φ) =∑z r(ζ|z) · q(z|x, φ) has a constant, connected support so long as 0 < q(z|x, φ) < 1. We further require that r(ζ|z) is continuous and differentiable except at the endpoints of its support, so the inverse conditional-marginal CDF of q(ζ|x, φ) is differentiable in Equations 3 and 4, as we discuss in Appendix A.
As shown in Figure 1b, we correspondingly augment the prior with ζ: p(ζ, z|θ) = r(ζ|z) · p(z|θ),
where r(ζ|z) is the same as for the approximating posterior. Finally, we require that the conditional distribution over the observed variables only depends on ζ:
p(x|ζ, z, θ) = p(x|ζ, θ). (7) The smoothing distribution r(ζ|z) transforms the model into a continuous function of the distribution over z, and allows us to use Equations 2 and 3 directly to obtain low-variance stochastic approximations to the gradient.
Given this expansion, we can simplify Equations 3 and 4 by dropping the dependence on z and applying Equation 16 of Appendix A, which generalizes Equation 3:
∂
∂φ Eq(ζ,z|x,φ) [log p(x|ζ, z, θ)] ≈
1
N ∑ ρ∼U(0,1)n ∂ ∂φ log p ( x|F−1q(ζ|x,φ)(ρ), θ ) . (8)
4We always use a variant of z for latent variables. This is zeta, or Greek z. The discrete latent variables z can conveniently be thought of as English z.
If the approximating posterior is factorial, then each Fi is an independent CDF, without conditioning or marginalization.
As we shall demonstrate in Section 2.1, F−1q(ζ|x,φ)(ρ) is a function of q(z = 1|x, φ), where q(z = 1|x, φ) is a deterministic probability value calculated by a parameterized function, such as a neural network. The autoencoder implicit in Equation 8 is shown in Figure 1c. Initially, input x is passed into a deterministic feedforward network q(z = 1|x, φ), for which the final nonlinearity is the logistic function. Its output q, along with an independent random variable ρ ∼ U [0, 1], is passed into the deterministic function F−1q(ζ|x,φ)(ρ) to produce a sample of ζ. This ζ, along with the original input x, is finally passed to log p (x|ζ, θ). The expectation of this log probability with respect to ρ is the autoencoding term of the VAE formalism, as in Equation 2. Moreover, conditioned on the input and the independent ρ, this autoencoder is deterministic and differentiable, so backpropagation can be used to produce a low-variance, computationally-efficient approximation to the gradient.
2.1 SPIKE-AND-EXPONENTIAL SMOOTHING TRANSFORMATION
As a concrete example consistent with sparse coding, consider the spike-and-exponential transformation from binary z to continuous ζ:
r(ζi|zi = 0) = {∞, if ζi = 0 0, otherwise
Fr(ζi|zi=0)(ζ ′) = 1
r(ζi|zi = 1) = { βeβζ
eβ−1 , if 0 ≤ ζi ≤ 1 0, otherwise
Fr(ζi|zi=1)(ζ ′) =
eβζ
eβ − 1
∣∣∣∣ζ ′
0
= eβζ ′ − 1 eβ − 1
where Fp(ζ ′) = ∫ ζ′ −∞ p(ζ) · dζ is the CDF of probability distribution p in the domain [0, 1]. This transformation from zi to ζi is invertible: ζi = 0⇔ zi = 0, and ζi > 0⇔ zi = 1 almost surely.5
We can now find the CDF for q(ζ|x, φ) as a function of q(z = 1|x, φ) in the domain [0, 1], marginalizing out the discrete z:
Fq(ζ|x,φ)(ζ ′) = (1− q(z = 1|x, φ)) · Fr(ζi|zi=0)(ζ ′) + q(z = 1|x, φ) · Fr(ζi|zi=1)(ζ ′) = q(z = 1|x, φ) · ( eβζ
′ − 1 eβ − 1 − 1
) + 1.
To evaluate the autoencoder of Figure 1c, and through it the gradient approximation of Equation 8, we must invert the conditional-marginal CDF Fq(ζ|x,φ):
F−1q(ζ|x,φ)(ρ) =
{ 1 β · log [( ρ+q−1 q ) · ( eβ − 1 ) + 1 ] , if ρ ≥ 1− q
0, otherwise (9)
where we use the substitution q(z = 1|x, φ) → q to simplify notation. For all values of the independent random variable ρ ∼ U [0, 1], the function F−1q(ζ|x,φ)(ρ) rectifies the input q(z = 1|x, φ) if q ≤ 1 − ρ in a manner analogous to a rectified linear unit (ReLU), as shown in Figure 2a. It is also quasi-sigmoidal, in that F−1 is increasing but concave-down if q > 1 − ρ. The effect of ρ on F−1 is qualitatively similar to that of dropout (Srivastava et al., 2014), depicted in Figure 2b, or the noise injected by batch normalization (Ioffe & Szegedy, 2015) using small minibatches, shown in Figure 2c.
Other expansions to the continuous space are possible. In Appendix D.1, we consider the case where both r(ζi|zi = 0) and r(ζi|zi = 1) are linear functions of ζ; in Appendix D.2, we develop a spikeand-slab transformation; and in Appendix E, we explore a spike-and-Gaussian transformation where the continuous ζ is directly dependent on the input x in addition to the discrete z.
5In the limit β →∞, ζi = zi almost surely, and the continuous variables ζ can effectively be removed from the model. This trick can be used after training with finite β to produce a model without smoothing variables ζ.
3 ACCOMMODATING EXPLAINING-AWAY WITH A HIERARCHICAL APPROXIMATING POSTERIOR
When a probabilistic model is defined in terms of a prior distribution p(z) and a conditional distribution p(x|z), the observation of x often induces strong correlations in the posterior p(z|x) due to phenomena such as explaining-away (Pearl, 1988). Moreover, we wish to use an RBM as the prior distribution (Equation 6), which itself may have strong correlations. In contrast, to maintain tractability, many variational approximations use a product of independent approximating posterior distributions (e.g., mean-field methods, but also Kingma & Welling (2014); Rezende et al. (2014)).
To accommodate strong correlations in the posterior distribution while maintaining tractability, we introduce a hierarchy into the approximating posterior q(z|x) over the discrete latent variables. Specifically, we divide the latent variables z of the RBM into disjoint groups, z1, . . . , zk,6 and define the approximating posterior via a directed acyclic graphical model over these groups:
q(z1, ζ1, . . . , zk, ζk|x, φ) = ∏
1≤j≤k r(ζj |zj) · q (zj |ζi<j , x, φ) where
q(zj |ζi<j , x, φ) = egj(ζi<j ,x,φ) >·zj∏ zι∈zj ( 1 + egzι (ζi<j ,x,φ) ) , (10) zj ∈ {0, 1}n, and gj(ζi<j , x, φ) is a parameterized function of the inputs and preceding ζi, such as a neural network. The corresponding graphical model is depicted in Figure 3a, and the integration of such hierarchical approximating posteriors into the reparameterization trick is discussed in Appendix A. If each group zj contains a single variable, this dependence structure is analogous to that of a deep autoregressive network (DARN; Gregor et al., 2014), and can represent any distribution. However, the dependence of zj on the preceding discrete variables zi<j is always mediated by the continuous variables ζi<j .
This hierarchical approximating posterior does not affect the form of the autoencoding term in Equation 8, except to increase the depth of the autoencoder, as shown in Figure 3b. The deterministic probability value q(zj = 1|ζi<j , x, φ) of Equation 10 is parameterized, generally by a neural network, in a manner analogous to Section 2. However, the final logistic function is made explicit in Equation 10 to simplify Equation 12. For each successive layer j of the autoencoder, input x and all previous ζi<j are passed into the network computing q(z = 1|ζi<j , x, φ). Its output qj , along with an
6The continuous latent variables ζ are divided into complementary disjoint groups ζ1, . . . , ζk.
independent random variable ρ ∼ U [0, 1], is passed to the deterministic function F−1q(ζj |ζi<j ,x,φ)(ρ) to produce a sample of ζj . Once all ζj have been recursively computed, the full ζ along with the original input x is finally passed to log p (x|ζ, θ). The expectation of this log probability with respect to ρ is again the autoencoding term of the VAE formalism, as in Equation 2.
In Appendix F, we show that the gradients of the remaining KL term of the ELBO (Equation 2) can be estimated stochastically using:
∂
∂θ KL [q||p] = Eq(z1|x,φ)
[ · · · [ Eq(zk|ζi<k,x,φ) [ ∂Ep(z, θ)
∂θ
]]] − Ep(z|θ) [ ∂Ep(z, θ)
∂θ
] and
(11) ∂
∂φ KL [q||p] = Eρ
[ (g(x, ζ)− b)> · ∂q
∂φ − z> ·W · ( 1− z 1− q
∂q
∂φ
)] . (12)
In particular, Equation 12 is substantially lower variance than the naive approach to calculate ∂ ∂φKL [q||p], based upon REINFORCE.
4 MODELLING CONTINUOUS DEFORMATIONS WITH A HIERARCHY OF CONTINUOUS LATENT VARIABLES
We can make both the generative model and the approximating posterior more powerful by adding additional layers of latent variables below the RBM. While these layers can be discrete, we focus on continuous variables, which have proven to be powerful in generative adversarial networks (Goodfellow et al., 2014) and traditional variational autoencoders (Kingma & Welling, 2014; Rezende et al., 2014). When positioned below and conditioned on a layer of discrete variables, continuous variables can build continuous manifolds, from which the discrete variables can choose. This complements the structure of the natural world, where a percept is determined first by a discrete selection of the types of objects present in the scene, and then by the position, pose, and other continuous attributes of these objects.
Specifically, we augment the latent representation with continuous random variables z,7 and define both the approximating posterior and the prior to be layer-wise fully autoregressive directed graphical models. We use the same autoregressive variable order for the approximating posterior as for the
7We always use a variant of z for latent variables. This is Fraktur z, or German z.
prior, as in DRAW (Gregor et al., 2015), variational recurrent neural networks (Chung et al., 2015), the deep VAE of Salimans (2016), and ladder networks (Rasmus et al., 2015; Sønderby et al., 2016). We discuss the motivation for this ordering in Appendix G.
The directed graphical model of the approximating posterior and prior are defined by: q(z0, . . . , zn|x, φ) = ∏
0≤m≤n q (zm|zl<m, x, φ) and
p(z0, . . . , zn|θ) = ∏
0≤m≤n p (zm|zl<m, θ) . (13)
The full set of latent variables associated with the RBM is now denoted by z0 = {z1, ζ1, . . . , zk, ζk}. However, the conditional distributions in Equation 13 only depend on the continuous ζj . Each zm≥1 denotes a layer of continuous latent variables, and Figure 4 shows the resulting graphical model.
The ELBO decomposes as: L(x, θ, φ) = Eq(z|x,φ) [log p(x|z, θ)]− ∑ m Eq(zl<m|x,φ) [KL [q(zm|zl<m, x, φ)||p(zm|zl<m, θ)]] .
(14)
If both q(zm|zl<m, x, φ) and p(zm|zl<m, θ) are Gaussian, then their KL divergence has a simple closed form, which is computationally efficient if the covariance matrices are diagonal. Gradients can be passed through the q(zl<m|x, φ) using the traditional reparameterization trick, described in Section 1.1.
5 RESULTS
Discrete variational autoencoders comprise a smoothed RBM (Section 2) with a hierarchical approximating posterior (Section 3), followed by a hierarchy of continuous latent variables (Section 4). We parameterize all distributions with neural networks, except the smoothing distribution r(ζ|z) discussed in Section 2. Like NVIL (Mnih & Gregor, 2014) and VAEs (Kingma & Welling, 2014; Rezende et al., 2014), we define all approximating posteriors q to be explicit functions of x, with parameters φ shared between all inputs x. For distributions over discrete variables, the neural networks output the parameters of a factorial Bernoulli distribution using a logistic final layer, as in Equation 10; for the continuous z, the neural networks output the mean and log-standard deviation of a diagonal-covariance Gaussian distribution using a linear final layer. Each layer of the neural networks parameterizing the distributions over z, z, and x consists of a linear transformation,
batch normalization (Ioffe & Szegedy, 2015) (but see Appendix H.2), and a rectified-linear pointwise nonlinearity (ReLU). We stochastically approximate the expectation with respect to the RBM prior p(z|θ) in Equation 11 using block Gibbs sampling on persistent Markov chains, analogous to persistent contrastive divergence (Tieleman, 2008). We minimize the ELBO using ADAM (Kingma & Ba, 2015) with a decaying step size.
The hierarchical structure of Section 4 is very powerful, and overfits without strong regularization of the prior, as shown in Appendix H. In contrast, powerful approximating posteriors do not induce significant overfitting. To address this problem, we use conditional distributions over the input p(x|ζ, θ) without any deterministic hidden layers, except on Omniglot. Moreover, all other neural networks in the prior have only one hidden layer, the size of which is carefully controlled. On statically binarized MNIST, Omniglot, and Caltech-101, we share parameters between the layers of the hierarchy over z. We present the details of the architecture in Appendix H.
We train the resulting discrete VAEs on the permutation-invariant MNIST (LeCun et al., 1998), Omniglot8 (Lake et al., 2013), and Caltech-101 Silhouettes datasets (Marlin et al., 2010). For MNIST, we use both the static binarization of Salakhutdinov & Murray (2008) and dynamic binarization. Estimates of the log-likelihood9 of these models, computed using the method of (Burda et al., 2016) with 104 importance-weighted samples, are listed in Table 1. The reported log-likelihoods for discrete VAEs are the average of 16 runs; the standard deviation of these log-likelihoods are 0.08, 0.04, 0.05, and 0.11 for dynamically and statically binarized MNIST, Omniglot, and Caltech-101 Silhouettes, respectively. Removing the RBM reduces the test set log-likelihood by 0.09, 0.37, 0.69, and 0.66.
We further analyze the performance of discrete VAEs on dynamically binarized MNIST: the largest of the datasets, requiring the least regularization. Figure 5 shows the generative output of a discrete VAE as the Markov chain over the RBM evolves via block Gibbs sampling. The RBM is held constant across each sub-row of five samples, and variation amongst these samples is due to the layers of continuous latent variables. Given a multimodal distribution with well-separated modes, Gibbs sampling passes through the large, low-probability space between the modes only infrequently. As a result, consistency of the digit class over many successive rows in Figure 5 indicates that the RBM prior has well-separated modes. The RBM learns distinct, separated modes corresponding to the different digit types, except for 3/5 and 4/9, which are either nearby or overlapping; at least tens of
8We use the partitioned, preprocessed Omniglot dataset of Burda et al. (2016), available from https://github.com/yburda/iwae/tree/master/datasets/OMNIGLOT.
9The importance-weighted estimate of the log-likelihood is a lower bound, except for the log partition function of the RBM. We describe our unbiased estimation method for the partition function in Appendix H.1.
thousands of iterations of single-temperature block Gibbs sampling is required to mix between the modes. We present corresponding figures for the other datasets, and results on simplified architectures, in Appendix J.
The large mixing time of block Gibbs sampling on the RBM suggests that training may be constrained by sample quality. Figure 6a shows that performance10 improves as we increase the number of iterations of block Gibbs sampling performed per minibatch on the RBM prior: p(z|θ) in Equation 11. This suggests that a further improvement may be achieved by using a more effective sampling algorithm, such as parallel tempering (Swendsen & Wang, 1986).
10All models in Figure 6 use only 10 layers of continuous latent variables, for computational efficiency.
Commensurate with the small number of intrinsic classes, a moderately sized RBM yields the best performance on MNIST. As shown in Figure 6b, the log-likelihood plateaus once the number of units in the RBM reaches at least 64. Presumably, we would need a much larger RBM to model a dataset like Imagenet, which has many classes and complicated relationships between the elements of various classes.
The benefit of the hierarchical approximating posterior over the RBM, introduced in Section 3, is apparent from Figure 6c. The reduction in performance when moving from 4 to 8 layers in the approximating posterior may be due to the fact that each additional hierarchical layer over the approximating posterior adds three layers to the encoder neural network: there are two deterministic hidden layers for each stochastic latent layer. As a result, expanding the number of RBM approximating posterior layers significantly increases the number of parameters that must be trained, and increases the risk of overfitting.
6 CONCLUSION
Datasets consisting of a discrete set of classes are naturally modeled using discrete latent variables. However, it is difficult to train probabilistic models over discrete latent variables using efficient gradient approximations based upon backpropagation, such as variational autoencoders, since it is generally not possible to backpropagate through a discrete variable (Bengio et al., 2013).
We avoid this problem by symmetrically projecting the approximating posterior and the prior into a continuous space. We then evaluate the autoencoding term of the evidence lower bound exclusively in the continous space, marginalizing out the original discrete latent representation. At the same time, we evaluate the KL divergence between the approximating posterior and the true prior in the original discrete space; due to the symmetry of the projection into the continuous space, it does not contribute to the KL term. To increase representational power, we make the approximating posterior over the discrete latent variables hierarchical, and add a hierarchy of continuous latent variables below them. The resulting discrete variational autoencoder achieves state-of-the-art performance on the permutation-invariant MNIST, Omniglot, and Caltech-101 Silhouettes datasets.
ACKNOWLEDGEMENTS
Zhengbing Bian, Fabian Chudak, Arash Vahdat helped run experiments. Jack Raymond provided the library used to estimate the log partition function of RBMs. Mani Ranjbar wrote the cluster management system, and a custom GPU acceleration library used for an earlier version of the code. We thank Evgeny Andriyash, William Macready, and Aaron Courville for helpful discussions; and one of our anonymous reviewers for identifying the problem addressed in Appendix D.3.
A MULTIVARIATE VAES BASED ON THE CUMULATIVE DISTRIBUTION FUNCTION
The reparameterization trick is always possible if the cumulative distribution function (CDF) of q(z|x, φ) is invertible, and the inverse CDF is differentiable, as noted in Kingma & Welling (2014). However, for multivariate distributions, the CDF is defined by:
F (x) = ∫ x1 x′1=−∞ · · · ∫ xn x′n=−∞ p(x′1, . . . , x ′ n).
The multivariate CDF mapsRn → [0, 1], and is generally not invertible.11
In place of the multivariate CDF, consider the set of conditional-marginal CDFs defined by:12
Fi(x) = ∫ x x′i=−∞ p (x′i|x1, . . . , xi−1) . (15)
That is, Fj(x) is the CDF of xj , conditioned on all xi such that i < h, and marginalized over all xk such the j < k. The range of each Fj is [0, 1], so F maps the domain of the original distribution to ρ ∈ [0, 1]n. To invert F, we need only invert each conditional-marginal CDF in turn, conditioning xj = F−1j (ρ) on x1 = F −1 1 (ρ), . . . , xj−1 = F −1 j−1(ρ). These inverses exist so long as the conditional-marginal probabilities are everywhere nonzero. It is not problematic to effectively define F−1j (ρ) based upon xi<j , rather than ρi<j , since by induction we can uniquely determine xi<j given ρi<j .
Using integration-by-substition, we can compute the gradient of the ELBO by taking the expectation of a uniform random variable ρ on [0, 1]n, and using F−1q(z|x,φ) to transform ρ back to the element of z on which p(x|z, θ) is conditioned. To perform integration-by-substitution, we will require the determinant of the Jacobian of F−1.
The derivative of a CDF is the probability density function at the selected point, and Fj is a simple CDF when we hold fixed the variables xi<j on which it is conditioned, so using the inverse function theorem we find:
∂F−1j (ρ)
∂ρj =
1
F ′j(F −1 j (ρ))
= 1 p ( xj = F −1 j (ρ)|xi<j , )
where ρ is a vector, and F ′j is ∂Fj ∂ρj . The Jacobian matrix ∂F∂x is triangular, since the earlier conditionalmarginal CDFs Fj are independent of the value of the later xk, j < k, over which they are marginalized. Moreover, the inverse conditional-marginal CDFs have the same dependence structure as F, so the Jacobian of F−1 is also triangular. The determinant of a triangular matrix is the product of the diagonal elements.
11For instance, for the bivariate uniform distribution on the interval [0, 1]2, the CDF F (x, y) = x · y for 0 ≤ x, y ≤ 1, so for any 0 ≤ c ≤ 1 and c ≤ x ≤ 1, y = c
x yields F (x, y) = c. Clearly, many different pairs
(x, y) yield each possible value c of F (x, y). 12The set of marginal CDFs, used to define copulas, is invertible. However, it does not generally map the original distribution to a simple joint distribution, such as a multivariate unform distribu-
tion, as required for variational autoencoders. In Equation 16, ∣∣∣∣det( ∂F−1q(z|x,φ)(ρ)∂φ )∣∣∣∣ does not cancel out
q ( F−1q(z|x,φ)(ρ)|x, φ ) . The determinant of the inverse Jacobian is instead [∏ i q ( zi = F −1 i (ρ) )]−1, which differs from [ q ( F−1q(z|x,φ)(ρ) )]−1 if q is not factorial. As a result, we do not recover the variational autoencoder formulation of Equation 16.
Using these facts to perform a multivariate integration-by-substitution, we obtain: Eq(z|x,φ) [log p(x|z, θ)] = ∫ z q(z|x, φ) · log p(x|z, θ)
= ∫ 1 ρ=0 q ( F−1q(z|x,φ)(ρ)|x, φ ) · log p ( x|F−1q(z|x,φ)(ρ), θ ) · ∣∣∣∣∣det ( ∂F−1q(z|x,φ)(ρ) ∂ρ )∣∣∣∣∣ =
∫ 1 ρ=0 q ( F−1q(z|x,φ)(ρ)|x, φ ) · log p ( x|F−1q(z|x,φ)(ρ), θ ) · ∏ j ∂F−1qj(zj |x,φ)(ρj) ∂ρj  =
∫ 1 ρ=0
q ( F−1q(z|x,φ)(ρ)|x, φ ) ∏ j q ( zj = F −1 j (ρ)|zi<j
) · log p(x|F−1q(z|x,φ)(ρ), θ) =
∫ 1 ρ=0 log p ( x|F−1q(z|x,φ)(ρ), θ ) (16)
The variable ρ has dimensionality equal to that of z; 0 is the vector of all 0s; 1 is the vector of all 1s.
The gradient with respect to φ is then easy to approximate stochastically:
∂
∂φ Eq(z|x,φ) [log p(x|z, θ)] ≈
1
N ∑ ρ∼U(0,1)n ∂ ∂φ log p ( x|F−1q(z|x,φ)(ρ), θ ) (17)
Note that if q(z|x, φ) is factorial (i.e., the product of independent distributions in each dimension zj), then the conditional-marginal CDFs Fj are just the marginal CDFs in each direction. However, even if q(z|x, φ) is not factorial, Equation 17 still holds so long as F is nevertheless defined to be the set of conditional-marginal CDFs of Equation 15.
B THE DIFFICULTY OF ESTIMATING GRADIENTS OF THE ELBO WITH
REINFORCE
It is easy to construct a stochastic approximation to the gradient of the ELBO that only requires computationally tractable samples, and admits both discrete and continuous latent variables. Unfortunately, this naive estimate is impractically high-variance, leading to slow training and poor performance (Paisley et al., 2012). The variance of the gradient can be reduced somewhat using the baseline technique, originally called REINFORCE in the reinforcement learning literature (Mnih & Gregor, 2014; Williams, 1992; Bengio et al., 2013; Mnih & Rezende, 2016):
∂
∂φ Eq(z|x,φ) [log p(x|z, θ)] = Eq(z|x,φ)
[ [log p(x|z, θ)−B(x)] · ∂
∂φ log q(z|x, φ) ] ≈ 1 N ∑ z∼q(z|x,φ) ( [log p(x|z, θ)−B(x)] · ∂ ∂φ log q(z|x, φ) ) (18)
where B(x) is a (possibly input-dependent) baseline, which does not affect the gradient, but can reduce the variance of a stochastic estimate of the expectation.
In REINFORCE, ∂∂φEq(z|x,φ) [log p(x|z, θ)] is effectively estimated by something akin to a finite difference approximation to the derivative. The autoencoding term is a function of the conditional log-likelihood log p(x|z, θ), composed with the approximating posterior q(z|x, φ), which determines the value of z at which p(x|z, θ) is evaluated. However, the conditional log-likelihood is never differentiated directly in REINFORCE, even in the context of the chain rule. Rather, the conditional log-likelihood is evaluated at many different points z ∼ q(z|x, φ), and a weighted sum of these values is used to approximate the gradient, just like in the finite difference approximation.
Equation 18 of REINFORCE captures much less information about p(x|z, θ) per sample than Equation 3 of the variational autoencoder, which actively makes use of the gradient. In particular, the change of p(x|z, θ) in some direction ~d can only affect the REINFORCE gradient estimate if a sample is taken with a component in direction ~d. In aD-dimensional latent space, at leastD samples are
required to capture the variation of p(x|z, θ) in all directions; fewer samples span a smaller subspace. Since the latent representation commonly consists of dozens of variables, the REINFORCE gradient estimate can be much less efficient than one that makes direct use of the gradient of p(x|z, θ). Moreover, we will show in Section 5 that, when the gradient is calculated efficiently, hundreds of latent variables can be used effectively.
C AUGMENTING DISCRETE LATENT VARIABLES WITH CONTINUOUS LATENT VARIABLES
Intuitively, variational autoencoders break the encoder13 distribution into “packets” of probability of infinitessimal but equal mass, within which the value of the latent variables is approximately constant. These packets correspond to a region ri < ρi < ri + δ for all i in Equation 16, and the expectation is taken over these packets. There are more packets in regions of high probability, so high-probability values are more likely to be selected. More rigorously, Fq(z|x,φ)(ζ) maps intervals of high probability to larger spans of 0 ≤ ρ ≤ 1, so a randomly selected ρ ∼ U [0, 1] is more likely to be mapped to a high-probability point by F−1q(z|x,φ)(ρ).
As the parameters of the encoder are changed, the location of a packet can move, while its mass is held constant. That is, ζ = F−1q(z|x,φ)(ρ) is a function of φ, whereas the probability mass associated with a region of ρ-space is constant by definition. So long as F−1q(z|x,φ) exists and is differentiable, a small change in φ will correspond to a small change in the location of each packet. This allows us to use the gradient of the decoder to estimate the change in the loss function, since the gradient of the decoder captures the effect of small changes in the location of a selected packet in the latent space.
In contrast, REINFORCE (Equation 18) breaks the latent represention into segments of infinitessimal but equal volume; e.g., zi ≤ z′i < zi + δ for all i (Williams, 1992; Mnih & Gregor, 2014; Bengio et al., 2013). The latent variables are also approximately constant within these segments, but the probability mass varies between them. Specifically, the probability mass of the segment z ≤ z′ < z + δ is proportional to q(z|x, φ). Once a segment is selected in the latent space, its location is independent of the encoder and decoder. In particular, the gradient of the loss function does not depend on the gradient of the decoder with respect to position in the latent space, since this position is fixed. Only the probability mass assigned to the segment is relevant.
Although variational autoencoders can make use of the additional gradient information from the decoder, the gradient estimate is only low-variance so long as the motion of most probability packets has a similar effect on the loss. This is likely to be the case if the packets are tightly clustered (e.g., the encoder produces a Gaussian with low variance, or the spike-and-exponential distribution of Section 2.1), or if the movements of far-separated packets have a similar effect on the total loss (e.g., the decoder is roughly linear).
Nevertheless, Equation 17 of the VAE can be understood in analogy to dropout (Srivastava et al., 2014) or standout (Ba & Frey, 2013) regularization. Like dropout and standout, F−1q(z|x,φ)(ρ) is an element-wise stochastic nonlinearity applied to a hidden layer. Since F−1q(z|x,φ)(ρ) selects a point in the probability distribution, it rarely selects an improbable point. Like standout, the distribution of the hidden layer is learned. Indeed, we recover the encoder of standout if we use the spike-andGaussian distribution of Section E.1 and let the standard deviation σ go to zero.
However, variational autoencoders cannot be used directly with discrete latent representations, since changing the parameters of a discrete encoder can only move probability mass between the allowed discrete values, which are far apart. If we follow a probability packet as we change the encoder parameters, it either remains in place, or jumps a large distance. As a result, the vast majority of probability packets are unaffected by small changes to the parameters of the encoder. Even if we are lucky enough to select a packet that jumps between the discrete values of the latent representation,
13Since the approximating posterior q(z|x, φ) maps each input to a distribution over the latent space, it is sometimes called the encoder. Correspondingly, since the conditional likelihood p(x|z, θ) maps each configuration of the latent variables to a distribution over the input space, it is called the decoder.
the gradient of the decoder cannot be used to accurately estimate the change in the loss function, since the gradient only captures the effect of very small movements of the probability packet.
To use discrete latent representations in the variational autoencoder framework, we must first transform to a continuous latent space, within which probability packets move smoothly. That is, we must compute Equation 17 over a different distribution than the original posterior distribution. Surprisingly, we need not sacrifice the original discrete latent space, with its associated approximating posterior. Rather, we extend the encoder q(z|x, φ) and the prior p(z|θ) with a transformation to a continuous, auxiliary latent representation ζ, and correspondingly make the decoder a function of this new continuous representation. By extending both the encoder and the prior in the same way, we avoid affecting the remaining KL divergence in Equation 2.14
The gradient is defined everywhere if we require that each point in the original latent space map to nonzero probability over the entire auxiliary continuous space. This ensures that, if the probability of some point in the original latent space increases from zero to a nonzero value, no probability packet needs to jump a large distance to cover the resulting new region in the auxiliary continuous space. Moreover, it ensures that the conditional-marginal CDFs are strictly increasing as a function of their main argument, and thus are invertible.
If we ignore the cases where some discrete latent variable has probability 0 or 1, we need only require that, for every pair of points in the original latent space, the associated regions of nonzero probability in the auxiliary continuous space overlap. This ensures that probability packets can move continuously as the parameters φ of the encoder, q(z|x, φ), change, redistributing weight amongst the associated regions of the auxiliary continuous space.
D ALTERNATIVE TRANSFORMATIONS FROM DISCRETE TO CONTINUOUS LATENT REPRESENTATIONS
The spike-and-exponential transformation from discrete latent variables z to continuous latent variables ζ presented in Section 2.1 is by no means the only one possible. Here, we develop a collection of alternative transformations.
D.1 MIXTURE OF RAMPS
As another concrete example, we consider a case where both r(ζi|zi = 0) and r(ζi|zi = 1) are linear functions of ζi:
r(ζi|zi = 0) = { 2 · (1− ζi), if 0 ≤ ζi ≤ 1 0, otherwise
Fr(ζi|zi=0)(ζ ′) = 2ζi − ζ2i ∣∣ζ′ 0 = 2ζ ′ − ζ ′2
r(ζi|zi = 1) = { 2 · ζi, if 0 ≤ ζi ≤ 1 0, otherwise
Fr(ζi|zi=1)(ζ ′) = ζ2i ∣∣ζ′ 0 = ζ ′ 2
where Fp(ζ ′) = ∫ ζ′ −∞ p(ζ) · dζ is the CDF of probability distribution p in the domain [0, 1]. The CDF for q(ζ|x, φ) as a function of q(z = 1|x, φ) is:
Fq(ζ|x,φ)(ζ ′) = (1− q(z = 1|x, φ)) · ( 2ζ ′ − ζ ′2 ) + q(z = 1|x, φ) · ζ ′2
= 2 · q(z = 1|x, φ) · ( ζ ′ 2 − ζ ′ ) + 2ζ ′ − ζ ′2. (19)
14Rather than extend the encoder and the prior, we cannot simply prepend the transformation to continuous space to the decoder, since this does not change the space of the probabilty packets.
We can calculate F−1q(ζ|x,φ) explicitly, using the substitutions Fq(ζ|x,φ) → ρ, q(z = 1|x, φ)→ q, and ζ ′ → ζ in Equation 19 to simplify notation:
ρ = 2 · q · (ζ2 − ζ) + 2ζ − ζ2
0 = (2q − 1) · ζ2 + 2(1− q) · ζ − ρ
ζ = 2(q − 1)±
√ 4(1− 2q + q2) + 4(2q − 1)ρ
2(2q − 1)
= (q − 1)±
√ q2 + 2(ρ− 1)q + (1− ρ)
2q − 1 if q 6= 12 ; ρ = ζ otherwise. F−1q(ζ|x,φ) has the desired range [0, 1] if we choose
F−1(ρ) = (q − 1) +
√ q2 + 2(ρ− 1)q + (1− ρ)
2q − 1
= q − 1 +
√ (q − 1)2 + (2q − 1) · ρ
2q − 1 (20)
if q 6= 12 , and F−1(ρ) = ρ if q = 12 . We plot F−1q(ζ|x,φ)(ρ) as a function of q for various values of ρ in Figure 7.
In Equation 20, F−1q(ζ|x,φ)(ρ) is quasi-sigmoidal as a function of q(z = 1|x, φ). If ρ 0.5, F−1 is concave-up; if ρ 0.5, F−1 is concave-down; if ρ ≈ 0.5, F−1 is sigmoid. In no case is F−1 extremely flat, so it does not kill gradients. In contrast, the sigmoid probability of z inevitably flattens.
D.2 SPIKE-AND-SLAB
We can also use the spike-and-slab transformation, which is consistent with sparse coding and proven in other successful generative models (Courville et al., 2011):
r(ζi|zi = 0) = {∞, if ζi = 0 0, otherwise
Fr(ζi|zi=0)(ζ ′) = 1
r(ζi|zi = 1) = { 1, if 0 ≤ ζi ≤ 1 0, otherwise
Fr(ζi|zi=1)(ζ ′) = ζi|ζ
′
0 = ζ ′
where Fp(ζ ′) = ∫ ζ′ −∞ p(ζ) · dζ is the cumulative distribution function (CDF) of probability distribution p in the domain [0, 1]. The CDF for q(ζ|x, φ) as a function of q(z = 1|x, φ) is: Fq(ζ|x,φ)(ζ
′) = (1− q(z = 1|x, φ)) · Fr(ζi|zi=0)(ζ ′) + q(z = 1|x, φ) · Fr(ζi|zi=1)(ζ ′) = q(z = 1|x, φ) · (ζ ′ − 1) + 1.
We can calculate F−1q(ζ|x,φ) explicitly, using the substitution q(z = 1|x, φ)→ q to simplify notation:
F−1q(ζ|x,φ)(ρ) =
{ ρ−1 q + 1, if ρ ≥ 1− q
0, otherwise
We plot F−1q(ζ|x,φ)(ρ) as a function of q for various values of ρ in Figure 8.
D.3 ENGINEERING EFFECTIVE SMOOTHING TRANSFORMATIONS
If the smoothing transformation is not chosen appropriately, the contribution of low-probability regions to the expected gradient of the inverse CDF may be large. Using a variant of the inverse function theorem, we find:
∂
∂θ F (F−1(ρ)) =
∂F
∂θ ∣∣∣∣ F−1(ρ) + ∂F ∂z ∣∣∣∣ F−1(ρ) · ∂ ∂θ F−1(ρ) = 0
p(z) · ∂ ∂θ F−1(ρ) = −∂F ∂θ ∣∣∣∣ z , (21)
where z = F−1(ρ). Consider the case where r(ζi|zi = 0) and r(ζi|zi = 1) are unimodal, but have little overlap. For instance, both distributions might be Gaussian, with means that are many standard deviations apart. For values of ζi between the two modes, F (ζi) ≈ q(zi = 0|x, φ), assuming without loss of generality that the mode corresponding to zi = 0 occurs at a smaller value of ζi than that corresponding to zi = 1. As a result, ∂F∂q ≈ 1 between the two modes, and ∂F −1
∂q ≈ 1r(ζi) even if r(ζi) ≈ 0. In this case, the stochastic estimates of the gradient in equation 8, which depend upon ∂F−1
∂q , have large variance.
These high-variance gradient estimates arise because r(ζi|zi = 0) and r(ζi|zi = 1) are too well separated, and the resulting smoothing transformation is too sharp. Such disjoint smoothing transformations are analogous to a sigmoid transfer function σ(c · x), where σ is the logistic function and c → ∞. The smoothing provided by the continuous random variables ζ is only effective if there is a region of meaningful overlap between r(ζ|z = 0) and r(ζ|z = 1). In particular,∑ zi r(ζi|zi = 0)+ r(ζi|zi = 1) 0 for all ζi between the modes of r(ζi|zi = 0) and r(ζi|zi = 1), so p(z) remains moderate in equation 21. In the spike-and-exponential distribution described in Section 2.1, this overlap can be ensured by fixing or bounding β.
E TRANSFORMATIONS FROM DISCRETE TO CONTINUOUS LATENT REPRESENTATIONS THAT DEPEND UPON THE INPUT
It is not necessary to define the transformation from discrete to continuous latent variables in the approximating posterior, r(ζ|z), to be independent of the input x. In the true posterior distribution,
p(ζ|z, x) ≈ p(ζ|z) only if z already captures most of the information about x and p(ζ|z, x) changes little as a function of x, since
p(ζ|z) = ∫ x p(ζ, x|z) = ∫ x p(ζ|z, x) · p(x|z).
This is implausible if the number of discrete latent variables is much smaller than the entropy of the input data distribution. To address this, we can define:
q(ζ, z|x, φ) = q(z|x, φ) · q(ζ|z, x, φ) p(ζ, z|θ) = p(ζ|z) · p(z|θ)
This leads to an evidence lower bound that resembles that of Equation 2, but adds an extra term:
LV AE(x, θ, φ) = log p(x|θ)− KL [q(z, ζ|x, φ)||p(z, ζ|x, θ)] = log p(x|θ)− KL [q(ζ|z, x, φ) · q(z|x, φ)||p(ζ|z, x, θ) · p(z|x, θ)]
= ∑ z ∫ ζ q(ζ|z, x, φ) · q(z|x, φ) · log [ p(x|ζ, θ) · p(ζ|z, θ) · p(z|θ) q(ζ|z, x, φ) · q(z|x, φ) ] = Eq(ζ|z,x,φ)·q(z|x,φ) [log p(x|ζ, θ)]− KL [q(z|x, φ)||p(z|θ)] − ∑ z q(z|x, φ) · KL [q(ζ|z, x, φ)||p(ζ|z)] . (22)
The extension to hierarchical approximating posteriors proceeds as in sections 3 and 4.
If both q(ζ|z, x, φ) and p(ζ|z) are Gaussian, then their KL divergence has a simple closed form, which is computationally efficient if the covariance matrices are diagonal. However, while the gradients of this KL divergence are easy to calculate when conditioned on z, the gradients with respect of q(z|x, φ) in the new term seem to force us into a REINFORCE-like approach (c.f. Equation 18):∑ z ∂q(z|x, φ) ∂φ ·KL [q(ζ|z, x, φ)||p(ζ|z)] = Eq(z|x,φ) [ KL [q(ζ|z, x, φ)||p(ζ|z)] · ∂ log q(z|x, φ) ∂φ ] . (23) The reward signal is now KL [q(ζ|z, x, φ)||p(ζ|z)] rather than log p(x|z, θ), but the effect on the variance is the same, likely negating the advantages of the variational autoencoder in the rest of the loss function.
However, whereas REINFORCE is high-variance because it samples over the expectation, we can perform the expectation in Equation 23 analytically, without injecting any additional variance. Specifically, if q(z|x, φ) and q(ζ|z, x, φ) are factorial, with q(ζi|zi, x, φ) only dependent on zi, then KL [q(ζ|z, x, φ)||p(ζ|z)] decomposes into a sum of the KL divergences over each variable, as does ∂ log q(z|x,φ)∂φ . The expectation of all terms in the resulting product of sums is zero except those
of the form E [ KL [qi||pi] · ∂ log qi∂φ ] , due to the identity explained in Equation 27. We then use the reparameterization trick to eliminate all hierarchical layers before the current one, and marginalize over each zi. As a result, we can compute the term of Equation 23 by backpropagating
KL [q(ζ|z = 1, x, φ||p(ζ|z = 1)]− KL [q(ζ|z = 0, x, φ)||p(ζ|z = 0)]
into q(z|x, φ). This is especially simple if q(ζi|zi, x, φ) = p(ζi|zi) when zi = 0, since then KL [q(ζ|z = 0, x, φ)||p(ζ|z = 0)] = 0.
E.1 SPIKE-AND-GAUSSIAN
We might wish q(ζi|zi, x, φ) to be a separate Gaussian for both values of the binary zi. However, it is difficult to invert the CDF of the resulting mixture of Gaussians. It is much easier to use a mixture of a delta spike and a Gaussian, for which the CDF can inverted piecewise:
q(ζi|zi = 0, x, φ) = δ(ζi) Fq(ζi|zi=0,x,φ)(ζi) = H(ζi) = { 0, if ζi < 0 1, otherwise q(ζi|zi = 1, x, φ) = N ( µq,i(x, φ), σ 2 q,i(x, φ) ) Fq(ζi|zi=1,x,φ)(ζi) = 1
2
[ 1 + erf ( ζi − µq,i(x, φ)√
2σq,i(x, φ)
)]
where µq(x, φ) and σq(x, φ) are functions of x and φ. We use the substitutions q(zi = 1|x, φ)→ q, µq,i(x, φ)→ µq,i, and σq,i(x, φ)→ σq,i in the sequel to simplify notation. The prior distribution p is similarly parameterized.
We can now find the CDF for q(ζ|x, φ) as a function of q(z = 1|x, φ)→ q: Fq(ζ|x,φ)(ζi) = (1− qi) ·H(ζi)
+ qi 2 · [ 1 + erf ( ζi − µq,i√
2σq,i )] Since zi = 0 makes no contribution to the CDF until ζi = 0, the value of ρ at which ζi = 0 is
ρstepi = qi 2
[ 1 + erf ( −µq,i√ 2σq,i )] so:
ζi =  µq,i + √ 2σq,i · erf−1 ( 2ρi qi − 1 ) , if ρi < ρ step i 0, if ρstepi ≤ ρi ≤ ρstepi + (1− qi) µq,i + √ 2σq,i · erf−1 ( 2(ρi−1) qi + 1 ) , otherwise
Gradients are always evaluated for fixed choices of ρ, and gradients are never taken with respect to ρ. As a result, expectations with respect to ρ are invariant to permutations of ρ. Furthermore,
2ρi qi − 1 = 2(ρ ′ i − 1) qi + 1
where ρ′i = ρi + (1− qi). We can thus shift the delta spike to the beginning of the range of ρi, and use
ζi = { 0, if ρi ≤ 1− qi µq,i + √ 2σq,i · erf−1 ( 2(ρi−1) qi + 1 ) , otherwise
All parameters of the multivariate Gaussians should be trainable functions of x, and independent of q. The new term in Equation 22 is:∑
z q(z|x, φ) · KL [q(ζ|z, x, φ)||p(ζ|z)] =∑ z,i q(zi = 1|x, φ) · KL [q(ζi|zi = 1, x, φ)||p(ζi|zi = 1)]
+ (1− q(zi = 1|x, φ)) · KL [q(ζi|zi = 0, x, φ)||p(ζi|zi = 0)] If zi = 0, then q(ζi|zi = 0, x, φ) = p(ζi|zi = 0, θ), and KL [q(ζi|zi = 0, x, φ)||p(ζi|zi = 0, θ)] = 0 as in Section 2. The KL divergence between two multivariate Gaussians with diagonal covariance matrices, with means µp,i, µq,i, and covariances σ2p,i and σ 2 q,i, is
KL [q||p] = ∑ i
( log σp,i − log σq,i + σ2q,i + (µq,i − µp,i)2
2 · σ2p,i − 1 2 ) To train q(zi = 1|x, φ), we thus need to backpropagate KL [q(ζi|zi = 1, x, φ)||p(ζi|zi = 1)] into it. Finally,
∂KL[q||p] ∂µq,i = µq,i − µp,i σ2p,i ∂KL[q||p] ∂σq,i = − 1 σq,i + σq,i σ2p,i
so
∑ z q(z|x, φ) · ∂ ∂µq,i KL [q||p] = q(zi = 1|x, φ) · µq,i − µp,i
σ2p,i∑ z q(z|x, φ) · ∂ ∂σq,i KL [q||p] = q(zi = 1|x, φ) · ( − 1 σq,i + σq,i σ2p,i )
For p, it is not useful to make the mean values of ζ adjustable for each value of z, since this is redundant with the parameterization of the decoder. With fixed means, we could still parameterize the variance, but to maintain correspondence with the standard VAE, we choose the variance to be one.
F COMPUTING THE GRADIENT OF KL [q(ζ, z|x, φ)||p(ζ, z|θ)]
The KL term of the ELBO (Equation 2) is not significantly affected by the introduction of additional continuous latent variables ζ, so long as we use the same expansion r(ζ|z) for both the approximating posterior and the prior:
KL [q||p] = ∑ z ∫ ζ  ∏ 1≤j≤k r(ζj |zj) · q(zj |ζi<j , x)  · log [∏1≤j≤k r(ζj |zj) · q(zj |ζi<j , x) p(z) ·∏1≤j≤k r(ζj |zj) ]
= ∑ z ∫ ζ  ∏ 1≤j≤k r(ζj |zj) · q(zj |ζi<j , x)  · log [∏1≤j≤k q(zj |ζi<j , x) p(z) ] . (24)
The gradient of Equation 24 with respect to the parameters θ of the prior, p(z|θ), can be estimated stochastically using samples from the approximating posterior, q(ζ, z|x, φ), and the true prior, p(z|θ). When the prior is an RBM, defined by Equation 6, we find:
− ∂ ∂θ KL [q||p] = − ∑ ζ,z q(ζ, z|x, φ) · ∂Ep(z, θ) ∂θ + ∑ z p(z|θ) · ∂Ep(z, θ) ∂θ
= −Eq(z1|x,φ) [ · · · [ Eq(zk|ζi<k,x,φ) [ ∂Ep(z, θ)
∂θ
]]] + Ep(z|θ) [ ∂Ep(z, θ)
∂θ
] (25)
The final expectation with respect to q(zk|ζi<k, x, φ) can be performed analytically; all other expectations require samples from the approximating posterior. Similarly, for the prior, we must sample from the RBM, although Rao-Blackwellization can be used to marginalize half of the units.
F.1 GRADIENT OF THE ENTROPY WITH RESPECT TO φ
In contrast, the gradient of the KL term with respect to the parameters of the approximating posterior is severely complicated by a nonfactorial approximating posterior. We break KL [q||p] into two terms, the negative entropy ∑ z,ζ q log q, and the cross-entropy − ∑ z,ζ q log p, and compute their gradients separately.
We can regroup the negative entropy term of the KL divergence so as to use the reparameterization trick to backpropagate through ∏ i<j q(zj |ζi<j , x):
−H(q) = ∑ z ∫ ζ  ∏ 1≤j≤k r(ζj |zj) · q(zj |ζi<j , x)  · log  ∏ 1≤j≤k q(zj |ζi<j , x)  = ∑ z ∫ ζ ∏ j r(ζj |zj) · q(zj |ζi<j , x)  · ∑ j log q(zj |ζi<j , x)
 = ∑ j ∑ z ∫ ζ ∏ i≤j r(ζi|zi) · q(zi|ζh<i, x)
 · log q(zj |ζi<j , x) = ∑ j Eq(ζi<j ,zi<j |x,φ) ∑ zj q(zj |ζi<j , x) · log q(zj |ζi<j , x)
 = ∑ j Eρi<j ∑ zj q(zj |ρi<j , x) · log q(zj |ρi<j , x)
 (26) where indices i and j denote hierarchical groups of variables. The probability q(zj |ρi<j , x) is evaluated analytically, whereas all variables zi<j and ζi<j are implicitly sampled stochastically via ρi<j .
We wish to take the gradient of −H(q) in Equation 26. Using the identity:
Eq [ c · ∂ ∂φ log q ] = c · ∑ z q · ( ∂q ∂φ /q ) = c · ∂ ∂φ (∑ z q ) = 0 (27)
for any constant c, we can eliminate the gradient of log qj|ρi<j in −∂H(q)∂φ , and obtain:
− ∂ ∂φ H(q) = ∑ j Eρi<j ∑ zj ( ∂ ∂φ q(zj |ρi<j , x) ) · log q(zj |ρi<j , x)  . Moreover, we can eliminate any log-partition function in log q(zj |ρi<j , x) by an argument analogous to Equation 27.15 By repeating this argument one more time, we can break ∂∂φq(zj |ρi<j , x) into its factorial components.16 If zi ∈ {0, 1}, then using Equation 10, gradient of the negative entropy reduces to:
− ∂ ∂φ H(q) = ∑ j Eρi<j ∑ ι∈j ∑ zι qι(zι) · ( zι · ∂gι ∂φ − ∑ zι ( qι(zι) · zι · ∂gι ∂φ )) · (gι · zι)  = ∑ j Eρi<j [ ∂g>j ∂φ · ( gj [ qj(zj = 1)− q2j (zj = 1)
])] where ι and zι correspond to single variables within the hierarchical groups denoted by j. In TensorFlow, it might be simpler to write:
− ∂ ∂φ H(q) = Eρi<j
[ ∂q>j (zj = 1) ∂φ · gj ] .
15∑ z c · ∂q ∂φ = c · ∂ ∂φ ∑ z q = 0, where c is the log partition function of q(zj |ρi<j , x). 16 ∂ ∂φ ∏ i qi = ∑ i ∂qi ∂φ · ∏ j 6=i qj , so the qj 6=i marginalize out of ∂qi ∂φ · ∏ j 6=i qj when multiplied by log qi.
When ∂qi ∂φ · ∏ j 6=i qj is multiplied by one of the log qj 6=i, the sum over zi can be taken inside the ∂ ∂φ
, and again ∂ ∂φ ∑ zi qi = 0.
F.2 GRADIENT OF THE CROSS-ENTROPY
The gradient of the cross-entropy with respect to the parameters φ of the approximating posterior does not depend on the partition function of the prior Zp, since:
− ∂ ∂φ ∑ z q log p = ∑ z ∂ ∂φ q · Ep + ∂ ∂φ q · logZp = ∑ z ∂ ∂φ q · Ep
by Equations 6 and 27, so we are left with the gradient of the average energy Ep. The remaining cross-entropy term is∑ z q · Ep = −Eρ [ z> ·W · z + b> · z ] . We can handle the term b> · z analytically, since zi ∈ {0, 1}, and Eρ [ b> · z ] = b> · Eρ [q(z = 1)] .
The approximating posterior q is continuous, with nonzero derivative, so the reparameterization trick can be applied to backpropagate gradients:
∂ ∂φ Eρ [ b> · z ] = b> · Eρ
[ ∂
∂φ q(z = 1)
] .
In contrast, each element of the sum z> ·W · z = ∑ i,j Wij · zi · zj
depends upon variables that are not usually in the same hierarchical level, so in general Eρ [Wijzizj ] 6=WijEρ [zi] · Eρ [zj ] .
We might decompose this term into Eρ [Wijzizj ] =Wij · Eρk≤i [ zi · Eρk>i [zj ] ] ,
where without loss of generality zi is in an earlier hierarchical layer than zj ; however, it is not clear how to take the derivative of zi, since it is a discontinuous function of ρk≤i.
F.3 NAIVE APPROACH
The naive approach would be to take the gradient of the expectation using the gradient of logprobabilities over all variables:
∂
∂φ E [Wijzizj ] = Eq
[ Wijzizj · ∂
∂φ log q ] = Eq1,q2|1,... [ Wijzizj ·
∑ k ∂ ∂φ log qk|l<k
] (28)
= Eq1,q2|1,... [ Wijzizj · ∑ k 1 qk|l<k · ∂qk|l<k ∂φ ] .
For ∂qk|l<k∂φ , we can drop out terms involving only zi<k and zj<k that occur hierarchically before k, since those terms can be pulled out of the expectation over qk, and we can apply Equation 27. However, for terms involving zi>k or zj>k that occur hierarchically after k, the expected value of zi or zj depends upon the chosen value of zk.
The gradient calculation in Equation 28 is an instance of the REINFORCE algorithm (Equation 18). Moreover, the variance of the estimate is proportional to the number of terms (to the extent that the terms are independent). The number of terms contributing to each gradient ∂qk|l<k∂φ grows quadratically with number of units in the RBM. We can introduce a baseline, as in NVIL (Mnih & Gregor, 2014):
Eq [ (Wijzizj − c(x)) · ∂
∂φ log q
] ,
but this approximation is still high-variance.
F.4 DECOMPOSITION OF ∂∂φWijzizj VIA THE CHAIN RULE
When using the spike-and-exponential, spike-and-slab, or spike-and-Gaussian distributions of sections 2.1 D.2, and E.1, we can decompose the gradient of E [Wijzizj ] using the chain rule. Previously, we have considered z to be a function of ρ and φ. We can instead formulate z as a function of q(z = 1) and ρ, where q(z = 1) is itself a function of ρ and φ. Specifically,
zi(qi(zi = 1), ρi) = { 0 if ρi < 1− qi(zi = 1) = qi(zi = 0) 1 otherwise.
(29)
Using the chain rule, ∂∂φzi = ∑ j ∂zi ∂qj(zj=1)
· ∂qj(zj=1)∂φ , where ∂zi∂qj(zj=1) holds all qk 6=j fixed, even though they all depend on the common variables ρ and parameters φ. We use the chain rule to differentiate with respect to q(z = 1) since it allows us to pull part of the integral over ρ inside the derivative with respect to φ. In the sequel, we sometimes write q in place of q(z = 1) to minimize notational clutter.
Expanding the desired gradient using the reparameterization trick and the chain rule, we find:
∂
∂φ Eq [Wijzizj ] =
∂
∂φ Eρ [Wijzizj ]
= Eρ [∑ k ∂Wijzizj ∂qk(zk = 1) · ∂qk(zk = 1) ∂φ ] . (30)
We can change the order of integration (via the expectation) and differentiation since
|Wijzizj | ≤Wij <∞ for all ρ and bounded φ (Cheng, 2006). Although z(q, ρ) is a step function, and its derivative is a delta function, the integral (corresponding to the expectation with respect to ρ) of its derivative is finite. Rather than dealing with generalized functions directly, we apply the definition of the derivative, and push through the matching integral to recover a finite quantity.
For simplicity, we pull the sum over k out of the expectation in Equation 30, and consider each summand independently. From Equation 29, we see that zi is only a function of qi, so all terms in the sum over k in Equation 30 vanish except k = i and k = j. Without loss of generality, we consider the term k = i; the term k = j is symmetric. Applying the definition of the gradient to one of the summands, and then analytically taking the expectation with respect to ρi, we obtain:
Eρ [ ∂Wij · zi(q, ρ) · zj(q, ρ)
∂qi(zi = 1) · ∂qi(zi = 1) ∂φ ] = Eρ [ lim
δqi(zi=1)→0 Wij · zi(q + δqi, ρ) · zj(q + δqi, ρ)−Wij · zi(q, ρ) · zj(q, ρ) δqi · ∂qi(zi = 1) ∂φ ] = Eρk 6=i [ lim
δqi(zi=1)→0 δqi · Wij · 1 · zj(q, ρ)−Wij · 0 · zj(q, ρ) δqi · ∂qi(zi = 1) ∂φ ∣∣∣∣ ρi=qi(zi=0) ]
= Eρk 6=i [ Wij · zj(q, ρ) · ∂qi(zi = 1)
∂φ
∣∣∣∣ ρi=qi(zi=0) ] .
The third line follows from Equation 29, since zi(q+ δqi, ρ) differs from zi(q, ρ) only in the region of ρ of size δqi around qi(zi = 0) = 1− qi(zi = 1) where zi(q + δqi, ρ) 6= zi(q, ρ). Regardless of the choice of ρ, zj(q + δqi, ρ) = zj(q, ρ).
The third line fixes ρi to the transition between zi = 0 and zi = 1 at qi(zi = 0). Since zi = 0 implies ζi = 0,17 and ζ is a continuous function of ρ, the third line implies that ζi = 0. At the same time, since qi is only a function of ρk<i from earlier in the hierarchy, the term ∂qi∂φ is not affected by the choice of ρi.18 As noted above, due to the chain rule, the perturbation δqi has no effect on other
17We chose the conditional distribution r(ζi|zi = 0) to be a delta spike at zero. 18In contrast, zi is a function of ρi.
qj by definition; the gradient is evaluated with those values held constant. On the other hand, ∂qi∂φ is generally nonzero for all parameters governing hierarchical levels k < i.
Since ρi is fixed such that ζi = 0, all units further down the hierarchy must be sampled consistent with this restriction. A sample from ρ has ζi = 0 if zi = 0, which occurs with probability qi(zi = 0).19 We can compute the gradient with a stochastic approximation by multiplying each sample by 1− zi, so that terms with ζi 6= 0 are ignored,20 and scaling up the gradient when zi = 0 by 1qi(zi=0) :
∂
∂φ E [Wijzizj ] = Eρ
[ Wij ·
1− zi 1− qi(zi = 1) · zj · ∂qi(zi = 1) ∂φ
] . (31)
The term 1−zi1−qi is not necessary if j comes before i in the hierarchy.
While Equation 31 appears similar to REINFORCE, it is better understood as an importanceweighted estimate of an efficient gradient calculation. Just as a ReLU only has a nonzero gradient in the linear regime, ∂zi∂φ effectively only has a nonzero gradient when zi = 0, in which case ∂zi ∂φ ∼ ∂qi(zi=1) ∂φ . Unlike in REINFORCE, we do effectively differentiate the reward, Wijzizj . Moreover, the number of terms contributing to each gradient ∂qi(zi=1)∂φ grows linearly with the number of units in an RBM, whereas it grows quadratically in the method of Section F.3.
G MOTIVATION FOR BUILDING APPROXIMATING POSTERIOR AND PRIOR
HIERARCHIES IN THE SAME ORDER
Intuition regarding the difficulty of approximating the posterior distribution over the latent variables given the data can be developed by considering sparse coding, an approach that uses a basis set of spatially locallized filters (Olshausen & Field, 1996). The basis set is overcomplete, and there are generally many basis elements similar to any selected basis element. However, the sparsity prior pushes the posterior distribution to use only one amongst each set of similar basis elements.
As a result, there is a large set of sparse representations of roughly equivalent quality for any single input. Each basis element individually can be replaced with a similar basis element. However, having changed one basis element, the optimal choice for the adjacent elements also changes so the filters mesh properly, avoiding redundancy or gaps. The true posterior is thus highly correlated, since even after conditioning on the input, the probability of a given basis element depends strongly on the selection of the adjacent basis elements.
These equivalent representations can easily be disambiguated by the successive layers of the representation. In the simplest case, the previous layer could directly specify which correlated set of basis elements to use amongst the applicable sets. We can therefore achieve greater efficiency by inferring the approximating posterior over the top-most latent layer first. Only then do we compute the conditional approximating posteriors of lower layers given a sample from the approximating posterior of the higher layers, breaking the symmetry between representations of similar quality.
H ARCHITECTURE
The stochastic approximation to the ELBO is computed via one pass down the approximating posterior (Figure 4a), sampling from each continuous latent layer ζi and zm>1 in turn; and another pass down the prior (Figure 4b), conditioned on the sample from the approximating posterior. In the pass down the prior, signals do not flow from layer to layer through the entire model. Rather, the input to each layer is determined by the approximating posterior of the previous layers, as follows from Equation 14. The gradient is computed by backpropagating the reconstruction log-likelihood, and the KL divergence between the approximating posterior and true prior at each layer, through this differentiable structure.
19It might also be the case that ζi = 0 when zi = 1, but with our choice of r(ζ|z), this has vanishingly small probability.
20This takes advantage of the fact that zi ∈ {0, 1}.
All hyperparameters were tuned via manual experimentation. Except in Figure 6, RBMs have 128 units (64 units per side, with full bipartite connections between the two sides), with 4 layers of hierarchy in the approximating posterior. We use 100 iterations of block Gibbs sampling, with 20 persistent chains per element of the minibatch, to sample from the prior in the stochastic approximation to Equation 11.
When using the hierarchy of continuous latent variables described in Section 4, discrete VAEs overfit if any component of the prior is overparameterized, as shown in Figure 9a. In contrast, a larger and more powerful approximating posterior generally did not reduce performance within the range examined, as in Figure 9b. In response, we manually tuned the number of layers of continuous latent variables, the number of such continuous latent variables per layer, the number of deterministic hidden units per layer in the neural network defining each hierarchical layer of the prior, and the use of parameter sharing in the prior. We list the selected values in Table 2. All neural networks implementing components of the approximating posterior contain two hidden layers of 2000 units.
On statically binarized MNIST, Omniglot, and Caltech-101 Silhouettes, we further regularize using recurrent parameter sharing. In the simplest case, each p (zm|zl<m, θ) and p (x|z, θ) is a function of ∑ l<m zl, rather than a function of the concatenation [z0, z1, . . . , zm−1]. Moreover, all p (zm≥1|zl<m, θ) share parameters. The RBM layer z0 is rendered compatible with this parameterization by using a trainable linear transformation of ζ, M · ζ; where the number of rows in M is
equal to the number of variables in each zm>0. We refer to this architecture as complete recurrent parameter sharing.
On datasets of intermediate size, a degree of recurrent parameter sharing somewhere between full independence and complete sharing is beneficial. We define the n group architecture by dividing the continuous latent layers zm≥1 into n equally sized groups of consecutive layers. Each such group is independently subject to recurrent parameter sharing analogous to the complete sharing architecture, and the RBM layer z0 is independently parameterized.
We use the spike-and-exponential transformation described in Section 2.1. The exponent is a trainable parameter, but it is bounded above by a value that increases linearly with the number of training epochs. We use warm-up with strength 20 for 5 epochs, and additional warm-up of strength 2 on the RBM alone for 20 epochs (Raiko et al., 2007; Bowman et al., 2016; Sønderby et al., 2016).
When p(x|z) is linear, all nonlinear transformations are part of the prior over the latent variables. In contrast, it is also possible to define the prior distribution over the continuous latent variables to be a simple factorial distribution, and push the nonlinearity into the final decoder p(x|z), as in traditional VAEs. The former case can be reduced to something analogous to the latter case using the reparameterization trick.
However, a VAE with a completely independent prior does not regularize the nonlinearity of the prior; whereas a hierarchical prior requires that the nonlinearity of the prior (via its effect on the true posterior) be well-represented by the approximating posterior. Viewed another way, a completely independent prior requires the model to consist of many independent sources of variance, so the data manifold must be fully unfolded into an isotropic ball. A hierarchical prior allows the data manifold to remain curled within a higher-dimensional ambient space, with the approximating posterior merely tracking its contortions. A higher-dimensional ambient space makes sense when modeling multiple classes of objects. For instance, the parameters characterizing limb positions and orientations for people have no analog for houses.
H.1 ESTIMATING THE LOG PARTITION FUNCTION
We estimate the log-likelihood by subtracting an estimate of the log partition function of the RBM (logZp from Equation 6) from an importance-weighted computation analogous to that of Burda et al. (2016). For this purpose, we estimate the log partition function using bridge sampling, a variant of Bennett’s acceptance ratio method (Bennett, 1976; Shirts & Chodera, 2008), which produces unbiased estimates of the partition function. Interpolating distributions were of the form p(x)β , and sampled with a parallel tempering routine (Swendsen & Wang, 1986). The set of smoothing parameters β in [0, 1] were chosen to approximately equalize replica exchange rates at 0.5. This standard criteria simultaneously keeps mixing times small, and allows for robust inference. We make a conservative estimate for burn-in (0.5 of total run time), and choose the total length of run, and number of repeated experiments, to achieve sufficient statistical accuracy in the log partition function. In Figure 10, we plot the distribution of independent estimations of the log-partition function for a single model of each dataset. These estimates differ by no more than about 0.1, indicating that the estimate of the log-likelihood should be accurate to within about 0.05 nats.
H.2 CONSTRAINED LAPLACIAN BATCH NORMALIZATION
Rather than traditional batch normalization (Ioffe & Szegedy, 2015), we base our batch normalization on the L1 norm. Specifically, we use:
y = x− x xbn = y/ ( |y|+ ) s+ o,
where x is a minibatch of scalar values, x denotes the mean of x, indicates element-wise multiplication, is a small positive constant, s is a learned scale, and o is a learned offset. For the approximating posterior over the RBM units, we bound 2 ≤ s ≤ 3, and −s ≤ o ≤ s. This helps ensure that all units are both active and inactive in each minibatch, and thus that all units are used.
I COMPARISON MODELS
In Table 1, we compare the performance of the discrete variational autoencoder to a selection of recent, competitive models. For dynamically binarized MNIST, we compare to deep belief networks (DBN; Hinton et al., 2006), reporting the results of Murray & Salakhutdinov (2009); importanceweighted autoencoders (IWAE; Burda et al., 2016); and ladder variational autoencoders (Ladder VAE; Sønderby et al., 2016).
For the static MNIST binarization of (Salakhutdinov & Murray, 2008), we compare to Hamiltonian variational inference (HVI; Salimans et al., 2015); the deep recurrent attentive writer (DRAW; Gregor et al., 2015); the neural adaptive importance sampler with neural autoregressive distribution estimator (NAIS NADE; Du et al., 2015); deep latent Gaussian models with normalizing flows (Normalizing flows; Rezende & Mohamed, 2015); and the variational Gaussian process (Tran et al., 2016).
On Omniglot, we compare to the importance-weighted autoencoder (IWAE; Burda et al., 2016); ladder variational autoencoder (Ladder VAE; Sønderby et al., 2016); and the restricted Boltzmann machine (RBM; Smolensky, 1986) and deep belief network (DBN; Hinton et al., 2006), reporting the results of Burda et al. (2015).
Finally, for Caltech-101 Silhouettes, we compare to the importance-weighted autoencoder (IWAE; Burda et al., 2016), reporting the results of Li & Turner (2016); reweighted wake-sleep with a deep sigmoid belief network (RWS SBN; Bornschein & Bengio, 2015); the restricted Boltzmann machine (RBM; Smolensky, 1986), reporting the results of Cho et al. (2013); and the neural adaptive importance sampler with neural autoregressive distribution estimator (NAIS NADE; Du et al., 2015).
J SUPPLEMENTARY RESULTS
To highlight the contribution of the various components of our generative model, we investigate performance on a selection of simplified models.21 First, we remove the continuous latent layers. The resulting prior, depicted in Figure 1b, consists of the bipartite Boltzmann machine (RBM), the smoothing variables ζ, and a factorial Bernoulli distribution over the observed variables x defined via a deep neural network with a logistic final layer. This probabilistic model achieves a log-likelihood of −86.9 with 128 RBM units and −85.2 with 200 RBM units. Next, we further restrict the neural network defining the distribution over the observed variables x given the smoothing variables ζ to consist of a linear transformation followed by a pointwise logistic nonlinearity, analogous to a sigmoid belief network (SBN; Spiegelhalter & Lauritzen, 1990; Neal, 1992). This decreases the negative log-likelihood to −92.7 with 128 RBM units and −88.8 with 200 RBM units.
We then remove the lateral connections in the RBM, reducing it to a set of independent binary random variables. The resulting network is a noisy sigmoid belief network. That is, samples are produced by drawing samples from the independent binary random variables, multiplying by an independent noise source, and then sampling from the observed variables as in a standard SBN. With this SBN-like architecture, the discrete variational autoencoder achieves a log-likelihood of −97.0 with 200 binary latent variables. Finally, we replace the hierarchical approximating posterior of Figure 3a with the factorial approximating posterior of Figure 1a. This simplification of the approximating posterior, in addition to the prior, reduces the log-likelihood to −102.9 with 200 binary latent variables.
21In all cases, we report the negative log-likelihood on statically binarized MNIST (Salakhutdinov & Murray, 2008), estimated with 104 importance weighted samples (Burda et al., 2016).
Figures 11, 12, and 13 repeat the analysis of Figure 5 for statically binarized MNIST, Omniglot, and Caltech-101 Silhouettes. Specifically, they show the generative output of a discrete VAE as the Markov chain over the RBM evolves via block Gibbs sampling. The RBM is held constant across each sub-row of five samples, and variation amongst these samples is due to the layers of continuous latent variables. Given a multimodal distribution with well-separated modes, Gibbs sampling passes through the large, low-probability space between the modes only infrequently. As a result, consistency of the object class over many successive rows in Figures 11, 12, and 13 indicates that the RBM prior has well-separated modes.
On statically binarized MNIST, the RBM still learns distinct, separated modes corresponding to most of the different digit types. However, these modes are not as well separated as in dynamically binarized MNIST, as is evident from the more rapid switching between digit types in Figure 11. There are not obvious modes for Omniglot in Figure 12; it is plausible that an RBM with 128 units could not represent enough well-separated modes to capture the large number of distinct character types in the Omniglot dataset. On Caltech-101 Silhouettes, there may be a mode corresponding to large, roughly convex blobs.
This paper presents a way of training deep generative models with discrete hidden variables using the reparameterization trick. It then applies it to a particular DBN-like architecture, and shows that this architecture achieves state-of-the-art density modeling performance on MNIST and similar datasets. 
The paper is well written, and the exposition is both thorough and precise. There are several appendices which justify various design decisions in detail. I wish more papers in our field would take this degree of care with the exposition!
The log-likelihood results are quite strong, especially given that most of the competitive algorithms are based on continuous latent variables. Probably the main thing missing from the experiments is some way to separate out the contributions of the architecture and the inference algorithm. (E.g., what if a comparable architecture is trained with VIMCO, or if the algorithm is applied to a previously published discrete architecture?)
I’m a bit concerned about the variance of the gradients in the general formulation of the algorithm. See my comment “variance of the derivatives of F^{-1}” below. I think the response is convincing, but the problem (as well as “engineering principles” for the smoothing distribution) are probably worth pointing out in the paper itself, since the problem seems likely to occur unless the user is aware of it. (E.g., my proposal of widely separated normals would be a natural distribution to consider until one actually works through the gradients — something not commonly done in the age of autodiff frameworks.)
Another concern is how many sequential operations are needed for inference in the RBM model. (Note: is this actually an RBM, or a general Boltzmann machine?) The q distribution takes the form of an autoregressive model where the variables are processed one at a time. Section 3 mentions the possibility of grouping together variables in the q distribution, and this is elaborated in detail in Appendix A. But the solution requires decomposing the joint into a product of conditionals and applying the CDFs sequentially. So either way, it seems like we’re stuck handling all the variables sequentially, which might get expensive. 
Minor: the second paragraph of Section 3 needs a reference to Appendix A.