PROXQUANT: QUANTIZED NEURAL NETWORKS VIA
Year: 2019

Abstract:
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 straightthrough 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 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.

heading: 1 INTRODUCTION
text: 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).

heading: 1.1 PRIOR WORK
text: 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).

heading: 2 PRELIMINARIES
text: 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.

heading: 2.1 THE STRAIGHT-THROUGH GRADIENT METHOD
text: 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).

heading: 2.2 STRAIGHT-THROUGH GRADIENT AS LAZY PROJECTION
text: 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.

heading: 2.3 PROJECTION AS A LIMITING PROXIMAL OPERATOR
text: 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.

heading: 3 QUANTIZED NET TRAINING VIA REGULARIZED LEARNING
text: We propose the PROXQUANT algorithm, which adds a quantization-inducing regularizer onto the loss and optimizes via the (non-lazy) prox-gradient method with a finite λ. The prototypical version of PROXQUANT is described in Algorithm 1.
Algorithm 1 PROXQUANT: Prox-gradient method for quantized net training Require: Regularizer R that induces desired quantizedness, initialization θ0, learning rates {ηt}t≥0, regularization strengths {λt}t≥0 while not converged do
Perform the prox-gradient step
θt+1 = proxηtλtR
( θt − ηt∇̃L(θt) ) . (4)
The inner SGD step in eq. (4) can be replaced by any preferred stochastic optimization method such as Momentum SGD or Adam (Kingma & Ba, 2014).
end while
Compared to usual full-precision training, PROXQUANT only adds a prox step after each stochastic gradient step, hence can be implemented straightforwardly upon existing full-precision training. As the prox step does not need to know how the gradient step is performed, our method adapts to other stochastic optimizers as well such as Adam.
In the remainder of this section, we define a flexible class of quantization-inducing regularizers through “distance to the quantized set”, derive efficient algorithms of their corresponding prox operator, and propose a homotopy method for choosing the regularization strengths. Our regularization perspective subsumes most existing algorithms for model-quantization (e.g.,(Courbariaux et al., 2015; Han et al., 2015; Xu et al., 2018)) as limits of certain regularizers with strength λ→∞. Our proposed method can be viewed as a principled generalization of these methods to λ < ∞ with a non-lazy prox operator.

heading: 3.1 REGULARIZATION FOR MODEL QUANTIZATION
text: LetQ ⊂ Rd be a set of quantized parameter vectors. An ideal regularizer for quantization would be to vanish on Q and reflect some type of distance to Q when θ /∈ Q. To achieve this, we propose L1 and L2 regularizers of the form
R(θ) = inf θ0∈Q ‖θ − θ0‖1 or R(θ) = inf θ0∈Q ‖θ − θ0‖22 . (5)
This is a highly flexible framework for designing regularizers, as one could specify any Q and choose between L1 and L2. Specifically, Q encodes certain desired quantization structure. By appropriately choosingQ, we can specify which part of the parameter vector to quantize1, the number of bits to quantize to, whether we allow adaptively-chosen quantization levels and so on. The choice between {L1, L2} will encourage {“hard”,“soft”} quantization respectively, similar as in standard regularized learning (Tibshirani, 1996).
In the following, we present a few examples of regularizers under our framework eq. (5) which induce binary weights, ternary weights and multi-bit quantization. We will also derive efficient algorithms (or approximation heuristics) for solving the prox operators corresponding to these regularizers, which generalize the projection operators used in the straight-through gradient algorithms.
Binary neural nets In a binary neural net, the entries of θ are in {±1}. A natural choice would be taking Q = {−1, 1}d. The resulting L1 regularizer is
R(θ) = inf θ0∈{±1}d
‖θ − θ0‖1 = d∑
j=1
inf [θ0]j∈{±1}
|θj − [θ0]j |
=
d∑
j=1
min {|θj − 1|, |θj + 1|} = ‖θ − sign(θ)‖1 .
(6)
This is exactly the binary regularizer Rbin that we discussed earlier in eq. (3). Figure 2 plots the W-shaped one-dimensional component of Rbin from which we see its effect for inducing {±1} quantization in analog to L1 regularization for inducing exact sparsity.
1Empirically, it is advantageous to keep the biases of each layers and the BatchNorm layers at full-precision, which is often a negligible fraction, say 1/ √ d of the total number of parameters
3 2 1 0 1 2 3
0.00
0.25
0.50
0.75
1.00
1.25
1.50
1.75
2.00
The prox operator with respect to Rbin, despite being a non-convex optimization problem, admits a simple analytical solution: proxλRbin(θ) = SoftThreshold(θ, sign(θ), λ)
= sign(θ) + sign(θ − sign(θ)) [|θ − sign(θ)| − λ]+. (7)
We note that the choice of the L1 version is not unique: the squared L2 version works as well, whose prox operator is given by (θ + λ sign(θ))/(1 + λ). See Appendix A.1 for the derivation of these prox operators and the definition of the soft thresholding operator.
(8) The squared L2 regularizer for this structure is
Rk−bit(θ) = inf α∈Rk,B∈{±1}d×k
‖θ −Bα‖22 , (9)
which is also the alternating minimization objective in (Xu et al., 2018).
We now derive the prox operator for the regularizer eq. (9). For any θ, we have
proxλRk−bit(θ) = arg min θ̃
{ 1
2
∥∥∥θ̃ − θ ∥∥∥ 2
2 + λ inf α∈Rk,B∈{±1}d×k
∥∥∥θ̃ −Bα ∥∥∥ 2
2
}
= arg min θ̃ inf α∈Rk,B∈{±1}d×k
{ 1
2
∥∥∥θ̃ − θ ∥∥∥ 2
2 + λ
∥∥∥θ̃ −Bα ∥∥∥ 2
2
} .
(10)
This is a joint minimization problem in (θ̃, B, α), and we adopt an alternating minimization schedule to solve it:
(1) Minimize over θ̃ given (B,α), which has a closed-form solution θ̃ = θ+2λBα1+2λ .
(2) Minimize over (B,α) given θ̃, which does not depend on θ0, and can be done via calling the alternating quantizer of (Xu et al., 2018): Bα = qalt(θ̃).
Together, the prox operator generalizes the alternating minimization procedure in (Xu et al., 2018), as λ governs a trade-off between quantization and closeness to θ. To see that this is a strict generalization, note that for any λ the solution of eq. (10) will be an interpolation between the input θ and its Euclidean projection to Q. As λ→ +∞, the prox operator collapses to the projection.
Ternary quantization Ternary quantization is a variant of 2-bit quantization, in which weights are constrained to be in {−α, 0, β} for real values α, β > 0. We defer the derivation of the ternary prox operator into Appendix A.2.

heading: 3.2 HOMOTOPY METHOD FOR REGULARIZATION STRENGTH
text: Recall that the larger λt is, the more aggressive θt+1 will move towards the quantized set. An ideal choice would be to (1) force the net to be exactly quantized upon convergence, and (2) not be too aggressive such that the quantized net at convergence is sub-optimal.
We let λt be a linearly increasing sequence, i.e. λt := λ · t for some hyper-parameter λ > 0 which we term as the regularization rate. With this choice, the stochastic gradient steps will start off close to full-precision training and gradually move towards exact quantizedness, hence the name “homotopy method”. The parameter λ can be tuned by minimizing the validation loss, and controls the aggressiveness of falling onto the quantization constraint. There is nothing special about the linear increasing scheme, but it is simple enough and works well as we shall see in the experiments.

heading: 4 EXPERIMENTS
text: We evaluate the performance of PROXQUANT on two tasks: image classification with ResNets, and language modeling with LSTMs. On both tasks, we show that the default straight-through gradient method is not the only choice, and our PROXQUANT can achieve the same and often better results.

heading: 4.1 IMAGE CLASSIFICATION ON CIFAR-10
text: Problem setup We perform image classification on the CIFAR-10 dataset, which contains 50000 training images and 10000 test images of size 32x32. We apply a commonly used data augmentation strategy (pad by 4 pixels on each side, randomly crop to 32x32, do a horizontal flip with probability 0.5, and normalize). Our models are ResNets (He et al., 2016) of depth 20, 32, 44, and 56 with weights quantized to binary or ternary.
Method We use PROXQUANT with regularizer eq. (3) in the binary case and eqs. (15) and (16) in the ternary case, which we respectively denote as PQ-B and PQ-T. We use the homotopy method λt = λ · t with λ = 10−4 as the regularization strength and Adam with constant learning rate 0.01 as the optimizer.
We compare with BinaryConnect (BC) for binary nets and Trained Ternary Quantization (TTQ) (Zhu et al., 2016) for ternary nets. For BinaryConnect, we train with the recommended Adam optimizer with learning rate decay (Courbariaux et al., 2015) (initial learning rate 0.01, multiply by 0.1 at epoch 81 and 122), which we find leads to the best result for BinaryConnect. For TTQ we compare with the reported results in (Zhu et al., 2016).
For binary quantization, both BC and our PROXQUANT are initialized at the same pre-trained fullprecision nets (warm-start) and trained for 300 epochs for fair comparison. For both methods, we perform a hard quantization θ 7→ q(θ) at epoch 200 and keeps training till the 300-th epoch to stabilize the BatchNorm layers. We compare in addition the performance drop relative to full precision nets of BinaryConnect, BinaryRelax (Yin et al., 2018), and our PROXQUANT.
Result The top-1 classification errors for binary quantization are reported in Table 1. Our PROXQUANT consistently yields better results than BinaryConnect. The performance drop of PROXQUANT relative to full-precision nets is about 1%, better than BinaryConnect by 0.2% on average and significantly better than the reported result of BinaryRelax.
Results and additional details for ternary quantization are deferred to Appendix B.1.

heading: 4.2 LANGUAGE MODELING WITH LSTMS
text: Problem setup We perform language modeling with LSTMs Hochreiter & Schmidhuber (1997) on the Penn Treebank (PTB) dataset (Marcus et al., 1993), which contains 929K training tokens, 73K validation tokens, and 82K test tokens. Our model is a standard one-hidden-layer LSTM with embedding dimension 300 and hidden dimension 300. We train quantized LSTMs with the encoder, transition matrix, and the decoder quantized to k-bits for k ∈ {1, 2, 3}. The quantization is performed in a row-wise fashion, so that each row of the matrix has its own codebook {α1, . . . , αk}.
Method We compare our multi-bit PROXQUANT (eq. (10)) to the state-of-the-art alternating minimization algorithm with straight-through gradients (Xu et al., 2018). Training is initialized at a pre-trained full-precision LSTM. We use the SGD optimizer with initial learning rate 20.0 and decay by a factor of 1.2 when the validation error does not improve over an epoch. We train for 80 epochs with batch size 20, BPTT 30, dropout with probability 0.5, and clip the gradient norms to 0.25. The regularization rate λ is tuned by finding the best performance on the validation set. In addition to multi-bit quantization, we also report the results for binary LSTMs (weights in {±1}), comparing BinaryConnect and our PROXQUANT-Binary, where both learning rates are tuned on an exponential grid {2.5, 5, 10, 20, 40}.
Result We report the perplexity-per-word (PPW, lower is better) in Table 2. The performance of PROXQUANT is comparable with the Straight-through gradient method. On Binary LSTMs, PROXQUANT-Binary beats BinaryConnect by a large margin. These results demonstrate that PROXQUANT offers a powerful alternative for training recurrent networks.

heading: 5 THEORETICAL ANALYSIS
text: In this section, we perform a theoretical study on the convergence of quantization algorithms. We show in Section 5.1 that our PROXQUANT algorithm (i.e. non-lazy prox-gradient method) converges under mild smoothness assumptions on the problem. In Section 5.2, we provide a simple example showing that the lazy prox-gradient method fails to converge under the same set of assumptions. In Section 5.3, we show that BinaryConnect has a very stringent condition for converging to a fixed point. Our theory demonstrates the superiority of our proposed PROXQUANT over lazy proxgradient type algorithms such as BinaryConnect and BinaryRelax (Yin et al., 2018). All missing proofs are deferred to Appendix D.
Prox-gradient algorithms (both lazy and non-lazy) with a fixed λ aim to solve the problem
minimize θ∈Rd L(θ) + λR(θ), (11)
and BinaryConnect can be seen as the limiting case of the above with λ =∞ (cf. Section 2.2).

heading: 5.1 A CONVERGENCE THEOREM FOR PROXQUANT
text: We consider PROXQUANT with batch gradient and constant regularization strength λt ≡ λ: θt+1 = proxηtλR(θt − ηt∇L(θt)). Theorem 5.1 (Convergence of ProxQuant). Assume that the loss L is β-smooth (i.e. has β-Lipschitz gradients) and the regularizer R is differentiable. Let Fλ(θ) = L(θ) + λR(θ) be the composite objective and assume that it is bounded below by F?. Running ProxQuant with batch gradient ∇L, constant stepsize ηt ≡ η = 12β and λt ≡ λ for T steps, we have the convergence guarantee
‖∇Fλ(θTbest)‖ 2 2 ≤ Cβ(Fλ(θ0)− F?) T where Tbest = arg min 1≤t≤T ‖θt − θt−1‖2 , (12)
where C > 0 is a universal constant. 2We thank Xu et al. (2018) for sharing the implementation of this method through a personal communication. There is a very clever trick not mentioned in their paper: after computing the alternating quantization qalt(θ), they multiply by a constant 0.3 before taking the gradient; in other words, their quantizer is a rescaled alternating quantizer: θ 7→ 0.3qalt(θ). This scaling step gives a significant gain in performance – without scaling the PPW is {116.7, 94.3, 87.3} for {1, 2, 3} bits. In contrast, our PROXQUANT does not involve a scaling step and achieves better PPW than this unscaled ALT straight-through method.
Remark 5.1. The convergence guarantee requires both the loss and the regularizer to be smooth. Smoothness of the loss can be satisfied if we use a smooth activation function (such as tanh). For the regularizer, the quantization-inducing regularizers defined in Section 3.1 (such as the W-shaped regularizer) are non-differentiable. However, we can use a smoothed version of them that is differentiable and point-wise arbitrarily close to R, which will satisfy the assumptions of Theorem 5.1. The proof of Theorem 5.1 is deferred to Appendix D.1.

heading: 5.2 NON-CONVERGENCE OF LAZY PROX-GRADIENT
text: The lazy prox-gradient algorithm (e.g. BinaryRelax (Yin et al., 2018)) for solving problem eq. (11) is a variant where the gradients are taken at proximal points but accumulated at the original sequence:
θt+1 = θt − ηt∇L(proxλR(θt)). (13) Convergence of the lazy prox-gradient algorithm eq. (13) is only known to hold for convex problems (Xiao, 2010); on smooth non-convex problems it generally does not converge even in an ergodic sense. We provide a concrete example that satisfies the assumptions in Theorem 5.1 (so that PROXQUANT converges ergodically) but lazy prox-gradient does not converge. Theorem 5.2 (Non-convergence of lazy prox-gradient). There exists L and R satisfying the assumptions of Theorem 5.1 such that for any constant stepsize ηt ≡ η ≤ 12β , there exists some specific initialization θ0 on which the lazy prox-gradient algorithm eq. (13) oscillates between two non-stataionry points and hence does not converge in the ergodic sense of eq. (12). Remark 5.2. Our construction is a fairly simple example in one-dimension and not very adversarial: L(θ) = 12θ 2 and R is a smoothed W-shaped regularizer. See Appendix D.2 for the details.

heading: 5.3 CONVERGENCE CHARACTERIZATION FOR BINARYCONNECT
text: For BinaryConnect, the concept of stataionry points is no longer sensible (as the target points {±1}d are isolated and hence every point is stationary). Here, we consider the alternative definition of convergence as converging to a fixed point and show that BinaryConnect has a very stringent convergence condition.
Consider the BinaryConnect method with batch gradients: st = sign(θt), θt+1 = θt − ηt∇L(st). (14)
Definition 5.1 (Fixed point and convergence). We say that s ∈ {±1}d is a fixed point of the BinaryConnect algorithm, if s0 = s in eq. (14) implies that st = s for all t = 1, 2, .... We say that the BinaryConnect algorithm converges if there exists t <∞ such that st is a fixed point. Theorem 5.3. Assume that the learning rates satisfy ∑∞ t=0 ηt = ∞, then s ∈ {±1}d is a fixed point for BinaryConnect eq. (14) if and only if sign(∇L(s)[i]) = −s[i] for all i ∈ [d] such that ∇L(θ)[i] 6= 0. Such a point may not exist, in which case BinaryConnect does not converge for any initialization θ0 ∈ Rd. Remark 5.3. Theorem 5.3 is in appearingly a stark contrast with the convergence result for BinaryConnect in (Li et al., 2017) in the convex case, whose bound involves a an additive error O(∆) that does not vanish over iterations, where ∆ is the grid size for quantization. Hence, their result is only useful when ∆ is small. In contrast, we consider the original BinaryConnect with ∆ = 1, in which case the error makes Li et al. (2017)’s bound vacuous. The proof of Theorem 5.3 is deferred to Appendix D.3.
Experimental evidence We have already seen that such a fixed point s might not exist in the toy example in Figure 1b. In Appendix C, we perform a sign change experiment on CIFAR-10, showing that BinaryConnect indeed fails to converge to a fixed sign pattern, corroborating Theorem 5.3.

heading: 6 CONCLUSION
text: 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.

heading: ACKNOWLEDGEMENT
text: We thank Tong He, Yifei Ma, Zachary Lipton, and John Duchi for their valuable feedback. We thank Chen Xu and Zhouchen Lin for the insightful discussion on multi-bit quantization and sharing the implementation of (Xu et al., 2018) with us. We thank Ju Sun for sharing the draft of (Sun & Sun, 2018) and the inspiring discussions on adversarial regularization for quantization. The majority of this work was performed when YB and YW were at Amazon AI.

heading: A ADDITIONAL RESULTS ON REGULARIZATION
text: A.1 PROX OPERATORS FOR BINARY NETS
Here we derive the prox operators for the binary regularizer eq. (6) and its squared L2 variant. Recall that
Rbin(θ) =
d∑
j=1
min {|θj − 1|, |θj + 1|}.
By definition of the prox operator, we have for any θ ∈ Rd that
proxλRbin(θ) = arg min θ̃∈Rd
   1 2 ∥∥∥θ̃ − θ ∥∥∥ 2 2 + λ d∑
j=1
min { |θ̃j − 1|, |θ̃j + 1| }   
= arg min θ̃∈Rd
   d∑
j=1
1 2 (θ̃j − θj)2 + λmin
{ |θ̃j − 1|, |θ̃j + 1| }   .
This minimization problem is coordinate-wise separable. For each θ̃j , the penalty term remains the same upon flipping the sign, but the quadratic term is smaller when sign(θ̃j) = sign(θj). Hence, the solution θ? to the prox satisfies that sign(θ?j ) = sign(θj), and the absolute value satisfies
|θ?j | = arg min t≥0
{ 1
2 (t− |θj |)2 + λ|t− 1|
} = SoftThreshold(|θj |, 1, λ) = 1+sign(|θj |−1)[||θj |−1|−λ]+.
Multiplying by sign(θ?j ) = sign(θj), we have
θ?j = SoftThreshold(θj , sign(θj), λ),
which gives eq. (7).
For the squared L2 version, by a similar argument, the corresponding regularizer is
Rbin(θ) =
d∑
j=1
min { (θj − 1)2, (θj + 1)2 } .
For this regularizer we have
proxλRbin(θ) = arg min θ̃∈Rd
   d∑
j=1
1 2 (θ̃j − θj)2 + λmin
{ (θ̃j − 1)2, (θ̃j + 1)2 }   .
Using the same argument as in the L1 case, the solution θ? satisfies sign(θ?j ) = sign(θj), and
|θ?j | = arg min t≥0
{ 1
2 (t− |θj |)2 + λ(t− 1)2 } = |θj |+ λ 1 + λ .
Multiplying by sign(θ?j ) = sign(θj) gives
θ?j = θj + λ sign(θj)
1 + λ ,
or, in vector form, θ? = (θ + λ sign(θ))/(1 + λ).
A.2 PROX OPERATOR FOR TERNARY QUANTIZATION
For ternary quantization, we use an approximate version of the alternating prox operator eq. (10): compute θ̃ = proxλR(θ) by initializing at θ̃ = θ and repeating
θ̂ = q(θ̃) and θ̃ = θ + 2λθ̂
1 + 2λ , (15)
where q is the ternary quantizer defined as
q(θ) = θ+1{θ ≥ ∆}+ θ−1{θ ≤ −∆}, ∆ = 0.7 d ‖θ‖1 , θ + = θ|i:θi≥∆, θ− = θ|i:θi≤−∆. (16)
This is a straightforward extension of the TWN quantizer (Li & Liu, 2016) that allows different levels for positives and negatives. We find that two rounds of alternating computation in eq. (15) achieves a good performance, which we use in our experiments.

heading: B ADDITIONAL EXPERIMENTAL RESULTS
text: B.1 TERNARY QUANTIZATION FOR CIFAR-10
Our models are ResNets of depth 20, 32, and 44. Ternarized training is initialized at pre-trained full-precision nets. We perform a hard quantization θ 7→ q(θ) at epoch 400 and keeps training till the600-th epoch to stabilize the BatchNorm layers.
Result The top-1 classification errors for ternary quantization are reported in Table 3. Our results are comparable with the reported results of TTQ,3 and the best performance of our method over 4 runs (from the same initialization) is slightly better than TTQ.

heading: C SIGN CHANGE EXPERIMENT
text: We experimentally compare the training dynamics of PROXQUANT-Binary and BinaryConnect through the sign change metric. The sign change metric between any θ1 and θ2 is the proportion of their different signs, i.e. the (rescaled) Hamming distance:
SignChange(θ1, θ2) = ‖sign(θ1)− sign(θ2)‖1
2d ∈ [0, 1].
In Rd, the space of all full-precision parameters, the sign change is a natural distance metric that represents the closeness of the binarization of two parameters.
Recall in our CIFAR-10 experiments (Section 4.1), for both BinaryConnect and PROXQUANT, we initialize at a good full-precision net θ0 and stop at a converged binary network θ̂ ∈ {±1}d. We are interested in SignChange(θ0, θt) along the training path, as well as SignChange(θ0, θ̂), i.e. the distance of the final output model to the initialization.
Our finding is that PROXQUANT produces binary nets with both lower sign changes and higher performances, compared with BinaryConnect. Put differently, around the warm start, there is a good binary net nearby which can be found by PROXQUANT but not BinaryConnect, suggesting that BinaryConnect, and in general the straight-through gradient method, suffers from higher optimization instability than PROXQUANT. This finding is consistent in all layers, across different warm starts, and across differnent runs from each same warm start (see Figure 3 and Table 4 in Appendix C.1). This result here is also consistent with Theorem 5.3: the signs in BinaryConnect never stop changing until we manually freeze the signs at epoch 400.
3We note that our PROXQUANT-Ternary and TTQ are not strictly comparable: we have the advantage of using better initializations; TTQ has the advantage of a stronger quantizer: they train the quantization levels (θ+, θ−) whereas our quantizer eq. (16) pre-computes them from the current full-precision parameter.
C.1 RAW DATA FOR SIGN CHANGE EXPERIMENT

heading: D PROOFS OF THEORETICAL RESULTS
text: D.1 PROOF OF THEOREM 5.1
Recall that a function f : Rd → R is said to be β-smooth if it is differentiable and∇f is β-Lipschitz: for all x, y ∈ Rd we have ‖∇f(x)−∇f(y)‖2 ≤ β ‖x− y‖2 . For any β-smooth function, it satisfies the bound
f(y) ≤ f(x) + 〈∇f(x), y − x〉+ β 2 ‖x− y‖22 for all x, y ∈ R d.
Convergence results like Theorem 5.1 are standard in the literature of proximal algorithms, where we have convergence to stataionarity without convexity on either L or R but assuming smoothness. For completeness we provide a proof below. Note that though the convergence is ergodic, the best index Tbest can be obtained in practice via monitoring the proximity ‖θt − θt−1‖2. Proof of Theorem 5.1 Recall the ProxQuant iterate
θt+1 = arg min θ∈Rd
{ L(θt) + 〈θ − θt,∇L(θt)〉+ 1
2η ‖θ − θt‖22 + λR(θ)
} .
By the fact that θt+1 minimizes the above objective and applying the smoothness of L, we get that
Fλ(θt) = L(θt) + λR(θt) ≥ L(θt) + 〈θt+1 − θt,∇L(θt)〉+ 1
2η ‖θt+1 − θt‖22 + λR(θt+1)
≥ L(θt+1) + ( 1
2η − β 2
) ‖θt+1 − θt‖22 + λR(θt+1) = Fλ(θt+1) + β
2 ‖θt+1 − θt‖22 .
Telescoping the above bound for t = 0, . . . , T − 1, we get that T−1∑
t=0
‖θt+1 − θt‖22 ≤ 2(Fλ(θ0)− Fλ(θT )) β ≤ 2(Fλ(θ0)− F?) β .
Therefore we have the proximity guarantee
min 0≤t≤T−1
‖θt+1 − θt‖22 ≤ 2(Fλ(θ0)− F?)
βT . (17)
We now turn this into a stationarity guarantee. The first-order optimality condition for θt+1 gives
∇L(θt) + 1
η (θt+1 − θt) + λ∇R(θt+1) = 0.
Combining the above equality and the smoothness of L, we get
‖∇Fλ(θt+1)‖2 = ‖∇L(θt+1) + λ∇R(θt+1)‖2 = ∥∥∥∥ 1
η (θt − θt+1) +∇L(θt+1)−∇L(θt) ∥∥∥∥ 2
≤ ( 1
η + β
) ‖θt+1 − θt‖2 = 3β ‖θt+1 − θt‖2 .
Choosing t = Tbest − 1 and applying the proximity guarantee eq. (17), we get
‖∇Fλ(θTbest)‖ 2 2 ≤ 9β 2 ‖θTbest − θTbest−1‖ 2 2 = 9β 2 min 0≤t≤T−1
‖θt+1 − θt‖22 ≤ 18β(Fλ(θ0)− F?)
T .
This is the desired bound.
D.2 PROOF OF THEOREM 5.2
Let our loss function L : R → R be the quadratic L(θ) = 12θ 2 (so that L is β-smooth with β = 1). Let the regularizer R : R → R be a smoothed version of the W-shaped regularizer in eq. (3), defined as (for ∈ (0, 1/2] being the smoothing radius)
R(θ) =    − 1 2 θ2 + 1− , θ ∈ [0, ) − θ + 1− 2 , θ ∈ [ , 1− ) 1 2 (θ − 1)2, θ ∈ [1− , 1 + )
θ − 1− 2 , θ ∈ [1 + ,∞)
and R(−θ) = R(θ) for the negative part. See Figure 4 for an illustration of the loss L and the regularizer R (with = 0.2).
It is straightforward to see that R is piecewise quadratic and differentiable on R by computing the derivatives at and 1 ± . Further, by elementary calculus, we can evaluate the prox operator in closed form: for all λ ≥ 1, we have
proxλR(θ) = θ + λ sign(θ)
+ λ sign(θ) for all |θ| ≤ 1.
Now, suppose we run the lazy prox-gradient method with constant stepsize ηt ≡ η ≤ 12β = 1 2 . For the specific initialization
θ0 = ηλ
2λ+ (2− η) ∈ (0, 1),
we have the equality proxλR(θ0) = 2 η θ0 and therefore the next lazy prox-gradient iterate is
θ1 = θ0 − η∇L(proxλR(θ0)) = θ0 − η∇L ( 2
η θ0
) = θ0 − η · 2
η θ0 = −θ0.
As both R and L are even functions, a symmetric argument holds for θ1 from which we get θ2 = −θ1 = θ0. Therefore the lazy prox-gradient method ends up oscillating between two points:
θt = (−1)tθ0. On the other hand, it is straightforward to check that the only stationary points of L(θ) + λR(θ) are 0 and ± λ +λ , all not equal to ±θ0. Therefore the sequence {θt}t≥0 does not have a subsequence with vanishing gradient and thus does not approach stationarity in the ergodic sense.
D.3 PROOF OF THEOREM 5.3
We start with the “⇒” direction. If s is a fixed point, then by definition there exists θ0 ∈ Rd such that θt = θ for all t = 0, 1, 2, .... By the iterates eq. (14)
θT = θ0 − T∑
t=0
ηt∇L(st).
Take signs on both sides and apply st = s for all t on both sides, we get that
s = sT = sign(θT ) = sign ( θ0 −∇L(s) T∑
t=0
ηt
)
Take the limit T → ∞ and apply the assumption that ∑ t ηt = ∞, we get that for all i ∈ [d] such that [∇L(θ)]i 6= 0,
s[i] = lim T→∞ sign
( θ0 −∇L(s) T∑
t=0
ηt
) [i] = − sign(∇L(s))[i].
Now we prove the “⇐” direction. If θ obeys that sign(∇L(s)[i]) = −s[i] for all i ∈ [d] such that ∇L(s)[i] 6= 0, then if we take any θ0 such that sign(θ0) = s, θt will move in a straight line towards the direction of −∇L(s), which does not change the sign of h0. In other words, st = sign(θt) = sign(θ0) = s for all t = 0, 1, 2, .... Therefore, by definition, s is a fixed point.

Review summaries:
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.

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.