Title: Bridging Discrete and Backpropagation: Straight-Through and Beyond
Venue: NeurIPS 2023
Accepted: yes
Keywords: discrete random variables

back-propagation

straight through
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-Through and Beyond

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.

1 Introduction

There has been a persistent pursuit to build neural network models with discrete or sparse variables (Neal, 1992). However, backpropagation (Rumelhari et al., 1986), the cornerstone of deep learning, is restricted to computing gradients for continuous variables. Correspondingly, many attempts have been made to approximate the gradient of parameters that are used to generate discrete variables, and most of them are based on the Straight-Through (ST) technique (Bengio et al., 2013).

The development of ST is based on the simple intuition that non-differentiable functions (e.g., sampling of discrete latent variables) can be approximated with the identity function in the backpropagation (Rosenblatt, 1957; Bengio et al., 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 hyper-parameter configurations or future algorithm development. Thus, researchers have to develop different ST variants for different applications in a trial-and-error manner, which is laborious and time-consuming (van den Oord et al., 2017; Liu et al., 2019; Fedus et al., 2021). To address these limitations, we aim to explore how ST approximates the gradient and how it can be improved.

E

0

5

10

15

20

25

30

35

40

Epoch

0.31

0.32

0.33

0.34

0.35

0.36

Loss

p = 1.5

0

5

10

15

20

25

30

35

40

Epoch

0.21

0.22

0.23

0.24

0.25

Loss

p = 2

0

10

20

30

40

Epoch

0.090

0.095

0.100

0.105

0.110

0.115

0.120

0.125

Loss

p = 3

ReinMax

STGS

GR-MCK

ST

GST-1.0

min

X

[

X

c

p

p

128

]

|

|

First, we adopt a novel perspective to examine ST and show that it works as a 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 hyper-parameters of ST and its variants, i.e., ST prefers to set the temperature τ ≥ 1 , and StraightThrough Gumbel-Softmax (STGS; Jang et al., 2017) prefers to set the temperature τ ≤ 1 .

Our analyses not only shed insights on the underlying mechanism of ST but also lead us to develop a 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 al. (2017); Grathwohl et al. (2018); Pervez et al. (2020); Paulus et al. (2021), unsupervised generative modeling (Kingma &Welling, 2013), structured output prediction (Nangia & Bowman, 2018), and differentiable neural architecture search (Dong et al., 2020a) to demonstrate that ReinMax brings consistent improvements over the state of the art 1 .

Our contributions are two-fold:

We formally establish that ST works as a first-order approximation to the gradient in the general multinomial case, which provides valuable guidance for future research and applications.

We propose a 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 to 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 to sigmoid belief network and Helmholtz machines (Williams, 1992; Dayan et al., 1995). To keep things straightforward, we will focus on a simplified scenario. We refer to the tempered softmax as softmax τ ( θ ) i = exp( θ i /τ ) ∑ n j =1 exp( θ j /τ ) , where n is the number of possible outcomes, θ ∈ R n × 1 is the parameter, and τ is the temperature 2 . For i ∈ [1 , · · · , n ] , we mark its one-hot representation as I i ∈ R n × 1 , whose element equals 1 if it is the i -th element or equals 0 otherwise. Let D be a discrete random variable and D ∈ { I 1 , · · · , I n } , we assume the distribution of D is parameterized as: p ( D = I i ) = π i = softmax ( θ ) i , and mark softmax τ ( θ ) as π ( τ ) . Given a differentiable function f : R n →R , we aim to minimize (note that temperature scaling is not used in the generation of D ):

Here, we mark the gradient of θ as ∇ :

In many applications, it is usually too costly to compute ∇ , since it requires the computation of { f ( I 1 ) , · · · , f ( I n ) } and evaluating f ( I i ) is costly for typical deep learning applications. Correspondingly, many efforts have been made to estimate ∇ efficiently.

The ∇ REINFORCE (Williams, 1992) is unbiased (i.e., E [ ∇ REINFORCE ] = ∇ ) and only requires the distribution of the discrete variable to be differentiable (i.e., no backpropagation through f ):

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 al., 2016; Tucker et al., 2017; Grathwohl et al., 2018; Shi et al., 2022). Still, it has been found that the REINFORCE-style estimators fail to work well in many real-world applications. Empirical comparisons between ReinMax and REINFORCE-style methods are elaborated in Section 6.5.

Efficient Gradient Approximation. In practice, a popular family of estimators is StraightThrough (ST) estimators. They compute the backpropagation "through" a surrogate that treats the non-differentiable function (e.g., the sampling of D ) as an identity function. The idea of ST originates from the perceptron algorithm (Rosenblatt, 1957; Mullin & Rosenblatt, 1962), which leverages a modified chain rule and utilizes the identity function as the proxy of the original derivative of a binary output function. Bengio et al. (2013) improves this method by using non-linear functions like sigmoid or softmax, and Jang et al. (2017) further incorporates the Gumbel reparameterization. Here, we briefly describe Straight-Through (ST) and Straight-Through Gumbel-Softmax (STGS).

In the general multinomial distribution case, as in Algorithm 1, the ST estimator treats the sampling process of D as an identity function during the backpropagation 3 :

̂ The STGS estimator is built upon the Gumbel re-parameterization trick (Maddison et al., 2014; Jang et al., 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 (Gumbel, 1954):

In practice, ̂ ∇ ST is usually implemented with the tempered softmax, under the hope that the temperature hyper-parameter τ may be able to reduce the bias introduced by ∇ ST (Chung et al., 2017).

STGS treats the zero-temperature limit as identity function during the backpropagation:

3 Discrete Variable Gradient Approximation: a Numerical ODE Perspective

Both ̂ ∇ ST and ̂ ∇ 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 E [ ̂ ∇ ST ] or E [ ̂ ∇ STGS ] can be viewed as an approximation of ∇ . Thus, we aim to answer the following questions: How ̂ ∇ ST approximates ∇ and how it can be improved?

In numerical analysis, extensive studies have been conducted to develop numerical methods for solving ordinary differential equations. In this study, we leverage these methods to approximate ∇ with the gradient of f . To begin, we demonstrate that ST works as a first-order approximation of ∇ . Then, we propose ReinMax, which integrates Heun's method for a better gradient approximation and achieves second-order accuracy.

3.1 Straight-Through as a First-order Approximation

We start by defining a first-order approximation of ∇ as ̂ ∇ 1st-order . Definition 3.1. One first-order approximation of ∇ is ̂ ∇ 1st-order := ∑ i ∑ j π j ∂f ( I j ) ∂ I j ( I i -I j ) d π i d θ .

Comparing ̂ ∇ 1st-order and Equation 6, it is easy to notice that ̂ ∇ 1st-order approximates f ( I i ) -f ( I j ) as ∂f ( I j ) ∂ I j ( I i -I j ) . In numerical analyses, this approximation is known as the forward Euler method, which has first-order accuracy (we provide a brief introduction to the forward Euler method in Appendix E). Correspondingly, we know that ∇ 1st-order is a first-order approximation of ∇ .

̂ Now, we proceed to show ̂ ∇ ST works as a first-order approximation. Note that our analyses only apply to ̂ ∇ ST as defined in Equation 4 and may not apply to its other variants. Theorem 3.1.

The proof of Theorem 3.1 is provided in Appendix A.

It is worth mentioning that Tokui & Sato (2017) discussed this connection for the special case of D being a Bernoulli variable. However, their study is built upon a Bernoulli variable property (i.e., ∇ = ( f ( I 2 ) -f ( I 1 )) d π 1 dθ = ( f ( I 1 ) -f ( I 2 )) d π 2 dθ ), making their analyses not applicable to multinomial variables. Alternatively, the analyses in Gregor et al. (2014) and Pervez et al. (2020) are applicable to multinomial variables but resort to modify ̂ ∇ ST as 1 n · π D ̂ ∇ ST, in order to position it as a first-order approximation. We suggest that this modification would lead to unwanted instability and provide more discussions in Section 4.1 and Section 6.6. Here, our study is the first to formally established ∇ ST works as a first-order approximation in the general multinomial case.

̂ Besides revealing the mechanism of the Straight-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 to integrate a better approximation for f ( I i ) -f ( I j ) .

3.2 Towards Second-order Accuracy: ReinMax

The literature on numerical methods for differential equations shows that it is possible to 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 ∇ is

Comparing ̂ ∇ 2nd-order and Equation 6, we can observe that, ̂ ∇ 2nd-order approximates f ( I i ) -f ( I j ) as 1 2 ( ∂f ( I i ) ∂ I i + ∂f ( I j ) ∂ I j )( I i -I j ) . This approximation is known as the Heun's Method and has secondorder accuracy (we provide a brief introduction to Heun's method in Appendix E). Correspondingly, we know that ∇ 2nd-order is a second-order approximation of ∇ .

̂ Based on this approximation, we propose the ReinMax operator as ( π D refers to π + D 2 , I refers to the identity matrix, and ⊙ refers to the element-wise product):

Then, we show that ̂ ∇ ReinMax approximates ∇ to the second order. Or, formally we have: Theorem 3.2.

The proof of Theorem 3.2 is provided in Appendix B.

Computation Efficiency of ReinMax. Instead of requiring Hessian or other second-order derivatives, ̂ ∇ ReinMax achieves second-order accuracy with two first-order derivatives (i.e., ∂f ( I j ) ∂ I j and ∂f ( I i ) ∂ I i ). As observed in our empirical efficiency comparisons in Section 6, the computation overhead of ̂ ∇ ReinMax is negligible. At the same time, similar to ̂ ∇ ST (as in Algorithm 1), our proposed algorithm can be easily integrated with existing automatic differentiation toolkits like PyTorch (a simple implementation of ReinMax is provided in Algorithm 2), making it easy to be integrated with existing algorithms.

Applicability of Higher-order ODE solvers. Although it's possible to apply higher-order ODE solvers, they require more gradient evaluations, leading to undesirable computational overhead. To illustrate this point: The approximation used by ReinMax requires n gradient evaluations, i.e., { ∂f ( I i ) ∂ I i } . In contrast, the approximation derived by RK4 needs n 2 + n gradient evaluations, i.e., { ∂f ( I i ) ∂ I i } and { ∂f ( I ij ) ∂ I ij } , where I ij = I i + I j 2 . Therefore, while higher-order solvers are applicable, they may not be suitable in our case.

4 ReinMax and Baseline Subtraction

Equation 6 plays a crucial role in positioning ST as a 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, we first discuss the reason for choosing E [ f ( D )] as the baseline, and then reveal that the derivation of ReinMax is independent to baseline subtraction.

4.1 Benefits of Choosing E [ f ( D )] as the Baseline

The choice of baseline in reinforcement learning has been the subject of numerous discussions (Weaver & Tao, 2001; Rennie et al., 2016; Shi et al., 2022). Similarly, in our study, different baselines lead to different gradient approximations.

Remark 4.1. When ∑ i ϕ i f ( I i ) is used as the baseline and f ( I i ) -f ( I j ) is approximated as ∂f ( I j ) ∂ I j ( I i -I j ) , we mark the resulting first-order approximation of ∇ as ̂ ∇ 1st-order-avg-baseline . Then, we have E [ ϕ D π D ̂ ∇ ST ] = ̂ ∇ 1st-order-avg-baseline .

Here, we discuss the rationale for choosing E [ f ( D )] as the baseline. Considering ∑ i ϕ i f ( I i ) as the general form of the baseline ( ϕ i is a distribution over { I 1 , · · · , I n } , i.e., ∑ i ϕ i = 1 ), we have:

The derivations of Remark 4.1 are provided in Appendix C. Intuitively, since π D is the output of the softmax function, it could have very small values, which makes ϕ D π D to be unreasonably large and leads to undesired instability. Therefore, we suggest that E [ f ( D )] is a better choice of baseline when it comes to gradient approximation, since its corresponding gradient approximation is free of the instability ϕ D π D brought.

It is worth mentioning that, when setting ϕ as 1 n , the result of Remark 4.1 echoes some existing studies. Specifically, both Gregor et al. (2014) and Pervez et al. (2020) propose to approximate ∇ as E[ 1 n · π D ∇ ST], which matches the result of Remark 4.1 by setting ϕ = 1 n .

̂ In Section 6, we compared the corresponding second-order approximation when treating E [ f ( D )] and 1 n ∑ i f ( I i ) as the baseline, respectively. We observed that gradient estimators that use E [ f ( D )] as the baseline consistently outperform gradient estimators that use 1 n ∑ i f ( I i ) as the baseline, which verifies our intuition and demonstrates the importance of the baseline selection.

Method =STGS

Method = GST-1.0

Method = GR-MCK

Method = ST

0.1

206

206

206

202

0.3

0.5

Temperature

L'0

1

2

3

5

4×248×1616x1264×8

4×248×16 16×1264×8

4x248×1616x1264x8

4×24 8×1616×12 64×8

4×248×16 16×1264×8

Categorical #xLatent #

Categorical#xLatent#

Categorical#xLatent#

Categorical#xLatent#

Categorical #xLatent #

100

110

120

100

110

120

130

100

110

120

120

140

160

180

100

120

140

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 ∂ L ∂ θ k and mark it as ∇ k . Similar to Equation 2, we have:

It is worth mentioning that the derivation of Equation 8 leverages the derivative of the softmax function (i.e., for π = softmax ( θ ) , we have ∂ π i /∂ θ k = π k ( δ ik -π i ) ) and does not involve the baseline subtraction technology.

Remark 4.2. In Equation 8, we approximate f ( I k ) -f ( I i ) as 1 2 ( ∂f ( I i ) ∂ I i + ∂f ( I k ) ∂ I k )( I k -I i ) , and mark the resulting second-order approximation of ∇ k as ̂ ∇ 2nd-order-wo-baseline ,k = π k ∑ i π i 1 2 ( ∂f ( I i ) ∂ I i + ∂f ( I k ) ∂ I k )( I k -I i ) , Then, we have E [ ̂ ∇ ReinMax ] = ̂ ∇ 2nd-order-wo-baseline The proof of Remark 4.2 is provided in Appendix D.

As in Remark 4.2, applying the Heun's method on Equation 8 and Equation 6 lead to the same gradient estimator, which implies another benefit of using E [ f ( D )] as the baseline: the resulting gradient estimator does not rely on additional prior (i.e., its derivation can be free of baseline subtraction).

5 Temperature Scaling for Gradient Estimators

Here, we discuss how to apply temperature scaling, a technique widely used in gradient estimators, to our proposed method, ReinMax. While the typical practice is to set the temperature τ to small values for STGS, we show that ST and ReinMax need a different strategy.

Temperature Scaling for ∇ STGS. As introduced in Section 2, ∇ STGS conduct a two-step approximation: (1) it approximates min θ E [ f ( D )] as min θ E [ f ( softmax τ ( θ + G )))] ; (2) it approximates ∂f ( softmax τ ( θ + G )) ∂ softmax τ ( θ + G ) as f ( D ) ∂ D . Since the bias introduced in both steps can be controlled by τ , ∇ STGS prefers to set τ as a relatively small value.

Temperature Scaling for ∇ ST and ∇ ReinMax. As in Section 4, it does not involve temperature scaling to show ∇ ST and ∇ ReinMax work as the first-order and the second-order approximation to the gradient. Correspondingly, temperature scaling technology cannot help to reduce the bias for ∇ ST in the same way it does for ∇ STGS. As in Figure 2, STGS, GR-MCK, and GST-1.0 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 τ ( θ ) ) as ̂ ∇ ReinMax = 2 · ̂ ∇ π τ + D 2 -1 2 ̂ ∇ ST . It is worth emphasizing that τ in ̂ ∇ ReinMax is used to stabilize the gradient approximation (instead of reducing bias) at the cost of accuracy. Therefore, the value of τ should be larger or equal to 1 .

6 Experiments

Here, we conduct experiments on polynomial programming, unsupervised generative modeling, and structured output prediction. In all experiments, we consider four major baselines: StraightThrough (ST), Straight-Through Gumbel-Softmax (STGS), Gumbel-Rao Monte Carlo (GR-MCK), and Gapped Straight-Through (GST-1.0). For a more comprehensive comparison, we run a complete grid search on the training hyper-parameters for all methods. Also, we would reference results from the literature when their setting is comparable with ours. More details are elaborated in Appendix F.

6.1 Polynomial Programming

Following previous studies (Tucker et al., 2017; Grathwohl et al., 2018; Pervez et al., 2020; Paulus et al., 2021), we start with a simple problem. Consider L i.i.d. latent binary variables X 1 , · · · , X L ∈ { 0 , 1 } and a constant vector c ∈ R L × 1 , we parameterize the distributions of { X 1 , · · · , X L } with L softmax functions, i.e., X i iid ∼ Multinomial ( softmax ( θ i )) and θ i ∈ R 2 . Following previous studies, we set every dimension of c as 0 . 45 , i.e., ∀ i, c i = 0 . 45 , and use min θ E X [ ∥ X -c ∥ p p L ] 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 is visualized in Figure 1 for p = 1 . 5 , 2 , and 3 . In all cases, 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 a faster convergence speed.

6.2 ListOps

We conducted unsupervised parsing on ListOps (Nangia & Bowman, 2018) and summarized the average accuracy and the standard derivation in Table 1. We also visualized the accuracy and loss on the valid set in Figure 3. Although the ST algorithm performs poorly on polynomial programming, it achieves a reasonable performance on this task. Also, while all baseline methods perform similarly, 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.

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

20.0

Epoch

0.90

0.95

1.00

1.05

1.10

1.15

1.20

1.25

1.30

Loss on Valid Set

ReinMax

STGS

GR-MCK

ST

GST-1.0

-

■-

E

0.9

0.8

V) 0.6

U)

U 0.5

0.4

0.3

0

20

40

60

80

100

120

140

160

ReinMax

STGS

GR-MCK

ST

GST-1.0

0.7

0

Epoch

6.3 MNIST-VAE

We benchmark the performance by training variational auto-encoders (VAE) with categorical latent variables on MNIST (LeCun et al., 1998). As we aim to compare gradient estimators, we focus our discussions on training ELBO. We find that training performance largely mirrors test performance (Dong et al., 2020b, 2021; Fan et al., 2022) and briefly discussed test ELBO in Appendix F.

Biases of the Approximated Gradient. With 4 latent dimensions and 8 categorical dimensions, we iterate through the whole latent space (the size of the latent space is only 4096 ), compute the gradient as in Equation 2, and measured the cosine similarity between the gradient of latent variables and various approximations. As visualized in Figure 4, 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. Let us proceed to larger latent spaces. First, we consider 4 settings with the latent space of 2 48 . Then, following Fan et al. (2022), we also conduct experiments with 10 latent dimensions and 30 categorical dimensions (the size of the latent space is 10 30 ). As summarized in Table 2, ReinMax achieves the best performance on all configurations.

GST-1.0 Performance on Different Problems. It is worth mentioning that, despite GST-1.0 achieving good performance on most settings of MNIST-VAE, it fails to maintain this performance on polynomial programming and unsupervised parsing, as discussed before. Upon discussing with Fan et al. (2022), we suggest that this phenomenon is caused by the characteristic of GST-1.0, which behaves similarly to ST on problems with a near one-hot optimal distribution. In other words, GST1.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 a significant performance boost over GST-1.0 on polynomial programming. However, on MNIST-VAE and ListOps, GST-p achieves an inferior performance.

This observation verifies our intuition that, without understanding the mechanism of ST, different applications have different preferences on its configurations. Meanwhile, ReinMax achieves consistent improvements in all settings, which greatly simplifies future algorithms developments.

RLOO

∗

DisARM-Tree

∗

STGS

∗

GR-MCK

∗

GST-1.0

∗

ST

∗

ReinMax

Neg. ELBO

104.03

±

0.23

103.10

±

0.25

97.32

±

0.20

110.74

±

1.23

96.09

±

0.25

116

±

0.09

93.44

±

0.51

0.0

0.2

0.4

0.6

0.8

1.0

Update

1e6

95

100

105

110

115

120

Loss

ReinMax

RODEO

(a) MNIST when K=2.

0.0

0.2

0.4

0.6

0.8

1.0

Update

1e6

100

105

110

115

120

125

130

135

140

Loss

ReinMax

RODEO

(b) Omniglot when K=2.

0.0

0.2

0.4

0.6

0.8

1.0

Update

1e6

234

236

238

240

242

244

246

248

250

Loss

ReinMax

RODEO

(c) Fashion-MNIST when K=2.

0.0

0.2

0.4

0.6

0.8

1.0

Update

1e6

95

100

105

110

115

120

Loss

ReinMax

RODEO

(d) MNIST when K=3.

0.0

0.2

0.4

0.6

0.8

1.0

Update

1e6

100

105

110

115

120

125

130

135

140

Loss

ReinMax

RODEO

(e) Omniglot when K=3.

0.0

0.2

0.4

0.6

0.8

1.0

Update

1e6

234

236

238

240

242

244

246

248

250

Loss

ReinMax

RODEO

(f) Fashion-MNIST when K=3.

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 NATS-Bench benchmark (Dong et al., 2020a), and summarize the results in Table 4. 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 gradient estimator (configurations elaborated in Appendix F). ReinMax brings consistent performance improvements across all three datasets, demonstrating the great potential of ReinMax.

6.5 Comparisons with REINFORCE-style Methods

Here, we conduct experiments to discuss the difference between ReinMax and REINFORCE-style methods. First, following Fan et al. (2022), we conduct experiments on the setting with a larger batch size (i.e., 200), longer training (i.e., 5 × 10 5 steps), 32 latent dimensions, and 64 categorical dimensions (details are elaborated in Appendix F). As in Table 5, ReinMax outperforms all baselines, including two REINFORCE-based methods (Dong et al., 2020b, 2021).

We further conduct experiments to compare with the state of the art. Specifically we apply ReinMax to Bernoulli VAEs on MNIST, Fashion-MNIST (Xiao et al., 2017), and Omniglot(Lake et al., 2015), adhering closely to the experimental settings of Shi et al. (2022), including pre-processing, model architecture, batch size, and training epochs. As in Tables 6 and Figure 5, ReinMax consistently outperforms RODEO across all settings. To better understand the difference between RODEO and ReinMax, we conduct more experiments on polynomial programming (as elaborated in Appendix F.6).

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 on simpler problem settings, i.e., larger batch size, fewer latent variables, or simpler problems. 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 ( D ) ).

ReinMax, using more information (i.e., a vector ∂f ( D ) ∂ D ), handles challenging scenarios better. Meanwhile, as a consequence of its estimation bias, ReinMax leads to slower convergence in some simple scenarios.

6.6 Discussions

Choice of Baseline. As introduced in Section 4.1, the choice of subtraction baseline has a huge impact on the performance. Here, we demonstrate this empirically.

We use 1 n ∑ i f ( I i ) as the baseline and compare the resulting gradient approximation with ReinMax. As visualized in Figure 6, ReinMax, which uses E [ f ( D )] as the baseline, significantly outperforms the one that uses 1 n ∑ i f ( I i ) as the baseline. We suspect that the gradient approximation using 1 n ∑ i f ( I i ) as the baseline is very unstable as it contains the 1 n · p ( D ) term.

Temperature Scaling. On MNIST-VAE (four settings with the 2 48 latent space), we utilize heatmaps to visualize the final performance of all five methods under different temperatures, i.e., { 0 . 1 , 0 . 3 , 0 . 5 , 0 . 7 , 1 , 2 , 3 , 4 , 5 } . As in Figure 2, these methods have different preferences for the temperature configuration. Specifically, STGS, GST-1.0, and GR-MCK prefer to set the temperature τ ≤ 1 . Differently, ST and ReinMax prefer to set the temperature τ ≥ 1 . These observations match our analyses in Section 5 that a small τ can help reduce the bias introduced by STGS-style methods. Also, it verifies that ST and ReinMax work differently from STGS, GST-1.0, and GR-MCK.

Efficiency. As summarized in Table 3, we can observe that, since GR-MCK uses the Monte Carlo method to reduce the variance, it has larger time and memory consumption, which becomes less significant with fewer Monte Carlo samples (we use GR-MCK s to indicate GR-MCK with s 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-Expert training (Liu et al., 2023). Specifically, for gradient approximation of sparse expert routing, while ReinMax requires the network to be fully activated, Liu et al. (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.

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 or propagating gradients through stochastic neurons for conditional computation. ArXiv , abs/1308.3432, 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 Computation , 7:889-904, 1995.

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 architecture topology and size. TPAMI , 2020a.

Dong, Z., Mnih, A., and Tucker, G. Disarm: An antithetic gradient estimator for binary latent variables. In NeurIPS , 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. I., and Ramadge, P. J. Training discrete deep generative models via gapped straight-through estimator. In ICML , 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 ICML , 2014.

Gu, S. S., Levine, S., Sutskever, I., and Mnih, A. Muprop: Unbiased backpropagation for stochastic neural networks. In ICLR , 2016.

Gumbel, E. J. 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. P. and Ba, J. Adam: A method for stochastic optimization. In ICLR , 2015.

Kingma, D. P. and Welling, M. Auto-encoding variational bayes. CoRR , abs/1312.6114, 2013.

Lake, B. M., Salakhutdinov, R., and Tenenbaum, J. B. Human-level concept learning through probabilistic program induction. Science , 2015.

LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. Gradient-based learning applied to 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 , abs/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/1804.06028, 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 gumbelsoftmax gradient estimator. In ICLR , 2021.

Pervez, A., Cohen, T., and Gavves, E. Low bias low variance gradient estimates for boolean stochastic networks. In ICML , 2020.

Rennie, S. J., Marcheret, E., Mroueh, Y., Ross, J., and Goel, V. Self-critical sequence training for image captioning. In CVPR , 2016.

Rosenblatt, F. The perceptron, a perceiving and recognizing automaton Project Para . Cornell Aeronautical Laboratory, 1957.

Rumelhari, D. E., Hintont, G. E., Ronald, J., and Williams. Learning representations by backpropagating errors. Nature , 323:533536, 1986.

Shi, J., Zhou, Y., Hwang, J., Titsias, M., and Mackey, L. Gradient estimation with discrete stein operators. In NeurIPS , 2022.

Tokui, S. and Sato, I. Evaluating the variance of likelihood-ratio gradient estimators. In ICML , 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-based reinforcement learning. In UAI , 2001.

Williams, R. J. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning , 8:229-256, 1992.

Xiao, H., Rasul, K., and Vollgraf, R. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. ArXiv , abs/1708.07747, 2017.

A Theorem 3.1

Let us define the first-order approximation of ∇ as ̂ ∇ 1st-order = ∑ i ∑ j π j ∂f ( I j ) ∂ I j ( I i -I j ) d π i d θ , which approximates f ( I i ) -f ( I j ) in Equation 6 as ∂f ( I j ) ∂ I j ( I i -I j ) .

Theorem 3.1.

Proof. Based on the definition, we have:

Since ∑ i π i = 1 , we have ∑ i d π i d θ = 0 . Also, since π = ∑ i π i I i , we have d π d θ = ∑ i I i d π i d θ . Thus, together with Equation 9, we have:

B Theorem 3.2

Theorem 3.2.

Proof. Here, we aim to proof, ∀ k ∈ [1 , n ] , we have E [ ̂ ∇ ReinMax ,k ] = ̂ ∇ 2nd-order ,k . As defined in Equation 8, we have (note that δ DI k is the indicator function of the event D = I k ):

Since δ DI k ( I k -∑ i π i I i ) = δ DI k ( I D -∑ i π i I i ) , together with Equation 10 and 11, we have:

C Remark 4.1

Remark 3.1. When ∑ i ϕ i f ( I i ) is used as the baseline and f ( I i ) -f ( I j ) is approximated as ∂f ( I j ) ∂ I j ( I i -I j ) , we mark the resulting first-order approximation of ∇ as ̂ ∇ 1st-order-avg-baseline . Then, we have:

Proof. Using ∑ i ϕ i f ( I i ) as the baseline, we have:

Approximating f ( I i ) -f ( I j ) as ∂f ( I j ) ∂ I j ( I i -I j ) , we have:

D Remark 4.2

Remark 3.2. In Equation 8, we approximate f ( I k ) -f ( I i ) as 1 2 ( ∂f ( I i ) ∂ I i + ∂f ( I k ) ∂ I k )( I k -I i ) , and mark the resulting second-order approximation of ∇ k as ̂ ∇ 2nd-order-wo-baseline ,k = π k ∑ i π i 1 2 ( ∂f ( I i ) ∂ I i + ∂f ( I k ) ∂ I k )( I k -I i ) , Then, we have:

Proof. Here, we aim to proof, ∀ k ∈ [1 , n ] , we have E [ ̂ ∇ ReinMax ,k ] = ̂ ∇ 2nd-order-wo-baseline ,k .

E Forward Euler Method and Heun's Method

For simplicity, we consider a simple function g ( x ) : R → R that is three times differentiable on [ t 0 , t 1 ] . Now, we proceed to a simple introduction to approximate ∫ t 1 t 0 g ′ ( x ) dx with the Forward Euler 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 1 ) with the first-order Taylor expansion, i.e., g ( t 1 ) = g ( t 0 ) + g ′ ( t 0 ) · ( t 1 -t 0 ) + O (( t 1 -t 0 ) 2 ) , then we have ∫ t 1 t 0 g ′ ( x ) dx ≈ g ′ ( t 0 )( t 1 -t 0 ) . Since we used the first-order Taylor expansion, this approximation has first-order accuracy.

Heun's Method. First, we approximate g ( t 1 ) with the second-order Taylor expansion:

Then, we show that we can match this approximation by combining the first-order derivatives of two samples. Taylor expanding g ′ ( t 1 ) to the first-order, we have:

Therefore, we have:

It is easy to notice that the right-hand side of the above equation matches the second-order Taylor expansion of g ( t 1 ) as in Equation 12. Therefore, the above approximation (i.e., approximating g ( t 1 ) -g ( t 0 ) as g ′ ( t 0 )+ g ′ ( t 1 ) 2 ( t 1 -t 0 ) ) has second-order accuracy.

Connection to f ( I i ) -f ( I j ) in Equation 6. By setting g ( x ) = f ( x · I i +(1 -x ) · I j )) , we have g (1) -g (0) = f ( I i ) -f ( I j ) . Then, it is easy to notice that the forward Euler Method approximates f ( I i ) -f ( I j ) as ∂f ( I j ) ∂ I j ( I i -I j ) and has first-order accuracy. Also, the Heun's Method approximates f ( I i ) -f ( I j ) as 1 2 ( ∂f ( I i ) ∂ I i + ∂f ( I j ) ∂ I j )( I i -I j ) and has second-order accuracy.

F Experiment Details

F.1 Baselines

Here, we consider four methods as our major baselines:

Straight-Through (ST; Bengio et al., 2013) backpropagate through the sampling function as if it had been the identity function.

Straight-Through Gumbel-Softmax (STGS; Jang et al., 2017) integrates the Gumbel reparameterization trick to approximate the gradient.

Gumbel-Rao Monte Carlo (GR-MCK; Paulus et al., 2021) 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 to 1000 in most experiments. Except in our discussions of efficiency, we set the number of Monte Carlo samples to 100, 300, and 1000 for a more comprehensive comparisons.

Gapped Straight-Through (GST-1.0; Fan et al., 2022) aims to reduce the variance of STGS and constructs a deterministic term to replace the Monte Carlo samples used in GR-MCK. Here, as suggested in (Fan et al., 2022), we set the gap (a hyper-parameter) as 1 . 0 .

GST-1.0 Performance. Despite GST-1.0 achieving good performance on most settings of MNISTVAE, it fails to maintain this performance on polynomial programming and unsupervised parsing, as discussed before. At the same time, a different variant of GST (i.e., GST-p) achieves a significant performance boost over GST-1.0 on polynomial programming. However, on MNIST-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 applications have different preferences on its configurations. Meanwhile, ReinMax achieves consistent improvements in all settings, which greatly simplifies future algorithms developments.

F.2 Hyper-Parameters

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 on ListOps). The hyper-parameter search space is summarized in Table 7. The search results for Table 2 and Table 1 are summarized in Table 8.

Polynomial Programming. As this problem is relatively simple, we set the learning rate to 0.001 and the optimizer to Adam, and only tune the temperature hyper-parameter.

MNIST-VAE. Following the previous study (Dong et al., 2020b, 2021; Fan et al., 2022), we used 2-layer MLP as the encoder and the decoder. We set 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 32 latent dimensions and 64 categorical dimensions, we set the batch size to 200, training steps to 5 × 10 5 , and activation function to LeakyReLU, in order to be consistent with the literature. For other experiments, we set the batch size to 100, the activation function to ReLU, and training steps to 9 . 6 × 10 4 (i.e., 160 epochs).

Differentiable Neural Architecture Search. We adopt most of the hyper-parameter setting from Dong et al. (2020a). Since GDAS employs a temperature schedule (decaying linearly from 10 to 0.1), and temperature scaling works differently in ReinMax and STGS (as discussed in Section 5 and Section 6.6), we removed the temperature scaling (i.e., set the temperature to a constant 1 . 0 ) and increased the weight decay (i.e., from 0 . 001 to 0 . 09 ).

ListOps. Wefollowed the same setting of Fan et al. (2022), i.e., used the same model configuration as in Choi et al. (2017) and set the maximum sequence length to 100 .

F.3 Hardware and Environment Setting

Most experiments (except efficiency comparisons) are conducted on Nvidia P40 GPUs. For efficiency comparisons, we measured the average time cost per batch and peak memory consumption on quadratic programming and MNIST-VAE on the same system with an idle A6000 GPU. Also, to better reflect the efficiency of gradient estimators, we skipped all parameter updates in efficiency comparisons.

F.4 Additional Results on Polynomial Programming

We visualized the training curve for polynomial programming with various batch sizes and latent dimensions in Figure 8 (for p = 1 . 5 ), Figure 9 (for p = 2 ), and Figure 10 (for p = 3 ).

F.5 Additional Results on MNIST-VAE

In our discussions in Section 6, we focused on the training ELBO only. Here, we provide a brief discussion on the test ELBO.

Choosing Hyper-parameter Based on Training Performance. Similar to Table 2, for each method, we select the hyper-parameter based on its training performance. The Test -ELBO in this setting is summarized in Table 9. Despite the model being trained without dropout or other overfitting reduction techniques, ReinMax maintained the best performance in this setting.

(a)

p

= 3

and

c

= [

0

.

5

L

,

1

.

5

L

,

· · ·

,

L

-

0

.

5

L

]

.

(b)

p

= 3

and

c

= [0

.

45

,

· · ·

,

0

.

45]

.

(c)

p

= 2

and

c

= [

0

.

5

L

,

1

.

5

L

,

· · ·

,

L

-

0

.

5

L

]

.

(d)

p

= 2

and

c

= [0

.

45

,

· · ·

,

0

.

45]

.

(e)

p

= 1

.

5

and

c

= [

0

.

5

,

1

.

5

,

,

L

-

0

.

5

]

.

(f)

p

= 1

.

5

and

c

= [0

.

45

,

,

0

.

45]

.

L

L

· · ·

L

· · ·

Choosing Hyper-parameter Based on Test Performance. We also conduct experiments by selecting hyper-parameters directly based on their test performance. In this setting, the test -ELBO is summarized in Table 10, and the training -ELBO is summarized in Table 11. ReinMax achieves the best performance in all settings except the test performance of the setting with 10 categorical dimensions and 30 latent dimensions.

F.6 More Comparisons with RODEO

To better understand the difference between RODEO and ReinMax, we conduct more experiments on polynomial programming, i.e., min θ E X [ | X -c | p p L ] . Specifically, we consider polynomial programming under two different settings that define c differently:

In setting A, we have c = [0 . 45 , · · · , 0 . 45] . This is the setting we used in the submission.

In setting B, we have c = [ 0 . 5 L , 1 . 5 L , · · · , L -0 . 5 L ] .

As to the difference between the Setting A and the Setting B, we would like to note:

In setting A, since ∀ i, c i = 0 . 45 and θ i ∼ Uniform ( -0 . 01 , 0 . 01) at initialization, E X i ∼ softmax ( θ i ) [ | X i -c i | p p L ] would have similar values. Therefore, the optimal control variates for θ 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 θ i are different across different i . Therefore, Setting A is a simpler setting for applying control variate to REINFORCE.

As in Figure 7, 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, REINFORCEand RODEOachieve better performance on simpler problem settings, i.e., larger batch size, fewer latent variables, or simpler problems (Setting A).

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nIt 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.\n\nLIMITATIONS:\nThere 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\na) 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.\n\nQUESTIONS:\nWhy the proposed method is called ReinMax?\n\nFor the other questions, see the weakness section.\n\nLIMITATIONS:\nNo potential negative social impact.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- **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.\n\nWEAKNESSES:\nMy 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe authors do not explicitly discuss the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n- (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\nWEAKNESSES:\n- (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\nQUESTIONS:\n- 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\nLIMITATIONS:\nSee  (W1)", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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\n\nWEAKNESSES:\n- 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?\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nAs 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.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-20"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n**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.\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\nCould 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”?\n\nLIMITATIONS:\nAs mentioned above, I would appreciate a clearer discussion of which existing forms of straight-through estimator the ReinMax estimator can replace.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-20"}