Understanding Straight-Through Estimator in Training Activation Quantized Neural Nets

1 INTRODUCTION . Deep neural networks ( DNN ) have achieved the remarkable success in many machine learning applications such as computer vision ( Krizhevsky et al. , 2012 ; Ren et al. , 2015 ) , natural language processing ( Collobert & Weston , 2008 ) and reinforcement learning ( Mnih et al. , 2015 ; Silver et al. , 2016 ) . However , the deployment of DNN typically require hundreds of megabytes of memory storage for the trainable full-precision floating-point parameters , and billions of floating-point operations to make a single inference . To achieve substantial memory savings and energy efficiency at inference time , many recent efforts have been made to the training of coarsely quantized DNN , meanwhile maintaining the performance of their float counterparts ( Courbariaux et al. , 2015 ; Rastegari et al. , 2016 ; Cai et al. , 2017 ; Hubara et al. , 2018 ; Yin et al. , 2018b ) . Training fully quantized DNN amounts to solving a very challenging optimization problem . It calls for minimizing a piecewise constant and highly nonconvex empirical risk function f ( w ) subject to a discrete set-constraint w ∈ Q that characterizes the quantized weights . In particular , weight quantization of DNN have been extensively studied in the literature ; see for examples ( Li et al. , 2016 ; Zhu et al. , 2016 ; Li et al. , 2017 ; Yin et al. , 2016 ; 2018a ; Hou & Kwok , 2018 ; He et al. , 2018 ; Li & Hao , 2018 ) . On the other hand , the gradient ∇f ( w ) in training activation quantized DNN is almost everywhere ( a.e . ) zero , which makes the standard back-propagation inapplicable . The arguably most effective way around this issue is nothing but to construct a non-trivial search direction by properly modifying the chain rule . Specifically , one can replace the a.e . zero derivative of quantized activation function composited in the chain rule with a related surrogate . This proxy derivative used in the backward pass only is referred as the straight-through estimator ( STE ) ( Bengio et al. , 2013 ) . In the same paper , Bengio et al . ( 2013 ) proposed an alternative approach based on stochastic neurons . In addition , Friesen & Domingos ( 2017 ) proposed the feasible target propagation algorithm for learning hard-threshold ( or binary activated ) networks ( Lee et al. , 2015 ) via convex combinatorial optimization . 1.1 RELATED WORKS . The idea of STE originates to the celebrated perceptron algorithm ( Rosenblatt , 1957 ; 1962 ) in 1950s for learning single-layer perceptrons . The perceptron algorithm essentially does not calculate the “ gradient ” through the standard chain rule , but instead through a modified chain rule in which the derivative of identity function serves as the proxy of the original derivative of binary output function 1 { x > 0 } . Its convergence has been extensive discussed in the literature ; see for examples , ( Widrow & Lehr , 1990 ; Freund & Schapire , 1999 ) and the references therein . Hinton ( 2012 ) extended this idea to train multi-layer networks with binary activations ( a.k.a . binary neuron ) , namely , to backpropagate as if the activation had been the identity function . Bengio et al . ( 2013 ) proposed a STE variant which uses the derivative of the sigmoid function instead . In the training of DNN with weights and activations constrained to ±1 , ( Hubara et al. , 2016 ) substituted the derivative of the signum activation function with 1 { |x|≤1 } in the backward pass , known as the saturated STE . Later the idea of STE was readily employed to the training of DNN with general quantized ReLU activations ( Hubara et al. , 2018 ; Zhou et al. , 2016 ; Cai et al. , 2017 ; Choi et al. , 2018 ; Yin et al. , 2018b ) , where some other proxies took place including the derivatives of vanilla ReLU and clipped ReLU . Despite all the empirical success of STE , there is very limited theoretical understanding of it in training DNN with stair-case activations . Goel et al . ( 2018 ) considers leaky ReLU activation of a one-hidden-layer network . They showed the convergence of the so-called Convertron algorithm , which uses the identity STE in the backward pass through the leaky ReLU layer . Other similar scenarios , where certain layers are not desirable for back-propagation , have been brought up recently by ( Wang et al. , 2018 ) and ( Athalye et al. , 2018 ) . The former proposed an implicit weighted nonlocal Laplacian layer as the classifier to improve the generalization accuracy of DNN . In the backward pass , the derivative of a pre-trained fullyconnected layer was used as a surrogate . To circumvent adversarial defense ( Szegedy et al. , 2013 ) , ( Athalye et al. , 2018 ) introduced the backward pass differentiable approximation , which shares the same spirit as STE , and successfully broke defenses at ICLR 2018 that rely on obfuscated gradients . 1.2 MAIN CONTRIBUTIONS . Throughout this paper , we shall refer to the “ gradient ” of loss function w.r.t . the weight variables through the STE-modified chain rule as coarse gradient . Since the backward and forward passes do not match , the coarse gradient is certainly not the gradient of loss function , and it is generally not the gradient of any function . Why searching in its negative direction minimizes the training loss , as this is not the standard gradient descent algorithm ? Apparently , the choice of STE is non-unique , then what makes a good STE ? From the optimization perspective , we take a step towards understanding STE in training quantized ReLU nets by attempting these questions . On the theoretical side , we consider three representative STEs for learning a two-linear-layer network with binary activation and Gaussian data : the derivatives of the identity function ( Rosenblatt , 1957 ; Hinton , 2012 ; Goel et al. , 2018 ) , vanilla ReLU and the clipped ReLUs ( Cai et al. , 2017 ; Hubara et al. , 2016 ) . We adopt the model of population loss minimization ( Brutzkus & Globerson , 2017 ; Tian , 2017 ; Li & Yuan , 2017 ; Du et al. , 2018 ) . For the first time , we prove that proper choices of STE give rise to training algorithms that are descent . Specifically , the negative expected coarse gradients based on STEs of the vanilla and clipped ReLUs are provably descent directions for the minimizing the population loss , which yield monotonically decreasing energy in the training . In contrast , this is not true for the identity STE . We further prove that the corresponding training algorithm can be unstable near certain local minima , because the coarse gradient may simply not vanish there . Complementary to the analysis , we examine the empirical performances of the three STEs on MNIST and CIFAR-10 classifications with general quantized ReLU . While both vanilla and clipped ReLUs work very well on the relatively shallow LeNet-5 , clipped ReLU STE is arguably the best for the deeper VGG-11 and ResNet-20 . In our CIFAR experiments in section 4.2 , we observe that the training using identity or ReLU STE can be unstable at good minima and repelled to an inferior one with substantially higher training loss and decreased generalization accuracy . This is an implication that poor STEs generate coarse gradients incompatible with the energy landscape , which is consistent with our theoretical finding about the identity STE . To our knowledge , convergence guarantees of perceptron algorithm ( Rosenblatt , 1957 ; 1962 ) and Convertron algorithm ( Goel et al. , 2018 ) were proved for the identity STE . It is worth noting that Convertron ( Goel et al. , 2018 ) makes weaker assumptions than in this paper . These results , however , do not generalize to the network with two trainable layers studied here . As aforementioned , the identity STE is actually a poor choice in our case . Moreover , it is not clear if their analyses can be extended to other STEs . Similar to Convertron with leaky ReLU , the monotonicity of quantized activation function plays a role in coarse gradient descent . Indeed , all three STEs considered here exploit this property . But this is not the whole story . A great STE like the clipped ReLU matches quantized ReLU at the extrema , otherwise the instability/incompatibility issue may arise . Organization . In section 2 , we study the energy landscape of a two-linear-layer network with binary activation and Gaussian data . We present the main results and sketch the mathematical analysis for STE in section 3 . In section 4 , we compare the empirical performances of different STEs in 2- bit and 4-bit activation quantization , and report the instability phenomena of the training algorithms associated with poor STEs observed in CIFAR experiments . Due to space limitation , all the technical proofs as well as some figures are deferred to the appendix . Notations . ‖ · ‖ denotes the Euclidean norm of a vector or the spectral norm of a matrix . 0n ∈ Rn represents the vector of all zeros , whereas 1n ∈ Rn the vector of all ones . In is the identity matrix of order n. For any w , z ∈ Rn , w > z = 〈w , z〉 = ∑ i wizi is their inner product . w z denotes the Hadamard product whose ith entry is given by ( w z ) i = wizi . 2 LEARNING TWO-LINEAR-LAYER CNN WITH BINARY ACTIVATION . We consider a model similar to ( Du et al. , 2018 ) that outputs the prediction y ( Z , v , w ) : = m∑ i=1 viσ ( Z > i w ) = v > σ ( Zw ) for some input Z ∈ Rm×n . Here w ∈ Rn and v ∈ Rm are the trainable weights in the first and second linear layer , respectively ; Z > i denotes the ith row vector of Z ; the activation function σ acts component-wise on the vector Zw , i.e. , σ ( Zw ) i = σ ( ( Zw ) i ) = σ ( Z > i w ) . The first layer serves as a convolutional layer , where each row Z > i can be viewed as a patch sampled from Z and the weight filter w is shared among all patches , and the second linear layer is the classifier . The label is generated according to y∗ ( Z ) = ( v∗ ) > σ ( Zw∗ ) for some true ( non-zero ) parameters v∗ and w∗ . Moreover , we use the following squared sample loss ` ( v , w ; Z ) : = 1 2 ( y ( Z , v , w ) − y∗ ( Z ) ) 2 = 1 2 ( v > σ ( Zw ) − y∗ ( Z ) ) 2 . ( 1 ) Unlike in ( Du et al. , 2018 ) , the activation function σ here is not ReLU , but the binary function σ ( x ) = 1 { x > 0 } . We assume that the entries of Z ∈ Rm×n are i.i.d . sampled from the Gaussian distribution N ( 0 , 1 ) ( Zhong et al. , 2017 ; Brutzkus & Globerson , 2017 ) . Since ` ( v , w ; Z ) = ` ( v , w/c ; Z ) for any scalar c > 0 , without loss of generality , we take ‖w∗‖ = 1 and cast the learning task as the following population loss minimization problem : min v∈Rm , w∈Rn f ( v , w ) : = EZ [ ` ( v , w ; Z ) ] , ( 2 ) where the sample loss ` ( v , w ; Z ) is given by ( 1 ) .

Summary: The paper examines the use of STE for learning simple one-layer convolutional networks with binary activations and non overlapping patches. In this setting, the gradients are 0 almost everywhere (gradient of sign(x)) hence it is not clear how to use gradient descent. The approach studied here is to instead use the gradient of an alternative function such as ReLU or identity which is not always 0. The authors prove that if the ReLU's gradient is used then under gaussian distribution, the algorithm will converge to the local minimas/saddle points of the expected squared loss. they also show that the same does not hold for the identity's gradient.