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 proposes a new approach to learning quantized deep neural networks, which overcome some of the drawbacks of previous methods, namely the lack of understanding of why straight-through gradient works and its optimization instability. The core of the proposal is the use of quantization-encouraging regularization, and the derivation of the corresponding proximity operators. Building on that core, the rest of the approach is reasonably standard, based on stochastic proximal gradient descent, with a homotopy scheme.