Understanding Straight-Through Estimator in Training Activation Quantized Neural Nets

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