Mirror Descent View For Neural Network Quantization

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 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 . 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 : 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 . 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 = { 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 . 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 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 .

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