SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Temperature Scaling for  $\nabla_{\mathrm{STGS}}$ . As introduced in Section 2,  $\nabla_{\mathrm{STGS}}$  conducts a two-step approximation: (1) it approximates  $\min_{\theta} E[f(D)]$  as  $\min_{\theta} E[f(\operatorname{softmax}_{\tau}(\theta + G)))$ ; (2) it approximates  $\frac{\partial f(\operatorname{softmax}_{\tau}(\theta + G))}{\partial D}$  as  $\frac{f(D)}{\partial D}$ . Since the bias introduced in both steps can be controlled by  $\tau$ ,  $\nabla_{\mathrm{STGS}}$  prefers to set  $\tau$  as a relatively small value.
Temperature Scaling for  $\nabla_{\mathrm{ST}}$  and  $\nabla_{\mathrm{ReinMax}}$ . As in Section 4, it does not involve temperature scaling to show  $\nabla_{\mathrm{ST}}$  and  $\nabla_{\mathrm{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  $\nabla_{\mathrm{ST}}$  in the same way it does for  $\nabla_{\mathrm{STGS}}$ . As in Figure 2, STGS, GR-MCK, and GST-1.0 work better when setting the temperature  $\tau \leq 1$ . ST and ReinMax work better when setting the temperature  $\tau \geq 1$ .
Thus, we incorporate temperature scaling to smooth the gradient approximation  $(\pi_{\tau} = \mathrm{softmax}_{\tau}(\theta))$  as  $\widehat{\nabla}_{\mathrm{ReinMax}} = 2\cdot \widehat{\nabla}^{\frac{\pi_{\tau} + D}{2}} - \frac{1}{2}\widehat{\nabla}_{\mathrm{ST}}$ . It is worth emphasizing that  $\tau$  in  $\widehat{\nabla}_{\mathrm{ReinMax}}$  is used to stabilize the gradient approximation (instead of reducing bias) at the cost of accuracy. Therefore, the value of  $\tau$  should be larger or equal to 1.
Table 1: Performance on ListOps.
STGS	GR-MCK	GST-1.0	ST	ReinMax
Valid Accuracy	66.95±3.05	66.53±0.58	66.28±0.52	66.51±0.76	67.65±1.25
Test Accuracy	67.30±2.50	66.53±0.86	66.30±0.62	66.26±0.48	68.07±1.18
Table 2: Training -ELBO on MNIST ( $N \times M$  refers to  $N$  categorical dim. and  $M$  latent dim.).
AVG	8 × 4	4 × 24	8 × 16	16 × 12	64 × 8	10 × 30
STGS	105.20	126.85±0.85	101.32±0.43	99.32±0.33	100.09±0.32	104±0.41	99.63±0.63
GR-MCK	107.06	125.94±0.71	99.96±0.25	99.58±0.31	102.54±0.48	112.34±0.48	102.02±0.18
GST-1.0	104.25	126.35±1.24	101.49±0.44	98.29±0.66	98.12±0.57	102.53±0.57	98.64±0.33
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.66±0.88	99.77±0.45	97.70±0.39	98.06±0.53	100.71±0.70	98.37±0.44
Here, we conduct experiments on polynomial programming, unsupervised generative modeling, and structured output prediction. In all experiments, we consider four major baselines: Straight-Through (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},\dots ,X_{L}\in \{0,1\}$  and a constant vector  $c\in \mathcal{R}^{L\times 1}$ , we parameterize the distributions of  $\{X_1,\dots ,X_L\}$  with  $L$  softmax functions, i.e.,  $X_{i}\stackrel{\mathrm{iid}}{\sim}$  Multinomial(softmax  $(\pmb {\theta}_i)$ ) and  $\pmb {\theta}_i\in \mathcal{R}^2$ . Following previous studies, we set every dimension of  $c$  as 0.45, i.e.,  $\forall i, c_i = 0.45$ , and use  $\min_{\pmb{\theta}}E_{\pmb{X}}[\frac{\|\pmb{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.
Figure 3: The accuracy (left) and loss (right) on the valid set of ListOps.
Figure 4: The training -ELBO (left) and the cos similarity between the gradient and its approximations (right) on MNIST-VAE (with 4 latent dimensions and 8 categorical dimensions).
Table 3: 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	ST	STGS	GST-1.0	GR-MCK100	GR-MCK300	GR-MCK1000
QP	0.2s / 6.5Mb	0.2s / 5.0Mb	0.2s / 5.5Mb	0.2s / 8.0Mb	0.8s / 0.3Gb	2.2s / 1Gb	6.6s / 3Gb
MNIST-VAE	5.2s / 13Mb	5.2s / 13Mb	5.2s / 13Mb	5.2s / 13Mb	5.2s / 76Mb	5.2s / 0.2Gb	5.4s / 0.6Gb
Table 4: Performance on NATS-Bench. * Baseline results are referenced from Dong et al. (2020a).
CIFAR-10	CIFAR-100	ImageNet-16-120
validation	test	validation	test	validation	test
GDAS + STGS*	89.68±0.72	93.23±0.58	68.35±2.71	68.17±2.50	39.55±0.00	39.40±0.00
GDAS + ReinMax	90.01±0.12	93.44±0.23	69.29±2.34	69.41±2.24	41.47±0.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 (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, 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 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.
Table 5: Training -ELBO on MNIST. * All baseline results are referenced from Fan et al. (2022)
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
(a) MNIST when K=2.	(b) Omniglot when K=2.	(c) Fashion-MNIST when K=2.
(d) MNIST when K=3.	(e) Omniglot when K=3.	(f) Fashion-MNIST when K=3.
Figure 5:  $2 \times 200$  VAE training curves on MNIST, Omniglot, and Fashion-MNIST when  $K = 2$  or 3.
Table 6: Train -ELBO of  $2 \times 200$  VAE on MNIST, Fashion-MNIST, and Omniglot. * Baseline results are referenced from Shi et al. (2022). K refers to the number of evaluations.
RELAX*	ARMS*	DisARM*	Double CV*	RODEO*	ReinMax
K=3	MNIST	101.99±0.04	100.84±0.14	/	100.94±0.09	100.46±0.13	97.83±0.36
Fashion-MNIST	237.74±0.12	237.05±0.12	/	237.40±0.11	236.88±0.12	234.53±0.42
Omniglot	115.70±0.08	115.32±0.07	/	115.06±0.12	115.01±0.05	107.51±0.42
K=2	MNIST	/	/	102.75±0.08	102.14±0.06	101.89±0.17	98.17±0.29
Fashion-MNIST	/	/	237.68±0.13	237.55±0.16	237.44±0.09	234.89±0.21
Omniglot	/	/	116.50±0.04	116.39±0.10	115.93±0.06	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 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 \times 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  $\frac{\partial f(D)}{\partial 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  $\frac{1}{n}\sum_{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  $\frac{1}{n}\sum_{i}f(I_{i})$  as the baseline. We suspect that the gradient approximation using  $\frac{1}{n}\sum_{i}f(I_{i})$  as the baseline is very unstable as it contains the  $\frac{1}{n\cdot p(D)}$  term.
Figure 6: Training -ELBO on MNIST-VAE when using  $\frac{1}{n}\sum_{i}f(I_{i})$  and  $E[f(D)]$  as baselines respectively.
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  $\tau \leq 1$ . Differently, ST and ReinMax prefer to set the temperature  $\tau \geq 1$ . These observations match our analyses in Section 5 that a small  $\tau$  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.
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(\mathbf{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.
We would like to thank all reviewers for their constructive comments, the engineering team at Microsoft for providing computation infrastructure support, Alessandro Sordoni, Nicolas Le Roux, and Greg Yang for their helpful discussions.
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 fourgpu 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 gumbel-softmax 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  $\nabla$  as  $\widehat{\nabla}_{\mathrm{1st-order}} = \sum_{i}\sum_{j}\pi_{j}\frac{\partial f(\boldsymbol{I}_{j})}{\partial\boldsymbol{I}_{j}} (\boldsymbol{I}_{i} - \boldsymbol{I}_{j})\frac{d\boldsymbol{\pi}_{i}}{d\boldsymbol{\theta}_i}$ , which approximates  $f(I_i) - f(I_j)$  in Equation 6 as  $\frac{\partial f(\boldsymbol{I}_j)}{\partial\boldsymbol{I}_j} (\boldsymbol {I}_i - \boldsymbol {I}_j)$ .
$$
E [ \widehat {\nabla} _ {S T} ] = \widehat {\nabla} _ {I s t - o r d e r}.
$$
Proof. Based on the definition, we have:
$$
\begin{array}{l} \widehat {\nabla} _ {\text {l s t - o r d e r}} = \sum_ {i} \sum_ {j} \pi_ {j} \frac {\partial f (\boldsymbol {I} _ {j})}{\partial \boldsymbol {I} _ {j}} \left(\boldsymbol {I} _ {i} - \boldsymbol {I} _ {j}\right) \frac {d \pi_ {i}}{d \theta} \\ = \sum_ {j} \pi_ {j} \frac {\partial f (\boldsymbol {I} _ {j})}{\partial \boldsymbol {I} _ {j}} \sum_ {i} \boldsymbol {I} _ {i} \frac {d \pi_ {i}}{d \boldsymbol {\theta}} - \sum_ {j} \pi_ {j} \frac {\partial f (\boldsymbol {I} _ {j})}{\partial \boldsymbol {I} _ {j}} \boldsymbol {I} _ {j} \sum_ {i} \frac {d \pi_ {i}}{d \boldsymbol {\theta}} \tag {9} \\ \end{array}
$$
Since  $\sum_{i}\pi_{i} = 1$ , we have  $\sum_{i}\frac{d\pi_{i}}{d\theta} = 0$ . Also, since  $\pi = \sum_{i}\pi_{i}I_{i}$ , we have  $\frac{d\pi}{d\theta} = \sum_{i}I_{i}\frac{d\pi_{i}}{d\theta}$ . Thus, together with Equation 9, we have:
$$
\begin{array}{l} \widehat {\nabla} _ {\text {I s t - o r d e r}} = \sum_ {j} \pi_ {j} \frac {\partial f (\boldsymbol {I} _ {j})}{\partial \boldsymbol {I} _ {j}} \sum_ {i} \boldsymbol {I} _ {i} \frac {d \pi_ {i}}{d \boldsymbol {\theta}} \\ = E \left[ \frac {\partial f (\boldsymbol {D})}{\partial \boldsymbol {D}} \frac {d \boldsymbol {\pi}}{d \boldsymbol {\theta}} \right] = E \left[ \widehat {\nabla} _ {\mathrm {S T}} \right]. \\ \end{array}
$$
B Theorem 3.2
Proof. Here, we aim to proof,  $\forall k\in [1,n]$ , we have  $E[\widehat{\nabla}_{\mathrm{ReinMax},k}] = \widehat{\nabla}_{2\mathrm{nd - order},k}$ . As defined in Equation 8, we have (note that  $\delta_{D}\pmb{I}_k$  is the indicator function of the event  $\pmb{D} = \pmb{I}_k$ ):
$$
\begin{array}{l} \widehat {\nabla} _ {2 \mathrm {n d} - \mathrm {o r d e r}, k} = \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} _ {k}} \\ = \sum_ {i} \sum_ {j} \frac {\pi_ {j} \pi_ {i} \left(\delta_ {i k} - \pi_ {k}\right)}{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) \\ = \sum_ {j} \frac {\pi_ {j} \pi_ {k}}{2} \left(\frac {\partial f (\boldsymbol {I} _ {j})}{\partial \boldsymbol {I} _ {j}} + \frac {\partial f (\boldsymbol {I} _ {k})}{\partial \boldsymbol {I} _ {k}}\right) (\boldsymbol {I} _ {k} - \boldsymbol {I} _ {j}) \\ = \frac {\pi_ {k}}{2} \frac {\partial f (\boldsymbol {I} _ {k})}{\partial \boldsymbol {I} _ {k}} (\boldsymbol {I} _ {k} - \sum_ {j} \pi_ {j} \boldsymbol {I} _ {j}) + \sum_ {j} \frac {\pi_ {j} \pi_ {k}}{2} \frac {\partial f (\boldsymbol {I} _ {j})}{\partial \boldsymbol {I} _ {j}} (\boldsymbol {I} _ {k} - \boldsymbol {I} _ {j}) \\ = \frac {1}{2} E _ {D \sim \pi} \left[ \delta_ {D I _ {k}} \frac {\partial f (D)}{\partial D} \left(I _ {D} - \sum_ {j} \pi_ {j} I _ {j}\right) \right] + \frac {1}{2} E _ {D \sim \pi} \left[ \pi_ {k} \frac {\partial f (D)}{\partial D} \left(I _ {k} - I _ {D}\right) \right] \\ = \frac {1}{2} E _ {D \sim \pi} \left[ \frac {\partial f (D)}{\partial D} \left(\boldsymbol {\pi} _ {k} \left(\boldsymbol {I} _ {k} - \boldsymbol {I} _ {D}\right) + \delta_ {D I _ {k}} \left(\boldsymbol {I} _ {D} - \sum_ {i} \boldsymbol {\pi} _ {i} \boldsymbol {I} _ {i}\right)\right) \right] \tag {10} \\ \end{array}
$$
At the same time, based on the definition of  $\widehat{\nabla}_{\mathrm{ReinMax}}$ , we have: