Riemannian Manifold Embeddings for Straight-Through Estimator

1 INTRODUCTION . Neural networks can handle many complex tasks due to their large number of trainable parameters and strong nonlinear capabilities ( Krizhevsky et al. , 2012 ) . However , the massive amount of models and calculations hinder the application of neural networks on mobile and miniaturized devices , which naturally comes with constraints on computing power and resources . Neural network quantization is considered an efficient solution in the inference that alleviates the number of parameters and optimizes the computation by reducing the bit width of weights or activations ( Courbariaux et al. , 2016 ; Li et al. , 2016 ; Zhu et al. , 2016 ) . Existing neural network quantization methods can be roughly divided into two categories : “ STE ” and “ Non-STE ” methods . Most of the quantization methods adopted by the QNNs belong to the former , i.e . there is always a non-differentiable quantization function during training . The role of STE is to penetrate this non-differentiable quantization function and pass the gradients in backpropogation ( Hinton , 2012 ) , e.g . DeepShift ( Elhoushi et al. , 2019 ) , INT8 ( Zhu et al. , 2020 ) , AQE ( Chen et al. , 2020 ) , etc . “ Non-STE ” methods refer to maintain feasible quantization during training , which does not need to apply STE directly to all full-precision weights . For example , Zhou et al . ( Zhou et al. , 2017 ) divided the weights into two groups until all parameters are quantized , where the first group is directly quantized and fixed ; the second group needs to be retrained to make up for the decrease of accuracy caused by quantization of the first group . Louizos et al . ( Louizos et al. , 2019 ) introduced a differentiable quantizer that can transform the continuous distributions of weights and activations to categorical distributions with gradient-based optimization . However , “ Non-STE ” methods face the setting and influence of heavy hyper-parameters in the training process . Relatively , “ STE ” methods are wider choices for quantized models from simplicity and versatility . To approximate ∇W w.r.t . W and complete stable quantization training , Courbariaux et al . ( Courbariaux et al. , 2016 ) binarized the neural networks by the approximated gradients using STE : ∇W = ∇Ŵ ◦ I , where ∂Ŵ ∂W = I = { 1 if |W | ≤ 1 0 otherwise . However , it will inevitably bring the gradient mismatch problem raised by just simply application of STE . To overcome this challenge , Zhou et al . ( Zhou et al. , 2016 ) proposed to transformW to W̃ : W̃ = tanh ( W ) max ( | tanh ( W ) | ) , and then quantize it using a quantization function Q ( · ) . During backpropogation , the gradients∇W w.r.t . W can be further computed using the chain rule : ∇W = ∇Q ( W̃ ) 1− tanh2 ( W ) max ( | tanh ( W ) | ) . Based on this work , Chen et al . ( Chen et al. , 2019 ) proposed to learn ∇W by fully-connected neural networks or LSTM , which replaces this gradients via a meta quantizer Mφ parameterized by φ across layers : ∇W =Mφ ( ∇Q ( W̃ ) , Q ( W̃ ) ) ∂W̃ ∂W . However , such methods not only add many additional parameters , but also increase the difficulty of the training of quantized models . In this paper , we introduce the Manifold Quantization ( ManiQuant ) to train a quantized model via embedding Riemannian manifolds for STE . In Figure 1 , we treat the parameter space in a quantized model as a Riemannian manifold , which will alleviate the gradient mismatch problem caused by non-differentiable quantization and assist quantized models in achieving a more stable convergence and better performance . The main contributions of this work are three-fold : First , we propose to use Fisher Information Matrix embedding to alleviate the gradient mismatch problem . Second , we define a novel Hyperbolic divergence by a convex function with geometric structure . With the constraint of Hyperbolic divergence , we introduce a weak curvature manifold that forms the background of ManiQuant . Third , based on the Second , we propose to use weak curvature metric embeddings for STE as the weak curvature gradient to approximate ∇W without extra parameters and training complexity . 2 RELATED WORK . 2.1 GENERAL GRADIENT . For a neural network with l layers , the full-precision weight matrix of each layer is just marked as Wi . By defining the operation vec ( · ) that vectorizes matrices by stacking their columns together , the total parameter vector of the neural network can be denoted as θ = [ vec ( W1 ) > , vec ( W2 ) > , . . . , vec ( Wl ) > ] > . The parameter vector is given as a column vector . Let θ ∈ Rn be a parameter space on which a loss function L associated with weight is well-defined . It is relatively easy to express the general gradient in training : ∇θL = ∂L ∂θ , ( 1 ) which is the steepest descent method in Euclidean space with an orthonormal coordinate . Note that the negative gradient represents the direction of the steepest descent . When the Euclidean space is considered , the Euclidean divergence between two sufficiently closed points θ and θ′ is actually defined by default : DE [ θ : θ ′ ] = 1 2 ∑ i ( θi − θ′i ) 2 , ( 2 ) which is a half of the square of the Euclidean distance identified by the Euclidean metric δij , so that ds2 = 2DE [ θ : θ + dθ ] = ∑ i ( dθi ) 2 = ∑ i , j δijdθidθj . ( 3 ) 2.2 NATURAL GRADIENT . The general gradient only considers the parameter update along the gradient direction that does not involve the metric tensor for the parameter space of the problem . Now , we attach a Riemannian metric Gij ( θ ) for the parameter space to form a Riemannian manifold ( M , Gij ) ( Amari , 1998 ; 2016 ) . In that case , the steepest descent direction also depends on the quadratic form introduced by a small incremental vector dθ that connects θ and θ + dθ , whose form is given by ds2 = ∑ i , j Gij ( θ ) dθidθj , ( 4 ) where dθi is the component of dθ . Under the constraint ds2 , the steepest descent direction is toward the optimization goal of L ( θ + dθ ) . Intuitively , it measures the Kullback-Leibler ( KL ) divergence between two distributions p ( x|θ ) and p ( x|θ + dθ ) of network model , which is equivalent to the interval of two adjacent points on a Riemannian manifold . The KL divergence under the Riemannian metric is well approximated through the Fisher Information Matrix ( FIM ) with the second-order Taylor expansion of the KL divergence ( see Appendix B.3 ) ( Ba et al. , 2016 ) : ds2 = 2DKL [ p ( x|θ ) : p ( x|θ + dθ ) ] ≈ −dθ > Ep ( x|θ ) [ ∇ log p ( x|θ ) ∇ log p ( x|θ ) > ] dθ . ( 5 ) We see that FIM is equal to the negative expected Hessian of log likelihood . Furthermore , Amari deduced that the Riemannian metric is given by the FIM ( Amari , 1998 ) . By the Lagrangian form , we have the natural gradient ( see Appendix C.1 ) ∇̃θL = F−1 ( θ ) ∂L ∂θ , whereF ( θ ) = Ep ( x|θ ) [ ∇ log p ( x|θ ) ∇ log p ( x|θ ) > ] , ( 6 ) which is the steepest descent method in a Riemannian space and F ( θ ) is the FIM with a parameter vector ( the natural gradient also known as the Riemannian gradient ) . Empirically , the immediate application of FIM is as drop-in replacement of Hessian in second order optimization algorithm . Note that ∇̃θL will return to∇θL when Gij ( θ ) is equal to the Euclidean metric δij , i.e . the identity matrix I . 3 MANIFOLD QUANTIZATION . Considering a quantized model of each layer , we notate the quantized weight matrix using Ŵi for differentiating full-precision weight matrix Wi . Note that we define the quantized parameter vector as θ̂ , which is similar to the full-precision parameter vector θ . During training a QNN , the quantization function Q ( · ) is a one-to-one mapping from full-precision values to quantized values , which can be expressed as θ̂ = Q ( θ ) . By involving the practice to train a neural network with low bit-width , the process of the quantization needs to be further designed that plays a vital role in the final performance , as for the other parts , e.g . via mapping from an input pattern âi−1 to output âi can still be imitated in the same way as full-precision neural networks : si = Ŵiâi−1 , âi = Q ( fi ◦ si ) , ( 7 ) where fi is non-linear function acted on element-wise . To distinguish the full-precision model , we mark the notation “ ˆ ” to represent the operation through a quantization function in a quantized model , whether for weight matrix Ŵi or activation vector âi . 3.1 FISHER INFORMATION MATRIX EMBEDDING FOR STE . The concept of the natural gradient is closely related to the FIM and KL divergence ( see Appendix B.3 ) . Since the KL divergence is intrinsic , the natural gradient is also intrinsically invariant under the parameter transformation . By viewing the FIM as an l by l block matrix where l denotes the number of layers in a neural network , the natural gradient formula can be introduced to alleviate the gradient mismatch problem when training a QNN and updating its parameters . Lemma 1 In the Riemannian manifold defined by KL divergence , Fisher Information Matrix embedding for Straight-Through Estimator is written by : ∇̃θL STE = F−1 ( θ̂ ) ∇θ̂L ◦ I|θ|≤1 = E−1 [ d log p ( x|θ̂ ) dθ̂ d log p ( x|θ̂ ) dθ̂ > ] ∇θ̂L ◦ I|θ|≤1 = E−1 [ vec ( ∂L ∂θ̂ ) vec ( ∂L ∂θ̂ ) > ] ∇θ̂L ◦ I|θ|≤1 . ( 8 ) Furthermore , F−1 ( θ̂ ) can be expressed as E [ vec ( ∂L ∂Ŵ1 ) vec ( ∂L ∂Ŵ1 ) > ] · · · E [ vec ( ∂L ∂Ŵ1 ) vec ( ∂L ∂Ŵl ) > ] ... . . . ... E [ vec ( ∂L ∂Ŵl ) vec ( ∂L ∂Ŵ1 ) > ] · · · E [ vec ( ∂L ∂Ŵl ) vec ( ∂L ∂Ŵl ) > ] −1 ( 9 ) where vec ( ∂L ∂Ŵi ) can be represented by the gradient error ∂L∂ai . Proof . The proofs are the combination between STE and Appendix C.1 . Considering that the gradient propagation needs to span over the quantized neurons and graphs , we still need to use STE to update the gradient of QNNs in the learning procedure : vec ( ∂L ∂Ŵi ) = â > i−1 ⊗ ( ∂L ∂ai ◦ f ′i ( si ) ) STE = â > i−1 ⊗ ∂L ∂âi ◦ ( f ′i ( si ) I|ai|≤1 ) , ( 10 ) where⊗ denotes the Kronecker product 1 . Relying on the KL divergence , we indicate the parameter space as the Riemannian manifold rather than the Euclidean space in the quantization procedure . For neural networks with a scale of one million or more parameters , the time complexity of inverting the FIM , a component of natural gradients , is O ( n3 ) ( Povey et al. , 2014 ) . Previously , there were some works to calculate the natural gradient efficiently , e.g . Roux et al . ( Roux et al. , 2008 ) decomposed the FIM into multiple diagonal blocks , where each diagonal block is approximated by a low-rank matrix . Bastian et al . ( Bastian et al. , 2011 ) also used the idea of diagonal blocks by constructing a diagonal block that corresponds to a weight matrix . Martens et al . ( Martens & Grosse , 1Since the size of â > i−1 and the size of ∂L∂ai may be different in each layer , Kronecker product is necessary . 2015 ) proposed to approximate the FIM through the Kronecker product of two smaller matrices to improve computational efficiency . Even so , FIM embedding for STE with complex decomposition methods is still not suitable for large-scale tasks , and the computational amount is large compared to the standard STE .

Summary: This paper proposed to improve the gradient of quantization operator in quantization-aware training (QAT). Basically, it pointed out that the Euclidean assumption in general gradient (used by Straight-Through Estimator) can not reflect the curvature in the loss surface. Instead, it revised the gradient (natural gradient) based on mainfold learning. Due to the huge cost of Fisher Information Matrix (FIM) used to achieve natural gradient, it proposed to approximate it by weak curvature embedding.