Optimizing Quantized Neural Networks in a Weak Curvature Manifold

Abstract: Quantized Neural Networks (QNNs) have achieved an enormous step in improving computational efficiency, making it possible to deploy large models to mobile and miniaturized devices.
      In order to narrow the performance gap between low-precision and full-precision models, we introduce the natural gradient to train a low-precision model by viewing the parameter space as a Riemannian manifold.
      Specifically, we propose a novel Optimized Natural Gradient Descent (ONGD) method defined by the Hyperbolic divergence, which provides a perspective to calculate the optimized natural gradient in weak curvature and updates the parameters with an amount of computation comparable to Stochastic Gradient Descent (SGD).
      We conduct an ablation study and results show that the 4-bit quantized ResNet-32 trained with ONGD has a better result than SGD, i.e. 2.05\% higher in Top-1 accuracy on CIFAR100 dataset.
      Further comparison experiments illustrate that our method achieves state-of-the-art results in CIFAR and ImageNet datasets, where the 8-bit version of MobileNet achieves 0.25\%/0.13\% higher in Top-1/Top-5 accuracies than the full-precision version on ImageNet dataset.

Reviews and author responses:

We thank the reviewer for the comments, and would like to clarify several things to address the reviewer's concerns:

Q1: The target of the method

A1: The main goal of the QNN is 'In order to narrow the performance gap between low-precision and full-precision models' in the abstract. The proposed training algorithm goal is not a low computing resource but to achieve higher accuracy at a lower bitwidth in inference. So our experiments compare the accuracy with other quantizers under the same or lower bitwidth. Of course, Efficient inference is also part of our consideration. In fact, we test the quantization method ONGD without post-training model in the ablation study in Table 3, which is used to compare to post-training quantization method ONGD-D. As you said, if it is for efficient inference, then whether to ensure low-precision in training is not the key, as long as the inference process is guaranteed to be low-precision. Just like the baselines compared in RQ [1] and INQ [2] are also training methods with low precision compute during training directly, where INQ must rely on post-training model, and RQ does not use low-precision compute directly during training.

[1] Christos Louizos, Matthias Reisser, Tijmen Blankevoort, Efstratios Gavves, and Max Welling. Relaxed quantization for discretized neural networks. In International Conference on Learning Representations, 2019. URL https://openreview.net/forum?id=HkxjYoCqKX.

[2] Aojun Zhou, Anbang Yao, Yiwen Guo, Lin Xu, and Yurong Chen. Incremental network quantization: Towards lossless cnns with low-precision weights. arXiv preprint arXiv:1702.03044, 2017.

Q2: Straight through estimator

A2: We are sorry for our negligence. We have added the n-bit case in Appendix 6.1, which can apply STE to quantized values symmetrically distributed on both sides of zero. Therefore, we can still use STE at 4bit fixed point number.

Q3: The pretrained model

A3: This doesn't seem right. The full precision validation accuracy numbers are directly from the reference [3], which does not use our proposed training method. Actually, our method ONGD is trained from scratch without the post-training full-precision model, which is reported in Table 1 and Table 3. Only the method ONGD-D uses the distillation that utilizes the smooth label information of the full-precision model. For the consideration of the fairness of the comparison without the pretrained model to jump start, we have added the ablation study on ImageNet for ONGD and ONGD-D in Table 3. The experimental result shows that both optimized natural gradients and distillation can improve the accuracy of neural networks based on these ablation studies.

[3] Mostafa Elhoushi, Zihao Chen, Farhan Shafiq, Ye Henry Tian, and Joey Yiwei Li. Deepshift: Towards multiplication-less neural networks. arXiv preprint arXiv:1905.13298, 2019.

Q4: Minor problems

A4: We have already added details about eq.10 in a footnote: Note that we use $\textbf{W}$ or $\hat{\textbf{W}}$ to denote a large matrix of $l\times n$ that is composed by each layer of weight tensor $\textbf{w}$ or $\hat{\textbf{w}}$ as a row vector, where $n$ is the largest size of rows and other insufficient rows are filled with zero. The weight vector $w_i$ or $\hat{w}_i$ is obtained by straightening $i$-th layer of the weight tensor $\textbf{w}$ or $\hat{\textbf{w}}$ into a row vector. The weight vector $w$ or $\hat{w}$ is obtained by straightening the weight tensor $\textbf{w}$ or $\hat{\textbf{w}}$ into a row vector. Then, the citation formatting has switched. Finally, we have defined the training phase for eq.21 and eq.22, where eq.21 is only used in the first training phase of distillation, and eq.22 is used in the second training phase of distillation.


We thank the reviewer for the comments, and would like to clarify several things to address the reviewer's concerns:

Q1: Declaration of the optimized natural gradient

A1: We are sorry for not making the hypothesis very clear. The previous manuscript was misunderstood and did not give the conditions of use and proof. We have added two new chapters 3.2 and 3.3 in modification to prove in detail how we have the optimized natural gradient in weak curvature condition. And we changed the 'natural gradient' to 'weak curvature manifold' in the title and revised the inaccurate statement you mentioned in the introduction and conclusion. Let us explain why it is possible to define the optimized natural gradient in a weak curvature manifold and link it with smooth quantization (eq.19). Firstly, we use the exponential mapping to connect $w$ and $\hat{w}$ on the manifold based on definition 2. Secondly, we symmetrize it to get a convex function with a geometric structure in eq.14. On this basis, we define a novel Hyperbolic divergence based on definition 3 that satisfies the criteria of divergence in Appendix 6.2. By the Taylor expansion of Hyperbolic divergence in Appendix 6.5, we derive a Riemannian metric that is close to the Euclidean metric in chapter 3.2, i.e. the background of the parameter space is a weak curvature manifold. Finally, imitating the natural gradient under the constraint of KL divergence in Appendix 6.4 and chapter 2.2, we derive the optimized natural gradient under the constraint of Hyperbolic divergence in weak curvature in chapter 3.3. Since FIM is obtained by KL divergence (in arbitrary curvature) in Appendix 6.4 and we newly define the Hyperbolic divergence in weak curvature, FIM is no longer needed. The optimized natural gradient with the constraint Hyperbolic divergence is consistent with our designed smooth quantization (eq.19) using STE. In a word, by defining a novel Hyperbolic divergence to replace the KL divergence, we obtain a method to update parameters in weak curvature instead of arbitrary curvature. It is still a natural gradient because the metric representing this manifold is a Riemannian metric based on eq.16, eq.12 and eq.4. 

Q2: Mathematical notations

A2: We are sorry for our confusing indicators. We define 4 forms of weight in the footnote of page 4. Note that we use $\textbf{W}$ or $\hat{\textbf{W}}$ to denote a large matrix of $l\times n$ that is composed by each layer of weight tensor $\textbf{w}$ or $\hat{\textbf{w}}$ as a row vector, where $n$ is the largest size of rows and other insufficient rows are filled with zero. The weight vector $w_i$ or $\hat{w}_i$ is obtained by straightening $i$-th layer of the weight tensor $\textbf{w}$ or $\hat{\textbf{w}}$ into a row vector. The weight vector $w$ or $\hat{w}$ is obtained by straightening the weight tensor $\textbf{w}$ or $\hat{\textbf{w}}$ into a row vector. Based on the above definition, we can easily find that $dL/dw_i$ is a row vector. In addition, since the input $\hat{a}_{i-1}^\top$ and output $\frac{\partial L}{\partial a_i}$ sizes of each layer may be different based on eq.11, Kronecker product is necessary. Finally, we have marked $Q$ for forced quantization (eq.7) and $\widetilde{Q}$ for smooth quantization (eq.19) to prevent ambiguity. At the same time, we replaced $\alpha$ in eq.19 with a constant factor $s$.

Q3: Reference

A3: We are sorry for our negligence. We have added the reference for definition 1 and 2, that comes from [1,2].

[1] Robert M Wald. General relativity. University of Chicago press, 2010.

[2] Sigurdur Helgason. Differential geometry and symmetric spaces, volume 341. American Mathematical Soc., 2001.

Q4: Recent quantization techniques

A4: There are some misunderstandings here. In fact, we are not compared to Binaryconnect on CIFAR. BinaryConnect is for 1-bit weights and 32-bit activations, and our method is for 4-bit weights and 4-bit activations. In addition, we have added the references you mentioned for the comparative experiments in Table 1 and Table 2. For example, the test error with ONGD-D (1 bit weight and 1 bit activation) achieves 2.23% lower compared to Bi-Real (1 bit weight and 1 bit activation) on ImageNet.

Q5: Minor problems

A5: We have added the equation for FIM in eq.6. As seen from Appendix 6.4, eq.5 and eq.6, the natural gradient is derived by KL divergence of two distributions with very close distance, which happens to be FIM after the Taylor expansion of the KL divergence. Thus, we say the natural gradient, FIM and KL divergence are closely related. In addition, we have added the experimental details for ImageNet experiments in Appendix 6.7: we use the initialization of Xavier, and quantize the weights and activations of all layers except the first layer where the low-precision model is trained from scratch. Finally, we deleted the inaccurate statement you mentioned in the experiments.


We thank the reviewer for the comments, and would like to clarify several things to address the reviewer's concerns:

Q1: Declaration of curvature information

A1: We are sorry for not making the hypothesis very clear. In fact, our work is divided into two parts, the first part is to introduce the work of others into the QNN, which is applied to arbitrary curvature, but the amount of calculation is huge. In the second part, we have added two new chapters 3.2 and 3.3 in modification to prove our hypothesis in detail how we have the optimized natural gradient in a weak curvature manifold. The core part (FIM) in eq.10 is defined by the constraint KL divergence in Appendix 6.4, which has nothing to do with smooth quantization (eq.19). The smooth quantization is consistent with the optimized natural gradient that is defined by the constraint Hyperbolic divergence in eq.17 and eq.18. As for the curvature information captured by eq.19, we use the exponential mapping to connect $w$ and $\hat{w}$ on the manifold based on definition 2. Then, we symmetrize it to get a convex function with a geometric structure in eq.14. On this basis, we define a novel Hyperbolic divergence based on definition 3 that satisfies the criteria of divergence in Appendix 6.2. By the Taylor expansion of Hyperbolic divergence in Appendix 6.5, we derive a Riemannian metric that is close to the Euclidean metric in chapter 3.2, i.e. the background of the parameter space is a weak curvature manifold. In a word, eq.19 captures the curvature information as long as there is a Riemannian metric instead of the Euclidean metric.

Q2: Ablation study

A2: We are sorry for our negligence. We have added the fair comparison experiments to further demonstrate the superiority of optimized natural gradient and distillation in Table 1 and Table 3, where ONGD only uses the optimized natural gradient, while ONGD-D uses both the optimized natural gradient and distillation. Compared with general gradients, both optimized natural gradients and distillation can improve the accuracy of neural networks based on these ablation studies.

Q3: Comparative experiment for ADAM

A3: This is a very attractive suggestion. We have added the comparison of ADAM in Table1, Table3, Figure 2 and Figure 3. For ADAM, regardless of whether the learning strategy is a constant learning rate or a step-down learning rate, it is not as effective as SGD based on Figure 2 and Figure 3.

Q4: Comparative experiment

A4: We have added the method MetaQuant [1] you mentioned in Table 2. The experimental result shows that the test error with ONGD-D (1 bit weight and 8 bits activation) achieves 3.73\% lower compared to MetaQuant (1 bit weight and 32 bits activation).

[1] 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.

Q5: Repeatability

A5: Of course, our experiments are all obtained through repeated experiments. We will publish the code later, and the experimental results can be reproduced.


We thank the reviewer for the comments, and would like to clarify several things to address the reviewer's concerns:

Q1: Hypothetical argument

A1: We are sorry for not making the hypothesis very clear. We have added two new chapters 3.2 and 3.3 in modification to prove our assumption in detail how we have the optimized natural gradient in a weak curvature manifold. Firstly, we use the exponential mapping to connect $w$ and $\hat{w}$ on the manifold based on definition 2. Secondly, we symmetrize it to get a convex function with a geometric structure in eq.14. On this basis, we define a novel Hyperbolic divergence based on definition 3 that satisfies the criteria of divergence in Appendix 6.2. By the Taylor expansion of Hyperbolic divergence in Appendix 6.5, we derive a Riemannian metric that is close to the Euclidean metric in chapter 3.2, i.e. the background of the parameter space is a weak curvature manifold. Finally, imitating the natural gradient under the constraint of KL divergence in Appendix 6.4 and chapter 2.2, we derive the optimized natural gradient under the constraint of Hyperbolic divergence in weak curvature in chapter 3.3. Since FIM is obtained by KL divergence (in arbitrary curvature) in Appendix 6.4 and we newly define the Hyperbolic divergence in weak curvature, FIM is no longer needed.

Q2: Experiment supplement

A2: For the misleading experiments in Table 2, we have added new comparison methods to evaluate the performance fairly in Table 2. With the same or less bitwidth, our method can achieve higher accuracy. For ResNet50, we add the comparison with 8 bits weight and 8 bits activation (the test error with ONGD-large achieves 0.23\% lower compared to INT8 [1]). At the same time, we design a new ablation study simultaneously in Table 3 to further demonstrate the superiority of our method, where ONGD only uses the optimized natural gradient, while ONGD-D uses both the optimized natural gradient and distillation.

[1] 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 on Computer Vision and Pattern Recognition, pp. 1969–1979, 2020.

This paper proposed to incorporate local curvature information in the parameter space to assist the training of quantized model. Specially, it utilized natural gradient, with the complex Finsher Information Matrix replaced by a smooth quantization. It further used distillation from full-precision model and an alternative update between quantized/full-precision model to improve performance.

Cons:
1. Experiments comparison in Table.2 is strange: comparsion should be made based on same bit width of weights/activation.

I have the following questions:
1. Author spent the whole Sec.3.1 to describe usage of natural gradient in quantized training. However, it replaced the core part (FIM) in Eq.(8) with a smooth quantization (Eq.10), which is published before by another work. Is there any intrinsic connection between smooth quantization and natural gradient? Author mention that such replacement is a hypothesis, which quite confused me, since it seems that Eq.(10) can hardly capture curvature information as claimed by the author.
2. This work used distillation and training of full-precision model. It is better to conduct ablation study to validate the effectiveness of proposed gradient modification and distillation.
3. ONGD is compared against SGD. However, quantized model can also be trained via other optimization method, such as Adam, especially in large model. How the performance of ONGD compared to Adam?
4. According to Eq.(11), gradient of full-precision parameters is a multiplication of gradient of quantized parameter and a tanh function. If so, author should compare with other quantization training method which also modify the graident, such as [1]
5. Since performance gap in CIFAR is narrow. Is repeated experiments conducted to reduce variance?

[1] https://papers.nips.cc/paper/8647-metaquant-learning-to-quantize-by-learning-to-penetrate-non-differentiable-quantization

In this paper, the authors proposed a new method to train DNN (partially) with fixed point weights and approximated natural gradients. The contribution of this paper comes in three folds. First, they propose to use an exponential map to encode the curvature information instead of using the inverse of Fisher information matrix. Secondly they propose to use this technique to compute an approximation of the natural gradient. Lastly, the authors proposed to train low precision DNN using a combination of this approximated natural gradient and an additional loss term which enforces consistency across a full precision and the low precision model weights. Empirically, the authors demonstrate that the proposed method can attain better validation accuracy than plain SGD when training fixed point low precision convolutional models.

Strong points: empirically the proposed partially low precision training method in general attains meaningfully higher validation accuracy than plain SGD and other baseline methods (the plain SGD and baseline methods use a pure low precision fixed point model representation).

Concerns: I have 3 major technical concerns and questions which lead to my initial rating to reject.

1. In the high level, is the proposed method targeting at low precision training with less training compute resource, or is it just aiming at to attain a low precision model that can be efficiently deployed? This is not clearly elaborated in the paper. 
     
If it is for training with low compute resource, the proposed training algorithm will involve forward and backward evaluation of a full precision model. The additional compute from the low precision model weights can only increase the training overhead compared to a pure full precision training algorithm. 

If the proposed method is just for attaining a low precision model for efficient inference, then the experiment protocol is not convincing to me. Specifically, the authors does not compare to post-training quantization after full precision training methods. Instead, the major baselines in the paper is training methods with low precision compute during training directly.

2. The definition of straight through estimator in this paper is not correct for general low precision fixed point training. The definition in the paper is only for 1 bit (binary) model weights; however in the experiment the authors considers 4 bit fixed point number. The straight through estimator for non binary model weights will involve indicator functions depending on where the nearest quantization step is.

3. In the ImageNet experiment, the authors use an already well trained full precision model to jump start the proposed partially low precision training algorithm. Is the full precision validation accuracy numbers directly fro them model weights which jump-starts the proposed method? If yes, then the fairness of the comparison needs to be enhanced. This is because the proposed methods is using additional training steps to build on top of the well trained full precision model weights.


Comments for improving the paper:
1. The vector notations in equation (8) needs to be clarified in the paper.
2. Citation formatting needs to switch to \citep instead of \cite in the second paragraph in intro section.
4. In equation (14) and (15), it might need to explicitly define which weights is updated in each training phase.







The authors propose an interesting method to train a quantized neural network with a natural gradient, which considers the curvature information of the model. They also provide an efficient way to approximate Fisher Information Matrix using the connection between quantized and full-precision numbers. Although this idea is interesting and novel, I am concerned about some assumptions made by the authors and the empirical performance of the proposed method. 

1. The authors assume that the connection between w and \hat{w} on the manifold can partially reflect the curvature distribution. Because this assumption is curial for the approximation of Fisher Information Matrix, it would be great if authors could provide more justifications on why such an assumption holds.

2. The results for ImageNet seems misleading. As shown in Table 2, The authors didn't report the same number of baselines for each model. Also, the bitwidth of the baselines are different, which is hard to evaluate the performance fairly. For example, they show the results of proposed ONGD-large with 8 bits weight and 8 bits activation for ResNet50, but they report some other baselines with the setting of 4 bits weight and 4 bits activation. 

In general, I think this paper provides an interesting connection between quantization-aware training and natural gradient optimization. I would champion this paper if the authors can provide fair and detailed empirical results.

## Summary 

The paper proposes a neural network quantization approach with a tanh-based quantization function and claims that it is a good approximation to the natural gradient method. In addition, they augment the cross-entorpy loss with a knowledge distillation type loss (between the logits of the real network and quantized ones while training) to improve the performance. The results seem to show marginal improvements over compared methods in cifar and imagenet datasets.

## Strengths

1. The idea of using a knowledge distillation type loss between the logits of the real network and quantized ones while training is interesting and could be one of the reasons for performance improvement.
2. Overall the paper is clearly written.
3. Reasonable experiments are conducted (even though comparison with some recent methods are missing) and the results show marginal improvements over compared methods.

## Weaknesses

The main weaknesses of the manuscript in my opinion are as follows:
1. The final method seems to be not based on natural gradient and the claimed connection is not explained: 
	- In fact, the final method simply uses a tanh-based quantization function (Eq. 10) and simply uses SGD with STE approximation. There is no explanation of why this approach closely approximates the natural gradient method, given that this the main claim of the paper and even the title contains "natural gradient". Definition 2 might be coming from a textbook (even though no reference is given) and it simply provides an existence condition and it does not say anything about the proposed exponential map corresponds to the manifold induced by the Fisher Information Matrix. This is critical for the paper and without this connection the whole discussion about natural gradient is redundant and could be deemed as "decorative maths". Please explain (with supporting theory) how Eq. 10 relates to FIM.
	- Given that the method is not based on natural gradient, there are many unsubstantiated claims. Examples include: 1) "unique limitaions of the QNNs to construct curvature information" 2nd last para in the introduction, 2) "demonstrated feasibility of training low-precision models using natural gradients" first para of Sec. 3.2, and 3) "our method .. avoids the risk of neural networks falling into the local minimum" in conclusion.

2. Misinterepretations and inconsistencies in mathematical writing/notations:
	- Eq. 8 has many issues. First, what is $vec(dL/dw_i)$? If $w_i$ is a scalar, $dL/dw_i$ is a scalar. So what does it mean to vectorize it?
	- To my understanding $F_{ij} = g_i g_j$ where $g_i = dL/dw_i$. There is no vector representation and no need for Kronecker product etc. Please clarify.
	- Please clearly mention the dimensions of each variable and define them appropriately for all the equations. Also, do not use same symbol to denote two things. Examples include $Q$ (Eq. 5, 10) and $\alpha$ (Eq. 10, 14).

3. No references given for the statements related to Riemannian geometry:
	- For both definitions 1 and 2 no reference is given. It seems like there are existing results and they should be cited so that someone can check/verify their applicability in the stated problem setting.

4. Recent quantization techniques are not compared:
	- For cifar experiments only comparison is done against the first quantization method BinaryConnect (Courbariaux, 2015) which is very old. Especially, comparison against other relaxed quantization methods [a,b] are required.
	- Even for imagenet experiments important comparisons are missing such as [c,d].

## Minor Comments

1. Please provide the equation for FIM before Eq. 4.
2. Second para in page 4: what is meant by "closely related"?
3. Please provide all the experimental details required for reproduction (eg, any layers kept floating point in imagenet experiment?, what initialization used?, starting from scratch or pretrained network, etc.)
4. After fig. 2, please tone down. The figures only show a marginal improvement in test error not "significantly preventing overfitting".

## References

- [a] Bai, Y., Wang, Y.X. and Liberty, E., 2019. Proxquant: Quantized neural networks via proximal operators. ICLR.
- [b] Ajanthan, T., Dokania, P.K., Hartley, R. and Torr, P.H., 2019. Proximal mean-field for neural network quantization. ICCV.
- [c] Yang, J., Shen, X., Xing, J., Tian, X., Li, H., Deng, B., Huang, J. and Hua, X.S., 2019. Quantization networks. CVPR.
- [d] Liu, Z., Wu, B., Luo, W., Yang, X., Liu, W. and Cheng, K.T., 2018. Bi-real net: Enhancing the performance of 1-bit cnns with improved representational capability and advanced training algorithm. ECCV.

Meta-review: The reviewers unanimously raised concerns over the clarity and technical correctness of the theory and the Imagenet experiments during the first round. The authors submitted a highly revised version during the rebuttal which allayed concerns for multiple reviewers, however all the reviewers raised the concern that the paper has gone through a very significant change, almost becoming "a new paper" in the rebuttal phase and should go through another cycle of resubmission and review to correctly judge the contributions and claims.  The reviewers were aligned in their judgement that the original manuscript with minor changes was not ready for publication. The authors are strongly encouraged to polish the version submitted during rebuttal and resubmit.  

Decision: withdrawn-rejected-submissions