Custom Gradient Estimators are Straight-Through Estimators in Disguise

Authors: Matthew Schoenbauer, Daniele Moro, Lukasz Lew, Andrew G. Howard
Primary area: optimization
Keywords: quantization, deep learning, optimization

Abstract:
Quantization-aware  training  comes with a fundamental challenge: the derivative of quantization functions such as rounding are zero almost everywhere and nonexistent elsewhere. Various differentiable approximations of quantization functions have been proposed to address this issue. In this paper, we prove that a large class of weight gradient estimators is approximately equivalent with the straight through estimator (STE). Specifically, after swapping in the STE and adjusting both the weight initialization and the learning rate in SGD, the model will train in almost exactly the same way as it did with the original gradient estimator. Moreover, we show that for adaptive learning rate algorithms like Adam, the same result can be seen without any modifications to the weight initialization and learning rate. These results reduce the burden of hyperparameter tuning for practitioners of QAT, as they can now confidently choose the STE for gradient estimation and ignore more complex gradient estimators. We experimentally show that these results hold for both a small convolutional model trained on the MNIST dataset and for a ResNet50 model trained on ImageNet.

Decision: reject
Reviewer scores: 5, 5, 6, 5

Reviews and discussion:
<thread>
creation date:
1730040007116

thread_type:
official_review

<thread_object>
creator: reviewer_kqap

<content>
summary:
The authors examine how different gradient estimators affect quantization-aware training. They theoretically demonstrate that, in the limit of small learning rates and with minor adjustments to the initialization and the learning rate magnitude, most gradient estimators yield equivalent weight movements. Consequently, they suggest that the Straight-Through Estimator, treating the gradient as if no quantization occurred in the backward pass, performs as well as any other more sophisticated alternative. Their theoretical claims are complemented by an empirical investigation.

strengths:
Equations (6) and (7) provide a rigorous theoretical contribution regarding the small difference in weight movement for different gradient estimators.

The presentation and writing style is very clear with helpful intuition such as the analogy of the "funhouse mirror". Additionally, the learning rate tweak in the experiment provides a practical comparison for the magnitude of the differences.

weaknesses:
The authors themselves acknowledge in Section 8 that many publications introduce more than just a new gradient estimator, raising questions about the broader practical impact of their findings. Novel gradient estimators are often proposed with other techniques to enable any benefits. As a result, the findings may have limited applicability beyond specific configurations.

It is somewhat difficult to draw clear conclusions about practical applications from the experiments. While the learning rate tweak provides a useful comparison, the differences in Tables 4 and 5 are challenging to interpret, particularly without comparisons across different initializations. A convincing experiment could be, to find a custom gradient estimator in the literature, which has been shown to improve validation accuracy over STE, replicate the results and then demonstrate that both perform equally with proper adjustments. The authors do not provide such a comparison, raising questions about whether the custom estimator was applied correctly in their experiments or if its potential advantages were overlooked.


Minor comments:

line 245: Typo: $Q(w_{\hat{Q}}^{(t)})=Q(w_{\hat{Q}}^{(t)})$

lines 329-340: Equations (8) to (12): "$\nabla$" missing before $f^{(t)}_{STE}$

lines 502-504: Missing figure number for Figure 3

questions:
Lines 400-404: Why was the initial 10% of training for the ImageNet-ResNet setup kept identical? What would the results be without this measure?

Lines 527-528: Regarding the statement that high learning rates might be the reason the equivalence is not observed in other studies, the authors write, "we expect that this counter-argument will not stand the test of time, since by our main results, the higher learning rate masks the fact that models with novel $\hat{Q}$ and the STE are still approximating the same process." Could you elaborate on what you mean by "will not stand the test of time" given that equations (6) and (7) indicate learning-rate-squared errors? It seems that higher learning rates would increase these errors. Why should these differences not enable advantages for novel $\hat{Q}$?


Lines 530-539: Can the authors comment on the potential implications of their results for the studies mentioned in the last paragraph, which propose additional innovations alongside novel gradient estimators?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730229272742

thread_type:
official_review

<thread_object>
creator: reviewer_xprq

<content>
summary:
This paper presents a theoretical analysis of gradient estimators used in quantization-aware training. The authors show that in the case of quantized weights (but full precision activations) many extant gradient estimators for the quantization operation are approximately equivalent, if the learning rate and weight initialization are adjusted, and the learning rate is small. They then verify empirically their theoretical results, on image classification benchmarks, demonstrating that models trained with different gradient estimators indeed show high weight agreement and similar accuracies.

strengths:
- The theoretical insights are interesting, unexpected and (to my knowledge) novel. They offer better understanding and insight into how gradient estimators work, which appeals to me.
- The paper is generally well written and easy to read. I appreciate how the authors lead their result with an intuitive explanation and illustrative graphic. This makes the following theory much easier to intuit.
- The experiments shown in the paper provide good evidence for the theoretical results.

weaknesses:
Major

1. The claims relating to practical impact feel overstated ("practitioners can now confidently choose the STE"). The problem setting that the authors explore (full precision activations, quantized weights, uniform fixed point quantization, small learning rate) is rather specific, and practitioners may be interested in quantized activations, or low-precision floating point or larger learning rates etc. I would prefer if the authors tempered their claims.
2. The experiments, although they demonstrate the theory well, are limited. They do not show finetuning from a pretrained full precision checkpoint, as is common for QAT in practice (I would expect this setting to match well with the theory since 1. QAT finetuning is done with a lower learning rate typically and 2. the gradient norm is likely to be low after initializing from a pretrained model). They do not show results other than 2-bit weight quantization even though they say results are similar. They do not show the practical limits of their theory, e.g. how weight alignment degrades/evolves over training or how much the learning rate needs to be increased for the error terms to start having a large impact.

Minor
1. Presentation could be improved in a number of ways. 
    1. Use of booktabs for tables. Place all table captions above the tables.
    1. Tables are hard to parse when skimming -- would benefit from more descriptive captions/grouping table 3 with 4
    2. All quotation marks are incorrectly rendered by LaTeX.
    2. Fig. 3 would look better with the bins.
2. The choice of training recipes are not explained -- it is unclear why the first 10 epochs are done using the same gradient estimator. Is it because the gradient norm is too high at the start of training resulting in the weights quickly diverging?
3. I think the point made in line 305 should be made more prominent. I think it is quite important that the reader is made aware that the gradient error is small/zero since Q-net and STE net will quantize to similar weights.

questions:
See weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1730297392794

thread_type:
official_review

<thread_object>
creator: reviewer_oppr

<content>
summary:
This paper studies behavior of gradient estimators, including straight through
estimator (STE), for weight quantization. It is shown that a large class of
weight gradient estimators is approximately equivalent to the STE during training using SGD and Adam.

strengths:
1. The paper is overall well presented.
2. The concept of mirror effect is interesting.

weaknesses:
A primary concern is that the key claims and several major concepts lack mathematical rigor. Additionally, the main theoretical results provided are too limited to substantiate the claims:

1. Contribution 1 states that '... all nonzero weight gradient estimators lead to approximately equivalent weight movement for non-adaptive learning rate optimizers ...'. However, the term 'approximately equivalent weight movement' lacks a precise mathematical definition. It would be helpful to formalize this concept, perhaps by specifying the conditions under which these movements are considered 'approximately equivalent.' 

2. According to Section 6.2, 'approximately equivalent weight movement' appears to refer to high 'Quantized Weight Agreement' or a small 'Normalized Weight Alignment Error ($\bar{E}$)'. Again, these metrics require explicit mathematical expressions for each. Additionally, this interpretation is not fully supported by the main theoretical results (e.g., Theorem 5.1), which only derive the increment in weight alignment error between two consecutive iterations, rather than a direct measure of agreement or alignment over the entire optimization trajectory. 

3. For the error bounds in Eq. (6) and (7), there is insufficient justification for why the gradient error terms should be small, nor any clear indication of how small these terms are. It is insufficient to merely claim that a term is 'small' and then disregard it. These errors could accumulate significantly over iterations, potentially undermining the main conclusions.

4. The use of mathematical notation is poor, which possibly lead to incorrect derivations. For example:
   - The Euclidean norm should be denoted by $\|| \cdot \||$ rather than $\| \cdot \|$ as in Eq. (5)-(12) and other instances.
   - It should be explicitly stated that  $Q$ and $M$ are applied *element-wise* to the weight vector. Additionally, it would be preferable to use bold letters to represent vectors and to distinguish them from scalars.
   - In Eq. (10), you have three vectors,  $\nabla f_{Q}^{(t)}$, $\hat{Q}^\prime$, $M^\prime$, how are they multiplied together? The manner in which they are multiplied together is unclear. Furthermore, the residual term in Eq. (10) should not be a scalar $O(\eta^2)$, but rather a vector. I believe that the second-order error term also depends on the *model size*, i.e., the dimension of $w$.

5. The experiments is only conducted for one instance of $\hat{Q}$ (HTGE Pei et al.), which is insufficient. Additionally, the expression of $\hat{Q}$ from HTGE is missing.

Minor comments:

1. A key reference on the theoretical analysis of STE is missing, specifically: *Yin et al., Understanding Straight-Through Estimator in Training Activation Quantized Neural Networks, ICLR2019.*

questions:
1. In Eq. (5), why the $E^{(t)}$ is defined differently for SGD and Adam?

2. Line 200, "Most multi-bit gradient estimators proposed in the literature are cyclical." Can you specify exactly which estimators are cyclical?

3. Typos in Eq. (9)-(12), $f_{STE}^{(t)}$ should be $\nabla f_{STE}^{(t)}$.

4. In assumption 5.1.1, should $| \hat{Q}^\prime (w) |$ be $\hat{Q}^\prime (w)$ without the absolute value?

5. Does Table 4 suggest that more than 95% quantized weights from $\hat{Q}$-net and STE-net are *identical*? This finding seems counterintuitive.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730681414308

thread_type:
official_review

<thread_object>
creator: reviewer_6p1r

<content>
summary:
The authors theoretically analyze the weight difference in QAT when trained with different gradient estimators. Under certain conditions, the authors show that the weight difference is small which means that there's no need to try another gradient estimator other than STE. Empirical results show that the weight difference is small when adopting the proposed weight initialization.

strengths:
1. The claim is strong that other gradient estimators works similar as STE in QAT.
2. Experiments show that the weight difference is small to support the claim.

weaknesses:
1. The mirror room story does not appear closely connected to the theoretical analysis.
2. Assumption 5.1.1 violates Figure 1 where the gradient could be zero.
3. From Table 4, Adam leads to larger weight difference.
4. For more complicated task like ImageNet, the weight difference is much larger than MNIST.

questions:
1. Could you provide unadjusted(A) in Table 4?
2. Besides average error, could you draw a histogram which can better validate the claims?
3. How do you empirically decide when the weight difference is going to affect the conclusion? E.g., Adam shows larger weight difference, while ImageNet shows larger weight difference. It does not seem clear to me they all supports the claim that STE works the same as other gradient estimators on various settings.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors analyze the impact of gradient estimators on quantization-aware training, concluding that with small learning rates and minor adjustments, most estimators result in similar weight updates. They also claim that the Straight-Through Estimator performs comparably to more complex methods, supporting this claim with theoretical and empirical evidence.

The main issue of this paper is that some main claims are not strongly supported by its theoretical results and its numerical experiments. For example, the authors claim that the weight movement differences will be small if the learning rate is small, so $STE$-net will behave similarly to $\hat Q$-net in SGD and Adam. However, as one of the reviewer pointed out, the error $E^t$ keep increasing as $t$ increases, no matter how small the learning rate is. For a large number of iteration, $E^t$ can still be very large. The review requested some analysis on the cumulative error. However, the authors just provide the formula of the cumulative error on page 7 which is basically $O(t\eta)$ and is not necessarily small for a large $t$. Moreover, $\eta$ is not necessarily always small when training a model. It may decrease with $t$ so it can be large for the first few iterations. 

Also, the authors claimed that practitioners only need to use the Straight-Through Estimator because other estimators are similar to STE after modifying the learning rate and initialization. However, one reviewer pointed out that this claim is not supported by the current experiments. It needs a direct comparison with ""custom gradient estimators in the settings where they were originally proposed"".

I also think this paper is not well written with many abuses of math notations and missing definitions. The review team has pointed out some but there are still many others. For example, what is the definition of "$round()$" function? Also, $STE$-net and $\hat Q$-net are defined in words rather than using clear math formula.