Title: MIRROR DESCENT VIEW FOR NEURAL NETWORK QUANTIZATION

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.

1. INTRODUCTION:

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. (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:

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. (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, 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 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.

2. PRELIMINARIES:

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.

2.1. MIRROR DESCENT:

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:

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 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 x k ∈ B (quantized space, X ) to the dual space B * (unconstrained space, IR r ), then performing gradient descent in the dual space, and finally 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. 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. 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), pq .

(3) 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

x k+1 = argmin x∈X ∩C D Φ (x, y k+1 ) .

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) ):

x k+1 = argmin x∈X ∩C η g k , x + D Φ (x, x k ) . (5) Note, if Φ(x) = 1 2 x 2 2 , then D Φ (x, x k ) = 1 2 x -x k 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.2. NEURAL NETWORK QUANTIZATION:

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.1. PARAMETER SPACE FORMULATION:

Given a dataset D = {x i , y i } n i=1 , NN quantization can be written as:

EQUATION

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:

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.2. LIFTED PROBABILITY SPACE FORMULATION:

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 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 finding 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.

3. MIRROR DESCENT FRAMEWORK FOR NETWORK QUANTIZATION:

Before introducing the MD formulation, we first 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 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.

lim x→∂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.

Figure 1 : MD formulation where mirror map is derived from the projection P . Note, g k is computed in the primal space (X ) but it is directly used to update the auxiliary variables in the dual space.

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]:

EQUATION

where β > 0 is the annealing hyperparameter and when β → ∞, tanh approaches the step function.

The inverse of the tanh is:

EQUATION

)

Note that, P -1 is monotonically increasing for a fixed β. 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 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:

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 = [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 (ũ) := softmax(β ũ) 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

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 ṽ = ũ + c1, for a particular scalar c, such that u = softmax(ṽ). Now, by choosing c = -log( µ=Q exp(ũ µ )), softmax can be written as:

EQUATION

)

Now, the inverse of the projection can be written as:

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, λ 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 Figure 2 : Plots of tanh and shifted tanh projections, and their inverses corresponding to the tanh projection. Note that, the inverses are monotonically increasing and the mirror map is strictly convex. 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 define our shifted tanh projection P : IR → C as:

EQUATION

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 set Q. 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 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 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 xk 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.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:

EQUATION

x k+1 = P (x k+1 ) ∈ X , projection to the primal space where η > 0, g k ∈ ∂f (x k ) and xk = 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. (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 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:

xk+1 = x k -η g k , assumes

x k and g k are in the same space (23)

x k+1 = prox(x k+1 ) , prox : IR r → IR r is the proximal mapping defined in (Bai et al. (2019) )

where η > 0, and g k ∈ ∂f (x k ). Note that, as opposed to MD (refer to Eq. ( 22)), 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 first mapping the primal point x k to a point xk 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.

4. RELATED WORK:

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 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) ).

5. EXPERIMENTS:

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, CIFAR-100 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. 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.1. RESULTS:

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 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.

6. DISCUSSION:

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.

Appendices:

Here, we first provide the proof of the theorem and the technical derivations. Later we give additional experiments and the details of our experimental setting.

A DERIVING MIRROR MAPS FROM PROJECTIONS:

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. lim x→∂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 x 0 .

B MD UPDATE DERIVATION FOR THE TANH PROJECTION:

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

= 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

EQUATION

)

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) = ηg k w + 1 2β w log (1 + w)(1 -w k ) (1 -w)(1 + w k ) + log(1 -w)(1 + w) -log(1 -w k )(1 -w k ) (27) + y(-w -1) + z(w -1) .

Now, setting the derivatives with respect to w to zero:

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 simplifies to:

EQUATION

(1 + w)(1 -w k ) (1 -w)(1 + w k ) = exp(-2βηg k ) , 1 + w 1 -w = 1 + w k 1 -w k exp(-2βηg k ) , w = 1+w k 1-w k exp(-2βηg k ) -1 1+w k 1-w k exp(-2βηg k ) + 1

.

The pseudocodes of original (MD-tanh) and numerically stable versions (MD-tanh-S) for tanh are presented in Algorithms 1 and 2 respectively.

C ADDITIONAL EXPERIMENTS:

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 * ∈ Q m 1: w 0 ∈ IR m , β ← 1 Initialization 2: w 0 ← tanh(βw 0 ) Projection 3: for k ← 0, . . . , K do for j ← 1, . . . , m do 7:

w k+1 j ← 1+w k j 1-w k j exp(-2βη k g k j )-1 1+w k j 1-w k j exp(-2βη k g k j )+1

MD update 8:

end for 9:

β ← ρβ Increase β 10: end for 11: w * ← sign( wK ) Quantization wk+1 ← wk -η k g k Gradient descent on w 7:

Algorithm 2 MD-tanh-S Require: K, b, {η k }, ρ > 1, D, L Ensure: w * ∈ Q m 1: w0 ∈ IR m , β ← 1 Initialization 2: for k ← 0, . . . , K do 3: w k ← tanh(β wk ) Projection 4: D b = {(x i , y i )} b i=1 ∼ D Sample a mini-batch 5: g k ← ∇ w L(w; D b ) w=w k Gradient w.

β ← ρβ Increase β 8: end for 9: w * ← sign( wK ) 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 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 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. , 200, 500, 1000, 2000] Table 4 : The hyperparameter search space for all the experiments. Chosen parameters are given in Tables 5 and 6 . 

:

A Banach space is a complete normed vector space where the norm is not necessarily derived from an inner product.

https://tiny-imagenet.herokuapp.com/

https://github.com/itayhubara/BinaryNet.pytorch

https://github.com/allenbai01/ProxQuant

https://github.com/tajanthan/pmf

Figures and Tables:

shifted tanh and its inverse

Figure 3: Training curves for binarization for CIFAR-10 (first two) and CIFAR-100 (last two) with ResNet-18. Compared to original MD variants, stable MD variants are less noisy and after the initial exploration phase (up to 60 in CIFAR-10 and 25 epochs CIFAR-100), the validation accuracies rise sharply and show gradual improvement afterwards.

b = {(x i , y i )} b i=1 ∼ DSample a mini-batch 5:g k ← ∇ w L(w; D b ) w=w k Gradient w.r.t. w at w k (Adam based gradients) 6:

Figure4: Training curves for binarization for CIFAR-10 (first two columns) and CIFAR-100 (last two columns) with ResNet-18. Compared to BC, our MD-tanh-S and PMF are less noisy and after the initial exploration phase (up to 60 in CIFAR-10 and 25 epochs CIFAR-100), the validation accuracies rise sharply and closely resembles the floating point network afterwards. This steep increase is not observed in regularization methods such as PQ.

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.

Table content: <table/>

Classification accuracies on the test set for different methods for binary quantization. PQ* denotes performance with biases, fully-connected layers, and shortcut layers in floating-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.

Table content: <table><tr><td/><td>Algorithm</td><td>Space</td><td colspan="4">CIFAR-10 VGG-16 ResNet-18 VGG-16 ResNet-18 CIFAR-100</td><td>TinyImageNet ResNet-18</td></tr><tr><td/><td>REF (float)</td><td>w</td><td>93.33</td><td>94.84</td><td>71.50</td><td>76.31</td><td>58.35</td></tr><tr><td/><td>BC</td><td>w</td><td>89.04</td><td>91.64</td><td>59.13</td><td>72.14</td><td>49.65</td></tr><tr><td/><td>PQ</td><td>w</td><td>85.41</td><td>90.76</td><td>39.61</td><td>65.13</td><td>44.32</td></tr><tr><td/><td>PQ*</td><td>w</td><td>90.11</td><td>92.32</td><td>55.10</td><td>68.35</td><td>49.97</td></tr><tr><td/><td>PMF</td><td>u</td><td>90.51</td><td>92.73</td><td>61.52</td><td>71.85</td><td>51.00</td></tr><tr><td/><td>PMF*</td><td>u</td><td>91.40</td><td>93.24</td><td>64.71</td><td>71.56</td><td>51.52</td></tr><tr><td/><td>GD-tanh</td><td>w</td><td>91.47</td><td>93.27</td><td>60.67</td><td>71.46</td><td>51.43</td></tr><tr><td/><td>MD-softmax</td><td>u</td><td>90.47</td><td>91.28</td><td>56.25</td><td>68.49</td><td>46.52</td></tr><tr><td>Ours</td><td>MD-tanh MD-softmax-S</td><td>w u</td><td>91.64 91.30</td><td>92.97 93.28</td><td>61.31 63.97</td><td>72.13 72.18</td><td>54.62 51.81</td></tr><tr><td/><td>MD-tanh-S</td><td>w</td><td>91.53</td><td>93.18</td><td>61.69</td><td>72.18</td><td>52.32</td></tr></table>

r.t. w at w k (Adam based gradients)

Table content: <table/>

Classification accuracies on the test set for ternary quantization. PQ* denotes performance with fully-connected layers, first convolution layer and shortcut layers in floating point whereas PQ represent results with all layers quantized. Also, PQ* optimize for the quantization levels as well (different for each layer), in contrast we fix it to Q = {-1, 0, 1}. GD-tanh denotes results without using STE and actually calculating the gradient through the projection.

Table content: <table/>

Experiment setup. Here, b is the batch size and K is the total number of iterations for all the methods.

Table content: <table><tr><td>Hyperparameters</td><td>Fine-tuning grid</td></tr><tr><td>learning_rate</td><td>[0.1, 0.01, 0.001, 0.0001]</td></tr><tr><td>lr_scale</td><td>[0.1, 0.2, 0.3, 0.5]</td></tr><tr><td>beta_scale</td><td>[1.01, 1.02, 1.05, 1.1, 1.2]</td></tr><tr><td colspan="2">beta_scale_interval [100</td></tr></table>

Hyperparameter settings used for the binary quantization experiments. Here, the learning rate is multiplied by lr_scale after every 30k iterations and annealing hyperparameter (β) is multiplied by beta_scale after every beta_scale_interval iterations. We use Adam optimizer with zero weight decay. For PQ, beta_scale denotes regularization rate.

Table content: <table><tr><td/><td/><td colspan="3">CIFAR-10 with ResNet-18</td><td/><td/></tr><tr><td/><td colspan="4">REF (float) MD-softmax-S MD-tanh-S GD-tanh</td><td>PQ</td><td>PQ*</td></tr><tr><td>learning_rate</td><td>0.1</td><td>0.001</td><td>0.01</td><td>0.01</td><td>0.01</td><td>0.01</td></tr><tr><td>lr_scale</td><td>0.3</td><td>0.3</td><td>0.2</td><td>0.5</td><td>0.3</td><td>-</td></tr><tr><td>beta_scale (ours)</td><td>-</td><td>1.05</td><td>1.2</td><td>1.02</td><td colspan="2">0.0001 0.0001</td></tr><tr><td>beta_scale_interval</td><td>-</td><td>500</td><td>1000</td><td>500</td><td>-</td><td>-</td></tr><tr><td>weight_decay</td><td>0.0001</td><td>0</td><td>0</td><td>0</td><td>0</td><td>0.0001</td></tr><tr><td/><td/><td colspan="3">CIFAR-100 with ResNet-18</td><td/><td/></tr><tr><td/><td colspan="4">REF (float) MD-softmax-S MD-tanh-S GD-tanh</td><td>PQ</td><td>PQ*</td></tr><tr><td>learning_rate</td><td>0.1</td><td>0.001</td><td>0.001</td><td>0.01</td><td>0.01</td><td>0.001</td></tr><tr><td>lr_scale</td><td>0.1</td><td>0.1</td><td>0.5</td><td>0.5</td><td>0.2</td><td>-</td></tr><tr><td>beta_scale (ours)</td><td>-</td><td>1.1</td><td>1.1</td><td>1.02</td><td colspan="2">0.0001 0.0001</td></tr><tr><td>beta_scale_interval</td><td>-</td><td>100</td><td>500</td><td>1000</td><td>-</td><td>-</td></tr><tr><td>weight_decay</td><td>0.0001</td><td>0</td><td>0</td><td>0</td><td>0</td><td>0.0001</td></tr><tr><td/><td/><td/><td colspan="2">CIFAR-10 with VGG-16</td><td/><td/></tr><tr><td/><td colspan="4">REF (float) MD-softmax-S MD-tanh-S GD-tanh</td><td>PQ</td><td>PQ*</td></tr><tr><td>learning_rate</td><td>0.1</td><td>0.001</td><td>0.01</td><td>0.01</td><td>0.01</td><td>0.1</td></tr><tr><td>lr_scale</td><td>0.2</td><td>0.3</td><td>0.3</td><td>0.3</td><td>-</td><td>-</td></tr><tr><td>beta_scale (ours)</td><td>-</td><td>1.05</td><td>1.1</td><td>1.01</td><td>1e-07</td><td>0.0001</td></tr><tr><td>beta_scale_interval</td><td>-</td><td>500</td><td>1000</td><td>500</td><td>-</td><td>-</td></tr><tr><td>weight_decay</td><td>0.0001</td><td>0</td><td>0</td><td>0</td><td>0</td><td>0.0001</td></tr><tr><td/><td/><td/><td colspan="2">CIFAR-100 with VGG-16</td><td/><td/></tr><tr><td/><td colspan="4">REF (float) MD-softmax-S MD-tanh-S GD-tanh</td><td>PQ</td><td>PQ*</td></tr><tr><td>learning_rate</td><td>0.1</td><td>0.0001</td><td>0.001</td><td>0.01</td><td>0.01</td><td>0.0001</td></tr><tr><td>lr_scale</td><td>0.2</td><td>0.3</td><td>0.5</td><td>0.2</td><td>-</td><td>-</td></tr><tr><td>beta_scale (ours)</td><td>-</td><td>1.05</td><td>1.1</td><td>1.05</td><td colspan="2">0.0001 0.0001</td></tr><tr><td>beta_scale_interval</td><td>-</td><td>100</td><td>500</td><td>2000</td><td>-</td><td>-</td></tr><tr><td>weight_decay</td><td>0.0001</td><td>0</td><td>0</td><td>0</td><td>0</td><td>0.0001</td></tr><tr><td/><td/><td colspan="3">TinyImageNet with ResNet-18</td><td/><td/></tr><tr><td/><td colspan="4">REF (float) MD-softmax-S MD-tanh-S GD-tanh</td><td>PQ</td><td>PQ*</td></tr><tr><td>learning_rate</td><td>0.1</td><td>0.001</td><td>0.01</td><td>0.01</td><td>0.01</td><td>0.01</td></tr><tr><td>lr_scale</td><td>0.1</td><td>0.1</td><td>0.1</td><td>0.5</td><td>-</td><td>-</td></tr><tr><td>beta_scale (ours)</td><td>-</td><td>1.2</td><td>1.2</td><td>1.05</td><td>0.01</td><td>0.0001</td></tr><tr><td>beta_scale_interval</td><td>-</td><td>500</td><td>2000</td><td>2000</td><td>-</td><td>-</td></tr><tr><td>weight_decay</td><td>0.0001</td><td>0</td><td>0</td><td>0</td><td>0</td><td>0.0001</td></tr></table>

Hyperparameter settings used for the ternary quantization experiments. Here, the learning rate is multiplied by lr_scale after every 30k iterations and annealing hyperparameter (β) is multiplied by beta_scale after every beta_scale_interval iterations. We use Adam optimizer except for REF for which SGD with momentum 0.9 is used. For PQ, beta_scale denotes regularization rate.

Table content: <table/>

Peer review (iclr2020):
A good paper that uses the Mirror Descent paradigm for learning quantized networks. Though Mirror Descent is not their original idea, but using it in the context of learning quantized network is novel and interesting. Empirically, they showed better results than existing method, with comparisons with reasonable baselines including using relaxed projected gradient descent. Overall, I don’t have much concerns, but here are some more specific comment/questions (most relates to writing) In the intro, it would be great to mention some past success on using MD, as opposed to just saying it’s well-known. Also you mention MD can be used for more than quantization, but compression in general, it’d be better to add that discussion, or remove this sentence. In the beginning of Section 2.1, it*d be easier for the readers to make clear that the primal space corresponds to the quantized weights and the dual space corresponds to the unconstrained space in the rest of the paper. At the top of page 3 you describe MD for the first time, but it’s unclear to me how y^0 is handled. The end of section 3 and section 4 talk quite a bit about STE, maybe it*d be clear if the authors can provide a concise description. As someone not super familiar with NN quantization, this work seems like a good contribution. My only possible concerns would be somehow comparisons to existing methods are not comprehensive enough (if this will be pointed out by the other reviewers)