Rao-Blackwellizing the Straight-Through Gumbel-Softmax Gradient Estimator
Authors: Max B Paulus
Chris J. Maddison
Andreas Krause
Venue: ICLR 2021 Oral (2021)
Keywords: gumbel
softmax
gumbel-softmax
straight-through
straightthrough
rao
rao-blackwell

Abstract: Gradient estimation in models with discrete latent variables is a challenging problem, because the simplest unbiased estimators tend to have high variance. To counteract this, modern estimators either introduce bias, rely on multiple function evaluations, or use learned, input-dependent baselines. Thus, there is a need for estimators that require minimal tuning, are computationally cheap, and have low mean squared error. In this paper, we show that the variance of the straight-through variant of the popular Gumbel-Softmax estimator can be reduced through Rao-Blackwellization without increasing the number of function evaluations. This provably reduces the mean squared error. We empirically demonstrate that this leads to variance reduction, faster convergence, and generally improved performance in two unsupervised latent variable models.

Published as a conference paper at ICLR 2021
RAO-BLACKWELLIZING THE STRAIGHT-THROUGH GUMBEL-SOFTMAR GRADIENT ESTIMATOR
Max B. Paulus
Chris J. Maddison"
ETH Zürich
University of Toronto
MPI for Intelligent Systems, Tühingen max. haaulus@iff.etta..c wtw.f.cte.cc...
Vector Institute
cmaddi sedadissess.. toronto. .dd:
Andreas Krause
ETH Zürich
krauseaßethz.ch
ABSTRACT
Gradient estimation models with discrete latent variables a challenging prob- lem, because the simplest unbiased estimators tend to have high variance. To counteract this, modern estimators either introduce bias, rely on multiple function evaluations,
learned, input-dependent baselines. Thus, there is need for estimators that require minimal tuning, are computationally cheap, and have low squared error. In this paper, we show that the variance of the trraight-througe variant of the popular Gumbel-Softmax estimator can be reduced through Rao- Blackwellization without increasing the number of function evaluations. This provably reduces the mean squared error. We empirically demonstrate that this leads to variance reduction, faster convergence, and generally improved perfor- mance in two unsupervised latent variable models.
INTRODUCTION
Models with discrete latent variables 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 minimize expected loss using gradient-based op- timization, so the problem of gradient estimation for discrete latent variable models has received considerable attention over recent years. Existing estimation techniques can broadly categorized
into two groups, based on whether they require one loss evaluation (Glynn 1990; Williams 1992
Bengio et 2013 Mnih Gregor 2014, Chung et al., 2017: Maddison al. [211]] 2017|ang Jang al. 2017 Grai|waat et al.) 2018) or multille loss evaluations (Gu et al.) [016 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. 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
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 particular need for single evaluation estimators that do not modify the loss computation.
Work done partly the Institute for Advanced Study, Princeton, NJ.

---

Published as a conference paper at ICLR 2021
In this paper we introduce such method. In particular, we propose Rao- Saoo-Blackwellization scheme for the straight- variant of the Gumbel-Softma 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-Softmat estimator (ST-GS, Jang et al., 2017) is a lightweight state- of-the-art single-evaluation estimator based on the Gumbel-Max trick (see Maddison 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. computes derivatives via backpropagation through tempered softmax of the same Gumbel sample. Our cho-Blackkellizatteee 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 our estimator is linear only in the number of Gumbels needed for single forward pass. As we show, the Rao-Blackwell theorem implies that our estimator has lower mcan squared error than the vanilla ST-GS. We demonstrate the effectiveness of our estimator in unsupervised parsing the ListOps dataset (Nangia & Bowman 2018) and on a variational autoencoder loss (Kingma & Welling, 2013] Rezende et al. (2044). 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.
scheme is based on the
BACKGROUND
For clarity, we consider the following simplified scenario. Let D pe be discrete random variable D E {0, 1 in a one-hot encoding, D; = 1, with distribution given by po (D) xx exp(D"0) where 0 RR. Given continuously differentiable ʃ 1T R, wish minimize,
min ())
(1)
where the expectation is taken over all of the randomness. In general 0 may be computed with some neural network, so aim is to derive estimators of the total derivative of the expectation with respect to 0 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:
a Olog Po D)
af(D,0)
V REINF == f(D, 0)
00
00
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 fDD 8) = f(D) does not depend on 6, because the dependence of f 8) on 0 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 relaxation (Jang al. 2017 Maddison et al. 2017). Define the tempered softmax softmax, IR" + R" by softmax, (x); exp(#x/T)) // exp(æj/7). The relaxations are based on the observation that the sampling of 1) can be reparameterized using Gumbel random variables and the zero-temperature limit of the tempered softmax under the coupling:
D lim S+:
S, softmax, (0 G)
where G is a vector of i.i.d. G3 Gumbel random variables. At finite temperatures S, is known as a Gumbel-Softman (GS) (Jang et al., 2017) or concrete (Maddison et al. (017)) random variable, and the relaxed loss (S 8) admits the following reparameterizzation gradient estimator for > 0
af(S,) dsoftmax, (0 + G)
V GS = 0s,
(4)
d0
For function f(x1, 22), (=1, /02 dzi is the partial derivative (e.g.. gradient vector) of f in the first variable evaluated at 21, 22. For function g()), dg. d0 is the total derivative of in 0. For example, d softmaxx (0 + G)/ d0 is the Jacobian of the tempered softmax evaluated the random variable 0 + G.
2

---

Published as a conference paper at ICLR 2021
This is unbiased estimator of the gradient of (f,, 8)], but biased estimator of our original problem Q For this to be well-defined 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 S, and G are coupled almost surely through (3)
Another popular family of estimators are the so-called straight-throug estimators (c.f., Bengio et al.] 2013; Chung et al. (017). 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-throug estimator (ST):
DD( d dsoftmax,((0
ST
(5)
0D
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-throug GumbellSoftman (ST- GS, Jang et al. 2017). The surrogate for ST-GS is S,, whose Gumbels are coupled to D through
AJ(D) dsoftmaxx + G)
STGSS
(6)
OD
d0
The straight-through family has the advantage that they tend 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 on 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
3 GUMBEL-RAO GRADIENT ESTIMATOR
3.1
RAO-BLACKWELLIZATIIN OF ST-GUMBEL-SOFTMA
We now derive our Rao-Blackwelizationn scheme for the ST-GS estimator. Our approach is based on the observation that there many-to-one relationship between realizations of G and D in the coupling described by (3) and that the variance introduced by 0 + G can be marginalized out. The resulting estimator, which we call the Gumbel-Re (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 33. D is one-hot indicator of the index of arg max; {0, + Gi). Because this argmax operation is non-inver,ible, there are many configurations of f + G that correspond to single D outcome. Consider an alternate factorization of the joint distribution of (0 + G, D): first sample D ~ pa. and then + G given D. this view, the Gumbels are auxillary random variables, 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,
Of(D) aD
dsoftmaa, (0+ G) D
|
VGR
d0
is not too difficult to see that VGR = E | STGGs| D). By the tower rule of expectation, GR has the same expected value ST-GS and is 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 lower mean squared error (not lower bias) than ST-GS.
3

---

Published as a conference paper at ICLR 2021
Proposition 1. Let VerGs and VeR be the estimators defined in ( and dB[f d0 be the true gradient that we are trying to estimate. We have
Let
Vo
B |VE Voll E [IIVVTTSS Voll".
(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 guarantees that, 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 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.
3.2 MONTE CARLO APPROXIMATION
The GR estimator requires computing the expected value of the Jacobian of the tempered softmax over the distribution + G D. Unfortunately, an analytical expression for this is only available in the simplest cases P In this section we provide simple Monte Carlo (MC) estimator with sample size K for E[dS, d0) DJ. which we call the Gumbel-Roe Monte Carlo Estimator (GR-MCK). This estimator can be computed locally cost that only scales like nKK (the arity of D times K).
They key property exploited by GR-MCK is that 0 + G can be reparameterized in the fol- >> -==1 exp(0.) and E; exponential i..dd., we have the following equivalence in distribution Maddison et al. 2014
lowing closed form. Given realization of D such that D, 2
Maddison 2016
Tucker et al. (017).
log (E₃) + log Z(0) E, + log exp(U,]
if = O.W.
f, + G, &
(9)
E.
Z(0))
With this in mind, we define the GR-MCK estimator:
VGRMCK: == Df(D)
1
K
d dsoftmax,
(10)
3D
K
d0
where G$ ~ using the reparameterization For the case = 1, our estimator reduces to the standard ST-GS estimator. The cost for drawing multiple samples G ~ + G scales only linearly in the arity of D and is usually negligible in modern applications, where the bulk of computation accrues from the computation of f. Moreover, drawing multiple samples of 0 - D can easily be parallelized on modern workstations (GPUs, etc.). Our estimator remains single-evaluation estimator under this scheme, because the loss function f is still only evaluated at D. Finally, with GR, the GR-MCK is guaranteed to improve in MSE over ST-GS for any K as confirmed Proposition 2
Proposition 2. Let VSGGS and VGRMCK be the estimators defined in (6) and 10
Let Vo =
dBIS(D)) do be the true gradient that we are trying estimate. For all K have
VGRMCK - voll?] < |IVSTGS - Voll
(11)
Proof. The proposition follows from Jensen's inequality and the linearity of expectations, see C.2
For example, in the case of 2 (binary) and 1 an analytical expression for the GR estimator available.
4

---

Published as a conference paper at ICLR 2021
3.3 VARIANCE REDUCTION IN MINIBATCHES
The variance of GR-MCK can be reduced by increasing K or by averaging B i..d. samples of the GR-MCK estimator. An average of i.i.d. samples V GRRCK for E (l...., B) is generalization of minibatching by sampling data points with replacement. In particular, 0 may depend on an additional P. In this case, VSTGS a random variable that depends not only D and G. but also on X. In this subsection, we consider the effect of increasing K and B separately. Expectations are taken over all the randomness.
source of randomness, i.c., = h(X) for
Let V GRRCCK be ii.d. as VGRMCK for b {1...., B} and define the following "minibatched" GR-MCK estimator:
B V
V V.R
GRMCK ==
GRMCK
(12)
B
b=1
Proposition 3 summarizes the scaling of the variance of [2), and is an elementary application of the law of total variance.
L:B
Proposition 3. Let STGGS, VeR and GRMCK be the estimators defined in 6 have
and
12 We
var
V
GRMCK
L-B
E [var [Vrass D,
BK
var [oRl
B
(13)
where var is the trace of the covariance matrix.
Proof. The proposition follows directly from the law of total variance, see C33
I:B
As expected the total variance of VGRMCK decreases like I/B. The key point of Proposition is that the component of the variance that K reduces also be reduced by increasing the batch size B. This suggests that the effect of GR-MCK will be most pronounced at small batch sizes. Proposition diminishing returns increasing K for fixed batch size B, such
3 also indicates that there that the variance of GR-MCK will eventually be dominated by the right-hand term of (13). In our experimental section, we explore various K and study the effect gradient estimation in more detail.
Finally, we note that the choice of a Monte Carlo scheme to approximate E [dS,/ D] permits the of additional well-known variance reduction methods improve the estimation properties of gradient estimator. For example, antithetic variates or importance sampling are sensible methods to explore in this setting (Kroese al. 2013). For low-dimensional discrete random variables, Gaussian quadrature or other numerical methods could be employed. However, we found the simple Monte Carlo scheme described above effective practice and report results based on this procedure in the experimental section.
RELATED WORK
The idea of using Bao-Blackwellizatte to reduce the variance of gradient estimators for discrete latent variable models has been explored in machine learning. For example, Mauple,Liuee al. (2018) describe a sum-and-sample style estimator that analytically computes part of the expectation to reduce the variance of the gradient estimates. The favorable properties of their estimator are due the Rao- Blackwell theorem. Kool et al. (2020) describe gradient estimator based on sampling without replacement. Their estimator emerges naturally as the Rao- Bao-Blackwell estimator of the importance- weighted estimator (Vieira 2017) and the estimator described by Liu et al. (2018). Both of these estimators rely on multiple function evaluations to compute a gradient estimate. In contrast, our work is the first to consider Rao-Blase Beoo--lleeeellsatiiion in the context of single-evaluation estimator.
Recently, Paulus et al. 2020) extend the Gumbel-Softmax gradient estimator to other discrete structures. Our approach can be used to reduce the variance of the corresponding ttraighttthrough variants, when an efficient reparameterization of the perturbation conditional on the discrete structure is available (Gane al., 2014).
5

---

Published as a conference paper at ICLR 2021
(a) Ohjective function
(b)
A (var) (c) (r=0.5)
A (var)
(d) (r=1.0)
(var)
(r=0.1)
Figure 1: Our estimator (GR-MCK) effectively reduces the variance over the entire simplex and is particularly effective at low temperatures. Contours for the quadratic programme in three dimensions (1a) and difference in log10-tras 10-trace of the covariance matrix between ST-GS and GR-MC1000 at Warmer means difference is larger.
different temperatures
[b
1c.
1d
5 EXPERIMENTS
PROTOCOL
In this section, study the effectiveness of gradient estimator in practice. In particular, evaluate its performance with respect to the temperature т, the number MC samples and the batch size B. We measure the variance reduction and improvements in MSE our estimator achieves in practice, and assess whether its lower variance gradient estimates accelerate the convergence on the objective improve final test set performance. Our focus is single-evaluation gradient estimation and we compare against other non non-relaxing estimators (ST, FouST, ST-GS and REINFORCE with running mean bascline) and relaxing estimators (GS), where permissible. Experimental details are given in Appendix D
First, we consider a toy example which allows us explore and visualize the variance of our estimator and suggests that it is particularly effective at low temperatures. Next, we evaluate the effect of and K in latent parse tree task which does not permit the use of relaxed gradient estimators. Here, our estimator facilitates training at low temperatures to improve overall performance and is effective even with few MC samples. Finally, we train variational auto-encoders with discrete latent variables (Kingma & Welling 2013 Rezende al. 2014). Our estimator yields improvements at small batch sizes and obtains competitive or better performance than the GS estimator at the largest arity.
5.2 QUADRATIC PROGRAMMING ON THE SIMPLEX
As toy problem. consider the problem of minimizing quadratic program (p-c) over the probability simplex " {p E R" p; 2 0, Σ = 1} for Q Rax= positive-defintc and R This problem may be reframed the following stochastic optimization problem,
min B[(D T A(p) - cc].
-
where D ~ Discrete(p) and (p(p)) P. 2 Qs and Ag = (n- (n--))(p-c;)) Qiy for i +
(p:-c:)2
evey -ip While solving the above problem is simple using standard methods, provides useful testbed to evaluate the effectiveness of our variance reduction scheme. For this purpose, we consider Qi exp (- jl) and c; = 2/2 in three dimensions.
Our estimator reduces the variance in the gradient estimation over the entire simplex and is particularly effective at low temperatures in this problem. Figure [ we compare the log10-trace of the covariance matrix of ST-GS and GR-MC1000 at three different temperatures and display their difference over the entire domain. The improvement universal. The pattern not always intuitive (oval bull's eyes), despite the simplicity of the objective function. Compared with ST-GS, our estimator on this example appears more effective closer to the corners and edges, which is important for learning discrete distributions. At lower temperatures, the difference between the two estimators becomes particularly acute. This suggests that our estimator may train better at lower temperatures and be more responsive optimizing the temperature successfully trade off bias and variance.

---

Published as a conference paper at ICLR 2021
Table 1: Our estimator (GR-MCK) facilitates training lower temperatures with improved perfor-
mance on the latent parse tree task. Best test classification accuracy on the ListOps dataset selected on the validation set. Best estimator at given temperature in bold, best estimator across temperatures in italics. Higher is better.
L S 10
L 25
L L 50
ESTIMATOR
0.01
0.1
7 1.0
0.01
T= 0.1
T 1.0
== 0.01
== 0.1
1.0
ST-GS
38.8
59.3
65.8
41.2
57.1
60.2
46.8
56.8
59.6
GR-M ER-MCIO
66.4
66.9
66.7
60.7
60.8
60.9
58.7
59.1
59.6
GR-MC100
65.6
66.3
65.9
60.0
61.3
61.2
59.6
59.1
59.6
GR-MC1000
66.5
67.1
67.0
60.2
60.9
61.2
60.0
59.8
59.9
5.3 UNSUPERVISED PARSING ON LISTOPS
Straight-througs estimators feature prominently in NLP (Martins et al. 2019) where latent discrete structure arises naturally, but the use of relaxations is often infeasible. Therefore, we evaluate our estimator in a latent parse tree task on subsets of the ListOps dataset (Nangia & Bowman 2018).
This dataset contains sequences of prefix arithmetic expressions x (e.g., max min
}])
that evaluate to an integer E {0, 1. 9}. The arithmetic syntax induces latent parse tree T. We consider the model by (Choi al., 2017) that learns distribution over plausible parse trees of given sequence to maximize
E log Poglo(u)T,a))
qo ar)
Both the conditional distribution over parse trees qo (T'(r) and the classifier p& (v]T, are param- eterized using neural networks. In this model, a parse tree T ~ qo (T((())) for a given sentence is sampled bottom-up by successively combining the embeddings of tokens that appear in given sequence until single embedding for the entire sequence remains. This is then used for performing subsequent classification. Because it computationally infeasible to marginalize over trees, Choi et al. (2017) rely on the ST-GS estimator for training. We compare this estimator against our estimator GR-MCA with E {10, 100, 1000]. We consider temperatures т E {0.01, 0.1, 1.03 and experiment with shallow and deeper trees by considering sequences of length up 10, 25 and 50. All models are trained with stochastic gradient descent with a batch size equal to the maximum Because are interested in controlled setting to investigate the effect of and K, our experimental set-up is significantly simpler than elsewhere (e.g.. Havrylov et al.) 2019). We give details and highlight important differences in Appendix D.I
Our estimator facilitates training at lower temperatures and achieves better final test set accuracy than ST-GS (Table [] Increasing K improves the performance at low temperatures, where the differences between the estimators are most pronounced. Overall, across all temperatures this results modest improvements, particularly for shallow trees and small batch sizes. We also find evidence for diminishing returns: The differences between ST-GS and GR-MC10 are larger than between GR-MC100 or GR-MC1000, suggesting that our estimator is effective even with few MC samples.
5.4 GENERATIVE MODELING WITH DISCRETE VARIATIONAL AUTO-ENCODERS
Finally, we train variational auto encodesss (Kingma & Welling 2013 Rezende et al.) 2014) with discrete latent random variables on the MNIST dataset handwritten digits (LeCun & Cortes 2010). We used the fixed binarization of (Salakhutdinov & Murray 2008) and the standard split into train, validation and test sets. Our objective is to maximize the following variational lower bound on the eggllikllhhodd.
(z D*)
log p(x)
E do (Di
log
M
9o
(Di (x)
Di
where x denotes the input image and D variables. This objective takes form in equation (1. For training, the bound is approximated using
go (D*(z) (x) denotes a vector of discrete latent random
7

---

Published as a conference paper at ICLR 2021
10.10
379
345
2000
309
241
19.00
380
0902
2018
6.00
17 665
(a) 1o vs Iterations
(b) Jog (msee VS
(c) Iterations vs ELBO
Figure 2: Our estimator (GR-MCK) effectively reduces the variance over the entire training trajectory 2a achieves a lower mean squared error at a lower temperature (26) and converges faster than ST and ST-GS on the discrete VAE objective (2c). Log10-trace of the covariance matrix over training trajectory (2a) and log 10-MSE (2b) at different temperatures during training, average number of iterations and standard error reach various thresholds of the objective on the validation set (2c).
only a single sample (M = 1). For final validation and testing, we use 5000 samples (M 5000). Both the generative model pa(ar, D) and the variational distributions go D(ar) were parameterized using neural networks. We experiment with different batch sizes and discrete random variables of arities in {2, 4, 8, 16] in Maddison et al. (2017). To facilitate comparisons, total dimension of the latent space and train all models for 50,000 iterations using stochastic gradient descent with momentum. Hyperparameters are optimised for each estimator using random search (Bergstra & Bengio 2012) over twenty independent runs. More details are given in Appendix D.2
do not alter the
Our estimator effectively reduces the variance over the entire training trajectory (Figure 2ap. Even a small number of MC samples (K' = 10) results in sizable variance reductions. The variance reduction compares favorably the magnitude of the minibatch variance (Appendix) Ep. Empirically, we find that lower temperatures tend to reduce bias. Our estimator facilitates training at lower temperatures and thus features lower MSE (Figure 2b). During training our estimator can trade off bias and variance to improve the gradient estimation. Empirically, we observed that on this task, the best models using ST-GS trained at an average temperature of 0.65, while the best models using GR-MC 1000 trained at an average temperature of 0.35. This is interesting, because it indicates that our estimator may make the use of temperature annealing during training more effective. We find lower variance gradient estimates improve convergence of the objective (Figure 2c). GR-MCI 000 reaches various performance thresholds on the validation set with reliably fewer iterations than ST or ST-GS. This effect is observable at different arities and persistent over the entire training trajectory.
For final test set performance, our estimator outperforms REINFORCE and all other straight-through estimators (Table The improvements over ST-GS extend up to nats (for batch size 20, 16-ary) at small batch sizes and are more modest at large batch sizes as expected (also see Appendix [E). This confirms that estimator might be particularly effective in settings, where training high batch sizes is prohibitively expensive. The improvements from increasing the number of MC samples tend saturate K 100 on this task. Further, our results suggest that relaxed estimators may be preferred (if they can be used) for discrete random variables of smaller arity. For example, the GS estimator outperforms all straight-through through estimators for binary variables for both batch sizes. For large arities however, we find that straight-through estimators can perform competitively: Our estimator GR-MC1000 achieves the best performance overall and outperforms the GS estimator for 16-ary variables.
CONCLUSION
We introduced the Gumbel-Rao estimator, new single-evaluation non-relaxing gradient estimator for models with discrete random variables. Our estimator a aaolllackwellizationn of the state-of- the-art straight-through Gumbel-Softmaa estimator. It enjoys lower variance and can be implemented efficiently using Monte Carlo methods. In particular and in contrast to most other work, it does not require additional function evaluations. Empirically, our estimator improved final test set performance in an unsupervised parsing task and on a variational auto-encoder loss. It accelerated convergence on the objective and compared favorably to other standard gradient estimators. Even though the
8

---

Published as a conference paper at ICLR 2021
Table Our estimator (GR-MCK) outperforms other straight-through estimators for discrete-latent- space VAE objectives on the MNIST dataset and is competitive with the Gumbel-Softmax (GS) large arities. Best bound on the test negative log-likelihood selected on the validation set. Best straight-through estimator in bold, best estimator in italics. Lower is better.
BINARY
4-ARY
8-ARY
16-ARY
ESTIMATOR
20
200
B
20
200
200
20
B 200
GS
98.2
96.4
95.7
93.8
95.5
92.3
96.8
94.3
REINFORCE
202.6
121.4
173.7
122.2
203.9
124.9
169.4
129.5
105.5
103.1
106.2
104.5
107.2
105.
108.2
104.5
FOUST
101.5
97.8
ST-GS
100.7
97.1
99.1
93.7
98.0
92.8
98.8
92.6
GR-MC10
100.7
97.4
97.8
93.8
97.4
93.1
97.9
92.4
GR-ME MC-10000
100.6
96.8
97.5
94.0
96.8
92.2
97.3
92.4
GR-MCION
100.5
97.0
97.6
93.5
96.5
92.5
96.8
92.2
gains were sometimes modest, they were persistent and particularly pronounced when models must be trained at low temperatures or with small batch sizes. We expect that our estimator will be most effective in such settings and that further gains may be uncovered when combining our Rao- Blackwellisation scheme with annealing schedule for the temperature. Finally, hope that our work inspires further exploration of the use of Rao-Blad Blackwellisation for gradient estimation.
ACKNOWLEDGEMEET
MBP gratefully acknowledges support from the Max Planck ETH Center for Learning Systems. CJM grateful for the support the James D. Wolfensohn Fund at the Institute Advanced Studies in Princeton, NJ. Resources used in preparing this research were provided, in part, by the Sustainable Chemical Processes through Catalysis (Suchcat) National Center Competence in Research (NCCR), the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute.
REFERENCES
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation. arXiv e-prints, art. MEXiv:1308.3441 Aug 2013.
James Bergstra and Yoshua Bengio. Random search for hyper- hyper-paramede parameter optimization. Journal of machine learning research, 13(Feb):281-3055 2012.
David Blackwell. Conditional expectation and unbiased sequential estimation. Ann. Math. Statist., 18(1):105-110, 03 1947. doi: (0.1214/amm/1117773044
Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. arXiv preprint arXiv: 2015.
Jihun Choi, Kang Min Yoo, and Sang-goo goo Lee. Unsupervised learning of spec-ipe tree structures with tree-Istms. In CoRR, 2017.
Junyoung Chung, Sungjin Ahn, and Yoshua Bengio. Hierarchical multiscale recurrent neural networks. In International Conference on Learning Representations, 2017.
Andrew Davis and Itamar Arel. Low-rank approximations for conditional feedforward computation in deep neural networks. arXiv preprint 131. 2013.

---

Published as a conference paper at ICLR 2021
Andreea Gane, Tamir Hazan, and Tommi Jaakkola. Learning with maximum a-posteriori perturbation models. In Artificial Intelligence and Statistics, pp. 247-2566 PMLR, 2014.
Peter W Glynn. Likelihood ratio gradient estimation for stochastic systems. Communications of the ACM, (3(10):75-88) 1990.
Will Grathwohl, Dami Choi, Yuhuai Wu, Gcoffrey Roeder, and David Duvenaud. Backpropagation through the void: Optimizing control variates for black-box gradient estimation. In International Conference Learning Representations, 2018.
Shixiang Gu, Sergey Levine, Ilya Sutskever, and Andriy Mnih. Muprop: Unbiased backpropagation for stochastic neural networks. In International Conference on Learning Representations, 2016.
Serhii Havrylov, Germán Kruszewski, and Armand Joulin. Cooperative learning of disjoint syntax and semantics. arXiv preprint arXiv:1902.099: 2019.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical Reparamctrization with Gumble-Softmax. In International Conference on Learning Representations (ICLR 2017), 2017.
Diederik P Kingma and Max Welling. Auto-Encoding Variational Bayes. arXiv e-prints, art. arXiv:1312.6114 Dec 2013.
Wouter Kool, Herke van Hoof, and Max Welling. Estimating gradients for discrete random variables by sampling without replacement. In International Conference on Learning Representations, 2020.
Dirk P Kroese, Thomas Taimre, and Zdravko I Botev. Handbook of monte carlo methods, volume 706. John Wiley & Sons, 2013.
Yann LeCun and Corinna Cortes. MNIST handwritten digit database. URL http://yann. lecun. com/exdb/mnist, 2010. URL http:://aan lecun. ex/p/ ins
Runjing Liu, Jeffrey Regier, Nilesh Tripuraneni, Michael Jordan, and Jon McAuliffe. Rao- blackwellized stochastic gradients for discrete distributions. arXiv preprint arXiv:1810.0477 2018.
Ryan Lowe, Yi Wu, Aviv Tamar, Jean Harb, Picter Abbeel, and Igor Mordatch. Multi-agent actor- critic for mixed cooperaaive--oomettiivee environments. CoRR, abs/ 2017. URL http://arxiv,
02275
Chris J. Maddison. A Poisson process model for Monte Carlo. Tamir Hazan, George Papandreou, and Daniel Tarlow (eds.), Perturbation, Optimization, and Statistics. MIT Press, 2016.
Chris J. Maddison, Daniel Tarlow, and Tom Minka. A Sampling. In Advances in Neural Information Processing Systems 27, 2014.
Chris J. Maddison, Andriy Mnih, and Yee Whye Teh. The Concrete Distribution: A Continuous Re-
laxation of Discrete Random Variables. In International Conference Learning Representations, 2017.
André F. T. Martins, Tsvetomila Mihaylova, Nikita Nangia, and Vlad Niculae. Latent structure models for natural language processing. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics: Tutorial Abstracts, pp. 1-5, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/v1/P19-440 wwww www.aclwee. http://oo....tttoooo//oeeeeeeee
URL https
Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks. In Proceedings of the 3 International Conference International Conference Machine Learning- Volume 32, pp. I1-1791, 2014.
Andriy Mnih and Danilo J Rezende. Variational inference for monte carlo objectives. In Proceedings of the 33rd International Conference International Conference Machine Cearning-Vollumee 48, pp. 2188-2196 2016.
Igor Mordatch and Pieter Abbeel. Emergence of grounded compositional language in multi-agent populations. CoRR, abs/I 703.04908, 2017. h http::httt:/ttte httt::://aaxi..orrrab////33 /ttp:/111333999000300 04908
10

---

Published as a conference paper at ICLR 2021
Nikita Nangia and Samuel R. Bowman. Listops: A diagnostic dataset for latent tree learning, 2018.
Max Paulus, Dami Choi, Daniel Tarlow, Andreas Krause, and Chris Maddison. Gradient estimation with stochastic softmax tricks. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 5691-5704. Curran As- sociates, Inc., 2020, URL httppps://poccccpdiigggoooeo neurips cc wwtpppppee//2//1LL...
pdf
Adeel Pervez, Taco Cohen, and Efstratios Gavves. Low Bias Low Variance Gradient Estimates for Hierarchical Boolean Stochastic Networks. In ICML, 2020. URL
icml hct.l..c//stattcppppp....eee////
Par002/03626-pppr...oom Paper.
C. Radhakrishna Rao. Information and the accuracy attainable in the estimation of statistical parame- ters. In Samuel Kotz and Norman L. Johnson (eds.), Breakthroughs in Statistics: Foundations and Basic Theory, pp. 235-247, New York, NY, 1992. Springer New York.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine Learning, pp. 1778-2866, 2014.
Ruslan Salakhutdinov and lain Murray, On the quantitative analysis of deep belief networks. In Proceedings of the 25th international conference on Machine learning, pp. 872-879, 2008.
John Schulman, Nicolas Heess, Theophane Weber, and Pieter Abbeel. Gradient estimation using stochastic computation graphs. In Advances in Neural Information Processing Systems, pp. 3528- 3536, 2015.
George Tucker, Andriy Mnih, Chris J. Maddison, and Jascha Sohl-Dickstein. REBAR Low-variance, unbiased gradient estimates for discrete latent variable models. Neural Information Processing Systems, 2017.
Tim Vieira. Estimating means in finite universe, 2017. URL https://iimviiira. github. 2/hboppps220077777788sstmmaiiin iantaan-------iit-eee
2017.
Théophane Weber, Nicolas Heess, Lars Buesing, and David Silver. Credit assignment techniques in stochastic computation graphs. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 2650-26600 200022660222 2019.
Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8(3-4):229-256, 1992.
Improved variational Zichao Yang, Zhiting Hu, Ruslan Salakhutdinov, and Taylor abs/ 2017. URL autoencoders for text modeling using dilated convolutions. CoRR, http://arxiv.co wttparwwwwww.orrr///a//2//0


Reviews:
[{"review_id": "df7KKqW0LRD", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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).\n\nClarity: The paper is well written.\n\nOriginality: The use of Rao-blackwellization in the proposed way is, up to the best of my knowledge, novel.\n\nPros of the paper and significance: \n- Relaxation-based gradient estimators are widely used, and the proposed method may their variance quite significantly.\n- The proposed algorithm has a clear justification from a theoretical perspective, and admits a simple implementation.\n- The proposed algorithm does not require additional model evaluations, and thus may lead to large reductions in variance without incurring a high computational cost.\n- The proposed method leads to more efficient optimization at lower temperatures (lower temperature translates to lower bias, but often higher variances).\n\nCons: I'd say one thing that could be included are additional baselines in the experimental section. There are other estimators that may be use. For instance, you could compare against VIMCO. While this is a different type of estimator (non single evaluation, not based on relaxations), it could be interesting to see how the results compare using this estimator too.\n\nRecommendation: Accept (reasons in the \"pros\" list above).", "extra_scores": {}, "extra_text": {}}, {"review_id": "c621Ks-ECY", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n* This paper proposes a Rao-Blackwellized version of the straight-through gumbel-softmax gradient (STGS) estimator.\n* The Gumbel-Rao estimator remains single-evaluation (but multiple sample), does not have higher variance than the original straight-through estimator.\n* The estimator exhibits lower variance at lower temperatures in the experiments.\n\nContributions:\n* Proposes a single-evaluation estimator that cannot have higher variance than the STGS gradient estimator.\n* Demonstrates effectiveness of proposed estimator in terms of the variance of the gradient estimator and the ELBO on a toy task, a simple parsing task (ListOps), and a mixture model for MNIST.\n\nStrengths:\n* The method is simple and the computational overhead is very small compared to the original STGS estimator.\n* The empirical results support lower variance claims and effectiveness at lower temperature.\n\nWeaknesses:\n* I am not convinced that the relative gains from training at lower temperatures are significant.\n* The overall gains over ST-GS seem to be modest on MNIST as well as the L <= 50 setting in ListOps.\n* In the ListOps experiments, lower temperatures barely achieved better accuracy.\n\nDecision: Marginally below acceptance threshold\n* Improving gradient estimators for discrete latent variable models is an important problem.\n* The method is straightforward and the claims of performing better at lower temperatures are supported by empirical evidence.\n* However, the overall performance on the ListOps dataset is lower than related work [1], and there does not appear to be a large gain from low temperatures.\n\nQuestions:\n* The main argument of this paper hinges on the claim that lower temperatures result in lower bias of the gradient estimator. This claim seems reasonable, and is supported by figure 2b. Is there a proof or citation for it, and do we know more? It would be nice to know how variance and bias are traded off, as that would tell us how much (or how little) we could gain from training at lower temperatures.\n* Is there an explanation for the difference in performance between the 99% accuracy obtained in Havrylov et. al. 2019 [1] and the performance obtained at low temperatures in this paper?\n* How does this method perform versus the estimator proposed in Pervez et. al. [2], which is also single-evaluation?\n\nSuggestions:\n* The GR estimator is not guaranteed to have lower variance than ST-GS, just not higher.\n* Is there an application where lower temperatures are necessary for training? That would strengthen the argument.\n\n[1] Serhii Havrylov, German Kruszewski, and Armand Joulin. Cooperative Learning of Disjoint Syntax and Semantics. In Proceedings of NAACL 2019.\n\n[2] Pervez, A., Cohen, T., & Gavves, E. 2020. Low Bias Low Variance Gradient Estimates for Hierarchical Boolean Stochastic Networks. ICML 2020.\n\nEdited score after author comments.", "extra_scores": {}, "extra_text": {}}, {"review_id": "DVdOEf4LS7G", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\n\nThis paper introduces the Rao-blackwellization technique to reduce the variance of the straight-through gumbel-softmax gradient (STGS) estimator wrt the parameters of discrete distributions. The proposed method introduces almost trivial computational costs (relative to function evaluations) and is empirically and theoretically shown to systematically improve STGS. \n\nI don’t have a lot of nitpicking to make for this paper, as it is quite well executed. The proposed method is very clean and the improvement over the STGS baseline is very consistent, and makes it even competitive with concrete-relaxation in the discrete latent variable model experiment. \n\nDetails:\n\nWhy not show the curve of ELBO during training, but the arrival-time-at-certain-thresholds in Fig 2-c? \nLast paragraph of sec 5.4: The larger batch size here also reduces the variance of minibatch SGD, not just the variance of $\\nabla_{GR}$ in (13). In fact each instance is a different approximate posterior, which has different base variance. This makes the discussion in 3.3 a bit misleading. \n\nSuggestions:\nFor figure 1, perhaps visualize the variance of both separately. An improvement by 2 is not that meaningful if the variances of both are >> 2. \n\n-- After rebuttal\n\nThank you for revising the paper. I've read the revised section, and stand by my original evaluation. ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents a variance reduction technique to the Straight-Through version of the Gumbel-Softmax estimator. The technique is relying on the truncated Gumbel of Maddison et al. I share the excitement of the reviewers about this work and I expect this technique to further influence the field. 
Decision: Accept (Oral)