Understanding Straight-Through Estimator in Training Activation Quantized Neural Nets

By Penghang Yin, Jiancheng Lyu, Shuai Zhang, Stanley J. Osher, Yingyong Qi, Jack Xin. 2019

440 citations recorded by the source corpus.

Training activation quantized neural networks involves minimizing a piecewise constant function whose gradient vanishes almost everywhere, which is undesirable for the standard back-propagation or chain rule. An empirical way around this issue is to use a straight-through estimator (STE) (Bengio et al., 2013) in the backward pass only, so that the "gradient" through the modified chain rule becomes non-trivial. Since this unusual "gradient" is certainly not the gradient of loss function, the following question arises: why searching in its negative direction minimizes the training loss? In this paper, we provide the theoretical justification of the concept of STE by answering this question. We consider the problem of learning a two-linear-layer network with binarized ReLU activation and Gaussian input data. We shall refer to the unusual "gradient" given by the STE-modifed chain rule as coarse gradient. The choice of STE is not unique. We prove that if the STE is properly chosen, the expected coarse gradient correlates positively with the population gradient (not available for the training), and its negation is a descent direction for minimizing the population loss. We further show the associated coarse gradient descent algorithm converges to a critical point of the population loss minimization problem. Moreover, we show that a poor choice of STE leads to instability of the training algorithm near certain local minima, which is verified with CIFAR-10 experiments.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

Citations in the corpus

Dynamic Model Pruning with Feedback (2020, background): “Our stronger assumptions on the structure of the perturbation allow to derive non-asymptotic convergence rates in the next section, though our analysis could also be extended to the setting in (Yin et al., 2019) if the perturbations can be bounded.”
Explaining Question Answering Models through Text Generation (2020, background): “For details on the GS-ST estimator, please see Bengio et al. (2013); Jang et al. (2016); Maddison et al. (2016); Yin et al. (2019).”
A Survey of Quantization Methods for Efficient Neural Network Inference (2021, neutral): “The work of [263] provides a theoretical justiﬁcation for this phenomena, and it ﬁnds that the coarse gradient approximation of STE can in expectation correlate with population gradient (for a proper choice of STE).”
AutoPrune: Automatic Network Pruning by Regularizing Auxiliary Parameters (2019, background): “5, Lemma4 and Lemma10 from Yin et al. [2019], where the ReLU and linear STE form acute angle with population gradient.”
Improving Inference for Neural Image Compression (2020, methodology): “All methods were tuned on a best-effort basis to ensure convergence, except that STE consistently encountered convergence issues even with a tiny learning rate (see [Yin et al., 2019]).”
Single-Path Mobile AutoML: Efficient ConvNet Design and NAS Hyperparameter Optimization (2019, methodology): “A deeper study on the STE is an interesting direction for future NAS work, similar to recent STE analysis in the context of hardware-aware quantization [49].”
Model Compression and Hardware Acceleration for Neural Networks: A Comprehensive Survey (2020, methodology): “Using pulse train to approximate the derivative of discontinuous points [151], designing proper STE formats [154], or gradually rather than immediately closing the gradient BP path to the high-precision data [155] also works.”
Overcoming Oscillations in Quantization-Aware Training (2022, methodology): “Earlier work in QAT (Cai et al., 2017; Yin et al., 2019) investigated different variations of the STE and demonstrate experimentally that the derivative of the clipped ReLU is a better approximation than the vanilla STE (derivative of identity).”
DKM: Differentiable K-Means Clustering Layer for Neural Network Compression (2021, background): “This approach might degrade model quality because it cannot formulate weight-cluster assignment during gradient back propagation [28].”
Mixed Precision DNNs: All you need is a good parametrization (2019, background): “This leads to which is non-zero in the interesting region | x | ≤ q max and which turned out to be very useful to train quantized DNNs in practice (Yin et al., 2019).”
MERF: Memory-Efficient Radiance Fields for Real-time View Synthesis in Unbounded Scenes (2023, methodology): “This use of a stop-gradient allows us to obtain gradients for the non-differentiable rounding function by treating 𝑞 as the identity function during the backward pass, which is referred to as the straight-through estimator, as used in [Bengio et al. 2013; Yin et al. 2019].”
Learning N: M Fine-grained Structured Sparse Neural Networks From Scratch (2021, methodology): “However, perturbations are introduced during the back-propagation (Yin et al., 2019; Bengio et al., 2013).”
Semi-Supervised Neural Text Generation by Joint Learning of Natural Language Generation and Natural Language Understanding Models (2019, methodology): “Through (ST) estimator, which simply consists of backpropagating through the argmax function as if it had been the identity function (Bengio et al., 2013; Yin et al., 2019).”