Summary
The paper shows that the straight-through estimator is a first-order approximation of the gradient. The authors then propose a method, called ReinMax, which provides a second-order approximation with negligible computational overhead. Experiments are performed in several settings involving discrete variables (polynomial programming, structured output prediction, and discrete latent variable generative models), showing that ReinMax is more accurate and stable.
Strengths
**Originality**
The 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.
**Quality**
The 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.
**Clarity**
Overall, 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.
**Significance**
While many papers have mentioned that the straight-through estimator is an approximation of the gradient, it appears that none of these papers have formally shown it. If, indeed, this paper is the first to do so, then it is a significant contribution to our theoretical understanding. The proposed improved estimator, ReinMax, also appears to offer some performance improvement over previously proposed (first-order) estimators. Considering that ReinMax has similar computational overhead as straight-through, then this could serve as a drop-in replacement for the straight-through estimator throughout various applications. However, it is unclear to me whether this serves as a drop-in replacement for all existing instances of the straight-through estimator, or merely those that operate on Multinomial distributions.
Weaknesses
I see two relatively minor weaknesses in the paper as-is: larger-scale empirical evaluation and slight improvements to the presentation. These are discussed below.
The 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.
Additionally, several minor aspects of the presentation could be improved.
* 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.
* As far as I can tell, the name “ReinMax” is never actually explained.
* 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.
* The shaded regions in Figure 5 (right) are rather jagged — perhaps there’s an issue with the evaluation interval or the plotting setup.
* 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.
* The citation in the second sentence of the introduction is incorrect.
Questions
Could you please elaborate on the applicability of the ReinMax estimator? Does this estimator (or the insights developed in the theoretical sections of the paper) alleviate the “laborious and time-consuming” work of developing “different ST variants for different applications in a trial-and-error manner”?
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Limitations
As mentioned above, I would appreciate a clearer discussion of which existing forms of straight-through estimator the ReinMax estimator can replace.