Year: 2021

Abstract:
Quantized Neural Networks (QNNs) aim at replacing full-precision weights W with quantized weights Ŵ , which make it possible to deploy large models to mobile and miniaturized devices easily. However, either infinite or zero gradients caused by non-differentiable quantization significantly affect the training of quantized models. In order to address this problem, most training-based quantization methods use Straight-Through Estimator (STE) to approximate gradients ∇W w.r.t. W with gradients ∇Ŵ w.r.t. Ŵ where the premise is that W must be clipped to [−1,+1]. However, the simple application of STE brings with the gradient mismatch problem, which affects the stability of the training process. In this paper, we propose to revise an approximated gradient for penetrating the quantization function with manifold learning. Specifically, by viewing the parameter space as a metric tensor in the Riemannian manifold, we introduce the Manifold Quantization (ManiQuant) via revised STE to alleviate the gradient mismatch problem. The ablation studies and experimental results demonstrate that our proposed method has a better and more stable performance with various deep neural networks on CIFAR10/100 and ImageNet datasets.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK

heading: 2.1 GENERAL GRADIENT
text: 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)

heading: 2.2 NATURAL GRADIENT
text: 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 .

heading: 3 MANIFOLD QUANTIZATION
text: 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.

heading: 3.1 FISHER INFORMATION MATRIX EMBEDDING FOR STE
text: 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.

heading: 3.2 WEAK CURVATURE METRIC EMBEDDING FOR STE

heading: 3.2.1 INTRODUCTION OF WEAK CURVATURE MANIFOLD
text: We can use the Euclidean coordinates to calculate the natural gradient on the manifold because of the local homeomorphism of a manifold (Wald, 2010). Note that the homeomorphic mapping φU satisfies U ∈ M 7→ φU (U) ∈ Rn. For any point x ∈ U , we can define φU (x) as the Euclidean coordinate absolutely.
The FIM embedding for STE poses a significant challenge for computing. Intuitively, the computing F−1(θ̂) attached to the natural gradient will take on massive computation, which can only be numerically estimated. The inversion is mostly unrealistic when deep neural networks are very redundant, with tens of thousands of neural connections.
Inspired by information geometry (Amari, 2016) and mirror descent (Bubeck, 2015), we propose a weak curvature manifold with the weak curvature metric, which implies that the gradient mismatch problem can be alleviated by embedding the weak curvature manifold into the STE, while keeping a low computational complexity. Geometrically, the Riemannian manifold is nearly flat, where the Riemannian metric is an approximation of the Euclidean metric. In practice, we develop a linearized Riemannian metric from the Euclidean metric δij , which is systematically defined in general relativity (Wald, 2010).
Definition 1 (Weak Curvature Metric) Let a linearized Riemannian metric deviates from the Euclidean metric. The deviated metric ij(θ̂) is much smaller than 1 in global Euclidean coordinate system of δij ,
Gij(θ̂) = δij + ij(θ̂), where ∣∣ ij∣∣ 1. (11)

heading: It is an adequate definition of “weakness” in this context and ensuresGij(θ̂) to be a positive-definite metric.

heading: 3.2.2 HYPERBOLIC DIVERGENCE
text: In practice, we can determine a unique geodesic through exponential map expθi(τθ) that maps τθ back toM where τ is a small constant, by defining τθ ∈ TθiM as the tangent vector. Note that the definition of exponential map is developed by (Wald, 2010; Helgason, 2001).
Definition 2 (Exponential Map) LetM be a Riemannian manifold, for the tangent vector v ∈ TxM in a point x ∈ M where TxM is the tangent space, there is a unique geodesic γv(t) locally that satisfies γv(0) = x and γ′v(0) = v. The exponential map expx : TxM 7→ M corresponding to γv(t) is defined as expx(v) = γv(1). When constraining expx to a neighbourhood U , this mapping is one-to-one.
Empirically, the exponential map expθi(τθ) is to map a tangent vector τθ in the tangent bundle to the point where the arc length from point θi is equal to |τθ| on geodesic with the initial condition (θi, τθ). In order to make exp(τθ) and exp(−τθ) have the same effect in the training process, we symmetrize them and derive a convex function
ψ(θ) = ∑ i 1 τ2 log [ exp(τθi) + exp(−τθi) 2 ] = ∑ i 1 τ2 log(cosh(τθi)). (12)
Geometrically, the convex function with geometric structure is introduced into the Bregman divergence (see Appendix B.2) (Bregman, 1967), and we obtain a novel Hyperbolic divergence that satisfies the criteria of divergence (see Appendix B.1).
Definition 3 (Hyperbolic Divergence) For a convex function ψ defined by Eq. 12, the Hyperbolic divergence between θ′ and θ is
DH [θ ′ : θ] = ∑ i [ 1 τ2 log cosh(τθ′i) cosh(τθi) − 1 τ (θ′i − θi) tanh(τθi) ] . (13)

heading: 3.2.3 THE GRADIENT FLOW IN WEAK CURVATURE MANIFOLD
text: Lemma 2 Let dθ → 0, the square of an infinitesimal distance defined by Hyperbolic divergence is
ds2 = 2DH [θ ′ : θ] ≈ dθ> [ δij − tanh(τθ) tanh(τθ)> ] dθ. (14)
Furthermore, the defined metric [ δij − tanh(τθ) tanh(τθ)> ] can be consistent with the weak curvature metric.
Proof. The proofs can be found in Appendix B.4. Empirically, we can give a small constant τ to satisfy the weak curvature metric based on Definition 1. Comparing Eq. 4 and Eq. 11, we can deduce the deviated metric ij = − [ tanh(τθ) tanh(τθ)> ] ij .
Lemma 3 In the Riemannian manifold defined by Hyperbolic divergence, a weak curvature metric embedding for Straight-Through Estimator is approximated by:
∇̃θL STE ≈ [ I + tanh(τθ) tanh(τθ)> ] ∇θ̂L ◦ I|θ|≤1 (15)
Proof. The proofs can be found in Appendix C.2.

heading: 3.3 HIERARCHICAL EMBEDDINGS OF THE WEAK CURVATURE METRIC
text: When embedding a weak curvature manifold, there is an unnatural gap between the calculation of weak curvature metric and the layer-by-layer gradient update in back-propagation. Specifically, the metric is not a block diagonal matrix (each block corresponds to a layer of the neural network).
For the process of the quantization, we have the hierarchical embeddings for STE through decoupling operation:
∇vec(Wi)L STE ≈ [ I + τ2i · vec(tanh(Wi)) vec(tanh(Wi))> ] ∇vec(Ŵi)L ◦ I| vec(Wi)|≤1
= ( I + χ vec(tanh(Wi) vec(tanh(Wi)) >
tr [vec(tanh(Wi)) vec(tanh(Wi))>]
) ∇vec(Ŵi)L ◦ I| vec(Wi)|≤1
=  1 + χ tanh 2(wi1) tr[Pi] χ tanh(wi1) tanh(wi2)tr[Pi] · · · χ tanh(wi2) tanh(wi1)tr[Pi] 1 + χ tanh2(wi2) tr[Pi]
· · · ... ... . . . ∇vec(Ŵi)L ◦ I| vec(Wi)|≤1. (16)
In the first line of the above formula, we can put τ2i outside because the elements of the deviated metric is very small based on Lemma 2. In the second line, we use a normalized formula to uniformly represent τ2i : τ 2 i = χ/ tr [ vec(tanh(Wi)) vec(tanh(Wi)) >] where χ is a constant factor and the matrix [ vec(tanh(Wi)) vec(tanh(Wi))
>] is re-expressed by Pi. By limiting the value range of full-precision weights to [−1,+1], we obtain the independent weak curvature metric embedding in each layer. The algorithm of this weak curvature gradient with STE is illustrated in Appendix E.

heading: 4 EXPERIMENT
text: In this section, we implement experiments to demonstrate the effectiveness of our proposed methods on benchmark datasets that are CIFAR10/100 and ImageNet mainly here. Intuitively, experiments on CIFAR10/100 and ImageNet are ablation studies to validate the advantages of weak curvature manifold embeddings. Comparisons with other training-based quantization on quantitative indicates will be carried out on CIFAR10 and ImageNet. All experiments are conducted with PyTorch.

heading: 4.1 EXPERIMENT SETUP
text: We use a weight decay of 1e-4, a batch size of 128, and SGD optimization method with nesterov momentum of 0.9. For CIFAR, we set total training epochs as 200 where the learning strategy is lowered by 10 times at epoch 80, 150, and 190, with the initial 0.1. For ImageNet, we set total training epochs as 50 where the learning strategy is lowered by 10 times at epoch 15, 30, and 45, with the initial 0.1. Note that we set χ as the reciprocal of output channel in a layer. All experiments are conducted for 5 times, and the statistics of the last 10/5 epochs’ test accuracy are reported as a fair comparison.

heading: 4.2 ABLATION STUDIES
text: In order to illustrate the superiority of the weak curvature gradient (WCG), we train QNNs with our gradient from scratch compared to other gradients, i.e. Dorefa (Zhou et al., 2016), MultiFCG (Chen et al., 2019) and FCGrad (Chen et al., 2019). We choose 1-bit weight quantization for all layers in the networks, which means that each layer contains just −1 or +1. Note that the initialization is Xavier (Glorot & Bengio, 2010). Specifically, we experiment ResNet20/32/44 models (He et al., 2016) on CIFAR10, ResNet56/110 models (He et al., 2016) on CIFAR100 and ResNet18 model (He et al., 2016) on ImageNet. The accuracy of full-precision baseline is reported by (Chen et al., 2019).
The results of ablation studies on CIFAR10, CIFAR100 and ImageNet are shown in Table 1, Table 2 and Table 3 respectively. Various quantized models applying weak curvature gradients show significant improvement in challenging classification tasks. For fairness, we use the same forward quantization process and optimizations in experiments where the only difference is the gradients of the backward process.

heading: 4.3 CONVERGENCE AND STABILITY ANALYSIS
text: In this experiment, we compare the convergence and stability of Dorefa and WCG during the training process. As Figure 2 shows, the quantized model trained with WCG has better performance and more stable convergence than the quantized model trained with Dorefa, both for training and testing. Here, we state that better performance means larger mean, and more stable convergence means smaller variance. As for why WCG performs better, we conjecture that WCG has increased the Fisher information of the original quantized model compared to other methods, although WCG does not use additional neural network information for learning like MultiFCG or FCGrad. Moreover, this is where the advantage of manifold learning lies.

heading: 4.4 COMPARISONS WITH OTHER TRAINING-BASED QUANTIZATION
text: In this experiment, our ManiQuant is initialized with the trained full-precision model, and continues to use WCG to learn the quantized model. ManiQuant focuses on the weak curvature manifold embeddings for STE, without any extra training tricks and the deformation of the quantization. We list the comparison results of 1-bit weight and activation quantization in Table 4. More experimental results can be found in Appendix D.

heading: 4.5 TRAINING TIME ANALYSIS
text: Our hardware environment is conducted with an Intel(R) Xeon(R) Silver 4214 CPU(2.20GHz), GeForce GTX 1080Ti GPU. We test the training time per iteration as for MetaQuant using MultiFCG, ManiQuant using WCG, and Dorefa with ResNet20 in CIFAR10. For finishing one iteration of training, ManiQuant costs 18.44s and MetaQuant costs 25.34s while Dorefa uses 17.81s.

heading: 5 CONCLUSION
text: Empirically, training-based quantization methods tend to use Straight-Through Estimator to penetrate the non-differentiable quantization. However, we throw out the gradient mismatch problem in training quantized neural networks. In this paper, we introduce manifold learning to assist StraightThrough Estimator penetrating the non-differentiable quantization, which can be considered to alleviate the gradient mismatch problem. Specifically, we embed the Riemannian manifold into the Straight-Through Estimator to increase Fisher information. On the one hand, we try to embed Fisher Information Matrix for Straight-Through Estimator as the natural gradient defined by KL divergence, but encounter a complicated calculation dilemma. On the other hand, we propose a weak curvature manifold defined by Hyperbolic divergence to jump out the complicated calculation dilemma and embed weak curvature metric for Straight-Through Estimator as the weak curvature gradient.

heading: A STRAIGHT-THROUGH ESTIMATOR IN 1-BIT QUANTIZATION
text: The key of the proof is to divide the loss function into two parts based on conditional probabilities P (a > |a) and 1− P (a > |a):
E [ ∂
∂a L
] = ∂
∂a E [L]
= ∂
∂a [L(â = 1)P (a > |a) + L(â = −1)(1− P (a > |a))]
= ∂
∂a P (a > |a)[L(â = 1)− L(â = −1)],
(17)
where we divide the conditional probability P (a > |a) into two parts |a| ≤ 1 and |a| > 1 based on the distribution of the noise source:
∂ ∂a P (a > |a) = ∂ ∂a
( P (a > |a) ∣∣ |a|>1 + P (a > |a) ∣∣ |a|≤1 ) = ∂
∂a (∫ 1 −1 1 2 d ∣∣∣∣ |a|>1 + ∫ a −a 1 2 d ∣∣∣∣ |a|≤1 ) = I|a|≤1.
(18)
Further, ignoring high-order derivative of a single neuron (â = ±1) that typically has only a small impact on the total loss function, the loss function can be approximated by Taylor expansion w.r.t. â = 0
L(â = +1) = L(â = 0) + ∂L
∂â ∣∣∣∣ â=0 + 1 2 ∂2L ∂â2 ∣∣∣∣ â=0 +O ( ∂3L ∂â3 ∣∣∣∣ â=0 ) ,
L(â = −1) = L(â = 0)− ∂L ∂â ∣∣∣∣ â=0 + 1 2 ∂2L ∂â2 ∣∣∣∣ â=0 +O ( ∂3L ∂â3 ∣∣∣∣ â=0 ) .
(19)
Eq. 17 can be simply expressed as E [ ∂L
∂a
] = ( 2 ∂L
∂â ∣∣∣∣ â=0 ) I|a|≤1 = ∂L ∂â I|a|≤1.

heading: B DIVERGENCE
text: B.1 DEFINITION OF DIVERGENCE IN A MANIFOLD
D[P : Q] is called a divergence when it satisfies the following criteria:
1) D[P : Q] ≥ 0. 2) D[P : Q] = 0 when and only when P = Q. 3) When P and Q are sufficiently close, by denoting their coordinates by ξP and ξQ = ξP + dξ, the Taylor expansion of D is written as
D[ξP : ξP + dξ] = 1
2 ∑ i,j Gij(ξP )dξidξj +O(|dξ|3), (20)
and Riemannian metric Gij is positive-definite, depending on ξP .
B.2 BREGMAN DIVERGENCE
Bregman divergence DB [ξ : ξ′] is defined as the difference between a convex function ψ(ξ) and its tangent hyperplane z = ψ(ξ′)+(ξ−ξ′)∇ψ(ξ′), depending on the Taylor expansion at the point ξ′:
DB [ξ : ξ ′] = ψ(ξ)− ψ(ξ′)− (ξ − ξ′)∇ψ(ξ′). (21)
B.3 KL DIVERGENCE AND FISHER INFORMATION MATRIX
The KL divergence can be defined between p(x|θ) and p(x|θ′): DKL[p(x|θ) : p(x|θ′)] = ∫ p(x|θ) log p(x|θ)dx− ∫ p(x|θ) log p(x|θ′)dx. (22)
The first derivative is:
∇θ′DKL[p(x|θ) : p(x|θ′)]
= ∫ p(x|θ)∇θ′ log p(x|θ)dx− ∫ p(x|θ)∇θ′ log p(x|θ′)dx
= − ∫ p(x|θ)∇θ′ log p(x|θ′)dx.
(23)
The second derivative is:
∇2θ′DKL[p(x|θ) : p(x|θ′)]
= ∫ p(x|θ)∇2θ′ log p(x|θ)dx− ∫ p(x|θ)∇2θ′ log p(x|θ′)dx
= − ∫ p(x|θ)∇2θ′ log p(x|θ′)dx.
(24)
We deduce the Taylor expansion of the KL divergence at θ = θ′:
DKL[p(x|θ) : p(x|θ′)] ≈ DKL[p(x|θ) : p(x|θ)] − (∫ p(x|θ)∇θ log p(x|θ)|θ=θ′dx )>
dθ − 1 2 dθ>
(∫ p(x|θ)∇2θ log p(x|θ)|θ=θ′dx) ) dθ
= 0− (∫
p(x|θ)∇p(x|θ) p(x|θ) dx
)> dθ − 1
2 dθ>
(∫ p(x|θ)∇ [ ∇p(x|θ) p(x|θ) ] dx ) dθ
= − ( ∇ ∫ p(x|θ)dx )> dθ − 1
2 dθ>
(∫ p(x|θ)∇ 2p(x|θ)p(x|θ)−∇p(x|θ)∇p(x|θ)>
p(x|θ)p(x|θ) dx
) dθ
= −(∇1)dθ − 1 2 dθ>
( ∇2 ∫ p(x|θ)dx− ∫ p(x|θ) [ ∇p(x|θ) p(x|θ) ] [ ∇p(x|θ) p(x|θ) ]> dx ) dθ
= 1
2 dθ>Ep(x|θ)[∇ log p(x|θ)∇ log p(x|θ)>]dθ
= 1
2 dθ>F (θ)dθ.
(25) The Taylor expansion of the KL divergence at θ = θ′ is related to the Fisher Information Matrix.
B.4 HYPERBOLIC DIVERGENCE AND WEAK CURVATURE METRIC
The first derivative of the Hyperbolic divergence is:
∇θ′DH [θ′ : θ]
= ∑ i [ ∇θ′ 1 τ2 log cosh(τθ′i)−∇θ′ 1 τ2 log cosh(τθi)− 1 τ ∇θ′(θ′i − θi) tanh(τθi) ] = ∑ i ∇θ′ 1 τ2 log cosh(τθ′i)− 1 τ tanh(τθ).
(26)
The second derivative of the Hyperbolic divergence is:
∇2θ′DH [θ′ : θ] = ∑ i ∇2θ′ 1 τ2 log cosh(τθ′i). (27)
We deduce the Taylor expansion of the Hyperbolic divergence at θ = θ′:
DH [θ ′ : θ] ≈ DH [θ : θ] + (∑ i ∇θ′ 1 τ2 log cosh(τθ′i)− 1 τ tanh(τθ) )> ∣∣∣∣ θ′=θ dθ
+ 1
2 dθ> (∑ i ∇2θ′ 1 τ2 log cosh(τθ′i) )∣∣∣∣ θ′=θ dθ
= 0 + 0 + 1
2τ2 dθ>∇ [ ∇ cosh(τθ) cosh(τθ) ] dθ
= 1 2τ2 dθ> ∇2 cosh(τθ) cosh(τθ)−∇ cosh(τθ)∇ cosh(τθ)> cosh2(τθ) dθ
= 1
2τ2 dθ> ( ∇2 cosh(τθ) cosh(τθ) − τ2 [ sinh(τθ) cosh(τθ) ] [ sinh(τθ) cosh(τθ) ]>) dθ
= 1
2 ∑ i,j δij − [ tanh(τθ) tanh(τθ)> ] ij dθidθj .
(28)

heading: C THE STEEPEST DESCENT DIRECTION IN A RIEMANNIAN MANIFOLD
text: C.1 PROOF OF LEMMA 1
We would to know in which direction minimizes the loss function with the constraint of the KL divergence, so that we do the minimization:
dθ∗ = argmin dθ s.t. DKL[p(x,θ):p(x,θ+dθ)]=c L(θ + dθ)
where c is the constant. The loss function descends along the manifold with constant speed, regardless the curvature.
Now, we write the minimization in Lagrangian form. Combined with Appendix B.3, the KL divergence can be approximated by its second order Taylor expansion. Approximating L(θ+ dθ) with it first order Taylor expansion, we get:
dθ∗ = argmin dθ L(θ + dθ) + λ (DKL[p(x,θ) : p(x,θ + dθ)]− c)
≈ argmin dθ
L(θ) +∇θL(θ)>dθ + λ
2 dθ>Fijdθ − cλ.
To solve this minimization, we set its derivative w.r.t. dθ to zero:
0 = ∂
∂dθ L(θ) +∇θL(θ)>dθ +
λ 2 dθ>Fijdθ − cλ
= ∇θL(θ) + λ
2 Fijdθ
λFdθ = −2∇θL(θ)
dθ = − 2 λ F−1∇θL(θ)
where a constant factor 2/λ can be absorbed into learning rate. Up to now, we get the optimal descent direction, i.e., the opposite direction of gradient which takes into account the local curvature defined by F−1.
C.2 PROOF OF LEMMA 3
Now we can deduce the steepest descent direction while taking into account the weak curvature manifold defined by the Hyperbolic divergence. With the constraint of Hyperbolic divergence in a
constant c, we do the minimization of the loss function L(θ) in Lagrangian form:
dθ∗ = arg min s.t.DH [θ′:θ]=c L(θ + dθ)
= arg min dθ
L(θ + dθ) + λ (DH [θ ′ : θ]− c)
≈ arg min dθ
L(θ) +∇θL(θ)>dθ + λ
2 dθ>
[ δij − tanh(τθ) tanh(τθ)> ] dθ − cλ.
(29)
To solve the above minimization, we set its derivative with respect to dθ to zero and obtain the opposite direction of gradients in weak curvature manifold:
0 = ∇θL(θ) + λ
2
[ δij − tanh(τθ) tanh(τθ)> ] dθ
→ −λ 2
[ δij − tanh(τθ) tanh(τθ)> ] dθ = ∇θL(θ)
→ dθ ≈ − 2 λ
[ I + tanh(τθ) tanh(τθ)> ] ∇θL(θ).
(30)
Where a constant factor 2/λ can be absorbed into learning rate. The last line of the above formula is true because each element in [ tanh(τθ) tanh(τθ)> ] is much smaller than 1 to satisfy the weak
curvature metric. Specifically, [ I − tanh(τθ) tanh(τθ)> ] is a strictly diagonally-dominant matrix which satisfies that the inverse matrix must exist.

heading: D COMPARISONS WITH OTHER TRAINING-BASED QUANTIZATION
text: We first list the comparison results of 1-bit weight quantization in Table 5.
Next, we extend the bit width to more than 1-bit, and the quantization function of k-bit used is
Q(W ) = 1
2k−1 − 1 round
( (2k−1 − 1)W ) . (31)
We quantize the weights of all layers, and activations of all layers except the first layer, i.e, the image itself. Note that the QNN is initialized by the trained full-precision model. The comparison results are listed in Table 6.

heading: E THE ALGORITHM DESIGN
text: Algorithm 1 An algorithm for computing the gradient of the loss function L and training a quantized neural network. Note that we use the gradient measured by the Hyperbolic divergence here. Norm(∗) specifies how to normalize the activations and BackNorm(∗) specifies how to backpropagate through the normalization, which is consistent with (Courbariaux et al., 2016). Note that mat(∗) is the reverse operator of vec(∗), which turns the column vector back into a matrix. Input: A minibatch of inputs and targets (x = a0,y), θ mapped to (W1,W2, . . . ,Wl), θ̂ mapped
to ( Ŵ1, Ŵ2, . . . , Ŵl ) , a nonlinear function f , a constant factor χ and a learning rate η.
Output: The updated discrete parameters θ̂. 1: {Forward propagation} 2: for i = 1; i ≤ l; i++ do 3: Compute Ŵi = Q(Wi); 4: Compute si = Ŵiâi−1; 5: Compute āi = Norm (f si); 6: Compute âi = Q(āi); 7: end for 8: {Loss derivative} 9: Compute ∇alL = ∂L(y,x) ∂z ∣∣ z=âl
I|āl|≤1; 10: {Backward propagation} 11: for i = l; i ≥ 1; i−− do 12: Compute ∇siL = BackNorm(∇aiL f ′(si)); 13: Compute Pi = vec (tanh (Wi)) vec (tanh (Wi)) >
14: Compute ∇̃WiL = mat (( I + χ · Pitr[Pi] ) vec ( ∇siL â>i−1 )) ; 15: Compute ∇ai−1L = Ŵ>i ∇siL; 16: end for 17: {The parameters update} 18: for i = 1; i ≤ l; i++ do 19: UpdateWi =Wi − η∇̃WiL I|Wi|≤1; 20: end for

Review summaries:
This paper proposes an alternative to straight-through estimator by considering the geometry of the likelihood. The proposed method can be viewed as a natural gradient descent algorithm on the Riemannian manifold. Compared to STE, the proposed estimator seems to penalize the quantities that are far from the quantization boundary. Slightly better accuracy results are reported for training 1-bit weight and activation neural networks. 

The main motivation for this paper is the gradient mismatch problem, which emerges from using the Straight-Through Estimator (STE) in training quantized neural networks and leads to an unstable training process. To deal with that gradient mismatch problem, the authors introduce the Manifold Quantization (ManiQuant) that embeds Riemannian manifolds into the STE. Specifically, the ManiQuant associates the gradient mismatch problem with Fisher information, which can then be exploited to alleviate that problem. Considering the high cost when inverting the Fisher information, the authors present an alternative simpler method related to Hyperbolic divergence and weak curvature manifold.  

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.