Custom Gradient Estimators are Straight-Through Estimators in Disguise
Matt Schoenbauer, Daniele Moro, Lukasz Lew, Andrew Howard
TL;DR
This work tackles the core challenge of differentiating through quantization in QAT by proving that, at small learning rates, a broad class of gradient estimators yields training dynamics nearly identical to STE after a specific weight reinitialization and learning-rate mapping. For SGD-type optimizers, the authors show a rigorous equivalence after applying a weight map $M$ and scaling the learning rate by $\alpha$, while Adam achieves the same equivalence without reinitialization. They extend the theory to adaptive optimizers and provide empirical validation on MNIST and ImageNet, demonstrating that STE can be a robust surrogate for complex gradient estimators under appropriate conditions. The findings reduce hyperparameter tuning burden and offer practical guidance for deploying QAT, while also clarifying why numerous gradient-estimator proposals may yield improvements primarily through accompanying initialization and LR adjustments. Overall, the paper argues that gradient error concerns in QAT may be less critical than previously thought, given proper initialization and learning-rate design.
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 when the learning rate is sufficiently small, a large class of weight gradient estimators is 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. 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.