Year: 2018

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 = {x1, ...,xN} (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).

heading: 1 THE VARIATIONAL AUTO ENCODER
text: 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 = {x1, ...,xN} (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).
During training, the VAE maximizes the probability p(x) to observe the data x. Therefore, the negative evidence lower bound (ELBO)
L(θ) = −Eq(z|x) [ln p(x|z)] + DKL(q(z|x)||p(z)) (1) ≥ − ln p(x) + DKL(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 DKL(q(z|x)||p(z|x)) is the Kullback-Leibler (KL) divergence between p(z|x) and q(z|x). Because DKL(q(z|x)||p(z|x)) > 0, minimizing L(θ) means to maximize the probability p(x) the model assigns to observations x. Therefore, DKL(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.
The conventional VAE proposed in (Kingma & Welling (2013); Kingma et al. (2016)) learns continuous latent representations z ∈ Rc. It uses i.i.d. Gaussian distributed z, meaning that η = [µ1/σ 2 1 ,−1/(2σ21), ..., µc/σ2c ,−1/(2σ2c )]T , T (z) = [z1, z21 , ..., zc, z2c ]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 ∈ {e1, ..., ec} 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
∂
∂θ LKL(θ) =
∂
∂θ DKL(q(z|x)||p(z)) (4)
∂
∂θ LL(θ) =
∂
∂θ Eq(z|x) [ln p(x|z)] , (5)
where LKL(θ) only depends on the encoder parameters and the expected log likelihood term LL(θ) depends on both encoder and decoder parameters. For a suited choice of p(z) and q(z|x), LKL(θ) can be calculated in closed form. However, LL(θ) contains an expectation over z ∼ q(z|x) that has to be estimated during training. A good estimator L̂L(θ) for LD(θ) that is unbiased, differentiable with respect to θ and that has low variance is the key to train VAEs. SGVB uses an estimator L̂RL(θ) 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 L̂IL(θ) for LL(θ) 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 LL(θ) (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.

heading: 2 ESTIMATING THE EXPECTED LOG-LIKELIHOOD WITH REPARAMETRIZED SAMPLING
text: The standard estimator L̂RL(θ) proposed in Kingma & Welling (2013) is based on reparametrized sampling
∂
∂θ LL(θ) =
∂
∂θ Eq(z|x) [ln p(x|z)] (6)
= ∂
∂θ Ep( ) [ln p(x|z = f( ,θ))] (7)
≈ ∂ ∂θ 1 M M∑ m=1 ln p(x|z = f( m,θ)) (8)
= ∂
∂θ L̂RL(θ) (9)
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
∂
∂θ L̂RL(θ) =
1
M M∑ m=1 ∂ ∂z ln p(x|z = f( m,θ)) ∂ ∂θ z, (10)
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.

heading: 3 ESTIMATING THE EXPECTED LOG-LIKELIHOOD WITH IMPORTANCE SAMPLING
text: We propose an estimator L̂IL(θ) 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
∂
∂θ LL(θ) =
∂
∂θ
∫ ln p(x|z)q(z|x)dz (11)
= ∂
∂θ
∫ ln p(x|z)q(z|x)
qI(z) qI(z)dz (12)
≈ ∂ ∂θ 1 M
( M∑
m=1
ln p(x|zm) q(zm|x) qI(zm)
) (13)
= ∂
∂θ L̂IL(θ), (14)
where qI(z) is an arbitrary distribution that is of the same form as q(z|x) which is independent from the parameters θ. zm ∼ qI(z) are samples from this distribution. The estimator computes a weighted sum of the log likelihood ln p(x|zm) with the weighting q(zm|x)/qI(z). The benefit is that the log likelihood ln p(x|zm) depends on the decoder parameters θD only and not on the encoder parameters θE whereas the weighting q(zm|x)/qI(z) depends only on θD and not on θE . Therefore, calculation of the gradient of L̂IL(θ) can be separated
∂
∂θ L̂IL(θ) =
[ ∂
∂θE L̂IL(θ),
∂
∂θD L̂IL(θ)
] , (15)
with
∂
∂θE L̂IL(θ) =
1
M M∑ m=1 ln p(x|zm) ∂ ∂θE q(zm|x) qI(zm)
(16)
∂
∂θD L̂IL(θ) =
1
M M∑ m=1 q(zm|x) qI(zm) ∂ ∂θD ln p(x|zm). (17)
As shown in Fig. 4, gradient backpropagation is split into two separate parts. ∂∂θD L̂ I L(θ) backpropagates the error ln p(x|z) from the output of the VAE to the sampling operator and ∂∂θE L̂ I L(θ) backpropagates the error q(z|x)qI(z) from the sampling operator to the input layer of the VAE.
Compared to L̂RL(θ), 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, L̂IL(θ) 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 qI(z) have the form
q(z|x) = exp(ηT z−A(η)) (18) qI(z) = exp(ξ T z−A(ξ)), (19)
where z ∈ {0, 1}c, η = [ln(q1/(1 − q1)), ..., ln(qc/(1 − qc))] is the output vector of the encoder that contains the logits of the independent Bernoulli distributions and A(η) = 1T ln(1+ eη) are the corresponding log-partition functions.
Hence, Eq. 17 is
∂
∂θE L̂IL(θ) =
1
M M∑ m=1 ln p(x|zm) ∂ ∂η exp ( (η − ξ)T zm − (A(η)−A(ξ)) )( ∂ ∂θE η )
= 1
M M∑ m=1 ln p(x|zm) exp ( (η − ξ)T zm − (A(η)−A(ξ)) ) (z− q)T ( ∂ ∂θE η ) ,
where q = [q1, ..., qc]T = 11+e−η contains the probabilities q(zi = 1|x). The variance of the estimator L̂IL(θ) heavily depends on the choice of the natural parameters ξ of the distribution qI(z). We choose ξ = η, leading to a gradient of the very simple form
∂
∂θE L̂IL(θ) =
1
M M∑ m=1 ln p(x|zm) (zm − q)T ( ∂ ∂θE η ) . (20)
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
EqI(z)
[ ∂
∂θE L̂IL(θ)
] = EqI(z) [ 1
M M∑ m=1 ln p(x|z) (zm − q) ( ∂ ∂θE η )] (21)
= q(1− q)(ln p(x|z = 1)− ln(p(x|z = 0)) ( ∂
∂θE η
) (22)
= ∂
∂θE LL(θ), (23)
meaning that the estimator is unbiased. Further, the variance of ∂∂θE L̂ I 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 Eq(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 Eq(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 Eq(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.
5 VAE WITH CATEGORICALLY DISTRIBUTED z (CVAE)
For Categorically distributed z both the variational posterior distribution q(z|x) and qI(z), again have the form
q(z|x) = exp(ηT z−A(η)) (24) qI(z) = exp(ξ T z−A(ξ)), (25)
but now z ∈ {e1, ..., ec} can assume only c different values. The vector of natural parameters is η = [ln(p1/pc), ..., ln(pc−1/pc), 0)] and the log partition function is A(η) = ln1T eη .
With the formulas above, we arrive at the same easy form of the expected gradient of the log likelihood
∂
∂θE L̂IL(θ) =
1
M M∑ m=1 ln p(x|zm) (zm − q)T ( ∂ ∂θE η ) , (26)
but now with q = softmax(η) that consists of the probabilities qi = q(z = ei), where ∑c i=1 qi = 1.

heading: 6 EXPERIMENTS
text: 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 ∈ {e1, ..., ec}.
To train both architectures, the estimator L̂IL(θ) 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 ∈ {z1, ..., z100}. We train them for 300 epochs on the MNIST dataset, using
SGVB with our proposed estimator L̂IL(θ), 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 L̂IL(θ), 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 L̂IL(θ) 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.
The performance of the FC CVAE is worse than the performance of the FC BVAE. Training converges to a lower log likelihood ln p(x|z), because the maximal information content HCV AE(z) ≤ ln(100) of the latent variables of the FC CVAE is much less than the maximal information content HBV AE(z) ≤ c ln(2) of the latent variables of the FC BVAE. The FC CVAE can at maximum learn to generate 100 different handwritten digits, what is a small number compared to the 250 different images that the FC CVAE can learn to generate.
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
digits by traveling through the latent space Kingma & Welling (2013). However, in comparison to VAEs with continuous latent variables, we can only generate discrete mixtures for VAEs if the latent variables z are Bernoulli.
Fig. 8b shows 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. 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.

heading: 7 CONCLUSION
text: 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.

heading: B RESULTS FOR CNN BVAE ON FASHION MNIST
text: As shown in Fig. 12, the gradient variance approaches 0, the closer we get to the optimum. This is the same behaviour as for the MNIST dataset. As shown in Fig. 13, the FC BVAE can correctly reconstruct the shape of the given clothes with high accuracy. However, details like texture are lost. This is due to the limited model capacity, i.e. the latent representation z ∈ {0, 1}50 of the given VAE can at most encode 50Bits of information.

Review summaries:
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.

This paper proposes training VAEs with discrete latent variables by importance sampling the expected log likelihood (ELL) term in the ELBO, which is the problematic term since it is not amenable to reparametrization gradients.  For the importance sampling distribution, they choose the variational distribution itself, making the ELL gradient E[(d q(z|x) / d \theta) \log p(x|z) / q(z|x)].  Experiments are reported for MNIST and Fashion-MNIST using Bernoulli and categorical latent variables.