Title: TRAINING VARIATIONAL AUTO ENCODERS WITH DIS-CRETE LATENT REPRESENTATIONS USING IMPORTANCE SAMPLING

Abstract: The Variational Auto Encoder (VAE) is a popular generative latent variable model that is often applied for representation learning. Standard VAEs assume continuous valued latent variables and are trained by maximization of the evidence lower bound (ELBO). Conventional methods obtain a differentiable estimate of the ELBO with reparametrized sampling and optimize it with Stochastic Gradient Descend (SGD). However, this is not possible if we want to train VAEs with discrete valued latent variables, since reparametrized sampling is not possible. Till now, there exist no simple solutions to circumvent this problem. In this paper, we propose an easy method to train VAEs with binary or categorically valued latent representations. Therefore, we use a differentiable estimator for the ELBO which is based on importance sampling. In experiments, we verify the approach and train two different VAEs architectures with Bernoulli and Categorically distributed latent representations on two different benchmark datasets.

1. THE VARIATIONAL AUTO ENCODER:

The variational auto encoder (VAE) is a generative model which it is trained to approximate the true data generating distribution p(x) of an observed random vector x from a given training set D = {x 1 , ..., x N } (Kingma & Welling (2013) ; Kingma et al. (2016) ). It is an especially suited model if x is high dimensional or has highly nonlinear dependent elements. Therefore, the VAE is oftenly used for tasks like density estimation, data generation, data interpolation (White (2016) ), outlier and anomaly detection (An & Cho (2015) ; Xu et al. (2018) ) or clustering (Jiang et al. (2016) ; Dilokthanakul et al. (2016) ).

As shown in Fig. 1 , the VAE is an easy latent variable model, where the observations x ∼ p(x|z) are dependent on latent variables z ∼ p(z).

x z During training, the VAE maximizes the probability p(x) to observe the data x. Therefore, the negative evidence lower bound (ELBO) L(θ) = -E q(z|x) [ln p(x|z) ] + D KL (q(z|x)||p(z))

(1) ≥ -ln p(x) + D KL (q(z|x)||p(z|x))

(2) is minimized, where p(z|x) = p(x|z)p(z)/ p(x|z)p(z)dz is the true but intractable posterior distribution the model assigns to z, q(z|x) is the corresponding tractable variational approximation and D KL (q(z|x)||p(z|x)) is the Kullback-Leibler (KL) divergence between p(z|x) and q(z|x). Because D KL (q(z|x)||p(z|x)) > 0, minimizing L(θ) means to maximize the probability p(x) the model assigns to observations x. Therefore, D KL (q(z|x)||p(z|x)) must be as as close as possible to 0, meaning that after training q(z|x) is a very good approximation of the true posterior p(z|x). Kingma & Welling (2013) proposed to minimize L(θ), using stochastic gradient descent on a training data set, which they called Stochastic Gradient Variational Bayes (SGVB).

The VAE uses parametric distributions that are parametrized by an encoder network with parameters θ E and a decoder network with parameters θ D for both q(z|x) and p(x|z), respectively. This leads to the well known encoder-decoder structure in Fig. 2 . The data likelihood is a distribution with mean x, that is the output of the decoder network. Further, we assume in this paper, that the variational posterior q(z|x) is a distribution from the exponential family q(z|x) = exp(η T (x; θ E )T (z) -A(η(x; θ E )))

(3) with natural parameters η(x; θ E ), sufficient statistic T (z) and log partition function A(η(x; θ E )). This gives us the flexibility to study training with different q(z|x) in the same mathematical framework. As shown in Fig. 2 , the natural parameters η are the output of the encoder network, where we drop the arguments x, θ E for shorter notations in the remainder of the paper.

θ E θ D x x z η

Encoder Decoder:

Figure 2 : The encoder-decoder structure of a VAE. The encoder parametrizes q(z|x) as an exponential family distribution with natural parameters η and the decoder parametrizes p(x|z) with mean x.

The conventional VAE proposed in (Kingma & Welling (2013) ; Kingma et al. ( 

EQUATION

T and A(η) is chosen such that q(z|x) integrates to one. The likelihood is also Gaussian, with p(x|z) ∼ N (x, 1). But in many applications learning discrete rather than continuous representations is advantageous. Binary representations z ∈ {0, 1} c can for example be used very efficiently for hashing, what is a powerful method for large-scale visual search (Liong et al. (2015) ). Learning Categorical representations z ∈ {e 1 , ..., e c } is interesting, because this naturally lead to clustering of the data x, as shown in the experiments. Further, for both binary and categorical z it is easy to find entropy based heuristics to choose the size of the latent space, because the entropy is bounded for discrete z.

However, training VAEs with discrete latent representations is problematic, since standard SGVB can not be applied for optimization. Because SGVB is a gradient based method, we need to calculate the derivative of the two cost terms with respect to the encoder and decoder parameters

EQUATION

where L KL (θ) only depends on the encoder parameters and the expected log likelihood term L L (θ) depends on both encoder and decoder parameters. For a suited choice of p(z) and q(z|x), L KL (θ) can be calculated in closed form. However, L L (θ) contains an expectation over z ∼ q(z|x) that has to be estimated during training. A good estimator LL (θ) for L D (θ) that is unbiased, differentiable with respect to θ and that has low variance is the key to train VAEs. SGVB uses an estimator LR L (θ) that is based on reparametrization of q(z|x) and sampling (Kingma & Welling (2013) ). However, as described in section 2, this method places many restrictions on the form of q(z|x) and fails if q(z|x) can not be reparametrized. This is the case if z is discrete, for example.

In this paper, we propose a simple and differentiable estimator LI L (θ) for L L (θ) that is based on importance sampling. Because no reparametrization is needed, it can be used to train VAEs with binary or categorical latent representations. Compared to previously proposed methods like the Vector Quantised-Variational Auto Encoder (VQ-VAE) (van den Oord et al. ( 2017)), which is based on a straight-through estimator for the gradient of L L (θ) (Bengio et al. (2013) ), our proposed estimator has two advantages. It is unbiased and its variance approaches zero the closer we are to the optimum.

2. ESTIMATING THE EXPECTED LOG-LIKELIHOOD WITH REPARAMETRIZED:



SAMPLING:

The standard estimator LR L (θ) proposed in Kingma & Welling (2013) is based on reparametrized sampling

EQUATION

where is a random variable with the distribution p( ), m are samples from this distribution and f ( , θ) is a reparametrization function, such that z = f ( , θ) ∼ q(z|x). This estimator can be used to train VAEs with SGVB if two conditions are fulfilled: I) There exists a distribution p( ) and a reparametrization function f ( , θ), such that z = f ( , θ) ∼ q(z|x). II) The derivative of Eq. 5 must exist. With Eq. 8, we obtain

EQUATION

meaning that both the reparametrization function f ( , θ) and ln p(x|z) must be differentiable with respect to z and θ, respectively, to allow direct backpropagation of the gradient through the reparametrized sampling operator. If these conditions are fulfilled, the gradient can flow directly from the output to the input layer of the VAE, as shown in Fig. 3 . Distributions over discrete latent representations z can not be reparametrized this way. Therefore, this estimator can not be used to train VAEs with such representations.

θ E θ D x x z η Encoder Decoder ∂ ∂z LR L (θ) ∂ ∂θ z

Figure 3 : The gradient flow through the VAE, using LR L (θ) based on reparametrized sampling. The gradient is propagated directly through the reparametrized sampling operator.

3. ESTIMATING THE EXPECTED LOG-LIKELIHOOD WITH IMPORTANCE:



We propose an estimator LI L (θ) which is based on importance sampling and can also be used to train VAEs with binary or categorical latent representations z. Expanding Eq. 4 leads to

EQUATION

where q I (z) is an arbitrary distribution that is of the same form as q(z|x) which is independent from the parameters θ. z m ∼ q I (z) are samples from this distribution. The estimator computes a weighted sum of the log likelihood ln p(x|z m ) with the weighting q(z m |x)/q I (z).

The benefit is that the log likelihood ln p(x|z m ) depends on the decoder parameters θ D only and not on the encoder parameters θ E whereas the weighting q(z m |x)/q I (z) depends only on θ D and not on θ E . Therefore, calculation of the gradient of LI L (θ) can be separated

EQUATION

with

EQUATION

)

As shown in Fig. 4 , gradient backpropagation is split into two separate parts. ∂ ∂θ D LI L (θ) backpropagates the error ln p(x|z) from the output of the VAE to the sampling operator and ∂ ∂θ E LI L (θ) backpropagates the error q(z|x) q I (z) from the sampling operator to the input layer of the VAE.

θ E θ D x x z η Encoder Decoder ln p(x|z) ∂ ∂θ E q(z|x) q I (z) q(z|x) q I (z) ∂ ∂θ D ln p(x|z)

Figure 4 : Gradient flow through the VAE when using LI L (θ), based on importance sampling.

Compared to LR L (θ), we do not need to find a differentiable reparametrization for q(z|x), because we do not propagate the gradient through the sampling operator. Therefore, LI L (θ) can also be used if no reparametrization function exists for q(z|x), e.g. if it is a Bernoulli or a Categorical distribution.

4. VAE WITH BERNOULLI DISTRIBUTED z (BVAE):

Assume the latent representation z has i.i.d. components that are Bernoulli distributed, i.e. both the variational posterior distribution q(z|x) and q I (z) have the form

EQUATION

where z ∈ {0, 1} c , η = [ln(q 1 /(1 -q 1 )), ..., ln(q c /(1 -q c ))] is the output vector of the encoder that contains the logits of the independent Bernoulli distributions and A(η) = 1 T ln(1 + e η ) are the corresponding log-partition functions.

Hence, Eq. 17 is

∂ ∂θ E LI L (θ) = 1 M M m=1 ln p(x|z m ) ∂ ∂η exp (η -ξ) T z m -(A(η) -A(ξ)) ∂ ∂θ E η = 1 M M m=1 ln p(x|z m ) exp (η -ξ) T z m -(A(η) -A(ξ)) (z -q) T ∂ ∂θ E η ,

where q = [q 1 , ..., q c ] T = 1 1+e -η contains the probabilities q(z i = 1|x). The variance of the estimator LI L (θ) heavily depends on the choice of the natural parameters ξ of the distribution q I (z). We choose ξ = η, leading to a gradient of the very simple form

EQUATION

This estimator of the gradient has two desirable properties for training, which can be easily seen in the one dimensional case with z ∈ {0, 1}. The mean of the estimator is

EQUATION

meaning that the estimator is unbiased. Further, the variance of ∂ ∂θ E LI L (θ) reduces to zero, the closer q is to 0 or 1, because q(1 -q) → 0 and hence the variance of the estimator approaches 0.

That is desirable, since there are only three interesting cases during training that are shown in Fig. 5: (a) ln p(x|z = 0) = ln p(x|z = 1): In this case the expected log likelihood E q(z|x) [ln p(x|z = 0)] = ln p(x|z = 0) we want to maximize is independent of the choice of q.

(b) ln p(x|z = 0) > ln p(x|z = 1): In this case the expected log likelihood E q(z|x) [ln p(x|z = 0)] = (1 -q) ln p(x|z = 0) + q ln p(x|z = 1) is maximized for q = 0.

(c) ln p(x|z = 0) < ln p(x|z = 1): In this case the expected log likelihood E q(z|x) [ln p(x|z = 0)] = (1 -q) ln p(x|z = 0) + q ln p(x|z = 1) is maximized for q = 1.

This means, that the only candidate points that maximize the log likelihood are q = 0 or q = 1 lie near q(z = 1|x) = 0/1. Therefore, the longer we train, the more accurate the gradient estimate will be. ) ] is maximized for q = 1.

z ln p(x|z) z = 0 z = 1 (a) z ln p(x|z) z = 0 z = 1 (b) z ln p(x|z) z = 0 z = 1 (c)

E q(z|x) [ln(x|z)] is independent of q. (b) ln(x|z = 0) > ln(x|z = 1), meaning E q(z|x) [ln(x|z)] is maximized for q = 0. (c) ln(x|z = 0) < ln(x|z = 1), meaning E q(z|x) [ln(x|z

5. VAE WITH CATEGORICALLY DISTRIBUTED z (CVAE):

For Categorically distributed z both the variational posterior distribution q(z|x) and q I (z), again have the form With the formulas above, we arrive at the same easy form of the expected gradient of the log likelihood

EQUATION

EQUATION

but now with q = softmax(η) that consists of the probabilities q i = q(z = e i ), where c i=1 q i = 1.

6. EXPERIMENTS:

In the following section, we show our preliminary experiments on the MNIST and Fashion MNIST datasets LeCun & Cortes (2010); Xiao et al. (2017) . Two different kinds of VAEs have been evaluated:

1. The BVAE with Bernoulli distributed z ∈ {0, 1} c . 2. The CVAE with Categorically distributed z ∈ {e 1 , ..., e c }.

To train both architectures, the estimator LI L (θ) derived in Sec. 5 is used. Both BVAE and CVAE are tested with two different architectures given in Tab. 1. The fully connected architecture has 2 dense encoder and decoder layers. The encoder and decoder networks of the convolutional architecture consist of 4 convolutional layers and one dense layer each.

In our first experiment we train a FC BVAE with c = 50, i.e. z ∈ {0, 1} 50 and a FC CVAE with c = 100, i.e. z ∈ {z 1 , ..., z 100 }. We train them for 300 epochs on the MNIST dataset, using 

∈ R c ∼ Ber( 1 1+e -η ) FC 1024/ReLu or FC CVAE FC c/linear or z ∈ R c ∼ Cat( e η e 1 T η ) FC 784/sigmoid CNN BVAE 28x28x1 Conv 3x3x32/2/ReLu z ∈ R c ∼ Ber( 1 1+e -η ) FC c/ReLu or CNN CVAE Conv 3x3x64/2/ReLu or z ∈ R c ∼ Cat( e η e 1 T η ) reshape Conv 3x3x64/2/ReLu Conv -1 3x3x64/2/ReLu Conv 3x3x128/2/ReLu Conv -1 3x3x64/2/ReLu flatten Conv -1 3x3x32/2/ReLu FC c/linear

Conv -1 3x3x1/2/sigmoid SGVB with our proposed estimator LI L (θ), to estimate the expected log likelihood, and ADAM as optimizer. Fig. 6 shows the convergence of the loss, the log likelihood the VAEs assign to the training data ln p(x|z) and the variance of the estimator LI L (θ), for a learning rate of 1e -3 and a batch size of 2048. During training, the loss decreases steadily without oscillation. We observe that the variance of the estimator LI L (θ) decreases the longer we train and the closer we get to the optimum. This is consistent with our theoretically considerations in Sec. 4. The results of the corresponding simulations with the CNN BVAE and the CNN CVAE are shown in Fig. 9 , in the appendix. Fig. 7 shows handwritten digits that are generated by the FC BVAE and the FC CVAE if we sample z from the variational posterior q(z|x). To draw samples from q(z|x), we feed test data which has not been seen during training to the encoders. The test data is shown in Fig. 7a and Fig. 7c . The corresponding reconstructions generated by the decoders are shown in Fig. 7b and Fig. 7d . Both input and reconstructed images are very similar in case of the FC BVAE, meaning that it can approximate the data generating distribution p(x) well. However, in case of the FC CVAE, the generated are blurry and look very different than the input of the encoder. In some cases, the class of the generated digit is even flipped. This happens because of the the very limited model capacity.

Similar results for the CNN BVAE and CNN CVAE are shown in Fig. 10 in the appendix.

Fig. 8a shows generated images of the FC BVAE if we sample z ∼ p(z) from the prior distribution. A few generated images look like templates of handwritten digits and the remaining generated images seem to resemble mixtures of different digits. This is similar to the behaviour of a VAE with continuous latent variables, where we can interpolate between or generate mixtures of different Since the FC CVAE can only learn to generate 100 different images, its decoder learns to generate template images that fit well to all the training images. We observe that some latent representations are decoded to meaningless patterns that just fit well to the data in avarage. However, the decoder also learned to generate at least one template image for each class of handwritten digits. Hence, the categorical latent representation can be interpreted as the cluster affiliation and the encoder of the FC CVAE automatically learns to cluster the data. Similar results for the CNN BVAE and CNN CVAE are shown in Fig. 11 in the appendix.

A major drawback of the FC CVAE is, that the latent space of the FC CVAE can encode only very little information and thus its generative capabilites are poor. However, we think that they can be increased considerably if we allow a hybrid latent space with some continuous latent variables, as proposed in Chen et al. (2016) . This could lead to a powerfull model for nonlinear clustering. 

7. CONCLUSION:

In this paper, we derived an easy estimator for the ELBO, which does not rely on reparametrized sampling and therefore can be used to obtain differentiable estimates, even if reparametrization is not possible, e.g. if the latent variables z are Bernoulli or Categorically distributed. We have shown theoretically and in experiments, close to the optimal parameter configuration, the variance of the estimator approaches zero. This is a very desirable property for training. 

:



Figures and Tables:

Figure 1: The latent variable model of a VAE with latent variables z and observations x.

Figure 5: Three different cases how the log-likelihood can vary over z. The decoder network defines ln(x|z) over z ∈ R. However, only the points z = 0 and z = 1 are of interest. (a) ln(x|z = 0) = ln(x|z = 1), meaning E q(z|x) [ln(x|z)] is independent of q. (b) ln(x|z = 0) > ln(x|z = 1), meaning E q(z|x) [ln(x|z)] is maximized for q = 0. (c) ln(x|z = 0) < ln(x|z = 1), meaning E q(z|x) [ln(x|z)] is maximized for q = 1.

25) but now z ∈ {e 1 , ..., e c } can assume only c different values. The vector of natural parameters is η = [ln(p 1 /p c ), ..., ln(p c-1 /p c ), 0)] and the log partition function is A(η) = ln 1 T e η .

Figure 6: Convergence of the loss, log-likelihood and gradient variance over 300 epochs of training on the MNIST dataset.

Figure 7: Test input images and generated handwritten images of the FC BVAE and FC CVAE

Fig.8bshows generated images of the FC CVAE if we sample z ∼ p(z) from the prior distribution. Since the FC CVAE can only learn to generate 100 different images, its decoder learns to generate template images that fit well to all the training images. We observe that some latent representations are decoded to meaningless patterns that just fit well to the data in avarage. However, the decoder also learned to generate at least one template image for each class of handwritten digits. Hence, the categorical latent representation can be interpreted as the cluster affiliation and the encoder of the FC CVAE automatically learns to cluster the data. Similar results for the CNN BVAE and CNN CVAE are shown in Fig.11in the appendix.

Figure 8: Handwritten digits generated by the FC BVAE and FC CVAE, if z is sampled from the prior distribution p(z).

Figure 9: Convergence of the loss and the log-likelihood over 300 epochs, while training on the MNIST dataset.

Figure 10: Test input images and generated handwritten images of the CNN BVAE and CNN CVAE

The architectural details of the trained VAEs. FC, Conv and Conv -1 are the fully connected, convolutional and deconvolutional layer, respectively. The shape of the convolutional layers is given in the form DimA × DimB × Channel/Stride/Activation.

Table content: <table><tr><td>Architecture</td><td>In/Out</td><td>Encoder</td><td>Latent</td><td>Decoder</td></tr><tr><td>FC BVAE</td><td>784</td><td>FC 1024/ReLu</td><td>z</td><td/></tr></table>

Peer review (iclr2019):
 In computing the gradient of the ELBO, the main challenge lies in computing the gradient of the reconstruction loss with respect to the encoder parameters. VAEs traditionally rely on reparameterization in order to obtain a low-variance estimate, but there are a number of other gradient estimators that one can apply. The authors here proprose to use a trick that is known, but perhaps not widely known: If we introduce an importance sampling distribution, then we can use samples from this distribution to compute an importance-weighted estimate of the gradient. The idea is now that we can compute the gradient w.r.t. the encoder parameters as a simple importance-sampling estimate, which obviates then need for reparameterization, or likelihood-ratio estimators. The authors then apply this trick to train VAEs with discrete latent variables. While I think that the idea that the authors present in this paper is worth further exploration, the paper in its current form is not sufficiently mature to appear at ICLR. The two areas where this paper would benefit from improvement are 1. Discussion of related work. While the authors seem to suggest that there has been no work on VAEs with discrete latent variables, there has in fact been quite a lot of work, including work on VAEs that contain both discrete and continuous variables (e.g. [8-10], but I*m almost certainly missing further references). There has also been a large body of work on continuous relaxations of discrete variables that are amenable to reparameterization (e.g. [6-7], and references therein). There has also been a line of work relating importance sampling to variational objectives (see [1-3] as key references). Finally, there is also related work on reweighted-wake-sleep style objectives (see [4]) which similarly don*t require reparameterization. From what I can tell, none of these references are cited or discussed as related work. In order to place this work in context, I would rewrite 2 to discuss approaches to gradient estimation in this space, which then makes it much easier to explain how this approach differs. 2. Empirical evaluation. The authors only evaluate on MNIST and F-MNIST, and don*t compare to any existing approaches. More than a couple of reconstructions, what I would like to see is an analysis of gradient variances, asymptotic ELBO estimates. I would also like to see a larger set of problems. Finally I would like to see a clear comparison to other methods based on, e.g., continuous relaxations. References [1] Y. Burda, R. Grosse, and R. Salakhutdinov, “Importance Weighted Autoencoders,” arXiv:1509.00519 [cs, stat], Sep. 2015. [2] T. Rainforth et al., “Tighter Variational Bounds are Not Necessarily Better,” arXiv:1802.04537 [cs, stat], Feb. 2018. [3] G. Tucker, D. Lawson, S. Gu, and C. J. Maddison, “Doubly Reparameterized Gradient Estimators for Monte Carlo Objectives,” arXiv:1810.04152 [cs, stat], Oct. 2018. [4] T. A. Le, A. R. Kosiorek, N. Siddharth, Y. W. Teh, and F. Wood, “Revisiting Reweighted Wake-Sleep,” arXiv:1805.10469 [cs, stat], May 2018. [5] A. Mnih and D. J. Rezende, “Variational inference for Monte Carlo objectives,” arXiv:1602.06725 [cs, stat], Feb. 2016. [6] G. Tucker, A. Mnih, C. J. Maddison, J. Lawson, and J. Sohl-Dickstein, “REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models,” in Advances in Neural Information Processing Systems, 2017, pp. 2624–2633. [7] W. Grathwohl, D. Choi, Y. Wu, G. Roeder, and D. Duvenaud, “Backpropagation through the Void: Optimizing control variates for black-box gradient estimation,” arXiv preprint arXiv:1711.00123, 2017. [8] J. T. Rolfe, “Discrete Variational Autoencoders,” arXiv:1609.02200 [cs, stat], Sep. 2016. [9] E. Dupont, “Learning Disentangled Joint Continuous and Discrete Representations,” arXiv:1804.00104 [cs, stat], Mar. 2018. [10] B. Esmaeili et al., “Structured Disentangled Representations,” arXiv:1804.02086 [cs, stat], Apr. 2018.