Cold Rao-Blackwellized Straight-Through Gumbel-Softmax Gradient Estimator
Authors: Alexander Shekhovtsov
Venue: Submitted to ICLR 2023 (2023)
Keywords: Gumbel-Softmax
categorical variables
Concrete distribution
gradient
straight-through
VAE
quantization

Abstract: The problem of estimating the gradient of an expectation in discrete random variables arises in many applications: learning with discrete latent representations, training neural networks with quantized weights, activations, conditional blocks, etc. 
This work contributes to the development of the popular Gumbel-Softmax family of estimator, which is based on approximating argmax with a temperature-parametrized softmax. The state-of-the art in this family, the Gumbel-Rao estimator uses internal MC samples to reduce the variance.
We show that in the limit of zero temperature the internal integration has a closed form solution. The limit estimator, called ZGR, has a favorable bias and variance, is simple to implement and computationally inexpensive and is obviously free of the temperature hyperparameter. Furthermore, ZGR is unbiased for the class of quadratic functions of categorical variables and can be decomposed into a sum of two simple but not very well performing on their own estimators: the straight through estimator and the DARN estimator. Experiments thoroughly validate the method.

Under review as a conference paper at ICLR 2023
COLD RAO-BLACKWELLIZED STRAIGHT-THROUGH GUMBEL-SOFTMAL GRADIENT ESTIMATOR
Anonymous authors Paper under double-blind review
ABSTRACT
The problem of estimating the gradient of an expectation in discrete random vari- ables arises in many applications: learning with discrete latent representations, training neural networks with quantized weights, activations, conditional blocks, etc. This work is motivated by the development of the Gumbel-Softuna fam- ily of estimators, which are based on approximating argmax with temperature- parametrized softmax. The state-of-the art in this family, the Gumbel-Ran esti- mator uses internal MC samples reduce the variance and appars to improve with lower temperatures. We show that it possesses a zero temperature limit with a surprisingly simple closed form defining our new estimator called ZGR. has favorable bias and variance properties, is easy to implement, computationally in- expensive and is obviously free of the temperature hyperparameter. Furthermore, it decomposed estimator two basic but not very well performing on their own estimators. Un- like them, ZGR is shown to be unbiased for the class of quadratic functions of categorical variables. Experiments thoroughly validate the method.
the average of the straight through estimator and the DARN
INTRODUCTION
Discrete variables and discrete structures are important in machine learning. For example, the se- mantic hashing idea (Salakhutdinov & Hinton, 2009) is to learn a binary vector representation such that semantically similar instances (e.g. images or text documents) would have similar bit repre- sentations. This allows for a quick retrieval via lookup table or via nearest neighbor search with respect to the Hamming distance of binary encodings. Recent work employs variational autoen- coders (VAFs) with binary latent states (Shen et al., 2018; Dadaneh al., 2020; Nanculef et al., 2020). Another example is neural networks with discrete (binary or quantized) weights and activa- tions. They allow for a low-latency and energy efficient inference, particularly important for edge devices. Recent results indicate that quantized networks can achieve competitive accuracy with better efficiency in various applications (Nie et al., 2022).
VAEs and quantized networks are two diverse examples that motivate development and the experimental benchmarks. Other potential applications include conditional computation (Bengio et al., 2013a; Yang et al., 2019; Bulat et al., 2021) reinforcement learning (Yin et al., 2019), learning task-specific tree structures for agglomerative neural networks (Choi et al., 2018), neural architecture search (Chang et al., 2019) and more.
The learning problem in the presence of stochastic variables is usually formulated as minimization of the expected loss. The gradient-based optimization requires gradient of the expectation in the probabilities of random variables (or parameters of networks inferring those probabilities), Unbiased gradient estimators have been developed (Williams, 1992; Grathwohl et al., 2018: Tucker et al., 2017: Gu et al., 2016). These estimators work even for non-differentiable losses, however their high variance the main limitation. More recent advances (Yin et al., 2019; Kool al., 2020; Dong et al., 2020; 2021; Dimitriev & Zhou, 2021b;a) reduce the variance by using several cleverly coupled samples. However, the hierarchical (or deep) case is not addressed satisfactory. There is experimental evidence that the variance grows substantially with the depth of the network leading to very poor performance training (e.g., Shekhovtsov al. 2020, Fig.C.6, C.7). Furthermore, existing extensions of coupled sampling methods to networks with L dependent layers (Dong al., 2020: Yin et al., 2019) apply their base method in every layer, requiring several complete forward

---

Under review as a conference paper at ICLR 2023
Table 1: Computation complexity of estimators in hierarchical network with L dependency layers
of K Kwway categorical variables. All methods require backward pass. educated guess.
Method
Unbiased
Forward passes
Cost per cat. variable per pass
ST. GS, GS-ST, ZGR
X
O(K)
GR-MC(M)
X
O(MK)
REINFORCE
O(K)
RF(M) (Kool et al., 2019)
2
O(K)
ARSM (Yin ct al., 2019)
O(
0(1)
DisARM-* (Dong et al., 2020; 2021)
O
O(K)
CARMS(M) Dimitriev & Zhou (b;2lb;a)
O(M
O(M
samples per layer. The computation complexity thus grows quadratically with the number of layers, as summarized in Table 1.
A different family of methods, fitting practical needs for deep models better, exploits continuation arguments. It includes ST variants (Bengio et al., 2013b; Shekhovtsov & Yanush, 2021; Pervez et al., 2020) and Gumbel-Softmax variants (to be discussed below). These methods assume the loss function to be differentiable and try estimate the derivative with respect to parameters of a discrete distribution from the derivative of the loss function. Such estimators can be easily incorporated into mack-propagation by adjusting the forward and backward passes locally for every discrete variable. They are, in general, biased because the derivative of the loss need not be relevant for the discrete expectation. The rationale though is that may be possible to obtain a low variance estimate at price of small bias, e.g. for sufficiently smooth loss function (Shekhovtsov & Yanush, 2021).
Gumbel Softmax (Jang et al., 2017) and the concurrently developed Concrete relaxation (Maddi- son et al., 2017) enable differentiability through discrete variables by relaxing them to real-valued variables with distribution approximating the original discrete distribution. The tightness of the relaxation is controlled by the temperature parameter t >> 0. The bias can be reduced by decreasing the temperature, but the variance grows as (1/tt) (Shekhovtsov, 2021). Gumbell-Softaa Straight- Through (GS-ST) heuristic (Jang et al., 2017) uses discrete samples on the forward pass, but relaxed ones on the backward, reducing side biases (see below). The Gumbel-Ran (GR) estimator (Paulus et al., 2021) is recent improvement of GS-ST, which can substantially reduce its variance by local expectation. However the local expectation results in an intractable integration in multiple variables. which is approximated by sampling. The experiments (Paulus et al., 2021) suggest that this estimator performs better at lower temperatures, which requires more MC samples. Therefore the computation cost appears to be major limitation.
Inspired by the performance of GR at low temperatures, we analyze its behavior for temperatures close to and in the limit of the absolute zero. Note that because the underlying relaxation becomes discrete in the cold limit, leading to explosion of variance in both GS and GS-ST, they do not have a zero temperature limit. is not obvious therefore that GR would have one. We prove that it does and denote this limit estimator as ZGR. In the case of binary variables we give an asymptotic se- ries expansion of GR around the already known DARN( }) estimator by Gregor et al. (2014) (also re-discovered as importance reweighed ST by Pervez et al. 2020). the general categorical case, we obtain the analytic expres- sion of ZGR and show that it is new estimator with simple formula (ST+DARN) We show that ZGR is unbiased for all quadratic functions categorical variables and experimentally show that achieves a useful bias-variance tradeoff. We also contribute refined experimental comparison of GS family with unbiased estimators.
0 and show that ZGR has a simple analytic expression, matching
2 BACKGROUND
Let x be a categorical random variable taking values in IC {0.... 1} with probabilities p(z: parametrized by nT. Let Φ(x) E Rd be embedding of the discrete state in the vector space. Categorical variables are usually represented using 1-hot embedding, in this case d = K. index. For binary and quantized variables we will adopt
The value of itself will still be used the embedding o(x) = æ, in this case d = 1.
2

---

Under review as a conference paper at ICLR 2023
Let C: Rd + IR be differentiable loss function. We will denote the Jacobian of L with respect to an input Φ as Is = dc is the transposed gradient L with respect to Φ, nevertheless, we will informally refer to such Jacobians as gradients. For brevity, let us also use a shorthand [(x) (()(x)) The goal is estimate the gradient of the expected loss
dø
Jn = di [[[()) = dn Σ, ((x)p(2:n)
(1)
In a network with many (dependent) categorical variables has proven efficient (e.g., Bengio et al. 2013a; Jang et al. 2017; Paulus et al. 2021; Shekhovtsov & Yanush 2021; Pervez et al. 2020) to con- sider estimators that make an estimate of Jn based on the gradient Jø ac()(x)) at sample from dø
p(a; 17) in (1). Such estimators can be easily extended to losses in multiple discrete variables (e.g. defined by stochastic computation graph) by simply applying the elementary estimator whenever the respective Jacobian is needed in backpropagation. In this case J, is the gradient in a specific variable or intermediate activation 7) at joint sample.
Let us review the basic estimators with which we will work theoretically.
REINFORCE We can the log-derivative trick rewrite
d di E. ((x)p(x;))) -. [xx)
dp(a:n) de
Σ. ((x)p((; 7) p(am)) du
= E [(xx) d (ogg((a))
(2)
dy
and define the REINFORCE estimate (Williams, 1992) to be JRF p(a; 7)) This estimator is clearly unbiased as E[J = Jnp but may have high variance.
((x) log p(ain)) where ~
ST Let ((7) = [[)(r) = E. (())p(;;;))))))))))())))<) v) - the mean embedding in I& under the current distri- bution of x. The Suraight-Throute estimator
JS = J do(r))
(3)
Note, ST estimators of different empirical forms exist. The present definition for categorical vari- ables is the same as, e.g., by Gu et al. (2016) and is consistent with Hinton (2012) and Shekhovtsov & Yanush (2021) in the binary case.
DARN Gregor et al. (2014) use LR with baseline b(x) f(x) + [ - (())). which is a first order Taylor approximation of L about (xx) evaluated at some point ø. This results the estimator
)ARR( Jo(((r) - .) d p(zw))
(4)
dry
When æ binary, the choice e- - (xx) the mean embedding under the uniform distribution, ensures that the estimator is unbiased for any quadratic function. However, for categorical variables no with such property exist. Gu al. (2016) have experimentally tested several heuristic choices for 6, including ő = )'n), and found that none performed well in the categorical case.
GS The Humbel-Soffmax (GS) estimator (Jang et al., 2017) is relaxation of the Gumbel-argna -argmax sampling scheme. Let 6₂ log (æ==;;,) n)) Let G5 ~ Gumbel(O, 1), & E KC where Gumbel(O. 1) Gumbel distribution with edf F(u) - Then
x arg maxx((( + G+)
(5)
a sample from p(z; 17). The relaxation is obtained by using a softmax instead of arg max. This construction assumes one-hot embedding ø and creates relaxed (continuous) samples in the simplex K Formally, introducing temperature hyperparameter I, it can be written as
g softmax((B + G)) =: softmax, + G);;
(6a)
JGS
dc(d) dc(²) dē de
(6b)
drj
There are two practical concerns. First, the loss function evaluated at relaxed sample, which in a large computation graph can offset estimates of all other gradients even not related to the discrete variable we relax. This effect can be mitigated by using a smaller temperature, causing relaxed samples ò to concentrate in a corner of the simplex. However, and this the second concern, the variance of the estimator grows as (()) if is decreased towards zero (Shekhovtsov, 2021).
3

---

Under review as a conference paper at ICLR 2023
GS-ST The Striihh-Thhouug Gumbel-Softme estimator (Jang al., 2017) is empirical modi- fication of GS, addressing the first concern above. It uses discrete samples in the forward pass but swaps in the Jacobian of the continuous relaxation in the backward pass:
G. Gumbel 1), k E KG;
(7a)
arg max; (0x + Gk);
(7b)
softmax, (8 + G);
(7c)
JGS-ST ac(o(z)) dé dø
(7d)
ab
Notice that the hard sample ((z) and the relaxed sample e are entangled through G. Although, has the law of p(æ; v) as desired, not biasing other variables, there is still bias in estimating the gradient in 0, which typically larger than that of GS (see Fig. 1). To make the bias smaller the temperature t should be decreased, however, the variance still grows as (11/t) (Shekhovtsov, 2021). Values of t between 0.1 and 1 are used in practice (Jang et al., 2017).
GR Notice that the forward pass in GS-ST is fully determined by alone and the value of G that generated that x is needed only in the backward pass. Paulus et al. (2021) proposed that the variance of GS- St can be reduced by computing the conditional expectation in G|æ,, leading to the Gumbel-Rao estimator:
JØ = EGI JJT--S
(c)
aC(o(z)) de
BG
de
(8)
Because the value of the loss ((x) and its gradient do not depend on the specific realization of G|r, enabling the equality above, the expectation is localized and can be computed in the backward pass. However, this expectation is in multiple variables and is not analytically tractable. Paulus et al. (2021) Monte Carlo integration with M samples from G|a. In their experiments they report improvement of the mean squared error of the estimator when the temperature was decreasing from 1 down to 0.1. The trend suggests that it would improve even further below t 0.1 provided that the conditional expectation is approximated accurately enough, indicating that the variance does not grow with the decrease of temperature in contrast (L//t) asymptote for GS and GS-ST and so there might be meaningful cold limit.
3 METHOD
Given the experimental evidence about the GR estimator, we took the challenge to study its cold asymptotic behavior, i.e. for -> O. The temperatured softmax in (7c) approaches non- differentiable arg max indicator this limit and we have to handle the limit of the GR estimator with care obtain correct results. We first analyze the binary case, where derivations are substan- tially simpler. Proofs of all formal claims can be found in Appendix A.
3.1 BINARY CASE
In the case with two categories we can simplify the initial ST-GS estimator as follows. We assume x E {0. 1} and (xx) = x. The argmax trick can be expressed as x = [6, + G1 2 Bu + Gul, where I-I is the Iverson bracket. It is convenient assume that the distribution of is parametrized so that 1(x=1:1)) = (n)) the logistic sigmoid function. This is without loss of generality, because any other parametrization will result in just extra deterministic Jacobian. Recalling that 0₂ logp(==</ k; 7), we have 0 0g 17. Next, denoting Z = G1 Go. we can write the argmax trick compactly as x In + Z Z oj. Being the difference of Gumbel(O.1) variables Z follows the standard logistic distribution (with cdf (z)). The GR estimator of gradient in simplifies as
Z Logistic(O.1);; -- O]; ã= +o Z);
JGK
deix)
di Ez1 day
where (u) = (u/t) is the temperatured logistic sigmoid function. Although there is no closed form, can compute, with careful limit-integral exchange, the series expansion around t 0.
Proposition 1.
J - C'(x) Clw (1) (2r 1)c, log(2))) + 0(())
p(z) pz
(10)
where pz is the logistic density: (n( a()) and e 2p
4

---

Under review as a conference paper at ICLR 2023
Corollary 1. In the limit 0 the GR estimator becomes the DARN()) estimator.
This establishes an interesting theoretical link, but discovers no new method for practical appli- cations, as the DARN estimator is already known. Using the same expansion, can study the asymptotic bias and variance of GR around
Corollary 2. The mean and variance of the GR estimator (9) in the asymptote t + 0 are:
_[ p(1 p) (1)) +()))(0))))))))))))))))))))))))) (C'(1) ((())= O(1)),
(Ila)
M[ (p(1 p)) [ -
b)2 +) ) - b))ēzt) + (tt))
(11b)
where p p= (()), (1))) b= ( 1-p and é (2p 1) (2g
p
This allows make some predictions, in particular in the case of a linear objective L. In this case the bias (t and the squared bias is O(t*). Therefore the MSE is determined by the variance alone up O(t)) The dependence of variance for a linear objective is negative in the first order term. Therefore the temperature corresponding to the minimum MSE will be non-zero.
3.2 GENERAL CATEGORICAL CASE
In the general categorical case, the analysis is more complicated (exchange of the limit and multi- variate integral over Gl but gives novel result:
Theorem (ZGR). The Gumbel-Rar estimator for one-hot embedding & the limit of zero tem- perature is given by
2 (Ja. - Jo 1p(x==1;1)
x; if - x.
JZGR
(12)
2 ) (Js Jo (p(x=j:n)
Proposition ZGR estimator decomposes as
JEGR (JST + (DARN(Ó(H))
(13)
i... with the choice ((1) in DARN.
This is consistent with the expression for the binary case given by Corollary 1 by noting that in this exactly DARN(1/2). As we have defined ST and DARN estimators for a general embedding. the form (2) is a valid expression of ZGR for any embedding (a change of the embedding is just linear transform). This form shows surprising connection. Neither ST nor DARN estimators perform particularly well in categorical VAEs on their own (Gu et al., 2016; Paulus et al., 2021). However ZGR, being effectively their average, appears superior. It has the following property, truly extending the design principle of binary ( to the categorical case:
case (ST + D-RN(
Theorem 2. ZGR is unbiased for quadratic loss functions L with arbitrary coefficients.
Both ST and (āR)( are unbiased for linear functions and the thcorem shows their biases for quadratic functions exactly opposite and cancel in the average. Because the bias of ST-GS, GR and GR-MC the same and we have shown that ZGR is the limit of GR at f - 0, the following is straightforward.
Corollary 3. ST-GS and GR are asymptotically unbiased for quadratic functions with t -
These results extend to multiple independent discrete variables follows.
Corollary 4.
Let Zi...., Tn be independent categorical variables and C(x V.,) be such that for all and all configurations x the restriction L(x:) is quadratic function. Then ZGR is unbiased.
The unbiased property for quadratic function gives us some intuition about applicability limits of ZGR. Namely. if the loss function is reasonably smooth, such that it can be approximated well by quadratic function, we expect gradient estimates to be accurate. Compared to ST, which is unbiased for multilinear functions only, hypothesize that ZGR capture interactions more accurately.
5

---

Under review as a conference paper at ICLR 2023
4 EXPERIMENTS
We will compare ZGR with Gumbel-Softmax (GS), Straight-through Gumbel-Softmax (GS- ST) (Jang et al., 2017), Gumbel-Rao with MC samples (GR-MC) (Paulus et al., 2021) and the ST estimator (3). We also compare to the REINFORCE with the leave-one-out baseline (Kool et al., 2019) using M 2 inner samples, denoted RF(M), which is strong baseline amongst unbiased estimators. some tests we include ARSM (Yin et al., 2019), which requires more computation than RF(4) but performs worse. See Appendix B.2 for details of implementations.
4.1 DISCRETE VARIATIONAL AUTOENCODERS
We follow established benchmarks for evaluating gradient estimators in discrete VAEs. We use a MNIST data with a fixed binarization (Yin et al., 2019) and Omniglot data with dynamic binariza- tion (Burda et al., 2016; Dong et al., 2021). We use the encoder-decoder networks same as (Yin al., 2019; Dong al., 2021) up to the following difference. We embed categorical variables with 2 states as {-1, vectors. This allows to vary the number of categorical variables and categories while keeping the decoder size the same, Please see full details in Appendix B.
4.1.1 ZERO TEMPERATURE LIMIT AND ESTIMATION ACCURACY IN VAE
First, measure the gradient estimation accuracy at particular point of VAE training, comparing GS family of estimators at different temperatures as in (Paulus et al., 2021, Fig 2b.). This referenced plot shows steady decrease of MSE with the decrease of temperature down to 0.1 and we were expecting ZGR to achieve the lowest MSE.
As evaluation point take VAE model after 100 epochs of training with RF(4). We then measure the bias and variance of all gradient estimators as follows. The loss function is the average ELBO of a fixed random mini-batch of size 200. Let X R.d be the reference unbiased estimator and Y E IRd be the tested estimator. We want to measure the average over parameters (= dimensions of the gradient) squared bias, which can be written 62 1/ [[[[]] - E[Y]]| We obtain nı 10a independent samples X, from RF(4) and n2 10 independent samples Y; from the tested estimator and compute an unbiased estimate of b²:
8- À
frll2
(14)
where ji the sample mean of X and V1 is the average (over dimensions) sample variance of X and (O and 1/2 are likewise for Y. The variance of the evaluated estimator is just V2.
The results are shown in Fig. All of ST-GS and GR estimators share the same bias according to the theory but differ in variance. GS estimator is asymptotically unbiased but the variance grows as O(1/t). We observe that the variance is by several orders larger than the squared bias. Respectively, the mean squared eπτоr, which is the sum of the variance and the squared bias, is dominated by the variance alone for all methods. This is in high contradiction with MSE analysis in (Paulus et al., 2021, Fig 2b.), which deem incorrect.
Note however, that the effect of the bias-variance trade-off in learning is not straightforward: com- mon optimization methods use momentum as effective way of variance reduction while the bias can be correlated across iterations and may potentially accumulate. Therefore MSE is not necessar- ily indicative of performance and ZGR still fulfills our expectations of the zero limit estimator: has the limiting bias, which is the lowest in the GR/GS-S ST family, and the limiting variance which moderate. In particular the variance is better than that of GR for temperatures below about 0.1 and is comparable to the variance of RF(4). More tests at different stages of training are provided in Appendix B.4.
4.1.2 TRAINING PERFORMANCE IN VAE
Next compare training performance of several methods, in the same setting prior work. In particular we use the same Adam optimizer, batch size, learning rate and training duration as (Dong et al., 2021; Dimitriev & Zhou, 2021b). Full details given in Appendix B.3.
Table 2 presents results for two datasets and different splitting of latent bits into discrete variables (from binary to 64-way categorical). We observe the following: ST performs the worst and

---

Under review as a conference paper at ICLR 2023
Table VAE training negative ELBO for binary MNIST 0.5 and the down-sampled and dy- namically binarized Omniglot. Each value is the mean over random initializations and confidence intervals are (mmxx m of the 3 runs. Bold results are the three best ones per configuration.
Number of Categories Categorical Variables
Method
C2 V192
C4 V96
C16 V48
C64 V32
GS(t=0.1)
91.2+0.1
84.4=0.5
02.8+0.5
86.9=0.9
()-
92.0+t.4
85,4=0.2
84.1
90.0=0.5
MNIST-B
GR-MC((=011M=100
88.810.2
82.7.0.1
81.0.10.1
82.4
GR-MC(-0.1LM=100)
88.0.10.5
82.410.2
80.5.10.3
81.7.10.5
ZGR
88.6110.2
83.0
80.610.44
81.9-10.1
ST
105.0.10.2
105.4.0.1
106.0.1..1
106.7:0.1
RF(M=2)
92.1+t.44
86.3
88.6:1..5
96.7=0.1
RF(M=4)
88.241..3
82.7=0.1
82,2:0..5
87.2=0.7
GS((=0.1)
117.6:0.2
116.0=0.2
117.0.0.1
120.6=0.2
Ommiglot-22-D
(S-)T(:=0.1)
119.410.1
116.7=0.2
118.520.1
123.2=0.33
(R-M
116.6:10.2
114.6-0.11
115.710.11
117.8=0.2
CRR-MC(tO.1,M-II1
116.5+00.1
114.4-0.1
115.2+00.1
116.9=0.3
ZGR
116.6+0.44
114.5-0.1
115.4+0.1
117.0-0.2
ST
130.240.1
130.5=0.1
131.5+002
132.0=0.1
RF(M=2)
120.5
118.810.1
122.0.1..3
127.510.44
RF(M=4)
117.110.1
115.8
117.910.1
120.7
GS
ST-G5
10
10ah
GR-MC10
GR MC100
RFIM-21
Avenge Squared
RFIM=4]
20-3
Average
ANSH
65
ST-C5
6R MC10
MCL00
10a2
10
10-
Figure 1: Gradient estimation accuracy in VAE on MNIST-B. Average (per parameter) squared bias (left) and variance (right) of gradient estimators versus temperature at model snapshot after 100 epochs of training with RF(4). Confidence intervals are 95% empirical intervals of 100 bootstrap samples of the estimates (very narrow for variances).
GS GR-MCIM-10)
65
0.82
Fiial training loss Inagative BLBOL
92
MCIM
ZGR
GR-NC10
.44
90
1.36
26R
b.98
ARSM
D1.11
82-
0.0
0.1
0.2
0.3
0.4
05
0.6
0.7
0.8
0.9
1.0
Time [ms]
Tomporature
Figure 2: Left: time [ms] of forward-backward pass GPU (16 variables, 16 categories, single batch of size 200, MNIST-B). GS and GS-ST estimators are shipped with Pytorch while other esti- mators are implemented by using tensor operations in Pytorch (no loops). ZGR significantly faster than GR-MC and one of the fastest methods overall. Right: Training ELBO after 500 epochs of training with different choices of temperature (MNIST-B. variables 64 categories). The regret of ZGR w.r.t. GR-MC at the optimal temperature is insignificant. The plot shows the mean from random initializations with confidence intervals = min)/2 of the runs.

---

Under review as a conference paper at ICLR 2023
blame its high bias (cff. Fig. 1) 2) ZGR performs no worse than GR-MC variants. In Fig. 2 additionally verify that at no other temperature GR-MC can achieve significantly better results; 3) ZGR outperforms RF(2) and RF(4), significantly so with more categories. Finally, we measure the (orientational) computation time in Fig. 2 and observe that ZGR is faster than both GR-MC and RF(2), consistently with theoretical expectations in Table 1.
Finally, according to the published results in a similar setup, the recent unbiased methods (Dong al., 2021; Dimitriev & Zhou, 2021b) appear improve only marginally over RF with an equal number of samples, i.e. the difference is much smaller than between ZGR and RF(2). We leave a peer-to-peer comparison to future work.
4.2 QUANTIZED NEURAL NETWORKS
The mainstream progress in training quantized and binary neural networks, following Hubara et al. (2017), has been achieved so far using empirical variants of ST estimator (with different clamping rules, etc.) applied to deterministically quantized models, where there is no gradient to be estimated in the first place. A sound training approach is to consider a stochastic relaxation, replacing all discrete weights and activations by discrete random variables, leading in the binary case to stochastic binary networks (Peters & Welling, 2018; Roth et al., 2019; Shekhovtsov & Yanush, 2021).
We will consider sarameter-efffcient stochastic relaxation for quantization of Louizos al. (2019). In this model the distribution of quantized weight or activation æ is defined by a single real-valued input 7) via: = [1 + z] z]xc. where L-] rounds to the nearest integer in K. and z is an injected noise, such as logistic noise. Therefore x' is a discrete integer variable with distribution determined by 7.. In Appendix B.5 we give an comprehensive evaluation of bias and variance of estimators for single such quantization unit. In deep network, the pre-activation input depends on the weights of the current layer as well as on the preceding activations (both stochastically quantized), causing a hierarchical dependence.
We train a convolutional network, closely following Louizos et al. (2019) and test on MNIST and FashionMNIST. We do not quantize the input (it has 8 bit resolution in the dataset), the first and last weight matrices are quantized 4 bits. All inner layer weights and activations are quantized to bits or below. We form two real-valued baselines applying ReLU or Clamp (z H> min m 0). K 1)) as activations instead of quantization. We expect that with enough bits, quantized training should achieve performance of at least the Clamp variant. We evaluate training with logistic injected noise (as in Louizos et al. (2019)) and triangular noise with the density p(=) max(0, |=|). For GS-ST variants we enable high temperatures (0.5, 1, 2) as recommended by Louizos et al. (2019). See Appendix B.5 for details of the experimental setup.
The results are presented in Table 3. We see that the best results for are obtained with estimators hav- ing the least variance, prominently ST and GS (S- It suggests that the bias is less detrimental in this application. In Fig. 3 we measured bias and variance along the training trajectory of ZGR following the same methodology as in Section 4.1.1 with 10a samples from the reference RF(4) estimator and 102s samples from candidate estimators. The two methods with the lowest variance are exactly ST and GS-ST(t=2) while the bias was hard to measure accurately to draw any conclusions. More generally, the ranking of results in Table is quite similar to the ranking of variance in Fig. 3. In particular, variance of RF(4) is several orders larger than that of biased estimators and its test accuracy is completely out of the competition.
Regarding performance of ZGR observe the following: 1) it outperforms GR-MC with temper- ature 0.I, more clearly so on Fashion-MNIST (Table B.2), while being cheaper and simpler; 2) is close in performance the best results obtained with more biased estimators.
DISCUSSION / CONCLUSION
On the theoretical side, we showed that GR estimator has zero temperature limit, computed this limit, studied its properties and connected to the existing estimators. Despite we derived ZGR from the Gumbel-Softmax family, we do not consider it to be a proper member of this family. The straight- through heuristic in GS-ST disposes of relaxed samples on the forward pass. The zero temperature limit disposes of them also on the backward pass, leaving essentially nothing form the Gumbel-
8

---

Under review as a conference paper at ICLR 2023
Table 3: MNIST classification test error[%] in deterministic mode (no injected noises at time) for different bit-width per weigh and activation (T denotes temaary.. Hyperparameters are selected on the validation set. Reference test errors: ReLU 0.69%. Clamp 0.64%.
Weights [bits] Activations [bits]
Method
2/2
T/T
T/1
1/1
GS-ST((=2)
0.75 + 0.03
0.73 = 0.05
0.74 + 0.04
0.83 + 0.05
GS-ST(t=1)
0.75 0.03
0.78 0.08
0.88
0.11
0.86 + 0.04
Logistic Noise
GS-ST((=0.S)
0.84 1 0.05
1
0.91
0.11
1.02 0.08
GR-MC((=0.5M=11))
0.01
0.78 1 0.06
0.88 0.10
0.99
GR-
0.82 0.04
0.85 0.007
0.96
1.04
ZGR
0.84 0.09
0.86 0.07
0.06
1.07 1
ST
0.76 = 0.04
0.77 = 0.05
0.06
0.81 0.02
RF(M=2) L
38.9
35.9
51.0
52.3
RF(M=4) L
31.6
39.8
35.6
55.4
GS-ST(:=2)
0.73 = 0.05
0.72 = 0.07
0.76 0.05
0.85 0.06
Triangular Noise
GS (t=))
0.71
0.77 0.03
0.74 0.03
0.90 + 0.07
()))
0.73
0.84
0.81 + 0.04
0.88 + 0.11
10
0.70
0.06
0.71 0.05
0.83 0.05
GR-MC(==0.1M==1)))
0.78 0.01
0.75
0.80 + 0.03
0.98 - 0.02
ZGR
0.75 0.05
0.78
0.80 1 0.06
0.90 1 0.08
ST
0.70 0.10
0.73 1 0.05
0.71 0.04
0.76 1 0.06
RF(M=2) T
37.8
43.2
68.3
69.7
RF(M=4) T
31.2
41.8
56.5
63.9
21
1,6-11
203
1011
70rk
003020
GR2 6 10
002015
60201
602005
10-4
127
45
002800
20
102 120
140
200
1.0
100 140 166
102
1.2
Timelma
Figure 3: MNIST classification gradient estimation accuracy and computation cost. Left. middle: average (per parameter) squared bias (rsp. variance) in the first layer the network of different estimators along the training trajectory of ZGR (ternary weights/activattinns, triangular noise). Right: Time of gradient estimate per batch (of size 128) in our implementation. Our RF implementation apparently not efficient.
softmax relaxation design. On the other side, we showed that it is unbiased for quadratic functions, generalizing the key property of DARN(?) ) to the categorical case. We believe that such rationale can be put forward for obtaining improved biased estimators.
On the practical side, ZGR is extremely simple, versatile and computationally inexpensive. In VAE can replace GR-MC family completely, reducing the computational burden and hyperparameter tuning. It outperforms state-of-the-an unbiased estimators of comparable computation complexity and ST by large margin. In quantized training it performs close to ST and can fairly replace low-temperature GR-MC variants, while unbiased estimators completely out of the competition. Thus, across the two corner applications, ZGR is the only estimator which is computationally cheap and performnr....... While VAE unbiased estimators performs well (and we not need to worry about the bias) and in quantization simple ST performs well, the above results suggest that there should cases where ZGR would perform significantly better than both RF and ST. This may be the case for example when considering a different learning formulation such as Bayesian learning of quantized weights or learning with the multi samlle objective Raiko et al. (2015).

---

Under review as a conference paper at ICLR 2023
REPRODUCIBILITY STATEMENT
Proofs of all formal claims are presented in Appendix A. Details of the experiments are described in Appendix B. The source code of our implementation will be made publicly available upon pub- lication. During the review period, we will be happy to answer questions and share the code with reviewers confidentially through the OpenReview platform.
REFERENCES
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochas- tic neurons for conditional computation. arXiv preprint arXiv:1308.34:: 3432, 2013a.
Yoshua Bengio, Nicholas Léonard, and Aaron C. Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. CoRR, 2013b.
Adrian Bulat, Brais Martinez, and Georgios Tzimiropoulos. High-capacity expert binary networks. In ICLR, 2021.
Yuri Burda, Roger B. Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. In ICLR, 2016.
Jianlong Chang, xinbang zhang. Yiwen Guo, Gaofeng Meng. Shiming Xiang, and Chunhong Pan. DATA: Differentiable architecture approximation. In NewrlPS, 2019.
Jihun Choi, Kang Min Yoo, and Sang-goo Lee. Learning compose task-specific tree structures. In AAAI, 2018. ISBN 777-1157735--8 800-8.
Siamak Zamani Dadaneh, Shahin Boluki, Mingzhang Yin, Mingyuan Zhou, and Xiaoning Qian. Pairwise supervised hashing with Bernoulli variational auto-encoder and self-control gradient estimator. ArXiv, 2020.
Alek Dimitriev and Mingyuan Zhou. CARMS: Cangoriialaatiiheii--EEFFOOOCO multi-sample gradient estimator. In NeurlPS, 2021a.
Aleksandar Dimitriev and Mingyuan Zhou. ARMS: Amiitheii-REIINOOOCC--mll----sa- gradient for hinary variables. In ICMI, pp. 2717-2727, 2021b.
Zhe Dong, Andriy Mnih, and George Tucker. DisARM: An antithetic gradient estimator for binary latent variables. In NeurlPS, pp. 18637-18647, 2020.
The Dong, Andriy Mnih, and George Tucker. Coupled gradient estimators for discrete latent variables. In NeurlPS, 2021.
Will Grathwohl, Dami Choi, Yuhuai Wu, Geoff Roeder, and David Duvenaud. Backpropagation through the void: Optimizing control variates for black-box gradient estimation. In ICLR, 2018.
Karol Gregor, Ivo Danihelka, Andriy Mnih, Charles Blundell, and Daan Wicrstra. Deep autoregressive net- works. In JCML, 2014.
Shixiang Gu, Sergey Levine, Ilya Sutskever, and Andriy Mnih. Muprop: Unbiased backpropagation for stochas- tic neural networks. In ICLR, May 2016.
Geoffrey Hinton. Lecture 15d Semantic hashing 3:05 3:35, 2012.
Itay Hubara, Matthieu Courbariaux, Daniel Soudry, Ran El-Yaniv, Yoshua Bengio. Quantized neural networks: Training neural networks with low precision weights and activations. J. Mach. Learn, Res., (I), jan 2017.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In ICML, 448-456,77 2015.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with Gumbel-softmax, In ICLR, 2017.
Diederik Kingma and Max Welling. Auto-encoding variational Bayes. CoRR, 2013.
Wouter Kool. Herke Hoof. and Max Welling, Buy 4 REINFORCE samples, get baseline for free!, 2019.
Wouter Kool, Herke van Hoof, and Max Welling. Estimating gradients for discrete random variables by sam- pling without replacement. In ICLR, 2020.
10

---

Under review as a conference paper at ICLR 2023
Christos Louizos, Matthias Reisser, Tijmen Blankevoort, Efstratios Gavves, and Max Welling. Relaxed quan- tization for discretized neural networks. In ICLR, 2019.
Chris Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. ICLR, 2017.
Henrick J. Malik and Bovas Abraham. Multivariate logistic distributions. The Annals of Statistics, (3):588-590, 1973.
Ricardo Nanculef, Francisco Alejandro Mena, Antonio Macaluso, Stefano Lodi, and Claudio Sartori. Self- supervised bernoulli autoencoders semi-supervised hashing. CoRR, 2020.
Guhong Nie, Lirui Xiao, Menglong Zhu, Dongliang Chu, Yue Shen, Peng Li, Kang Yang. Li Du, and Bo Chen. Binary neural networks general propose compute paradigm for on-device computer vision, 2022.
general-propose
Max B Paulus, Chris Maddison, and Andreas Krause. Rao-Blackwellizing the straight-through Gumbel- softmax gradient estimator. ICLR, 2021.
Adeel Pervez, Taco Cohen, and Eistratios Gavves. Low bias low variance gradient estimates for boolean stochastic networks. ICML, pp. 7632-7640. 13-18 Jul 2020.
Jorn WT Peters and Max Welling. Probabilistic binary neural networks. arXiv preprint 1rX0v:1809.03338, 2018.
Tapani Raiko, Mathias Berglund, Guillaume Alain, and Laurent Dinh. Techniques for learning binary stochastic feedforward neural networks. In ICLR, 2015.
Wolfgang Roth, Günther Schindler, Holger Fröning, and Franz Perkopf. Training discrete-valued neural networks with sign activations using weight distributions. 2019.
Ruslan Salakhutdinov and Geoffrey Hinton. Semantic hashing. Int. J. Approx. Reasoning, (7):969-978, 2009.
Alexander Shekhovtsov. Bias-nariance tradeoffs in single-sample binary gradient estimators. In GCPR, 2021.
Alexander Shekhovtsov and Viktor Yanush. Reintroducing ttraigh--throuhh estimators as principled methods for stochastic binary networks. In GCPR, 2021.
Alexander Shekhovtsov, Viktor Yanush, and Boris Flach. Path sample-analytic gradient estimators for stochas- tic binary networks. In NeurlPS, 2020.
Dinghan Shen, Qinliang Su, Paidamoyo Chapfuwa, Wenlin Wang. Guoyin Wang, Rícardo Henao, and Lawrence Carin. NASH: Toward end-to-end neural architecture for generative semantic hashing. In An- nual Meeting of the Association for Computational Linguistics, 2018.
George Tucker, Andriy Mnih, Chris J Maddison, John Lawson, and Jascha Sohl-Dickstein. REBAR: Low- variance, unbiased gradient estimates for discrete latent variable models. In NewrlPS, 2017.
Ronald Williams. Simple statistical cradient-following algorithms for connectionist reinfrmmmen learning. Machine Learning, (3):229-256, May 1992.
Inc. Wolfram Research. Mathematica, Version 13.1, 2021. hammatii. Champaign. IL, 2022.
URL https://www.woll com/
Brandon Yang. Gabriel Bender. Quoc V Le, and Jiquan Ngiam. CondConv: Conditionally parameterized convolutions for efficient inference. In NeurlPS, 2019.
Mingzhang Yin. Yuguang Yue, and Mingyuan Zhou. ARSM: REINENRtEEEFOOCE--SSS- -warp-rr--ssa-mmrreee estimator for gradient backpropagation through categorical variables. ICML, pp. 7095-7104, 2019.
11

---

A PROOFS
A. BINARY CASE
Proposition 1.
C(w) J p(z) (z)) DZ)) (2r - 1)cr log(2)t) + O(f))
(10)
where pz is the logistic density: pz( = =(n)o(-1) and c₁ = 2p-
Proof. The conditional density p(z(z) is
if 1; p(=|2)) [z( 0j/p(x=0), if 0.
2
(15)
Let us denote p(x=1) as just p. The GR estimator expands
(1) S d((n+ 2)p(2d22 if 1;
GR
(0)
(16)
Jnd =((n+))p(2dd) if O.
1-p
Using the change of variables v = (r + 2), with the inverse = = (logit(v) we have
de oi(n + 2)d=
(17)
and can write the estimator as
JGR
(2(2) p(x)
x pz flogitz)dr + (1 - 2) for pz(n - flogitv)du)
(18)
Note that Pz flogitv) is bounded above by a constant sup, (2) = 1.4 A constant is integrable on [0, 1]. By dominated convergence theorem take the limits -> under the integral. In particular we can use
lim pz (n (v) p2()))
(19)
under the integral. In order to get more detailed view, we make the Taylor series expansion of (n flogit(u)) and substitute under the integral. With the help of Mathematica (Wolfram Re- search, 2021) we obtain:
JGR ('(x) p(a) pla)
+ (2x 1)e, log(2)t + C2 f O(t))
(20)
where ej = (1nh(1122) = 2p and 62 1 - 3/ (cosh(())) + 1)).
Corollary 1. In the limit + the GR estimator becomes the DARN ) estimator.
Proof. From the series expansion, the limit (- - is
JGR ('()) 7 2 p(z) P2(1) p(a) p(1 p),
(()
(21)
where = (n). It remains to show that matches DARN as defined in (4). Note that (ppp((=l))) (1) (n)) p(1 p) and log p(a-O dy
p(1 p). By expanding the
dry
cases for and we verify that
x)
d log p(z= 1mn)
1pp(1 (p).-) p),
(22)
dy
where 7 = 2.
Corollary 2. The mean and variance of the GR estimator (9) in the asymptote i - 0 are:
E[ = =p(1 =CC
((())=t O(t2),
(Ila)
V[JGR| (p(1 - ()a - b3)ett) + Ott2)
-)a
(11b)
where y p= (1), a = C'(1) = L'()) and غ = (2p 1) log(2).
I-p
12

---

Proof. The mean of the estimator is computed from the series expansion up to the first order
p - (z( { + aat) +
+(1 p) -. Pz (2( + -11t) +
px (2) +() +0(((())) + 0(12).
(23)
Since the GR estimator Jo (2) is Bernoulli variable with values J and J with proba- bilities p and p. respectively, we can compute is variance simply as
(J |GR (1) J (0)) p(1 p)-
(24)
Using that pz(1) = p(1 - p), the asymptotic expansion of variance up to first order in t is
(p(1 - p))² (a( aat) - -63 - -at)) +(P)
(25a)
= =p((1 - p)) ( (a-a00)) - b) (a (a) + 0(t))
(25b)
where a = (1) b
('(0) 1-p ē = 1og(()c1- The first order term is
[ - (1{11-))><a'' -a))
(26)
It could be positive negative depending the values of the derivatives and of Let us expand a,b and 01 = tanh(n/2) = 2p - 1. We obtain, up to positive constants,
- () - p)² (()))²²) 1)t.
(27)
We see that for the comer points, where p approaches either or 1, this linear term is negative. In particular for linear objective have (1( (0 and the linear term becomes
pp(1 (1) 1
(28)
which is non-positive for any and is zero for 1
A.2 GENERAL CATEGORICAL CASE
This case is significantly more difficult, as we are dealing with multivariate integration in K Gumbel variables. We will make use of the following statistical relationship.
Lemma A.1. Let GI.... Gk be independent standard Gumbel random variables. Then with components Z; = G, - GG for = 1... K - -1 has the multivariate logistic distribution (Malik & Abraham, 1973) with cdf
Fz(2) I
(29)
Proof. The edf and density of Gumbel(0, 1) distribution are given respectively by
Fe(x) = = PG(x) = = e (xte
(30)
The conditional distribution of Z; given Gx is F2,IGH (z)l)) e- (tsto) The conditional joint distribution Z given GK is respectively
K II
K Σ
K-1
F e
exp(
( exp(-e
Σ
-e (31)
1=1
where S =EK= #i
The edf of Z is obtained by computing the expectation of Fz|Gx in Gk:
F2(2) =
2
e - " (yte (yte = e -- ------ -
8
I+3 f (1 S)e -- "dy It3 I e-v- du 1+5
(32)
where v = y - lo + 1) and the last equality is by recognizing the Gumbel density under the integral.
n
13

---

Theorem 1 (ZGR). The Gumbel-Ran estimator for one-hot embedding in the limit of zero tem- perature is given by
(J.. Joz Jo )p(==;;))
if x; ifi = x.
JAGR
(12)
--2 -e.
out of the conditional expectation since it does not depend on G for fixed x. Furthermore, X distributed as p(æ) by the sampling procedure and therefore Po(X=r) p(x).. We can thus rewrite the conditional expectation (8) as
dc(d(a)) Proof. We can take do
de
Jø [x() =1] ) de (Fc(w)/Po(X==)) - Jo J
dalu) dFa(u),
(33)
arg max
where FG is the joint edf of G. The condition arg max, (0k + GR) can be expressed as
8 + Gj (02 + G+) 00 Vj # z.
(34)
Let us define 33 = 0, 0g and Z = G G. for j = 1 K. Note that 3, = Zz = 0 by this definition. Then the constraint can be written as
Z < -3
(35)
The integrand ē expresses in variables 8, Z as
softmax, (0 + G) softmax, (3 Z).
(36)
Let us denote Zrr (Z f x). The joint distribution of Z is the (K ---riiae multivariate logistic distribution (Malik & Abraham, 1973), as detailed Lemma A.1, with edf:
(z I-E
(37)
To simplify notation, let Ze have the discrete law with mass 1 at single point 2g 0 and extend Fz to the full joint Fz accordingly. We then can rewrite the integral as
-3x
softmax, (3 =))F2(2)) as 80 ( -5-B 03
(38)
The Jacobian 23 softmax: + 2) is K K matrix with indices (k, where the column j zero by definition. Let consider one component of the above integral for x:
Ik.j
a 08, softmax, + 22))F.
(39)
=5-8
We want to evaluating its limit for / + O. We cannot push the limit under the integral in this form, we need to transform it first. To shorten the notation, let us denote a, A We change the variable 2j by the mapping
a; T: 2, + 2, (2 + S2) 1+a,+S,"
(40)
where Sy = -Eizs aj. This mapping is monotone increasing and one-to one from (-oc, -B) to (0, 1), therefore the constraint 2j S B3 will trivialize.
Let A₃ = softmax, (3 + z)x- We can rewrite the integrand 3 Abdzz as follows:
dAs d2, ap,
dAi dzz dz,
(AA depends on By in the same way on 2y)
dA da ds dr dA dv,
(41)
dz, dv
We will thus need to evaluate Cj= = d4g
dAg day
dAA de, day
For j = k we simply have
du,
de,
day
A Ita, 88
a
245, Uji C 2455
(42)
14

---

For + / we have
dAc day
at daj 1ta++8; (Ita)+S5)*
(43a)
dep da,
(2 + S)) Itay+3, (Ita+t8,))
(245)
(43b)
(I+a+t8,)*
C (2+5))(1+.)))
(43c)
The integral 1x with the change of variable 2j 7 U expresses
S So (d (2n5)
(44)
#5-B: vidj
where ž-j = (zilk + 5) and
PZ, :)
(45)
The dependance of ft on t is through T, while C2 depends on t and 2 y Note that is a squashed density and is itself not a density.
Next we show dominated convergence of he z-z) Cyf((vylz-)) in / + O. If he(uz, z-z) converges point-wise and bounded above by an integrable function, then the limit - 0 can be taken under the integral.
We show a constant bound on he (-)) as follows. Note that |C₁| S 1. We then have
[,,( =-j)- S sup f(Dj|z=)) sup 2 I,, (2,,,,,)))))))))))
(46)
vE(0.1),a
which is the supremum of the conditional density of the standard multivariate distribution and is equal to some constant independent tf The integral of a constant function c over (0, 1) x K-1 with respect to the measure du (z-j) c.
The point-wise limit is as follows. For 2 satisfying the constraints 2; + 3 S 0 strictly for i + j, x. we have
lim a lim ((i++2)/ 0 and lim Sy 0. [-
(47)
100
t-0
Therefore we have
if if k A k # if k
lim Cj
(48)
The inverse of mapping T' is given by the relations
a 245 (1+)
2j -3, + l log(a;))
(49)
It is seen that the limit of log(a,) is finite and therefore
lim T (w;) = -3g.
(50)
+0
and
lim fe(vg(z-5) lim P2,12 t0
(T;)l)) = P2212 , (lim T (1=))=)))
2040
-P,,1Z,, (-3g (2z))-
(51)
By dominated convergence theorem, we can now claim
lim Iugj = 0 if + k and k + x.
(52)
And elsewise, if = k or kr == =x, I.
lim L8.j
カzソ 1 ( (2-1))) + 8 Fz F2((2)
(53)
-.S-Bi Vişj
2=-8
== 8
7 8
= 0j, softmax
(54)
Di, I+I epr
80, Σ.
==}p(#)p(()
(55)
15

---

where the upper sign corresponds to the case k and the lower to k = 2.
Let us denote lim 1. Multiplying it with the incoming derivative J on the left, we obtain:
(J = (J,, - JΦ, (o))p(())())
(56)
And finally, multiplying (56) with the Jacobian per (33), we obtain
da on the right per (38) and with the factor p(z) 2e
2 (Jg. Jo))p(t) 1-2 fe (Jo, Jon ((()
if x; if x.
JZGR
(57)
Proposition 2. ZGR estimator decomposes as
JZGR (JST JDARNIGON))
(13)
Le., with the choice ş = (()) in DARN.
Proof. Let p denote the vector of probabilities (p(a=k; n)) O,..., K - 1). Recall that have derived ZGR under the assumption of on-hot embedding d inherited from GS. In this case (+)(1) Jo, and 6e >; o(j) Pk-
Note that ZGR (57) defines the gradient in the parametrization used in Gumbel Rao and initially in Gumbel-SSfttax, while ST and DARN estimators are given by us with respect v, We need to bring these two to a common basis. We chose to reconstruct JZZR because both SST and pp particularly simple:
JDARN are
JST Job(t) Josn
(58)
jDARN Joló: - )>=iii//p(x>) (Jo. -ca-l-))l-))
(59)
Note, because p lies in the simplex, gradients in p are defined up additive constant all coordinates. In other words any such additive constant is irrelevant and will not affect the gradient in ?).
In order to reconstruct JEGR we represent JZCR = JZGR P, where P is the Jacobian of softmax, given by
P diag(p) - pp = diag(p) 1p)).
(60)
We first note that JEGR satisfies E; JRGR 0 (as any gradient should, but not necessarily a stochas- tic estimator) and therefore
J7C
J7KR (I 1pT) = JJCR diag(p)-12 PP
(61)
We obtained:
[J - Jo.) - 10 (Js, J())v/((( if
xi
JZGR
(62)
up constant, iee. adding the same number to all components. We further add the constant JJo. and obtain
ZGR
if x; Jor E Jon ((1)//(()
1
Jo:
(63)
if = x,
Subtracting 1 JST the reminder JRE with 2 p
JRE
= [i- píσs Σ Ejja( (J. Joo))(i).
(64)
pi
Simplifying
Ejte(Jo. Jo, p(i) Eg(Je, -om)(()) -Jo- - -E, Josp(i)
(65)
we obtain
JRE = |==]| p((]
1 (Ja--En - Joppii))
(66)
and we that JRB JDARN
JDARN with क p= T(1)-
16

---

Theorem 2. ZGR is unbiased for quadratic loss functions C with arbitrary coefficients.
Proof. Since ZGR estimator is linear in (estimate for a linear combination of two loss functions is the linear combination of estimates), it is sufficient to prove the claim for one-hot embedding o and some elementary functions forming a basis for all quadratic functions. With one-hot embedding we have (()); ())=;;)) P:-
Let us start with a linear monomial [(z) = x(x);- The expected loss is E[C(x)] = p(æ=i; n)) The true gradient is
Jn d p(==;;;
(67)
Substituting Jø, = [K=ii] in ST have
JS
d((d(=)) do(n) de
dé(n). dn
Jn-
(68)
This may come surprise for someone, but ST for single categorical variable is exact (zero bias and zero variance). The expectation of JARRN simplifies as follows for any & and a linear loss function, ensuring that Ja is constant in
E[JDARN] = E. p(x))o(((())
x) dp(z:n) = J6 (0(x)) - 3) dm(ann) p(z)
dn
Jo Ex 6(x)
dp(a:)
Joh d dn
dn Eg p(x; = Jo dō(n)
(69)
du
Substituting Jax = [k==]] and ((1) we obtain
B[J,DANN dp(z=in) Jny
(70)
dy
reconfirming that DARN is unbiased for linear function of categorical variables as expected. follows that {JSTT + J.AANN is also unbiased.
Let us now consider the clementary quadratic function L(ő(x)) (x)); c(x); For all discrete assignments it is zero, therefore the true gradient of its expected value is zero. We have
(2((x) Jor (x) = {
- k ==
(71)
k ti i.
Therefore
JST pa
for + and
E E[JST] = ==[2p(x); = 2p
(72)
JAARN we have For p.
DARNN (Jo, - J p(a)
(73)
(1 (2)(((( p(x)
(74)
Its expectation is
(1 - (1)( 1) 1 ps.
(75)
For J.h
JDARN- 0 for k # we have
JDARN-
-6 pla) [==
(76a)
Pk
--Jo(z)p [==k]
(76b)
= --P((20(z) 1) p(zz [==
(76c)
Its expectation is
- (kk( 1) = Pp:
(77)
By subtracting P: from all ordinates of JPARN, we obtain an equivalent (having identical derivative in 12) form where E[JDARN] for all k * i and R[JDARN] 1 2pi,, which cancels with E(JST)
17

---

Next consider bilinear monomial in o: L()(x)) o(z)1 where we have taken indices
and 2, without loss of generality. Its is zero for all discrete assignments and therefore the gradient of its expectation is zero. We have
Jphi, = =x((x) = 2x-2]
(78a)
Jphig = 101 (x) = ===
(78b)
For ST we have _TT = J and
E[ P2.
(79a)
E| |[
Pz = P1N
(79b)
| = 0, k # 1,2.
(79c)
For DARN part we have:
DARN [x= (Jo, Jo P1 JorP2);
(80a)
P1
p1
J JARN = [-=2] 1 (Joz JØn P1 JppP2);
(80b)
p2
JAARN = [r--k] p&u (-Jo, 1 - Jo2D2), + 1,2.
(80c)
In the expectation, substituting Jai
[[JDANN] = [ - ( (1) 62(1)R1 (1()p)2) = P2
(81a)
pi
E DARRN = [x=2] يير (2) -2(()p) 01 (2)p2) = -PI,
(81b)
E| [ARNN] = [-= (-d2(()) (1(k))2) = 0, k x 1,2.
(81c)
This exactly cancels with ST.
The elementary functions we have considered form a basis in the space of all quadratic functions. By linearity argument, J2GR (JST + JDARN) unbiased for all quadratic functions.
B DETAILS OF EXPERIMENTS
Here we give detailed specifications our experiments. The implementation of all experiments will be made publicly available upon publication. During the review period, we will be happy to answer questions and share the code with reviewers confidentially through the OpenReview platform.
B. DATASET
In quantized training MNIST' and FashionMNIST' datasets. Each contains 60000 training and 10000 test images. We used 54000 images for training and 6000 for validation.
In VAE training, following the prior work, we use decoder with Bernoulli output layer, which requires binary datasets. MNIST-B is binarized MNIST with a fixed threshold of 0.5, same as in Yin et al. (2019). The original Omniglot dataset of the size 105 105 and contains binary images. However the established benchmarks use its down-sampled version (to size 28 x 28). which then dynamically sampled: binary pixel values are generated with probabilities proportional to the original pixel values (Burda et al., 2016; Dong et al., 2021). which we denote as Omniglot- 28-D. The down-scaled dataset published by Burda et al. (2016)² was used, same as in the public implementation of Dong et al. (2021). It contains about 24000 training images, which were split into training (90%) and validation (10%) parts and currently we are not using the validation part.
htttp://yann.1 http:///yyam.eecn.coomo//eee/e 1ecun, www.://w....//c//////mm..ma
mat
18

---

def ZGR (ensor))))---ooooooooooo -XTensor::
"* ***Reuur categorica sample from P [x,C] (over axis=-1) one-hot vector, with ZGR gradient.
index = Categorical (probs=))
x = Fr.nneeeheeeee (index, num num_classes- shape . (t)(p) (p)
logpx = p. p.log() gathhgratterrrrre (-1, index. index. unsguezzeuuueeeeeee (-1)) log p(x)
dx_ST = P
dx_RE (x p.detach ()) * legpx
dx (dx ST dx RE) /
return (dx dx.detach ()) value x with backprop through
Figure B.1: ZGR implementation in Pytorch for general categorical variable.
B.2 METHODS
ZGR for categorical variables can be implemented as shown in Fig. B.1. It is a plug-in estimator, meaning that it is sufficient to use it for every tensor of categorical variables in a hierarchical model and the gradient in all parameters will computed by back-propagation automatically.
Gumbel-Softnn (GS) and Straight-through Gumbel-Softmax (GS-ST) (Jang et al., 2017) are shipped with pytorch'.
For Gumbel-Rao with MC samples (GR) we adopted the public reimplementation by nshepperd? 5,
RF(M) we implemented according to (Kool et al., 2019, Eq. 8). The part of the computation relevant to the encoder is propagated forward and backward only once. In the decoder we perform as many backward passes as forward, as this reduces variance of the gradient decoder parameters. In quantization our implementation performs backward pass for each forward pass, and is not well organized.
For ARSM (Yin et al., 2019) made own reimplementation, cross-checked with the authors tensor- flow implementation As with RF(M), we also performed backward pass for each forward pass.
B.3 VAE
Model In our model each categorical variable is encoded vector of +1, corresponding the bit representation of x, similar to Paulus et al. (2021). There is a fixed number of total hidden bits (192), which are split into several categorical variables. For example 192 lb variables or 32 6-bit variables. This way the number of weights in the network stays constant. The network architecture adopted from Yin et (2019):
Linear(784,512) > LReLU Linear(512,256) - LReLU Linear(256, D*K),
where in the last layer we have D of K-way categorical units and LReLU has leaky coefficient of 0.2 (same as in Dong et al. (2021), default in tensorflow). The output of the encoder defines logits of the encoder Bernoulli model ((===1/r)) where is the input binary image and is the latent discrete state.
The decoder has exactly the reverse Linear-LReLL architecture and outputs logits of conditionally independent Bernoulli generative model p(x,=112). We optimize the standard evidence lower bound (ELBO) Kingma& & Welling (2013) with prior distribution p(=) uniform and not learned.
We do not perform any special data-based initializations like subtracting data mean in the encoder
Dong et al. (2021).
Optimization In the forward pass all methods produce a sample, from which a stochastic esti- mate of the gradient with respect the decoder parameters is readily computed by backpropagation through decoder. We compute the KL term in ELBO analytically for a mini-batch and use its exact
Pytorch function torch. wwnetiionaaii httpr//ooooooommeeeooooootteeee
http:///hh.hhhccoo/////peee////r pytorch : http://hhhhh....ma//aagaaaaaaaaa
19

---

gradient. The estimation problem (1) occurs for the gradient of the data term with respect to the encoder parameters, where the estimators through discrete variables are applied.
All methods, including GS that optimizes ELBO with relaxed samples, are evaluated by the correct ELBO with discrete samples.
In the VAE experiments we measure the gradient accuracy and the training performance and do not make use of validation or test sets. First, this is reasonable when comparing quality of gradient estimators, regardless generalization. Second, the prior work Dong et al. (2021) has verified that improvement in the training ELBO translates into improvement of the test ELBO and IWAE bounds.
Following Dong al. (2021) we train with Adam with learning rate 10-1 using batch size 50. Furthermore we tried to match the training time that of Dong et al. (2021). For MNIST we perform 500 epochs, and for Omniglot-28-D we perform 1000 epochs, roughly equivalent in booth cases to their 500K iterations with batch size 50.
B.4 ADDITIONAL VAE EXPERIMENTS
We include some extended results.
Bias-variance analysis We conducted the bias-variance analysis for VAE at different training stages. Namely, we trained the model using RF(4) for 1, 10, 100, and 200 epochs and at each stage evaluated bias and variance of all gradient estimators. The model used 16 categorical variables of 16 categories (64 total latent bits). The results are displayed in Fig. B.2. At the very beginning of training, the picture looks substantially different in that there is some bias reversal in GS-ST and derived estimators. However from epoch 10 and on the trends and relative ordering of methods stabilizes, with only RF(4) slightly overtaking ZGR in variance. The different picture after 1 epoch suggests that it would be beneficial in practice to warm-up the training with a few epochs of GS-ST 1 or ST. We left such tuning to future work.
B.5 QUANTIZED NEURAL NETWORKS
Experimental Setup In this experiment we train convolutional network 32C5-MP2-64C5-MMP 512FC-I0FC, closely replicating the model evaluated by Louizos et al. (2019) for MNIST. Each activation quantization is preceded by batch normalization (Ioffe & Szegedy, 2015).
All gradient estimators are working with the same network, parametrization and initialization. In the case of logistic noise the noise standard deviation is learnable and is initialized to 1/3, All methods are applied with Adam optimizer for 200 epochs. For every method we select the bast validation performance with the grid search for the learning rate from {10-3,35 3.3 3,33 10 10-4) We used the step-wise learning rate schedule decreasing the learning rate 10 times at epochs 100 and 150. The whole procedure repeated for different initialization seeds and we report the mean test error over seeds and + =(max- - min)/² 2 over seeds.
For validation and testing, we evaluate the network in the 'deterministic' mode, turning off all in- jected noises. This corresponds to a simple deterministic quantized model to be deployed.
Single Unit Quantization We include the following toy experiment that well illustrates properties of different estimators. We evaluate bias and variance of all estimators on a simple function of a single quantized variable. Let 7) real-valued parameter. Let p(æ: 17) be given by the stochastic quantization model with K = 4 states and a particular noise type. Given a test function L(x) we can compute the true gradient of [x)( For each estimator we draw 10ah samples to compute its mean and standard deviation for each value of n) The results are presented in Fig. B.3. In this plot we show several combinations of loss functions and noises. The test functions are: linear C(x) = 2; quadratic L(x) 1 (x c)² and sigmoid L(x) (2 c)), where
(K 1)/2 is chosen for centering. The noises shown refer to the logistic noise with std = 1/3 as used at initialization by Louizos et al. 2019 and the triangular noise with the density p(z) max(0, 1 |=]).
The bias of the GS family quickly decreases with the temperature. ZRG estimator achieves the same expected value as GS-ST in the limit of small temperature illustrated by GS-ST(=0.1) and the variance comparable to that of GR(1= M=100). We also verify that ZGR has zero bias for quadratic objectives as have shown theoretically.
20

---

After epoch of training
DA
22-4
GRAM2208
20R
12.4
33"0
After epoch of training
20
05
23
5345
MC.10
22-
A154
25
27-25
23
04 MC310
30
22
After 100 epochs of training
05
10
30
R0M-2
23*
25
G4 MC300
10
After 200 epochs of training
05
5TC5
23
18
035 23300300
10
23.4
Figure B.2: Gradient estimation accuracy VAE MNIST-B. Average (per parameter) squared bias (left) and variance (right) of gradient estimators versus temperature for model snapshot at particular iteration of training with RF(4). VAE network with 16 categorical variables with 16 categories.
Additional Comparisons In Table 3 left out GS method after preliminary testing. These results (without confidence bounds) are available in Table B.I. we also performed a full comparison on FMNIST dataset in Table B.2, which qualitatively agrees with the results for MNIST Table 3 and confirms that ZGR improves over GR-MC((=011, M=10).
21

---

GS-ST(t=1)
GS-ST(t=0.1)
GR((=0.1.M100)
ZGR
RF(M=2)
ST
Linear / Loogsstic
m
Quadratic / Logisttie
-1
Signoid Logistic
22
Quadratic Trinugllar
-2
Figure B.3: Performance of estimators in stochastic quantization with selected combinations of test functions and injected noise for varied input of the stochastic quantizer. The dashed blue line the exact gradient. Red line the mean of the estimator and the red shaded area shows =1 std. GR with 100 is able to reduce the variance of GS-ST substantially. ZGR reduces the variance by edge further while keeping the bias equal to the theoretical bias of GR with zero temperature. The variance ARSM substantially higher. Plain ST estimator has a yet smaller variance but a larger bias, which may accumulate during taring.

---

Table B.1: MNIST test error[%] complementing Table 3 with GS variants, which are found to be inferior.
Logistic Noise Relaxation
Weights [bits] / Activations [bits]
method
2/2
T/T
T/1
1/1
GS(t=2)
0.80
0.78
0.87
1.11
GS((=))
0.86
0.72
0.98
0.86
GS((00.5)
0.94
0.82
0.99
1.18
Table B.2: FMNIST error [%] in deterministic mode (no injected noises test time) for dif- ferent bit-width weigh and activation (T' denotes ternary). Hyperparameters are selected on the validation set. Reference test errors: ReLU 9.73%. Clamp 9.6%.
Weights [bits] Activations [bits]
Method
2/2
T/T
T/1
1/1
GS-ST((=2)
8.15 + 0.30
8.32 = 0.13
8.58 + 0.18
8.92 = 0.08
Logistic Noise
GS-ST(t=1)
8.56 0.23
8.59 = 0.30
8.99 + 0.22
9.07 = 0.05
GS-ST((=0.5)
9.11 + 0.32
9.60 = 0.19
9.85
10.99 0.266
GR-MC((0.5.M=10)
9.04 0.23
8.84 + 0.10
0.11
9.97 - 0.08
GR-MC((=011,--I))
9.35
10.10 + 0.19
0.16
11.61 + 0.35
ZGR
0.33
9.33 + 0.22
9.25
10.67 0.13
ST
0.11
8.79 0.10
9.09
9.28 0.17
Triangular Noise
GS-ST(1=2)
8.57
8.76 + 0.14
10.19
10.92 + 0.30
GS-ST(==1)
0.00
8.94 0.21
9.95
10.56 + 0.26
GS-ST((=0.5)
9.35 0.299
9.63 0.20
0.13
11.34 +
GR-MC(=0.5.M-10)
8.11
8.62 0.07
9.14
9.76 - 0.21
GR-MC(=O.1.M-10)
8.85
9.32 0.04
1..16
10.83 1 0.19
ZGR
8.67
8.96
9.73
10.09 0.38
ST
0.29
8.61
0.15
9.98 0.19
23


Reviews:
[{"review_id": "pMMjojKThX", "reviewer": "Reviewer_wSJz", "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": "The paper presents a variance reduction technique that improves the Straight-Through version of the Gumbel-Softmax estimator by taking the zero-temperature limit. The paper further shows that the proposed estimator could be decomposed into a sum of straight-through estimator and DARN. The superior performance of the proposed estimator using 1 or 2 sample(s) is demonstrated through experiments on MNIST, by comparing against an unbiased categorical estimator ARSM, and Gumbel-Rao with large number of Monte-Carlo samples. \n\n", "summary_of_the_review": "The paper did thorough theoretical analysis on the proposed gradient estimator, with good explanation. The author(s) did good evaluations in experiment section. The reviewer thinks that the experiment section could be improved with better baselines and more analysis on different benchmarks and ablations as mentioned in \"Weaknesses\" section.  \nI would raise the score, if the author(s) could help address some of the concerns in Weaknesses section.", "strength_and_weaknesses": "Strength:\nThe paper idea is simple but sharp: inspired by the analysis of Gumbel-Rao (GR) estimator (Paulus et al., 2021) that the estimator performs better at lower temperature, they push the estimator to the zero temperature limit, and found superior performance. Another major contribution is that the paper identifies the connection between GR with Straight-Through (ST) and DARN at the zero temperature limits, for both binary and categorical cases. The insights and analysis are incremental but very important, in terms of both theoretical understanding and improving applicability of the estimator. \n\nWeakness:\nThere are some minor problems / questions:\n\n1. Questions:\n  - In Figure 1, the author compares the computational time across different estimators, is this comparison based on the numbers from different implementations, or are all the estimators implemented in the same code base? Instead of shown the comparison of time, it would be great to include a discussion about complexity analysis on each algorithms, which is not subject to the difference in the implementation and helps the readers / reviewers to better understand how the estimator scales with increasing number of categories. E.g. ARSM requires $\\mathcal{O}(C^2)$ evaluations, where $C$ is number of categories.  \n  - The author(s) discussed the multi-sample version, ZGR s = 2, in experiment section. How does it formulated? Taking average of two estimations from each samples?\n  - Is there any analysis for scaling the number of categories? Would the relative performances of different estimators hold for larger / smaller number of categories?\n  - MNIST is a standard yet simple task. It would be good to extend the analysis on other benchmarks as well, e.g. FashionMNIST and Omniglot.\n  - The author(s) claimed that the bias is very small comparing to the variance, based on the bias and variance estimations at \"an evaluation point\" \"after 100 epochs of training with ARSM\". It would be nice if this analysis could be done at an early stage of the training. \n  - A side question: instead of setting the temperature to zero or a finite number, how performant is the model if one makes the temperature learnable?\n\n2. Missing baselines:\n  - Binary case: as the paper considers the binary case in 3.1. In this case, it's worth to compare against the latest unbiased gradient estimators for binary variables, ARM (Yin et al. 2019) and DisARM (Dong et al. 2020). Both are using antithetical sampling to improve REINFORCE, and DisARM is Rao-Blackwellized ARM by analytically integrating out the randomness. \n\n  - Categorical case: ARSM is a categorical version of ARM. However, previous research (Dong et al. 2021) found ARSM underperform simpler REINFORCE Leave-one-out (RLOO, Kool et al. 2019) baselines by significant amount with comparable number of samples. Also, Dong et al. 2021 proposed Rao-Blackwellized ARSM with binary reparameterization, further improves RLOO. \n\n3. Missing References and things to be fixed:\n  - line below Eq (1) \"It has proven efficient in practice to ...\" Please add reference for this. Also the sentence should be \"It has been proven ...\"\n \n  - Eq (2), in general, especially for VAE, as the paper considered in experiments, \"x\" the latent variable, therefore it is usually sampled from posterior distribution. In this case, $\\mathcal{L}(x)$ will also depends on $\\eta$. Correct derivation should include two terms, e.g. Dong el at. 2021 Eq (2) and relevant discussions, but the term with gradients of $\\mathcal{L}(x)$ can typically be estimated with a single Monte Carlo sample and neglected from discussion. \n\n  - Broken sentence: line below Theorem 2: \"By quadratic functions we understand quadratic functions .. with arbitraty coefficients\"\n\n  - paragraph above 4.2, \"GR-MC with K = 10 or K=100 MC samples\", it's better to use \"M\" instead of \"K\", as the paper used \"K\" for number of categories. \n\n  - Please try to make the main text self-contained. E.g. Eq (10) mentioned $p_Z(\\eta)$, and it is the logistic density. Writing an expression for it, would help the reader to easily reproduce the results.\n\n\nReference:\n  - Yin, M. and Zhou, M. (2019). ARM: Augment-REINFORCE-merge gradient for stochastic binary networks. ICLR 2029\n  - Dong, Z., Mnih, A., and Tucker, G. (2020). DisARM: An antithetic gradient estimator for binary latent variables. NeurIPS 2020\n  - Kool, W., van Hoof, H., and Welling, M. (2019). Buy 4 reinforce samples, get a baseline for free! Deep RL Meets Structured Prediction ICLR Workshop.\n  - Dong, Z., Mnih, A., and Tucker, G. (2021). Coupled gradient estimators for discrete latent variables. NeurIPS 2021\n\n\n ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "iVdf9USDmx2", "reviewer": "Reviewer_Waws", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "The paper proposes a gradient estimator for distributions over discrete variables, following in the tradition of the Gumbel-softmax estimator. In this paper, the proposed estimator is the zero-temperature limiting value of the Gumbel-Rao estimator. Interestingly, this is shown to be the same as an average of the DARN and classic Straight-Through (ST) estimator. The authors show that their proposed estimator has zero bias on quadratic objectives, and compare its performance on several downstream tasks.", "summary_of_the_review": "This paper is a nice piece of work, containing a derivation of the zero-temperature limit of the Gumbel-Rao estimator. However, the experimental comparisons don't currently justify the novelty of the work for me. As it stands, the experiments demonstrate that some gradient estimators are better than others at different tasks, and different combinations of bias and variance are useful for different tasks. This paper provides a new option on the bias-variance curve, but it's not obvious to me that this actually translates to real gains on any downstream task. \n\nUpdate after the rebuttal: \nI'm impressed by the authors' additional experiments and will raise my recommendation to 6. ", "strength_and_weaknesses": "Strengths\n+ The theoretical results are nice. It is very interesting to see that the zero-temperature limit of the GR estimator is given by such a simple combination of the DARN and ST estimator.\n+ The experimental results are fairly comprehensive, and presented quite fairly, as far as it is possible to judge.\n+ The bias and variance properties of the new estimator are clearly described, and illustrated in the experiments. \n\nWeaknesses\n+ The experimental results ultimately do not clearly show superiority of the proposed ZGR estimator. For instance, across table 1 and Figure 2, the Gumbel-Rao estimator performs similarly or better than the ZGR. Similarly in figure 2, the GS-ST(t=2) strictly dominates the ZGR. The experiments do show that the ZGR is a decent enough estimator in practice, but there is still enough variation between all the different estimators that the best bet for the practitioner is to just do a hyperparameter sweep including all the available estimators.\n+ Similarly to the previous point, I think there needs to be more discussion about the connection between the bias and variance of the estimator itself, and the error in the optimization process, particularly when combined in with the particular problem of training a neural network with an adaptive optimizer. As the authors point out, simply considering the MSE in the gradient is not necessarily that relevant for the actual question of how well an NN will convergence. For instance, an interesting experiment could be to train the VAE with all the different estimators, using both an adaptive optimizer and regular SGD, to see if the bias might be worse than variance as the authors hypothesise.\n+ Minor points: \n  + In page 3, shouldn't the CDF of the Gumbel be $e^{-e^{-u}}$?\n  + 'unique amongst single-sample categorical estimators' isn't REINFORCE unbiased? Or you need to define what you mean by single-sample estimator\n  + Can you not just get in touch with people to get the MNIST dataset you need?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "h-d4M-UaCq", "reviewer": "Reviewer_PG9N", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 2, "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": "The paper introduces a new estimator called ZGR aimed for the stable and accurate use of discrete random variables as part of approximate inference pipelines.", "summary_of_the_review": "This is a promising paper, but gives the impression of unfinished work. The proposed method is a very close variant of what already exists. The results are also not strong, although the experiments are comprehensive.", "strength_and_weaknesses": "Strengths:\n  i) The studied problem is important for the probabilistic machine learning community\n  ii) The paper reports comprehensive experiments that span many critical cases where such an estimator could be useful.\n\nWeaknesses:\n i) The technical novelty is limited.\n ii) Experiment results are weak. ZGR does not appear to show an improvement over for instance ST. I understand that it outperforms the GR family, but what that family can do ST cannot is not clear from the experiments.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "Y6gn6yFLsQ", "reviewer": "Reviewer_GHDv", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "This paper proposes Zeroed Gumbel-Rao (ZGR) stochastic gradient estimators for the discrete random variables. The authors prove that the limiting the temperature to zero in the Gumbel-Rao results in closed-from solution. They also show that the proposed ZGR estimator is in the middle of two well-known stochastic gradient estimators, namely Straight-Through estimator and DARN estimator. They argue that the proposed ZGR has merits both the low bias & variance perspective (when the temperature is less than 0.1), and the low computational complexity, based on their experimental results.\n", "summary_of_the_review": "I am very neutral on this paper. The work is theoretic and seems to be correct. However, the concern on the fair comparison in the experiment section needs to be resolved, since the GS variants are sensitive to the temperature tuning. ", "strength_and_weaknesses": "The paper is theoretically grounded and the proposed ZGR surprisingly filling the gap between ST and DARN.\n\nError bars are missing in some tables and figures. Also, the authors did not explore various temperatures (and temperature scheduling) for the baselines, which makes unfair comparison. Finally, the experimental result in Table 1 is not very promising.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject