Year: 2019

Abstract:
Quantizing large Neural Networks (NN) while maintaining the performance is highly desirable for resource-limited devices due to reduced memory and time complexity. NN quantization is usually formulated as a constrained optimization problem and optimized via a modified version of gradient descent. In this work, by interpreting the continuous parameters (unconstrained) as the dual of the quantized ones, we introduce a Mirror Descent (MD) framework (Bubeck (2015)) for NN quantization. Specifically, we provide conditions on the projections (i.e., mapping from continuous to quantized ones) which would enable us to derive valid mirror maps and in turn the respective MD updates. Furthermore, we discuss a numerically stable implementation of MD by storing an additional set of auxiliary dual variables (unconstrained). This update is strikingly analogous to the popular Straight Through Estimator (STE) based method which is typically viewed as a “trick” to avoid vanishing gradients issue but here we show that it is an implementation method for MD for certain projections. Our experiments on standard classification datasets (CIFAR-10/100, TinyImageNet) with convolutional and residual architectures show that our MD variants obtain fully-quantized networks with accuracies very close to the floating-point networks.

heading: 1 INTRODUCTION
text: Despite the success of deep neural networks in various domains, their excessive computational and memory requirements limit their practical usability for real-time applications or in resource-limited devices. Quantization is a prominent technique for network compression, where the objective is to learn a network while restricting the parameters to take values from a small discrete set (usually binary). This leads to a dramatic reduction in memory (a factor of 32 for binary quantization) and inference time – as it enables specialized implementation using bit operations.
Neural Network (NN) quantization is usually formulated as a constrained optimization problem minx∈X f(x), where f(·) denotes the loss function by abstracting out the dependency on the dataset and X ⊂ IRr denotes the set of all possible quantized solutions. Majority of the works in the literature (Hubara et al. (2017); Yin et al. (2018); Ajanthan et al. (2019)) convert this into an unconstrained problem by introducing auxiliary variables (x̃) and optimize via (stochastic) gradient descent. Specifically, the objective and the update step take the following form:
min x̃∈IRr
f(P (x̃)) , x̃k+1 = x̃k − η ∇x̃f(P (x̃))|x̃=x̃k , (1)
where P : IRr → X is a mapping from the unconstrained space to the quantized space (sometimes called projection) and η > 0 is the learning rate. In cases where the mapping P is not differentiable, a suitable approximation is employed (Hubara et al. (2017)).
In this work, by noting that the well-known Mirror Descent (MD) algorithm, widely used for online convex optimization (Bubeck (2015)), provides a theoretical framework to perform gradient descent in the unconstrained space (dual space, IRr) with gradients computed in the quantized space (primal space, X ), we introduce an MD framework for NN quantization. In essence, MD extends gradient descent to non-Euclidean spaces where Euclidean projection is replaced with a more general projection defined based on the associated distance metric. Briefly, the key ingredient of MD is a concept called mirror map which defines both the mapping between primal and dual spaces and the exact form of
the projection. Specifically, in this work, by observing P in Eq. (1) as a mapping from dual space to the primal space, we analytically derive corresponding mirror maps under certain conditions on P . This enables us to derive different variants of the MD algorithm useful for NN quantization.
Furthermore, as MD is often found to be numerically unstable (Hsieh et al. (2018)), we discuss a numerically stable implementation of MD by storing an additional set of auxiliary variables similar to the existing methods. As will be shown later, this update is strikingly analogous to the popular Straight Through Estimator (STE) based gradient method (Hubara et al. (2017); Bai et al. (2019)) which is typically viewed as a “trick” to avoid vanishing gradients issue but here we show that it is an implementation method for MD under certain conditions on the mapping P . We believe this connection sheds some light on the practical effectiveness of STE.
We evaluate the merits of our MD variants on CIFAR-10/100 and TinyImageNet classification datasets with convolutional and residual architectures. Our experiments show that the quantized networks obtained by the MD variants yield accuracies very close to the floating-point counterparts while outperforming directly comparable baselines. Finally, we would like to emphasize that even though our formulation does not necessarily extend the theory of MD, we believe showing MD as a suitable framework for NN quantization with superior empirical performance opens up new ways of designing MD-inspired update rules for NNs.

heading: 2 PRELIMINARIES
text: We first provide some background on the MD algorithm and NN quantization. Then we discuss the link between them and provide our MD framework for NN quantization.

heading: 2.1 MIRROR DESCENT
text: The Mirror Descent (MD) algorithm is first introduced in (Nemirovsky & Yudin (1983)) and it has been extensively studied in the convex optimization literature ever since. In this section we provide a brief overview and we refer the interested reader to Chapter 4 of (Bubeck (2015)). In the context of MD, we consider a problem of the form:
min x∈X f(x) , (2)
where f : X → IR is a convex function and X ⊂ IRr is a compact convex set. The main concept of MD is to extend gradient descent to a more general non-Euclidean space (Banach space1), thus overcoming the dependency of gradient descent on the Euclidean geometry. The motivation for this generalization is that one might be able to exploit the geometry of the space to optimize much more efficiently. One such example is the simplex constrained optimization where MD converges at a much faster rate than the standard Projected Gradient Descent (PGD).
To this end, since the gradients lie in the dual space, optimization is performed by first mapping the primal point xk ∈ B (quantized space, X ) to the dual space B∗ (unconstrained space, IRr), then performing gradient descent in the dual space, and finally mapping back the resulting point to the primal space B. If the new point xk+1 lie outside of the constraint set X ⊂ B, it is projected to the set X . Both the primal/dual mapping and the projection are determined by the mirror map. Specifically, the gradient of the mirror map defines the mapping from primal to dual and the projection is done via the Bregman divergence of the mirror map. We first provide the definitions for mirror map and Bregman divergence and then turn to the MD updates. Definition 2.1 (Mirror map). Let C ⊂ IRr be a convex open set such that X ⊂ C̄ (C̄ denotes the closure of set C) and X ∩ C 6= ∅. Then, Φ : C → IR is a mirror map if it satisfies:
1. Φ is strictly convex and differentiable. 2. ∇Φ(C) = IRr, i.e., ∇Φ takes all possible values in IRr. 3. limx→∂C ‖∇Φ(x)‖ =∞ (∂C denotes the boundary of C), i.e.,∇Φ diverges on the boundary of C. Definition 2.2 (Bregman divergence). Let Φ : C → IR be a continuously differentiable, strictly convex function defined on a convex set C. The Bregman divergence associated with Φ for points p,q ∈ C is the difference between the value of Φ at point p and the value of the first-order Taylor expansion of Φ around point q evaluated at point p, i.e.,
DΦ(p,q) = Φ(p)− Φ(q)− 〈∇Φ(q),p− q〉 . (3) 1A Banach space is a complete normed vector space where the norm is not necessarily derived from an inner
product.
Notice, DΦ(p,q) ≥ 0 with DΦ(p,p) = 0, and DΦ(p,q) is convex on p.
Now we are ready to provide the mirror descent strategy based on the mirror map Φ. Let x0 ∈ argminx∈X∩C Φ(x) be the initial point. Then, for iteration k ≥ 0 and step size η > 0, the update of the MD algorithm can be written as:
∇Φ(yk+1) = ∇Φ(xk)− η gk , where gk ∈ ∂f(xk) and yk+1 ∈ C , (4) xk+1 = argmin
x∈X∩C DΦ(x,y
k+1) .
Note that, in Eq. (4), the gradient gk is computed at xk ∈ X ∩ C (solution space) but the gradient descent is performed in IRr (unconstrained dual space). Moreover, by simple algebraic manipulation, it is easy to show that the above MD update (4) can be compactly written in a proximal form where the Bregman divergence of the mirror map becomes the proximal term (Beck & Teboulle (2003)):
xk+1 = argmin x∈X∩C 〈η gk,x〉+DΦ(x,xk) . (5)
Note, if Φ(x) = 12 ‖x‖ 2 2, then DΦ(x,x k) = 12 ∥∥x− xk∥∥2 2 , which when plugged back to the above problem and optimized for x, leads to exactly the same update rule as that of PGD. However, MD allows us to choose various forms of Φ depending on the problem at hand.

heading: 2.2 NEURAL NETWORK QUANTIZATION
text: Neural Network (NN) quantization amounts to training networks with parameters restricted to a small discrete set representing the quantization levels. Here we review two constrained optimization formulations for NN quantization: 1) directly constrain each parameter to be in the discrete set; and 2) optimize the probability of each parameter taking a label from the set of quantization levels.

heading: 2.2.1 PARAMETER SPACE FORMULATION
text: Given a dataset D = {xi,yi}ni=1, NN quantization can be written as:
min w∈Qm L(w;D) := 1 n n∑ i=1 `(w; (xi,yi)) . (6)
Here, `(·) denotes the input-output mapping composed with a standard loss function (e.g., crossentropy loss), w is the m dimensional parameter vector, and Q with |Q| = d is a predefined discrete set representing quantization levels (e.g., Q = {−1, 1} or Q = {−1, 0, 1}). The approaches that directly optimize in the parameter space include BinaryConnect (BC) (Courbariaux et al. (2015)) and its variants (Hubara et al. (2017); Rastegari et al. (2016)), where the constraint set is discrete. In contrast, recent approaches (Bai et al. (2019); Yin et al. (2018)) relax this constraint set to be its convex hull: conv(Qm) = [qmin, qmax]m , (7) where qmin and qmax represent the minimum and maximum quantization levels, respectively. In this case, a quantized solution is obtained by gradually increasing an annealing hyperparameter.

heading: 2.2.2 LIFTED PROBABILITY SPACE FORMULATION
text: Another formulation is based on the Markov Random Field (MRF) perspective to NN quantization recently studied in (Ajanthan et al. (2019)). It treats Eq. (6) as a discrete labelling problem and introduces indicator variables uj:λ ∈ {0, 1} for each parameter wj where j ∈ {1, . . . ,m} such that uj:λ = 1 if and only if wj = λ ∈ Q. For convenience, by denoting the vector of quantization levels as q, a parameter vector w ∈ Qm can be written in a matrix vector product as:
w = uq , where u ∈ Vm = { u ∑ λ uj:λ = 1, ∀ j
uj:λ ∈ {0, 1}, ∀ j, λ
} . (8)
Here, u is a m × d matrix (i.e., each row uj = {uj:λ | λ ∈ Q}), and q is a column vector of dimension d. Note that, u ∈ Vm is an overparametrized (i.e., lifted) representation of w ∈ Qm. Now, similar to the relaxation in the parameter space, one can relax the binary constraint in Vm to form its convex hull:
∆m = conv(Vm) = { u ∑ λ uj:λ = 1, ∀ j
uj:λ ≥ 0, ∀ j, λ
} . (9)
The set ∆m is in fact the Cartesian product of the standard (d− 1)-probability simplexes embedded in IRd. Therefore, for a feasible point u ∈ ∆m, the vector uj for each j (j-th row of matrix u) belongs to the probability simplex ∆. Hence, we can interpret the value uj:λ as the probability of assigning the discrete label λ to the weight wj . This relaxed optimization can then be written as:
min u∈∆m L(uq;D) := 1 n n∑ i=1 `(uq; (xi,yi)) . (10)
In fact, this can be interpreted as finding a probability distribution u ∈ ∆m such that the cost L(u) is minimized. Note that, the relaxation of u from Vm to ∆m translates into relaxing w from Qm to the convex region conv(Qm). Even in this case, a discrete solution u ∈ Vm can be enforced via an annealing hyperparameter or using rounding schemes.

heading: 3 MIRROR DESCENT FRAMEWORK FOR NETWORK QUANTIZATION
text: Before introducing the MD formulation, we first write NN quantization as a single objective unifying (6) and (10) as:
min x∈X f(x) , (11)
where f(·) denotes the loss function by abstracting out the dependency on the dataset D, and X denotes the constraint set (Qm or Vm depending on the formulation). Note that, as discussed in Sec. 2.2, many recent NN quantization methods optimize over the convex hull of the constraint set. Following this, we consider the solution space X in Eq. (11) to be convex and compact. To employ MD, we need to choose a mirror map (refer Definition 2.1) suitable for the problem at hand. In fact, as discussed in Sec. 2.1, mirror map is the core component of an MD algorithm which determines the effectiveness of the resulting MD updates. However, there is no straightforward approach to obtain a mirror map for a given constrained optimization problem, except in certain special cases.
To this end, we observe that the usual approach to optimize the above constrained problem is via a version of projected gradient descent, where the projection is the mapping from the unconstrained auxiliary variables (high-precision) to the quantized space X . Now, noting the analogy between the purpose of the projection operator and the mirror maps in the MD formulation, we intend to derive the mirror map analogous to a given projection. Precisely, we prove that if the projection is invertible and strictly monotone, a valid mirror map can be derived from the projection itself. This does not necessarily extend the theory of MD as finding a strictly monotone map is as hard as finding the mirror map itself. However, this derivation is interesting as it connects existing PGD type algorithms to their corresponding MD variants. For completeness, we now state it as a theorem for the case X ⊂ IR and the multidimensional case can be easily proved with an additional assumption that the vector field P−1(x) is conservative. Theorem 3.1. Let X be a compact convex set and P : IR → C be an invertible function where C ⊂ IR is a convex open set such that X = C̄ (C̄ denotes the closure of C). Now, if
1. P is strictly monotonically increasing. 2. limx→∂C ‖P−1(x)‖ =∞ (∂C denotes the boundary of C).
Then, Φ(x) = ∫ x x0 P−1(y)dy is a valid mirror map. Proof. This can be proved by noting that ∇Φ(x) = P−1(x). Please refer to Appendix A.
The MD update based on the mirror map derived from a given projection is illustrated in Fig. 1. Note that, to employ MD to the problem (11), in theory, any mirror map satisfying Definition 2.1 whose domain (the closure of the domain) is a superset of the constraint set X can be chosen. However, the above theorem provides a method to derive only a subset of all applicable mirror maps, where the closure of the domain of mirror maps is exactly equal to the constraint set X .
We now give some example projections useful for NN quantization (tanh for w-space and softmax for u-space) and derive their corresponding mirror maps. Given mirror maps, the MD updates are straightforward based on Eq. (5). Even though we consider differentiable projections, Theorem 3.1
does not require the projection to be differentiable. For the rest of the section, we assume m = 1, i.e., consider projections that are independent for each j ∈ {1, . . . ,m}. Example 3.1 (w-space, binary, tanh). Consider the tanh function, which projects a real value to the interval [−1, 1]:
w = P (w̃) := tanh(βw̃) = exp(2βw̃)− 1 exp(2βw̃) + 1 , (12)
where β > 0 is the annealing hyperparameter and when β →∞, tanh approaches the step function. The inverse of the tanh is:
P−1(w) = 1
β tanh−1(w) =
1
2β log
1 + w 1− w . (13)
Note that, P−1 is monotonically increasing for a fixed β. Correspondingly, the mirror map from Theorem 3.1 can be written as:
Φ(w) = ∫ P−1(w)dw = 1
2β
[ (1 + w) log(1 + w) + (1− w) log(1− w) ] . (14)
Here, the constant from the integration is ignored. It can be easily verified that Φ(w) is in fact a valid mirror map. The projection, its inverse and the corresponding mirror map are illustrated in Fig. 2a. Consequently, the resulting MD update (5) takes the following form:
wk+1 = argmin w∈(−1,1)
〈η gk, w〉+DΦ(w,wk) = 1+wk 1−wk exp(−2βηg k)− 1
1+wk 1−wk exp(−2βηgk) + 1 . (15)
The update formula is derived using the KKT conditions (Boyd & Vandenberghe (2009)). For the detailed derivation please refer to Appendix B. A similar derivation can also be performed for the sigmoid function, where C̄ = X = [0, 1]. Note that the sign function has been used for binary quantization in (Courbariaux et al. (2015)) and tanh can be used as a soft version of sign function as pointed out by (Zhang et al. (2015)). Mirror map corresponding to tanh is used for online linear optimization in (Bubeck et al. (2012)) but here we use it for NN quantization.
Example 3.2 (u-space, multi-label, softmax). Now we consider the softmax projection used in Proximal Mean-Field (PMF) (Ajanthan et al. (2019)) to optimize in the lifted probability space. In this case, the projection is defined as P (ũ) := softmax(βũ) where P : IRd → C with C̄ = X = ∆. Here ∆ is the (d− 1)-dimensional probability simplex and |Q| = d. Note that the softmax projection is not invertible as it is a many-to-one mapping. In particular, it is invariant to translation, i.e.,
u = softmax(ũ + c1) = softmax(ũ) , where uλ = exp(ũλ)∑ µ∈Q exp(ũµ) , (16)
for any scalar c ∈ IR (1 denotes a vector of all ones). Therefore, the softmax projection does not satisfy Theorem 3.1. However, one could define the inverse of softmax as follows: given u ∈ ∆, find a unique point ṽ = ũ + c1, for a particular scalar c, such that u = softmax(ṽ). Now, by choosing c = − log( ∑ µ=Q exp(ũµ)), softmax can be written as:
u = softmax(ṽ) , where uλ = exp(ṽλ) , ∀λ ∈ Q . (17)
Now, the inverse of the projection can be written as:
ṽ = P−1(u) = 1
β softmax−1(u) , where ṽλ =
1 β log(uλ) , ∀λ . (18)
Indeed, log is a monotonically increasing function and from Theorem 3.1, by summing the integrals, the mirror map can be written as:
Φ(u) = 1
β [∑ λ uλ log(uλ)− uλ ] = − 1 β H(u)− 1/β . (19)
Here, ∑ λ uλ = 1 as u ∈ ∆, and H(u) is the entropy. Interestingly, as the mirror map in this case is the negative entropy (up to a constant), the MD update leads to the well-known Exponentiated
Gradient Descent (EGD) (or Entropic Descent Algorithm (EDA)) (Beck & Teboulle (2003); Bubeck (2015)). Consequently, the update takes the following form:
uk+1λ = ukλ exp(−βηgkλ)∑
µ∈Q u k µ exp(−βηgkµ)
∀λ . (20)
The derivation follows the same approach as in the tanh case above. It is interesting to note that the MD variant of softmax is equivalent to the well-known EGD. Notice, the authors of PMF hinted that PMF is related to EGD but here we have clearly showed that the MD variant of PMF under the above reparametrization (17) is exactly EGD.
Example 3.3 (w-space, multi-label, shifted tanh). Note that, similar to softmax, we wish to extend the tanh projection beyond binary. The idea is to use a function that is an addition of multiple shifted tanh functions. To this end, as an example we consider ternary quantization, with Q = {−1, 0, 1} and define our shifted tanh projection P : IR→ C as:
w = P (w̃) = 1
2
[ tanh (β(w̃ + 0.5)) + tanh (β(w̃ − 0.5)) ] , (21)
where β > 0 and w = C̄ = X = [−1, 1]. When β → ∞, P approaches a stepwise function with inflection points at −0.5 and 0.5 (here, ±0.5 is chosen heuristically), meaning w move towards one of the quantization levels in the setQ. This behaviour together with its inverse is illustrated in Fig. 2b. Now, one could potentially find the functional form of P−1 and analytically derive the mirror map corresponding to this projection. Note that, while Theorem 3.1 provides an analytical method to derive mirror maps, in some cases such as the above, the exact form of mirror map and the MD update might be nontrivial. In such cases, as will be shown subsequently, the MD update can be easily implemented by storing an additional set of auxiliary variables w̃.
Effect of Annealing. Note that, to ensure a discrete solution, projection P is parametrized by a scalar β and it is annealed throughout the optimization. This annealing hyperparameter translates into a time varying mirror map (refer Eqs. (14) and (19)) in our case. Such an adaptive mirror map gradually constrains the solution spaceX to its boundary and in the limit enforces a quantized solution. Since this adaptive behaviour can affect the convergence of the algorithm, in our implementation β is capped at an arbitrarily chosen maximum value, and empirically, the algorithm converges to fully quantized solutions in all tested cases. We leave the theoretical analysis of annealing for future work.
Numerically Stable form of Mirror Descent. We showed a few examples of valid projections, their corresponding mirror maps, and the final MD updates. Even though, in theory, these updates can be used directly, they are sometimes numerically unstable due to the operations involving multiple logarithms, exponentials and divisions (Hsieh et al. (2018)). To this end, we provide a numerically stable way of performing MD by storing an additional set of auxiliary parameters during training.
A careful look at the Fig. 1 suggests that the MD update with the mirror map derived from Theorem 3.1 can be performed by storing auxiliary variables x̃ = P−1(x). In fact, once the auxiliary variable x̃k is updated using gradient gk, it is directly mapped back to the constraint set X via the projection. This is mainly because of the fact that the domain of the mirror maps derived based on the Theorem 3.1 are exactly the same as the constraint set. Formally, with this additional set of variables, one can write the MD update (4) corresponding to the projection P as:
x̃k+1 = x̃k − η gk , update in the dual space (22) xk+1 = P (x̃k+1) ∈ X , projection to the primal space
where η > 0, gk ∈ ∂f(xk) and x̃k = P−1(xk). Experimentally we observed these updates to show stable behaviour and performed remarkably well for both the tanh and softmax.
Note, above updates can be seen as optimizing the function f(P (x̃)) using gradient descent where the gradient through the projection (i.e., Jacobian) JP = ∂P (x̃)/∂x̃ is replaced with the identity matrix. This is exactly the same as the Straight Through Estimator (STE) for NN quantization (following the nomenclature of (Bai et al. (2019); Yin et al. (2018))). Despite being a crude approximation, STE has shown to be highly effective for NN quantization with various network architectures and datasets (Yin et al. (2018); Zhou et al. (2016)). However, a solid understanding of the effectiveness of STE is lacking in the literature except for its convergence analysis in certain special cases (Li et al. (2017); Yin et al. (2019)). In this work, by showing STE based gradient descent as an implementation method of MD under certain conditions on the projection, we provide a justification on the effectiveness of STE. Besides, as shown in Example 3.3, in cases where the MD formulation is nontrivial, the STE based implementation can be used. The pseudocodes of original and numerically stable versions of our MD algorithm for tanh are presented in Appendix B.
Comparison against ProxQuant. The connection between the dual averaging version of MD and STE was recently hinted in ProxQuant (PQ) (Bai et al. (2019)). However, no analysis of whether an analogous mirror map exists to the given projection is provided and their final algorithm is not based on MD. In particular, following our notation, the final update equation of PQ can be written as:
x̃k+1 = xk − η gk , assumes xk and gk are in the same space (23) xk+1 = prox(x̃k+1) , prox : IRr → IRr is the proximal mapping defined in (Bai et al. (2019))
where η > 0, and gk ∈ ∂f(xk). Note that, as opposed to MD (refer to Eq. (22)), PQ assumes the point xk and gradient gk are in the same space. Then only the formula xk − η gk is valid. This would only be true for the Euclidean space. However, as discussed in Sec. 2.1, MD allows gradient descent to be performed on a more general non-Euclidean space by first mapping the primal point xk to a point x̃k in the dual space via the mirror map. Such an ability enabled theoretical and practical research on MD for the past three decades.
Convergence of MD in the Nonconvex Setting. We would like to point out that MD is originally developed for convex optimization, however, in this paper we directly apply MD to NN quantization where the loss is highly nonconvex and gradient estimates are stochastic, and empirically evaluate its convergence behaviour and performance. Theoretical analysis of MD for nonconvex, stochastic setting is an active research area (Zhou et al. (2017a;b)) and MD has been recently shown to converge in the nonconvex stochastic setting under certain conditions (Zhang & He (2018)). We believe convergence analysis of MD for NNs could constitute to a completely new theoretical paper.

heading: 4 RELATED WORK
text: In this work we consider parameter quantization, which is usually formulated as a constrained problem and optimized using a modified projected gradient descent algorithm, where the methods (Courbariaux et al. (2015); Carreira-Perpinán & Idelbayev (2017); Yin et al. (2018); Bai et al. (2019); Ajanthan et al. (2019)) mainly differ in the constraint set, the projection used, and how backpropagation through the projection is performed. Among them, STE based gradient descent is the most popular method as it enables backpropagation through nondifferentiable projections and it has shown to be highly effective in practice (Courbariaux et al. (2015)). In fact, the success of this approach lead to various extensions by including additional layerwise scalars (Rastegari et al. (2016)), relaxing the solution space (Yin et al. (2018)), and even to quantizing activations (Hubara et al. (2017)), and/or gradients (Zhou et al. (2016)). Moreover, there are methods focusing on loss aware quantization (Hou et al. (2017)), quantization for specialized hardware (Esser et al. (2015)), and quantization based on the variational approach (Achterhold et al. (2018); Louizos et al. (2017; 2019)). We have only provided a brief summary of relevant methods and for a comprehensive survey we refer the reader to (Guo (2018)).

heading: 5 EXPERIMENTS
text: Due to the popularity of binary neural networks (Courbariaux et al. (2015); Rastegari et al. (2016)), we mainly consider binary quantization and set the quantization levels as Q = {−1, 1}. We would
like to point out that we quantize all learnable parameters, meaning our approach results in 32 times less memory compared to the floating-point counterparts.
We evaluate two MD variants corresponding to tanh and softmax projections, on CIFAR-10, CIFAR100 and TinyImageNet2 datasets with VGG-16 and ResNet-18 architectures. We also evaluate the numerically stable versions of our MD variants (denoted with suffix “-S”) performed by storing auxiliary parameters during training as explained in Eq. (22). The results are compared against parameter quantization methods, namely BinaryConnect (BC) (Courbariaux et al. (2015)), ProxQuant (PQ) (Bai et al. (2019)) and Proximal Mean-Field (PMF) (Ajanthan et al. (2019)). In addition, for completeness, we also compare against a standard PGD variant corresponding to the tanh projection (denoted as GD-tanh), i.e., minimizing f(tanh(x̃)) using gradient descent. The only difference of this to our MD-tanh-S is that, in Eq. (22), the Jacobian of tanh is directly used in the updates. Note that, numerous techniques have emerged with BC as the workhorse algorithm by relaxing constraints such as the layer-wise scalars (Rastegari et al. (2016)), and similar extensions are straightforward even in our case. Briefly, our results indicate that the binary networks obtained by the MD variants yield accuracies very close to the floating-point counterparts while outperforming the baselines.
For all the experiments, standard multi-class cross-entropy loss is used. We crossvalidate the hyperparameters such as learning rate, learning rate scale, rate of increase of annealing hyperparameter β, and their respective schedules for all tested methods including the baselines. This extensive crossvalidation improved the accuracies of previous methods by a large margin, e.g., up to 3% improvement for PMF. We provide the hyperparameter tuning search space and the final hyperparameters in Appendix C. Our algorithm is implemented in PyTorch (Paszke et al. (2017)) and the experiments are performed on NVIDIA Tesla-P100 GPUs. Our code will be released upon publication.

heading: 5.1 RESULTS
text: The classification accuracies of binary networks obtained by both variants of our algorithm, namely, MD-tanh and MD-softmax, their numerically stable versions (suffix “-S”) and the baselines BC, PQ, PMF, GD-tanh and the floating point Reference Network (REF) are reported in Table 1. Both the numerically stable MD variants consistently produce better or on par results compared to other binarization methods while narrowing the performance gap between binary networks and floating point counterparts to a large extent, on multiple datasets.
Our stable MD-variant perform slightly better than MD-softmax, whereas for tanh, MD updates either perform on par or sometimes even better than numerically stable version of MD-tanh. We
2https://tiny-imagenet.herokuapp.com/
believe, the main reason for this empirical variation in results for our MD-variants is due to numerical instability caused by the floating-point arithmetic of logarithm and exponential functions in Eq. (15) and Eq. (20). Furthermore, even though our two MD-variants, namely MD-softmax and MD-tanh optimize in different spaces, their performance is similar in most cases. This may be explained by the fact that both algorithms belong to the same family where a “soft” projection to the constraint set (in fact the constraints sets are equivalent in this case, refer Sec. 2.2.2) is used and an annealing hyperparameter is used to gradually enforce a quantized solution.
Note, PQ (Bai et al. (2019)) does not quantizee the fully connected layers, biases and shortcut layers. For fair comparison, we crossvalidate PQ with all layers binarized and original PQ settings, and report the results denoted as PQ and PQ* respectively in Table 1. Our MD-variants outperform PQ consistently on multiple datasets in equivalent experimental settings. This clearly shows that entropic or tanh-based regularization with our annealing scheme is superior to a simple “W” shaped regularizer and emphasizes that MD is a suitable framework for quantization.
Furthermore, the superior performance of MD-tanh against GD-tanh and on par or better performance of MD-softmax against PMF for binary quantization empirically validates that MD is useful even in a nonconvex stochastic setting. This hypothesis along with our numerically stable form of MD can be particularly useful to explore other projections which are useful for quantization and/or network compression in general.
The training curves for our MD variants for CIFAR-10 and CIFAR-100 datasets with ResNet-18 are shown in Fig. 3. The original MD variants show unstable behaviour during training which is attributed to the fact that it involves logarithms and exponentials in the update rules. In addition, we believe, the additional annealing hyperparameter also contributes to this instability. Regardless, by storing auxiliary variables, the MD updates are demonstrated to be quite stable. This clear distinction between MD variants emphasizes the significance of practical considerations while implementing MD especially in NN optimization. For more experiments such as training curves comparison to other methods and ternary quantization results please refer to the Appendix C.

heading: 6 DISCUSSION
text: In this work, we have introduced an MD framework for NN quantization by deriving mirror maps corresponding to various projections useful for quantization. In addition, we have discussed a numerically stable implementation of MD by storing an additional set of auxiliary variables and showed that this update is strikingly analogous to the popular STE based gradient method. The superior performance of our MD formulation even with simple projections such as tanh and softmax is encouraging and we believe, MD would be a suitable framework for not just NN quantization but for network compression in general. Finally, some theoretical aspects such as the use of time-varying mirror maps and the combination of MD and a stochastic optimizer such as Adam are left unattended in this paper, which we intend to analyze in a future work.

heading: A DERIVING MIRROR MAPS FROM PROJECTIONS
text: Theorem A.1. Let X be a compact convex set and P : IR → C be an invertible function where C ⊂ IR is a convex open set such that X = C̄ (C̄ denotes the closure of C). Now, if
1. P is strictly monotonically increasing. 2. limx→∂C ‖P−1(x)‖ =∞ (∂C denotes the boundary of C).
Then, Φ(x) = ∫ x x0 P−1(y)dy is a valid mirror map.
Proof. From the fundamental theorem of calculus, the gradient of Φ(x) satisfies, ∇Φ(x) = P−1(x). Since P is strictly monotonically increasing and invertible, P−1 is strictly monotonically increasing. Therefore, Φ(x) is strictly convex and differentiable. Now, from the definition of projection and since it is invertible (i.e., P−1 is one-to-one and onto), ∇Φ(C) = P−1(C) = IR. Therefore, together with condition (2), we can conclude that Φ(x) = ∫ x x0 P−1(y)dy is a valid mirror map (refer Definition 2.2 in the main paper). For the multi-dimensional case, we need an additional condition that the vector field P−1(x) is conservative. Then by the gradient theorem (gra), there exists a mirror map Φ(x) = ∫ x x0 P−1(y)dy for some arbitrary base point x0.

heading: B MD UPDATE DERIVATION FOR THE TANH PROJECTION
text: We now derive the MD update corresponding to the tanh projection below. From Theorem A.1, the mirror map for the tanh projection can be written as:
Φ(w) = ∫ P−1(w)dw = 1
2β
[ (1 + w) log(1 + w) + (1− w) log(1− w) ] . (24)
Correspondingly, the Bregman divergence can be written as:
DΦ(w, v) = Φ(w)− Φ(v)− Φ′(v)(w − v) , where Φ′(v) = 12β log 1+v 1−v , (25)
= 1
2β
[ w log
(1 + w)(1− v) (1− w)(1 + v)
+ log(1− w)(1 + w)− log(1− v)(1− v) ] .
Now, consider the proximal form of MD update
wk+1 = argmin x∈(−1,1) 〈η gk, w〉+DΦ(w,wk) . (26)
The idea is to find w such that the KKT conditions are satisfied. To this end, let us first write the Lagrangian of Eq. (26) by introducing dual variables y and z corresponding to the constraints w > −1 and w < 1, respectively:
F (w, x, y) = ηgkw + 1
2β
[ w log
(1 + w)(1− wk) (1− w)(1 + wk)
+ log(1− w)(1 + w)− log(1− wk)(1− wk) ]
(27) + y(−w − 1) + z(w − 1) .
Now, setting the derivatives with respect to w to zero:
∂F ∂w = ηgk + 1 2β log (1 + w)(1− wk) (1− w)(1 + wk) − y + z = 0 . (28)
From complementary slackness conditions,
y(−w − 1) = 0 , since w > −1 ⇒ y = 0 , (29) z(w − 1) = 0 , since w < 1 ⇒ z = 0 .
Therefore, Eq. (28) now simplifies to:
∂F ∂w = ηgk + 1 2β log (1 + w)(1− wk) (1− w)(1 + wk) = 0 , (30)
log (1 + w)(1− wk) (1− w)(1 + wk) = exp(−2βηgk) ,
1 + w 1− w =
1 + wk 1− wk exp(−2βηgk) ,
w = 1+wk 1−wk exp(−2βηg k)− 1
1+wk 1−wk exp(−2βηgk) + 1 .
The pseudocodes of original (MD-tanh) and numerically stable versions (MD-tanh-S) for tanh are presented in Algorithms 1 and 2 respectively.

heading: C ADDITIONAL EXPERIMENTS
text: We first give training curves of all compared methods, and provide ternary quantization results as a proof of concept. Later, we provide experimental details.
Convergence Analysis. The training curves for CIFAR-10 and CIFAR-100 datasets with ResNet-18 are shown in Fig. 4. Notice, after the initial exploration phase (due to low β) the validation accuracies of our MD-tanh-S increase sharply while this steep rise is not observed in regularization methods such as PQ. The training behaviour for both our stable MD-variants (softmax and tanh) is quite similar.
Algorithm 1 MD-tanh
Require: K, b, {ηk}, ρ > 1,D, L Ensure: w∗ ∈ Qm
1: w0 ∈ IRm, β ← 1 . Initialization 2: w0 ← tanh(βw0) . Projection 3: for k ← 0, . . . ,K do 4: Db = {(xi,yi)}bi=1 ∼ D . Sample a mini-batch 5: gk ← ∇wL(w;Db) ∣∣ w=wk
. Gradient w.r.t. w at wk (Adam based gradients) 6: for j ← 1, . . . ,m do 7: wk+1j ← 1+wkj 1−wk j exp(−2βηkgkj )−1
1+wk j 1−wk j
exp(−2βηkgkj )+1 . MD update
8: end for 9: β ← ρβ . Increase β 10: end for 11: w∗ ← sign(w̃K) . Quantization
Algorithm 2 MD-tanh-S
Require: K, b, {ηk}, ρ > 1,D, L Ensure: w∗ ∈ Qm
1: w̃0 ∈ IRm, β ← 1 . Initialization 2: for k ← 0, . . . ,K do 3: wk ← tanh(βw̃k) . Projection 4: Db = {(xi,yi)}bi=1 ∼ D . Sample a mini-batch 5: gk ← ∇wL(w;Db) ∣∣ w=wk
. Gradient w.r.t. w at wk (Adam based gradients) 6: w̃k+1 ← w̃k − ηkgk . Gradient descent on w̃ 7: β ← ρβ . Increase β 8: end for 9: w∗ ← sign(w̃K) . Quantization
Ternary Quantization. As a proof of concept for our shifted tanh projection (refer Example 3.3), we also show results for ternary quantization with quantization levels Q = {−1, 0, 1} in Table 2. Note that the performance improvement of our ternary networks compared to their respective binary networks is marginal as only 0 is included as the 3rd quantization level. In contrast to us, the baseline method PQ (Bai et al. (2019)) optimizes for the quantization levels (differently for each layer) as well in an alternating optimization regime rather than fixing it to Q = {−1, 0, 1}. Also, PQ does ternarize the first convolution layer, fully-connected layers and the shortcut layers. We cross-validate hyperparameters for both the original PQ setup and the equivalent setting of our MD-variants where we optimize all the weights and denote them as PQ* and PQ respectively.
Our MD-tanh variant performs on par or sometimes even better in comparison to tanh projection results where gradient is calculated through the projection instead of performing MD. This again empirically validates the hypothesis that MD yields in good approximation for the task of network quantization. The better performance of PQ in their original quantization setup, compared to our approach in CIFAR-10 can be accounted to their non-quantized layers and different quantization levels. We believe, similar explorations are possible with our MD framework as well.
Experimental Details. As mentioned in the main paper the experimental protocol is similar to (Ajanthan et al. (2019)). To this end, the details of the datasets and their corresponding experiment setups are given in Table 3. For CIFAR-10/100 and TinyImageNet, VGG-16 (Simonyan & Zisserman (2015)) and ResNet-18 (He et al. (2016)) architectures adapted for CIFAR dataset are used. In particular, for CIFAR experiments, similar to (Lee et al. (2019)), the size of the fully-connected (FC) layers of VGG-16 is set to 512 and no dropout layers are employed. For TinyImageNet, the stride of the first convolutional layer of ResNet-18 is set to 2 to handle the image size (Huang et al. (2017)). In all the models, batch normalization (Ioffe & Szegedy (2015)) (with no learnable parameters) and
ReLU nonlinearity are used. Only for the floating point networks (i.e., REF), we keep the learnable parameters for batch norm. Standard data augmentation (i.e., random crop and horizontal flip) is used.
For both of our MD variants, hyperparameters such as the learning rate, learning rate scale, annealing hyperparameter β and its schedule are crossvalidated from the range reported in Table 4 and the chosen parameters are given in the Tables 5 and 6. To generate the plots, we used the publicly available codes of BC3, PQ4 and PMF5.
All methods are trained from a random initialization and the model with the best validation accuracy is chosen for each method. Note that, in MD, even though we use an increasing schedule for β to enforce a discrete solution, the chosen network may not be fully-quantized (as the best model could be obtained in an early stage of training). Therefore, simple argmax rounding is applied to ensure that the network is fully-quantized.
3https://github.com/itayhubara/BinaryNet.pytorch 4https://github.com/allenbai01/ProxQuant 5https://github.com/tajanthan/pmf
Dataset Image # class Train / Val. b K
MNIST 28× 28 10 50k / 10k 100 20k CIFAR-10 32× 32 10 45k / 5k 128 100k CIFAR-100 32× 32 100 45k / 5k 128 100k TinyImageNet 64× 64 200 100k / 10k 128 100k
Table 3: Experiment setup. Here, b is the batch size and K is the total number of iterations for all the methods.
Hyperparameters Fine-tuning grid
learning_rate [0.1, 0.01, 0.001, 0.0001] lr_scale [0.1, 0.2, 0.3, 0.5]
beta_scale [1.01, 1.02, 1.05, 1.1, 1.2] beta_scale_interval [100, 200, 500, 1000, 2000]
Table 4: The hyperparameter search space for all the experiments. Chosen parameters are given in Tables 5 and 6.

Review summaries:
This paper proposes a Mirror Descent (MD) framework for the quantization of neural networks, which, different with previous quantization methods, enables us to derive valid mirror maps and the respective MD updates. Moreover, the authors also provide a stable implementation of MD by storing an additional set of auxiliary dual variables. Experiments on CIFAR-10/100 and TinyImageNet with convolutional and residual architectures show the effective of the proposed model. 

This paper proposes a neural network (NN) quantization based on Mirror Descent (MD) framework. The core of the proposal is the construction of the mirror map from the unconstrained auxiliary variables to the quantized space. Building on that core, the authors derive some mapping functions from the corresponding projection, i.e. tanh, softmax and shifted tanh. The experimental result on benchmark datasets (CIFAR & TinyImageNet) and basic architectures (VGG & ResNet-18) showed that the proposed method is suitable for quantization. The proposed method is a natural extension of ProxQuant, which adopted the proximal gradient descent to quantize NN (a.k.a $\ell_2$ norm in MD). Different projections in NN quantization lead to different Bregman divergences in MD.