ISAAC Newton: Input-based Approximate Curvature for Newton's Method

Abstract: We present ISAAC (Input-baSed ApproximAte Curvature), a novel method that conditions the gradient using selected second-order information and has an asymptotically vanishing computational overhead, assuming a batch size smaller than the number of neurons. We show that it is possible to compute a good conditioner based on only the input to a respective layer without a substantial computational overhead. The proposed method allows effective training even in small-batch stochastic regimes, which makes it competitive to first-order as well as quasi-Newton methods.

Reviews and author responses:

 I have read other reviews and the rebuttal. I still have reservations about the significance of the one-parameter update proposed but I have no strong reasons to reject the revised paper.

 We wish to thank you for helping us improve the paper. Hopefully, you have had a chance to take a look at our responses. Please let us know whether we have addressed your concerns and questions.


 We wish to thank you for helping us improve the paper. Hopefully, you have had a chance to take a look at our responses. Please let us know whether we have addressed your concerns and questions.

 Thank you for the replies to my questions and comments. While some of my concerns are addressed, it is not enough to increase my score for the paper (which is already a 7). 

 Thank you for your fast reply and for increasing your score. We really appreciate it.

Regarding Figure 8, the wall-clock time for the blue line (gradient descent) is 2099.5 s while the time wall-clock time for the orange plot (ISAAC only for the last layer) is 2195.7 s, so it only takes 4.6% longer. 
We would like to note that, in Figure 8, we apply our method only to the last layer and use the regular gradient for all other layers. This means that any effect of our method is naturally limited to effects arising from the final layer; thus, the effects are rather small.

In this context, we would like to mention that we also conducted large-scale experiments where the entire networks are trained with ISAAC on the BERT tasks in Table 1.

*"training and test loss on MNIST are very similar"*: this only relates to the case of training an auto-encoder with an MSE loss. To clarify this, we *now also added Figure 9*, which displays the test accuracy for the experiment of Figure 5. Here, we can observe that, e.g., the purple plot ($\zeta^*$ for layers 3,4,5) has a training loss very close to the gradient baseline, while performing best on test accuracy. Overall, each configuration of using ISAAC for at least one layer performed better than gradient descent.

We hope that we were able to successfully resolve your concerns regarding the empirical evidence. 

 Thanks for providing additional experiments, given the limited time window. Overall, my concern regarding $\lambda_a$, $\lambda_g$ are resolved. But my concerns regarding the empirical evidence are still not resolved. 

I appreciate figures 7 and 8 the authors added to the appendix. I am happy to see that the proposed method achieves faster convergence in the MNIST test set (Figure 7). However, on the other side, I did not see similar performance gains in CIFAR-10 (Figure 8). With momentum and weight decay, ISAAC and SGD have very close iteration-wise convergence performance. In terms of wall-clock time, I expect ISAAC is even slower. As this is a commonly used setting in training neural networks, my concerns on convergence performance and the generalization ability of ISAAC are still not resolved. Furthermore, as observed that the training and test loss on MNIST are very similar, I am not sure MNIST is a good dataset to validate the convergence and generalization performance of ISAAC.

I know asking the authors to provide results on every benchmark dataset is unrealistic. However, since KFAC was proposed several years ago to train neural networks on CIFAR-10, ImageNet, etc.,  the authors should demonstrate the improvements, at least in small-scale settings such as CIFAR-10. Unfortunately, from Figure 8 in the appendix, I did not see obvious benefits from ISAAC. 



 Thank you following up and motivating us to strengthen our experimental section. *We have conducted the experiments as you requested and you will find the plots in Supplementary Material D.*

* **2.1**: We have added results for Adam and SGD with momentum for the experiment of Figure 2 (see Figure 7). The results show that the performance for SGD and SGD w/ Momentum (momentum=0.9) are almost indistinguishable on the experiment. Further, we note that Adam performs worse than SGD. We also note that we ran a grid search for the optimal learning rate in each case.
* **2.3**: We have added test loss plots for the experiments of Figure 2 (see the second row in Figure 7). We find that the performance relations between methods on the training dataset are also very similar on the test dataset.
* **2.4**: We have adapted the hyperparameters of a public CIFAR-10 implementation that lead to a 92% accuracy and present the results in Figure 8. The primary difference here to our original CIFAR-10 implementation is the use of image augmentation and the use of a learning rate schedule (where we follow the public implementation in using a cosine schedule). When comparing $\zeta^*$ on the last layer to the public CIFAR-10 implementation, we obtain a small improvement over the baseline. To do an ablation study wrt.~the optimizer, we also considered training without weight decay and without momentum. Here, we find that our advantage over the baseline is substantially increased when training without weight decay and/or without momentum. We note that we use a batch size of 128, following the public implementation. 
* **3**: We have added new experiments to empirically compare our method to SGD with a larger batch size of 600, as well as to K-FAC with a larger batch size of 600. These results are also presented in Figure 7. Here, we can observe that both SGD and K-FAC do not improve for larger batch sizes. We also tested other batch sizes, where the effect remains the same, and we report 600 as a representative. Here, we also provide train and test loss.

We will include these new experiments in the main paper for the camera-ready version of the paper.

Please let us know whether we have successfully addressed your concerns and whether you have any other questions regarding our paper.


 Thanks for the clarification. I would like to clarify the reason why I reject the paper. Overall, I am not convinced that this work presents a fair comparison between the proposed method and other baseline methods. Most importantly, some important metrics such as convergence on validation dataset are missing. As a result, I think contributions in the work are questionable. Below is my further response:

For 2.1, Figure 1 and Figure 4 do not show the convergence performance of Adam and SGD with momentum. I believe the authors should provide the result of Adam and SGD with momentum in Figure 2 (, also Figure 6). In Figure 2, the authors show KFAC and the proposed method has a better convergence performance compared to SGD. The author should really include other baseline so that we can have a direct comparison. KFAC is not the only baseline when showing the convergence performance. I believe the goal is still to show the advantage over commonly used optimizers.

For 2.3, I understand that the training loss shows the convergence performance. However, convergence on the validation dataset is also an important metric to evaluate the optimizer’s performance (convergence speed and generalization ability). As the authors stated in 3, one of the motivations for using smaller batch size is to improve the generalization performance. However, the authors do not provide convergence on the validation dataset to help the reviewers verify the claim and most importantly the generalization performance of the proposed method.

For 2.4, Tuning the hyper parameters indeed needs a lot of time, which I also think is cumbersome. However, I suggest the authors first adopt some standard hyperparameter settings for SGD. As ResNet on CIFAR-10 is a standard task, obtaining the optimal hyperparameter is not difficult. My big concern for this one is that the authors did not compare the baseline in a fair and optimal setting. If they are not in an optimal setting, the conclusion is not convincing.

For 3, I hope the authors provide a good motivation for small batches. In the paper, the batch size is small (e.g., 100 in CIFAR-10).  To clarify, the reason for using KFAC or other second-order methods is to improve the convergence performance. However, compared to large batches, training with small batches usually converges slower. As a result, it seems training with small batches offsets the benefits of using second-order methods. For instance, how does ISAAC with small batches compare SGD with large batches? If the proposed method is suitable for small batches, the authors should find a use case to justify it. 

 We thank the reviewer for the review and for providing valuable feedback, which improves our paper.


*main contribution*: Using two separate regularizers allows to study the influence that the components of the GGN have and to determine whether they are equally important. It turns out that they are not, but that actually one can be set to infinity with small harm, allowing for a more efficient training method. So, the main contribution is the special case mentioned in Remark 3, i.e., the efficient training method. We will make a clarification in the revision.

Reviewer 7xsl also mentioned having made some numerical validations and observing that $aa^\top$ is more important than $\bar g^\top \bar g$ in the Kronecker approximation of the GGN. This aligns with our findings.

*Theoretical results*: They may not require deep mathematical background, but we still believe that they convey many relevant insights about the method. Moreover, other reviewers clearly appreciate our theoretical results.

*Eq.21*: The empirical results presented in Figure 1 indicate that $\lambda_g$ has (relatively) little influence and can be set to infinity, allowing for the simpler and much more efficient update step in Theorem 3, Eq. 21.



Please let us know if you have any other questions regarding our paper or response. If we have successfully addressed your questions, we would strongly appreciate an increased score. 

 We thank the reviewer for the review and for providing valuable feedback, which improves our paper. Given the comments by the reviewer, we cannot really understand what led to the low grade.


1. Yes, if both $\lambda_a$ and $\lambda_g$ are increased/decreased, the method indeed reduces to SGD/KFAC. However, the novel discovery is that the method does not reduce to SGD if only $\lambda_g$ is increased to infinity, but instead maintains an accelerated convergence similar to KFAC, while requiring substantially less computational resources. 
 
We would also like to refer to reviewer 7xsl: reviewer 7xsl also mentioned having made some numerical validations and observing that $aa^\top$ is more important than $\bar g^\top \bar g$ in the Kronecker approximation of the GGN. This aligns with our findings. 

2.1. *Adam, SGD with momentum*: We would like to point out that we do provide experiments with Adam in Figure 4 and Table 1, and results for SGD with momentum in Figure 1 (c, g) and Table 1. 

2.2. *Figure 1*: The method does not reduce to SGD when only $\lambda_g$ is increased to infinity, but maintains some of the accelerated convergence of KFAC. It is not the same as SGD. We also note that when letting $\lambda_a$ go to infinity, we can observe that it does not work anymore. Figure 1 serves the purpose of highlighting the relative influence that the two components have, based on which we obtain the insight that it is harmless to let one of the regularizers (bot not both) go to infinity.

2.3. *Figure 2*: We did not use weight decay and hence there is no corresponding parameter. For Figure 2, we were interested in the convergence properties of the methods. We'll clarify this in the revised version. 
The main focus of this paper, similar to that of KFAC and other 2nd order method papers, is on how 
the 2nd order information can accelerate the optimization process, hence training loss curves are standard. 
However, if you would like to see that, we propose to also include a validation curve in the final paper.

2.4. *Figure 6*: We discuss the reasons for using a vanilla SGD setup in the paper. We note that vanilla SGD is only a secondary reason for the performance, the more important reason is the constant learning rate. However, tuning a learning rate scheduler would cause a combinatorial explosion of the search space, and without hyperparameter search the comparison would not be meaningful. If you would like to see that, we propose to also include results with SGD with momentum in the final paper, but we note that this has only a minor effect to the accuracy.

3. *small / large batch regimes*: 

With small-batch regime, we refer to mini-batch training, contrasting full batch training. In our experiments we use batch sizes that are typical for the respective experiments.

Small batch size yields faster convergence per epoch and is also known to help with generalization via implicit regularization. (Smith et al., 2021, "On the Origin of Implicit Regularization in Stochastic Gradient Descent", https://openreview.net/forum?id=rq_Qr0c1Hyo)
    
Reliable estimation of second order information usually requires larger batch sizes. But our experiments demonstrate that with our approach effective training is possible also in small-batch stochastic regimes.



Please let us know if you have any other questions regarding our paper or response. If we have successfully addressed your questions, we would strongly appreciate an increased score. 

 We thank the reviewer for the review and for providing valuable feedback, which improves our paper.

Thank you for pointing our the references, we will include and cite all of them in the revision. We will also polish the paper for the final version.

*Convergence result*:
Regarding a convergence result, we would like to point out that convergence results in the literature are given for the convex case of a linear model. In the linear case, there is no difference in formulation between our method and the Newton method, and hence, there is no difference in convergence compared to the Newton method. Only in the deep case, the difference between the methods arises. The linear case does not have direct implications for the deep neural network case. To the best of our knowledge, convergence guarantees for the Newton method for non-linear deep neural networks are not available. Hence, we cannot provide convergence guarantees for the Newton method nor our method for non-linear deep neural networks.

*ImageNet*:
We agree that results on ImageNet would be valuable, but we would like to point out that we do not have the hardware to train on ImageNet within a reasonable amount of time (even for SGD).

We will add a discussion between the GGN matrix and the Fisher matrix.

*Brackets in Eq.16:*
They would indeed make the formula easier to read and will be added.

*Eq.19:*
The equality holds regardless of a division by b. We will fix it in line 119. Yes, they could be 
from different batches, and the theory would still apply.

We agree that the method can be extended to convolutional neural networks.

You observation regarding $aa^\top$ being more important than $\bar g^\top \bar g$ in the Kronecker approximation of the GGN is very interesting, and it perfectly aligns with our observations and the motivation for this work. 
Thank you for bringing up NG+, we will cite it in the final paper.



Please let us know if you have any other questions regarding our paper or response. If we have successfully addressed your questions, we would strongly appreciate an increased score. 

 We thank the reviewer for the review and for providing valuable feedback, which improves our paper.

*H in line 70*: Yes, it should be $\nabla^2$.

*intuitive interpretation of the preconditioner*: What Remark 2 intuitively says is that $\zeta^*$ (the proposed method) lies somewhere between the Gauss-Newton step and the SGD step. The preconditioner $M$ translates the Gauss-Newton step to $\zeta^*$. Including $M$ does not improve performance of the Gauss-Newton update. But what we find is that the form including $M$ (i.e., computing $\zeta^*$ and not actually evaluating $M\zeta$) is very cheap to compute while impacting performance only slightly. We included the remark to elaborate the relationship between our method and the Gauss-Newton update. We will add a sentence to clarify the purpose of Remark 2.


Please let us know if you have any other questions regarding our paper or response. If we have successfully addressed your questions, we would strongly appreciate an increased score. 

 The paper addresses the problem of using second-order information in large-scale machine learning tasks (e.g. training neural networks). It proposes a Kronecker-factored Approximate Curvature (KFAC)-approximated, Tikhonov regularized Gauss-Newton (GN) update step. Instead of a single regularization parameter ($\lambda$), the update step in the paper uses two independent regularization parameters ($\lambda_g$ and $\lambda_a$). The paper relates different settings of $\lambda_a$ and $\lambda_g$ to other algorithms (i.e. classical GN and gradient descent). It shows that allowing $\lambda_g$ grows to infinity gives a computationally efficient update step. It provides empirical results on their proposed method. Strengths:
1. The update step $\xi^*$ is interesting and practical. The paper provides experiments to show that using $\xi^*$ can lead to lower computation time without sacrificing performance. Experiments are thorough and tested on some fairly large neural network architectures.
2. The update step $\xi$ is a nice, compact way to describe an algorithm for training neural networks. The analysis of $\xi$ provides various ways of interpreting the algorithm in terms of different regularization parameter settings. The properties of $\xi$ are interesting. The analysis is relatively thorough (e.g. showing that any combination of $\lambda_g$ and $\lambda_a$ will still result in a descent direction).
3. Paper is well-written and clear. 

Weaknesses
1. Some of the experiments are not benchmarked to SOTA (which is acknowledged by the paper). This might affect people's willingness to adopt the method in practice.   1. Line 70 refers to $H$ in Equation 6. Typo?
2. Is there an intuitive interpretation of the preconditioner $M$ in Remark 2? Does this preconditioning help in terms of reducing the number of iterations needed to reach a solution? This work does not appear to me to have negative social impact.

 Inspired by the KFAC approximation and the Tikhonov regularization of the Newton method, this paper proposes a regularized K-FAC approximated Gauss-Newton type method called ISAAC. The proposed method is based on only the input to a respective layer, therefore does not suffer from a substantial computational overhead. Numerical experiments show the advantages of the proposed method compared to both first-order type and second-order type methods. 
 Strengths:

Regularization parameters is crucial for second-order type methods.  The authors first approximate the GGN matrix as a kroncker product of two matrices, and then view two regulation parameters independently. By letting the two regularization parameters vanish or go to infinity, one can obtain new directions and therefore new methods. The proposed method is more efficient than KFAC since only one matrix should be inverted and no need extra propagation to compute $g$.  

Weaknesses:

This paper should be polished to make it easier to read. 

Some related literatures on second-order methods should be mentioned, such as Shampoo[1], KBFGS[2], NG+[3] and SENG[4]. 

Some theoretical properties is natural and trivial. It is better to formalize an algorithm and give a convergence result.

It is better to add some large numerical examples, such as Resnet50 with Imagenet-1k and make a comparison with other second-order methods Shampoo[1], KBFGS[2], NG+[3] and SENG[4]. 

[1] Gupta,V.,Koren,T., Singer,Y.:, Shampoo: Preconditioned stochastic tensor optimization. In: International Conference on Machine Learning.

[2] Goldfarb, D., Ren, Y., Bahamou, A.:, Practical quasi-Newton methods for training deep neural networks. In:Advances inNeural Information Processing Systems.

[3] Yang, M., Xu, D., Cui, Q., Wen, Z., & Xu, P. (2021). NG+: A Multi-Step Matrix-Product Natural Gradient Method for Deep Learning. arXiv preprint arXiv:2106.07454.

[4] Yang, M., Xu, D., Wen, Z., Chen, M., & Xu, P. (2020). Sketch-based empirical natural gradient methods for deep learning. Journal of Scientific Computing. line 60: change [14],[15] to [14.15]. It is better to put multiple references in the same brackets. Please check other places.

It is better to add the discussion between GGN matrix and Fisher matrix, which can refer to [1,2] and related literatures.

It is better to add a left brackets after $\eta$ and a right bracket before the gradient in eq.(16). Please also add brackets in a suitable place in eq. (17).

In eq.(19), the gradient should be $\mathrm{g}^\top a /b$ where the $b$ is missed. In eq.(19), the authors use the same batch to approximate the GGN matrix and compute the gradient. Can the batches for approximation of the GGN and the computation of the gradient be different? In that case, what the related theoretical properties will be?

When the objective function is convex, the GGN matrix is positive semi-definite matrix. Adding a regularization parameter in GGN, It is t easy to show that $\zeta$ is an ascent direction, which means Theorem 2 is trivial.

The authors consider the fully-connected neural network in this paper. I believe this method can be extended to convolutional neural networks. And in numerical part, the authors have done some experiments. I wonder can we obtain some explanation from theory for this case? 

For the kronecker approximation of the GGN matrix, $\mathrm{a}\mathrm{a}^\top$ is more important than $\overline{g}^\top\overline{g}$ part. I once validate it numerically. One reason is that each value of the matrix $\overline{g}^\top\overline{g}$  is very small so that the $\lambda_g$ domains. I have noticed one related work recently. The work is named NG+[3].

[1] Botev, Aleksandar, Hippolyt Ritter, and David Barber. "Practical gauss-newton optimisation for deep learning." International Conference on Machine Learning. PMLR, 2017.

[2] Martens, James. "New insights and perspectives on the natural gradient method." The Journal of Machine Learning Research 21.1 (2020): 5776-5851.

[3] Yang, M., Xu, D., Cui, Q., Wen, Z., & Xu, P. (2021). NG+: A Multi-Step Matrix-Product Natural Gradient Method for Deep Learning. arXiv preprint arXiv:2106.07454. Yes

 This paper proposed a KFAC variant to improve the computation efficiency. It introduces two Tikhonov regularization terms into KFAC, and makes gradient conditioning more efficient by force $\lambda_g$ go to infinity. The authors present a theoretical guarantee that the optimization still goes in the right direction even with such an adaptation. The authors also show ISAAC is more computation efficient in small-batch stochastic regimes.

However, I feel contributions in this paper are rather limited, and some motivations might need further justification.
 Strength:

The paper is well written. Preliminaries and methods are clearly explained. 

Weaknesses:

1, Contributions are limited. Reducing the compute costs of KFAC is a good problem. However, I am not convinced that separating the regularization term and using a large regularization is a novel and effective solution. It only makes the optimizer behave more and more like SGD. As the authors stated in Theorem 1, small $\lambda_a, \lambda_g$ leads to KFAC, while large $\lambda_a, \lambda_g$ leads to vanilla gradient descent. 

2, Empirical results need to be enhanced. The authors use several empirical results to support the claim that ISAAC still behaves like KFAC even though using a strong regularization factor. However, I do not think current results are strong enough to support the claim. The issues in the empirical evaluations also make me doubt the contributions in the paper. 

- 2.1, The author should provide more baselines (Adam, SGD with momentum) so that we can really have a fair evaluation. 

- 2.2, For Figure 1, my understanding is that the authors are trying to show ISAAC still obtains good results (loss and accuracy) even with large regularization. If so, I do not think it is necessary. As the authors stated in Theorem 1, larger regularization makes ISAAC behave more like SGD. Therefore, getting similar loss and accuracy is no surprise.
    
- 2.3, For Figure 2, there are several issues. First, hyper parameters are not fully revealed. For instance, what is the weight decay for each optimizer since weight decays large affects training convergence?  Second, only the training curve is reported. The author really should present both a training and validation curve. In many cases, we can have faster convergence on training dataset, but not for validation dataset. 

- 2.4, For Figure 6, the reported accuracy is obviously lower compared to results we usually have on CIFAR-10. The reason might be that the authors use SGD without momentum. However, I am not convinced under such a setting. SGD with momentum is a default setting in many training tasks. Without momentum, I do not think the results will provide meaningful comparisons and conclusions.

3, Training with small batches should be further justified. I understand the authors stress the compute benefits of ISAAC in a small-batch regime (Theorem 3-ii). However, considering more general settings, the authors should well justify it: is it worth accelerating the training using small batches given the fact that large batches can also improve convergence performance even using SGD. Therefore, I think not only should the author provide results in small-batch regimes, training using large batches should also be evaluated and compared with other baselines. See my comments above. No

 In unconstrained optimisation, Newton's method uses inverse Hessian to determine the step size of the update. When applied to deep learning, the sheer size of the Hessian matrix makes Newton's method intractable. Among existing attempts to address this challenge, K-FAC is a promising approach. In practice, Tikhonov regularization/damping is used to avoid divergence.

In this paper, the authors introduced an extension to the standard Tikhonov damping for K-FAC. The main idea is to use two independent damping parameters (hyper-parameters for regularisation) instead of one damping parameter, one for the gradient factor and one for the activity factor, as shown in Equation (17). The authors explored special cases where the two parameters are either 0 or infinity. They also discussed the special case of setting one parameter to infinity, effectively removing the gradient factor in K-FAC. One major claim was that this activity-only form has an advantage in reduced computational complexity, although I don't find a theoretical comparison with K-FAC in this aspect. 

The claim for computational efficiency was mostly supported via empirical evaluation. Strengths:

1. The paper addresses an important problem: efficient approximation of Gaussian-Newton approximation
2. The exposition is mostly clear, especially if you are already familiar with the existing literature.
3. The experiments seem comprehensive.

Weaknesses:

1. It is not clear what is the main contribution of this paper. Is it the two-parameter damping mechanism (What advantage do we gain)? Or is it the special case mentioned in Remark 3 (Why this is a still good approximation without $g^{\top}g$)? 
2. Many theoretical results presented look rather trivial. What is the intuition behind Eq (21)? Compared with Eq (13), why is it a good approximation? N/A

Meta-review: The paper adds two regularization parameters that creates a method that interpolates between gradient descent and KFAC. The authors then prove the method interpolates, then several numerical experiments are presented on MNIST, and training BERT, and resnets. But the reviewers were not convinced by the experimental results, which did not benchmark the methods against SOTA, or lacking large scale experiments, and also lacking comparisons against other second order methods, for which there are now many.

In addition to this, though the paper cites broadly, it does not review, cite or cover the recent efforts on improving KFAC.  For example the missing reference [1]. Though this is a quasi-Newton approach, they also employ a type of regularization to aa^T term. Finally, I think the introduction of these two regularization parameters needs to be slightly better motivated. For now, the motivation is to interpolate between SGD and KFAC. But this is not enough. One can always generate new methods by interpolating between existing methods by adding a parameter. I recommend exploring the LM (Levenberg Marquardt) viewpoint of this type of regularization. This may gives other viewpoints and motivation for using these two regularization parameters.

[1] Goldfarb, D., Ren, Y., Bahamou, A.:, Practical quasi-Newton methods for training deep neural networks. In:Advances inNeural Information Processing Systems.

Decision: Reject