Training Variational Auto Encoders with Discrete Latent Representations using Importance Sampling

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 ) . 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 . 2 ESTIMATING THE EXPECTED LOG-LIKELIHOOD WITH REPARAMETRIZED SAMPLING . 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 . 3 ESTIMATING THE EXPECTED LOG-LIKELIHOOD WITH IMPORTANCE SAMPLING . 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 .

Summary: 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.