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: This paper proposes combining the Gumbel-max trick and "direct loss optimization" for variance reduction in VAEs with discrete latent variables. This is a natural combination (in hindsight), since the Gumbel-max trick turns sampling into non-differentiable optimization, and direct loss optimization provides a way to optimize the expected value of a non-differentiable loss. The paper is well-written for the most part and is backed by good experimental results. However it like some of the mathematical details and some of the exposition could be greatly improved.