To make deep neural networks feasible in resource-constrained environments (such as mobile devices), it is beneficial to quantize models by using low-precision weights. One common technique for quantizing neural networks is the straight-through gradient method, which enables back-propagation through the quantization mapping. Despite its empirical success, little is understood about why the straight-through gradient method works.
 Building upon a novel observation that the straight-through gradient method is in fact identical to the well-known Nesterov’s dual-averaging algorithm on a quantization constrained optimization problem, we propose a more principled alternative approach, called ProxQuant , that formulates quantized network training as a regularized learning problem instead and optimizes it via the prox-gradient method. ProxQuant does back-propagation on the underlying full-precision vector and applies an efficient prox-operator in between stochastic gradient steps to encourage quantizedness. For quantizing ResNets and LSTMs, ProxQuant outperforms state-of-the-art results on binary quantization and is on par with state-of-the-art on multi-bit quantization. We further perform theoretical analyses showing that ProxQuant converges to stationary points under mild smoothness assumptions, whereas variants such as lazy prox-gradient method can fail to converge in the same setting. Deep neural networks (DNNs) have achieved impressive results in various machine learning tasks BID6 . 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 BID7 BID4 BID19 BID13 BID26 BID27 . In addition, the structuredness of the quantized weight matrix can often enable faster matrixvector product, thereby also accelerating inference BID13 .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 BID4 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 FIG7 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 BID4 .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.rL(q(✓t)) DISPLAYFORM0 0 B N 6 N u 6 N R + P F e F 2 O 5 o z V z j H 6 A e P t E 7 v e l p c = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " k 8 o W Y E M 4 F 9 J U h g q 5 4 4 1 T N Q W 4 X m E = " > A A A B / n i c d V D L S g M x F M 3 U V 6 2 v U X H l J l g E V 2 W m t e 1 0 V 3 D j s o J 9 Q K e U T C Z t Q z M P k j t C G Q r + i h s X i r j 1 O 9 z 5 N 2 b a C i p 6 I O R w z r 3 k 5 H i x 4 A o s 6 8 P I r a 1 v b G 7 l t w s 7 u 3 v 7 DISPLAYFORM1 0 B N 6 N u 6 N R + P F e F 2 O 5 o z V z j H 6 A e P t E 7 v e l p c = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " k 8 o W Y E M 4 F 9 J U h g q 5 4 4 1 T N Q W 4 X m E = " > A A A B / n i c d V D L S g M x F M 3 U V 6 2 v U X H l J l g E V 2 W m t e 1 0 V 3 D j s o J 9 Q K e U T C Z t Q z M P k j t C G Q r + i h s X i r j 1 O 9 z 5 N 2 b a C i p 6 I O R w z r 3 k 5 H i x 4 A o s 6 8 P I r a 1 v b G 7 l t w s 7 u 3 v 7 DISPLAYFORM2 0 B N 6 N u 6 N R + P F e F 2 O 5 o z V z j H 6 A e P t E 7 v e l p c = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " k 8 o W Y E M 4 F 9 J U h g q 5 4 4 1 T N Q W 4 X m E = " > A A A B / n i c d V D L S g M x F M 3 U V 6 2 v U X H l J l g E V 2 W m t e 1 0 V 3 D j s o J 9 Q K e U T C Z t Q z M P k j t C G Q r + i h s X i r j 1 O 9 z 5 N 2 b a C i p 6 I O R w z r 3 k 5 H i x 4 A o s 6 8 P I r a 1 v b G 7 l t w s 7 u 3 v 7 DISPLAYFORM3 0 B N 6 N u 6 N R + P F e F 2 O 5 o z V z j H 6 A e P t E 7 v e l p c = < / l a t e x i t > q(✓t) DISPLAYFORM4 a E H 9 I S e j X v j 0 X g x X h e l K 8 a y 5 w j 9 g P H 2 C a M i n A U = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " DISPLAYFORM5 a E H 9 I S e j X v j 0 X g x X h e l K 8 a y 5 w j 9 g P H 2 C a M i n A U = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " DISPLAYFORM6 a E H 9 I S e j X v j 0 X g x X h e l K 8 a y 5 w j 9 g P H 2 C a M i n A U = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " DISPLAYFORM7 a E H 9 I S e j X v j 0 X g x X h e l K 8 a y 5 w j 9 g P H 2 C a M i n A U = < / l a t e x i t > rL(✓t) < l a t e x i t s h a 1 _ b a s e 6 4 = " k M y 6 L o 8 X P y through method computes the gradient at the quantized vector and performs the update at the original real vector; PROXQUANT performs a gradient update at the current real vector followed by a prox step which encourages quantizedness. (b) A two-function toy failure case for BinaryConnect. The two functions are f1(x) = |x + 0.5| − 0.5 (blue) and f−1(x) = |x − 0.5| − 0.5 (orange). The derivatives of f1 and f−1 coincide at {−1, 1}, so any algorithm that only uses this information will have identical behaviors on these two functions. However, the minimizers in {±1} are x 1 = −1 and x −1 = 1, so the algorithm must fail on one of them. DISPLAYFORM8 DISPLAYFORM9 DISPLAYFORM10 DISPLAYFORM11 DISPLAYFORM12 DISPLAYFORM13 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 a W -shaped non-smooth regularizer, which shrinks parameters towards either −1 or 1 in the same way that the L 1 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 FIG7 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 BID25 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). In this paper, we propose and experiment with the PROXQUANT method for training quantized networks. Our results demonstrate that PROXQUANT offers a powerful alternative to the straightthrough gradient method and has theoretically better convergence properties. For future work, it would be of interest to propose alternative regularizers for ternary and multi-bit PROXQUANT and experiment with our method on larger tasks.
A principled framework for model quantization using the proximal gradient method, with empirical evaluation and theoretical convergence analyses. Proposes ProxQuant method to train neural networks with quantized weights. Proposes solving binary nets and its variants using proximal gradient descent.