M IRROR
D ESCENT
V IEW
FOR
N EURAL
N ETWORK
Q UANTIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
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 modiﬁed 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.
Speciﬁcally, 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 (continuous).
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 classiﬁcation 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 ﬂoating-point networks.
I NTRODUCTION
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 min x ∈X f ( x ) , where f ( · ) denotes the loss function by abstracting out the dependency on the dataset and X ⊂
IR r denotes the set of all possible quantized solutions.
Majority of the works in the literature (Hubara et al.
; Yin et al.
; Ajanthan et al.
(2019)) convert this into an unconstrained problem by introducing auxiliary variables ( ˜ x ) and optimize via (stochastic) gradient descent.
Speciﬁcally, the objective and the update step take the following form:
[Equation]
where P : IR r → 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.
In this work, by noting that the well-known Mirror Descent ( MD ) algorithm (Bubeck (2015)) provides a theoretical framework to perform gradient descent in the unconstrained space (dual space, IR r ) 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 deﬁned based on the associated distance metric.
Brieﬂy, the key ingredient of MD is a concept called mirror map which deﬁnes both the mapping between primal and dual spaces and the exact form of the projection.
Speciﬁcally, 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.
, 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.
; 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 classiﬁcation datasets with convolutional and residual architectures.
Our experiments show that the quantized networks obtained by the MD variants yield accuracies very close to the ﬂoating-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 NN s.
P RELIMINARIES
We ﬁrst 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.
2. M IRROR D ESCENT
The Mirror Descent ( MD ) algorithm is ﬁrst 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:
[Equation]
where f : X → IR is a convex function and X ⊂ IR r is a compact convex set.
The main concept of MD is to extend gradient descent to a more general non-Euclidean space (Banach space 1 ), 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 efﬁciently.
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 ﬁrst mapping the primal point x k ∈ B (quantized space, X ) to the dual space B ∗ (unconstrained space, IR r ), then performing gradient descent in the dual space, and ﬁnally mapping back the resulting point to the primal space B .
If the new point x k +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 .
Speciﬁcally, the gradient of the mirror map deﬁnes the mapping from primal to dual and the projection is done via the Bregman divergence of the mirror map.
We ﬁrst provide the deﬁnitions for mirror map and Bregman divergence and then turn to the MD updates.
Deﬁnition 2.1 (Mirror map) .
Let C ⊂
IR r be a convex open set such that X ⊂ ¯ C ( ¯ C denotes the closure of set C ) and X ∩ C (cid:54) = ∅ .
Then, Φ : C → IR is a mirror map if it satisﬁes:
Φ is strictly convex and differentiable.
∇ Φ( C ) =
IR r , i
e ., ∇ Φ takes all possible values in IR r . 3.
lim x → ∂ C (cid:107)∇ Φ( x ) (cid:107) = ∞ ( ∂ C denotes the boundary of C ), i e ., ∇ Φ diverges on the boundary of C .
Deﬁnition 2.2 (Bregman divergence) .
Let Φ : C → IR be a continuously differentiable, strictly convex function deﬁned 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 ﬁrst-order Taylor expansion of Φ around point q evaluated at point p , i.e .,
[Equation]
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 x 0 ∈ argmin x ∈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:
[Equation]
where g k ∈ ∂f ( x k ) and y k +1 ∈ C ,
[Equation]
[Equation]
Note that, in Eq.
(4), the gradient g k is computed at x k ∈ X ∩ C (solution space) but the gradient descent is performed in IR r (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)): [Equation] k
[Equation]
Note, if Φ( x ) = 12 (cid:107) x (cid:107) 22 , then D Φ ( x , x k ) =
12 (cid:13)(cid:13) x − x k (cid:13)(cid:13) 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.
2. N EURAL N ETWORK Q UANTIZATION
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.
2.2. P ARAMETER S PACE F ORMULATION
Given a dataset D = { x i , y i } ni =1 , NN quantization can be written as:
[Equation]
Here, (cid:96) ( · ) 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 predeﬁned 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.
and its variants (Hubara et al.
; Rastegari et al.
(2016)), where the constraint set is discrete.
In contrast, recent approaches (Bai et al.
; Yin et al.
(2018)) relax this constraint set to be its convex hull:
[Equation]
where q min and q max represent the minimum and maximum quantization levels, respectively.
In this case, a quantized solution is obtained by gradually increasing an annealing hyperparameter.
2.2. L IFTED P ROBABILITY S PACE F ORMULATION
Another formulation is based on the Markov Random Field ( MRF ) perspective to NN quantization recently studied in (Ajanthan et al.
It treats Eq.
(6) as a discrete labelling problem and introduces indicator variables
u j : λ ∈ { 0 , 1 } for each parameter w j where j ∈ { 1 , . . .
, m } such that u j : λ = 1 if and only if w j = λ ∈ Q .
For convenience, by denoting the vector of quantization levels as q , a parameter vector w ∈ Q m can be written in a matrix vector product as:
[Equation]
Here, u is a m × d matrix ( i.e ., each row
u
j = { u j : λ | λ ∈ Q} ), and q is a column vector of dimension d .
Note that, u ∈ V m is an overparametrized ( i.e ., lifted) representation of w ∈ Q m .
Now, similar to the relaxation in the parameter space, one can relax the binary constraint in V m to form its convex hull:
[Equation]
The set ∆ m is in fact the Cartesian product of the standard ( d − 1) -probability simplexes embedded in IR d .
Therefore, for a feasible point u ∈ ∆ m , the vector u j for each j ( j -th row of matrix u ) belongs to the probability simplex ∆ .
Hence, we can interpret the value u j : λ as the probability of assigning the discrete label λ to the weight w j .
This relaxed optimization can then be written as:
[Equation]
In fact, this can be interpreted as ﬁnding a probability distribution u ∈ ∆ m such that the cost L ( u ) is minimized.
Note that, the relaxation of u from V m to ∆ m translates into relaxing w from Q m to the convex region conv( Q m ) .
Even in this case, a discrete solution u ∈ V m can be enforced via an annealing hyperparameter or using rounding schemes.
M IRROR D ESCENT F RAMEWORK FOR N ETWORK Q UANTIZATION
Before introducing the MD formulation, we ﬁrst write NN quantization as a single objective unifying (6) and (10) as:
[Equation]
where f ( · ) denotes the loss function by abstracting out the dependency on the dataset D , and X denotes the constraint set ( Q m or V m 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 Deﬁnition 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 ﬁnding a strictly monotone map is as hard as ﬁnding 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 ﬁeld
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.
lim x → ∂ C (cid:107) P − 1 ( x ) (cid:107) =
∞ ( ∂ C denotes the boundary of C ).
Then, Φ( x ) = (cid:82) x x
P
−
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.
Note that, to employ MD to the problem (11), in theory, any mirror map satisfying Deﬁnition 2.
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.
Even though we consider differentiable projections, Theorem 3.
Example 3.1 ( w -space, binary, tanh ) .
Consider the tanh function, which projects a real value to the interval [ − 1 , 1] :
[Equation]
where β > 0 is the annealing hyperparameter and when β → ∞ , tanh approaches the step function.
The inverse of the tanh is:
[Equation] [Equation]
Note that, P − 1 is monotonically increasing for a ﬁxed β .
Correspondingly, the mirror map from Theorem 3.1 can be written as:
[Equation]
Here, the constant from the integration is ignored.
It can be easily veriﬁed 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:
[Equation] [Equation]
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
=
Note that the sign function has been used for binary quantization in (Courbariaux et al.
and tanh can be used as a soft version of sign function as pointed out by (Zhang et al.
Mirror map corresponding to tanh is used for online linear optimization in (Bubeck et al.
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.
to optimize in the lifted probability space.
In this case, the projection is deﬁned as P (˜ u ) :=
softmax( β ˜ u ) where P : IR d → 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 .,
[Equation]
where
[Equation]
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 deﬁne the inverse of softmax as follows: given u ∈ ∆ , ﬁnd a unique point ˜ v = ˜ u
+ c 1 , for a particular scalar c , such that u =
softmax(˜ v ) .
Now, by choosing c = − log( (cid:80) µ = Q exp(˜ u µ )) , softmax can be written as:
[Equation]
where
[Equation] [Equation]
Now, the inverse of the projection can be written as:
[Equation]
where
[Equation]
Indeed, log is a monotonically increasing function and from Theorem 3.1, by summing the integrals, the mirror map can be written as:
[Equation]
Here, (cid:80) λ 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
(a) tanh , its inverse, and mirror map --  2 - 1.
- 1.0 - 0.5 0.0 0.5 1.0 --   (b) shifted tanh and its inverse
Moreover, when β → ∞ , the projections approaches their respective hard versions.
Gradient Descent ( EGD ) (or Entropic Descent Algorithm ( EDA ))
(Beck & Teboulle (2003); Bubeck (2015)).
Consequently, the update takes the following form:
[Equation]
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 deﬁne our shifted tanh projection P :
IR → C as:
[Equation]
where β > 0 and w = ¯ C = X
=
[ − 1 , 1] .
When β → ∞ , P approaches a stepwise function with inﬂection points at − 0 .
5 and 0 .
5 (here, ± 0 .
5 is chosen heuristically), meaning w move towards one of the quantization levels in the set Q .
This behaviour together with its inverse is illustrated in Fig. 2b.
Now, one could potentially ﬁnd 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 space X 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 ﬁnal 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.
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. can be performed by storing auxiliary variables
˜ x
= P − 1 ( x ) .
In fact, once the auxiliary variable ˜ x k is updated using gradient g k , 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. 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:
[Equation]
projection to the primal space where η > 0 , g k ∈ ∂f ( x k ) and
˜ x k
=
P − 1 ( x k ) .
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) J P = ∂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.
; Yin et al.
Despite being a crude approximation, STE has shown to be highly effective for NN quantization with various network architectures and datasets (Yin et al.
; Zhou et al.
We believe that our MD view of STE provides a justiﬁcation for 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.
However, no analysis of whether an analogous mirror map exists to the given projection is provided and their ﬁnal algorithm is not based on MD .
In particular, following our notation, the ﬁnal update equation of PQ can be written as:
[Equation]
assumes x k and g k are in the same space (23) [Equation] prox : IR r → IR r is the proximal mapping deﬁned in (Bai et al.
where η > 0 , and g k ∈ ∂f ( x k ) . Note that, as opposed to MD (refer to Eq. , PQ assumes the point x k and gradient g k are in the same space. Then only the formula x k − η g k 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 ﬁrst mapping the primal point x k 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 NN s could constitute to a completely new theoretical paper.
R ELATED W ORK
In this work we consider parameter quantization, which is usually formulated as a constrained problem and optimized using a modiﬁed projected gradient descent algorithm, where the methods (Courbariaux et al. (2015); Carreira-Perpinán & Idelbayev (2017);
Yin et al.
; Bai et al.
; 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.
In fact, the success of STE 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.
and/or gradients (Zhou et al.
Furthermore, there are theoretical works on analyzing the convergence properties of STE in certain special cases (Li et al.
; Yin et al.
In this work, by showing that STE based gradient descent as an implementation method of MD under certain conditions on the projection, we provide a justiﬁcation on the effectiveness of STE .
Moreover, there are methods focusing on loss aware quantization (Hou et al.
, quantization for specialized hardware (Esser et al.
(2015)), and quantization based on the variational approach (Achterhold et al.
; Louizos et al.
We have only provided a brief summary of relevant methods and for a comprehensive survey we refer the reader to (Guo (2018)).
E XPERIMENTS
Due to the popularity of binary neural networks (Courbariaux et al.
; Rastegari et al.
, 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 ﬂoating-point counterparts.
We evaluate two MD variants corresponding to tanh and softmax projections, on CIFAR-10 , CIFAR100 and TinyImageNet 2 datasets with VGG-16 and ResNet-18 architectures.
We also evaluate the numerically stable versions of our MD variants (denoted with sufﬁx “- S ”) performed by storing auxiliary parameters during training as explained in Eq.
The results are compared against parameter quantization methods, namely BinaryConnect ( BC )
(Courbariaux et al.
(2015)), ProxQuant ( PQ ) (Bai et al.
and Proximal Mean-Field ( PMF ) (Ajanthan et al.
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.
Brieﬂy, our results indicate that the binary networks obtained by the MD variants yield accuracies very close to the ﬂoating-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 [Table]
g ., up to 3% improvement for PMF .
We provide the hyperparameter tuning search space and the ﬁnal hyperparameters in
Appendix C. Our algorithm is implemented in PyTorch (Paszke et al.
and the experiments are performed on NVIDIA Tesla-P100 GPUs.
Our code will be released upon publication.
[Table]
[Table] denotes performance with biases, fully-connected layers, and shortcut layers in ﬂoating-point (original PQ setup) whereas PQ represents full quantization.
PMF * denotes the performance of PMF after crossvalidation similar to our MD -variants and the original results from the paper are denoted as PMF .
Note our MD variants obtained accuracies virtually the same as the best performing method and it outperformed the best method by a large margin in much harder TinyImageNet dataset.
5. R ESULTS
The classiﬁcation accuracies of binary networks obtained by both variants of our algorithm, namely,
MD - tanh and MD - softmax , their numerically stable versions (sufﬁx “- S ”) and the baselines BC ,
PQ , PMF , GD - tanh and the ﬂoating 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 [Figure] binarization methods while narrowing the performance gap between binary networks and ﬂoating 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 believe, the main reason for this empirical variation in results for our MD -variants is due to numerical instability caused by the ﬂoating-point arithmetic of logarithm and exponential functions in Eq.
(15) and Eq.
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.
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 signiﬁcance 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.
D ISCUSSION
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 is left unattended in this paper, which we intend to analyze in a future work.
Appendices
Here, we ﬁrst provide the proof of the theorem and the technical derivations.
Later we give additional experiments and the details of our experimental setting.
A D ERIVING M IRROR M APS FROM P ROJECTIONS
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
P is strictly monotonically increasing.
lim x → ∂ C (cid:107) P − 1 ( x ) (cid:107) =
∞ ( ∂ C denotes the boundary of C ).
Then, Φ( x ) = (cid:82) x x 0 P − 1 ( y ) dy is a valid mirror map.
Proof.
From the fundamental theorem of calculus, the gradient of Φ( x ) satisﬁes, ∇ Φ( 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 deﬁnition 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 ) =
(cid:82) x x 0 P − 1 ( y ) dy
is a valid mirror map (refer
Deﬁnition 2.2 in the main paper).
For the multi-dimensional case, we need an additional condition that the vector ﬁeld P − 1 ( x ) is conservative.
Then by the gradient theorem (gra), there exists a mirror map Φ( x )
=
(cid:82) x x 0 P − 1 ( y )
d y for some arbitrary base point x 0 .
B MD U PDATE D ERIVATION FOR THE TANH P ROJECTION
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:
[Equation]
Correspondingly, the Bregman divergence can be written as:
[Equation]
Now, consider the proximal form of MD update
[Equation]
The idea is to ﬁnd w such that the KKT conditions are satisﬁed.
To this end [Equation] let us ﬁrst write the
Lagrangian of Eq.
(26) by introducing dual variables y and z corresponding to the constraints w > −and w < 1 , respectively:
[Equation]
Now, setting the derivatives with respect to w to zero: ∂F
[Equation]
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 simpliﬁes to:
[Equation]
The pseudocodes of original ( MD - tanh ) and numerically stable versions ( MD - tanh - S ) for tanh are presented in Algorithms 1 and 2 respectively.
C A DDITIONAL E XPERIMENTS
We ﬁrst 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 ResNetare shown in Fig.
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.
Require: K, b, { η k } , ρ > 1 , D , L
Ensure: w ∗ ∈ Q m 1: w 0 ∈ IR m , β ←(cid:46) Initialization [Table] : w 0 ← tanh( β w 0 )
[Table]
: for k ← 0 , [Table] . . , K do 4: [Table]
[Table]
i , y
[Table]
[Table]
[Table]
[Table]
[Table]
: for j ← 1 , .
[Table]
. , m do
[Table]
: end for
[Table]
[Table]
: end for [Table] : w ∗ ← sign( ˜ w [Table]
[Table] Require: K, b, { η k } , ρ > 1 , D , L Ensure: w ∗ ∈ Q m 1: ˜ w 0 ∈ IR m , β ← [Table]
Initialization 2: for k ← 0 , . . , K do 3: w k ← tanh( β ˜ w k )
[Table]
D b = { ( x i , y i ) } bi [Table] ∼ D [Table]
[Table] ← [Table] [Table]
: ˜ w k +
[Table]
← ˜ w k
[Table]
η k g k (cid:46)
Gradient descent on ˜ w 7: β ← ρβ (cid:46) Increase β 8: end for 9: w ∗ ← sign( ˜ w K ) (cid:46)
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 3 rd 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 ﬁxing it to Q = {− 1 , 0 , 1 } .
Also, PQ does ternarize the ﬁrst 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.
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.
architectures adapted for CIFAR dataset are used.
In particular, for CIFAR experiments, similar to (Lee et al.
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 ﬁrst convolutional layer of ResNet-18 is set to 2 to handle the image size (Huang et al.
In all the models, batch normalization (Ioffe & Szegedy (2015))
(with no learnable parameters) and
[Figure]
Algorithm
Space
CIFARCIFAR-
TinyImageNet VGG
[Table] ResNetVGG- ResNetResNet- REF
[Table]
{+w 93 .
+}
{+PQ w 
+}
{+PQ * w 92 .93 .
+}
{+GD - tanh w 91 .
+}
[Table]
ReLU nonlinearity are used.
Only for the ﬂoating point networks ( i.e ., REF ), we keep the learnable parameters for batch norm.
Standard data augmentation ( i.e ., random crop and horizontal ﬂip) 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 BC 3 , PQ 4 and PMF 5 .
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.
[Table] [Table]
space for all the experiments.
Chosen parameters are given in Tables 5 and 6.
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Figure]
[Figure]
beta
[Figure]
[Figure]
interval
[Figure]
with
[Figure]
BC
[Figure]
learning_rate
[Figure]
[Figure]
BC
[Figure]
learning
[Figure]
rate
[Figure]
[Figure]
[Figure]
[Table]
[Table]