Riemannian Manifold Embeddings for Straight-Through Estimator
Authors: Jun Chen
Hanwen Chen
Jiangning Zhang
yuang Liu
Tianxin Huang
Yong Liu
Venue: ICLR 2022 Submitted (2022)
Keywords: Neural network quantization
Riemannian manifold
Information geometry
Mirror descent

Abstract: Quantized Neural Networks (QNNs) aim at replacing full-precision weights $\boldsymbol{W}$ with quantized weights $\boldsymbol{\hat{W}}$, 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 $\nabla_{\boldsymbol{W}}$ w.r.t. $\boldsymbol{W}$ with gradients $\nabla_{\boldsymbol{\hat{W}}}$ w.r.t. $\boldsymbol{\hat{W}}$ where the premise is that $\boldsymbol{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.

Under review as a conference paper at ICLR 2022
RIEMANNIAN MANIFOLD EMBEDDINGS FOR ESTIMATOR
Anonymous authors Paper under double-blind review
ABSTRACT
Quantized Neural Networks (QNNs) aim at replacing full-precision weights W with quantized weights W, which make possible to deploy large models to mobile and miniaturized devices easily. However, either infinite or zero gradi- ents caused by non-differentiable quantization significantly affect the training of quantized models. order to address this problem, most training- zation methods use Straight-Throug Estimator (STE) to approximate gradients
quanti-
Vw w.r.f. W with gradients Vw w.r.t W where the premise is that W must be clipped to [-],+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 param- eter space as a metric tensor in the Riemannian manifold, we introduce the Mani- fold Quantization (ManiQuant) via revised STE to alleviate the gradient mismatch problem. The ablation studies and experimental results demonstrate that our pro- posed method has better and more stable performance with various deep neural networks on CIFAR10/100 and ImageNet datasets.
INTRODUCTION
Neural networks can handle many complex tasks due to their large number of trainable parameters and strong nonlinear capabilities (Krizhevsky al. 2012). However, the massive amount of mod- els 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 quan- tization 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 non-differentiable quantization function during training. The role of STE is to penetrate this non-differentiable quantization function and pass the gradients in backpro- pogation (Hinton 2012). e.g. DeepShift (Elhoushi al., 2019), INT8 (Zhu et al. 2020). AQE (Chen et al. (020). 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. (017)) 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 ppper-parrmeters in the training process.
Relatively, "STE" methods are wider choices for quantized models from simplicity and versatil- ity. To approximate Vw w.r.t. W and complete stable quantization training, Courbariaux et al. (Courbariaux al. 2016) binarized the neural networks by the approximated gradients using STE:
Vw = V T,, where aw OW
{
1
if W WI <
otherwise

---

Under review as a conference paper at ICLR 2022
061
M
General Gradiant
Naturel Gracient
Weak Cunatura Graciont
The general gradient is intuitively defined in Euclidean space, which means that the di- rection of the general gradient will not be affected by the curvature. On the contrary, the natural gradient and weak curvature gradient intrinsically contain the curvature manifolds, although the definitions of the manifolds constructing these two gradients are different.
Figure 1:
However, will inevitably bring the gradient mismatch problem raised by just simply application of STE. To overcome this challenge, Zhou al. (Zhou al. 2016) proposed transform W W:
tanh (W)
max{ Lanh W)
and then quantize it using a quantization function Q(-). During backpropogation, the gradients Vw w.r.t. W can be further computed using the chain rule:
1- tanh (W)
V V V Q(W) max( lanh W)])
Based on this work, Chen et al. (Chen et al. neural networks LSTM, which replaces this gradients via meta quantizer Ms parameterized by across layers:
2019) proposed to learn Vw by fully-comnected
OW
Vw Ma (VQww) Q(w) Ow
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 we treat the parameter space in 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, define a novel Hyperbolic divergence by convex function with geometric structure. With the con- straint 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 STE as the weak curvature gradient to approximate w without extra parameters and training complexity.
RELATED WORK
2.1 GENERAL GRADIENT
For neural network with layers, the full-precision weight matrix of each layer is just marked as W, By defining the operation vec(.) that vectorizes matrices by stacking their columns together, the total parameter vector of the neural network can be denoted
vec (Wi) vec (2)) ...., vec (Wi) The parameter vector is given as a column vector.
2

---

Under review as a conference paper at ICLR 2022
Let 0 R be a parameter space which loss function L associated with weight is well-defined. is relatively easy to express the general gradient in training:
OT.
Vol
00
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 considered, the Euclidean divergence between two sufficiently closed points O and f' is actually defined by default:
DE 07) - 1 - -0.
which is half of the square of the Euclidean distance identified by the Fuclidean metric diz, so that
ds² 2 + d0)
Σ (d0.)²
Σ d0.
(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 Gzy(O) ()) for the parameter space to form a Riemannian manifold (M, Gig) (Amari 1998
2016). In that case, the steepest descent direction also depends on the quadratic form introduced by a small incremental vector d0 that connects O and 0 + d0, whose form is given by
ds² = Σ Giy(O)d0,dt) i.
(4)
where d8, is the component of d0. Under the constraint ds? the steepest descent direction is toward the optimization goal of L(O + d0). Intuitively, it measures the Kullback-Leibler (KL) divergence between two distributions p()[0) and p(æ - d0) of network model, which equivalent to 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):
d (p(=()) p( + d0)] xs -d0 E
Ep V logp(æ 0) logp(æ 0) Jdd.. (5)
We see that FIM is equal to the negative expected Hessian log likelihood. Furthermore, Amari deduced that the Riemannian metric given by the FIM (Amari 1998). By the Lagrangian form, we have the natural gradient (see Appendix C.I
aL (0) 00 where F(0) Ep( (V ( logp(z=0))
(6)
which is the steepest descent method in a Riemannian space and F(O) is the FIM with a parameter vector (the natural gradient also known the Riemannian gradient). Empirically, the immediate application of FIM is as drop-in replacement of Hessian in second order optimization algorithm. Note that V L will return to VoL when Gij is equal to the Euclidean metric d,j, i.c. the identity matrix I.
MANIFOLD QUANTIZATION
Considering a quantized model of each layer, we notate the quantized weight matrix using W for differentiating full-precision weight matrix W. Note that we define the quantized parameter vector as Ó, which similar to the full-precision parameter vector During training a QNN, the quantization function Q(-) is one-to-one mapping from full-precision values to quantized values, which can be expressed as θ Q()).
3

---

Under review as a conference paper at ICLR 2022
By involving the practice train 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 input pattern -;-1 to output a can still be imitated in the same way as full-precision neural networks:
8; W,
à. = Q o o..);
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 quantization function in a quantized model, whether for weight matrix W, activation vector a,
3.1 FISHER INFORMATION MATRIX EMBEDDING FOR STE
The concept of the natural gradient is closely related to the FIM and KL divergence (see Ap- pendix B.3). Since the KI. divergence intrinsic, the natural gradient is also intrinsically invariant under the parameter transformation. By viewing the FIM an / by / block matrix where / denotes the number of layers in a neural network, the natural gradient formula can introduced to alleviate the gradient mismatch problem when training QNN and updating its parameters.
Lemma In the Riemannian manifold defined by KL. divergence, Fisher Information Matrix em- bedding for Straight- Thhhoughh Estimator is written by:
L
STE
V0VVooo
dlogp(æ0 d d ((
E
V o 1ols1
db
dه
(8)
TTL) ai
0L aē
E
vec
vec
V L 110/5-
Furthermore, F F0) can be expressed as
E
vec
aL ow,
vec
ow,
vec
oW
ØL
vec
BW
E
vec
OL aw.
vec
aL aw.
E
vCc
aL aaw.
vec
aL oow
where vec
OL ow,
can be represented by the gradient error
aL da."
Proof. The proofs are the combination between STE and Appendix
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 the learning procedure:
vec
aL
= a. =
2
(OL.
ØL Da.
o f((..)
STE a,T
X
aL oā,
(f((s)))aa)s))
(10)
aW.
where @ denotes the Kronecker product [ Relying the KL divergence, indicate the parameter space as the Riemannian manifold rather than the Euclidean space in the quantization procedure.
For neural networks with scale of one million more parameters, the time complexity of inverting the FIM, a component of natural gradients, is ((n3) (Povey et al. 2014). Previously, there were some works calculate the natural gradient efficiently, e.g. Roux et al. (Roux al. 2008) de- composed the FIM into multiple diagonal blocks, where each diagonal block is approximated by low-rank matrix. Bastian et al. (Bastian al., 201T) also used the idea of diagonal blocks by con- structing diagonal block that corresponds to weight matrix. Martens et al. (Martens & Grosse)
Since the size of a_ and the size of AL da,
may he different in each layer, Kronecker product necessary.
4

---

Under review as a conference paper at ICLR 2022
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.
2013)
3.2 WEAK CURVATURE METRIC EMBEDDING FOR STE
3.2.1
INTRODUCTION OF WEAK CURVATURE MANIFOLD
We can use the Euclidean coordinates to calculate the natural gradient on the manifold because of the local homeomorphism of manifold (Wald 2010). Note that the homeomorphic mapping de satisfies U M Φu (U) E R² For any point E U, we can define du (z) as the Euclidean coordinate absolutely.
The FIM embedding for STE poses a significant challenge for computing. Intuitively, the com- puting (0) 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 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 nearly flat, where the Riemannian metric is approximation of the Fuclidean metric. In practice, develop linearized Riemannian metric from the Euclidean metric dij, which is systematically defined in general relativity (Wald 2010)..
Definition (Weak Curvature Metric) Let a linearized Riemannian metric deviates from the Eu- ()) is much smaller than 1 in global Euclidean coordinate clidean metric. The deviated metric (i)(0) system of dig.
(0) dig Eij (0), where leyl <<<
1.
(11)
is an adequate definition of weakness" in this context and ensures Giş to be positive-defnnte metric.
3.2.2 HYPERBOLIC DIVERGENCE
In practice, we can determine unique geodesic through exponential map exPo, (T0) that maps 70 back to M where is small constant, by defining -0 E To as the tangent vector. Note that the definition of exponential map developed by (Wald 2010 Helgason 2001).
Definition 2 (Exponential Map) Let M be Riemannian manifold, for the tangent vector v E T.M a point x M where T.M the tangent space, there is unique geodesic (to(t) locally that satisfies (0) = I and % (0) = v. The exponential map exp, TTMM + M corresponding to Too is defined as exp,(v) = 7n (1). When constraining exp, to a neighbourhood U, this mapping one-to-one.
Empirically, the exponential map exPo, (70) is to map tangent vector 7 in the tangent bundle to the point where the are length from point 0; is equal to |r0] on geodesic with the initial condition (0,, 70). In order to make cxp(r@) and cxp(-70) have the same effect the training process, we symmetrize them and derive convex function
(0) Σ 1
log (xx((T0:) +
Σ log(cosh(TO:)).
(12)
Geometrically, the convex function with geometric structure is introduced into the Bregman di- vergence (see Appendix B.2) (Bregman 1967). and we obtain novel Hyperbolic divergence that satisfies the criteria divergence (see Appendix B.T.
5

---

Under review as a conference paper at ICLR 2022
Definition 3 (Hyperbolic Divergence) For convex function & defined by Eq. 12 the Hyperbolic divergence between 8' and is
cosh(r00)
DHlo's 0 Σ
log cosh(T7;))
(0, 0,) tanh((00:)
(13)
3.2.3 THE GRADIENT FLOW IN WEAK CURVATURE MANIFOLD
Lemma 2 Let d0 - the square of an infinitesimal distance defined by Hyperbolic divergence is ds² = 2D= 0) =2 d [ōsg - -aanh(70) Lanh(r0) do.
(14)
can be consistent with the weak cur- Furthermore, the defined metric [dis tanh(r0) (70) (a)n(a0))))))
vature metric.
Proof. The proof' can be found Appendix B Empirically, we can give a small constant T to satisfy the weak curvature metric based on Definition ! Comparing Eq Fand Eq. I1 can deduce
the deviated metric Ejj = [tanh(00) Sanh(T0) "Iyg"
Lemma In the Riemannian manifold defined by Hyperbolic divergence, a weak curvature metric embedding for Straight- Throughh Estimator is approximated by:
STE R
VOl.
+ tanh(r0)) tanh(r0) V T
(15)
Proof. The proofs can be found in Appendix
3.3 HIERARCHICAL EMBEDDINGS OF THE WEAK CURVATURE METRIC
When embedding weak curvature manifold, there is unnatural gap between the calculation of weak curvature metric and the layer-by-layer gradient update in fack-propagation Specifically, the metric not a block diagonal matrix (each block corresponds to layer of the neural network).
For the process of the quantization, we have the hierarchical embeddings for STE through decou- pling operation:
V vec(WW)
I -². vec(tanh W:) veo(tanh( (W)))
voo voo( I veo(W)) s1
(rc(Lanh(W.) veo(Lanh(W.)) (
Veo(W:) II veo(W.))s1
tr (vec(tanh(W]) po((tanh(Ws)))
(
tank(we)) tanhi tr(P:
+ X or PP]
tanh (w:a) (ws1)
+
tauh (wva)
vec((W)) vec(W.)ls
(16)
In the first line of the above formula, we can put 7? outside because the elements of the deviated metric is very small based Lemma 2 In the second line, we use normalized formula to uniformly represent 7/: 7" = x/ tr /ve W:)) vec(tanh(W)) W:)) where X is a constant factor and the matrix [vec(Lanh( W:)) vec(tanh, W] is re-expressed by P,- 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
EXPERIMENT
In this section, we implement experiments demonstrate the effectiveness of our proposed methods on benchmark datasets that are CIFAR10/109 and ImageNet mainly here. Intuitively, experiments LIFAR10/100 and ImageNet ablation studies to validate the advantages of weak curvature manifold embeddings. Comparisons with other training- based quantization on quantitative indicates will be carried out CIFAR10 and ImageNet. All experiments conducted with PyTorch.

---

Under review as a conference paper at ICLR 2022
Table 1: The experimental results CIFARIO with ResNe120/32/44.
Network
Forward
Backward
Test Ace (%)
Original Acc (%)
Dorefa MultiFCG
88.28:00.81
ResNet20
{ [-1,+1}
88.94=0.46
91.50
WCG
89.78++0.33
Dorefa
90,23:0.63
ResNet32
{-1++1}
MultiFCG
89.63=0.38
92.13
WCG
90.52++0.27
Dorefa
90.71
ResNet44
{-1,+1}
MultiFCG
90.54::.21
93.56
WCG
91.38+10.11
4.1 EXPERIMENT SETUP
We use weight decay of le-4, 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 X as the reciprocal of output channel layer. All experiments conducted for 5 times, and the statistics of the last 10/5 epochs' test accuracy are reported fair comparison.
Table 2: The experimental results CIFAR100 with ResNet56/1 10.
Network
Forward
Backward
Test Acc (%)
Original Ace (%)
Dorefa
66.71=2.32
MultiFGG
66.58=1037
ResNet56
1-1,11}
71.22
FCGrad
66.56=0035
WCG
68.85
Dorefa
18.05
MultiFGG
68.27-0.14
ResNetl 10
[-1,+1}
72.54
FCGrad
6..74-
WCG
0.26
4.2 ABLATION STUDIES
In order illustrate the superiority of the weak curvature gradient (WCG), train QNNs with our gradient from scratch compared to other gradients, i.e. Dorefa (Zhou et al. 2016). MultiFGG CCen al.] 2019) and FCGrad (Chen 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 Xavier (Glorot Bengio 2010). Specifically, we experiment ResNer20/32/44 models (He al.) 2016) CIFARIO, ResNet56/110 models (He al. 2016 CIFAR 100 and ResNet18 8 model (He ct al. 2016) on ImageNet. The accuracy of full- mill-precision baseline reported by (Chen et al. 2019).
The results of ablation studies CIFARIO, CIFARI00 and ImageNet shown in Table Ta- ble and Table 3 respectively. Various quantized models applying weak curvature gradients show significant improvement in challenging classification tasks. For faimess, quantization process and optimizations in experiments where the only difference is the gradients of the backward process.
the same forward

---

Under review as a conference paper at ICLR 2022
Table 3: The experimental results ImageNet with ResNet18.
Network
Forward
Backward
Test Top1/TopS ( (%)
Original Top1/Top5 (%)
Dorefa
58.3442.0788.44771.555
MultiFCG
59,47+00.028244100001
ResNet18
[-1,+1}
69.76/89.08
FCGrad
59,8344..36652.71-02223
WCG
61.02-0.3333344000116
100
Train Dorefa
Train Dorefa
85
Test acc: Dorefa
Test acc: Dorefa
WCG
acc: WCG
Test acc: WCG
Test WCG
75
Accurec (94)
Accuracy (%)
80
30
50
100
150
200
50
50
100
150
200
opech
apoch
(a) ResNet56
(b) ResNetl
Figure 2: Training and curves of ResNet56/1 10 on CIFAR 100 compared between Dorefa and WCG.
4.3 CONVERGENCE AND STABILITY ANALYSIS
In this experiment, we compare the convergence and stability of Dorefa and WCG during the training process. As Figure P 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, conjecture that WCG has increased the Fisher information the original quantized model compared to other methods, although WCG does not use additional neural network information for learning like MultiFCG FCGrad. Moreover, this is where the advantage manifold learning lies.
4.4 COMPARISONS WITH OTHER TRAINING-BANE BASED QUANTIZATION
In this experiment, our ManiQuant is initialized with the trained full pell-rrecision model, and continues to WCG learn the quantized model. ManiQuant focuses the weak curvature manifold embeddings for STE, without any cxtra 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 found in Appendix [
4.5 TRAINING TIME ANALYSIS
Our hardware environment conducted with an Intel(R) Xeon(R) Silver 4214 CPU(2.2GGH2), GeForce GTX 1080Ti GPU. We test the training time per iteration for MetaQuant using Mul- tiFGG, ManiQuant using WCG, and Dorefa with ResNet20 in CIFARIO. For finishing one iteration of training, ManiQuant costs 18.44s and MetaQuant costs 25.34s while Dorefa uses 17.81s.
8

---

Under review as a conference paper at ICLR 2022
Table 4: The experimental results CIFARIO with ResNer18 and VGGI6.
Network
W/A
Method
Test Acc (%)
Original Acc ( (%)
RBNN (inn al.
2020
92.20
MD
Ajanthan et a..
2021
91.28
ResNet18
1/1
94.84
PGD
Ajanthan et al. 2019
92.60
ManiQuant (WCG)
92.83
RBNN (Lin et al. 2020)
91.30
MD (Ajanthan et al..
2021
90.47
VGG16
1/1
93.33
PGD (Ajanthan et 2019)
88.48
ManiQuant (WCG)
91.77
CONCLUSION
Empirically, training-based quantization methods tend use Thriigh Estimator to pene- trate the non-differentiable quantization. However, we throw out the gradient mismatch problem in training quantized neural networks. In this paper, introduce manifold learning assist Straight- Through Estimator penetrating the non-differentiable quantization, which can be considered to al- leviate 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 complicated calculation dilemma. On the other hand, we propose weak curvature manifold defined by Hyperbolic divergence to jump out the complicated calculation dilemma and embed weak curvature metric for Straigh--Through Estimator as the weak curvature gradient.
REFERENCES
Thalaiyasingam Ajanthan, Puneet K Dokania, Richard Hartley, and Philip HS Torr. Proximal mean- field for neural network quantization. In Proceedings of the IEEE/CVP International Conference Computer Vision, pp. 4871 2019.
Thalaiyasingam Ajanthan, Kartik Gupta, Philip Torr, Richad Hartley, and Puneet Dokania. Mirror descent view for neural network quantization. In International Conference on Artificial Intelli- gence and Statistics, pp. 2809-2817. PMLR, 2021.
Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251- 276, 1998.
Shun-ichi Amari. Information geometry and its applications, volume 194. Springer, 2016.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint
arX.v:66070064500 2016.
Yu Bai, Yu-Xiang Wang, and Edo Liberty. Proxquant: Quantized neural networks via proximal operators. arXiv preprint arXiv:1810.00840 001008080,00 2018.
Michael R Bastian, Jacob H Gunther, and Todd K Moon. A simplified natural gradient learning algorithm. Advances in Artificial Neural Systems, 2011, 2011.
Lev M Bregman. The relaxation method of finding the common point convex sets and its applica- tion to the solution of problems in convex programming. USSR computational mathematics and mathematical physics, 7(3):200-217, 1967.
Sébastien Bubeck. Convex optimization: Algorithms and complexity. Foundations and Trends( in Machine Learming, 8(3-4)2231-355 2015.
J. Chen, Y. Liu, H. Zhang, S. Hou. and J. Yang. Propagating asymptotic-estinate gradients for low bitwidth quantized neural networks. IEEE Journal of Selected Topics in Signal Processing. 14(4): 848-859, 2020.

---

Under review as a conference paper at ICLR 2022
Shangyu Chen, Wenya Wang, and Sinno Jialin Pan. Metaquant: Learning to quantize by learning to penetrate non-differentiable quantization. In Advances in Neural Information Processing Systems, volume 32, pp. 3916-3926. Curran Associates, Inc., 2019.
Matthieu Courbariaux, Itay Hubara, Daniel Soudry, Ran EI- El-Yaniv, and Yoshua Bengio. Binarized neural networks: Training deep neural networks with weights and activations constrained to+ or-1. arXiv preprint ar2iv:160..028300 2016.
Harald Cramér. Mathematical Methods of Statistics (PMS-9), Volume 9. Princeton university press, 2016.
Mostafa Elhoushi, Zihao Chen, Farhan Shafiq, Ye Henry Tian, and Joey Yiwei Li. Deepshift: To- wards multiplication-les neural networks. arXiv preprint arXiv; 1905. 13298, 2019.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and stafistics, pp. 249-256, 2010.
Philipp Gysel, Jon Pimentel, Mohammad Motamedi, and Soheil Ghiasi. Ristretto: A framework for empirical study of resource-efficiess inference in convolutional neural networks. IEEE transac- fions on neural networks and learning systems, 29(1 1):5744-5789) 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Sigurdur Helgason. Differential geomelry and symmetric spaces, volume 341. American Mathe- matical Soc., 2001.
G Hinton. Neural networks for machine learning. coursera,(video lectures], 2012.
Andrew Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam. Mobilenets: Efficient convolutional neural networks for mobile vision applications. arXiv preprint arXiv: / 7704044661 2017.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton. Imagenet classification with deep convo- lutional neural networks. In Advances in neural information processing systems, pp. 1097-1105, 2012.
Cong Leng, Zesheng Dou, Hao Li, Shenghuo Zhu, and Rong Jin. Extremely low bit neural net- work: Squeeze the last bit out with admm. In Proceedings of the AAAI Conference Artificial Intelligence, volume 32, 2018.
Fengfu Li, Bo Zhang, and Bin Liu. Ternary weight networks. arXiv preprint arXiv: 10050:1605.04711, 2016.
Minghao Lin, Rongrong Ji, Zihan Xu, Baochang Zhang, Yan Wang, Yongjian Wu, Feiyue Huang, and Chia-War Lin. Rotated binary neural network. Advances in Neural Information Processing Systems, 33, 2020.
Christos Louizos, Matthias Reisser, Tijmen Blankevoort, Efstratios Gavves, and Max Welling. Re- laxed quantization for discretized neural networks. In International Conference on Learning Representations, 2019. URL
James Martens and Roger Grosse. Optimizing neural networks with kronecker-factoree approximate curvature. In International conference on machine learning, pp. 2408-2417, 2015.
Daniel Povey, Xiaohui Zhang, and Sanjeev Khudanpur. Parallel training of dnns with natural gradi- ent and parameter averaging. arXiv preprint 1rXiv:1410.7455. 2014.
C Radhakrishna Rao. Information and the accuracy attainable in the estimation of statistical param- eters. In Breakthroughs in statistics, pp. 235-247. Springer, 1992.
Nicolas L Roux, Pierre-Antoine Manzagol, and Yoshua Bengio. Topmoumoute online natural gra- dient algorithm. In Advances in neural information processing systems, pp. 849-856, 2008.
10

---

Under review as a conference paper at ICLR 2022
Oran Shayer, Dan Levi, and Ethan Fetaya. Learning discrete weights using the local repa- rameterization trick. In International Conference on Learning Representations 2018. URL
httts::/ net /tt.:/ooommnindyaaaGCCa
Robert M Wald. General relativity. University of Chicago press, 2010.
Dongqing Zhang, Jiaolong Yang, Dongqiangzi Ye, and Gang Hua. Lq-nets: Learned quantization for highly accurate and compact deep neural networks. In Proceedings of the European conference on computer vision (ECCV), pp. 365-382, 2018.
Aojun Zhou, Anbang Yao, Yiwen Guo, Lin Xu, and Yurong Chen. Incremental network quantiza- tion: Towards lossless cnns with low-precision weights. arXiv preprint 1702v 2017.
Shuchang Zhou, Yuxin Wu, Zekun Ni, Xinyu Zhou, He Wen, and Yuheng Zou. Dorefa-net: Train- ing low bitwidth convolutional neural networks with low bitwidth gradients. arXiv preprint arXiv:1606. 060.666666666. 2016.
Chenzhuo Zhu, Song Han, Huizi Mao, and William Dally. Trained ternary quantization. arXiv preprint arXiv:
2016.
Feng Zhu, Ruihao Gong, Fengwei Yu, Xianglong Liu, Yanfei Wang, Zhelong Li, Xiuqi Yang, and Junjie Yan. Towards unified int8 training for convolutional neural network. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 1969-1973 2020.

---

Under review as a conference paper at ICLR 2022
A STRAIGHI-THROUGH ESTIMATOR IN 1-BIT QUANTIZATION
The key of the proof is to divide the loss function into two parts based on conditional probabilities P(a > ela) and P(a ela):
0
i R. [T]
R2
T,
0a
da
0
[L = I)P(a ela) + L(ā = -1)(1 P(a > e|a))]
(17)
aa
0
aa P(a (La)[T((a I.ā -1)]
where we divide the conditional probability P(a > ela) into two parts |a| S 1 and |a] the distribution of the noise source:
based on
0 P(a > cla)
0
(P(a (P(a> P(a>>|a)|1a/>> + +P(a > cla)|
da
da
0
p (/"2dt de
+ $
(18)
de
= -laas--
da
2
|
|als1
Further, ignoring high-order derivative of single neuron (á = =1) that typically has only a small impact on the total loss function, the loss function can be approximated by Taylor expansion w.rt.
àL
L
8 L
L(a = +1) L(a = -) +
+ +0
dû a=0
2 dh
0a
a=0
á=0
(19)
aL
L
0
L(a -1) L(
+0
Oā
h=0
Eq. 17 can be simply expressed as
TL
aL
ØL
E,
=
2
1
10|51-
da
dā
6=0
B DIVERGENCE
B.1 DEFINITION OF DIVERGENCE IN MANIFOLD
D[P: Q) is called divergence when it satisfies the following criteria:
1) D[P: Q) 2 2) D[P: Q = when and only when Q. 3) When and Q are sufficiently close, by denoting their coordinates by &p and Eq = Ex + cC,, the Taylor expansion of D is written as
DEE Ep del
Σ
(Ep)dddE + O(Ide))) i,j
(20)
and Riemannian metric Gij is positive-definite depending on &P.
B.2 BREGMAN DIVERGENCE
Bregman divergence DBLE El] is defined as the difference between convex function (i(E) and its
( V (EE),, depending on the Taylor expansion at the point g": tangent hyperplane (b())) + (E -
DElE 57] = (E) - ((E') (E - (V'))((('')
(21)
B.3 KL DIVERGENCE AND FISHER INFORMATION MATRIX
The KL divergence can be defined between p( 0) and p( 0'):
Dkllp(æ(0) p(x(0')) f p(z(0) p( 0 ddx f p(x0(0) log
(22)
12

---

Under review as a conference paper at ICLR 2022
The first derivative is:
Vo DKLpp(x()) p(=]0']
/ p(x: logpla))dd ((zx0)\oo logg(a>)0'ddl
(23)
V 1ogp(z)00)daz
The second derivative is:
Vg DyDDlll(a ( p(x
/ 0) log logp(=))0))))
p( 0) logp(z00))
(24)
log
=
We deduce the Taylor expansion of the KI. divergence at A = 0':
DKL|p(a:00 0 p(x(0')) DKL (0) [z)
- (I 2o010)Vo o
-00
2 do
( or do
d0
0- (f p(z)0)
Vp(z dæ
(frtelory
0)
d0 1 do
dx d0
1(x(0)
p(x)0)
d0 マ do (ʃ (fD(alo)
V
V'p(z10)p( (0000)p(()))) (x(0) VV(2> dx d0
plz
p(<]0))(<(1))
Vp(z(0)
Vp(z 0)
(1)dd-
3 -0 22 p(x(0)dx - p(x10)
dæ d0
p(x]0)
p(x]0)
1
d Ep(a)o)) / log Vogg Td0
1d
2
de F(0)d0.
2
(25)
The Taylor expansion of the KL divergence 0 0' is related to the Fisher Information Matrix.
B.4 HYPERBOLIC DIVERGENCE AND WEAK CURVATURE METRIC
The first derivative of the Hyperbolic divergence is:
V DH 00' 0
1
Σ V -
log cosh(78f) - V 72
log cosh(r0,) Vo. (6% - 0,) tanh((T:))
(26)
1
Σ
log cosh (TO,) tanh(r0).
The second derivative of the Hyperbolic divergence is:
Vg 0' =
Σ
V/ -
log cosh(r8.))
(27)
13

---

Under review as a conference paper at ICLR 2022
We deduce the Taylor expansion of the Hyperbolic divergence 6':
O' :0' so DH10:01
Σ
Vo
log cosh(T0f)
1 tann((0))
d0
72
0'=0
1
+ doT
V 72
log cosh(T(0))
d0
0'=0
V cosh( (0)
+0+ 0 0+
d0
d0
2T
cosh(
(28)
2 cosh 0 cosh(r0)
V cosh(r0) cosh(r0))
d0
d0
27
cosh (o6)
4 cOsS
sinl((rO)
sinh((r))
T
d0
2T
cosh(TT)
cosh(rt)
coos (8)
Σ
[lanh(r0) Lanh(rO) "] d0,d0,.
2
C THE STEEPEST DESCENT DIRECTION IN A RIEMANNIAN MANIFOLD
C.I PROOF OF LEMMA T
We would to know in which direction minimizes the loss function with the constraint of the KI. divergence, so that we do the minimization:
d
arg min
L(O + d0)
do s.1. DE 0, (o
where e is the constant. The loss function descends along the manifold with constant speed, regard- less the curvature.
Now, we write the minimization in Lagrangian form. Combined with Appendix B.3 the KL diver-
gence can approximated by its second order Taylor expansion. Approximating L(e + d0) with first order Taylor expansion, we get:
d0 = arg min L(e + d0) A (DDrLlp(a,, 0) p(x., 0 d0)] c)
do
>
2 arg min L( + ()) do dd Fi cx.
2
To solve this minimization, we set its derivative w.r.t. d0 to zero:
0
X
0
0d0 [(0) + (oI do d0 Fi, cd
X
VoL(O) + Fi,do
2
XFd0 2V)
2 d0 X F-'V
where constant factor 2/A can be absorbed into learning rate. Up 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
14

---

Under review as a conference paper at ICLR 2022
constant e, we do the minimization of the loss function L(0) in Lagrangian form:
d0
arg min I.(0 d0)
s.t.Dz [|
arg min 1(0 X (D 0 - c)
(29)
de
arg min L( VoL(O) d0 + do [ōis - tanh(r0) tanh(r0) [ d0 ch.
de
To solve the above minimization, we set its derivative with respect to d0 to zero and obtain the opposite direction of gradients in weak curvature manifold:
0 VoL(O) + [dg tanh(r0) tanh(r0) d0
A [ōg - -tanh (Tnh( tanh(r0) | d0 VoT.(0)
(30)
2
2 [ Lanh(rO)) (anh(T7)) VaL(O).
- d0 ~
Where constant factor 2/A be absorbed into learning rate. The last line of the above formula true because each element in [tanh(r0) tanh(r0) [] is much smaller than 1 to satisfy the weak curvature metric. Specifically, [r Lanh(r0) lanh(r7) T] is strictly diagonally-dominant matrix which satisfies that the inverse matrix must exist.
D COMPARISONS WITH OTHER TRAINING-BASED QUANTIZATION
We first list the comparison results of 1-bit weight quantization Table 5
Table 5: The experimental results of ManiQuant compared to other quantization.
Dataset
Network
Method
Test Acc (%)
Original Ace (%)
ProxQuant (Bai et al. 2018)
90.21
MetaQuant CCenn et a
2019
90.80
CIFARIO
ResNet20
91.50
LQ-Nets (Zhang et al. 2018)
90.10
ManiQuant (WCG)
90.88
LR Net (Shayer et al. 2018)
59.90 $9.90/82.30
MetaQuant CChnn al. 2019)
63.44/84.77
ImageNet
ResNet18
69.76/89.08
ADMM (Leng et al.
2018
64.80/86.20
ManiQuant (WCG)
65.18/86.26
Next, we extend the bit width to more than 1-bit, and the quantization function of k-bit used is
Q(W) = 2 - round -12--1
(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 initialized by the trained full- precisoon model. The comparison results listed in T
15

---

Under review as a conference paper at ICLR 2022
Table 6: The classification accuracy results ImageNet and comparison with other training-based
quantizations, with AlexNet (Krizhevsky et al. 2012). ResNet-50 and MobileNet (Howard et al.
2017). Note that the accuracy of full-precision baseline reported by (Elhoushi 2
Top-1 (%)
Top-5 (%)
Method
W
A
Accuracy
Gap
Accuracy
Gap
AlexNet (Original)
32
32
56.52
79.07
ManiQuant (WCG)
6
32
56.39
-0.13
78.78
0.29
DeepShift-Q (Elhoushi et al.
2019)
6
32
54.97
-1.55
78.26
0.81
ResNet-50 (Original)
32
32
76.13
92.86
ManiQuant (WCG)
8
8
76.10
0.03
92.88
+0.02
INT8 (Zhu et al. 2020)
8
8
75.87
0.26
MobileNet (Original)
32
32
70.61
89.47
ManiQuant (WCG)
5
61.32
-9.29
84.08
5,39
SR+DR (Gysel et al.
2018
5
5
59.39
-1.22
82.35
-7.12
ST Lounzos et al.
2019)
5
5
56.85
133.76
80.35
-9.12
ManiQuant (WCG)
8
8
70.86
+0.25
89.60
+0.13
RQ (Louizos et al. 2019)
8
8
70.43
-0.18
89.42
-0.05
E THE ALGORITHM DESIGN
Algorithm An algorithm for computing the gradient of the loss function L and training a quan- tized neural network. Note that Norm(*) specifies how to normalize the activations and Back specifies how to back- propagate 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 matrix.
the gradient measured by the Hyperbolic divergence here.
Input: A minibatch of inputs and targets (ær ao y). O mapped to (W1, W Wi), mapped nonlinear function f. constant factor x and a learning rate Output: The updated discrete parameters Ô.
to (W,, w.,.,, W.)
1: (Forward propagation
2: for 1; I; + do
3: Compute W. Q(W);;
4: Compute 8; W;ā;-i;
5: Compute a. Norm (f o 0:);
6: Compute à, = Q(ā:):
7: end for
8: (Loss derivative]
9: Compute Va, (L(,
x=d; | 51i
(Backward propagation)
11: for = 1; 2 1; do
12:
Compute V L = BackNorm V L
13:
Compute P. = vec (tanh (W,) vee (tanh (W,))*
14: Compute V w. L = mat ++ x i.P.. vec (Ve, L a a.T))
15:
Compute V Va.,LL W. V L;
16: end for
17: (The parameters update}
18: for = 1; < i + + do
19: Update W, = W. 7 Iw; s/:
20: end for
16


Reviews:
[{"review_id": "wwZSA_W1fiH", "reviewer": "Reviewer_WCKm", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n- Investigating principled gradient estimators for the discrete quantization is an important problem to study. \n- The idea of utilizing manifold geometry makes sense.\n\nWeaknesses\n- the actual used manifold in Sec. 3.3 does not seem to use the Fisher information. Why is the weak curvature metric better than the Euclidean distance?\n- It looks like that the proposed estimator in Sec. 3.3 only applies to 1-bit quantization. How does the proposed approach generalize to multi-bit quantization?\n- Some convergence theorems are required to justify that learning with the proposed estimator can indeed converge to better models.\n- Is the vanilla straight-through estimator compared in the experiments?", "summary_of_the_paper": "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. ", "summary_of_the_review": "The idea of the paper makes sense, but the theoretical and experimental results cannot support the usefulness of the proposed method well.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "pOTTRmcklUY", "reviewer": "Reviewer_HJGf", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I find the paper clearly written. As a non-expert, I think that the questions that the gradient mismatch problem that the author aims to address is significant and the experiment results convincingly demonstrate the superior performance of the proposed gradient estimation method.\n\nMy main suggestion for the authors is to provide more context for the benefits of training QNN and position their contribution in this context. For instance, is the goal of training QNN to gain training speed, reduce GPU memory, or reduce the size of the network for storage purposes? How much can the proposed method improve upon these metrics (other than just accuracy), compared to full-precision networks? Indeed, the proposed method makes use of the natural gradients, which, to my knowledge, are much more computationally expensive to evaluate than first-order optimization methods such as the plain SGD. Can the proposed method improve training speed or memory usage upon a full-precision network trained with vanilla SGD? It will be great to have a table (similar to Table 2 and Table 3) to list the acceleration or storage gain of different methods and compare them with vanilla training of full-precision networks. There is a small subsection (subsection 4.5) that discusses the training time. This subsection, in my opinion, can be expanded, for instance, to include training time required for a full-precision network using SGD. \n\nMinor: In the first paragraph of Section 3 on page 3, the authors mentioned that the quantization function is a **one-to-one** mapping from full-precision values to quantized values. Why is the quantization function a one-on-one map? I supposed that a single quantized value may correspond to multiple pre-quantized, full-precision values. ", "summary_of_the_paper": "This paper proposed a new gradient estimation method for training quantized neural networks (QNN). It alleviates the gradient mismatch problem that occurred in previous quantization methods that use the Straight-Through Estimator.\n", "summary_of_the_review": "As a non-expert, I find that the method proposed in the paper is novel and empirically well-tested. However, the proposed approach for training QNNs seems to be computationally expensive on its own, and I am not sure if the required computational budget defeats the purpose of training QNN -- it will be great if the authors can help me understand in this regard.  ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "clSS65r2Z_A", "reviewer": "Reviewer_gTNq", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n1. The motivation of this paper is straight-forward and interesting: Basically, it first pointed out that quantization error is introduced by variance of STE, which can be upper bounded by the FIM. Then it related FIM with STE and kl-divergence.\n\nQuestions:\n1. Author mentioned that FIM should be increased to alleviate the gradient mismatch problem. But the natural gradient introduced in Eq.10 does not seems to contribute to this goal. How the increase of FIM is related to the training ojective of model quantization?\n2. How does optimizer interact with the proposed gradient? In experiments, author used SGD with momentum (SGD-M). Does that mean the natural gradient is firstly attained by Eq.18, then SGD-M is applied to the gradient to achieve gradient used in parameters update ? How if Adam is applied in the experiments since Adamm affects much to the initial gradient.\n3. Besides, is SGD-M still applied in ImageNet experiments? Accuracy in table 3 is too high for SGD-M in ResNet18.", "summary_of_the_paper": "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. ", "summary_of_the_review": "It is interesting to incorporate manifold learning in QAT, but more concerete analysis (experiments) in deep learning field should be conducted.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "YCpgxXuUlhZ", "reviewer": "Reviewer_Goby", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "In Theorem 1 (more specifically in the corresponding proofs in Appendix A), it seems Eq (8) is only $\\textbf{approximately}$ unbiased, because of the Taylor expansion in Eq (21) in Appendix A. It's not rigorous to state the estimator in Eq (8) is unbiased. At least, this should be discussed carefully.\n\nHonestly, I am confused about the underlying logic in Section 3.1. Of course, the gradient mismatch problem is caused by using the STE. But why does the root of that problem ONLY come from the STE variance (or its Cramer-Rao Lower Bound)?\n\nIn Eq (10), from the second row to the third row, you cannot simply replace the log-likelihood with the loss function L in general. Please elaborate on it.\n\nIn Figure 2, it seems the demonstrated results are from a single run. Please add the error bar that is shown in the tables.\n\nTheoretically, why using the weak curvature gradient could deliver better performance?", "summary_of_the_paper": "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.\n\n", "summary_of_the_review": "The underlying logic is not clear (see the detailed comments above). \nSome statements/derivations are not rigorous and thus are not convincing. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper seeks to improve straight-through estimators by combining them with the ideas for correcting the step direction to be closer to a natural gradient.
 
While some (modest) improvements are demonstrated experimentally, the paper critically lacks technical correctness and has quite some gaps when trying to derive the algorithm from the natural gradient and Rao-Cramer bound. See public comments by reviewers and AC. The algorithm ends up to be a mirror descent with a mirror map, which is cheap to compute but not particularly well motivated. Moreover application of mirror descent to the activations (unlike the weights) is not well justified. The paper is rather unclear and hard to read also language-wise. Please proofread _before_ submitting.
Decision: Reject