ProxQuant: Quantized Neural Networks via Proximal Operators

1 INTRODUCTION . Deep neural networks ( DNNs ) have achieved impressive results in various machine learning tasks ( Goodfellow et al. , 2016 ) . High-performance DNNs typically have over tens of layers and millions of parameters , resulting in a high memory usage and a high computational cost at inference time . However , these networks are often desired in environments with limited memory and computational power ( such as mobile devices ) , in which case we would like to compress the network into a smaller , faster network with comparable performance . A popular way of achieving such compression is through quantization – training networks with lowprecision weights and/or activation functions . In a quantized neural network , each weight and/or activation can be representable in k bits , with a possible codebook of negligible additional size compared to the network itself . For example , in a binary neural network ( k = 1 ) , the weights are restricted to be in { ±1 } . Compared with a 32-bit single precision float , a quantized net reduces the memory usage to k/32 of a full-precision net with the same architecture ( Han et al. , 2015 ; Courbariaux et al. , 2015 ; Rastegari et al. , 2016 ; Hubara et al. , 2017 ; Zhou et al. , 2016 ; Zhu et al. , 2016 ) . In addition , the structuredness of the quantized weight matrix can often enable faster matrixvector product , thereby also accelerating inference ( Hubara et al. , 2017 ; Han et al. , 2016 ) . Typically , training a quantized network involves ( 1 ) the design of a quantizer q that maps a full-precision parameter to a k-bit quantized parameter , and ( 2 ) the straight-through gradient method ( Courbariaux et al. , 2015 ) that enables back-propagation from the quantized parameter back onto the original full-precision parameter , which is critical to the success of quantized network training . With quantizer q , an iterate of the straight-through gradient method ( see Figure 1a ) proceeds Code available at https : //github.com/allenbai01/ProxQuant . as θt+1 = θt − ηt∇̃L ( θ ) |θ=q ( θt ) , and q ( θ̂ ) ( for the converged θ̂ ) is taken as the output model . For training binary networks , choosing q ( · ) = sign ( · ) gives the BinaryConnect method ( Courbariaux et al. , 2015 ) . Though appealingly simple and empirically effective , it is information-theoretically rather mysterious why the straight-through gradient method works well , at least in the binary case : while the goal is to find a parameter θ ∈ { ±1 } d with low loss , the algorithm only has access to stochastic gradients at { ±1 } d. As this is a discrete set , a priori , gradients in this set do not necessarily contain any information about the function values . Indeed , a simple one-dimensional example ( Figure 1b ) shows that BinaryConnect fails to find the minimizer of fairly simple convex Lipschitz functions in { ±1 } , due to a lack of gradient information in between . In this paper , we formulate the problem of model quantization as a regularized learning problem and propose to solve it with a proximal gradient method . Our contributions are summarized as follows . • We present a unified framework for defining regularization functionals that encourage binary , ternary , and multi-bit quantized parameters , through penalizing the distance to quantized sets ( see Section 3.1 ) . For binary quantization , the resulting regularizer is aW -shaped non-smooth regularizer , which shrinks parameters towards either −1 or 1 in the same way that the L1 norm regularization shrinks parameters towards 0 . • We propose training quantized networks using PROXQUANT ( Algorithm 1 ) — a stochastic proximal gradient method with a homotopy scheme . Compared with the straightthrough gradient method , PROXQUANT has access to additional gradient information at non-quantized points , which avoids the problem in Figure 1b and its homotopy scheme prevents potential overshoot early in the training ( Section 3.2 ) . • We demonstrate the effectiveness and flexibility of PROXQUANT through systematic experiments on ( 1 ) image classification with ResNets ( Section 4.1 ) ; ( 2 ) language modeling with LSTMs ( Section 4.2 ) . The PROXQUANT method outperforms the state-of-the-art results on binary quantization and is comparable with the state-of-the-art on ternary and multi-bit quantization . • We perform a systematic theoretical study of quantization algorithms , showing that our PROXQUANT ( standard prox-gradient method ) converges to stataionary points under mild smoothness assumptions ( Section 5.1 ) , where as lazy prox-gradient method such as BinaryRelax ( Yin et al. , 2018 ) fails to converge in general ( Section 5.2 ) . Further , we show that BinaryConnect has a very stringent condition to converge to any fixed point ( Section 5.3 ) , which we verify through a sign change experiment ( Appendix C ) . 1.1 PRIOR WORK . Methodologies Han et al . ( 2015 ) propose Deep Compression , which compresses a DNN via sparsification , nearest-neighbor clustering , and Huffman coding . This architecture is then made into a specially designed hardware for efficient inference ( Han et al. , 2016 ) . In a parallel line of work , Courbariaux et al . ( 2015 ) propose BinaryConnect that enables the training of binary neural networks , and Li & Liu ( 2016 ) ; Zhu et al . ( 2016 ) extend this method into ternary quantization . Training and inference on quantized nets can be made more efficient by also quantizing the activation ( Hubara et al. , 2017 ; Rastegari et al. , 2016 ; Zhou et al. , 2016 ) , and such networks have achieved impressive performance on large-scale tasks such as ImageNet classification ( Rastegari et al. , 2016 ; Zhu et al. , 2016 ) and object detection ( Yin et al. , 2016 ) . In the NLP land , quantized language models have been successfully trained using alternating multi-bit quantization ( Xu et al. , 2018 ) . Theories Li et al . ( 2017 ) prove the convergence rate of stochastic rounding and BinaryConnect on convex problems and demonstrate the advantage of BinaryConnect over stochastic rounding on non-convex problems . Anderson & Berg ( 2017 ) demonstrate the effectiveness of binary networks through the observation that the angles between high-dimensional vectors are approximately preserved when binarized , and thus high-quality feature extraction with binary weights is possible . Ding et al . ( 2018 ) show a universal approximation theorem for quantized ReLU networks . Principled methods Sun & Sun ( 2018 ) perform model quantization through a Wasserstein regularization term and minimize via the adversarial representation , similar as in Wasserstein GANs ( Arjovsky et al. , 2017 ) . Their method has the potential of generalizing to other generic requirements on the parameter , but might be hard to tune due to the instability of the inner maximization problem . Prior to our work , a couple of proximal or regularization based quantization algorithms were proposed as alternatives to the straight-through gradient method , which we now briefly review and compare with . ( Yin et al. , 2018 ) propose BinaryRelax , which corresponds to a lazy proximal gradient descent . ( Hou et al. , 2017 ; Hou & Kwok , 2018 ) propose a proximal Newton method with a diagonal approximate Hessian . Carreira-Perpinán ( 2017 ) ; Carreira-Perpinán & Idelbayev ( 2017 ) formulate quantized network training as a constrained optimization problem and propose to solve them via augmented Lagrangian methods . Our algorithm is different with all the aformentioned work in using the non-lazy and “ soft ” proximal gradient descent with a choice of either ` 1 or ` 2 regularization , whose advantage over lazy prox-gradient methods is demonstrated both theoretically ( Section 5 ) and experimentally ( Section 4.1 and Appendix C ) . 2 PRELIMINARIES . The optimization difficulty of training quantized models is that they involve a discrete parameter space and hence efficient local-search methods are often prohibitive . For example , the problem of training a binary neural network is to minimize L ( θ ) for θ ∈ { ±1 } d. Projected SGD on this set will not move unless with an unreasonably large stepsize ( Li et al. , 2017 ) , whereas greedy nearestneighbor search requires d forward passes which is intractable for neural networks where d is on the order of millions . Alternatively , quantized training can also be cast as minimizing L ( q ( θ ) ) for θ ∈ Rd and an appropriate quantizer q that maps a real vector to a nearby quantized vector , but θ 7→ q ( θ ) is often non-differentiable and piecewise constant ( such as the binary case q ( · ) = sign ( · ) ) , and thus back-propagation through q does not work . 2.1 THE STRAIGHT-THROUGH GRADIENT METHOD . The pioneering work of BinaryConnect ( Courbariaux et al. , 2015 ) proposes to solve this problem via the straight-through gradient method , that is , propagate the gradient with respect to q ( θ ) unaltered to θ , i.e . to let ∂L∂θ : = ∂L ∂q ( θ ) . One iterate of the straight-through gradient method ( with the SGD optimizer ) is θt+1 = θt − ηt∇̃L ( θ ) |θ=q ( θt ) . This enables the real vector θ to move in the entire Euclidean space , and taking q ( θ ) at the end of training gives a valid quantized model . Such a customized back-propagation rule yields good empirical performance in training quantized nets and has thus become a standard practice ( Courbariaux et al. , 2015 ; Zhu et al. , 2016 ; Xu et al. , 2018 ) . However , as we have discussed , it is information theoretically unclear how the straight-through method works , and it does fail on very simple convex Lipschitz functions ( Figure 1b ) . 2.2 STRAIGHT-THROUGH GRADIENT AS LAZY PROJECTION . Our first observation is that the straight-through gradient method is equivalent to a dual-averaging method , or a lazy projected SGD ( Xiao , 2010 ) . In the binary case , we wish to minimize L ( θ ) over Q = { ±1 } d , and the lazy projected SGD proceeds as { θ̃t = ProjQ ( θt ) = sign ( θt ) = q ( θt ) , θt+1 = θt − ηt∇̃L ( θ̃t ) . ( 1 ) Written compactly , this is θt+1 = θt − ηt∇̃L ( θ ) |θ=q ( θt ) , which is exactly the straight-through gradient method : take the gradient at the quantized vector and perform the update on the original real vector . 2.3 PROJECTION AS A LIMITING PROXIMAL OPERATOR . We take a broader point of view that a projection is also a limiting proximal operator with a suitable regularizer , to allow more generality and to motivate our proposed algorithm . Given any set Q , one could identify a regularizer R : Rd → R≥0 such that the following hold : R ( θ ) = 0 , ∀θ ∈ Q and R ( θ ) > 0 , ∀θ /∈ Q . ( 2 ) In the case Q = { ±1 } d for example , one could take R ( θ ) = Rbin ( θ ) = d∑ j=1 min { |θj − 1| , |θj + 1| } . ( 3 ) The proximal operator ( or prox operator ) ( Parikh & Boyd , 2014 ) with respect to R and strength λ > 0 is proxλR ( θ ) : = arg min θ̃∈Rd { 1 2 ∥∥∥θ̃ − θ ∥∥∥ 2 2 + λR ( θ̃ ) } . In the limiting case λ =∞ , the argmin has to satisfy R ( θ ) = 0 , i.e . θ ∈ Q , and the prox operator is to minimize ‖θ − θ0‖22 over θ ∈ Q , which is the Euclidean projection onto Q . Hence , projection is also a prox operator with λ = ∞ , and the straight-through gradient estimate is equivalent to a lazy proximal gradient descent with and λ =∞ . While the prox operator with λ =∞ correponds to “ hard ” projection onto the discrete set Q , when λ < ∞ it becomes a “ soft ” projection that moves towards Q . Compared with the hard projection , a finite λ is less aggressive and has the potential advantage of avoiding overshoot early in training . Further , as the prox operator does not strictly enforce quantizedness , it is in principle able to query the gradients at every point in the space , and therefore has access to more information than the straight-through gradient method .

Summary: This paper proposed ProxQuant method to train neural networks with quantized weights. ProxQuant relax the quantization constraint to a continuous regularizer and then solve the optimization problem with proximal gradient method. The authors argues that previous solvers straight through estimator (STE) in BinaryConnect (Courbariaux et al. 2015) may not converge, and the proposed ProxQuant is better.