Direct Optimization through $\arg \max$  for  Discrete Variational Auto-Encoder

1 INTRODUCTION . Models with discrete latent variables drive extensive research in machine learning applications , such as language classification and generation ( Yogatama et al. , 2016 ; Hu et al. , 2017 ; Shen et al. , 2018 ) , molecular synthesis ( Kusner et al. , 2017 ) , or game solving ( Mena et al. , 2018 ) . Compared to their continuous counterparts , discrete latent variable models can decrease the computational complexity of inference calculations , for instance , by discarding alternatives in hard attention models ( Lawson et al. , 2017 ) , they can improve interpretability by illustrating which terms contributed to the solution ( Mordatch & Abbeel , 2017 ; Yogatama et al. , 2016 ) , and they can facilitate the encoding of inductive biases in the learning process , such as images consisting of a small number of objects ( Eslami et al. , 2016 ) or tasks requiring intermediate alignments ( Mena et al. , 2018 ) . Finally , in some cases , discrete latent variables are natural choices , including when modeling datasets with discrete classes ( Rolfe , 2016 ; Jang et al. , 2016 ; Maddison et al. , 2016 ) . Models involving discrete latent variables are hard to train , with the key issue being to estimate the gradients of the resulting non-differentiable objectives . While one can use unbiased estimators , such as REINFORCE ( Williams , 1992 ) , their variance is typically high ( Paisley et al. , 2012 ; Mnih & Gregor , 2014 ; Titsias , 2015 ; Gu et al. , 2015 ; Mnih & Rezende , 2016 ; Tucker et al. , 2017 ) . In variational auto-encoders ( VAEs ) with continuous latent variables , the reparameterization trick provides a successful alternative ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) . However , it can not be directly applied to non-differentiable objectives . Recent work ( Maddison et al. , 2016 ; Jang et al. , 2016 ) uses a relaxation of the discrete VAE objective , where latent variables follow a Gumbel-Softmax distribution . The Gumbel-Softmax approach is motivated by the connection of the Gumbel-Max perturbation model and the Gibbs distribution ( Papandreou & Yuille , 2011 ; Tarlow et al. , 2012 ; Hazan et al. , 2013 ; Maddison et al. , 2014 ) . The Gumbel-Softmax distribution relaxes the non-differentiable Gumbel-Max trick for sampling from the categorical distribution by replacing the argmax operation with a softmax operation . The reformulation results in a continuous objective function , which allows the use of the reparameterization trick . ( Mena et al. , 2018 ) provide an extension of this approach for discrete latent structures , namely distributions over latent matchings . Our work proposes minimizing the non-differentiable objective , by extending the direct loss minimization technique to generative models ( McAllester et al. , 2010 ; Song et al. , 2016 ) . Since categorical variables are represented by the Gibbs distribution , we start from the argmax formulation of the Gibbs distribution . We subsequently derive an optimization method that can propagate ( biased ) gradients through reparameterized argmax . The gradient of the resulting expectation is estimated by the difference of gradients of two argmax-perturbations . We begin by introducing the notation , the VAE formulation and the equivalence between the Gibbs distribution and the Gumbel-Max perturbation model in Section 3 . In Section 4.1 , we use these components to reformulate the discrete VAE objective , with respect to the argmax prediction operation . We subsequently state and prove the main result that allows us to differentiate through the argmax function in Section 4.2 . In Section 5 , we extend this result to mixed discrete-continuous VAEs and to semi-supervised VAE objectives . Finally , we demonstrate the effectiveness of our approach on image generation . 2 RELATED WORK . Variational inference has been extensively studied in machine learning , see Blei et al . ( 2017 ) for a review paper . In our work , we consider variational Bayes bounds with a discrete latent space . Many approaches to optimizing the variational Bayes objective , that are based on samples from the distribution , can be seen as applications of the REINFORCE gradient estimator ( Williams , 1992 ) . These estimators are unbiased , but without a carefully chosen baseline , their variance tends to be too high for the estimator to be useful and considerable work has gone into finding effective baselines ( Paisley et al. , 2012 ) . Other methods use various techniques to reduce the estimator variance ( Ranganath et al. , 2014 ; Mnih & Gregor , 2014 ; Gu et al. , 2015 ) . Reparameterization is an effective method to reduce the gradient estimate variance in generative learning . Kingma & Welling ( 2013 ) have shown its effectiveness in auto-encoding variational Bayes ( also called variational auto-encoders , VAEs ) for continuous latent spaces . Rezende et al . ( 2014 ) demonstrated its effectiveness in deep latent models . The success of these works led to reparameterization approaches in discrete latent spaces . Rolfe ( 2016 ) and Vahdat et al . ( 2018 ) represent the marginal distribution of each binary latent variable as a continuous variable in the unit interval . This reparameterization allows the backpropogation of gradients through the continuous representation . These works are restricted to binary random variables , and as a by-product , it encourages highdimensional representations for which inference is exponential in the dimension size . In contrast , our work reparameterizes the discrete Gibbs latent model , using a Gumbel-Max perturbation model and directly propagates gradients through the reparameterized objective . Maddison et al . ( 2016 ) and Jang et al . ( 2016 ) recently introduced a novel distribution , the Concrete distribution or the Gumbel-Softmax , that continuously relaxes discrete random variables . Replacing every discrete random variable in a model with a Concrete random variable , results in a continuous model , where the reparameterization trick is applicable . These works are close to ours , with a few notable differences . They use the Gumbel-Softmax function and their model is smooth and reparameterization may use the chain rule to propagate gradients . Similar to our setting , the GumbelSoftmax operation results in a biased estimate of the gradient . Different from our setting , the softmax operation relaxes the variational Bayes objective and results in a non-tight representation . Our work uses the Gumbel-Max perturbation model , which is an equivalent representation of the Gibbs distribution . With that , we do not relax the variational Bayes objective , while our argmax prediction remains non-differentiable and we can not naively use the chain rule to propagate gradients . Instead , we develop a direct optimization method to propagate gradients through argmax operation using the difference of gradients of two max-perturbations . Our gradient estimate is biased , except for its limit . Differentiating through argmax prediction was previously done in discriminative learning , in the context of direct loss minimization ( McAllester et al. , 2010 ; Song et al. , 2016 ) . Unfortunately , direct loss minimization can not be applied to generative learning , since it does not have a posterior distribution around its argmax prediction . We apply the Gumbel-Max perturbation model to transform the argmax prediction to the Gibbs distribution . This also allows us to overcome the “ general position ” assumption in ( McAllester et al. , 2010 ; Song et al. , 2016 ) using our “ prediction generating function ” . 3 BACKGROUND . To model the data generating distribution , we consider samples S = { x1 , ... , xm } originating from some unknown underlying distribution . We explain the generation process of a parameterized model pθ ( x ) , by minimizing its log-loss when marginalizing over its hidden variables z . Using variational Bayes , we upper bound the log-loss of an observed data point∑ x∈S − log pθ ( x ) ≤ ∑ x∈S −Ez∼qφ log pθ ( x|z ) + ∑ x∈S KL ( qφ ( z|x ) ||pθ ( z ) ) ( 1 ) Typically , the model distribution pθ ( x|z ) follows the Gibbs distribution law pθ ( x|z ) = e−θ ( x , z ) . When considering a discrete latent space , i.e. , z ∈ { 1 , ... , k } , the approximated posterior distribution follows the Gibbs distribution law qφ ( z|x ) ∝ eφ ( x , z ) . The challenge in generative learning is to reparameterize and optimize −Ez∼qφ log pθ ( x|z ) = k∑ z=1 eφ ( x , z ) ∑ ẑ e φ ( x , ẑ ) θ ( x , z ) ( 2 ) In our work , we reparameterize the variational bound using the equivalence between Gibbs models and Gumbel-Max perturbation models . Gumbel-Max perturbation models allow an alternative representation of Gibbs distributions qφ ( z|x ) ∝ eφ ( x , z ) that is based on the extreme value statistics of Gumbel-distributed random variables . Let γ be a random function that associates random variables γ ( z ) for each z = 1 , ... , k. When the random perturbations follow the zero mean Gumbel distribution law , whose probability density function is g ( γ ) = ∏k z=1 e − ( γ ( z ) +c+e− ( γ ( z ) +c ) ) for the Euler constant c ≈ 0.57 , we obtain the following identity between Gibbs models and Gumbel-Max perturbation models1 ( cf . Kotz & Nadarajah ( 2000 ) ) eφ ( x , z ) ∑ ẑ e φ ( x , ẑ ) = Pγ∼g [ zφ+γ = z ] , where zφ+γ def = arg max ẑ=1 , ... , k { φ ( x , ẑ ) + γ ( ẑ ) } ( 3 ) For completeness , a proof for this statement appears in Appendix A . 4 REPARAMETERIZATION AND DIRECT OPTIMIZATION . In the following section , we reformulate the discrete VAE objective using the Gumbel-Max perturbation model . We then derive an optimization method that directly propagates gradients through the reparameterized argmax function . 4.1 VAES WITH GUMBEL-MAX PERTURBATION MODELS . Perturbation models allow an alternative representation of Gibbs distributions qφ ( z|x ) ∝ eφ ( x , z ) . Using the Gumbel-Max perturbation model in Equation ( 3 ) , the negative log likelihood in Equation ( 2 ) takes the form −Ez∼qφ log pθ ( x|z ) = k∑ z=1 Pγ∼g [ zφ+γ = z ] θ ( x , z ) = k∑ z=1 Pγ∼g [ zφ+γ = z ] θ ( x , zφ+γ ) ( 4 ) = k∑ z=1 Eγ∼g [ 1zφ+γ=zθ ( x , zφ+γ ) ] = Eγ∼g [ θ ( x , zφ+γ ) ] ( 5 ) These quantities result from applying the law of total expectation , while realizing the probability events for zφ+γ : The last equality in Equation ( 4 ) holds since we restrict to the space zφ+γ = z . The last equality in Equation ( 5 ) holds by the linearity of the expectation , i.e. , ∑k z=1 Eγ∼g [ 1zφ+γ=zθ ( x , zφ+γ ) ] = Eγ∼g [ ∑k z=1 1zφ+γ=zθ ( x , z φ+γ ) ] and the fact that∑k z=1 1zφ+γ=z = 1 . The gradient of the decoder log probability θ ( x , zφ+γ ) with respect to its parameters is derived by the chain rule . The main challenge is to evaluate the gradient of Eγ∼g [ θ ( x , zφ+γ ) ] with respect to the encoder parameters , since the chain rule does not propagate through the argmax function zφ+γ . 1The set argmaxẑ=1 , ... , k { φ ( x , ẑ ) + γ ( ẑ ) } is the set of all maximal arguments , and does not always consist of a single element . However , since the Gumbel distribution is continuous , the γ for which their set argmaxẑ=1 , ... , k { φ ( x , ẑ ) + γ ( ẑ ) } consists more than a single element has a measure of zero . For notational convenience , when we consider integrals ( or probability distributions ) , we ignore measure zero sets . 4.2 DIRECT OPTIMIZATION THROUGH argmax Our main result is presented in Theorem 1 and shows how to compute the gradient of the reparameterized discrete VAE , i.e. , Eγ∼g [ θ ( x , zφ+γ ) ] with respect to the encoder parameters v. In the following we omit γ ∼ g for brevity . To make the encoder parameters v explicit in our notation , we denote the encoder function by φv ( x , z ) . Instrumental to our approach , is a novel “ prediction generating function ” . G ( v , ) = Eγ [ max ẑ { θ ( x , ẑ ) + φv ( x , ẑ ) + γ ( ẑ ) } ] ( 6 ) The proof of Theorem 1 is composed from three steps : 1 . We prove that G ( v , ) is a smooth function of v , . Therefore , the Hessian of G ( v , ) exists and it is symmetric , namely ∂v∂ G ( v , ) = ∂ ∂vG ( v , ) . ( 7 ) 2 . We show that encoder gradient is apparent in the Hessian : ∂v∂ G ( v , 0 ) = ∇vEγ [ θ ( x , zφv+γ ) ] . ( 8 ) 3 . We derive our update rule as the complement representation of the Hessian : ∂ ∂vG ( v , 0 ) = lim →0 1 ( Eγ [ ∇vφv ( x , z θ+φv+γ ) −∇vφv ( x , zφv+γ ) ] ) ( 9 ) The reparameterized gradient computation appears in Equation ( 10 ) and it is formally derived in Theorem 1 . Theorem 1 . Assume φv ( x , z ) is a smooth function of v. Then ∇vEγ [ θ ( x , zφv+γ ) ] = lim →0 1 ( Eγ [ ∇vφv ( x , z θ+φv+γ ) −∇vφv ( x , zφv+γ ) ] ) ( 10 ) Proof . First , we prove that G ( v , ) is a smooth function . Recall , g ( γ ) is the zero mean Gumbel probability density function . Applying a change of variable γ̂ ( z ) = θ ( x , ẑ ) + φv ( x , ẑ ) + γ ( ẑ ) , we obtain G ( v , ) = ∫ ∞ −∞ g ( γ ) max ẑ { θ ( x , ẑ ) + φv ( x , ẑ ) + γ ( ẑ ) } dγ = ∫ ∞ −∞ g ( γ̂ − θ − φv ) max ẑ { γ̂ ( ẑ ) } dγ̂ . Since g ( γ̂ − θ− φv ) is a smooth function of and φv ( x , z ) and φv ( x , z ) is a smooth function of v , we conclude that G ( v , ) is a smooth function of v , . Therefore , the Hessian of G ( v , ) exists and symmetric , i.e. , ∂v∂ G ( v , ) = ∂ ∂vG ( v , ) . We thus proved Equation ( 7 ) . To prove Equations ( 8 ) and ( 9 ) we differentiate under the integral , both with respect to and with respect to v. We are able to differentiate under the integral , since g ( γ̂− θ−φv ) is a smooth function of and v and its gradient is bounded by an integrable function ( cf . Folland ( 1999 ) , Theorem 2.27 ) . We turn to prove Equation ( 8 ) . We begin by noting that maxẑ { θ ( x , ẑ ) + φv ( x , ẑ ) + γ ( ẑ ) } is a maximum over linear function of , thus by Danskin Theorem ( cf . Bertsekas et al . ( 2003 ) , Proposition 4.5.1 ) holds ∂ ( maxẑ { θ ( x , ẑ ) + φv ( x , ẑ ) + γ ( ẑ ) } ) = θ ( x , z θ+φv+γ ) . By differentiating under the integral , ∂ G ( v , ) = Eγ [ ∂ ( θ ( x , z θ+φv+γ ) + φv ( x , z θ+φv+γ ) + γ ( z θ+φv+γ ) ) ] = Eγ [ θ ( x , z θ+φv+γ ) ] . We obtain Equation ( 8 ) by differentiating under the integral , now with respect to v , and setting = 0 . Finally , we turn to prove Equation ( 9 ) . By differentiating under the integral ∂vG ( v , ) = Eγ [ ∇vφv ( x , z θ+φv+γ ) ] . Equation ( 9 ) is attained by taking the derivative with respect to = 0 on both sides . The theorem follows by combining Equation ( 7 ) when = 0 , i.e. , ∂v∂ G ( v , 0 ) = ∂ ∂vG ( v , 0 ) with the equalities in Equations ( 8 ) and ( 9 ) . The gradient estimate in Theorem 1 is unbiased in the limit → 0 . However , for small epsilon the gradient is either zero , when z θ+φ+γ = zφ+γ , or very large , since the gradients difference is multiplied by 1/ . For large we often obtain a moderate non-zero gradient . In practice we use ≥ 0.1 which means that the gradient estimate is biased . This bias-variance tradeoff is evaluated in Figure 1 , on an encoder with three layers of sizes ( 784 , 300 , 10 ) , a discrete latent structure of size 10 and a matching decoder . In this evaluation we sampled 200 gradients for a fixed to evaluate the bias and standard deviation for each coordinate of the gradient . We report in the graph these quantities . The above theorem closely relates to the direct loss minimization technique ( cf . McAllester et al . ( 2010 ) ; Song et al . ( 2016 ) ) , which , in our setting , can be used to compute the gradient of Exθ ( x , zφ ) . The direct loss minimization predicts a single zφ for a given x and , therefore , can not generate a posterior distribution on all z = 1 , ... , k , i.e. , it lacks a generative model that exists in Gumbel-Max perturbation models .

Summary: The authors propose a method to apply the reparametrization trick when the random variables of interest are discrete. Their technique is based on a formulation of the objective function in terms of Gumbel-Max operators. They propose a derivation of the gradient in terms of an auxiliary variable \epsilon, such that the resulting gradient estimate is biased but the bias is reduced as \epsilon approaches zero, at the cost of increasing variance. Experiments are performed with VAE including discrete latent variable models. The authors show how their method converges faster than other baselines formed by estimators of the gradient given by the REBAR, RELAX and Gumbel-soft-max methods. In experiments with semi-supervised VAEs, their method outperforms the Gumbel softmax method in terms of accuracy and objective function.