AdaFisher: Adaptive Second Order Optimization via Fisher Information

Authors: Damien MARTINS GOMES, Yanlei Zhang, Eugene Belilovsky, Guy Wolf, Mahdi S. Hosseini
Primary area: optimization
Keywords: Second Order Optimization, Fisher Information, Kronecker-factored Approximate Curvature, Deep Learning, Computer Vision, Natural Language Processing

Abstract:
First-order optimization methods are currently the mainstream in training deep neural networks (DNNs). Optimizers like Adam incorporate limited curvature information by employing the diagonal matrix preconditioning of the stochastic gradient during the training. Despite their widespread, second-order optimization algorithms exhibit superior convergence properties compared to their first-order counterparts e.g. Adam and SGD. However, their practicality in training DNNs is still limited due to increased per-iteration computations compared to the first-order methods. We present *AdaFisher*--an adaptive second-order optimizer that leverages a *diagonal block-Kronecker* approximation of the Fisher information matrix for adaptive gradient preconditioning. AdaFisher aims to bridge the gap between enhanced *convergence/generalization* capabilities and computational efficiency in second-order optimization framework for training DNNs. Despite the slow pace of second-order optimizers, we showcase that AdaFisher can be reliably adopted for image classification, language modeling and stands out for its stability and robustness in hyper-parameter tuning. We demonstrate that AdaFisher **outperforms the SOTA optimizers** in terms of both accuracy and convergence speed. Code is available from https://github.com/AtlasAnalyticsLab/AdaFisher.

Decision: accept (poster)
Reviewer scores: 6, 5, 6, 8

Reviews and discussion:
<thread>
creation date:
1730139524785

thread_type:
official_review

<thread_object>
creator: reviewer_k3hy

<content>
summary:
The authors introduce **AdaFisher**, a novel adaptive second-order optimizer that uses a block-diagonal approximation of the Fisher Information Matrix (FIM) to improve both convergence speed and computational efficiency in second-order optimization. The paper aims to address the limitations of current second-order methods, such as high computational costs and poor generalization, by proposing a diagonal block-Kronecker approximation of the FIM. The authors demonstrate that AdaFisher outperforms several SOTA optimizers on tasks such as image classification and language modeling.

strengths:
1. **Novel methodology:** The introduction of a diagonal block-Kronecker approximation applicable to various layers offers an interesting balance between computational efficiency and the use of curvature information typically employed in second-order optimization. 

2. **Empirical results:** The paper provides comprehensive experimental evidence showing that AdaFisher outperforms baseline methods (Adam, K-FAC, AdaHessian, etc.) on benchmark datasets such as CIFAR-10, CIFAR-100, ImageNet, and WikiText-2 across different network architectures.

3. **Stability:** AdaFisher shows strong stability across varying learning rates and batch sizes, reducing the need for extensive hyperparameter tuning, which is a common challenge when training deep models.

4. **Theoretical contribution:** The paper presents a rigorous theoretical convergence analysis for both convex and non-convex cases, asserting a convergence rate of  $O(\log T / \sqrt{T})$ , similar to Adam-type methods. The derivation of the update rules using the diagonal approximation of the FIM is clearly explained.

weaknesses:
1. The paper considers nonconvex optimization, but in reality DNNs are **nonsmooth**, due to operations such as ReLU and max-pooling. It would be valuable to discuss this limitation in Section 3.4. 

2. **Hyperparameter tuning:** It is unclear how many epochs were used for hyperparameter tuning, especially in grid search experiments (Appendix D). Providing more details on the tuning process would enhance the reproducibility of the results.

3. **Comparision with AdamW:** AdamW has become the most widely used optimizer due to its improved handling of weight decay. It would have been interesting to compare AdamW directly with AdaFisherW. (in particular, for image classification).

questions:
1.	**Figure 1:**  I try reproducing Figure 1 using your code. In my experiments, AdaFisher did not consistently converge to similar local minima across multiple runs. Could you discuss the sensitivity to initialization or provide further insights?

2.	**Algorithm 1:** How did you choose the values for $\gamma_1$ and $\gamma_2$? You mentioned that “the decay factors $\gamma_1$ and $\gamma_2$ for AdaFisher were tuned within $\(\{0.1, 0.2, \dots, 0.9, 0.99\}\)$, but you did not specify which values were  used. Providing these details would help practitioners apply AdaFisher effectively.

3.	**Related work:** The paper does not mention some other recent second-order optimizers (e.g., Sophia, INNA, etc.), which could be relevant in the context of adaptive second-order methods. Including these in the discussion would provide a more comprehensive overview of related work. 

4.	**Table 3:** I was surprised by the Top-1 accuracy of Adam for ResNet50 on ImageNet. Could you clarify which hyperparameters were used for Adam in this experiment?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730463619822

thread_type:
official_review

<thread_object>
creator: reviewer_jye4

<content>
summary:
This  paper  proposes  a new  second-order  optimizer,  AdaFisher,  which  ensures  better  convergence  while  maintaining  computational  efficiency.  Based  on  the  K-FAC optimizer,  AdaFisher  discovers  that  the  Kronecker  factor  is  diagonally  dominant  and  the  Fisher  information  matrix can be approximated  by  using  a diagonal block Kronecker approximation.  On image  classification  and  language  modeling  tasks,  AdaFisher  achieves  better  results  than  other  second-order  optimization  methods when applying Wall-Clock-Time method (training  for  the  same  amount  of  time).

strengths:
The  paper  innovatively  discovers  that  the  Kronecker  factor  is  diagonally  dominant  and  proposes  a diagonal  block-Kronecker  approximation  for  the  FIM. The  resulting  AdaFisher  optimizer  shows  good  performance.  And the  paper  is  well-organized  and  easy to follow.

weaknesses:
Please  refer  to  Questions.

questions:
- In Section  2,  it  is  stated  that  $F_i = \cal{H_{i-1}} \otimes \cal{S_{i}}$, but  in  the  algorithm  implementation,  it  becomes  $F_i = \cal{H_{i}} \otimes \cal{S_{i}}$. Why  is  this  the  case?
- Why are two parameters, $\gamma_1$ and $\gamma_2$, introduced for calculating the exponential moving averages of $\cal{H}$ and $\cal{S}$ in the AdaFisher optimizer? Would it not suffice to introduce only one parameter?
- It is known that the SGD optimizer performs better than the Adam optimizer in image classification tasks. Why is there no comparison with SGD?
- The experimental section compares performance under the same computation time. Could you provide a performance comparison under the same number of epochs? I suspect that second-order optimizers like AdaHessian may not have converged yet. When all optimizers have converged, does AdaFisher still outperform other second-order optimizers?
- The AdaFisher optimizer introduces two parameters, $\gamma_1$ and $\gamma_2$. I see that the authors conducted extensive searches for these parameters. Could you provide the values of $\gamma_1$ and $\gamma_2$ used in different tasks for the AdaFisher optimizer? If these parameters also require extensive searching to achieve good performance, then the AdaFisher optimizer might be very limited.

I'm willing to improve my score if you address my concerns.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730507632413

thread_type:
official_review

<thread_object>
creator: reviewer_5wtx

<content>
summary:
The authors use the Fisher information matrix as a precondition matrix to obtain a second order optimizer. They also provided some analysis and numerical results for the method. 

My concern is that this method is too similar to a previous method "Kronecker-Factored Second-Order Optimizers Perform First-Order Descent on Neurons" by Frederik Benzing.

strengths:
Using the Fisher imformation matrix as a precondition is really a good point.

weaknesses:
It seems that the author didn't realize a previous work "Kronecker-Factored Second-Order Optimizers Perform First-Order Descent on Neurons" by Frederik Benzing, which is very similar to this work.

questions:
1. The authors should specify the difference of the current work to the previous work "Kronecker-Factored Second-Order Optimizers Perform First-Order Descent on Neurons" by Frederik Benzing.

2. Empirically, the high-order optimizers are computationally more expensive than first order optimizers for each step and also have poorer generalizability. It is somehow strange that the numerical results in this work didn't show how the method accelerate the training process clearly but mainly about the strength of the method in generalization. The authors should explain why the method can help in enhance the generalization significantly. They should also include results obtained with SGD, since usually the generalization ability  of SGD is better than ADAM.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730704762795

thread_type:
official_review

<thread_object>
creator: reviewer_jcna

<content>
summary:
The paper proposes a new idea: using Kronecker factored preconditioners, with diagonal factors. The Kronecker product approximates empirical Fisher information matrix. The method demonstrates better performance compared to diagonal approaches such as Adam.

strengths:
1. Kronecker factored preconditioner with diagonal factors hasn't been used with empirical Fisher matrix before. 
2. Thorough experimentation.

weaknesses:
1. The preconditioner still requires layer inputs and gradients backpropaged through the layer, which is not always feasible for large training systems. 
2. The Adafactor already uses Kronecker factored preconditioner with diagonal factors. There is no comparison against Adafactor. 
3. sub-optimal regret bound $O(\log(T)\sqrt{T})$, compared to Shampoo - which is optimal $O(\sqrt{T})$.
4. There are low rank approaches with similar complexity as proposed methods  such as EVA [1]. There should be comparison against EVA.

If 2, 3, and 4 are addressed, I am willing to increase my score.

[1] Zhang, Lin, Shaohuai Shi, and Bo Li. "Eva: Practical Second-order Optimization with Kronecker-vectorized Approximation." The Eleventh International Conference on Learning Representations. 2022.

questions:
Does your shampoo baseline use grafting? 
CASPR is another paper which uses Kronecker sum based approach, what are your thoughts on Kronecker sum based combination to approximate empirical Fisher information matrix.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732064969692

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Dear Reviewers (rebuttal)

comment:
We would like to thank all reviewers for their time and valuable feedback when reviewing our paper. We appreciate their constructive criticisms that aided us in formulating our responses which will serve to improve the overall quality of our paper. We have received all comments and addressed them all in the revised manuscript as well the official comment section for each reviewer. Furthermore, we encourage all the reviewers to read official comments from our reply to the other reviewers to follow up more on the detailed questions. For the ease of the reviewers, we have highlighted the revised manuscript in blue for ease of visual navigation. 

Below, we summarize the strengths of our paper based on the reviewers' evaluations.

### Novel Methodology
- **`jcNA`**: *"Kronecker factored preconditioner with diagonal factors hasn't been used with empirical Fisher matrix before."*
- **`jye4`**: *"The paper innovatively discovers that the Kronecker factor is diagonally dominant and proposes a diagonal block-Kronecker approximation for the FIM."*
- **`K3HY`**: *"The introduction of a diagonal block-Kronecker approximation applicable to various layers offers an interesting balance between computational efficiency and the use of curvature information typically employed in second-order optimization."*

### Theoretical Contributions
- **`K3HY`**: *"The paper presents a rigorous theoretical convergence analysis for both convex and non-convex cases, asserting a convergence rate of [..], similar to Adam-type methods. The derivation of the update rules using the diagonal approximation of the FIM is clearly explained."*

### Empirical Performance of the Proposed Method
- **`K3HY`**: *"Comprehensive experimental evidence showing AdaFisher outperforms baseline methods."*
- **`jcNA`**: *"The method demonstrates better performance compared to diagonal approaches such as Adam."*
- **`jye4`**: *"On image classification and language modeling tasks, AdaFisher achieves better results than other second-order optimization methods when applying Wall-Clock-Time method (training for the same amount of time)."*
- **`jye4`**: *"The resulting AdaFisher optimizer shows good performance."*
- **`K3HY`**: *"Comprehensive experimental evidence showing that AdaFisher outperforms baseline methods (Adam, K-FAC, AdaHessian, etc.) on benchmark datasets such as CIFAR-10, CIFAR-100, ImageNet, and WikiText-2 across different network architectures."*
- **`jcNA`**: *"Thorough experimentation."*

### Stability
- **`K3HY`**: *"AdaFisher shows strong stability across varying learning rates and batch sizes, reducing the need for extensive hyperparameter tuning, which is a common challenge when training deep models."*

### Presentation
- **`jye4`**: *"The paper is well-organized and easy to follow."*

---
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes AdaFisher, a new adaptive second-order optimization method, which offers a strong balance between convergence efficiency and computational feasibility. The authors introduce a novel diagonal block-Kronecker approximation to the Fisher Information Matrix (FIM) for preconditioning gradients, making second-order methods more practical for deep neural network training. One of the key strengths is the robustness of AdaFisher across different tasks, including image classification and language modeling, as well as its ability to outperform state-of-the-art optimizers like Adam and K-FAC in terms of accuracy and convergence speed. The paper also provides a solid theoretical foundation, demonstrating convergence analysis for both convex and non-convex cases, which is a valuable contribution to the field of second-order optimization.

The empirical results are thorough and demonstrate clear improvements over baseline optimizers. The experiments are well-structured, showing that AdaFisher achieves better generalization and stability across a range of hyperparameters, which reduces the effort needed for tuning. Notably, the method is shown to perform well under both wall-clock time comparisons and fixed-epoch evaluations. Additionally, the authors have addressed concerns during the discussion phase effectively, clarifying the novelty compared to related works and providing further comparisons with optimizers like EVA and AdaFactor. Overall, this work presents meaningful advancements in optimization strategies for deep learning.