Understanding Straight-Through Estimator in Training Activation Quantized Neural Nets

Penghang Yin; Jiancheng Lyu; Shuai Zhang; Stanley Osher; Yingyong Qi; Jack Xin

Source-linked AI summary

Understanding Straight-Through Estimator in Training Activation Quantized Neural Nets

Penghang Yin, Jiancheng Lyu, Shuai Zhang, Stanley Osher, Yingyong Qi, Jack Xin

arXiv:1903.05662v4Mar 13, 2019cs.LGmath.OCstat.ML

TL;DR

Straight-through estimators enable training activation-quantized networks despite vanishing true gradients, but their coarse gradients lack a clear optimization justification. This paper analyzes three STE choices and proves that vanilla and clipped ReLU STEs yield descent directions and convergence, while identity STE can be unstable.

Problem
Because STE coarse gradients are generally not true gradients and STE choices are non-unique, their descent behavior and proper selection remain unclear.
Method
The paper analyzes identity, vanilla ReLU, and clipped ReLU STEs for a two-layer binarized-ReLU network with Gaussian data using population-loss theory.
Results
Negative expected coarse gradients from vanilla and clipped ReLU STEs are descent directions and their algorithms converge to critical points, unlike identity STE.
Takeaways & Limitations
Properly chosen STEs can support descent-based training of activation-quantized networks, while identity STEs may be incompatible with the loss landscape.
Takeaways & Limitations
The theoretical convergence results assume a two-layer binarized-ReLU network with Gaussian data, and identity STE training can become unstable near local minima.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper studies why straight-through estimators can train activation-quantized networks despite producing a coarse gradient that is not generally a true gradient. It theoretically analyzes when STE choices yield descent directions and empirically examines their convergence and instability.

Motivation: Fully quantized DNN training minimizes a piecewise-constant, highly nonconvex empirical risk over a discrete quantized-weight constraint.This optimization difficulty motivates constructing non-trivial search directions instead of relying on ordinary gradients.
Motivation: An STE replaces the almost-everywhere-zero derivative of a quantized activation with a surrogate derivative used only during back-propagation.The forward and backward passes therefore do not match, producing a coarse gradient that is generally not the gradient of any function.
Theory: For a two-linear-layer network with binary activation and Gaussian data, the paper analyzes identity, vanilla ReLU, and clipped ReLU STEs under population-loss minimization.It proves that properly chosen STEs produce expected coarse gradients positively correlated with the population gradient, making their negatives descent directions.
Theory: The associated coarse-gradient descent algorithm is shown to converge to a critical point, while poor STE choices can destabilize training near certain local minima.The paper emphasizes that identity-STE guarantees from perceptron and Convertron analyses do not generalize to this two-trainable-layer setting.
Experiments: Experiments on MNIST and CIFAR-10 find vanilla and clipped ReLUs effective for shallow LeNet-5, while clipped ReLU is arguably best for deeper VGG-11 and ResNet-20.CIFAR experiments observe identity- or ReLU-STE training becoming unstable at good minima and moving to an inferior one with substantially higher loss.

2 LEARNING TWO-LINEAR-LAYER CNN WITH BINARY ACTIVATION

The section formulates learning a two-linear-layer CNN with binary activation under Gaussian inputs and analyzes its population-loss landscape. Because the binary activation has an almost-everywhere-zero derivative, training uses a straight-through estimator to define a non-trivial coarse gradient and coarse gradient descent.

Model formulation: The model uses a shared first-layer filter across input patches and a second linear layer as classifier, with labels generated by non-zero true parameters.The activation is binary: σ(x) = 1{x>0}.
Model formulation: Entries of Z are i.i.d. N(0, 1), and scale invariance permits the normalization ∥w∗∥ = 1 for population-loss minimization.The loss satisfies ℓ(v, w; Z) = ℓ(v, w/c; Z) for any scalar c > 0.
Straight-through estimator: The binary activation’s derivative is zero almost everywhere, so standard back-propagation cannot provide a usable training gradient.The population gradient is not directly available during network training; only the expected sample gradient can be accessed.
Straight-through estimator: An STE replaces σ′ with the non-trivial derivative µ′ of a differentiable or subdifferentiable function µ, producing the surrogate coarse gradient.Using µ′ yields full-batch coarse gradient descent for the binary-activation CNN.
Population-loss landscape: The population-loss landscape has global minimizers at v = v∗ and θ(w, w∗) = 0, while stationary points can only be saddles and opposite-alignment points are potential spurious local minimizers.If the stated stationary or opposite-alignment cases do not occur, the model has no saddle points or spurious local minimizers.

3 MAIN RESULTS

The main results show that vanilla and clipped ReLU STEs yield coarse-gradient descent methods with positive alignment and convergence to population-loss critical points. The identity STE lacks these guarantees and can cause instability near spurious local minima.

Assumptions and empirical behavior: The convergence guarantee assumes infinitely many training samples; with few samples, empirical loss roughly follows negative coarse-gradient steps, becoming more monotone and smooth as sample size increases.This explains why properly chosen STEs work well with massive datasets.
ReLU and clipped ReLU STEs: With sufficiently small learning rate η, vanilla ReLU coarse-gradient descent monotonically decreases the population loss until convergence.The descent argument uses the positive-correlation estimate together with Lipschitz continuity of the population gradient.
ReLU and clipped ReLU STEs: When Algorithm 1 converges with ReLU or clipped ReLU STE, its coarse gradient vanishes only at population-loss critical points, including saddles and specified minimizers.Thus, both STE choices converge to a critical point rather than an arbitrary stationary point of the coarse-gradient dynamics.
Identity STE: The identity STE may fail near spurious local minima because its coarse gradient does not vanish there and the algorithm may never converge.For m > 1 and 1_m^⊤v∗ ≠ 0, the descent property fails near minimizers with θ(w, w∗) = π, causing training loss to increase and instability to arise.

4 EXPERIMENTS

Experiments compare identity, vanilla ReLU, and clipped ReLU STEs on MNIST and CIFAR-10 with 2-bit and 4-bit activations. Clipped ReLU performs best overall, while identity STE can destabilize training near good minima.

Experimental setup: The experiments compare identity, vanilla ReLU, and clipped ReLU STEs on MNIST and CIFAR-10 using 2-bit or 4-bit quantized activations with float weights.LeNet-5 is trained for 50 epochs on MNIST, while VGG-11 and ResNet-20 are trained for 200 epochs on CIFAR-10 using momentum 0.9.
Results: The derivative of clipped ReLU gives the best overall performance, followed by vanilla ReLU and then the identity function.On the relatively shallow LeNet-5 network, vanilla ReLU performs comparably to clipped ReLU; clipped ReLU is best for deeper networks.
Instability: On ResNet-20 with 4-bit activations, vanilla and clipped ReLUs reach validation accuracies of 86.59% and 91.24%, with training losses of 0.25 and 0.04, respectively.The corresponding coarse gradient descent algorithms converge to neighborhoods of the minima.
Instability: Using the identity STE on ResNet-20 with 4-bit activations repels coarse gradient descent from good minima, demonstrating the predicted instability.The same instability explains the poor performance of the ReLU STE on 2-bit activated ResNet-20, although divergence is slower.

5 CONCLUDING REMARKS · APPENDIX

The paper gives a theoretical justification for STE-based descent training in a two-linear-layer CNN with binary activation. It analyzes three STE choices and shows that vanilla and clipped ReLU yield descent directions through explicit expected coarse-gradient formulas.

5 CONCLUDING REMARKS: The paper provides the first theoretical justification that STE gives rise to a descent training algorithm.
5 CONCLUDING REMARKS: The analysis considers a two-linear-layer CNN with binary activation.
5 CONCLUDING REMARKS: Three STEs are studied: derivatives of the identity function, vanilla ReLU, and clipped ReLU.
5 CONCLUDING REMARKS: The paper derives explicit formulas for the expected coarse gradients associated with these STEs.
5 CONCLUDING REMARKS: The negative expected coarse gradients based on vanilla ReLU are descent directions for minimizing the population loss.
5 CONCLUDING REMARKS: The negative expected coarse gradients based on clipped ReLU are descent directions for minimizing the population loss.

A. THE PLOTS OF QUANTIZED AND CLIPPED RELUS

Figure 3 depicts a 2-bit quantized ReLU with four quantization levels, including 0, alongside its associated clipped ReLU. The resolution α is fixed before network training.

A. THE PLOTS OF QUANTIZED AND CLIPPED RELUS: Figure 3 plots the 2-bit quantized ReLU σα(x), with four quantization levels including 0, and the associated clipped ReLU ˜σα(x).The resolution α is determined in advance of network training.

B. THE SCHEDULE OF LEARNING RATE

This section presents the learning-rate schedule used in the study.

The learning-rate schedule is reported in Table 2.

C. INSTABILITY OF RELU STE ON RESNET-20 WITH 2-BIT ACTIVATIONS

On ResNet-20 with 2-bit activations, the ReLU STE coarse-gradient descent is unstable near weights found by the clipped ReLU STE, causing both classification and training errors to rise.

C. INSTABILITY OF RELU STE ON RESNET-20 WITH 2-BIT ACTIVATIONS: 88.38% validation accuracy was achieved by weights produced with the clipped ReLU STE before testing stability.
C. INSTABILITY OF RELU STE ON RESNET-20 WITH 2-BIT ACTIVATIONS: At a 10−5 learning rate, coarse gradient descent using the ReLU STE is not stable at those initialized weights.
C. INSTABILITY OF RELU STE ON RESNET-20 WITH 2-BIT ACTIVATIONS: Both classification and training errors begin to increase during the unstable optimization.

D. ADDITIONAL SUPPORTING LEMMAS

This section develops auxiliary lemmas for Gaussian expectations involving nonzero vectors and their angles, then establishes comparison and norm-difference bounds used in the analysis.

Gaussian expectation identities: Lemmas 11 and 12 derive identities for Gaussian-vector expectations parameterized by the angle between two nonzero vectors.Their proofs use coordinate reduction, polar representations of two-dimensional Gaussians, and symmetry of odd integrands.
Function comparison: Lemma 13 compares the functions p(θ, w) and q(θ, w) over θ ∈ [π/2, ...], using rearrangement and an inequality at θ = π/2.The supplied passage states that p(π/2, w) ≤ q(π/2, w) supports one comparison step.
Norm-difference bound: Lemma 14 bounds the relevant normalized vector difference when ∥˜w∥ ≥ ∥w∥ = c > 0 by 1/c^2 times ∥w − ˜w∥.The proof uses Cauchy–Schwarz, orthogonal projections, the angle between the vectors, and the identity relating their normalized inner product to ∥w − ˜w∥^2.

E. MAIN PROOFS

The proofs establish that suitable ReLU and clipped-ReLU straight-through estimators yield monotonically decreasing population loss and convergence to saddle points or local minimizers. They also show that the identity estimator can fail to descend or converge near certain local minimizers.

Stationary-point structure: The population loss has saddle points satisfying (8) and spurious local minimizers at θ(w,w*)=π with the specified v; otherwise, it has no saddle points or spurious local minimizers.The local-optimality proof identifies the relevant stationary points using an indefinite Hessian for saddle points and a strict local comparison for the spurious minima.
Gradient alignment: The expected coarse gradient with ReLU has positive alignment with the true population gradient away from degenerate angles, supporting descent of the population loss.Lemma 5 establishes this alignment when w≠0_n and θ(w,w*)∈(0,π), with a further lower bound under bounded ∥v∥ and nonzero ∥w∥.
Failure of identity STE: With the identity estimator, descent and convergence fail near local minimizers satisfying θ(w,w*)=π and the specified v when 1_m^T v*≠0 and m>1.The instability result is stated specifically for Algorithm 1 near those local minimizers.
Convergence: Coarse-gradient descent with ReLU converges subsequentially to a saddle point or minimizer, and analogous convergence holds for the clipped-ReLU STE.These conclusions follow from the gradient-alignment and stationary-point characterizations developed in the preceding lemmas.

F. CONVERGENCE TO GLOBAL MINIMIZERS

Under specified initialization conditions and a sufficiently small positive learning rate, vanilla or clipped ReLU STE training preserves the favorable region and converges to a global minimizer. The proof uses induction and convergence of the iterates.

Global convergence: Vanilla or clipped ReLU STE yields convergence to global optima under the initialization assumptions of Theorem 1.The initialization requires (v0)⊤v∗ > 0 and θ(w0, w∗) < π/2, along with the remaining stated condition.
Global convergence: For sufficiently positive learning rate η, the iterates preserve (vt)⊤v∗ > 0 and θ(wt, w∗) < π/2 for all t > 0.These invariants hold when the initialization satisfies the additional conditions in Theorem 2.
Proof strategy: The proof proceeds by induction and concludes that the convergent sequence {(vt, wt)} can only converge to a global minimizer.The argument uses the stated lemmas to characterize the expected gradients and establish the needed bounds.