Rao-Blackwellizing the Straight-Through Gumbel-Softmax Gradient Estimator

1 INTRODUCTION . Models with discrete latent variables are common in machine learning . Discrete random variables provide an effective way to parameterize multi-modal distributions , and some domains naturally have latent discrete structure ( e.g , parse trees in NLP ) . Thus , discrete latent variable models can be found across a diverse set of tasks , including conditional density estimation , generative text modelling ( Yang et al. , 2017 ) , multi-agent reinforcement learning ( Mordatch & Abbeel , 2017 ; Lowe et al. , 2017 ) or conditional computation ( Bengio et al. , 2013 ; Davis & Arel , 2013 ) . The majority of these models are trained to minimize an expected loss using gradient-based optimization , so the problem of gradient estimation for discrete latent variable models has received considerable attention over recent years . Existing estimation techniques can be broadly categorized into two groups , based on whether they require one loss evaluation ( Glynn , 1990 ; Williams , 1992 ; Bengio et al. , 2013 ; Mnih & Gregor , 2014 ; Chung et al. , 2017 ; Maddison et al. , 2017 ; Jang et al. , 2017 ; Grathwohl et al. , 2018 ) or multiple loss evaluations ( Gu et al. , 2016 ; Mnih & Rezende , 2016 ; Tucker et al. , 2017 ) per estimate . These estimators reduce variance by introducing bias or increasing the computational cost with the overall goal being to reduce the total mean squared error . Because loss evaluations are costly in the modern deep learning age , single evaluation estimators are particularly desirable . This family of estimators can be further categorized into those that relax the discrete randomness in the forward pass of the model ( Maddison et al. , 2017 ; Jang et al. , 2017 ; Paulus et al. , 2020 ) and those that leave the loss computation unmodified ( Glynn , 1990 ; Williams , 1992 ; Bengio et al. , 2013 ; Chung et al. , 2017 ; Mnih & Gregor , 2014 ; Grathwohl et al. , 2018 ) . The ones that do not modify the loss computation are preferred , because they avoid the accumulation of errors in the forward direction and they allow the model to exploit the sparsity of discrete computation . Thus , there is a particular need for single evaluation estimators that do not modify the loss computation . ∗Work done partly at the Institute for Advanced Study , Princeton , NJ . In this paper we introduce such a method . In particular , we propose a Rao-Blackwellization scheme for the straight-through variant of the Gumbel-Softmax estimator ( Jang et al. , 2017 ; Maddison et al. , 2017 ) , which comes at a minimal cost , and does not increase the number of function evaluations . The straight-through Gumbel-Softmax estimator ( ST-GS , Jang et al. , 2017 ) is a lightweight stateof-the-art single-evaluation estimator based on the Gumbel-Max trick ( see Maddison et al. , 2014 , and references therein ) . The ST-GS uses the argmax over Gumbel random variables to generate a discrete random outcome in the forward pass . It computes derivatives via backpropagation through a tempered softmax of the same Gumbel sample . Our Rao-Blackwellization scheme is based on the key insight that there are many configurations of Gumbels corresponding to the same discrete random outcome and that these can be marginalized over with Monte Carlo estimation . By design , there is no need to re-evaluate the loss and the additional cost of our estimator is linear only in the number of Gumbels needed for a single forward pass . As we show , the Rao-Blackwell theorem implies that our estimator has lower mean squared error than the vanilla ST-GS . We demonstrate the effectiveness of our estimator in unsupervised parsing on the ListOps dataset ( Nangia & Bowman , 2018 ) and on a variational autoencoder loss ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) . We find that in practice our estimator trains faster and achieves better test set performance . The magnitude of the improvement depends on several factors , but is particularly pronounced at small batch sizes and low temperatures . 2 BACKGROUND . For clarity , we consider the following simplified scenario . Let D ∼ pθ be a discrete random variable D ∈ { 0 , 1 } n in a one-hot encoding , ∑ Di = 1 , with distribution given by pθ ( D ) ∝ exp ( DT θ ) where θ ∈ Rn . Given a continuously differentiable f : R2n → R , we wish to minimize , min θ E [ f ( D , θ ) ] , ( 1 ) where the expectation is taken over all of the randomness . In general θ may be computed with some neural network , so our aim is to derive estimators of the total derivative of the expectation with respect to θ for use in stochastic gradient descent . This framework covers most simple discrete latent variable models , including variational autoencoders ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) . The REINFORCE estimator ( Glynn , 1990 ; Williams , 1992 ) is unbiased ( under certain smoothness assumptions ) and given by : ∇REINF : = f ( D , θ ) ∂ log pθ ( D ) ∂θ + ∂f ( D , θ ) ∂θ . ( 2 ) Without careful use of control variates ( Mnih & Gregor , 2014 ; Tucker et al. , 2017 ; Grathwohl et al. , 2018 ) , the REINFORCE estimator tends to have prohibitively high variance . To simplify exposition we assume henceforth that f ( D , θ ) = f ( D ) does not depend on θ , because the dependence of f ( D , θ ) on θ is accounted for in the second term of ( 2 ) , which is shared by most estimators and generally has low variance . One strategy for reducing the variance is to introduce bias through a relaxation ( Jang et al. , 2017 ; Maddison et al. , 2017 ) . Define the tempered softmax softmaxτ : Rn → Rn by softmaxτ ( x ) i = exp ( xi/τ ) / ∑n j=1 exp ( xj/τ ) . The relaxations are based on the observation that the sampling of D can be reparameterized using Gumbel random variables and the zero-temperature limit of the tempered softmax under the coupling : D = lim τ→0 Sτ ; Sτ = softmaxτ ( θ +G ) ( 3 ) where G is a vector of i.i.d . Gi ∼ Gumbel random variables . At finite temperatures Sτ is known as a Gumbel-Softmax ( GS ) ( Jang et al. , 2017 ) or concrete ( Maddison et al. , 2017 ) random variable , and the relaxed loss E [ f ( Sτ , θ ) ] admits the following reparameterization gradient estimator for τ > 0:1 ∇GS : = ∂f ( Sτ ) ∂Sτ d softmaxτ ( θ +G ) dθ . ( 4 ) 1For a function f ( x1 , x2 ) , ∂f ( z1 , z2 ) /∂x1 is the partial derivative ( e.g. , a gradient vector ) of f in the first variable evaluated at z1 , z2 . For a function g ( θ ) , dg/dθ is the total derivative of g in θ . For example , d softmaxτ ( θ +G ) /dθ is the Jacobian of the tempered softmax evaluated at the random variable θ +G . This is an unbiased estimator of the gradient of E [ f ( Sτ , θ ) ] , but a biased estimator of our original problem ( 1 ) . For this to be well-defined f must be defined on the interior of the simplex ( where Sτ sits ) . This estimator has the advantage that it is easy to implement and generally low-variance , but the disadvantage that it modifies the forward computation of f and is biased . Henceforth , we assume D , Sτ , and G are coupled almost surely through ( 3 ) . Another popular family of estimators are the so-called straight-through estimators ( c.f. , Bengio et al. , 2013 ; Chung et al. , 2017 ) . In this family , the forward computation of f is unchanged , but backpropagation is computed “ through ” a surrogate . One popular variant takes as a surrogate the tempered probabilities of D , resulting in the slope-annealed straight-through estimator ( ST ) : ∇ST : = ∂f ( D ) ∂D d softmaxτ ( θ ) dθ . ( 5 ) For binary D , a lower bias variant of this estimator ( FouST ) was proposed in Pervez et al . ( 2020 ) . The most popular straight-through estimator is known as the straight-through Gumbel-Softmax ( STGS , Jang et al. , 2017 ) . The surrogate for ST-GS is Sτ , whose Gumbels are coupled to D through ( 3 ) : ∇STGS : = ∂f ( D ) ∂D d softmaxτ ( θ +G ) dθ . ( 6 ) The straight-through family has the advantage that they tend to be low-variance and f need not be defined on the interior of the simplex ( although f must be differentiable at the corners ) . This family has the disadvantage that they are not known to be unbiased estimators of any gradient . These estimators are quite popular in practice , because they preserve the forward computation of f , which prevents the forward propagation of errors and maintains sparsity ( Choi et al. , 2017 ; Chung et al. , 2017 ; Bengio et al. , 2013 ) . All of the estimators discussed in this paper can be computed by any of the standard automatic differentiation software packages using a single evaluation of f on a realization of D or some underlying randomness . We present implementation details for these and our Gumbel-Rao estimator in the Appendix , emphasizing the surrogate loss framework ( Schulman et al. , 2015 ; Weber et al. , 2019 ) and considering the multiple stochastic layer case not covered by ( 1 ) . 3 GUMBEL-RAO GRADIENT ESTIMATOR . 3.1 RAO-BLACKWELLIZATION OF ST-GUMBEL-SOFTMAX . We now derive our Rao-Blackwelization scheme for the ST-GS estimator . Our approach is based on the observation that there is a many-to-one relationship between realizations of θ + G and D in the coupling described by ( 3 ) and that the variance introduced by θ + G can be marginalized out . The resulting estimator , which we call the Gumbel-Rao ( GR ) estimator , is guaranteed by the Rao-Blackwell theorem to have lower variance than ST-GS . In the next subsection we turn to the practical question of carrying out this marginalization . In the Gumbel-max trick ( 3 ) , D is a one-hot indicator of the index of arg maxi { θi +Gi } . Because this argmax operation is non-invertible , there are many configurations of θ +G that correspond to a single D outcome . Consider an alternate factorization of the joint distribution of ( θ +G , D ) : first sample D ∼ pθ , and then θ +G given D. In this view , the Gumbels are auxillary random variables , at which the Jacobian of the tempered softmax is evaluated and which locally increase the variance of the estimator . This local variance can be removed by marginalization . This is the key insight of our GR estimator , which is given by , ∇GR : = ∂f ( D ) ∂D E [ d softmaxτ ( θ +G ) dθ ∣∣∣∣D ] . ( 7 ) It is not too difficult to see that∇GR = E [ ∇STGS|D ] . By the tower rule of expectation , GR has the same expected value as ST-GS and is an instance of a Rao-Blackwell estimator ( Blackwell , 1947 ; Rao , 1992 ) . Thus , it has the same mean as ST-GS , but a lower variance . Taken together , these facts imply that GR enjoys a lower mean squared error ( not a lower bias ) than ST-GS . Proposition 1 . Let ∇STGS and ∇GR be the estimators defined in ( 6 ) and ( 7 ) . Let ∇θ : = dE [ f ( D ) ] /dθ be the true gradient that we are trying to estimate . We have E [ ‖∇GR −∇θ‖2 ] ≤ E [ ‖∇STGS −∇θ‖2 ] . ( 8 ) Proof . The proposition follows from Jensen ’ s inequality and the linearity of expectations , see C.1 . While GR is only guaranteed to reduce the variance of ST-GS , Proposition 1 guarantees that , as a function of τ , the MSE of GR is a pointwise lower bound on ST-GS . This means GR can be used for estimation at temperatures , where ST-GS has low bias but prohibitively high variance . Thus , GR extends the region of suitable temperatures over which one can tune . This allows a practitioner to explore an expanded set when trading-off of bias and variance . Empirically , lower temperatures tend to reduce the bias of ST-GS , but we are not aware of any work that studies the convergence of the derivative in the temperature limit . In our experiments , we observe that our estimator facilitates training at lower temperatures to improve in both bias and variance over ST-GS . Thus , our estimator retains the favourable properties of ST-GS ( single , unmodified evaluation of f ) while improving its performance .

Summary: The paper presents a new way algorithm to compute the straight-through variant of the Gumbel Softmax gradient estimator. The method does not change the estimator's bias, but provably reduces its variance (with a small overhead, using Rao-blackwellization). The new estimator shows good performance on different tasks, and appears to lead to more efficient optimization for lower temperatures (lower bias).