Bridging Discrete and Backpropagation: Straight-Through and Beyond
Authors: Liyuan Liu
Chengyu Dong
Xiaodong Liu
Bin Yu
Jianfeng Gao
Venue: NeurIPS 2023 oral (2023)
Keywords: discrete random variables
back-propagation
straight through
TL;DR: We show Straight-Through works as a first-order approximation of the gradient and propose ReinMax, which achieves second-order accuracy with negligible computation overheads.

Abstract: Backpropagation, the cornerstone of deep learning, is limited to computing gradients for continuous variables. This limitation poses challenges for problems involving discrete latent variables. To address this issue, we propose a novel approach to approximate the gradient of parameters involved in generating discrete latent variables. First, we examine the widely used Straight-Through (ST) heuristic and demonstrate that it works as a first-order approximation of the gradient. Guided by our findings, we propose ReinMax, which achieves second-order accuracy by integrating Heun’s method, a second-order numerical method for solving ODEs. ReinMax does not require Hessian or other second-order derivatives, thus having negligible computation overheads. Extensive experimental results on various tasks demonstrate the superiority of ReinMax over the state of the art.

Bridging Discrete and Backpropagation: Straight- Thrhuuhh and Beyond
Liyuan Liu Chengyu Dong
Xiaodong Liu Bin Yu Jianfeng Gao Microsoft Research
[lucliu, v-chedong, ziaodl, v-ybi, jfgaao)mmcroooot.coo
Abstract
Backpropagation, the comeestonee of deep learning, is limited to computing gra- dients for continuous variables. This limitation poses challenges for problems involving discrete latent variables. To address this issue, we propose a novel ap- proach to approximate the gradient of parameters involved in generating discrete latent variables. First, we examine the widely used Straight-Through (ST) heuris- tic and demonstrate that it works Guided by our findings, we propose ReinMax, which achieves second-order accu- racy by integrating Heun's 'ss method, second-order numerical method for solving ODEs. ReinMax does not require Hessian or other second-order derivatives, thus having negligible computation overheads. Extensive experimental results on vari- ous tasks demonstrate the superiority of ReinMax over the state of the art.
first-order approximation of the gradient.
Introduction
There has been persistent pursuit build neural network models with discrete or sparse vari- ables (Neal, (192). However, backpropagation (Kumelhari et all, (1966), the cornerstone of deep learning, is restricted computing gradients for continuous variables. Correspondingly, many at- tempts have been made to approximate the gradient of parameters that are used to generate discrete variables, and most of them are based on the SraiightTThrouge (ST) technique (Bengıo et al, 2013).
The development of ST is based on the simple intuition that non-differentiable functions (e.g., sampling of discrete latent variables) propagation (Rosenblatii. 1957% Bengio el all 2013). Due to the lack of theoretical underpinnings, there is neither guarantee that ST can be viewed as an approximation of the gradient, nor guidance on dyper-paramedd configurations or future algorithm development. Thus, researchers have to de- velop different ST variants for different applications in trial-and-error manner, which laborious and time-consun consuming (van den Oord et all 2017 Liu et all 201% Fedus et all (22)). To address these limitations, we aim explore how ST approximates the gradient and how it can be improved.
be approximated with the identity function in the back-
p
p
0.10
RenMax
STGS
0.25
0.22
0.22
0150
10
0.21
0.355
12
25
25
30
15
40
5
10
15
20
25
63
0.030
32
Epoch
Epoch
Epoch
Figure 1: Training curves of polynomial programming, i.e., ming Ex
c c
1128,
where E
R 128x2 X E {0, 1]228, and X; Multinomialsofttmax
iid
Details are elaborated in Section E
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Algorithm ST.
Algorithm 2: ReinMax.
Input: softmax input, т: temperature.
Input: 0: softmax input, 7: temperature.
Output: D: one-hot samples.
Output: D: one-hot samples.
To f softmax(0)
To softmax(0)
D < sample_one_ sacee__ne_hotmmmeg
2 D sample_one one_hot(TT)
3 T f softmax, (0)
π1 f- D+softmax.(0)
/* stop stop_gradient// duplicates
2
4 <- caannaxssoo_gaadennnniii
its input and detaches it
from backpropagation.
n t- 2 2/2
D f #1 stop stop_ (rp_gardddent.)) + D
6 D π22 stop_ (rop_gradien((T)) + D
return D
7 return
First, we adopt a novel perspective examine ST and show that works special case of the forward Euler method, approximating the gradient with first-order accuracy. Besides confirming that ST is indeed an approximation of the gradient, our finding provides guidance on how to optimize syper-parameters of ST and its variants, i.e., ST prefers to set the temperature т 2 1, and Straight- Through Humbel-Softman (STGS; Jang et al, 2017)) prefers to set the temperature T S
Our analyses not only shed insights on the underlying mechanism of ST but also lead us to develop novel gradient estimation method called ReinMax. ReinMax integrates Heun's Method and achieves second-order accuracy, i.e., its approximation matches the Taylor expansion of the gradient to the second order, without requiring the Hessian matrix or other second-order derivatives.
We conduct extensive experiments on polynomial programming Tucker et all (2017): Grathwohl all (2018); Pervez all (2020); Paulus all (2021)), unsupervised generative modeling (Kingma & Welling, 20E3), structured output prediction (Nangia & Bowman, 2018), and differentiable neural architecture search (Dong et all (020)) demonstrate that ReinMax brings consistent improve- ments over the state of the art".
Our contributions are two-fold:
We formally establish that ST works as first-order approximation to the gradient in the general multinomial case, which provides valuable guidance for future research and applications.
We propose novel and sound gradient estimation method ReinMax that achieves second-order accuracy without requiring the Hessian matrix or other second-order derivatives. ReinMax is shown outperform the previous state-of-the-art methods in extensive experiments.
2 Related Work and Preliminary
Discrete Latent Variables and Gradient Computation. The idea of incorporating discrete latent variables and neural networks dates back sigmoid belief network and Helmholtz ma- chines (Willram 1992: Dayan et all 1995). To keep things straightforward, we will focus on simplified scenario. We refer to the tempered softmax as softmax, (0), = Σ exp(U,/7) where
exp(O./T)
is the number of possible outcomes, 0 Rax1 is the parameter, and is the temperature For i E [1.... n2). we mark its one-hot representation as I; E R=x1, whose element equals it is the i-th element or equals 0 otherwise. Let D be discrete random variable and D ( I, }. we assume the distribution of D is parameterized as: p(D = 1,) π. softmax(O)s. and mark softmax, (0) as π (r), Given a differentiable function ʃ R - R, we aim to minimize (note that temperature scaling not used in the generation of D):
min L(O), where L(0) = softmaxx (/(( (D))
Here, we mark the gradient of 0 as V:
AL(0)
Σ f(I) dm,
V
(2)
00
do
'Implementations are available at stttps:/ggittn htipps////giiiooo//totoottoooaae http:////mmmmicooo////eieillaa Without specification, the temperature (i.e., 7) is set 1.

---

In many applications, it is usually too costly to compute V, since it requires the computation of
{f(I)), [I,)] and evaluating fI;) is costly for typical deep learning applications. Corre- spondingly, many efforts have been made to estimate V efficiently.
The VREIFORREE (Williams, (192)) is unbiased (i.e., FIV REINNORE distribution of the discrete variable to be differentiable (i.e., no backpropagation through f):
V) and only requires the
d log d0
V RRENFORCE := (DD)
Despite the REINFORCE estimator being unbiased, it tends to have prohibitively high variance, especially for networks that have other sources of randomness (i.e., dropout or other independent random variables). Recently, attempts have been made to reduce the variance of REINFORCE (Gu et ala 2016; Tucker et al, 2017; Grathwohl et all, 2018; Shi all, R022). Still, it has been found that the REINFORCE-styL estimators fail to work well in many real-world applications. Empirical comparisons between ReinMax and REINFORCE- methods are elaborated in Section 6
In practice, popular family of estimators is Straight- Through (ST) estimators. They compute the backpropagation "through" a surrogate that treats the non-differentiable funetion (e.g., the sampling of D) identity function. The idea of ST origi- nates from the perceptron algorithm (Rosenblatii (157t Mullin & Rosenblatt, (162),, which leverages modified chain rule and utilizes the identity function as the proxy of the original derivative of binary output function. Bengio et al. (2013) improves this method by using non-linear functions like sigmoid softmax, and Jang et al, (2(17) further incorporates the Gumbel reparameterization. Gumbel-Softmax (STGS).
Efficient Gradient Approximation.
Here, we briefly describe Straight-Through Through (ST) and Straight-
In the general multinomial distribution case, in Algorithm M. the ST estimator treats the sampling process of D as an identity function during the backpropagation
af(D) d V = 3D d0
In practice, ST is usually implemented with the tempered softmax, under the hope that the temper- ature hyper-parameterr may be able to reduce the bias introduced by V (Chung all 2017).
The STGS estimator built upon the Gumbel re- parameterization trick (Maddison et 2014 Jang all 2017). It is observed that the sampling of D can be reparameterized using Gumbel random variables at the zero-temperature limit of the tempered softmax (Gumbbel, [954)::
lim softmax, (0 + G) where G, are i.i.d. and G; ~ Gumbel(0, 1). H-0
STGS treats the zero-temperature limit as identity function during the backpropagation:
STGS == 0f(D) dsofttmx, + G) 3D
(5)
d0
Both ST and v STGS are clearly biased. However, since the mechanism of ST is unclear, it remains unanswered what the form of their biases are, how to configure their hyper-parameters for optimal performance, or even whether S or E sTes] can be viewed as an approximation of V. Thus, we aim to answer the following questions: How V sT approximates V and how it can be improved?
3 Discrete Variable Gradient Approximation: a Numerical ODE Perspective
In numerical analysis, extensive studies have been conducted to develop numerical methods for solving ordinary differential equations. this study, we leverage these methods to approximate V with the gradient of f. To begin, we demonstrate that ST works first-order approximation of V. Then, we propose ReinMax, which integrates Heun's 's method for a better gradient approximation and achieves second-order accuracy.
'We use the notation V indicate gradient approximations. Note that the generation of D not differen- tiable, and V does not have the term 0 /dx.
3

---

3.1 Straight- Throigh as First-order Approximation
We start by defining first-order approximation of V as V Ist-order
af())) Definition 3.1. One first-rder approximation of V Ist-order == Σ;Σ, ; 8I, (1-1) do.
V
To understand why V Ist-order is first-order approximation. we rewrite V in Equation Das
d do + =E dT, == -
do EET,(u - f(L,)
dT.
(6)
d0
Comparing V 1st-order and Equation E. it is easy to notice that V Ist-order approximates (I,) f(I)) as 8f(1, method, which has first-order accuracy (we provide brief introduction the forward Fuler method in Appendix E). Correspondingly, we know that V Ist-order is a first-order approximation of V.
01; (I, I In numerical analyses, this approximation is known the forward Euler
Now, we proceed to show VS works as first-order approximation. Note that our analyses only apply to VsT as defined in Equation !! and may not apply to its other variants. Theorem 3.1.
E V sr
1st-order-
The proof of Theorem [3 is provided in Appendix (A.
is worth mentioning that Tokun & Satd (2017) discussed this connection for the special case of D being a Bernoulli variable. However, their study is built upon Bernoulli variable property (i.e., V (f(12) - f(()) dr (f(I)) - f(2)) dwa
). making their analyses not applicable to de
multinomial variables. Alternatively, the analyses in Gregor et al. (2014) and Pervez eall (2021)
are applicable to multinomial variables but resort to modify VST as n.n ST. order to position as a first-order approximation. We suggest that this modification would lead to unwanted instability and provide more discussions in Section I and Section B.G. Here, our study is the first formally established ST works as first-order approximation in the general multinomial case.
Besides revealing the mechanism of the Straight-Thre Through estimator, our finding also shows that the bias of ST comes from using the first-order approximation (i.e., the forward Euler method). Accordingly, we propose integrate better approximation for f((+)- [(1))
3.2 Towards Second-order Accuracy: ReinMax
The literature on numerical methods for differential equations shows that possible achieve higher-order accuracy without computing higher-order derivatives. Correspondingly, we propose to integrate a second-order approximation to reduce the bias of the gradient estimator.
Definition 3.2. One second-order approximation of V is
T ΣΣΣ AS(1)
AJ(I) 2 0I, aI, )I: - I,) d n
2nd-order =
+
d0
Comparing v 2ndd-order and Equation E. we can observe that, 22dd-order approximates f(T,) - ((I,) as ( ( 01; + 01, [I; - -,). This approximation is known as the Heun's 's Method and has second- order accuracy (we provide a brief introduction to Heun's method in Appendix E)). Correspondingly, we know that 2ndd-ordr is second-order approximation of V.
Based on this approximation, we propose the ReinMax operator as (ND refers to #+D the identity matrix, and refers to the element-wise product):
refers to
2
0f(D)
ReinMMxx 2- V m+n
VST., where A ー==
((TD 1') 0 TD =D (7)
0D
PPlease note that E, E((()) dn d
EJ(())
E d0
ET ا =0

---

Then, we show that V ReinMax approximates V to the second order. Or, formally we have: Theorem 3.2.
F
ReinMax
2nd-order-
The proof of Theorem S is provided in Appendix Computation Efficiency of ReinMax. Instead requiring Hessian or other second-order derivatives, ReinMaa achieves second-order accuracy with two first-order derivatives (i.e., at,
af(4)
and @r(L)) ). As observed in empirical efficiency comparisons in Section the computation O1.
overhead of V ReinMaa is negligible. At the same time, similar to SST (as in Algorithm MI, pro- posed algorithm can be easily integrated with existing automatic differentiation toolkits like Py PyTocch (a simple implementation of ReinMax is provided in Algorithm E), making it easy be integrated with existing algorithms.
Although it's possible to apply higher-order ODE solvers, they require more gradient evaluations, leading to undesirable computational over- head. To illustrate this point: The approximation used by ReinMax requires gradient evaluations, @f(L.)) i... { S1, j. In contrast, the approximation derived by RK4 needs n²² + gradient evaluations, i.e., (8f())) ar, } and (
Applicability of Higher-order ODE solvers.
I,+, Therefore, while higher-order solvers are applicable, they may not suitable in our case.
0T,y }. where Iy 2
ReinMax and Baseline Subtraction
Equation f plays a crucial role in positioning ST as first-order approximation of the gradient and deriving our proposed method, ReinMax. This equation is commonly referred to as baseline subtraction, a common technique for reducing the variance of REINFORCE.
In this section, first diseuss the reason for choosing ( the baseline, and then reveal that the derivation of ReinMax is independent to baseline subtraction.
4.1 Benefits of Choosing (Df as the Baseline
The choice of baseline in reinforcement learning has been the subject of numerous discus- sions (Weaver & Tad. 2001: Rennie all 2016: Shi all 2022). Similarly, our study, different baselines lead to different gradient approximations.
Here, discuss the rationale for choosing (D) the baseline. Considering Σ, (I, as the general form of the baseline (Φ; is a distribution over (I,,.. In i.e., Σ, Φ; = 1), we have:
Remark 4.1. When E, (II,) is used as the baseline and f(ī,) f(I,) is approximated as OS(1,) 01, (I, I3), we mark the resulting first-order approximation of as V Ist-order Then, we have E 0D sr) = aiv-order-vg-baselli
The derivations of Remark I. provided in Appendix IC Intuitively, since æD the output of the softmax function, could have very small values, which makes éD be unreasonably large and leads undesired instability. Therefore, we suggest that EDf better choice of baseline when it comes to gradient approximation, since its corresponding gradient approximation is free of the instability &D brought. TD
TO
It is worth mentioning that, when setting Φ 1. the result of Remark I echoes some existing studies. Specifically, both Grreor et al] (2014) and Pervez et all (2020) propose to approximate V
x.
as El 1 Ssr,, which matches the result of Remark [ by setting Ø
In Section K. we compared the corresponding second-order approximation when treating ET(()) 1 " f(I) as the baseline, respectively. We observed that gradient estimators that use DD) as the baseline consistently outperform gradient estimators that use 1 EE, ((,,) as the baseline, which verifies our intuition and demonstrates the importance of the baseline selection.
and
5

---

Method = STGS
Nethod GST:1.0
Method GR MCK
Nethod 57
Nethod ReinMox
203 206
252
03
120
210
160 302
100
Figure 2: Training - ELBO on MNIST-VAE (lighter color indicates better performance). STGS, GST-1.0, and GR-MCK prefer to set the temperature t = ST and ReinMax prefer to set
4.2 Independence of ReinMax over Baseline Subtraction
To better understand the effectiveness of ReinMax, we further provide an alternative derivation that does not rely on the selection of the baseline. For simplicity, we only discuss Oc and mark it VR. Similar to Equation 2 have:
50,
Oc
d d TTk Σ क. (J(E)) - ((I)).
V
= - ((I)
00,
do
It is worth mentioning that the derivation of Equation B leverages the derivative of the softmax function (i.e., for π softmax(0), we have /00 गτ (8ax TT)) and does not involve the baseline subtraction technology.
Remark 4.2. In Equation E we approximate f(Ix) f(I,) a5()) af (T I,), and mark OL.
21.
the resulting second-order approximation of V k as V 2nd-orrder-wo-baeeliee T&k E, TT (5(.) @f(L) (Ix I(), Then, we have E ReinMax]
( 01.
v 2nd-oonder-wo-busellne
01s
The proof of Remark 2.2 is provided in Appendix D
As in Remark 2.21 applying the Heun's 's method on Equation and Equation lead the same gradi- ent estimator, which implies another benefit of using (D) as the baseline: the resulting gradient estimator does not rely on additional prior (i.e., its derivation can be free of baseline subtraction).
Temperature Scaling for Gradient Estimators
Here, we discuss how to apply temperature scaling, technique widely used in gradient estimators, our proposed method, ReinMax. While the typical practice to set the temperature 5 to small values for STGS, we show that ST and ReinMax need different strategy.
Temperature Scaling for VTTGS As introduced in Section A VTGG conduct two-step approx- imation: (1) it approximates ming E[] as ming F (sofimxx, (0 + G)))]: (2) it approximates
((D), Since the bias introduced in both steps can be controlled by т, V prefers to set as relatively small value.
affsofuaax (0+ (0+G))
asoftmax. (0+G)
as "DD
Temperature Scaling for Vs and V ReiiMax As in Section R it does not involve temperature scaling to show V and ReinMax work as the first-order and the second-order approximation to the gradient. Correspondingly, temperature scaling technology cannot help reduce the bias for Vst in the same way it does for STGSS As in Figure 2 STGS, GR-MCK, and GST-1. 10 work better when setting the temperature
1. ST and ReinMax work better when setting the temperature
1.
Thus, we incorporate temperature scaling to smooth the gradient approximation (π, = softmax, (0) as VReinMax = 2 -ッ z ST. is worth emphasizing that in V VeinMaax is used to stabilize the gradient approximation (instead of reducing bias) the cost of accuracy. Therefore, the value should be larger equal to 1.

---

Table 1: Performance on ListOps.
STGS
GR GR-MCK
GST-1.0
ST
ReinMax
Valid Accuracy
66.95+3..05
66.53+00.58
66.28
66.51-0.76
67.65+1.255
Test Accuracy
67.30.2.50
66.53.10.86
66.30100.2
66.26-0.48
68.07:11.18
Table 2: Training ELBO on MNIST (N x M refers to N categorical dim. and M latent dim.).
AVG
xx4
4 24
x 16
16 x 12
64 x 8
10 30
STGS
105.20
126.8510.85
101.32.10.33
99.32.10.33
100.0910.32
104.10.41
99.6310..3
GR-MCK
107.06
125.94=0.71
99,9610.25
99.58-0.31
102.54=0.48
112.34:10.48
102.02::0188
GST-1.0
104,25
1,2.35
101.49
9..29
98.12 08:5077
102.53++0.57
98.64
ST
116.72
135.53+0.31
112.03=0.03
112.94+0.32
113.31=0.43
113.90+-0.28
112.63+0.34
ReinMax
103.21
124.6610.88
99.77:10.45
97.70=0.39
98.06:10.53
100.71:1.700
98.37=0.44
6 Experiments
Here, we conduct experiments on polynomial programming, unsupervised generative modeling, and structured output prediction. In all experiments, we consider four major baselines: Straight- (STGS), Gumbel- Monte Carlo (GR-MCK) and Gapped Frraigh-TThroug (GST-1.0). For a more comprehensive comparison, we run complete grid search on the training hyper-parament for all methods. Also, we would reference results from the literature when their setting comparable with ours. More details are elaborated in Appendix
Through (ST), Straught-Throuh Gumbel-Softng
6.1 Polynomial Programming
Following previous studies (lucker et all 2017: Grathwohl all 2018 Pervez et all 2020 Paulus stall 2021),, we start with simple problem. Consider /, i.i.d. latent binary variables XI,... X
{0, 1} and constant vector R parameterize the distributions of (X,,,.. X2} with L
(0)) and 0, E R². Following previous studies, we set every dimension as 0.45, i.e., Vi, c; = 0.45, and use ming Ex |X
softmax functions, i.c., X, ~ Multinomial
3 as the objective.
Training Curve with Various p. We first set the number of latent variables (i.e., L.) as 128 and batch size as 256. The training curve visualized in Figure II for = 1.5, and ReinMax achieved near-optimal performance and the best convergence speed. Meanwhile, we can observe that ST and GST-1.0 do not perform well in all three cases. Although the final performance of STGS and GR-MCK is close to ReinMax, ReinMax has faster convergence speed.
all cases,
6.2 ListOps
We conducted unsupervised parsing on ListOps (Nangıa & Bowman, 2018) and summarized the average accuracy and the standard derivation in Table MI We also visualized the accuracy and loss the valid set Figure Although the ST algorithm performs poorly on polynomial programming, achieves reasonable performance our proposed method stands out and brings consistent improvements. This further demonstrates the benefits of achieving second-order accuracy and the effectiveness of our proposed method.
this task. Also, while all baseline methods perform similarly,
0.700
1.30
RenMaa
0.675
1.25
STE3
Set
1.20
G5T-16
0,000
Accuracy
6
1.25
0.575
1.00
0.550
0.35
0.525
0.96
0.500
Epoch
Epoch
Figure 3: The accuracy (left) and loss (right) on the valid set of ListOps.
7

---

140
Max
138
STC3
GR-CCK
136
Similority
134
0.7
132
130
CO5
128
0.5
126
0.4
124
40
100
240
1.50
0.3
20
60
20
1000
120
1.40
160
Epoch
Epoch
Figure 4: The training ELBO (left) and the cos similarity between the gradient and its approxima- tions (right) on MNIST-VAE (with 4 latent dimensions and 8 categorical dimensions).
Table Average time cost (per epoch) / peak memory consumption on quadratic programming (QP) and MNIST-VAE. QP is configured to have 128 binary latent variables and 512 samples per batch. MNIST-VAE is configured to have 10 categorical dimensions and 30 latent dimensions.
ReinMax
STGS
GST-1.0
GR-MCKj00
GR-MCKa
MR-M
QP
0.2s 0225/6.500 6.5Mb
0,0M/
0.25/5.5Mb
0.2s 8.0Mb
0.8s /033Gb
2.2s/IGb
6.6s 3Gb
VAE 5.2s / 13Mb
5.2s/ 13Mb
5.2s / 13Mb
5.2s/ / 13Mb
76Mb
5.2s
5.4s /0.6Gb
Table 4: Performance on BATS-Bench. Baseline results are referenced from Dong et al. (20207)
CIFAR-10
CIFAR- 100
ImageNet- 16--NNet---112 16-120
validation
test
validation
Test
validation
test
GDAS STGS
89.68=0.72
93.23+40.58
68.35+2.71
39.55++0.00 68.1742.50
39.40=0.00
GDAS ReinMax
90.01=0.12
93.44.10.23
69.2912.33
69.41:12.24
41.47.10.79
42.03=0.41
6.3 MNIST-VAE
We benchmark the performance by training variational auto-encoders (VAE) with categorical latent variables on MNIST (LeC'un et all 1198). As we aim to compare gradient estimators, focus our discussions on training ELBO. We find that training performance largely mirrors test perfor- mance (Dong et all 20206, 2021: Fan etll 2022) and briefly discussed test ELBO in Appendix IE
Biases of the Approximated Gradient. With 4 latent dimensions and 8 categorical dimensions, iterate through the whole latent space (the size of the latent space is only 4096), compute the gradient as in Equation D. and measured the cosine similarity between the gradient of latent vari- ables and various approximations. As visualized in Figure ReinMax achieves consistently more accurate gradient approximation across the training and, accordingly, faster convergence. Also, we can observe that, besides faster convergence, the performance of ReinMax is more stable.
Experiment with Larger Latent Spaces. consider 4 settings with the latent space of 248. Then, following at (2222), also conduct experiments with 10 latent dimensions and 30 categorical dimensions (the size of the latent space is 10)0). As summarized in Table E ReinMax achieves the best performance all configurations.
Let us proceed to larger latent spaces. First, we
GST-1.0 Performance on Different Problems. It is worth mentioning that, despite GST-1.0 achieving good performance on most settings of MNIST-VAF, it fails to maintain this performance on polynomial programming and unsupervised parsing, as discussed before. Upon discussing with all (2022), suggest that this phenomenon is caused by the characteristic of G.T- which behaves similarly to ST on problems with a near one-hot optimal distribution. In other words, GST- 1.0 has an implicit prior and prefers distributions that are not one-hot. At the same time, a different variant of GST (i.e., GST-p) would behave similarly to STGS on problems with a near one-hot optimal distribution, which achieves significant performance boost over GST-1.0 on polynomial programming. However, on MNIST- VAE and ListÖps, GST-p achieves an inferior performance.
This observation verifies our intuition that, without understanding the mechanism of ST, different ap- plications have different preferences on its configurations. Meanwhile, ReinMax achieves consistent improvements in all settings, which greatly simplifies future algorithms developments.
8

---

Table 5: Training ELBO on MNIST. "All baseline results are referenced from Fan all (2022)
R1.00
DiSARM- Tree"
STGS
GR-MCCK
GSTT-1.0"
ST
Max
Neg. ELBO 104.03=0.23
103.10:0.25
97.32:10.20
110.74=1.23
96.091:.255
116:0.099
93.44=0.51
herMec
NorMic
244
402
112
246
100
110
210
105
0.0
9.2
0.5
3.8
10
100
0.0
3.2
0.4
3.0
0.0
3.2
0.5
3B
10
Updare
Updare
(a) MNIST when K=2.
(b) Omniglot when K=2.
(c)
when K-2.
MM
115
246
110
244
204
666
206
226
146
0.2
100
1,0
02
Up302e
34
Up3ene
1p300e
14
(d) MNIST when K=3.
(e) Omniglot when K=3.
(I) Fashion-MNIST when K=3.
Figure 5: 2x200 VAE training curves on MNIST, Omniglot, and Fashion- when K=2 or 3.
Table 6: Train ELBO of 2 x 200 VAE on MNIST, MaShion-MIIST, and Omniglot.
Baseline
results referenced from et (2022). K refers to the number of evaluations.
RELAX*
ARMS
DisARM
Double CV"
RODEO
ReinMax
MNIST
100.84+1..14 0.009
100.941:.09
100.46
97.83 0.36
K=3
Fashion-MNIST
237.05.10.12 237.7410.11
237 10.10.1
236.8810.12
234.53.10.42
Omniglot
115.70.10.08
115.32.10.07
115.06:10.12
115.0110.00
107.51:10.42
MNIST
102.75.10.08
102.
101.8910.17
98.17.00.9
K-2
Fushion-MNIST
237, 137.6811.113
10..6
237.44:10.09
234.89-0..21
Omniglot
1 116.50+0.04
116.39+0.10
115.93+00.66
107.79+0.27
6.4
Applying ReinMax to Differentiable Neural Architecture Search
To demonstrate the applicability of ReinMax as a drop-in replacement, we conduct experiments following the topology search setting in the HATS-Benche benchmark (Dong et al, 20203), and sum- marize the results in Table g. GDAS is an algorithm that employs STGS to estimate the gradient of neural architecture parameters (Dong Yang. 2019). We replaced STGS with ReinMax as the gra- dient estimator (configurations elaborated in Appendix E). ReinMax brings consistent performance
improvements across all three datasets, demonstrating the great potential of ReinMax.
6.5 Comparisons with REityORCF.style Methods
Here, we conduct experiments to discuss the difference between ReinMax and REINFORCEE-tyle
methods. First, following Ean el all (21122), we conduct experiments on the setting with a larger batch size (i.e., 200), longer training (i.c.. 5 x 10y steps), 32 latent dimensions, and 64 categorical dimensions (details are elaborated in Appendix E.. As in Table E, ReinMax outperforms all basclines, including two DEINFORCE-based methods (Dong et all 20206 202I).
We further conduct experiments to compare with the state of the art. Specifically we apply ReinMax to Bernoulli VAEs on MNIST, Fashion- (Xiao et all, 2017), and Omniglot(Lake et all 2015). adhering closely to the experimental settings of Shi et all (2022)), including pre-processing. model architecture, batch size, and training epochs. As in Tables E and Figure EL ReinMax consistently out- performs RODEO across all settings. To better understand the difference between RODEO and Rein- Max, we conduct more experiments on polynomial programming (as elaborated in Appendix Ed).

---

Overall, ReinMax achieves better performance in more challenging scenarios, i.e. smaller batch size, more latent variables, or more complicated problems. Meanwhile, REINFORCE and RODEO achieve better performance simpler problem settings, i.e., larger batch size, fewer latent variables, or simpler problems. This observation matches our intuition:
REIFORCE-Spi algorithms excel as they provide unbiased gradient estimation but may fall short in complex scenarios, since they only utilize the zero-order information (i.c., a scalar f(D)).
ReinMax, using more information (i.e., a vector af(D)), ), handles challenging scenarios better. aD Meanwhile, as consequence of its estimation bias, ReinMax leads slower convergence in some simple scenarios.
6.6 Discussions
Choice of Baseline. As introduced in Section A.ll the choice of subtraction baseline has huge impact on the performance. Here, we demonstrate this empirically.
Megaltve ELBO
334
We use 1 Σ; f(T) as the baseline and compare the re- sulting gradient approximation with ReinMax. As visu- alized in Figure &. ReinMax, which uses (D] as the baseline, significantly outperforms the one that uses 4 >; (I.) as the baseline. We suspect that the gradi- ent approximation using E; the baseline is very unstable as it contains the m.p(D) term.
EInD!1
4..4
10.4 22
Categurical Dim Latent Dim
Figure Training -ELBO on MNIST: VAE when using E. f(I,) and E[f(D)] as baselines respectively.
On MNIST-VAE (four settings with the 218 latent space), we utilize heatmaps to visualize the final performance of all five methods under different temperatures, i.c., 00.1, 0.7, 3, 5.. As in Figure D. these methods have different preferences for the temperature configuration. Specifically, STGS, GST-1.0, and GR-MCK prefer to set the temperature 7 S 1. Differently, ST and ReinMax prefer to set the temperature 7 2 1. These observations match analyses in Section B that small can help reduce the bias introduced by STGS-styb methods. Also, it verifies that ST and ReinMax work differently from STGS, GST-1.0, and GR-MCK.
Temperature Scaling.
Efficiency.
As summarized in Table BL we can observe that, since GR-MCK uses the Monte Carlo method reduce the variance, has larger time and memory consumption, which becomes less significant with fewer Monte Carlo samples (we use GR-MCK, to indicate GR-MCK with Monte Carlo samples). Meanwhile, all remaining methods have roughly the same time and memory consumption. This shows that the computation overheads of ReinMax are negligible.
7 Conclusion and Future Work
In this study, we seek the underlying principle of the Straight-Through (ST) gradient estimator. We formally establish that ST works as a first-order approximation of the gradient and propose a novel method, ReinMax, which incorporates Heun's Method and achieves second-order accuracy without requiring second-order derivatives. We conduct extensive experiments on polynomial programming, unsupervised generative modeling, and structured output prediction. ReinMax brings consistent improvements over the state-of-the-art methods.
It is worth mentioning that analyses in this study further guided us to empower Mixture-of- training eL all 2023). Specifically, for gradient approximation of sparse expert routing, while ReinMax requires the network to be fully activated, Liu all (2023) uses f(0) as the baseline and only requires the network to be partially activated. In the future, we plan to conduct further analyses on the truncation error to stabilize and improve the gradient estimation.
Acknowledgement
We would like to thank all reviewers for their constructive comments, the engineering team at Mi- crosoft for providing computation infrastructure support, Alessandro Sordoni, Nicolas Le Roux, and Greg Yang for their helpful discussions.
10

---

References
Ascher, U. M. and Petzold, L. R. Computer methods for ordinary differential equations and differential-algebraic equations. 1998.
Bengio, Y., Léonard, N., and Courville, A. C. Estimating propagating gradients through stochastic neurons for conditional computation. ArXiv, abs/1 2013.
Choi, J., Yoo, K. M., and goo Lee, S. Learning to compose task-specific tree structures. In AAAI, 2017.
Chung. J., Ahn, S., and Bengio, Y. Hierarchical multiscale recurrent neural networks. In ICLR, 2017.
Dayan, P.. Hinton, G. E., Neal, R. M., and Zemel, R. S. The helmholtz machine. Neural Computa- tion, 7:889-904, 1995.
Dong, and Yang, Searching for robust neural architecture in four gpu hours. CVPR, 2019.
Dong, X., Liu, L., Musial, K., and Gabrys, B. Nats-bench: Benchmarking nas algorithms for architecture topology and size. TPAMI, 2020a.
Dong, Z., Mnih, A., and Tucker, G. Disarm: An antithetic gradient estimator for binary latent variables. NeurlPS, 2020b.
Dong, Z., Mnih, A., and Tucker, G. Coupled gradient estimators for discrete latent variables. In NeurIPS, 2021.
Fan, T.-H., Chi, T.-C., Rudnicky, A. and Ramadge, P. J. Training discrete deep generative models via gapped straight-through estimator. In /CML, 2022.
Fedus, W., Zoph, B., and Shazeer, N. M. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. ArXiv, abs/2101.03961 2021.
Grathwohl, W., Choi, D., Wu, Y., Roeder, G., and Duvenaud, D. K. Backpropagation through the void: Optimizing control variates for black-box gradient estimation. In ICLR, 2018.
Gregor, K., Danihelka, I., Mnih, A., Blundell, C., and Wierstra, D. Deep autoregressive networks. In ICMI, 2014.
Gu, S. S., Levine, S., Sutskever, I., and Mnih, A. Muprop: Unbiased backpropagation for stochastic neural networks. In ICLR, 2016.
Gumbel, E. Statistical theory of extreme values and some practical applications A series of lectures. 1954.
Jang, E., Gu, S. S., and Poole, B. Categorical reparameterization with gumbel-softmax. In ICLR, 2017.
Kingma, D. and Ba, J. Adam: A method stochastic optimization. In ICLR, 2015.
Kingma, D. and Welling, M. Auto-encoding variational bayes. CoRR, abs/1312.61 2013.
Lake, B. M., Salakhutdinov, R., and Tenenbaum, J. B. Human-level concept learning through prob- abilistic program induction. Science, 2015.
LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. Gradient-based learning applied document recognition. Proc. IEEE, 1998.
Liu, H., Simonyan, K., and Yang, Y. Darts: Differentiable architecture search. In ICLR, 2019.
Liu, L., Jiang, H., He, P. Chen, W. Liu, X., Gao, J., and Han, J. On the variance of the adaptive learning rate and beyond. In ICLR, 2020.
Liu, L.,, Gao, J., and Chen, W. Sparse backpropagation for moe training. ArXiv, 1b./2310.00811, 2023.

---

Maddison, C. J., Tarlow, D., and Minka, T. P. A sampling. In NIPS, 2014.
Mullin, A. A. and Rosenblatt, F. Principles of neurodynamics. 1962.
Nangia, N. and Bowman, S. R. Listops: A diagnostic dataset for latent tree learning. ArXiv, abs/ 10ss/880406028, 2018.
Neal, R. M. Connectionist learning of belief networks. Artif. Intell., 56:71-113, 1992.
Paulus, M. B., Maddison, C. J., and Krause, A. Rao-blackwellizing the straight-through gumbel- softmax gradient estimator. In ICLR, 2021.
Pervez, A., Cohen, T., and Gavves, E. Low bias low variance gradient estimates for boolean stochas- tic networks. In ICML, 2020.
Rennie, S. J.. Marcheret, E., Mrouch, Y., Ross, J., and Goel, V. Self-critical sequence training for image captioning. In CVPR, 2016.
Rosenblatt, The perceptron, a perceiving and recognizing automaton Project Para. Coreell Aeronautical Laboratory, 1957.
Rumelhari, D. E., Hintont, G. E., Ronald, J., and Williams. Learning representations by backpropa- gating errors. Nature, 323:533536, 1986.
Shi, J., Zhou, Y., Hwang. J., Titsias, M., and Mackey, L. Gradient estimation with discrete stein operators. In NeurlPS, 2022.
Tokui, and Sato, I. Evaluating the variance of likelihood-rati, gradient estimators. JCML, 2017.
Tucker, G., Mnih, A., Maddison, C. J., Lawson, J., and Sohl-Dickstein, J. N. Rebar: Low-variance, unbiased gradient estimates for discrete latent variable models. In NIPS, 2017.
van den Oord, A., Vinyals, O., and Kavukcuoglu, K. Neural discrete representation learning. In NIPS, 2017.
Weaver, L. and Tao, N. The optimal reward baseline for gradient-basee reinforcement learning. In UAI, 2001.
Williams, Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning, 8:229-256. 1992.
Xiao, H., Rasul, K., and Vollgraf, R. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms. ArXiv, abs/1708.07747, 2017.
12

---

A Theorem 31
Let define the first-order approximation of V as V Ist-order =i n, 01, which approximates f(I)) ((z) in Equation B as at, (I: 1j). Theorem 3.1.
a
(I, 1
dx:
d0
af(L)
B
Ist-order
Proof. Based on the definition, we have:
ΣΣ T if(1,)
(1, - (1:-1)) dı
Ist-order
01,
d0
Σ n 011
0f(L)
I
d
Σ ;
}f(4)
d
(9)
do
01,
do
Since E; we have Σ, du 0. Also, since π
Ei Nils we have
dm
Σ,
1
ao
Thus, together with Equation 9 we have:
Σ Of(L))
Σ I
d
Ist-order
01j
d0
El DD) d E sT)
El
3D d0
B Theorem +
Theorem 3.2.
E ReinMax]
V 2nd-order-
Proof. Here, we aim to proof, 4h: E [1, n). we have E ReinMax,k V 2nde- As defined in Equation & we have (note that 5DI, is the indicator function of the event D = Ix):
T 0f(L) 0f(L)
)I: 1;) dT:
ΣΣ
2nd-order.k
2
01,
01,
do.
TT TT (dik T&) (S(L) Of(1)
= ΣΣ
01, 0I; (I - 1,)
2
3
Σ TYT
a)(1,)
+ (
){ 11)
2
01,
0I,
Ty
0f(Ik)
(Ix Σ T1)
+
Σ
J 0f(1,)
(I: - -1;)
2 Oh,
2 01,
3
3
ED- (dDI
Df(D)
(In
Σ
T, T + 1
ED T 0D
0f(D)
(I. In)]
2
0D
OS(D)
= ED
T (Ix ID) + 5D TTTTi))
(10)
0D
At the same time, based on the definition of ReiiMMax we have:
E ReinMax,k]
= EDa
Df(
(2
π& åDI.
(Ds Σ n: + +0D
I.)
Th
(Dx
Σ
T. T.Li))
3D
2
2
2
0f(D)
ED
(T (Is - In) + pD (Ix - 16a-Ammn T;li))
(11)
0D
13

---

Since 8pI. - T(I,) = ÓDI, (ID - E, I together with Equation IT0 and I.I we have: ReVeRinax,.k 2add-order,e
C Remark 4.
Remark 3.1. When E, fII(I.) is used as the baseline and f(I,) -((;) is approximated as 81,
01(1 (I; 13}. we mark the resulting first-order approximation of
V Ist-order-arg-bant baseline
Then, we have:
DD STl V
Ist-order-avg ai--orderavg-baseeine baseline
TD
Proof. Using E, (I.) as the bascline, we have:
d
d d0
V [((()) )
do (((L))
Approximating ((I, ((T,) aful) al,
(I: Ij), have:
Of(I)
d
-tt-orderaagg-baseline
ΣΣ Φ 01,
(1: - 1,
d0
Li Pi.aj nj 0f(L) Σ 1: d T
01,
00
ΦD Srl
E
D Remark 4.
Remark 3.2. In Equation E we approximate f(L)) as
1
BILL) Of(TL) 21:
({ 1;), and mark
81,
the resulting second-order approximation of l as V ondl-rrder = T E. _T (1(1. of(Ts) (Ix -I;,, Then, we have:
81,
0T,
E V ReinMar] 2ndloonder-wo-busellnee
Proof. Here, aim to proof, Ak [1, n), we have E V VeeinMax,k V 2ndd
1 Of(L)
0f(Ik)
Vnnd-orddrr-o--baellnee
Σ π:
OI )(I 11)
0I,
1 (I,
10
n& Σ F: 01,
(Ir. 1;) T& Σ Riz 0l
(Ir I.)
E = (f(
म. [I In) SD (Ix E, TIIT)
E ReinMax,ki
aD
14

---

E Forward Euler Method and Heun's Method
For simplicity, consider simple function g(x) R + R that is three times differentiable [t]. Now, we proceed to a simple introduction approximate Jr, g' d'(x)rr with the Forward Fuler Method and the Heun's Method. For a detailed introduction to numerical ODE methods,
please refer to Ascher & Petzold (1998).
Forward Euler Method. Here, we approximate g(t,) with the first-order Taylor expansion, i.e., g(41) = g(4) + (1)) (4 (0) O((4 - (q))), then we have frg g' (x))dr s g (l)) - lo).
Since we used the first-order Taylor expansion, this approximation has first-order accuracy.
Heun's Method. First, we approximate g(t1) with the second-order Taylor expansion:
g((,) = g(to) + g(()) (t, -t) g"(to) (t, to)² O(t to)²))
(12)
2
Then, show that match this approximation by combining the first-order derivatives of two samples. Taylor expanding g g((1)) to the first-order, we have:
g((1,) = (1o) + g" g(o) (1 - (n) + +(((1 - 20)
Therefore, we have:
g(fo) + g' (to) + g (t1)) (t) - to) g(to) + (to) (t1 to) g'(t))
(t1 ( to)
2
It is easy to notice that the right-hand side of the above equation matches the second-order Taylor expansion of g(41) as in Equation I2. Therefore, the above approximation (i.e., approximating g((1) g(to) as g'(t))+g' (t:) (41 - fo)) has second-order accuracy.
By setting g(x) f(x (1- I,),,, we have f(,) f(1,). Then, it is easy to notice that the forward Euler Method approximates g(1) - (0) =
Connection to f(I))(())))))))) f(I,) in Equation &
f(I.)- as af())
01, (1)))) and has first-order accuracy. Also, the Heun's 's Method approximates
-(4)) - ((1, as ( (aL)) 2r()) (I, - I5) and has second-order accuracy.
01,
F Experiment Details
F1 Baselines
Here, we consider four methods as our major baselines:
Throgggh (ST; Bengio all 2013)) backpropagate through the sampling function if it had been the identity function.
Gumbel-Softmax (STGS; Jang al. 2017) integrates the Gumbel reparameter-
ization trick to approximate the gradient.
leverages the Monte Carlo method to reduce the variance introduced by the Gumbel noise in STGS. To obtain the optimal performance for this baseline, we set the number of Monte Carlo samples 1000 in most experiments. Except in our discussions of efficiency, we set the number of Monte Carlo samples to 100, 300, and 1000 for more comprehensive comparisons.
Gumbel-Rao Monte Carlo (GR-MCK; Paulns
Gapped Brraight-Througe (GST-1.0; Fan all 2022) aims to reduce the variance of STGS and constructs a deterministic term to replace the Monte Carlo samples used GR-MCK. Here, as suggested in (Ian all (12)), we set the gap (a hyper-parameter) 1.0.
GST-1.0 Performance. Despite GST-1.0 achieving good performance most settings of MNIST- VAE, it fails to maintain this performance on polynomial programming and unsupervised parsing, as discussed before. At the same time, different variant of GST (i.e., GST-p) achieves a signifi- cant performance boost over GST-1.0 on polynomial programming. However, on MNIST-VA VAE and ListOps, GST-p achieves an inferior performance. Upon discussing with the author of the GST-1.0, we suggest that this phenomenon is caused by different characteristics of GST-1.0 and GST-p.
This observation verifies our intuition that, without understanding the mechanism of ST, different ap- plications have different preferences on its configurations. Meanwhile, ReinMax achieves consistent improvements in all settings, which greatly simplifies future algorithms developments.
15

---

F.2
Without specifically, we conduct full grid search for all methods in all experiments, and report the best performance (averaged with 10 random seeds on MNIST-VAE and 5 random seeds ListOps). The hyper-parameter search space is summarized in Table / The search results for Table D and Table @ are summarized in Table &
Table 7: Hyper- search space.
Hyperparameters
Search Space
Optimizer Learning Rate
A Kingma & Ba, 2015), RAdam et 20200)
{0.001, 0.0007, 0.0005, 0.0003) (0.1, 0.3, 0.5, 0.7, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5)
Temperature
Table paper-parameters Search Result for Results in Table M and Table D.
STGS
GR-MCK
GST-1.0
ST
ReinMax
Optimizer
Adam
Adam
Adam
Adam
Adam
MNIST-VAB 8 x 4
Learning Rate
0.0003
0.0005
0.0005
0.001
0,0005
Temperature
0.5
0.5
0.7
1.3
1.3
Optimizer
RAdam
RAdam
RAdam
RAdam
RAdam
MNIST-VAE 4 x 24
Learning Rate
0.0005
0.0005
0.0005
0.001
0.0005
Temperature
0.3
0.3
0.5
1.5
1.5
Optimizer
RAdam
RAdam
RAdam
RAdam
RAdam
MNIST-VAE 8 16
Learning Rate
0.0005
0.0007
0.0007
0.001
0.0007
Temperature
0.5
0.7
0.5
1.5
1.5
Optimizer
RAdam
Adam
RAdam
Adam
RAdam
MNIST-VAL VAE 16 x 12
Learning Rate
0.0007
0.0005
0.0007
0.0005
0.0007
Temperature
0.7
1.0
0.5
1.5
1.5
Optimizer
RAdam
Adam
RAdam
Adam
RAdam
MNIST-VAE 64 8
Learning Rate
0.0007
0.0007
0.0007
0.0005
0.0005
Temperature
0.7
2.0
0.7
1.5
1.5
Optimizer
RAdam
RAdam
RAdam
RAdam
RAdam
MNIST-VAE 10 30
Learning Rate
0.0005
0.0005
0.0005
0.0007
0.0005
Temperature
0.5
1.0
0.5
1.4
1.3
Optimizer
RAdam
RAdam
RAdam
RAdam
RAdam
ListOps
Learning Rate
0.0005
0.0005
0.001
0.001
0.0007
Temperature
0.1
0.3
0.1
1.4
1.1
Polynomial Programming. As this problem is relatively simple, set the learning rate to 0.001 and the optimizer to Adam, and only tune the temperature haper-parameter.
Following the previous study (Dong et all, 2020R, 2021: atal 2022), we used 2-layer MLP the encoder and the decoder. We the hidden state dimension of the first-layer and the second-layer as 512 and 256 for the encoder, and 256 and 512 for the decoder. For our experiments on MNIST-VAE with
MNIST-VAE.
latent dimensions and 64 categorical dimensions, we set batch size to 200, training steps to 5 x 10 and activation function to LeakyReL.U, in order to be consistent with the literature. For other experiments, we set the batch size 100, the activation function to ReI.U, and training steps to 9.6 x 10a (i.e., 160 epochs).
16

---

Table 9:
Test ELBO on MNIST. Hyper-parameters are chosen based on Train ELBO.
AVG
x
4 x 24
8 x
x 30
STGS
106.89
128.09 0.79
103. 0,4.60:::.45
99.32.10.33
0.332
106.20:10.46
101.6110.54
GR-MCK
109.03
127.90
102.76.10.33
102.12:10.29
104.23
113,54.10.50
103.6210..3
GST-1-0
106.85
1.12
103.95++0499
101.44+0..2
101.28=0.59
105.44+0.62
100.78=0.44
ST
118.85
137.06-0..
03.441
114.25+
114,48
115,43-+:.29
118.46+0.18
ReinMax
105.74
126.89-0.99
102.40=0.43
100.63=0.41
100.85=0.50
102.91=0067
100.75=0.50
Table 10: Test ELBO on MNIST. Hyper-parameten are chosen based on Test ELBO.
AVG
8
2
8 16
16 12
64 x 8
10 x 30
STGS
107.15
128.09-00 0.79
103.25+0.22
101.44+0.32
102.29-0.39
101.61+0.54 106.20+0.46
GR-MCK
108.87
027.86-0.54
102.40=0.37
101.59+0.22
104.22=0.63
103.62=0.13 113.54:0.50
GST-1.0
106.55
128.03 1.02
103.63.10.24
100.6710.34
10..94.10.99
105.44.10.22
100.51:0.37
118.79
137.05=0.36
113.2310043
114.111:.311
115.48=0.56
115.43:10.29
118.46::0.18
ReinMax
105.60
126.29=0.32
102.40::0.43
100.45:10.66
1.5.84=0.56
102.91::0668
100.69=0.48
Table 11: Train ELBO MNIST. Hyper- Hyper-parameters are chosen based on Test
ELBO..
AVG
4 x 24
30
STGS
105.31
126.85
101,8110.14
99.32.10.33
100.22.10.47
104.02.10.41 .1
99 99.63-10.63
GR-MCK
107.37
025.53_0.55
100,47:0.31
99.7511.229
103.11=0.58
112.34110449 08
102.02-10118
GST-1.0
104.60
126.63 1.16
102.11+0.24
98.40+0.34
98.76++0.41
102.53++0.57
99.14+00.30
ST 117.76
136,75_0.. 0.22
0.2.09:0..50
113.06110.26
113.31
113.90.10.28
117.4610.099
ReinMax
103.40
124.92-0.38
99.77-0.45
98.06=0.31
98.51+0.54
100.71++0700
98.40-+0.48
Differentiable Neural Architecture Search. We adopt most of the hyper-parameter setting from Dong et al. (M(21T). Since GDAS employs temperature schedule (decaying linearly from 10 to 0.1), and temperature scaling works differently in ReinMax and STGS (as discussed in Section and Section b.f)), we removed the temperature scaling (i.e., set the temperature constant 1.0) and increased the weight decay (i.e., from 0.001 to 0.09).
ListOps. followedd the same setting of Ean (2022), i.c. used the same model configuration as in al (2 and set the maximum sequence length to 100.
F3 Hardware and Environment Setting
Most experiments (except efficiency comparisons) are conducted on Nvidia P40 GPUs. For effi- ciency comparisons, we measured the average time cost per batch and peak memory consumption quadratic programming and UNIST-VAE the same system with idle A6000 GPU. Also, to better reflect the efficiency of gradient estimators, we skipped all parameter updates in efficiency comparisons.
E4 Additional Results on Polynomial Programming
We visualized the training curve for polynomial programming with various batch sizes and latent dimensions in Figure & (for p 1.5), Figure g (for p 2), and Figure I[ (for p 3).
F.5
Additional Results on MNIST-VAE
In our discussions in Section E. we focused on the training ELBO only. Here, we provide brief discussion the test ELBO.
Choosing Hyper-parameter Based on Training Performance. method, we select the hyper-parameter based on its training performance. The Test -ELBO in this setting is summarized Table H Despite the model being trained without dropout or other overfitting reduction techniques, ReinMax maintained the best performance in this setting.
Similar to Table D for each
17

---

(a) p and
(b)
3 and
00.45 0.45).
1u
(c) 2 and
0.5
1.5
L-0.5
(d) 2 and
00:45 0.45].
(c) 1.5 and c
0.5 1.5
1.505
(l) p 1.5 and e [0.45, 0.45].
X - where X e
Figure 7: Training curves of polynomial programming, i.c. ming Ex
(0, 1)4, X, üdd (ultinomial(softmax 0 = [0,, Or.'²' Og E IR?, and is the number of latent dimensions.
Based on Test Performance. We also conduct experiments by selecting syper-parameters directly based their test performance. In this setting, the test -ELBO summarized in Table III and the training EL ELBO is summarized Table III ReinMax achieves the best performance in all settings except the test performance of the setting with 10 categorical dimensions and 30 latent dimensions.
Choosing
F.6 More Comparisons with RODEO
To better understand the difference between RODEO and ReinMax, we conduct more experiments on polynomial programming, i.e., ming Ex |X [ Specifically, we consider polynomial pro- gramming under two different settings that define e differently:
setting A, we have e = [0.45, ... 0.45). This the setting we used in the submission.
setting B, we have e = L L
10.5
L-0.5]
As to the difference between the Setting A and the Setting B, we would like to note: setting A, since Vi,c 0.45 and 0, Uniform(-0.01, 0.01) initialization,
Fr =x.-ssoftmxx00 (0:) - would have similar values. Therefore, the optimal control variates for 0, are similar across different
In setting B, we set e; to different values for different i, and thus the optimal control variate for O are different across different i. Therefore, Setting simpler setting for applying control variate to REINFORCE.
18

---

As in Figure ZI ReinMax achieves better performance in more challenging scenarios, i.c., smaller batch size, more latent variables, more complicated problems (Setting B VAEs). Meanwhile, REINFORCE and RODEO achieve better performance simpler problem settings, i.e., larger batch size, fewer latent variables, or simpler problems (Setting A).
M-Mpa-1
W-12
18
B-
B-H112-2
B-H
8-
8-210-11
8-210-0
8- 11
331
43
RaiMaa
STCS
GRAKK
ST
GST.0 ..... min 14
Figure& 8: Polynomial programming training curve, with different batch sizes and random vari-
able counts (L), i.e., ming E [ where RLx2 X E {0, 1)4, and X; iid Alltinommal((sotmax((e More details are elaborated in Section B.
4
19

---

1-119-
B-310 10+4
Ri 10+4
B-3 10-10
B-1 10-22
331
5-19--
E-41004
B-812-7
B-R13-4
8-210-1
2-110-2
E=
K-G
K-
X. 44 140
26 1410-11
83
13-
M M-
25
25
STOS
Quadratic programming training curve, with different batch sizes and random vari- able counts (L), i.e., ming E 1 where A castnommialssotmax//a More details are elaborated in Section B.
Figure
R4x2 X 10,1)4 and X; idd
20

---

8-119-1
B-119-7
Bi-3 H0
Ri 10+
B-3100-18
B-1 Dedl
K-
M-
M 1
W 100-0
Ma
Mi R
Ko f1
k31
k3
25
25
STOS
GRANKK
Figure 10:
Polynomial programming training curve, with different batch sizes and random where O R4x2 X 1), and X, ndd
variable counts (L), i.e., ming E X (ultinom More details are elaborated in Section B.
21


Reviews:
[{"review_id": "6VPsW4FqU7", "reviewer": "Reviewer_GD4h", "summary": "The paper shows that the straight-through estimator is a first-order approximation of the gradient. The authors then propose a method, called ReinMax, which provides a second-order approximation with negligible computational overhead. Experiments are performed in several settings involving discrete variables (polynomial programming, structured output prediction, and discrete latent variable generative models), showing that ReinMax is more accurate and stable.", "questions": "Could you please elaborate on the applicability of the ReinMax estimator? Does this estimator (or the insights developed in the theoretical sections of the paper) alleviate the “laborious and time-consuming” work of developing “different ST variants for different applications in a trial-and-error manner”?", "limitations": "As mentioned above, I would appreciate a clearer discussion of which existing forms of straight-through estimator the ReinMax estimator can replace.\n", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "**Originality**\n\nThe originality of the paper is fairly strong. As I see it, the main contributions of the paper are in 1) identifying the straight-through estimator as an instance of the forward Euler method, which is a first-order approximation of the gradient, and 2) using Heun’s method to derive a second-order gradient approximation. To someone mostly outside of the field of gradient approximation, these are non-trivial insights that, to the best of my knowledge, are unique to this paper.\n\n**Quality**\n\nThe paper is high quality. I particularly appreciated the effort that the authors put into the empirical evaluation, performing hyperparameter sweeps to demonstrate stability, as well as empirically verifying that the ReinMax gradient estimator provides improved estimates of the true gradient. Overall, the result is a paper that provides compelling evidence of 1) improved understanding of straight-through and 2) an improved gradient estimator.\n\n**Clarity**\n\nOverall, the paper is quite clear. The authors do an excellent job of presenting mathematical notation to help the reader see the similarities across gradient derivations and estimators. Theoretical results are presented in a clear, logical order. The empirical results are also generally presented well, with clear labeling of tables and plots.\n\n**Significance**\n\nWhile many papers have mentioned that the straight-through estimator is an approximation of the gradient, it appears that none of these papers have formally shown it. If, indeed, this paper is the first to do so, then it is a significant contribution to our theoretical understanding. The proposed improved estimator, ReinMax, also appears to offer some performance improvement over previously proposed (first-order) estimators. Considering that ReinMax has similar computational overhead as straight-through, then this could serve as a drop-in replacement for the straight-through estimator throughout various applications. However, it is unclear to me whether this serves as a drop-in replacement for all existing instances of the straight-through estimator, or merely those that operate on Multinomial distributions.", "weaknesses": "I see two relatively minor weaknesses in the paper as-is: larger-scale empirical evaluation and slight improvements to the presentation. These are discussed below.\n\nThe current empirical evaluation involves three settings: quadratic programming, structured output prediction, and latent variable generative modeling with discrete latent variables. While the authors demonstrate ReinMax in all three settings, much of the empirical evaluation revolves around the final setting (categorical VAE on MNIST). These results are a useful indicator of the benefits of ReinMax, however, the empirical setting itself is rather toy-ish compared with modern settings. Given that ReinMax is a drop-in replacement for the straight-through estimator, I would imagine that it should be trivial to replace the ST estimator in an existing scaled-up setting, e.g., vector-quantized VAEs or Hafner et al.’s discrete world models. The authors could alternatively / additionally explore categorical VAEs on more complex data. This would allow the authors to more definitively claim that their proposed gradient estimator leads to tangible empirical improvements.\n\nAdditionally, several minor aspects of the presentation could be improved.\n* The plots in Figure 1 are presented with fairly minimal context in the surrounding text. The caption states “Details are elaborated in Section 6.” Then it may make sense to place this figure closer to Section 6.\n* As far as I can tell, the name “ReinMax” is never actually explained.\n* The labels for the baseline methods, e.g., in Figure 1, 4, etc. are difficult to read; they are various dashed lines. Colors and/or larger lines would make this clearer. \n* The shaded regions in Figure 5 (right) are rather jagged — perhaps there’s an issue with the evaluation interval or the plotting setup.\n* In various tables, e.g., Tables 1, 2, …, the results for ReinMax are bolded, despite falling within the error bounds of the baseline methods. I find this to be somewhat misleading.\n* The citation in the second sentence of the introduction is incorrect.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback. We value your comments and will address the concerns regarding experimental design, applicability, and limitations of ReinMax in this rebuttal, with further elaborations to be included in the final paper.\n\n**Reply to Weakness 1 and Question 1:** In our submission, we adhered to the experiment design of the existing study, focusing on small-scale problems for controllability and resource efficiency, as our experiments were mainly conducted on P100/P40 GPUs.\n\nTo further demonstrate the generalizability of ReinMax, we conducted additional experiments (detailed in the general rebuttal), including (1) comparisons with the REINFORCE variant that employs the state-of-the-art variance reduction technology, specifically RODEO (SHI 2022), and (2) applications to a real-world scenario, i.e., differentiable neural architecture search on CIFAR10, CIFAR100, and ImageNet-16-120. \n\nIn the architecture search application, ReinMax consistently improved performance over the baseline in a plug-and-play manner by:\n- Replacing STGS with ReinMax as the gradient estimator,\n- Conducting a minor change to the temperature hyper-parameters (changing the minimal value of the temperature from 0.1 to 1.1), as guided by our findings (refer to Sections 5 and 6.2).\n\n**Reply to limitation 1:** While our main text focuses on multinomial distributions, by reparameterizing other categorical distributions into a multinomial distribution, ReinMax can be generally applied to all categorical random variables. \n\nIt is worth mentioning that ReinMax is not applicable to continuous random variables. While extending ReinMax to these cases is possible, it may not be necessary, given that many commonly used distributions can be reparameterized. For example, the normal distribution $z \\sim \\mathcal{N}(\\mu, \\sigma)$ can be re-written as $z = \\mu + \\sigma \\cdot \\mathcal{N} (0, 1)$, making it trivial to compute $\\partial z/\\partial \\mu$ and $\\partial z/\\partial \\sigma$. \n\nWe hope our responses have adequately addressed your concerns and further highlighted the innovations and potential impact of our study. If you have any further questions or need additional information, please do not hesitate to ask.\n"}, {"review_id": "3ZX0AoxMim", "reviewer": "Reviewer_WPn8", "summary": "This paper is about designing a new approach to approximating the gradient of parameters in generating discrete random variables.\nThe work first starts with discovering the connection of Straight-Through (ST) estimator and the first-order approximation of the true gradient.\nFrom the insight, the authors aim to improve the ST estimator, by applying second-order approximation of the true gradient. To apply second-order approximation without actually calculating second-order derivatives, the authors use Heun's Method, which approximates second-order derivatives using two first-order derivatives; thus no expensive calculation is needed.\nThe proposed estimator is coined ReinMax.\nThrough mathematical analyses, the effectiveness of using the expected value of function outputs is proven.\nFrom evaluations, the authors empirically prove ReinMax method outperforms other estimators, along with presenting other properties and insights e.g. sensitivity to the number of dimensions, batch size, convergence speed, memory usage, and running time.", "questions": "- In experiments using the most basic REINFORCE algorithm, not RLOO or DisARM-Tree, is there any baseline subtraction used? If not can the $E[f(I_i)]$ baseline be used also in REINFORCE?", "limitations": "As mentioned in the checklist, I suppose no potential negative societal impact will arise by this work. Experiments are done quite fairly, including standard deviations of scores and implementation code.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The paper is well-organized and easy to follow\n- Mathematical background is given thoroughly\n- The method can be applied with minimal change of code", "weaknesses": "- Equation 6 seems to be one of core findings from the paper; the derivation process may be written more comprehensibly (like in Appendix A)\n- More evaluations on benchmark datasets close to real-world distribution would be beneficial. For instance, how much improvement will be made when we apply the ReinMax estimator into other models for NLI or sentiment analysis instead of ListOps?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback. We value your comments and will address the concerns regarding experimental design in this rebuttal, with further elaborations to be included in the final paper.\n\n**Reply to weakness 1:** Thanks for the suggestions! We will add more elaborations in the revision. \n\n**Reply to weakness 2:** In our submission, we adhered to the experiment design of the existing study, focusing on small-scale problems for controllability and resource efficiency, as our experiments were mainly conducted on P100/P40 GPUs.\n\nTo further demonstrate the generalizability of ReinMax, we conducted additional experiments (detailed in the general rebuttal), including (1) comparisons with the REINFORCE variant that employs the state-of-the-art variance reduction technology, specifically RODEO (SHI 2022), and (2) applications to a real-world scenario, i.e., differentiable neural architecture search on CIFAR10, CIFAR100, and ImageNet-16-120. \n\nReinMax maintained outstanding performance throughout these expanded tests, showing consistent improvements over the baseline.\n\n**Reply to question 1:** In the general rebuttal, we detailed comparisons between ReinMax and RODEO, a REINFORCE variant employing a state-of-the-art variance reduction method. RODEO outperforms ReinMax on simple scenarios (e.g., large batch size, small number of latent variables, Setting A). Meanwhile, ReinMax achieves better performance on complex scenarios (e.g., small batch size, large number of latent variables, VAE, and Setting B). We will add more elaborations and make corresponding revisions in the final version. \n\nWe hope our responses have adequately addressed your concerns and further highlighted the innovations and potential impact of our study. If you have any further questions or need additional information, please do not hesitate to ask.\n"}, {"review_id": "nbYk2IQg3v", "reviewer": "Reviewer_ZCX1", "summary": "This work develops a new method, ReinMax, to compute gradients of parameters\nused to generate discrete random variables; specifically, the parameters of a\nmultinomial's softmax parameterization. It provides a new perspective onto the\nexisting straight-through (ST) gradient estimator as a first-order\napproximation. The new method is a second-order approximation and, in contrast\nto ST, relies on gradient differences rather than a single gradient. The ReinMax\ngradient estimator is compared to other approaches on polynomial programming toy\nproblems, unsupervised parsing on ListOps, and ELBO training of a VAE on MNIST.\n", "questions": "- Why is the method called ReinMax?\n\n- ReinMax is a second-order extension of ST. Can this be extended to even higher\n  orders? Would this be beneficial or are there diminishing returns in including\n  higher orders?\n\n- What is $\\phi_i$ in L150 and what are its constraints?\n", "limitations": "See  (W1)", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- (S1) The paper provides a new perspective onto the Straight-Through (ST) technique,\n  which simply replaces the backpropagation through a non-differentiable\n  operation by applying the identity. This adds a new theoretical justification\n  why such heuristics work. Framing ST as first-order approximation also naturally\n  suggests an extension to second-order, as presented by the paper.\n\n- (S2) The new method does not seem to add significant overhead to existing\n  approaches. The analysis includes a sensitivity analysis over different\n  problem hyper-parameters which suggests that the estimation works\n  consistently.\n", "weaknesses": "- (W1) Limited to softmax parameterization of multinomial: The main text focuses\n  on computing gradients for the parameters of a multinomial parameterized by a\n  softmax. The authors should comment on the (im)possibility to treat other\n  parameterizations or distributions.\n\n- (W2) Experiments: In the presented experiments, ReinMax seems to consistently\n  work well and perform comparably or more favourably than the competitors.\n  However, I do not have the expertise to judge whether the experiments\n  represent state-of-the-art tasks. I am willing to adapt my score and\n  confidence during the discussion phase with the authors and other reviewers.\n\n- (W3) Clarity: Some steps could be easier to follow by providing the required\n  mathematical properties. The figures and tables should be moved closer to\n  where they are referenced in the text. Here is a list of actionable\n  suggestions which I think would improve clarity:\n  - In Eq. (4), explicitly add the term $\\partial D / \\partial \\pi$ that is set\n    to the identity in ST\n  - In Eq. (6), mention that $\\sum_i \\partial \\pi_i / \\partial \\theta = \\partial\n    \\sum_i \\pi_i / \\partial \\theta = \\partial 1 / \\partial \\theta = 0$.\n  - I think it would be helpful to move parts of appendix E to the main text to\n    have a more formal description of first- and second-order approximation in\n    the specific context.\n  - Fix '2rd-order' into '2nd-order' in various places in the main text\n    and the appendix.\n  - Mention the derivative of a softmax around L171, that is $\\partial \\pi_i /\n    \\partial \\theta_k = \\pi_k (\\delta_{ik} - \\pi_i)$\n   - Clarify the notation $\\delta_{\\mathbf{D} k}$ in the appendix\n  - Minor suggestions: Use consistent symbol $\\mathcal{R}$ in L51, 'softmax' in\n    caption of Figure 1, remove 'on' in L84, 'computational efficiency' rather\n    than 'computation efficiency', 'approximation' in L185, use bold symbol for\n    $\\theta_i$ in L210 and add $\\in \\mathcal{R}^2$, 'phenomena' rather than\n    'phenomenon' in L218), add 'that' between 'one' and 'uses' in L264\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback. We value your comments and will address the concerns regarding experimental design, presentation, and limitations of ReinMax in this rebuttal, with further elaborations to be included in the final paper.\n\n**Reply to limitation 1:** While our main text focuses on multinomial distributions, by reparameterizing other categorical distributions into a multinomial distribution, ReinMax can be generally applied to all categorical random variables. \n\nIt is worth mentioning that ReinMax is not applicable to continuous random variables. While extending ReinMax to these cases is possible, it may not be necessary, given that many commonly used distributions can be reparameterized. For example, the normal distribution $z \\sim \\mathcal{N}(\\mu, \\sigma)$ can be re-written as $z = \\mu + \\sigma \\cdot \\mathcal{N} (0, 1)$, making it trivial to compute $\\partial z/\\partial \\mu$ and $\\partial z/\\partial \\sigma$. \n\nFurthermore, ReinMax differs from REINFORCE in that ReinMax requires both the $\\frac{\\partial f(D)}{\\partial D}$ and $\\frac{\\partial p(D)}{\\partial \\theta}$, while REINFORCE only requires $\\frac{\\partial p(D)}{\\partial \\theta}$. Thus, REINFORCE can be applied in cases where $f(D)$ is not differentiable.\n\n**Reply to weakness 2:** In our submission, we adhered to the experiment design of the existing study, focusing on small-scale problems for controllability and resource efficiency, as our experiments were mainly conducted on P100/P40 GPUs.\n\nTo further demonstrate the generalizability of ReinMax, we conducted additional experiments (detailed in the general rebuttal), including (1) comparisons with the REINFORCE variant that employs the state-of-the-art variance reduction technology, specifically RODEO (SHI 2022), and (2) applications to a real-world scenario, i.e., differentiable neural architecture search on CIFAR10, CIFAR100, and ImageNet-16-120. \n\nOur findings showed that ReinMax performs strongly and consistently in these settings. We will add more discussions and elaborations about these experiments in the final version of the paper.\n\n**Reply to weakness 3:** We appreciate your suggestions for improving clarity, and we will incorporate them into our revision.\n\n**Reply to question 1:** We used the paper title and the abstract as the prompt to query ChatGPT for several names, among which we selected ReinMax since it concisely implies the method's connection to softmax and REINFORCE, and we believe it is a fitting and appealing term. \n\n**Reply to question 2** Although it's possible to apply higher-order ODE solvers, they require more gradient evaluations, leading to undesirable computational overhead. To illustrate this point:\n- The approximation used by ReinMax (as described in Definition 3.2) requires N gradient evaluations, i.e., $\\\\{\\\\frac{\\\\partial f(I_i)}{\\\\partial I_i}\\\\}$.\n- In contrast, the approximation derived by RK4 needs $N^2+N$ gradient evaluations, i.e., $\\\\{\\\\frac{\\\\partial f(I_i)}{\\\\partial I_i}\\\\}$ and $\\\\{\\\\frac{\\\\partial f(I_{ij})}{\\\\partial I_{ij}}\\\\}$, where $I_{ij} = \\\\frac{I_i + I_j}{2}$. \n\nTherefore, while higher-order solvers are applicable, they may not be suitable in our case.\n\n**Reply to question 3** $\\phi$ is a distribution over $\\{I_1, \\cdots, I_n\\}$, i.e., $\\sum_i \\phi_i = 1$ and $\\phi_i = P(I_i)$.\n\nWe hope our responses have adequately addressed your concerns and further highlighted the innovations and potential impact of our study. If you have any further questions or need additional information, please do not hesitate to ask.\n"}, {"review_id": "inVKWRqDbv", "reviewer": "Reviewer_txhR", "summary": "The paper proves that the straight-through (ST) gradient estimator for the categorical distribution can be viewed as a first-order approaximation of the true gradients. Based on this point of view, a new gradient estimator, ReinMax is proposed based on Heun's method which is a second-order method. Experiments are conducted which show that ReinMax improves upon the state of the art methods.", "questions": "- What is the effect of temperature scaling on the proposed method? How does it influence the relative performance of ST-type methods when compared with unbiased REINFORCE-based methods?\n- Is there a complexity-efficiency trade-off here? For example, would using 4th-order Runge-Kutta instead of Heun's method yield much better gradient estimations and improve the overall result?", "limitations": "The authors do not explicitly discuss the limitations of their work.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- **The paper has good writing quality.** The background is introduced in a well-organized manner and contributions are stated clearly. Mathematical formulae are also well laided out and are easy to understand. Relationship with existing work (ST, REINFORCE, baseline subtraction) are also expanded upon in good detail and nuance.\n\n- **Solid contribution with a comprehensive set of experiments.** The proposed new method has solid theoretical underpinnings (Heun's method being more accurate than Euler's method), is still very efficient (not requiring the Hessian), and is proven to be effective under multiple different settings (polynomial programming, MNIST, etc.). The new method is compared against many commonly used methods (STGS, DisARM, ...) and there is also discussion of using a different baseline for subtraction. Extensive hyperparameter tuning is performed and advantage of the proposed method is demonstrated well.", "weaknesses": "My main concern with this paper is that it **oversells its theoretical contribution**. The view of the straight-through estimator as a first-order Taylor expansion of the unbiased gradients is not news to people who study gradient estimators. The authors are fair in pointing out that previous work (Tokui & Sato 2017) only dealt with the Bernoulli case, but one might argue that The categorical distribution is just a natural extension of the Bernoulli case. Regardless, I think it is **potentially misleading to say that \"ST as a first-order approximation to the gradient\" is a novel perspective**. Despite this, I still consider the proposed approach to be novel enough that the contribution of the paper as a whole is still significant.\n\nAnother concern is the lack of comparison with RODEO (Shi 2022) which as far as I know is the state of the art for REINFORCE methods with variance reduction. The paper only reported numbers on RLOO and DisARM in table 3 but these methods are known to be weaker than RODEO, which is reported to outperform DisARM significantly in terms of gradient variance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback. We value your comments and will address the concerns regarding the novelty of our study and experiment design, with further elaborations to be included in the final paper.\n\n**Reply to weakness argument 1:**  While it may be intuitive to some that the straight-through estimator functions as a gradient approximation, no prior work has formally established this for the general multinomial case. \n\nAlso, the derivation of Theorem 3.1, while not overly complicated, is more than a mere expansion of existing results: \n\n- In Tokui & Sato (2017), the authors positioned $\\hat{\\nabla}_{ST}$ as a first-order approximation, but their analysis is exclusively rooted in the properties of Bernoulli variables. As an example, let us consider a Bernoulli random variable $D \\\\in \\\\{ I_1, I_2 \\\\}$. Their approach depends on the property that $\\\\nabla = (f(I_2) - f(I_1)) \\frac{d\\\\pi_1}{d \\\\theta} = (f(I_1) - f(I_2)) \\\\frac{d\\\\pi_2}{d \\\\theta}$, and thus is not applicable to multinomial variables.\n\n- On the other hand, the analyses in Gregor et al. (2014) and Pervez et al. (2020) are applicable to multinomial variables but resort to adding the term $\\\\frac{1}{n \\\\cdot \\\\pi_D}$ in $\\\\hat{\\\\nabla}_{ST}$, an alteration that we believe could induce unwanted instability. This concern is discussed in Section 4.1 and Section 6.4. \n\nIn the revision, we will modify our claim as “our study is the first to formally establish $\\\\hat{\\\\nabla}_{ST}$ works as a first-order approximation in the general multinomial case\".\n\n**Reply to weakness argument 2:**\nIn the general rebuttal, we provide additional experimental results on  (1) comparisons with the REINFORCE variant that employs the state-of-the-art variance reduction technology, specifically RODEO (SHI 2022), and (2) applications to a real-world scenario, i.e., differentiable neural architecture search on CIFAR10, CIFAR100, and ImageNet-16-120. \n\nReinMax maintained outstanding performance throughout these expanded tests, showing consistent improvements over the baseline.\n\n\n**Reply to question 1:**\nIn our experiments, we observed that temperature scaling enhances the stability of both ST and ReinMax algorithms. We conjecture that this scaling acts as a variance reduction method. \n\nAcross all six VAE settings discussed in the general rebuttal, ReinMax achieved the best performance when the temperature was set to 1.1. Adjusting the temperature to 1.0 or 1.2 would lead to a slight performance drop (with an ELBO difference of approximately 1), while ReinMax still outperformed RODEO with these sub-optimal temperatures. \n\n**Reply to question 2.1**\nAs discussed in Section 6.4, the computation overhead of ReinMax is negligible. \n\n**Reply to question 2.2**\nAlthough it's possible to apply higher-order ODE solvers, they require more gradient evaluations, leading to undesirable computational overhead. To illustrate this point:\n- The approximation used by ReinMax (as described in Definition 3.2) requires N gradient evaluations, i.e., $\\\\{\\\\frac{\\\\partial f(I_i)}{\\\\partial I_i}\\\\}$.\n- In contrast, the approximation derived by RK4 needs $N^2+N$ gradient evaluations, i.e., $\\\\{\\\\frac{\\\\partial f(I_i)}{\\\\partial I_i}\\\\}$ and $\\\\{\\\\frac{\\\\partial f(I_{ij})}{\\\\partial I_{ij}}\\\\}$, where $I_{ij} = \\\\frac{I_i + I_j}{2}$. \n\nTherefore, while higher-order solvers are applicable, they may not be suitable in our case.\n\nWe hope our responses have adequately addressed your concerns and further highlighted the innovations and potential impact of our study. If you have any further questions or need additional information, please do not hesitate to ask.\n"}, {"review_id": "qqR8rzNh3Q", "reviewer": "Reviewer_fihA", "summary": "This paper works on gradient approximation for discrete latent variables. The problem is challenging because the discreteness hinders direct backpropagation through the neural networks; hence, an approximation is needed. The authors address the issues from the perspective of a second-order approximation of the gradient. Specifically, the authors proposed ReinMax and showed that ReinMax achieved second-order accuracy with negligible computation overheads. The experiments on Polynomial Programming, MNIST-VAE, and ListOps demonstrated the superiority of ReinMax over the prior methods such as Straight-Through Gumbel-Softmax, Gumbel-Rao Monte Carlo, and Gapped Straight-Through.", "questions": "Why the proposed method is called ReinMax?\n\nFor the other questions, see the weakness section.", "limitations": "No potential negative social impact.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The problem has been challenging for a long time. Most prior work addresses the issue from the perspective of Straight-Through Gumbel Softmax (STGS). To the reader's best knowledge, the perspective of second-order approximation is novel in this field and the performance improvement is believable. Below highlights the strengths of the paper.\n\na) The use of second-order approximation is novel and the theoretical derivation is believable. \n\nb) The performance gain is clear with similar computational overhead as in the prior work.\n\nc) The experimental analysis covers some relevant topics such as temperature scaling and the choice of baseline. The conclusion is clear from the analysis.", "weaknesses": "a) The design choice of ReinMax is unclear. \n\nAlthough ReinMax is well-motivated by the second-order approximation, the reader might be confused by the form of ReinMax at Eq. (7). For example, why mixing $\\pi$ and $D$ by introducing $\\pi_D=(\\pi+D)/2$? Why substracts by $\\nabla_{ST}$? Where are the coefficients (2 and 1/2) from? It would be nice if the authors can elaborate more on these.\n\nb) The implication of the experiments remains unclear.\n\nFor a fair comparison with the prior work, the authors follow the literature and run ReinMax on Polynomial Programming, MNIST-VAE, and ListOps. However, these are toy problems with limited implications for the applications. For example, does the good performance of ReinMax on Polynomial Program/MNIST-VAE/ListOps imply any possible applications? Conversely, if we have a reinforcement learning problem with discrete actions or a VQ-VAE model, can ReinMax be helpful? It would be nice if the authors can evaluate more on the application side so that the other practitioners can follow.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback. We value your comments and will address the concerns regarding experimental design and presentation in this rebuttal, with further elaborations to be included in the final paper.\n\n**Reply to weakness 1:** We understand the complexity of ReinMax's exact form may not be immediately intuitive, as it is a result of specific derivations. To alleviate confusion, we will dedicate additional space in the final version to elaborate on the exact form of ReinMax. \n\n**Reply to weakness 2:** In our submission, we adhered to the experiment design of the existing study, focusing on small-scale problems for controllability and resource efficiency, as our experiments were mainly conducted on P100/P40 GPUs.\n\nTo further demonstrate the generalizability of ReinMax, we conducted additional experiments (detailed in the general rebuttal), including (1) comparisons with the REINFORCE variant that employs the state-of-the-art variance reduction technology, specifically RODEO (SHI 2022), and (2) applications to a real-world scenario, i.e., differentiable neural architecture search on CIFAR10, CIFAR100, and ImageNet-16-120. \n\nReinMax maintained outstanding performance throughout these expanded tests, showing consistent improvements over the baseline. These additional experiments should provide a more comprehensive understanding of ReinMax's potential applications.\n\n**Reply to question 1:** We used the paper title and the abstract as the prompt and queried ChatGPT for several names, among which we selected ReinMax, since it concisely implies the method's connection to softmax and REINFORCE, and we believe it is a fitting and appealing term. \n\nWe hope our responses have adequately addressed your concerns and further highlighted the innovations and potential impact of our study. If you have any further questions or need additional information, please do not hesitate to ask.\n"}, {"review_id": "MRl1pw5PM5", "reviewer": "Reviewer_A6oc", "summary": "The presented paper tackles the optimization of the parameter of the distribution of a discrete random variable through stochastic gradient descent by addressing the problem of gradient estimation given samples from the discrete distribution.\nThe authors consider three popular estimators as a starting point for their investigation, namely, the reinforce estimator, the straight-through estimator and the straight-through gumbel-softmax estimator.\n\nThe authors first establish that the straight-through estimator is a first-order approximation of the true gradient, since the difference $f(I_i) - f(I_j)$ in the true gradient is replaced by $\\frac{\\partial f(I_j)}{\\partial I_j}(I_i - I_j)$.\nThe author thus replace the $\\frac{\\partial f(I_j)}{\\partial I_j}(I_i - I_j)$, which gives a first-order approximation of $f(I_i) - f(I_j)$, by $\\frac{1}{2} (\\frac{\\partial f(I_j)}{\\partial I_j} + \\frac{\\partial f(I_i)}{\\partial I_i})(I_i - I_j)$, which yield a second order approximation.\n\nThe authors then empirically evaluates both the baselines and the proposed estimator on a synthetic task (polynomial programming), unsupervised modeling (unsupervised sequence parsing), and structured output prediction (generative modeling).\n\n", "questions": "It is said that the derivation of equation (8) leverages the derivative of the softmax function, thus there may be a handy formula about these derivative that I am not aware of and which justify all the equations that do not seem valid to me. Could the authors elaborate on my concern about the validity of their equations ?\n\nI highly doubt that the equations are wrong, given that they coincide with previous work on the topic and the consistent improvements provided in the experiment section. I would be willing to increase my score if the author correctly addresses my concerns.", "limitations": "There is a decent discussion about the limitation of the proposed method througout the paper.\nAlbeit, given that high variance is the main limiting factor of the REINFORCE estimator, and given that the other estimators are biased, a focus on bias variance tradeoff would have been welcome.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The theoretical analysis is well written and intuitively explained.\nThe interpretation of the straight-through estimator as a first-order approximation of the true gradient allows the author to use numerical analysis results to improve the theoretical behavior of their proposed solution.\n- The experiment section is decent, with qualitatively different setups addressing different application domains.\n- The choice of the baseline for the baseline subtraction method as well as the tuning of the temperature parameter is appropriately discussed.", "weaknesses": "- Equation (6) does not seem to be valid. I would expect $\\sum_i (f(I_i) - E[f(D)]) \\frac{d \\pi_i}{d\\theta} = \\sum_i \\sum_j \\pi_j(f(I_i) - f(I_j)) \\frac{d \\pi_i}{d\\theta}$, thus I don't understand why the term $\\sum_i E[f(D)] \\frac{d \\pi_i}{d \\theta}$ vanishes.\n- The derivation of remark 4.1 as provided in appendix C does not seem to be valid. I don't understand how\n$\\sum_i \\sum_j \\pi_j \\frac{\\partial f(I_j)}{\\partial I_j}(I_i - I_j)) \\frac{d \\pi_i}{d\\theta} = \\sum_i \\frac{\\phi_D}{\\pi_D} \\pi_D \\frac{\\partial f(I_j)}{\\partial I_j} \\sum_i I_i \\frac{d \\pi_i}{d\\theta}$. Furthermore, there should be a typo at line 155 since it should be $\\pi_D$ who is the output of the softmax and could take very small values. This would be consistent with the fact that it is the denominator of the fraction $\\frac{\\phi_D}{\\pi_D}$.\n- I have no idea how equation (8) has been derived.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback. We value your comments and will address your concerns regarding the correctness of the derivation, with further elaborations to be included in the final paper.\n\n\n**Reply to weakness argument 1:** Since $\\sum_i \\pi_i = 1$, we have $\\sum_i E[f(D)] \\frac{d \\pi_i}{d \\theta} = E[f(D)] \\cdot \\frac{d\\sum_i \\pi_i }{d\\theta} = E[f(D)] \\cdot \\frac{d 1}{d\\theta} = 0$\nIn the revision, we will add a simple explanation of this in Eq. (6).\n\n**Reply to weakness argument 2:** Thanks for pointing out the typo, and we will fix the typo and add more elaborations in the revision. Please find the detailed derivation of the last-second equation in Appendix C as below. \n\n\n$\\sum_i \\sum_j \\phi_j \\frac{\\partial f(I_j)}{\\partial I_j} (I_i - I_j) \\frac{d \\pi_i}{d \\theta}$\n\n$= \\sum_j \\Large(\\normalsize \\phi_j\\frac{\\partial f(I_j)}{\\partial I_j} \\sum_i (I_i - I_j) \\frac{d \\pi_i}{d \\theta} \\Large)\\normalsize$ \n\n$= \\sum_j \\Large(\\normalsize \\phi_j\\frac{\\partial f(I_j)}{\\partial I_j} (\\sum_i I_i \\frac{d \\pi_i}{d \\theta} - \\sum_i I_j \\frac{d \\pi_i}{d \\theta} )\\Large)\\normalsize$\n\n$= \\sum_j \\Large(\\normalsize \\phi_j\\frac{\\partial f(I_j)}{\\partial I_j} (\\sum_i I_i \\frac{d \\pi_i}{d \\theta} - I_j \\frac{d \\sum_i \\pi_i}{d \\theta} )\\Large)\\normalsize$\n\n$= \\sum_j \\Large(\\normalsize \\phi_j\\frac{\\partial f(I_j)}{\\partial I_j} (\\sum_i I_i \\frac{d \\pi_i}{d \\theta} - I_j \\frac{d 1}{d \\theta}) \\Large)\\normalsize$\n\n$= \\sum_j \\phi_j\\frac{\\partial f(I_j)}{\\partial I_j} \\sum_i I_i \\frac{d \\pi_i}{d \\theta}$\n\n$= \\sum_j \\frac{\\phi_j}{\\pi_j}\\cdot \\pi_j \\cdot \\frac{\\partial f(I_j)}{\\partial I_j} \\sum_i I_i \\frac{d\\pi_i}{d\\theta}$.\n\n\n**Reply to weakness argument 3 and question 1:**\nIn the revision, we will mention the derivative of the softmax around L171, i.e., for $\\pi = \\mbox{softmax}(\\theta)$, we have $\\partial \\pi_i / \\partial \\theta_k = \\pi_k (\\delta_{ik} - \\pi_i)$. Please find the detailed derivation of equation (8) as below. \n\n$\\frac{\\partial \\mathcal{L}}{\\partial \\theta_k}= \\frac{\\partial \\sum_i \\pi_i f(I_i) }{\\partial \\theta_k}$\n\n$=\\sum_i f(I_i) \\frac{d \\pi_i}{d \\theta_k}$\n\n$= \\sum_i f(I_i) \\pi_k (\\delta_{ik} - \\pi_i)$\n\n$= \\sum_i f(I_i) \\pi_k \\delta_{ik} - \\sum_i f(I_i) \\pi_k \\pi_i$\n\n$= f(I_k) \\pi_k - \\sum_i f(I_i) \\pi_k \\pi_i$\n\n$= f(I_k) \\pi_k \\sum_i \\pi_i - \\sum_i f(I_i) \\pi_k \\pi_i$\n\n$= \\pi_k \\sum_i \\pi_i (f(I_k) - f(I_i))$.\n\nWe hope our responses have adequately addressed your concerns and further highlighted the innovations and potential impact of our study. If you have any further questions or need additional information, please do not hesitate to ask.\n"}]

Author rebuttal:
# General Rebuttal

We thank all reviewers for their thoughtful feedback. In this work, we tackle critical challenges in gradient estimation for discrete variables, and our contributions are notable in two primary areas:
- we formally established that the straight-through estimator is a first-order approximation of the gradient, 
- we proposed ReinMax, offering a second-order approximation with negligible computational overhead and consistent performance improvements.

In this general rebuttal, we provide additional experiment results on:
- Comparisons with the REINFORCE variant that employs the state-of-the-art variance reduction technology, namely RODEO (SHI 2022).
- Applications of ReinMax on real-world problems, specifically differentiable neural architecture search on CIFAR10, CIFAR100, and ImageNet-16-120.

## Comparisons with RODEO

### Bernoulli VAEs 

We utilized ReinMax to train Bernoulli VAEs on MNIST, Fashion-MNIST, and Omniglot, adhering closely to the experimental settings of RODEO (SHI et al., 2022),  including pre-processing, model architecture, batch size, and training epochs. As summarized in Tables A and B, ReinMax consistently outperforms RODEO across all settings.

**Train A: -ELBO of 2 x 200 VAE on MNIST, Fashion-MNIST, and Omniglot when K=3 (i.e., three evaluations per image). \* Baseline results are referenced from SHI et al. (2022).**
||ARMS$^*$|DoubleCV$^*$|RODEO$^*$|RELAX$^*$|ReinMax|
|-|-|-|-|-|-|
|MNIST |100.84±0.14|100.94±0.09|100.46±0.13|101.99±0.04|97.83±0.36|
|Fashion-MNIST|237.05±0.12|237.40±0.11|236.88±0.12|237.74±0.12|234.53±0.42|
|Omniglot|115.32±0.07|115.06±0.12|115.01±0.05|115.70±0.08|107.51±0.42|

**Train B: -ELBO of 2 x 200 VAE on MNIST, Fashion-MNIST, and Omniglot when K=2 (i.e., two evaluations per image). \* Baseline results are referenced from SHI et al. (2022).**

||DisARM$^*$|Double CV$^*$|RODEO$^*$|ReinMax|
|-|-|-|-|-|
|MNIST |102.75±0.08|102.14±0.06|101.89±0.17|98.05±0.29|
|Fashion-MNIST|237.68±0.13|237.55±0.16|237.44±0.09|234.86±0.33|
|Omniglot|116.50±0.04|116.39±0.10|115.93±0.06|107.79±0.27|

### Polynomial Programming

To better understand the difference between RODEO and ReinMax, we conduct more experiments on polynomial programming, i.e., $\min_\theta E_{X} \Large[\normalsize \frac{\|X - c\|_p^p}{L}\Large]\normalsize$. Specifically, we consider polynomial programming under two different settings that define $c$ differently (the difference between these two settings is elaborated at the end of the part):
- Setting A: $c = [0.45, \cdots, 0.45]$. This is the setting we used in the submission. 
- Setting B: $c = [\frac{0.5}{L}, \frac{1.5}{L}, \cdots, \frac{L-0.5}{L}]$. 

We visualized the training curve of polynomial programming in the attached pdf (Figure 13), together with the training curve of Bernoulli VAE (Figure 12). ReinMax achieves better performance in more challenging scenarios, i.e., smaller batch size, more latent variables, or more complicated problems (Setting B or VAEs). Meanwhile, REINFORCE and RODEO achieve better performance on simpler problem settings, i.e., larger batch size, fewer latent variables, or simpler problems (Setting A). This observation matches our intuition:
- REIFORCE-style algorithms excel as they provide unbiased gradient estimation but may fall short in complex scenarios, since they only utilize the zero-order information, i.e., a scalar $f(\cdot)$ for each training instance. 
- ReinMax, using more information (i.e., a vector $\frac{\partial f(D)}{\partial D}$ for each training instance), handles challenging scenarios better. Meanwhile, as a consequence of its estimation bias, ReinMax leads to slower convergence in some simple scenarios. 

As to the difference between the Setting A and the Setting B, we would like to note: 
- In Setting A, since $\forall i, c_i=0.45$ and $\theta_i\sim Uniform(-0.01, 0.01)$ at initialization, $E_{X_i \sim \mbox{softmax}(\theta_i)}\Large[\normalsize \frac{\|X_i - c_i\|_p^p}{L}\Large]\normalsize$ would have similar values. Therefore, the optimal control variates for $\theta_i$ are similar across different $i$. 
- In Setting B, we set $c_i$ to different values for different $i$, and thus the optimal control variate for $\theta_i$ are different across different $i$. 

Therefore, Setting A is a simpler setting for applying control variate to REINFORCE. 

## Differentiable Neural Architecture Search 

We demonstrate the applicability of ReinMax as a drop-in replacement in differentiable neural architecture search. 

GDAS (Dong & Yang, 2019) is an algorithm that employs STGS to estimate the gradient of neural architecture parameters with a temperature schedule (decaying linearly from 10 to 0.1). We replaced STGS with ReinMax as the gradient approximator and changed the minimal temperature from 0.1 to 1.1(as discussed in Section 5 and Section 6.2, temperature scaling plays a different role in ReinMax). 

We evaluate the resulting algorithm with the official implementation under the topology search setting in the NATS-Bench benchmark (Dong et al., 2020), and summarize the results in Table C as below. ReinMax brings consistent performance improvements over the baseline across all three datasets, demonstrating the great potential of ReinMax. We will add more analyses and discussions in the revision. 

**Train C: Performance in NATS-Bench.\* Baseline results are referenced from Dong et al. (2020).**

||CIFAR-10 DEV|CIFAR-10 TEST|CIFAR-100 DEV|CIFAR-100 TEST|ImageNet-16-120|
|-|-|-|-|-|-|
|GDAS-Straight-Through Gumbel-Softmax$^*$|89.68±0.72|93.23±0.58|68.35±2.71|68.17±2.50|39.55±0.00|
|GDAS-ReiMax|89.92±0.27|93.47±0.35|69.40±1.63|69.61±1.71|41.11±2.09|

Dong, X. and Yang, Y. Searching for a robust neural architecture in four GPU hours. *CVPR*, 2019.

Dong, X., Liu, L., Musial, K., and Gabrys, B. Nats-bench: Benchmarking nas algorithms for 307 architecture topology and size. *TPAMI*, 2020.

Meta-review: All six reviewers agreed this paper should be accepted: it is original, high quality, clear, and provides (to the best of my knowledge) the first formal result showing the straight-through estimator is a first-order approximation in the multinomial case. Given that ReinMax has similar computational complexity to straight-through and improves performance it will likely gain traction as a drop-in replacement. This is a clear accept. Authors: you've already indicated that you've updated the submission to respond to reviewer changes, if you could double check their comments for any recommendation you may have missed on accident that would be great! The paper will make a great contribution to the conference!
Decision: Accept (oral)