SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Liyuan Liu Chengyu Dong Xiaodong Liu Bin Yu Jianfeng Gao Microsoft Research
{lucliu, v-chedong, xiaodl, v-ybi, jfgao}@microsoft.com
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.
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.
Figure 1: Training curves of polynomial programming, i.e.,  $\min_{\pmb{\theta}} E_{\pmb{X}}[\| \pmb{X} - \pmb{c}\|_p^p / 128]$ , where  $\pmb{\theta} \in \mathbb{R}^{128 \times 2}$ ,  $\pmb{X} \in \{0,1\}^{128}$ , and  $\pmb{X}_i \stackrel{\mathrm{iid}}{\sim}$  Multinomial(softmax( $\theta_i$ )). Details are elaborated in Section 6.
Algorithm 1: ST.
Input:  $\theta$  : softmax input,  $\tau$  : temperature.
Output:  $D$ : one-hot samples.
1  $\pi_0\gets \mathrm{softmax}(\pmb {\theta})$
2  $D\gets$  sample_one-hot  $(\pi_0)$
3  $\pi_1\gets \mathrm{softmax}_{\tau}(\pmb {\theta})$
/\*stop_gradien(·) duplicates its input and detaches it from backpropagation. \*/
4  $D \gets \pi_1 - \text{stop\_gradient}(\pi_1) + D$
5 return  $D$
Algorithm 2: ReinMax.
1  $\pi_0\gets \mathrm{softmax}(\pmb {\theta})$
2  $D\gets$  sample_one HOT(  $\pi_0$
3  $\pi_1\gets \frac{D + \mathrm{softmax}_\tau(\theta)}{2}$
4  $\pi_1\gets$  softmax(stop_gradient  $(\ln (\pi_1) - \theta) + \theta)$
5  $\pi_2\gets 2\cdot \pi_1 - \frac{1}{2}\cdot \pi_0$
6  $D\gets \pi_{2}$  -stop_gradient  $(\pi_2) + D$
7 return  $D$
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  $\tau \geq 1$ , and Straight-Through Gumbel-Softmax (STGS; Jang et al., 2017) prefers to set the temperature  $\tau \leq 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<sup>1</sup>.
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  $\mathrm{softmax}_{\tau}(\theta)_i = \frac{\exp(\theta_i / \tau)}{\sum_{j=1}^n \exp(\theta_j / \tau)}$ , where  $n$  is the number of possible outcomes,  $\theta \in \mathcal{R}^{n \times 1}$  is the parameter, and  $\tau$  is the temperature<sup>2</sup>. For  $i \in [1, \dots, n]$ , we mark its one-hot representation as  $I_i \in \mathcal{R}^{n \times 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 \in \{I_1, \dots, I_n\}$ , we assume the distribution of  $D$  is parameterized as:  $p(D = I_i) = \pi_i = \mathrm{softmax}(\theta)_i$ , and mark  $\mathrm{softmax}_{\tau}(\theta)$  as  $\pi^{(\tau)}$ . Given a differentiable function  $f: \mathcal{R}^n \to \mathcal{R}$ , we aim to minimize (note that temperature scaling is not used in the generation of  $D$ ):
$$
\min  _ {\boldsymbol {\theta}} \mathcal {L} (\boldsymbol {\theta}), \text {w h e r e} \mathcal {L} (\boldsymbol {\theta}) = E _ {\boldsymbol {D} \sim \operatorname {s o f t m a x} (\boldsymbol {\theta})} [ f (\boldsymbol {D}) ]. \tag {1}
$$
Here, we mark the gradient of  $\theta$  as  $\nabla$ :
$$
\nabla := \frac {\partial \mathcal {L} (\boldsymbol {\theta})}{\partial \boldsymbol {\theta}} = \sum_ {i} f \left(\boldsymbol {I} _ {i}\right) \frac {d \pi_ {i}}{d \boldsymbol {\theta}}. \tag {2}
$$
1 Implementations are available at https://github.com/microsoft/ReinMax.
Without specification, the temperature (i.e.,  $\tau$ ) is set to 1.
In many applications, it is usually too costly to compute  $\nabla$ , since it requires the computation of  $\{f(I_1), \dots, f(I_n)\}$  and evaluating  $f(I_i)$  is costly for typical deep learning applications. Correspondingly, many efforts have been made to estimate  $\nabla$  efficiently.
The  $\nabla_{\mathrm{REINFORCE}}$  (Williams, 1992) is unbiased (i.e.,  $E[\nabla_{\mathrm{REINFORCE}}] = \nabla$ ) and only requires the distribution of the discrete variable to be differentiable (i.e., no backpropagation through  $f$ ):
$$
\nabla_ {\text {R E I N F O R C E}} := f (\boldsymbol {D}) \frac {d \log p (\boldsymbol {D})}{d \boldsymbol {\theta}}. \tag {3}
$$
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 Straight-Through (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<sup>3</sup>:
$$
\widehat {\nabla} _ {\mathrm {S T}} := \frac {\partial f (\boldsymbol {D})}{\partial \boldsymbol {D}} \cdot \frac {d \boldsymbol {\pi}}{d \boldsymbol {\theta}}. \tag {4}
$$
In practice,  $\widehat{\nabla}_{\mathrm{ST}}$  is usually implemented with the tempered softmax, under the hope that the temperature hyper-parameter  $\tau$  may be able to reduce the bias introduced by  $\widehat{\nabla}_{\mathrm{ST}}$  (Chung et al., 2017).
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):
$$
\boldsymbol {D} = \lim  _ {\tau \rightarrow 0} \operatorname {s o f t m a x} _ {\tau} (\boldsymbol {\theta} + \boldsymbol {G}) \text {w h e r e} \boldsymbol {G} _ {i} \text {a r e i . i . d . a n d} \boldsymbol {G} _ {i} \sim \operatorname {G u m b e l} (0, 1).
$$
STGS treats the zero-temperature limit as identity function during the backpropagation:
$$
\widehat {\nabla} _ {\mathrm {S T G S}} := \frac {\partial f (\boldsymbol {D})}{\partial \boldsymbol {D}} \cdot \frac {d \operatorname {s o f t m a x} _ {\tau} (\boldsymbol {\theta} + \boldsymbol {G})}{d \boldsymbol {\theta}}. \tag {5}
$$
Both  $\widehat{\nabla}_{\mathrm{ST}}$  and  $\widehat{\nabla}_{\mathrm{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[\widehat{\nabla}_{\mathrm{ST}}]$  or  $E[\widehat{\nabla}_{\mathrm{STGS}}]$  can be viewed as an approximation of  $\nabla$ . Thus, we aim to answer the following questions: How  $\widehat{\nabla}_{\mathrm{ST}}$  approximates  $\nabla$  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. In this study, we leverage these methods to approximate  $\nabla$  with the gradient of  $f$ . To begin, we demonstrate that ST works as a first-order approximation of  $\nabla$ . Then, we propose ReinMax, which integrates Heun's method for a better gradient approximation and achieves second-order accuracy.
3We use the notation  $\widehat{\nabla}$  to indicate gradient approximations. Note that the generation of  $D$  is not differentiable, and  $\widehat{\nabla}_{\mathrm{ST}}$  does not have the term  $\partial D / \partial \pi$ .
3.1 Straight-Through as a First-order Approximation
We start by defining a first-order approximation of  $\nabla$  as  $\widehat{\nabla}_{1\mathrm{st - order}}$ .
Definition 3.1. One first-order approximation of  $\nabla$  is  $\widehat{\nabla}_{Ist - order} \coloneqq \sum_{i}\sum_{j}\pi_{j}\frac{\partial f(I_{j})}{\partial I_{j}} (I_{i} - I_{j})\frac{d\pi_{i}}{d\theta}$ .
To understand why  $\widehat{\nabla}_{1\mathrm{st - order}}$  is a first-order approximation, we rewrite  $\nabla$  in Equation 2 as:
$$
\nabla = \sum_ {i} \left(f \left(\boldsymbol {I} _ {i}\right) - E [ f (\boldsymbol {D}) ]\right) \frac {d \pi_ {i}}{d \boldsymbol {\theta}} + \sum_ {i} E [ f (\boldsymbol {D}) ] \frac {d \pi_ {i}}{d \boldsymbol {\theta}} = \sum_ {i} \sum_ {j} \pi_ {j} \left(f \left(\boldsymbol {I} _ {i}\right) - f \left(\boldsymbol {I} _ {j}\right)\right) \frac {d \pi_ {i}}{d \boldsymbol {\theta}}. \tag {6}
$$
Comparing  $\widehat{\nabla}_{1\mathrm{st - order}}$  and Equation 6, it is easy to notice that  $\widehat{\nabla}_{1\mathrm{st - order}}$  approximates  $f(I_{i}) - f(I_{j})$  as  $\frac{\partial f(I_j)}{\partial 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  $\widehat{\nabla}_{1\mathrm{st - order}}$  is a first-order approximation of  $\nabla$ .
Now, we proceed to show  $\widehat{\nabla}_{\mathrm{ST}}$  works as a first-order approximation. Note that our analyses only apply to  $\widehat{\nabla}_{\mathrm{ST}}$  as defined in Equation 4 and may not apply to its other variants.
Theorem 3.1.
$$
E [ \widehat {\nabla} _ {S T} ] = \widehat {\nabla} _ {1 s t - o r d e r}.
$$
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.,  $\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}$ ), 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  $\widehat{\nabla}_{\mathrm{ST}}$  as  $\frac{1}{n \cdot \pi_D}\widehat{\nabla}_{\mathrm{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  $\widehat{\nabla}_{\mathrm{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  $\widehat{\nabla}_{\mathrm{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  $\nabla$  is
$$
\widehat {\nabla} _ {2 n d - o r d e r} := \sum_ {i} \sum_ {j} \frac {\pi_ {j}}{2} \left(\frac {\partial f (\boldsymbol {I} _ {j})}{\partial \boldsymbol {I} _ {j}} + \frac {\partial f (\boldsymbol {I} _ {i})}{\partial \boldsymbol {I} _ {i}}\right) \left(\boldsymbol {I} _ {i} - \boldsymbol {I} _ {j}\right) \frac {d \pi_ {i}}{d \boldsymbol {\theta}}.
$$
Comparing  $\widehat{\nabla}_{2\mathrm{nd - order}}$  and Equation 6, we can observe that,  $\widehat{\nabla}_{2\mathrm{nd - order}}$  approximates  $f(I_i) - f(I_j)$  as  $\frac{1}{2} (\frac{\partial f(I_i)}{\partial I_i} +\frac{\partial f(I_j)}{\partial I_j})(I_i - I_j)$ . This approximation is known as the Heun's Method and has second-order accuracy (we provide a brief introduction to Heun's method in Appendix E). Correspondingly, we know that  $\widehat{\nabla}_{2\mathrm{nd - order}}$  is a second-order approximation of  $\nabla$ .
Based on this approximation, we propose the ReinMax operator as  $(\pi_D$  refers to  $\frac{\pi + D}{2}$ ,  $|$  refers to the identity matrix, and  $\odot$  refers to the element-wise product):
$$
\widehat {\nabla} _ {\text {R e i n M a x}} := 2 \cdot \widehat {\nabla} ^ {\frac {\pi + D}{2}} - \frac {1}{2} \widehat {\nabla} _ {\mathrm {S T}}, \text {w h e r e} \widehat {\nabla} ^ {\frac {\pi + D}{2}} = \frac {\partial f (\boldsymbol {D})}{\partial \boldsymbol {D}} \cdot \left(\left(\boldsymbol {\pi} _ {\boldsymbol {D}} \cdot 1 ^ {T}\right) \odot \mathbf {I} - \boldsymbol {\pi} _ {\boldsymbol {D}} \cdot \boldsymbol {\pi} _ {\boldsymbol {D}} ^ {T}\right) \tag {7}
$$
4Please note that  $\sum_{i}E[f(\pmb {D})]\frac{d\pi_{i}}{d\theta} = E[f(\pmb {D})]\frac{d\sum_{i}\pi_{i}}{d\theta} = E[f(\pmb {D})]\frac{d\mathbf{1}}{d\theta} = 0$
Then, we show that  $\widehat{\nabla}_{\mathrm{ReinMax}}$  approximates  $\nabla$  to the second order. Or, formally we have:
Theorem 3.2.
$$
E [ \widehat {\nabla} _ {R e i n M a x} ] = \widehat {\nabla} _ {2 n d - o r d e r}.
$$
The proof of Theorem 3.2 is provided in Appendix B.
Computation Efficiency of ReinMax. Instead of requiring Hessian or other second-order derivatives,  $\widehat{\nabla}_{\mathrm{ReinMax}}$  achieves second-order accuracy with two first-order derivatives (i.e.,  $\frac{\partial f(I_j)}{\partial I_j}$  and  $\frac{\partial f(I_i)}{\partial I_i}$ ). As observed in our empirical efficiency comparisons in Section 6, the computation overhead of  $\widehat{\nabla}_{\mathrm{ReinMax}}$  is negligible. At the same time, similar to  $\widehat{\nabla}_{\mathrm{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.,  $\{\frac{\partial f(I_i)}{\partial I_i}\}$ . 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}$ . 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(\mathbf{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.
Here, we discuss the rationale for choosing  $E[f(D)]$  as the baseline. Considering  $\sum_{i} \phi_{i} f(I_{i})$  as the general form of the baseline ( $\phi_{i}$  is a distribution over  $\{I_1, \dots, I_n\}$ , i.e.,  $\sum_{i} \phi_{i} = 1$ ), we have:
Remark 4.1. When  $\sum_{i}\phi_{i}f(I_{i})$  is used as the baseline and  $f(I_{i}) - f(I_{j})$  is approximated as  $\frac{\partial f(I_j)}{\partial I_j} (I_i - I_j)$ , we mark the resulting first-order approximation of  $\nabla$  as  $\widehat{\nabla}_{1st}$ -order-avg-baseline. Then, we have  $E[\frac{\phi_D}{\pi_D}\widehat{\nabla}_{ST}] = \widehat{\nabla}_{1st}$ -order-avg-baseline.
The derivations of Remark 4.1 are provided in Appendix C. Intuitively, since  $\pi_D$  is the output of the softmax function, it could have very small values, which makes  $\frac{\phi_D}{\pi_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  $\frac{\phi_D}{\pi_D}$  brought.
It is worth mentioning that, when setting  $\phi$  as  $\frac{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  $\nabla$  as  $\operatorname{E}[\frac{1}{n\cdot\pi_D}\widehat{\nabla}_{\mathrm{ST}}]$ , which matches the result of Remark 4.1 by setting  $\phi = \frac{1}{n}$ .
In Section 6, we compared the corresponding second-order approximation when treating  $E[f(\pmb{D})]$  and  $\frac{1}{n}\sum_{i}f(\pmb{I}_{i})$  as the baseline, respectively. We observed that gradient estimators that use  $E[f(\pmb{D})]$  as the baseline consistently outperform gradient estimators that use  $\frac{1}{n}\sum_{i}f(\pmb{I}_{i})$  as the baseline, which verifies our intuition and demonstrates the importance of the baseline selection.
Figure 2: Training -ELBO on MNIST-VAE (lighter color indicates better performance). STGS, GST-1.0, and GR-MCK prefer to set the temperature  $\tau \leq 1$ . ST and ReinMax prefer to set  $\tau \geq 1$ .
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  $\frac{\partial\mathcal{L}}{\partial\theta_k}$  and mark it as  $\nabla_{k}$ . Similar to Equation 2, we have:
$$
\nabla_ {k} := \frac {\partial \mathcal {L}}{\partial \boldsymbol {\theta} _ {k}} = \sum_ {i} f (\boldsymbol {I} _ {i}) \frac {d \boldsymbol {\pi} _ {i}}{d \boldsymbol {\theta} _ {k}} = \boldsymbol {\pi} _ {k} \sum_ {i} \boldsymbol {\pi} _ {i} (f (\boldsymbol {I} _ {k}) - f (\boldsymbol {I} _ {i})). \tag {8}
$$
It is worth mentioning that the derivation of Equation 8 leverages the derivative of the softmax function (i.e., for  $\pi = \mathrm{softmax}(\theta)$ , we have  $\partial \pi_{i} / \partial \theta_{k} = \pi_{k}(\delta_{ik} - \pi_{i})$ ) and does not involve the baseline subtraction technology.
Remark 4.2. In Equation 8, we approximate  $f(\mathbf{I}_k) - f(\mathbf{I}_i)$  as  $\frac{1}{2} (\frac{\partial f(\mathbf{I}_i)}{\partial \mathbf{I}_i} + \frac{\partial f(\mathbf{I}_k)}{\partial \mathbf{I}_k})(\mathbf{I}_k - \mathbf{I}_i)$ , and mark the resulting second-order approximation of  $\nabla_k$  as  $\widehat{\nabla}_{2nd-order-wo-baseline,k} = \pi_k \sum_i \pi_i \frac{1}{2} (\frac{\partial f(\mathbf{I}_i)}{\partial \mathbf{I}_i} + \frac{\partial f(\mathbf{I}_k)}{\partial \mathbf{I}_k})(\mathbf{I}_k - \mathbf{I}_i)$ . Then, we have  $E[\widehat{\nabla}_{\text{ReinMax}}] = \widehat{\nabla}_{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  $\tau$  to small values for STGS, we show that ST and ReinMax need a different strategy.