AdaFisher: Adaptive Second Order Optimization via Fisher Information
Authors: Damien MARTINS GOMES
Yanlei Zhang
Eugene Belilovsky
Guy Wolf
Mahdi S. Hosseini
Venue: ICLR 2025 Poster (2025)
Keywords: Second Order Optimization
Fisher Information
Kronecker-factored Approximate Curvature
Deep Learning
Computer Vision
Natural Language Processing
Primary area: optimization
TL;DR: This manuscript introduces a robust second-order optimization algorithm for training deep neural networks by approximating Fisher Information from deep layers and incorporating it into an adaptive optimization framework.

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.

Published as a conference paper at ICLR 2025
ADAFISHER: ADAPTIVE SECOND ORDER OPTIMIZA- TION VIA FISHER INFORMATION
Damien Martins Gomes*
Yanlei Zhang
Concordia University and IPSA Toulouse damien. rt nagomes @mail concordia.co
Université de Montréal and Mila
yanlei zhangêmi la quebec
Eugene Belilovsky
Guy Wolf
Concordia University and Mila
Université de Montréal and Mila wo fgué@mi la. quebec
eugene be welllllovkyeee www.oowwwoooooorri.....oooo
Mahdi Hosseini!
Concordia University and Mila mahdi.hoseen poosseiniicoocordii..
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 lever- ages diagonal Bron-KKoneeker approximation of the Fisher information matrix for adaptive gradient preconditioning. AdaFisher aims bridge the gap between enhanced imvergeenceggeeenaliaatoin 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 ro- bustness in hyper-parameter tuning. We demonstrate that AdaFisher outperforms the SOTA optimizers in terms both accuracy and convergence speed. Code is available from https//ggthabcccm///mllsssaaaaaa
INTRODUCTION
Deep Neural Network (DNN) optimization often struggles with the challenge of generalizing across varied architectures and complex data distributions. Current methods such as Adam op- timizer (Kingma & Ba 2015) and its variants (AdamP (Heo et al.) (021). AdaInject (Dubey (et al., 2022), AdaBelief (Zhuang et al., 2020) and YOGI Zaheer et al. (2018)) require exten- sive Hlper-Paramett (HP) tuming and often fail to generalize efficiently. DNN training typically minimizes a highly non-convex loss function C(O), updating parameters 0 using the expression ((C()))-VV(((")) at time step t, where g(t) represents the curvature information. Here, is the identity matrix for first-order optimizations such SGD (Kiefer & Wolfowitz)) 1952),
g(++1)- = (t))
and Hessian or Fisher Information Matrix (FIM) for the second-order case (Amari & Nagaoka
2000 The Hessian matrix interfaces with the deterministic Newton-Rand Raphoon method (Holmgren
1996 whereas the FIM harmonizes with the statistical measure of the Natural Gradient Descent
(NGD) approach (Amar Nagaoka 2000). This curvature information crucially optimizes the
gradient's preconditionin by accurately rescaling and orienting it. This adjustment significantly
"To my father and grandmother, whose strength and love continue to inspire me, this work is dedicated. 1 "Correponndigg Author

---

Published as a conference paper at ICLR 2025
accelerates convergence by ensuring more direct progress towards minima, thus enhancing training efficiency and reducing the number of required iterations (Kashyap) 2022).
As mentioned, the second-order methods employ curvature matrices for & to enhance the optimiza- tion process. Although these matrices accelerate convergence by effectively navigating through saddle points and swiftly moving towards minima (Foret al. 202T), they require higher compu- tational resources for the inverse computation. fact, when the number of learnable parameters increases, the curse of dimensionality associated with curvature matrix G makes the entire training process completely intractable by a commodity hardware platform.
Noteworthy approaches, such Adagrad (Duchi al. 201T), Adadelta (Zeiler 2012 RMSProp (Hinton et al., 2012), and Adam family utilize a simple diagonal approximation of the empirical FIM, which often results in convergence suboptimal local minima and poor generalization (Wilson et al.
2017 Luo et al. 2019). Advanced methods like AdaHessian (Yao et al. 2021) and Shampoo
(Gupta al. 2018) improve this by integrating structured matrices such as the diagonal Hessian or tensor-based preconditioners to enhance optimization. However, these second-order approaches, including K-FAC (Martens & Grosse 2020 [Eschenhagen et al. 2024), still face challenges high computational demands, lacking generalization and needing extensive HP tuning when applied to large-scale models (Ma et al. 2020).
Weight Trajectories Loss Landscape
To address these challenges, present AdaFisher, adaptive second-order optimizer, as an innovative solu- tion to address the generalization challenges raised in training DNNs. Substituting the second moment of Adam by novel diagonal block-Kronecker approximation of the FIM, AdaFisher strikes a balance between simplicity and generalization by introducing an extra HP compared to Adam but fewer than K-FAC, AdaHessian or Shampoo. With memory and time requirements on par with first- order methods, AdaFisher remains a practical choice for achieving effective generalization DNN optimization. As illustrated in Figure AdaFisher not only converges more rapidly but also reaches a superior local minimum by effectively navigating through saddle points compared to its counterparts. Further details regarding the visual- ization can be found in Appendix C
Figure 1: Visualizing optimization trajec- tories for various optimizers overlaid loss landscape.
In summary, our contributions build upon these findings as follows: [C1] We empirically showcase the energy of the Kronecker Factors (KF) is mainly concentrated along the diagonal and provide fresh insights of FIM in optimization; [C2] We introduce a diagonal -rock-Kroeeeker approxima- tion of the FIM applicable various layers, including normalization layers, enhancing model adapt- ability: [C3] We demonstrate AdaFisher's robustness and stability across diverse settings, proving its effectiveness; [C4] We showcase AdaFisher's s empirical performance against SOTA optimizers in image classification and language modeling, highlighting its superior performance; [C5] We develop a new technique that visualizes trajectories across different optimizers for better understanding of model behavior in the loss landscape. Additionally, we introduce an explainable FIM measure from AdaFisher, enabling comparative analysis of optimizer behavior.
2 BACKGROUND
We consider supervised learning framework with a dataset D containing N i.i.d samples, D = (xn; yn ) where In IRd and yn E CRC Let fo Rd RCC be a L--layer neural network parametrized by 0 where 0, concat(W,, b.) E P. and P, = pout x (Pin + 1). Let L RC x C + IR be the loss function defined by the negative log-likelihood, i.e. Lly. fo(x)) = log po (y) where po(y)<>) is the likelihood of the neural network fo. The network computes its output hr. Sa (x:) according
a; 0,h4-1, h4 di(a:), Vie (L...., ho In.
where h4-1 = [4_,- T]_ E R 1 and Φ; an element-wiss wise nonlinearity applied at layer For given input target pair (x, y), the gradient of the loss L(y, fo(x:) concerning the weights
2

---

Published as a conference paper at ICLR 2025
computed by the backpropagation algorithm (Lecun 2001). For convenience, adopt the special symbol 8; = Va, L for the pre-activation derivative. Starting from Vha C = Ür. (yy h2), we perform
8 V.=VV.-VV, (a.), s V (V. 1}.
where O denotes the element-wise product. Finally, the gradient V&C is retrieved by: VeC [vec Vo, C'', v ..... vec(Vo, [ " vec(-) denotes the Kronecker vectorization oper- ator which stacks the columns problem of finding the parameter set 0 that maximizes the likelihood, equivalently, minimizes the negative log-likelihood of the observed data. This Maximum Likelihood Estimation can be ex- pressed as an unconstrained composite optimization problem: ming J(0) = > L(yn, fo(xn)) )), where J(0) denotes the objective function, corresponding to the negative log- egglikelihood of the data. For notation convenience, we define Ji = (0). The FIM, utilized in lieu of the Hessian for Newton-Raphson' 's method, approximates the curvature of the log-likelihood function (Amari 1998).
matrix into vector. Optimization of a DNN can be recast as
F
E By-plsllf(rr.) Vo log (o Vo Vo log E VoJ((o)) = Elgg 1, (1)
where F measures the expected information that an observable conveys about the parameter 0.
For brevity, we write E instead of Ey--puulfg(z..); The K-FAC approach further simplifies FIM cal- culation using a block-diagonal approximation DNNs, known as Empirical FIM (EFIM), denoted by F. In Eq. (1) F is construed block Fag is articulated by Fag = [E]vec(g))vec(()
block matrix with dimensions I, x L, where each (i,j)th From the Kronecker-vectorivatte
equality vec we v u. we express vec(g:) (g:) as h;- @ (Petersen & Pedersen 2008), where gi is defined 5 1- By segmenting the FIM into discrete layer-specifie blocks, systematic factorization of
each block yields as
Pr.j )evc())))e vec(()) E(h 1hJ-1 @ 818]] B(h,-1 [[sss]]],
where i,j span the layer indices from to L. Given the dimensionality
h4-1 1 and 8; E R the Fisher matrix dimension will be Fig E R.P.P, Initially, K-FAC estimates the ex- pectation of the Kronecker product under the pre- sumption that activations and pre-activation deriva- tives are mutually independent. This can be rep- resented the Kronecker product of the individ- ual expectations Peg Hi-15-1 Si.j. where H4-1,5-1 E 1h and Sig R(sts]] denoting the KFs. The assumption for the block- diagonal structure posits that weight derivatives across distinct layers are uncorrelated, expressed by diag .o...........E Fr...) .iag, Figure illustrates the EFIM computation via K-FAC for given layer i.
Backward Propagation
dL
h4-1
0
L
da;
[h
aj
hi
Forward Propagation
F: = E
[
X E[s
s
Hi-1
S
Figure 2: Illustration EFIM computation using
K-FAC for a given layer i.
3 METHODOLOGY
Our methodology consists of four primary components: (i) Analyzing the KFs' structure in Sec- tion 3.1 and showcase their diagonal dominance; (ii) Introducing novel approximation of the FIM that retains only the diagonals of the KFs, detailed in Section 3.2 (iii) Incorporating the diagonal FIM approximation within the adaptive optimization framework as an alternative to the conven- tional second moment used in Adam, described in Section 3.3] (iv) Providing theoretical proof of
AdaFisher's convergence under both convex and non-convex conditions in Section 3.4
3.1 DIAGONAL CONCENTRATION KFs
Inspired by Gershgorin circle theorem (Hom & Johnson 2012), empirically conclude the KFs are diagonally concentrated by studying their eigenvalue distribution and perturbation under Gaus- sian noise. For demonstration, focus the eigenvalue spectrum of weight matrices from the
3

---

Published as a conference paper at ICLR 2025
37th layer of ResNet-18 (He al. 2016) after training for 50 epochs CIFAR-10 (Krizhevsky let al. 2009). As illustrated in Figure the eigenvalues (denoted as red crosses) predominantly clus- ter within the Gershgorin discs, which are centered along the matrix's diagonal elements (denoted as black circles), signifying substantial diagonal dominance. This phenomenon is quantitatively sup- ported by the Gershgorin eircle theorem, which posits that every eigenvalue of complex square matrix A lies within at least one of the Gershgorin discs (aa R.), where Rt = Eij şt |a:yl rep- resents the radius computed as the sum of the absolute values of the off-diagonal entries of the ith row. Next, we add Gaussian noise (O0 )²), matrix M is then expressed as M A + E, where & [et}] and Cij NN(,, o?) for # j. The noise perturbation on KF eigen- values is critical to comprehend the dynamics and stability of the ma- trix. This demonstrates the introduc- tion of noise to the off-diagonal ele- ments yields minimal eigenvalue per- turbations, particularly those surpass- ing the Kaiser criterion (i.e. most dominant eigenvalues), which re- main virtually unchanged (Braeken & Van Assen 2017). Both the above understandings from Gershgorin disc analysis and eigenvalue perturbation corroborate the robustness of the ma- trix's diagonal dominance. Exten- sive discussions of these analyses, in- cluding the KFs Fourier analysis, are available in Appendix
10-3 on off-diagonal elements. The perturbed
1 Slap 2000
Mistta Mapp 8830
PRACTA SlAp 5200
Step
Vra take
Figure 3: Gershgorin disks and eigenvalue perturbations from the 37th Convolutional Layer of ResNet-18 steps 5200 (middle of training) and 9800 (end of training). Left:: Gershgorin circles; Right: Eigenvalue spectrum w/w-o noise.
3.2 EFFICIENT COMPUTATION OF THE FIM
In the realm of optimization, NGD offers geometrically nuanced adaptation of the classical steepest descent approach (in Euclidean space), transitioning the focus from parameter space to the model's distribution space underpinned by the adoption of Riemannian metric, (Amari & Nagaoka (2000). The formulation of the preconditioned gradient (()) given the NGD method is articulated as
g(+) (r())-1g))) (0
where F(t) denotes the FIM at time step t distinguished from F,, the FIM at layer i. One of the distinguishing features of NGD within this framework is its re- paraamerrizzaion invariance, a direct consequence of leveraging the model's distribution properties rather than its parameters. Neverthe- less, the direct FIM computation is highly demanding, and we solve this by adopting the diagonal approximation of the KFs as supported by our analyses from Section 3.1 In addition, a critical com- ponent of modern DNNs is known to be the normalization of the layers by introducing scale and shift parameters (e-g. nommaliannnnnnnoonn (Toffe 2015), layer-normalization (Lei Ba et al., 2016))- This is to adjust the network's s dynamics (e.g., reducing covariance shift) in a non-trivial way Huang al. (2023) where the lack of FIM approximation on such normalization layers can lead subop- timal preconditioning. Therefore, we introduce a method for calculating the KFss for normalization layers detailed in Proposition 3.T
Proposition 3.1 (EFIM for normalization layer). Let (v,, B.) RCC. be the scale and shift parame- ters of a normalization layer i. The empirical KFs for the FIM approximation are
H--1 =
Σ ZET,
h h Ho-1/ = 111, S,
1
Σ RET:
$ S,,
T
T
T
where h4-1,8; RCxXITI represent the pre-normalized activations and gradients, respectively. Here, To is the set of dimensions over which normalization statistics are computed, and C₁ is the channelsffeutures size.
The proof of this proposition and the extended computation of other type of layers given in Appendix A.2 and Section timization, the KFs for given layer can be estimated using Exponentially Moving Average
A.3 respectively. Note that in the context of online and stochastic op-
4

---

Published as a conference paper at ICLR 2025
(EMA) scheme across batches defined by
- = - + +1- 5."0 =15 + + (1-1))((-1)
(3)
where 0 < S 1 is the exponential decay factor step time t, H and S:a) in the right- side are new KFs calculated during each forward and backward pass computation. This EMA scheme is commonly used in methods involving diagonal
dlocoadiagonal approximations to the curvature matrix (e.g. [eCunn et al. (2012): : Park et al. (2000); |Schaul et al. (2013)). Such schemes have the desirable property that they allow the curvature estimation depend on much more data than what can be reasonably processed in single batihb
suggests that the FIM's s critical information predominantly resides along its diagonal. Building upon this, we propose novel approximation for the FIM, described in Propo- sition 3.2 that conceptualizes the KFs as diagonal matrices denoted as FD. for layer i.
Our study from Section 3.1
Proposition 3.2 (Efficient EFIM). Assume that Hi-1 and S, can be closely approximated by diag- onal matrices, denoted by HD.,1 and SD, respectively layer i, such that HDD,1- Diag(H; (H:-1). Sp. (Sa where Diag denote the diagonal of a matrix. Accordingly, the Empirical FIM is defined by
FD. 2 H'D & StD, + A1,
(4)
where HiD,_, and S,b, denote the Min-Max normalization of HD and Sp, Patro & Sahu 2015) and X is regularization parameter.
The proof of this proposition is given in Appendix A.2 This approximation strikes balance be- tween computational time and space complexity and the accuracy of performance, as discussed in Section 4 We set the regularization parameter À 0.001, which acts damping factor following the Tikhonov regularization principle Martens & Grosse (2015), enhancing computational stability and conditioning of the FIM. The closed-for -oom solution for the preconditioned gradient g(*) is derived
from the diagonal approximation of the FIM, given by g( (F-1)) - for time step t. This represents the AdaFisher augmented gradient and incorporates local loss curvature information. focuses the diagonal elements to reduce computational overhead while maintaining reasonable FIM approximation. This simplification enhances the efficiency of the optimization process, which is crucial for training DNNs where computational resources are limited.
Table I: Summary of the first, second moments, regret bound and Applicability used Adam Kingma & Ba (2015), AdaHessian Yao et al. (2021). K-FAC Martens & Grosse (2020). Shampoo Cupta et al. (2018). AdaFisher for updating model parameters g(t+T) = g( am (2) /V Here 33 and 8o are first and second moment HPs. I (2) and rº(*) refer to the preconditioning method used by Shampoo Gupta al. 2018
(d) ())(G(()), and T denotes the total number of steps. Note that Transf. denotes Transformers.
Optimizer
m!
Regret Bound
Applicability CNNS Transf.
Adam
(1-8) Er 8'-19
(1 Re) f a
1/2
OlloTTV)
1/2
AdaHessian
1-81) 8
(1-82) y
8
D
O(logTVT)
1-81
K-FAC
(f(10)-110) G(0
O(vT)
Shampoo
O(vT)
x
Adalisher
(1-4) 24 ,0
pile
0llogTvi
V
3.3 INTEGRATING FIM INTO ADAPTIVE OPTIMIZATION FRAMEWORK
Following the spirit of the adaptive optimization framework from Adam, which combines mo- mentum and RMSProp principles (Sutskever al., 2013). the parameters are updated via
g(++1) = g m9 first and second moment estimates, respectively, for time step Although Adam is widely
अरं.. Here, represents the learning rate, while ((() and n(() denote the
used, its approximation of the second moment using simple diagonal elements of second- order statisties through squared gradients can mirror stability challenges (Kunstner et al. 2019)
5

---

Published as a conference paper at ICLR 2025
We overcome these challenges by utilizing more refined di- agonal Brooc-Krronckerr approximation of the FIM introduced in Section 3 more precise approximation of the Hessian from Taylor series expansion viewpoint. This structured al- ternative to Adam's diagonal approximation enables precise curvature estimation, mitigating stability issues and improving convergence in non-convex settings. AdaFisher distinguishes itself from Adam by incorporating higher fidelity approxima- tion of the FIM, enhancing both optimization efficiency and model robustness in complex scenarios. As demonstrated in Figure AdaFisher's FIM values exhibit narrow variations and lower mean values during training, suggesting a conver- gence towards flatter minima and an implicit regulariza- tion effect. In contrast, Adam shows broader variations, in-
Histogram of Diagonal for
Frequencies
Histogrant am Bins
400
200
25p
dicating less efficient generalization (Wang et al., (024), for more details regarding the convergence behavior refer to Ap- pendix B.I] Moreover, AdaFishhr omits the square root and the traditional EMA applied over the second moment since the FIM naturally incorporates an EMA of its KFs (detailed in Eq.
Histogram of FIM Diagonal for AdaFisher
Rrequencies
B). ) The exclusion of the square root aligns with the theoreti-
cal definition of second order methods. as using square root deviates from the second-order Taylor expansion approxima- tion that these methods aim to follow. A comparative sum- mary of different moment estimates, (*) and u(9,, along with their regret bounds and applicability across various optimiz- ers, is presented in Table []] Building on the principles of AdamW (Loshchilov Hutter 2019b), which modifies Adam by integrating weight decay directly into the weight update step to counteract suboptimal decay behaviors and boost op- timization performance, we introduce AdaFisherW. This vari- ant adapts the Adam framework to further enhance the op- timizer by leveraging curvature information from AdaFisher. Finally, AdaFisher is compatible with multi-GPU environ- ments, with distributed version detailed in Appendix A..4 The implementation for both AdaFisher variants is delineated in the pscudo-code presented in Algorithm
Histogrant am Bins
Training
fraining Iteratioss
200
200
Figure 4: Comparison of FIM diago- nal histograms during ResNet18 train- ing on CIFAR10: The figure displays the FIM diagonal elements for the first convolutional layer with Adam and AdaFisher over 1,000 training itera- tions.
3.4 CONVERGENCE ANALYSIS
In this section, we provide a theoretical analysis of AdaFisher's convergence in both convex opti- mization and non-convex stochastic optimization. We first present a standard convergence behavior of Eq. (2) for simple strongly convex and strictly smooth function f(J).
Proposition 3.3 (Convergence in convex optimization). For FIM defined in Eq. { the updating
scheme g(++1)= g(t) at V converges. Further, if V J is Lipschitz, the convergence rate is bounded.
For non non-convex cases, we adopt the similar derivations of et al. 20199 since AdaFisher belongs the family of generalized Adam-tyre Stam-type methods.
Proposition 3.4 (Convergence in non-convex stochastic optimization). Under the assumptions:
(i) J is lower bounded and differentiable; [|V (O))l Ll10 0'1
|1 1l
L, 0,0' (ii) Both the true and stochastic gradient are bounded, i.e. [|V)(((())]ll
and
( la X. Vt for some X > 0, (ii) Unbiased and independent noise in git), iec g(t)
VJ(((t)) + ((t), E[c(0]] 0, and (()) 1 ((,), Vi # Assume 17(t) v 300) I B < 1is
non-increasing, F p
Vt [T],5 [d]. we then have
L
min Ell
|
(C + log T) + C2dy + Czdy² + C4)
VT

---

Published as a conference paper at ICLR 2025
where C1 C2 C3 constants independent of d and T, C4 is constant independent of T, the expectation is taken with respect to all the randomness corresponding to {g()).
The proofs of propositions 3.3 and 3.4 are given in Appendix A.2
3.4 implies the convergence rate for AdaFisher in the non-convex case at (log T/VT), which is similar to Adam-type optimizers. While DNNs often include non-smooth components like ReLU and max pooling, which create non-differentiable points in the loss landscape, optimizers like AdaFisher handle these cases effectively, as shown by our results in Section 4
Proposition
Algorithm AdaFisher optimization algorithm. Good default settings for the tested machine lcarning prob- lems are a 0.001 (Icarning rate), À 0.001 (Tikhonov damping parameter),7 = 0.8 (Exponentially
decaying factor). [Default parameters are: 3 0.9 (Exponentially decaying factor of Adam), (weight decay)
(Kingma & Ba 20155 Loshchilov & Hutter (011bb)].
Require:
Step size as Exponential decay rate for KFs 7 E [0, 1); Tikhonov damping parameter A; Exponential
decay rate for first moments in [0, 1); Initial parameters 0
Initialize 1st moment variable = 0; FIM FD, = I; time step = 0
I: while stopping criterion not met do
Sample minibatch of M examples from the training set {(Zn, Uo.) Jn-
]
3: Compute HD..1. for i (1,..., L}} using Section A.3 (notice that: HDg = 2)
4: Compute EMAS of HD.,, and SD, using Eq. 3
5: Compute FD; for i ( {L...., L using Eq.
6: (0 En V Voco ((f(mm; )00,, yo) (Compute gradient)
7: m (0+1) +(1 (Update and correct biased first moment) 1-gt
8: Case AdaFisher: 20 -((( -2m
Case AdaFisherW: Ag(t) = -a (F( + nol())
9: ((++1) + g00) + AO (Apply update)
10: t 1
11: end while
4 RESULTS
To evaluate AdaFisher, we conduct experiments on six benchmark datasets across Image Classifi- cation for Computer Vision (CV) and Language Modeling for Natural Language Processing (NIP) that are commonly used to evaluate optimization algorithms: CIFAR-10. CIFARIO0 (Krizhevsky al. 2009). Tiny ImageNet (Le & Yang. 2013), and ImageNet-Ik 1k (Deng et al. 2009) for image classification; Wikitext-2 (Merity et al. (017) and Penn Treebank (PTB) (Mareus et al. 1993) for language modeling. The six baseline methods we compare with are SGD, Adam/ W. K-FAC, AdaHessian, and Shampoo. For CIFAR experiments, we report the average over five runs. We also perform transfer learning task using the ImageNet-Ik weights from Paszke et al. (2019). De- tailed descriptions of the experimental setup (including HP tuning, datasets, and data augmentation). results, and analyses are provided in Appendix D
Table 2: Performance metrics (mean, std) of different networks and optimizers on CIFARIO and CIFAR100 using batch size 256 with 200-epoch AdaFisher training cutoff.
CIFAR10
CIFAR100
Network
SGD
Adam
[ddaHsssaad K-PAC [Shampoo|
Adam
[AdaHessian
K-FAC /Shampoo/Adaaishhee AdaFisher
ResNet18
95.64a.1 99. 85..
95.44..
85.17a 27.o299..88. 91.08...96. 25
76.56o.27..74.
71.79
76. 37..776.78.. 77 .8.
ResNe150
95.7lo.194.45,
95.54a.1
05.66o
94.59.. 96.34..2
78.01o. 74.65a.5
7..81
77.40
79.770..
ResNer101
95.98..294.5fn
95.29..
36.01o.
94 03..
96. 16.
78. 89.890.
73.38o 2
17.01o.488.88.
80 65
DansoNet11196.09,0
99..099..911.6,
96.11a.1
96.12p
96. 72..1
75.87a.
74.80a.s
79.79..28.221...
(o 81.36..3
MobileNeV3344.49. 44. 293.32o
92.86s.
94.34p
33.816. 25. 28
73.899.370.622
56.581.5
73.75..3
77.56p.
Tiny Swin 82.3 82.34. 37..
84.
81.796.
63.91o.6......
48. 74..4
51.899..
56.86p.
31.45p.45000
30.39,
66 05..5
FocalNet
82.03o.2/86 23o.
64.18g.
38..44..337 37. 877.377667777....
87.90
47.7663
52.714.5
32.33a.
9.966
9.18p. 53.69p.
2/T-2/3x22 78. 75.776..883..8 (83.89..4
233.08g,/35.1111
7.17
8.60p. 62.91o.6
"Note that Adam and AdaFisher were used for all CNN architectures, while AdamI and AdaFisberW were applicd for VIT experiments.
7

---

Published as a conference paper at ICLR 2025
4.1 IMAGE CLASSIHICATION
We commence our analysis by assessing the convergence and generalization capabilities of various models on image classification tasks. Specifically, we deploy ResNet architectures (ResNetX where 2017). MobileNetV3 (Howard et al. 2019). Tiny Swin (Liu et al., 202T), FocalNet (Yang et al., 2022) and CCT-2/3 (Hassani et al. 2021) on CIFARIO and CIFAR 1OO, while utilizing standard ResNet50 for Tiny ImageNet and ImageNet-Ik, Ik. Our empirical evaluation of AdaFisher optimizer across these models and datasets illustrates its efficiency in optimizing image classification, surpassing all SOTA optimizers. We employ the Wall-Clock- -ime (WCT) method with a cutoff of 200 epochs for AdaFisher's 's training, except for ImageNet- 1k, where we use 90- epoch WCT for Adam, which surprisingly matched AdaFisher's ' training duration. Results confirm AdaFisher's superior classification accuracy on both CNNs and ViTs. Note that the results for Tiny ImageNet are described in Appendix D.2.4
X E {18, 50, 101 D. DenseNet121 (Huang et al.
The performance outcomes for CIFAR datasets are detailed in Table 2
ImageNet-Ik Training. Training on ImageNet-Ik typically requires multiple GPUs and large batch sizes. Our study showcases that AdaFisher achieves superior validation accuracy on single GPU than its counterparts in scenarios marked by the light blue region. This performance outstrips traditional ap- proaches like SGD, LAMB (You et al.) 2020), and LARS Yoou et al.) 2017). which typically utilize batch sizes of 16K. While AdaFisher attains SOTA results on a single GPU, it further excels when scaled up in a distributed setting with larger batch sizes. The results, benchmarked using 256 batch size and a WCT of 90 Adam training epochs, are detailed in Table 3 and illustrated in Figure 5 Distributed AdaFisher curves are illustrated in Figure 15 The
Table 3: Validation of ImageNet- Ik ResNet50 by different optimizers reported on Top-1 I and Top-5 accuracy.
Optimizers
Batch size
Top-1
Top-5
Adam
256
67.78
88.37
K-FAC
256
70.96
89.44
Shampoo
256
72.82
91.42
AdaFiisher
256
76.95
93.39
Adalfisher
512
77.01
93.45
AdaFiiaerr
1021
77.09
93.56
SGD Goyal et (2017)
256
76.40
Adam Chtnn
(2024)
1024
76.34
LAMB You al. (0200)
16K
76.66
93.22
SGD You el al. 2020
16K
75.20
LARS Huo ct al. (202)
16K
75.1
Traisin less: Receetto lmagetet
Testing smer Ressette Imagaltet
light blue highlights in the table represent our exper- iments with a batch size of 256 on single GPU. The light green indicates results from a distributed version of AdaFisher employing larger batch sizes, whereas the orange reflects results from SOTA meth- ods using a higher batch size of 16K, SGD with a batch size of 256 and Adam W with a batch size of 1024. It is important to note, however, that the train- ing setups and augmentation techniques for the re- sults highlighted in orange, taken from the literature, may differ from those in our study. These results included to provide broader context and intu- ition regarding AdaFisher's 's performance compared to other experiments. Overall, by balancing curvature- updates with parameter efficiency, AdaFisher maintains strong generalization even under constrained computational budgets, a critical advantage over methods like K-FAC that struggle with over-parameterized models.
20000 stice
XB:00 10:00 310003 00:01
Figure 5: Training loss and validation error of ResNet-50 on ImageNet- AdaFisher consis- tently achieves lower test error as compared to its counterparts.
Table 4: Performance comparison of different networks and optimizers on CIFARIO and CIFAR100 using ImageNet- Ik pretrained weights. Evaluation is based on wall clock time of 50 training epochs with AdaFisher.
CIFARIO
CIFAR100
Network SGD
Adam AdaHessian K-FAC [Shampoo AdnFisher
SGD
Adam
AdaHession
K-FAC [Shampoo] AdaFisher
ResNet50
p6.50.
266.4..6
96.456. 96.366..
166
99. 13 |2. 2g.
(2.00
80.64o.a
80.55a.488 (81.70p 2 2/2.230.
ReaNett01 07.07o.266. 96.70....
96.65, 66..2
90.8 96.84o. (66.166663, 36.63..1 1 99.222..
84.01a. S2.43
81.36o.a
82 22.266.a/82. 05 84.47. 47
DenseNet1219.. 94.80 141.77.
93.08..1
91.41o.
94176p
95.03,
75.
71.06o.a
76.
78.92..3
MobilNeeV331.700
90.922
86.45
91.72
91.39
92.78.
66.86...6..111
11..s
59.69..3
69.85 68.877,.
(7...87..77...333 72.388.
4.2 TRANSFER LEARNING
Following a more sustainable practice of training DNNs, we employ pretrained models from ImageNet- 1k in PyTorch on datasets like CIFAR1O and CIFAR1 100 to showcase AdaFisher's gener-
8

---

Published as a conference paper at ICLR 2025
alization capability for transfer learning. We applied these pretrained weights across various CNN architectures to train on these datasets. The results, presented in Table highlight the significant ad- vantages of using AdaFisher, consistently achieving top accuracy across both datasets. More details can be found in Appendix D.2.3
4.3 LANGUAGE MODEL.
We employ the WikiText-2 dataset, which encompasses approxi- mately 100 million tokens derived from over million words ex- tracted from a curated set of "Good" and "Featured" articles on Wikipedia. Additionally, we utilize the PTB dataset, renowned for its extensive collection of English words with part-of-speech tags, which has been widely used in NLP tasks for training and bench- marking language models. Our experiments utilize a scaled-down version of GPT-1 (Radford et al. 2019), featuring four self-attention layers with masking capabilities with more than 28 million learn- able parameters. More details about tuning HPss and models be found Appendix D.3] The perplexity (PPL) on the test set, corresponding the hest-performing model during validation, is documented in Table 5) Símilar to approaches in image classification, we apply the WCT method
Table 5: Language Modeling performance (PPL) on Wikitext- and PTB test dataset (lower is better).
Optimizer
Text PPI.
WikiTest-2
FIB
AdamW
173.06
41.70
AdaHessian
407.69
59.43
Shampon
1727.75
AdalisherW
41.15
with 50 epochs training time of AdaFisher the cutoff period. Notice that Shampoo did not achieve
convergence despite using optimal HPs, and the K-FAC was unable to train with ASDL library (Os-
awa al. 2023).
(A)
B
(C)
(D)
00
AdaHessian
300
40
30
Batch 512
B4tch 1024
10
Batch Size
Adam
AdeRsher
Adalesssia
Kfac
Shampoo
Figure 6: Performance comparison of AdaFisher and other optimizers using the ResNet50 network the CIFAR 100 dataset. (A) Test accuracy by batch size. (B) Accuracy learning rates. (C) Accuracy related to epoch time across batch sizes. (D) Epoch time for different optimizers with batch size of 256.
4.4 STABILITY ANALYSIS
In this section, assess AdaFisher's stability under varying learning rates and batch sizes ing ResNet50 on CIFAR 100 and compare its performance to other optimizers. Improved stability indicates reduced need for HP tuning while maintaining high performance. To ensure a fair com- parison, all methods were evaluated using a consistent experimental setup, with parameters tailored to each optimizer's 's strengths. However, we exclude AdaHessian results for a batch size of 1024 due to its significant computation cost.
Batch Size Analysis. We examine the impact of batch size on AdaFisher's performance, as shown in Panels (A) and (C) of Figure 6 AdaFisher maintains high test accuracy across various batch sizes, excelling particularly at smaller sizes despite some sensitivity to larger ones. Panel (C) high- lights AdaFisher's 's efficiency, achieving high accuracy with shorter epoch times compared to Adam, detailed further in Panel (D), where AdaFisher shows competitive epoch durations against other op- timizers. These results, discussed in Appendix [D2.7 underscore AdaFisher's effective performance across batch size variations without adjusting other HPs.
Learning Rate Stability. This analysis evaluates the impact of learning rate variations AdaFisher's performance, as depicted in Panel (B) of Figure 6 AdaFisher demonstrates superior stability, particularly at lower learning rates, maintaining consistent performance across broad

---

Published as a conference paper at ICLR 2025
spectrum. This stability alleviates the need for meticulous learning rate adjustments, thereby stream- lining model training in various computational environments. Additionally, AdaFisher's 's stability across various learning rates can be attributed to its effective approximation of the curvature matrix.
Ablation Studies. We further conduct extensive ablation studies on additional components of novel approximation of the FIM, the signif- icance of EMA for Kronecker factors, the impact of the square root, the stability across learning rate schedulers and the updated computation of the FIM for normalization layers. These analyses are thoroughly detailed in Appendix B
AdaFisher, including the convergence efficiency,
RELATED WORK
Tractable Approximation of the FIM. Efficient approximations of the FIM for neural network optimization have evolved significantly. beginning with block-diagonal strategies exemplified by TONGA (Roux al. 2007) and extending to the Kronecker-factoree approaches like K-FAC. Fur- ther innovations have emerged, such as SK-FA -KFFAC (Tagg et al. ). EVA (Zhang et al. 2023), which accelerates the computation of the FIM, and |Eschenhagen
al. (202)), who propose a generalized framework for FIM computation that enhances preconditioning methods like Shampoo. More re- cently, Liu et al. (2024). Huang et al. (2024) and Duvvuri et al. (2024) have introduced tractable solutions for computing the FIM. AdaFisher distinguishes itself by integrating enhanced FIM com- putations with novel diagonal Kronecker factors, enriching the Adam optimization framework. This integration, outlined in Proposition 3.2 and detailed in Appendix A.3] advances the fusion of second- order optimization principles with first-order methods. This builds upon innovations like AdaHes- sian, which incorporates Hessian diagonals into the Adam framework.
Adaptive First-Order Methods. Building upon the diagonal approximation heritage of the FIM, AdaFisher extends traditional diagonally-scaled first-order methods such AdaGrad (Duchi al., 2011), AdamP, Adalnject, AdaBelief, and Adam. These methods have inspired advancements like AMSGrad (Reddi et al. 2018). AdaBound (Luo al.) 2019), RAdam (Liu et al. (020). and en- hanced AdaBelief, FOOF (Benzing, 2022), improving both theoretical rigor and practical effective- ness. A recent study by Diang al. (2024a) illustrates that first-order adaptive methods can bias training trajectories and effectively navigate through saddle points. In response, MMshhhenno & Stich (2023) propose an empirical solution involving the addition of noise mitigate these biases. Leplat et al. (2022) introduces a novel method accelerating convergence via Gauss-Seide type dis- cretization. AdaFisher differentiates itself by eliminating the conventional square root in the second moment calculation, with benefits underscored by Lin et al. (2024) and Malladi et al. (2022) in CNN architectures, and Zhang et al. (2024) demonstrates the critical role of Adam family optimizers in Transformer models. Its unique preconditioning. based on the Fisher Information, is elaborated in Algorithm
CONCLUSION, LIMITATIONS AND FUTURE RESEARCH
In this work, introduced AdaFisher-an adaptive optimizer that utilizes the Fisher Information Matrix (FIM) with a new diagonal improve descent directions. Integrated within the adaptive optimization framework, AdaFisher not only accelerates training but also minimizes the need for higher accuracy and stability in tasks like image classification and language modeling. Empirical and theoretical analyses confirm its superiority over existing optimizers, and its ease of implementation, along with modest space and time requirements, allows to adapt across various tasks. Notably, AdaFisher outperforms SOTA optimizers on ImageNet- 1k when trained using both single and multi- GPU configurations. AdaFisher is optimized for statistical learning tasks where the deep neural network outputs parameterize distributions within the exponential family, and the loss function is the negative likellhooo (or its derívatives). Outside these conditions, where the FIM no longer coincides with the Generalized Gauss-Newton matrix, the method may not capture the true curvature information of the loss landscape.
approximation to enhance gradient rescaling and
tuning. thereby achieving
Looking ahead, extensive testing across a wider range of models and domains, including generative modeling and graph neural networks, will further validate AdaFisher's 's effectiveness. Additionally, implementing CUDA kemnels for efficient computation of Kronecker factors could significantly en- hance AdaFisher's 's scalability and performance.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENT
This research is funded by Fonds de recherche du Québec (FRQNT) Masters (B1X) Scholarship [D.M.G]; Natural Sciences & Engineering Research Council (NSERC)-Discovery Grant RGPIN- 2022-05378 [M.S.H.J; FRONT-NSERC grant 2023-NOVA-329125 [E.B.& G.W.J.
IMPACT STATEMENT
AdaFisher represents significant advancement in training efficiency, achieving superior accuracy on the ImageNet- IK dataset using only a single GPU. This optimization is particularly beneficial for academia and students who may not have access to extensive computational resources. By enabling effective training with fewer GPUs, AdaFisher offers an accessible yet powerful solution, reduc- ing hardware costs and making advanced machine learning more attainable for those with limited resources. This capability underscores AdaFisher's potential as a valuable tool in democratizing machine learning technology.
REFERENCES
iris, 2018.
Shun-ichi Amari. Natural gradient works efficiently in learning. Neural Computation, 10(2):251-
276, Feb 1998. ISSN 0899-7667.
doi: 10.1 101116//88977003300011166
Shun-ichi Amari and Hiroshi Nagaoka. Methods of information geometry. 2000. URL https:
7/api semanticscholar htt.sccccoaaa.cccooooopaaaaaaaa. 1169776622111
Frederik Benzing. Gradient descent neurons and its link to approximate second-order optimiza- tion. In International Conference on Machine Learning, pp. 1817-1853, PMLR, 2022.
Stephen Boyd and Lieven Vandenberghe. Convex Optimization. Cambridge University Press, 2004.
Johan Bracken and Marcel ALM Van Assen. An empirical kaiser criterion. Psychological methods, 22(3):450, 2017.
Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Scunghyun Park, Yunsung Lee, and Sungrae Park. Swad: Domain generalization by seeking flat minima. Advances in Neural Infor- mation Processing Systems, 01122405-22118 2021.
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong. Thang Luong, Cho-Jui Hsieh, Yifeng Lu, et al. Symbolic discovery of optimization algorithms. Advances in neural information processing systems, 36, 2024.
Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong. On the convergence of a class of adam-type algorithms for non-convex optimization. In International Conference on Learning Representa- tions, 2019. URL C httttttt//pppeeeeeeesee
net forum? 2dw-hxxxx3=
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hier- archical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp. 248-255, 2009. doi: 10. 101100///P22000 200.5.688.0.
Li Deng. The mnist database of handwritten digit images for machine learning research [best of the web]. IEEE Signal Processing Magazine, 29(6):141-14,, 2012. doi: 10.1100// 2211477.
Terrance Vriei and Graham W Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint arXiv: 04552, 2017.
Shiy Ram Dubey, SH Shabbeer Basha, Satish Kumar Singh, and Bidyut Baran Chaudhuri. Adain- ject: Injection based adaptive gradient descent optimizers for convolutional neural networks. IEEE Transactions on Artificial Intelligence, 2022.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(61):2111-2159 2011. URL http:// //mp: r. wttprrr.rrg..p//p.eapeeeap 2/duchi html]

---

Published as a conference paper at ICLR 2025
Sai Surya Duvvuri, Fnu Devvrit, Rohan Anil, Cho-Jui Hsieh, and Inderjit S Dhillon. Combining axes preconditioners through kronecker upproximation for deep learning. The Twelfth International Conference Learning Representations, 2024.
Muhammad EINokrashy, Badr AlKhamissi, and Mona Diab. west attention (dwatt): A layer fusion method for data-efficient classification. arXiv preprint arXiv:2209.1516. 2022.
Runa Eschenhagen, Alexander Immer, Richard Tumnr, Frank Schneider, and Philipp Hennig. Kronecker--actoreed approximate curvature for modern neural network architectures. Advances in Neural Information Processing Systems, 36, 2024.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimiza- tion for efficiently improving generalization. In International Conference on Learning Represen- fations, 2021. URL htt.hhtttt///opotie
net/ forum?irooo d1 mposs rM
Karl Pearson F.R.S. Lii.. on lines and planes of closest to systems of points in space. The London, Edinburgh, and Dublin Philosophical Magazine and Journal of Science, 2(11):559-572, 1901. doi: 10. 10880000777440009497777777
Thomas George. NNGeometry: Easy and Fast Fisher Information Matrices and Neural Tangent Ker- nels in PyTorch, February 2021. URL
45325977
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint arXiv: 170.00277777 2017.
Roger Grosse and James Martens. A kronecker-factoonn approximate fisher matrix for convolution layers. International Conference on Machine Learning. PMLR, 2016.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor opti- mization. In International Conference on Machine Learning, pp. 1852-1850, PMLR, 2018.
Ali Hassani, Steven Walton, Nikhil Shah, Abulikemu Abuduweili, Jiachen Li, and Humphrey Shi. Escaping the big data paradigm with compact transformers. 2021. URL phttps://arde org/ abs/ 2104. 05704
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.
Byeongho Heo, Sanghyuk Chun, Seong Joon Oh, Dongyoon Han, Sangdoo Yun, Gyuwan Kim, Youngjung Uh, and Jung-Woo Ha. Adamp: Slowing down the slowdown for momentum optimiz- on scale-invariant weights. In International Conference Learning Representations, 2021.
URL
net
Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning lecture 6a overview of mini-batch gradient descent. 2012.
Richard A. Holmgren. Newton's 's Method, pp. 127-151. Springer New York, New York, NY, 1996. ISBN 978-1-4411--733277
doi: 10. 111177/978881111-877---7-71-7 UtL/tttps://do..orr 10. 10.1007// 9788 1178----411-872-----111111-
Roger A. Horn and Charles R. Johnson. Matrix Analysis. Cambridge University Press, 2 edition, 2012.
Andrew Howard, Mark Sandler, Grace Chu, Liang-Chieh Chen, Bo Chen, Mingxing Tan, Weijun Wang, Yukum Zhu, Ruoming Pang, Vijay Vasudevan, et al. Searching for mobilenetv3. In Pro- ceedings of the IEEE/CVF international conference on computer vision, pp. 1314-1324, 2019.
Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger. Densely connected convolutional networks. Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4700-4708, 2017.
12

---

Published as a conference paper at ICLR 2025
Keke Huang, Ruize Gao, Bogdan Cautis, and Xiaokui Xiao. Scalable continuous-ting diffusion framework for network inference and influence estimation. In The Web Conference 2024, 2024. URL httpttss//ppeeeeeiieetttee net/
Lei Huang et al. Normalization techniques in training dnns: Methodology, analysis and application. IEEE Transactions on Patter Analysis and Machine Intelligence, 2023, doi: 10.1 109. 2023.3250241.
Zhouyuan Huo, Bin Gu, and Heng Huang. Large batch optimization for deep learning using new complete layer-wise adaptive rate scaling. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pp. 7883-7890, 2021.
Sergey loffe. Batch normalization: Accelerating deep network training by reducing internal covari- ate shift. arXiv preprint urXiv:/502.03366) 2015.
Kaiqi Jiang, Dhruv Malik, and Yuanzhi Li. How does adaptive optimization impact local neural network geometry? Advances in Neural Information Processing Systems, 36, 2024a.
Zixuan Jiang, Jiaqi Gu, Hanqing Zhu, and David Pan. Pre-rmsnorm and pre-crmsnorn transformers: equivalent and efficient pre-In transformers. Advances in Neural Information Processing Systems, 36, 2024b.
Rohan V Kashyap. A survey of deep learning optimizers-fiss and second order methods. arXiv preprint arXiv:2211.15596, 2022.
Jack Kiefer and Jacob Wolfowitz. Stochastic estimation of the maximum of a regression function. The Annals of Mathematical Statistics, pp. 462-466, 1952.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, In International 142..
Conference Learning Representations, 2015. URL
6980
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Frederik Kunstner, Philipp Hennig, and Lukas Balles, Limitations of the empirical fisher approxima- tion for natural gradient descent, H. Wallach, H. Larochelle, A. Beygelzimer, d' d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL
ccpppp ppd
Ya Le and Xuan Yang. Tiny imagenet visual recognition challenge. CS 231N, 7(7):3, 2015.
Yann Lecun. A theoretical framework for back-propagation. 08 2001.
Yann A. LeCun, Léon Bottou, Genevieve B. Orr, and Klaus Robert Müller. Efficient backprop, pp. 9-48. Lecture Notes in Computer Science (including subseries Lecture Notes in Artificial Intelli- gence and Lecture Notes in Bioinformaties). Springer Verlag, 2012. ISBN 978364235288. 1. doi: 10.1007/978-3-66 3---2---222------- 3. Copyright: Copyright 2021 Elsevier B. V. All rights reserved.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. ArXiv e-prints, pp. arXiv-1607, 2016.
Valentin Leplat, Daniil Merkulov, Aleksandr Katrutsa, Daniel Bershatsky, Olga Tsymhoi, and Ivan Oseledets. Nag-gs: Semi-implicit, accelerated and robust stochastic optimizer. arXiv preprint 0rXiv:2209.149371 2022.
Wu Lin, Felix Dangel, Runa Eschenhagen, Juhan Bac, Richard E. Turner, and Alireza Makhzani. Can remove the square-root in adaptive gradient methods? second-order perspec- tive. In Forty-first International Conference on Machine Learning, 2024. URL. https:
/ //operreiiww wwt.//tttoruae wtw.forrmmmm...um......
d=vuMD7 71
Hong Liu, Zhiyuan Li, David Leo Wright Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre- pre-training. In The Twelfth International Conference Learning Representations, 2024.
13

---

Published as a conference paper at ICLR 2025
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. In International Conference on Learning Representations, 2020. URL www./pperrrntatt..wvvee
03366
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 10012-10022, 2021.
Iyy Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Confer-
ence on Learming Representations, 2019a. URL https:: httpeeiwwwwwwwweeeeeeeeeeeeeeeee fortt/form?eeee
EikgRicqYY
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Confer- ence on Learning Representations, 2019b.
Liangchen Luo, Yuanhao Xiong, and Yan Liu. Adaptive gradient methods with dynamic bound of learning rate. In International Conference on Learning Representations, 2019. URL https:
//openreviee.nnt//orr.... BBgg3mm222dd333333FX Bkg3g2R9FX
Linjian Ma, Gabe Montague, Jiayu Ye, Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael Ma- honey. Inefficiency of k-fac for large batch size training. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 5053-5060, 2020.
Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the sdes and scaling rules for adaptive gradient algorithms. Advances in Neural Information Processing Systems, 35: 7697-7711, 2022.
Mitchell Marcus, Beatrice Santorini, and Mary Ann Marcinkiewicz. Building a large annotated corpus of English: The Penn Treebank. Computational Linguistics, 19(2):313-330, 1993. URL https://aclantholle
org/ (a3gg/J932000
James Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21(146):1-76, 2020.
James Martens and Roger Grosse Optimizing neural networks with kronecker-factored approxi- mate curvature. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 2408-24177 Lille, France, 07-09 Jul 2015, PMLR. URL /ttpttts///prcceeinngs mllr. press/v3 maammartaaa htt./m/ateennem..mmmmmmmm htmm
James Martens and Roger Grosse. Optimizing neural networks with kronecker-factorde approximate curvature, 2020.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mix- ture models. In International Conference Learning Representations, 2017. URL https:
/ //openreview. wwweeenrvveee.ee//ffrrreee www./eetromma d=
Konstantin Mishchenko and Sebastian U Stich. Noise injection irons out local minima and saddle points. In OPT 2023: Optimization for Machine Learning, 2023.
Alan V. Oppenheim. Ronald W. Schafer, and John R. Buck. Discrete-Time Signal Processing. Prentice-hall Englewood Cliffs, second edition, 1999.
Kazuki Osawa, Satoki Ishikawa, Rio Yokota, Shigang Li, and Torsten Hoefler. Asdl: A unified interface for gradient preconditioning in pytorch, 2023.
Park, S.II Amari, and Fukumizu. Adaptive natural gradient learning algorithms for var- ious stochastic models. Newral Networks, 13(7):755 2000. ISSN 0893-6080. doi: https//doaayy/////.oossss3....00
URL https:: http///wcciinceireetttttt
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, 32, 2019.
14

---

Published as a conference paper at ICLR 2025
SGOPAL Patro and Kishore Kumar Sahu. Normalization: A preprocessing stage. arXiv preprint arXivv1503.06442. 2015.
K. B. Petersen and M. S. Pedersen. The matrix cookbook, October 2008, URL hattp:/wwet
imm. dtw..tt..co ck wtw.k/ppuuutee http:///wwwpbbbbppppp php?3274 Version 20081110.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence adam and beyond. In International Conference Learning Representations, 2018. URL /tppsttttp////eeee/e/eee
feet/forum?id=zyyu722 RZ
Nicolas Roux, Pierre-Antoine Manzagol, and Yoshua Bengio. Topmoumoute online natural gradient algorithm. Advances in neural information processing systems, 20, 2007.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (JCV), 115(3):211 21-22,212122,..11 2015. doi: 10. 10010/ 1263-015-0816-yy.
Ernest Ryu and Stephen Boyd. A primer on monotone operator methods survey. Applied and computational mathematics, 15:3-43, 01 2016.
Tom Schaul, Sixin Zhang, and Yann LeCun. No more pesky learning rates. In International confer- ence on machine learning, pp. 343-351. PMLR, 2013.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596-4604. PMLR, 2018.
lya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initial- ization and momentum in deep learning. In Sanjoy Dasgupta and David McAllester (eds.), Pro- ceedings of the 30th International Conference on Machine Learning, volume 28 of Proceedings of Machine Learning Research, pp. 1139-1147, Atlanta, Georgia, USA, 17-19 Jun 2013. PMLR. URL https::/ppoceeeeeeee http:////cccceeeigs.../e mlr www.prresssvvv http://vv//ttttteeeeeeeeee html
Ryo Takahashi, Takashi Matsubara, and Kuniaki Uehara. Data augmentation using random image cropping and patching for deep cnns. IEEE Transactions on Circuits and Systems for Video Tech- nology, 3019):2917-2931, September 2020. ISSN 1558-2205, doi: 10. (09/esst.209922955122
URL httL: http:::/wwxxdd..ooooooo wtt.:/...000...111111/////////// 0
201 9 2935128
Zedong Tang, Fenlong Jiang, Maoguo Gong, Hao Li, Yue Wu, Fan Yu, Zidong Wang, and Min Wang. Skfac: Training neural networks with faster kronecker-factored approximate curvature. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp-
13479-13487, 2021.
Mingze Wang, Jinbo Wang. Haotian He, Zilin Wang. Guanhua Huang, Feiyu Xiong, Zhiyu li, Weinan E, and Lei Wu. Improving generalization and convergence by en- haneing implicit regularization.
A. Globerson, Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomezak, and C. Zhang (eds.), Advances in Neural Information Process- ing Systems, volume 37, pp. 118701-118744, Curran Associates, Inc., 2024. URL
https: /ttpp::/pxocceeings.oo neurips cttessccccpaae ww/paper.ppppppppppptpp. wttp://alessep/e//22//// file/ d
Ross Wightman, Hugo Touvron, and Herve Jegou. Resnet strikes back: An improved training procedure in timm. In NewrlPS 202/ Workshop ImageNet: Past, Present, and Future, 2021.
URL httpp:://pperevieee wwt.:eetttfooo htt.://www....eetoooooieneeeeaaa
Ashia C Wilson, Rebecca Roclofs, Mitchell Stern, Nati Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. Advances in neural information process- ing systems, 30, 2017.
Jianwei Yang, Chunyuan Li, Xiyang Dai, and Jianfeng Gao. Focal modulation networks. Advances in Neural Information Processing Systems, 35:4203-4217, 2022.
15

---

Published as a conference paper at ICLR 2025
Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael Mahoney. Adahessian: An adaptive second order optimizer for machine learning. In proceedings of the AAAl conference artificial intelligence, volume 35, pp. 10665-10673, 2021.
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks, 2017.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference on Learning Representations, 2020.
net wett://ororrm...dys.tt.eeteeeeee
Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive meth- ods for nonconvex optimization. Advances in neural information processing systems, 31, 2018.
Matthew D. Zeiler. Adadelta: An adaptive learning rate method, 2012.
Lin Zhang, Shaohuai Shi, and Bo Li. Eva: Practical second-order optimization with kronecker- vectorized approximation. In The Eleventh International Conference on Learning Representa- lions, 2023. URL
net
weet/foormiii i Mi c8V96Voy
Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi Zhi-Quan Luo. Why trans- formers need adam: A hessian perspective. In The Thirty-eighth Annual Conference on Neu- ral Information Processing Systems, 2024. URL
X6raEpbnj3
Zhilu Zhang and Mert Sabuncu. Generalized cross entropy loss for training deep neural networks with noisy labels. Advances in neural information processing systems, 31, 2018
Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar Tatikonda, Nicha Dvornek, Xenophon Pa- pademetris, and James Duncan. Adabelief optimizer: Adapting stepsizes by the belief in ob- served gradients. H. Larochelle, M. Ranzato, Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 18795 18806. Curran Associates, Inc., 2020. URL https: paper/2020/ aal.a//ceoooco11112220a2aaaaaaa1
files/
16

---

Published as a conference paper at ICLR 2025
APPENDIX
CONTENTS
A Theory
18
A.I KFs: A Structural Examination (Continue)
18
A.2 Proofs
21
A.3 Computation KFss
A.4 Distributed AdaFisher
25
Ablation Studies
26
B.1
Evaluating Stability Across Learning Rate Schedulers, and Assessing Convergence Efficiency
26
B.2 Component Analysis: Evaluating the Significance of AdaFisher Elements
28
Visualization
29
Experiments
30
D I Hardware
30
D.2
Image Classification
30
D.2.1
HP HP Tuning
30
D.2.2
Dataset Details
32
D.2.3
Transter Learning
32
D.2.4
Results
33
D.2.5
Comparison with Other Relevant Methods
39
D.2.6
Comparison with Consistent Epoch Counts
39
D.2.7
Comparison of Training Speed and Memory Utilization
40
D.3
Language Modeling
41
D.3.1
Dataset Details
D.3.2
Network Details
42
D.3.3
HPs
42
D.3.4
Results
42
17

---

Published as a conference paper at ICLR 2025
A THEORY
A.1 KFs: A STRUCTURAL EXAMINATION (CONTINUE)
In the realm of matrix theory, the Gershgorin circle theorem offers a principle for localizing the eigenvalues of complex square matrix, asserting that each eigenvalue is situated within at least one Gershgorin disk. These disks are defined by the matrix's 's diagonal elements and the sum of the absolute values of the respective off-diagonal row entries. Formally, the theorem stated as follows:
Theorem A.1 Gershgorin Circle Theorem). Let A be complex square matrix with eigenvalues X. For each there exists an index such that
As S
Aigl
J
where the summation excludes the diagonal entry Ass
For a detailed proof of Theorem A.I the reader is referred to the seminal work by Horn and Johnson Horn& & Johnson (2012). Extending the application of the Gershgorin circle theorem to the study of KFs within deep neural networks, we analyze these factors from both convolutional (37th) and linear (41st) layers of ResNet-18 network through different training phases on CI- FAR10 dataset. As elucidated in Section 3.1] leveraging Theorem A. demonstrates that the eigen- values of the KFs from the convolutional layer are predominantly concentrated along the diago- nal. This observation is analogously applicable to the linear layer. Figure 3 showcases the Ger- shgorin disks for the 41st (linear) layer, with the eigenvalues (red crosses) significantly clustered within these disks (centered at the black circles), underscoring pronounced diagonal dominance. Moreover, upon introducing Gaussian noise to the off off-diagonal elements following this scheme:
M A & where [esy] and Eij (N0 o for i + j. the perturbation analysis elucidates that such stochastic variations engender only marginal displacements in the eigenvalues. Notably, those eigenvalues fulfilling the Kaiser criterion are minimally affected, substantiating the resilience of the diagonal dominance against woise-induces perturbations. Our next analysis focus centers on
With Noise
Without Noise
Kaiser Rule
Matrix X Step 5200
Matrix Step 9800
Matrix H Step 5200
Matrix
Step 9800
10a
Imaginary Part
15
10
10
10
0
10
10-3
10
-10
-10
5
10-5
10-5
-10
0
10
-10
0
10
100
10
102
100
10
102
Matrix S Step 5200
Matrix S| Step 9800
Matrix S Step 5200
Matrix s Step 9800
Imaginary Part
1
1.6
10-1
10-1
10-3
10-3
-1
-1
100
10a
100
10a
Part
Real Part
Figure 7: Gershgorin disks and eigenvalue perturbation analysis for matrices H. and training steps 5200 (middle training) and 9800 (end of training) in ResNet-18 18 Network's Linear Layer (41st Layer). left panel depicts Gershgorin's 's circles in the complex plane, while the right panel illustrates the magnitude spectrum of eigenvalues with and without the influence of Gaussian noise.
elucidating the behaviors of matrices through consecutive steps in the frequency domain, thereby highlighting the intricate patterns and transformations emergent from the training process. By de- ploying Fast Fourier Transform (FFT) 7 and S, along with their noise-infused variants 7/ and S, aim dissect the spectral nuances of these factors. The deliberate addition of noise to the
18

---

Published as a conference paper at ICLR 2025
off off-diagonal serves probe to validate hypothesis that the pivotal information of the KFs predominantly concentrated along their diagonals. The minimal impact such noise perturba- tions observed empirically underscores this diagonal dominance. Our analysis aims to juxtapose the frequency domain representations of both the uncontaminated and the noise-affected matrices at assorted iterative phases, thereby illuminating the inherent stability and tenacity of the Kronecker structures amidst stochastic disturbances.
Let A be two-dimensional X n matrix. The FFT A, denoted as F(A), is computed as
m-1n-
F(A)
Σ p=0 9=0
Apg
e 2 +++))
(5)
where F(A)aa is the value of the FFT at the kth row and /th column of the transformed matrix, A pq
the value of the original matrix at the pth row and gth column, and i is the imaginary unit (Op
penheim et al. T999). Figure 8
demonstrates the Fourier spectral analysis of the KFs H and S
A]
Matrix 2r Step Matrix
Step
Matrix s Step 5200
Matrix s Step 9800
0
0-
0-
-2
-2
-2
-2
2
-2
0
0
Matrix i/ Step 5200
Matrix 2 Step 9800
Matrix Š Step 5200
Matrix 5 Step 9800
-0
-25
2-
2
-50
0-
-100
-75
-2-
-2
-2
100
SNR: 18.85dB
SNR: 19.10dB
SNR: 9.81
SNR: 9.17dB
B)
Matrix // Step 5200
Matrix 2/ Step 9800
Matrix Step 5200
Matrix S Step 9800
2
2-
2-
150
0-
100
-2
25
50
-2
Matrix 7 Step
Metrix Step 9800
Matrix S Step 5200
Matrix $ Step 9800
-2
2-
2
0-
0
0-
-50
-2
-2
-2
SNR: 9.00dB
SNR: 9.51dB
SNR: 5.98dB
SNR: 6.43dB
Figure 8: Comparative Visualization of FFT Outputs for KFs in a ResNet-18 18 Network's Convolutional and
Linear Layers. (A) FFT results KF's and 7l from 37th convolutional layer under noise-fiee conditions (top) and with Gaussian noise (bottom) at iterations 5200 (middle of training) and 9800 (end of training). (B) Analogous FFT results for the KFs and & from the 41st linear layer, contrasted between noise- free (top) noisy conditions (bottom) at same iterations.
distinct iterative stages of training- 5200 and 9800 for convolutional and linear layers
19

---

Published as a conference paper at ICLR 2025
(37th and 41st of ResNet-18 network respectively). Each KF is analyzed via FFT in both a pris- tine, noise-free condition and a Gaussian noise-affected state, with the associated Signal-to-Noiss Ratios (SNRs) detailed in Eq. 6 In the noise-free FFT spectra, pronounced diagonal energy concentration is manifest in the H and S factors of the convolutional layer, indicative of significant informational preservation along the diagonal. In contrast, the linear layer exhibits less pronounced but still discernible diagonal energy distribution, suggesting a more diffuse yet still noteworthy di- agonal information structure. With the addition of noise, the matrices H and S still display notable diagonal pattern, indicating minimal SNR deterioration. This observation supports the proposition that the KFs primarily encode their information along the diagonal, and the introduction of noise into the off-diagonal elements has a limited impact. The SNR between matrix M and M is computed using the formula:
IMsi Σnst MMs
SNR 10 logio
where Ma denotes the diagonal elements of M, and May represents the upper triangular elements of M excluding the diagonal (Oppenheim et al. 1999). The observed reduction in SNR from step 5200 to step 9800 for the KF S the convolutional layer, under noisy conditions, could suggest an incremental integration of noise effects across iterations. Conversely, for the remaining factors, an increase in SNR throughout the training process is detected, which may indicate an enhancement in signal clarity. Nevertheless, the integrity of the diagonal concentration of energy remains predomi- nantly intact, demonstrating the underlying robustness of the network's s feature extraction capability against noise perturbations. Ultimately, the spectral analyses validate the hypothesis that the KFs' informational content is predominantly diagonal and resistant to the effects of off-diagonal Gaussian noise. This durability is sustained through successive iterations, maintaining the primary spectral characteristics of the KFs. Figure 9 offers a visual exposition of the Kronecker Product Factors H
A)
Matrix Step 5200
Matrix Step 9800
Matrix 5200
Matrix Step 9800
3393
1399
0.0
106
3000
3000
152
3.00
1000
2000
.000
4200
2000
2006
3000
4000
30
20
60
13
B)
Matrix X/ Step 5200
Matrix X Step 9800
Matrix S Step 5200
Matrix S Step 9800
0.00
100
132
0.06
200
235
300
202
0.04
400
470
0.02 52
50
622
500
02
Figure 9: Visualization of KFs H. and S for convolutional (A) and linear (B) layers at different iteration steps within ResNet-18 network. For the convolutional layer (37th layer), the first two plots in (A) represent factor H steps 5200 (middle of training) and 9800 (end of training), clucidating the matrix's structure these stages. The subsequent two plots display factor S, highlighting changes in granularity and contrast with iteration progression. Similarly, in (B) for the linear layer (11st position), we observe the structural evolution of factor H and S over the same iterations. with variations in pattern density and clarity. These visualizations collectively underscore the dynamic nature of the KFFs' architecture training advances.
and S at progressive iteration junctures-speeiinne linear layers (37th and 41st ResNet-18 network respectively). The initial duo of plots in each (A) and (B), delineate the KF H at the aforementioned steps, elucidating the matrix's 's structure at two distinct evolutionary stages. The next duo plots in (A) and (B) represent the KF S at differ- ent steps of training. This visual examination, in conjunction with the preceding spectral analyses, articulates an integrated story of the developmental trajectory of the KFs. The enduring diagonal salience observed in both H and S underscores the notion that the informational energy of the KFs
steps 5200 and 9800 for convolutional and
20

---

Published as a conference paper at ICLR 2025
is predominantly concentrated along the diagonal. This persistent feature accentuates the structural stability and the focused nature of information encoding within the network's 's layers.
A.2 PROOFS
Proposition A.1 (FIM for normalization layer). Let (vi, 3.) E &R be the scale and shift parameters of normalization layer The empirical KFs for the FIM approximation are
Hi-1
Σ TET.
h1-1-2h-11,0 Hi-1 = =11,, S;
Σ DET.
81
T
|Ts
T
where 8 RCXIITA represent the pre-normalized activations and gradients, respectively, Here, To is the set of dimensions over which normalization statistics are computed, and Cs is the channelsffeatures size.
Proof. Let (vs, Bi) E CCC, be the scale and shift parameters normalization layer i, with trans- formation
h; a, v4 O h4-1 3
where 1i-1 E RC contains normalized activations and @ denotes element-wise multiplication. Let V. = Σ, 1,2 O 8i,z and (() >n Si,e where s; Vn, (0).
For parameters:
2 8i,ar I
oV, C = E
Σ -i 2 88
Σ him
Rv E
Σ (hi (1i- - @ (Sime
(K-FAC independence assumption)
(17
Σ
h hi /
2 (Tm
Σ
Si,z
Ti
x
x
= H4-1
@ S;
For & parameters:
Σ Si.z T
E 3, LV
1 E
Σ Si.
8
-11
Σ
$
(Bias term factorization)
|T
H4-1
X
S
9.
Cross-terms between Vi and 3; are excluded under the diagonal block assumption.
Proposition A.2 (Efficient FIM). Let Hi-1 and S₁ represent the KFs for a given layer index i within a neural network, where these factors exhibit semi-diagonal characteristics indicating energy concentration predominantly along the diagonal, as elaborated in Section 3.1 Define gi as the gradient obtained through backpropagation at layer i. Assume that H4-1 and S; can be closely approximated by diagonal matrices, denoted by H.D.-, and Sp. respectively layer i, such that H.D... = Diag(H,-1). Sp. Diag(S; (S,) where Diag(M, denote the diagonal approximation of matrix M, which retains only the main diagonal. Therefore, we define the Empirical FIM as
Fp. 2 H'D_, & &SD. + Al,
(7)
21

---

Published as a conference paper at ICLR 2025
where M' denotes the Min-Max normalization technique Patro & Sahu (2015) for M H.D., or SD.- The regularization parameter X set to 0.001 serves as damping factors, in alignment with the principles of Tikhonov regularization, enhance computational stability and improve the ditioning of the matrix. The foundational aspects of the K-FAC optimization approach are detailed in Martens & Grosse 2015). Then, the closed-form solution for the preconditioned gradient gD),
derived from the diagonal approximation of the FIM, is given by: g( = (
Proof. The justification of our approach comprises two principal components: the rationale for adopting diagonal approximation of the KFss and the methodology for normalization and regular- ization of these factors.
Part 1: Diagonalization of KFss
The assumption of independent neuronal activity within layers is foundational to our approach. This assumption posits that the covariance matrices H. and S, encapsulating the second-order statistics of activations and sensitivities, respectively, are diagonal. This diagonal nature arises because in- dependence among random variables implies covariance of zero for any pair of distinct variables. thereby nullifying all off-diagonal elements of these covariance matrices.
Consider matrices A and B, each being diagonal with elements au and bjj. respectively. The Kro- necker product A @ B, by definition, generates elements رnbyy at the corresponding (i,j) positions. For diagonal A and В, this product maintains non-zero values exclusively at diagonal positions where i = j. resulting in:
A & B = diag(an Oryn Dmm))
yielding a purely diagonal matrix. Moreover, we have empirically demonstrated that the energy of the KFs is concentrated along the diagonal, detailed in Sections 3.1 and A.I These arguments support our initial premise.
Part 2: Normalization and Regularization
Let M E (HD, Sp. be a diagonal matrix with entries my > 0. The min-max normalized matrix M' satisfies
D Imin D diag( V Mmaax Mmiin
where mmiin ming My. mmax max; my. This affine transformation ensures that O z M' < T where & denotes Loewner ordering. Combined with Tikhonov regularization, the modified FIM, Fr. = -HD.., @ Sb. + AI admits eigenvalue bounds
S Au AFD.) +1+> Vk:
which guarantees numerical stability for inversion. This approach ensures that all elements are scaled uniformly, preserving their relative magnitudes and distances. Compared to other normaliza- tion methods, such as z-score normalization (Patro & Sahu 2013), Min-Max normalization offers several advantages such the normalization Stability, where for M' (M //min///mmm (m 1] where o(-) denotes matrix spectrum. Morcover, the Kronccker
mmin) we have (MM' S [0. product satisfies (H @ Sb,)) S [0, 1]. thus (FD.) 2 > 0, guaranteeing invertibility.
And the relative error satisfies [1 Fp.1 11
< O(c+ A) where measures diagonal approximation
|
eтror. Therefore, the preconditioned gradient can be written has g( = (H' &Sp, +A (11)) Proposition A.3 (Convergence in convex optimization). For the FIM defined in Eq. 7 the up- dating scheme git+1) g(t) az - converges. Moreover, if JJ is Lipschitz, i.e [() VJ(0')l2 L| 0'll any 0 and 0', then for the k-step iteration with fixed step size a < L,L, then
- - J(0) < || - 1011
2ac
where ((0) is the optimal value.
22

---

Published as a conference paper at ICLR 2025
Proof. We follow the same proof as in Yao et al. (2021). Assume that J()) is a strongly convex and strictly smooth function in R", such that there exist positive constants a and B so that al I (²)(0)= & 31 for all w. We can show that the update formulation A0(+) = (F(tt)- g((t) converges by showing that with the proper learning rate:
Ag == J(B(t+1) - J(B()) 2,32 |
a
Note that when A: = 0 or 1, the convergence rate is the same as gradient descent or Newton method, respectively. Our proof is similar to Boyd & Vandenberghe (2004) for the Newton method. We denote X(((t)) (o (F(1)) ( ((0)))) Since J(O) is strongly convex, we have
17 17
J( n( s I(6(0) 7/ (TT()))_
(
2
B 11 2(0()))2
- +
2xx
The second inequality comes from the fact that
)()())2 (0(1)) p(t) ( (0(0) 2 all
Therefore, the step size i = a/ B will make f decrease as follows,
J( k00(0)) - J(0(0) - 2 7
Since al 3 F(1 3 3I, we have
=(0(0))=(1)" ( > 3 11
Therefore,
1
J(O( -) J((()) S 23 u
(0)
2,32 11
Since p(()) is positive definite, hence Eq. (8) holds true. For the bound on convergence rate, we refer to Ryu & Boyd (2016) for the details of the complete proof.
Proposition A.4 (Convergence in non-convex stochastic optimization), Under the assumptions:
(i) f is lower bounded and differentiable; [[V](0) (O))l) L|10 0'112. 11
L. vt,0 0,0'.
(ii) Both the true and stochastic gradient are bounded, i.e. ||V S À and 119cl12 S X, Vt for
some X >> O.
(ūii) Unbiased and independent noise in ( Lee. g = VJ(((0)+ Blc = O, (andç()) (())
Vi + j.
Assume (1 3/. g00) S S 1 is non-increasing,
F F 1pt-11
F0 Φi
Vt E [T], E [d]. we then
have
L min 1O S [ LE[T]
+ log T) + C2dn + Cydy²² + C.)
(9)
where C₁ O2, C3 are constants independent of d and T, CA is constant independent of the
expectation is taken w.r.t all the randommess corresponding to (1)²)
Proof. Follow Chen et al. (201)), as AdaFisher is an Adam-type method with the condition
mm (1) / ll2 I G for some G (which can be obtained by r)()) n. (lg l|2 s and
2 2 1), we have
() (t)
2
(t())
(t-1)
E
(1) (V)((())) (T
(00))
<E
C
+ C2
7.(0)
P 11-1)
T
17(1)
17
2
+ C33
Σ
+ CC+
(10)
70
23

---

Published as a conference paper at ICLR 2025
We first bound non-constant terms in RHS of Eq. we have
10
For the term with C1 since || 1/2 1,
n
g)
22
E
E
In
00 g
10
F
T
7 (0)
E
T
S
X1 + log
For the term with C2,
we have
T
(1))
17
(t--) 1)
d
T
n
77())
E
E
1) 1)
[i]
1 bil
D
/=1
(7
E
n
E
dn
A
where the first equality is due
[T].]
[d]
(1
For the term with C3, we have
n
27
()(1)
1
E
E
F ((--)
1)
F Li)
T
v)(1)
1)(1)
1 (
E
P Lil
6 b
F
5
T
(1
1
11
7
<
E
F
F 15-"6o 6 [
VIFD [}]
V1- D
1/(111
E
7
F [
A " b
(tt-1)
E
F6-11 1]
D
< dip
Hence
o (t)
2
(
(t-1)
(t)
n 1)
2
E C
C2
C3
+ C4
F
P(O)
-
F F
2
S C,n X1 + log T) + C2dy + Czdr2² + Cs
(11)
Now we lower bound the LHS of Eq.
g With the assumption |10 S L. we have
(1) / Z "
LVi
Thus
VT
E
E Σ V.J(OL)))
min
(()
Lt
7
(E[?]
L=l
(12)
Combining Eq. 11 and 12 gives the desired result.
24

---

Published as a conference paper at ICLR 2025
A.3 COMPUTATION OF KFs
The KFs H. and S, which are integral to the AdaFisher optimizer, are computed following method- ologies similar to those described in Grosse & Martens (2016): Martens & Grosse (201)). This section revisits the key equations used for this computation. For a given layer in neural network, the empirical KFs are computed as follows:
For fully connected layers, the KFs are:
HD.., diag(h:-
Sp, diag(s;s;
For convolutional layers, the computation accounts for the spatial positions within the layer, denoted as T:
--=") Sp. diag
[h [
8:8 [
H.D.,. diag
The algorithm employs the expansion operation denoted by [-] Grosse & Martens 2016
This operation essentially takes the patches surrounding spatial locations, stretches them into vectors, and compiles these vectors into a matrix.
For Normalization layers (BatchNorm & LayerNorm) refer to Proposition. 3.1
For all other type of layers the KFss are:
HD,,:
+1g
SD. Ippes
Table 6: AdaFisher training time per epoch (s) across various numbers of GPUs ResNet-50 ImageNet- Ik.
GPU amount
Batch Size
AdaFisher training time per epoch (s)
256
2882
512
1438
768
9633
1024
720
A.4 DISTRIBUTED ADAFISHER
The efficacy of AdaFisher hinges on its innovative approximation of the FIM, denoted as F. which leverages KFs for computation. In a distributed setting, it is crucial to aggregate these KFs across multiple GPUs before updating the model parameters. Consider training environment consisting of GPUs. For any given layer i, the KFs are computed and aggregated across GPUs as
H (SSM
Σ
D
(UM DD.
1 K
kk) DD
(13)
K
The theoretical justification for this aggregation lies the linearity of expectation and the unbi- asedness of the local KF estimates. Specifically, if each (k) D.-, and S are unbiased estimators of their respective true factors H.p... and Sp, for k = I...., K, then the averaged factors 7. and S remain unbiased estimators of HDD and Sp,- Consequently, using Eq. (3). the aggregated EFIM for layer can be calculated as
H (SUM)
(SUM) S + AI
D.
where À is a regularization parameter added to ensure numerical stability. This methodology ensures that each GPU contributes to comprehensive update the model, enhancing both convergence and performance in large-scale distributed training environments. We assessed the distributed version of AdaFisher on ImageNet- 1k, utilizing batch sizes of 512 and 1024 (refer to Table 3 and Figure [5] for details). Our findings indicate that AdaFisher scales nearly linearly with the number of GPUs, as evidenced in Table 6 There remains scope for additional low-level optimizations within the implementation further enhance performance.
25

---

Published as a conference paper at ICLR 2025
B ABLATION STUDIES
Building on the ablative studies detailed in Section 4.4 this section extends our stability analysis to explore the impact of various learning rate schedulers and convergence efficiency, as discussed in Section B. Additionally, we conduct an in-depth examination of the key components of AdaFisher. This includes analyzing the effects of the EMA, the of square root transformations, approximation the FIM, and the critical role of computing the FIM for normalization layers, all of which are detailed in Section B.2
novel
We have consolidated the key findings of each ablation study
in Table
Table 7: Summary of Ablation Studies for AdaFisher Optimizer.
Ablation Study Learning rate sched- ulers
Component Studied
Key Findings
Impact of Cosine Annealing. StepLR. and scheduler on AdaFisher
AdaFisher maintains stable and efficient performance across various schedulers, demonstrating its robustness and adaptability in diverse training environments. Fur- ther details are given in Section B.1
Convergence Efficiency
Performance and alignment of FIM with Hessian
AdaFisher shows marked performance improvements to- wards the end of training, with FIM alignment to the Hes- sian enhancing rapid convergence and stable generaliza- tion across training and testing phases. Further details are provided in Section B.I
Square Root Uti- lization
Effect of omitting square root in update rules
Eliminating the square root enhances AdaFisher's perfor- mance and stability, outperforming both its own version with the square root and Adam without the square root, while also improving computational efficiency. Further details are listed Section [.2]
EMA
of Utilization of EMA for curvature estima- tion
Using EMA on KF's enhances AdaFisher's curvature es- timation, leveraging data from multiple mini-batchess for continuous updates, demonstrating significant benefits in methods with diagonal or block-diagonal curvature approximations. Further analysis included in Sec- tion B.2
KFs
Importance of Fisher
Impact of EFIM in normalization layers
Incorporating Fisher computation in normalization lay- ers significantly improves AdaFisher's generalization and stability by enhancing parameter sensitivity and gradient variability insights, crucial for optimizing training dy- namics and model convergence. Further details are given in Section B.2
Compu-
tation for Normaliza- tion Layers
New Approxi- mation of the FIM
Diagonal approxima- tion of the FIM
Our novel method focuses on the diagonal elements of the FIM, enhaneing computation efficiency without los- ing critical information. Validation shows our approxi- mation closely aligns with the true Fisher, confirming its efficacy. Further details are contained in Section B.2
B. EVALUATING STABILITY ACROSS LEARNING RATE SCHEDULERS, AND ASSESSING CONVERGENCE EFFICIENCY
Learning rate schedulers. This analysis evaluates the impact different learning rate schedulers- Cosine Annealing, StepL.R, and no scheduler- the performance of AdaFisher, as depicted in Figure 10 AdaFisher exhibits remarkable robustness across these scheduling strategies. Notably, its performance remains stable and efficient, whether it is paired with the gradual adjustments of Cosine Annealing, the abrupt changes StepLR, or even in the absence of any scheduler. This underscores AdaFisher's adaptability and effectiveness in diverse training environments.
Convergence Efficiency. As training progresses, AdaFisher optimizer demonstrates a significant enhancement in performance compared to its counterparts, especially evident towards the end of the training period (see Appendix D.2.4). This rapid convergence is attributed to AdaFisher's ap- proach by incorporating the FIM. Early and mid-training, the FIM serves approximation to
26

---

Published as a conference paper at ICLR 2025
the Hessian matrix, equivalent the Generalized Gauss Newton Matrix (Eschenhagen al. 2024) However, as the model approaches local minimum, the FIM increasingly aligns precisely with the Hessian (Martens] 2020). This precise alignment accelerates convergence, markedly improving the optimizer's efficiency in the final phases of training. Additionally, AdaFisher's tendency to con- verge to flat local minima leads more stable generalization when transitioning from training to lesting distributions (Cha et al. 202I), contrasting sharply with other optimizers. To support these points, we analyze the training distribution of our diagonal block-Kronecke FIM during the train- ing of ResNet18 on CIFARIO. Specifically, we examine the FIM distribution for the first (Panel A), middle (Panel B) convolutional layers and the last linear layer (Panel C), as shown in Figure [T] is evident that for each layer, the FIM distribution with AdaFisher narrows to smaller values with fewer variations compared to that with Adam. This pattern demonstrates AdaFisher's convergence toward flatter local minima, as the Fisher Information, approximation of the Hessian, contains crucial curvature information.
Training Resnet50
Testing
10
Schedsler: Nome
Scheduler: None
Schediler: Codinekenealing
Schedultn. Cositolaneallng
Scheddder:: Stepth
Schecblln:: Stepl.
60
30
10-
100
150
200
150
200
170
150
100
200
Figure 10: Performance comparison of AdaFisher using the ResNer50 on the CIFARIO with a batch size of 256 with different learning rate schedulers.
(A) First Layer
B Middle Layer
(C) Final layer
Histogram of FIM Diagonal for Adam
Histogram FIM Diagonal Adam
Histogram of FIM Dingonal for Adam
weopgen Aice
Hiltogram of TIM Dingonal for Adallshhr
Histogram of FIM Diagonal for AdaFisher
Histegram of FIM Diagonal for AdaFisher
Eisr
1884
Figure 11: Comparison of FIM Diagonal Histograms during ResNet18 Training CIFARIO with Adam and AdaFisher over 1,000 training iterations. Panel (A) displays the FIM diagonal elements for the first convo- lutional layer; Panel (B) illustrates the FIM diagonal elements for the middle convolutional layer; Panel (C) shows the FIM diagonal elements for the last Linear layer.
27

---

Published as a conference paper at ICLR 2025
(A)
Training
CHARB1
Testing RecnettO
B
Training loss: Recnetto on CHARRD
Teating error: Becnetto on CIARR1
uqgare spuo
146
596 No
Idam kqui Rooo kg
scuire sozt
Spuacc
50
54
30"14
50
300
300
200
159
53
100
189
200
Cauch
Spech
(C)
(D)
Training Tay Adaftshe
compuraiio other
compurttion ooher
MAE
Computation 3o:
Computation Stc
Lapn 102
0:53
6.020
0.02
18.
0:31
200
200
250
300
30
1pech
Spcch
Figure 12: AdaFisher Component Analysis. (A) Comparison of MAE between the true FIM Fb and our approximation FD across convolutional and dense layers. (B) Performance comparison of AdaFisher with and without the EMA of KFs. (C) Assessment of AdaFisher's performance with and without the computation of EFIM for Batch Normalization (BN) layers.
B.2 COMPONENT ANALYSIS: EVALUATING THE SIGNIFICANCE OF ADAFISHER'S s ELEMENTS
AdaFisher incorporates several key components, including novel approximation of the FIM, the EMA of the KFs, the omission of the square root in the update rule, and a new EFIM formula for normalization layers. In this part, we elucidate each component and its significance within the AdaFisher optimizer.
Square Root Utilization. Recent studies, such as (Lin et al. 2024), have reevaluated the necessity of the square root operation in the Adam family's 's update rules. These studies suggest that eliminating the square root does not affect convergence and may even narrow the generalization gap compared to SGD in CNN models. Our analysis, shown in panel (A) of Figure [2] investigates this aspect by comparing the performance of AdaFisher and Adam, both with and without the square root operation. The findings reveal that removing the square root not only boosts the performance and stability of both optimizers but also significantly enhances computational efficiency. Specifically, AdaFisher without the square root not only outperforms the version with the square root but also surpasses Adam without the square root. However, Adam without the square root typically requires an additional scaling factor proportional to the batch size, denoted f x batch size, to function correctly. Without this factor, Adam, without the square root, fails to learn effectively, making direct comparisons with AdaFisher invalid.
EMA of KFs. As elucidated in Section 3.2 employing an EMA over the KFs facilitates a more sophisticated curvature estimation. This technique leverages data across multiple mini-batches, en- abling continuous updates to the Fisher information rather than relying solely on the data from a single batch. Panel (B) of Figure 12 underscores, using ResNet-50 on CIFARIO over 200 epochs, the benefits of using EMA on KFs,, a strategy particularly advantageous in methods that utilize di- agonal or block- bock-diagonal approximations of the curvature matrix.
Importance of Fisher Computation for Normalization Layers. The integration of the EFIM in normalization layers, as detailed in Proposition 3.1 significantly enhances the generalization pro- cess. Panel (C) of Figure 12] illustrates the impact of incorporating Fisher computation in these layers during the training of AdaFisher with ResNet-50 on CIFARIO over 200 epochs. In contrast, the identity matrix is employed when Fisher's computation is omitted. The superior performance of AdaFisher when incorporating Fisher computation can be attributed to the critical role normalization layers play in adjusting the input distribution for each mini-batch This adjustment substantially en- hances the neural network's 's learning stability (Jing et al., 20246). By quantifying the information each output y carries about the parameters 0 under the model distribution p(y|ar; 0), the computa-
28

---

Published as a conference paper at ICLR 2025
tion of the FIM in these layers provides valuable insights into parameter sensitivity and gradient variability. This insight crucial for optimizing training dynamics and enhancing model conver- gence- areas that are often inadequately addressed by existing optimizers.
New Approximation the FIM. In Proposition 3.2, we introduce new methodology for approx- imating the FIM that diverges from the K-FAC optimizer. Unlike K-FAC, which utilizes the full Kronecker product, our approach focuses solely on the diagonal elements of the FIM, where, as demonstrated in Section 3.1, the energy of the KFs predominantly concentrated. This method enables more efficient computation of the FIM without sacrificing critical information. To validate our approach, we compare the true FIM diagonal with our approximation convolutional and dense layers using a toy model composed of 2 convolutional layers and two linear layers on a subset of the MNIST dataset (Deng 2012) over 50 epochs. Specifically, we calculate the true Fisher using the NNgeometry Python package (George 2021), which facilitates the computation of the FIM, Gauss-Newton Matrix, or Neural Tangent Kernels applied to neural networks. We estimate p(y(æ) through Monte-Carlo sampling. During each epoch, we collected both the empirical and true Fisher information and calculated the Mean Absolute Error (MAE) between these two measures. Panel (D) of Figure [2] showcases the close approximation of AdaFisher's empirical diagonal to the true Fisher, thus validating the efficacy of our approximation method.
C VISUALIZATION
The convergence rate of optimizer is crucial, serving an indicator of its robustness against sad- dle points and its ability to generalize effectively. In this section, we introduce a novel methodology for visualizing the convergence behavior of optimizers through statistical model, depicted in Figure [ Initially, our process employs Principal Component Analysis (PCA) for dimensionality reduction, reducing the dataset dimensions from D E Rmx* to E mRxx2, following the protocol established in F.R.S. (1901). We then apply this reduced dataset layer multi-layer perceptron (MLP). Notably, we focus on the first weight matrix W of this MLP, which resides in R where e denotes the epoch number. For consistency and ensure comparabil- ity, all layers' weights are initialized identically across different optimizers. Following the training
toy model composed of an L-
Yered
D R
PCA
Wf R?
Wi
w
wr
U
em1
Pipeline for visualization of optimization paths for various algorithms on loss surface, comparing their convergence efficiency.
Figure 13:
phase with various optimizers where denote of optimizer results O, we analyze both the collection of first-layer weights, W, and the evolution of the loss function, & defined as:
(W))
(W2)
W-
2= -. CTIT
TTF)T
where (WF)T represents the weight vector at the eth epoch, and Li represents the loss at the eth epoch, extracted from the optimization results O. We construct a grid (X, Y) spanning the range of weight parameters, discretized into 200 linearly spaced points along each axis:
X, Y meshgrid (Win), max(W..1), min max 200)
29

---

Published as a conference paper at ICLR 2025
Finally, we interpolate the loss values L over the grid using cubic interpolation to obtain smooth loss surface Z:
Z = griddata L, (X, Y), method = cubic)
These elements are integral to the visualization process, which elucidates the optimizer's trajectory through the parameter space across training epochs. It is important to note that while focus the first layer's s weight matrix for clarity, the methodology can be adapted to visualize the weights of any layer within the network. Figure 13 summarizes the pipeline.
In the experiment depicted in Figure T we selected the IRIS dataset rzz widespread recognition and compatibility with PCA application. Our model employs 2-layer MLP architecture. We specifically attend to the weight matrix of the first layer, denoted by W E R2. This particular focus is informed by the empirical observation that the parameters of the first layer tend to exhibit faster convergence rate compared to those of subsequent layers in the network. Such phenomenon can be attributed to the more direct influence of the input features on the first layer's weights, which often results in more pronounced and expedited learning dynamic. Given the classification nature the task, we employed the Cross-Entropy loss function (Zhang Sabuneu, 2018). The network was trained over 20 epochs using a suite of optimizers: Adam, AdaHessian, K- FAC, Shampoo, and AdaFisher. We standardized the learning rate across all optimizers at 1 x 10-3 to ensure comparability of results. Examination of Figure reveals that AdaFisher's convergence markedly superior to that of its counterparts, achieving rapid convergence to the local minimum of the loss landscape concerning the first weight parameter within few iterations. Conversely, the alternative optimizers demonstrate convergence to less optimal local minima. Note that while the re- sults may vary due the stochastic nature of parameter initialization, AdaFisher typically converges to better local minimum compared to its counterparts.
2018). owing to its
D EXPERIMENTS
D.1 HARDWARE
In total, had server with 6 NVIDIA RTX 6000 Ada Generation GPUS with 48 gigabytes of VRAM and 128 gigabytes of RAM available for all experiments, All experiments described in this report were conducted system equipped with a single NVIDIA RTX 6000 Ada Generation GPU and 64 gigabytes of RAM, except for training AdaFisher on ImageNet- Ik with batch sizes of 512 and 1024, where four GPUs were utilized.
D.2 IMAGE CLASSIFICATION
We provide further results and detailed descriptions of our image classification experiments in this section. We conducted five trials with random initializations for the CIFAR experiments and one trial each for Tiny ImageNet and ImageNet-Ik. We present the mean and standard deviation of the results for these trials.
Note training time. Given that various optimizers demonstrate significantly different epoch durations, we have standardized our comparisons by restricting training to the total WCT consumed by 200 epochs using AdaFisher for both CIFAR and Tiny ImageNet experiments. Conversely, for ImageNet- 1k, we report the results based on 90 WCT training epochs using Adam, as, surprisingly. AdaFisher and Adam exhibited the same duration in this experiment. The final selected number of epochs for each optimizer is detailed in Table 8 Note that we were unable to train AdaHessian on ImageNet- Ik due to the significant computational resources required by this optimizer.
Table 8: Comparison of the final epoch counts for various optimizers across different datasets.
CIFARIO/100 Tiny ImageNet
ImagoNet-1k
Optimizers
SGD Adam/AdomW
AdaHessian K-FAC Shampoo|
AdaFisher/AddFiihheeW
Adam K-FAC)
Shampoo AdaFisher
Ppochs
226
210
89
107
36
200
90
60
26
90
D.2.1 HP TUNING
Effective HP tuning is crucial for optimizing the performance of deep learning models. In this study, systematically explored various HPs for both CNNs and VIT's across multiple image clas-
30

---

Published as a conference paper at ICLR 2025
sification tasks. The following subsections detail the tuning strategies employed for each model architecture and dataset.
CNNs. For all image classification tasks involving CNNs, utilized ResNet18 the backbone architecture and evaluated its performance on the CIFAR-10 10 dataset with a fixed batch size of 256 trained on 50 epochs. The HP tuning process encompassed the following components:
Optimizer Selection and Learning Rate Tuning: Each optimizer was fine-tuned using ResNet18 on CIFAR-IO. We performed grid search identify the optimal learning rate from the set {0.0001, 0.0003, 0.0005, 0.0009, 0.1, 0.3, 0.5, 0.9}.
Learning Rate Scheduling: A cosine annealing learning rate decay strategy was em- ployed, aligning with the number of training epochs specified for each optimizer in Table This approach follows the methodology proposed by Loshchilo Hutter (2019a) and was determined be optimal for our experimental setup.
Weight Decay: We applied a uniform weight decay of 5 x 10 across all optimizers for CIFAR-IO and Tiny ImageNet. An exception was made for MobileNetV3, where the weight decay was set to 1 10-5, For experiments on ImageNet-Ik the weight decay was established at 1 x 10
Damping Parameter Tuning:
AdaFisher, K-FAC, and Shampoo:
K-FAC and AdaFisher: The damping parameter searched within
{0.0001, 0.0003, 0.0005, 0.0009, 0.001, 0.003, 0.005, 0.009, 0.01, 0.03, 0.05, 0.09). This range was chosen based on prior research (Martens & Grosse 2015) and
own experiments, which indicated optimal damping values around 1 x 10 -3
* Shampoo: The damping parameter was tuned within {1 x 10 3. x 10 6 x 10 9 10 x 10 3.3 x 10 5,5 10 9 10 1 10 4.3 3. 10-4 1.-4.x 10 10-4.9 x 10 4). as optimal values typically reside around 1 x 10-3.
AdaHessian: The Hessian power was tuned within the range
1.0).
SGD: The momentum of SGD was tuned within the range (0.1, 0.2, 0.9, 1.03. AdaFisher Decay Factors: The decay factor y for AdaFisher was tuned within {0.1, o.,.,........0000 0.99]. The optimal value is: 7 = 0.8.
Implementation Details: For the Shampoo and K-FAC optimizers, we utilized the ASDL library as implemented in PyTorch provided by Osawa et al. (2023).
VITs. For ViT-based image classification tasks, we employed the Tiny Swin Transformer on the CIFAR-10 dataset with a batch size of 256. The HP tuning strategy for ViTs included the following elements:
Weight Decay: Weight decay values were set as indicated in the respective original publi- cations for each model:
Tiny Swin: 10-2
FocalNet: 5 10-2
CCT-2/3x2: x 10-2
Learning Rate Tuning: For SGD, AdaFisher, AdaHessian, K-FAC, and Shampoo optimizers, we conducted grid search over the learning rates {0.3, 0.15,, 0.1, 0.05, 0.03, 0.015, 0.01, 0.005, 0.003, 0.0015, 0.001}. as these optimiz- ers typically operate with higher learning rates compared to Adam-based optimizers. For AdamW, the learning rates were adopted from the original publications:
-Tiny Swin and FocalNet: x 10-4
-CCT-2/332: 5.5 10-5
Damping Parameter Tunning: We performed the same grid search over the damping parameter for K-FAC, Shampoo and AdaFisher, the Hessian power for AdaHessian, the momentum for SGD, and the decay factors for AdaFisher explained in the CNNs part.
This meticulous HP tuning process ensures that each optimizer is optimally configured for the re- spective model architectures and datasets, thereby facilitating a fair and comprehensive comparison
31

---

Published as a conference paper at ICLR 2025
of their performance across different image classification tasks. The final learning rates for all opti- mizers and models are detailed in Table
Table 9: Final selected learning rates for each optimizer, tuned using ResNet18 (for CNN) and Tiny Swin (for ViT) on CIFARIO using a batch size 256. selected based on final validation top-1 accuracy.
Architecture
SGD
Adam
AdamW
AdafTessian
K-FAC
Shampoo
AdaFisher
AdaFFisher
CNNs
0.1
0.001
0.15
0.3
0.3
0,001
VITS
0.01
0.000100000005
0.01
0.003
0.003
0.001
D.2.2 DATASET DETAILS
The training/test sets for Cifar10/100 dataset contain 50k/10k images, respectively. We consider batch size of 256. For CleaR-related experiments. perform 32 32 random-resize cropping, random horizontal flipping and cutout (DeVries & Taylor 2017) as data augmentations. Refer to Takahashi al. (2020) for more details.
CIFAR.
Tiny ImageNet. The training/test sets for Tiny ImageNet Le Yang (0015) contains 100k/10k images. We perform 64 x 64 random-resize cropping and random horizontal flipping. The batch size is set to be 256.
ImageNet-Ik. The training/test sets for ImageNet- 1k Russakovsky et (0015) contains images. We consider a batch size of 256. as we performed experiments on a single GPU instance without any GPU parallelism. We follow He al. (2016) and perform random resized cropping to 224 x 244 and random horizontal flipping on the train set and 256 x 256 resizing with 224 224 center cropping on the test set.
Table 10: Final selected learning rates for each optimizer with ImageNet- Ik pretrained weights, tuned using ResNet50 on CIFARIO using batch size of 256. We tuned by completing a full WCT epoch training cycle and selected based on final validation top-1 accuracy.
SGD
Adam
AdaHeessin
K-FAC
Shampoo
AdaFisher
0.01
0.0001
0.15
0.3
0.03
0.001
Table 11: Final selected epoch counts for various optimizers across transfer learning task.
SGD
Adam/Adamm
AdaHessian
K-FAC
Shampoo
AddFFiher/AAdaFicceit
D.2.3 TRANSFER LEARNING
For transfer learning, weights are initialized to the values provided by the publicly available check- points by PyTorch, except the first convolutional for the ResNet architecture and last dense layers for all networks, which change size to accommodate the new kemel size and number of classes, respectively, that are randomly initialized. We train all models with weight decay 1e as sug- gested in Wightman et al. (2021). except for MobileNetV3, where weight decay is set to be le Moreover, we did grid search for each optimizer for selecting the best learning rate of the range (0.3, 0.15, 0. 0.03, 0.015, 0.01. le - 5} where we tabulate the selected learning rate for each optimizer in Table I0 We use a batch size of 256 and cosine learning rate decay. We use the same augmentation policy (without Cutout) in the previous experiments. The results were obtained using the WCT technique over training epochs of AdaFisher, with the final epoch count detailed in Table II
All other parameters remained unchanged.
Table 12: Performance of various networks and optimizers on Tiny ImageNet using batch size 256. Reported using wall clock time 200 AdaFisher training epochs as cutoff.
Network
Adam
AdaHcssian
K-FAC
Shampoo
AdaFisher
RecNe150
53.06
50.21
50.05
53.53
57.41 48 86
Big Swin
48.11
8.89
4.11
32

---

Published as a conference paper at ICLR 2025
D.2.4 RESULTS
[2 the results for the Tiny ImageNet dataset using ResNer50 and Big Swin networks, AdaFisher and AdaFisherW consistently outperform current SOTA optimizers. Notably, Figure 14 illustrates that although AdaFisher converges slower than K-FAC during ResNet50 training. it achieves superior generalization. This is evidenced by lower testing errors, suggesting that AdaFisher tends to converge to a flatter local minimum, enabling smoother transitions between training and testing datasets with minimal generalization loss. For further explanation, see|Cha et al. (202T). Note that due to AdaHessian's 's high memory consumption, we were unable to train it on Big Swin. Table 13 presents the performance of various networks on CIFARION 10/00 datasets using different optimizers, both with and without the cutout augmentation technique. AdaFisher and AdaFisherW consistently outperform their counterparts in both scenarios, demonstrating stable training and robustness to the augmentation techniques. The training losses and test errors for the CIFAR experiments, both with and without cutout, are visually represented in Figures 16. 17 18 and 19 Moreover, Figure 20 shows the training and validation error of transfer learning task. Finally, Figure ( [l]itusttttes the training and validation error of the distributed version of AdaFisher on ImageNet- Ik across various batch sizes. AdaFisher not only outperforms its counterparts with smaller batch sizes (256), but it also continues to achieve superior generalization as batch sizes increase. Furthermore, these results reinforce the stability analysis concerning batch sizes presented in Section 4.4 extending it to a more challenging dataset.
with visualizations provided in Figure T4
Table 13: Performance metrics (mean, std) of different networks and optimizers on CIFARIO 10 and CIFAR 100 using batch size 256 (a) without Cutout and (b) with Cutout. Reported using WCT of 200 AdaFisher training epochs as the cutoff.
(a) Without Cutout
CIFARIO
CIFARIO0 10
Network
SGD
Adam
AdaHessian K-FAC
[Shampoo AdaFisher
SGD
Adam
AdaHeasian
K-FAC
[Shampoo) AdaFisher
ResNet18
13.89%..33.64p...
94.056.1
24.044.44.522 99.002.
(8.422.:/22.1o..
73.646.
74.79u.a/76.
76.533. 77.1177 17.
ResNer50
2
94.26p.1
9..92o.
95.
27.50a.33.122.
75.29v.s
75.19
2 77..91..
ResNet101
13.174..33.144.1
94.736.a
94.236
78.76.
72.196.2
75.466.3
82. 79.76..3
DemsNNet/2185.11,
33..74..2
94.54..1
94.97
14.97
34.99.. 95. 29.
75.388..
72.54o. 9
77.09 a7.. 77....70./... (70. 09..33.
MobileNetV332.115
91.4g. 91.9999
91.92..1
91.91 92..8%
73.81
65.
60.78
39.877.a68.011
00..0.0....7. 77. 150.2
Tiny Swin 10.088..877.477
78.34o.2
60.844.
68,444
29.4..888.... 08
57.430
32..0.
64.126.3
33.
FocalNet
80.87o. 25.
71.03o.3
12.92o.2
2
92.
15.66p. 52..8............
880.3
38.05o.3
11.233.
06o.3 52 9
CCT-2/xx2 233.22.333995..
34.63,
35, 84.63n $2.12, o/00.14..)
88066..
9.76m 60. 60.633..
"Note that Adam and Adalishhr were used for all CNN architectures. while AdamW and AdaFFsherW were applied for all VIT experiments.
(b) With Cutout
CI
CIFARIO0
Network
SGD
Adam
AAdaHessian K-FAC
/Shampoo/Adafishhrr AdaFisher
SGD
Adam
AdaHessiar
K-FAC
SShampoo
ResNet18
25.646.1
95.44..1
17.
76.56g.275.744
71.79o.
76.03..af
777..78..777 77.28%..
ResNet50
26.71o..44.45m.22
95.54o.1
95.66..14.590,
95.66o.1
74.60..
76.81o.s
77.4 4..00.488..70.. 79. 770.4
ReaNer101
25.980.229..70...
95.29..s
96.01a. 91..3......g..333333
78.89
73.38
77.01p.477 83..2 65..4
96.11o..
96.12. 1
74.80%.
79.79. 20.0022..8816.,
92.86
94.34o.1
0 28.
73.89p
77..6...
56.584.5
73.75g. a
(500.855..37 77.56.77..66..
Tiny Swin
82.340.
156.2
64.79o.
63..91...3....... 88.. 748.774...
14.890
50.866.5
34.455.
FocalNet
26.
68..8
38.94l..
97. 87..66...7790..0 90.
47.76us
52.71o.a
32.33o.3
9.98p.s
9. 9.18 53.690.3
XCT-2/3x2 78..66. S99.
35.16p $3.0082.455.166
03..082..
84.944.a
(54.05g.4580 (59.78p.s
7.176.2
8.60g. 662.9l..5
that Adam and Adal were used for all CNN architectures, while Adamm and AdaFisher were applied for VIT experimments.
Adam
AdamW
Adafisher
AdafisherW
AdaHessian
kfac
Shampoo
mamning Tlayinapeleet
Teating Big Thsytnageleet
Training
Tinyinagelet
Thyptnageliet
1.0
15
Vnll dlscc
Click Time 20
Figure 14: WCT training loss and testing error curves of several optimizers on Tiny ImageNet dataset, ResNet- 50 and Big Swin with batch size of 256. AdaFisher consistently achieves lower test error as compared to Adam, AdaHessian, K-FAC and Shampoo. The final accuracy results are reported in Table 12
33

---

Published as a conference paper at ICLR 2025
Training loss: Resnet50 on ImageNett
Testing error: Resnet50 on ImageNet
Batch 1024 GPUS
00
Batch 1024 GPUs)
24.0
23.8
70
23.6
23.4
23.2
23.0
107
22.8
22.6
20
40
60
50
20
40
60
50
82
86
spach
spoch
Figure 15: Performance of distributed AdaFisher using ResNet50 ImageNet-Ik with different batch sizes for 90 epochs. The final accuracy results are reported in Table
Adam
AdamW
AdaFicher
AdaFisherW
AdaHessian
kfac
Shampoo
Training lear Recnetil CIAR10
Tacting Recnetil CIFAR1
Training loss: Racretso GFAR10
Teeting emarr Reenetti CIRAR10
06-
13"
000
64-
Ipiie
20-1
3000
win
1004
3013
240
2500
2500
29.1
16670
CloC 1009 13
112
1rem 34
Thatring Resrat101 CIRAR10
Testing Resnett D1 CIFAR10
Training loaa: Denseret121 CIARRO
Testing emer: Denaenet122 CITAR10
10-
63-
16-
30'1
03-
0.-
30
4000
0005
12000
3000
4000
3000
12000
2000
KD91
1783
18003
8200
1700
19201
INT Clcck Tine OF
Wr.Chlck Timc 5)
Trus 34
Traising loss: Mottlenets on CIFAR1O
Teeting emer: Nobilenettä en CHARR0
Training loss: Thay sain CITAR
Testing emor Thay sath on FARLD
2011
10-
1574
30"1
2500
5000
7582
20000
2593
3000
7500
1300
2000
B0X
15000
20003
3000
10800
15000
20000
Time
OFA510
10-1
V50 Clock
sedi Cluck Tame--i
Figure 16: WCT training loss, test eтror, for CNNs and ViTs CIFARIO experiments, without Cutout. A batch size of 256 was used, and all networks were tuned using ResNet18 applied on CIFARIO. The final accuracy results are reported Table (a).
34

---

Published as a conference paper at ICLR 2025
Adam
AdamW
AdaFisher
AdaFisherVV
AdaHessian
kfac
Shampoo
500
Training Becnettio CIRAR
Recserdo CiRAR100
20
10
Tima
Wal Dbsc
WalCCssh
Trainin leas: Denacmeti21 CIAA100
Teating crer: Denscneti21
10'
12"0
4000
A0C3
1100E
16600
4009
1330
13900
39000
Trainig Mobileneted CIRAR100
Testing Mobilenett's CBAR100
OIFR1009
OFAR100
28
Oloce
Clock
Training loos: Focalnet on CIAAI100
Teating cmron: Feceinet on CIFARD0
Training oae Ccb-2/222 GRAR100
Terting ener: Gct-2002 or OFRR100
0.5
5.5
2000
1500
10008
27009
T500
30000
2500
3043
Hell Clock Thee-
WMrCCrck Time-
Figure 17: WCT training loss, test eтror, for CNNs and ViT's on CIFAR 100 experiments, without Cutout, batch size of 256 was used, and all networks were tuned using ResNet18 applied on CIFAR10. The final accuracy results reported Table (a).
35

---

Published as a conference paper at ICLR 2025
AdamW
AdaFisher
AdaFisherW
AdaHessian
kfac
Shampoo
SGD
Ire men losa: Hesneet18 on CITAN10
Textip croo: Reanet19 CIFARLS
Training Recnetto CiFAR10
Recnetsc CHRAR10
16
200
X.09
1300
3000
160 Clic Time
1008 Clock Timee
Theniing le: Resnwt201 en CITAA1
Teating emorr: Besret201 en CINAR10
Traising loes: Demeret322 en CIARR1
Testing error: Domsenett33 on CIFAREO
12"
02
7082
4230
8030
1000
2002
$200
8300
2002
42:3
8300
2000
16P
0.-
0.4
30
WMT Cbct 1:0-11 100
Mt11 Cbclh Time
Cep Trms 011
Vtall Clocc Time- 11l
0.
WalC 7in--12
and Clock Tima- 31
Clcc Time- On
Vidl Cloce Times (t)
Figure 18: WCT training loss, test for CNNs and ViTs CIFARIO experiments, with Cutout. A batch size of 256 used, and all networks were tuned using ResNet18 applied on CIFARIO. The final accuracy results are reported Table (b).
36

---

Published as a conference paper at ICLR 2025
Adam
AdamW
AdaFisher
AdaFisherW
AdaHessian
kfac
Shampoo
Recnetis CIARR100
Testing Ressetis CIAARU00
Trairing Recnettee CIRA2199
Reanecso CHARL10
0.3
1011
30'11
1300
1830
1.10
1903
1tt Cloc TIme
158 Clock Timee 111
MBFC 61 Time
CloC Time 61
Training Resect101 CIFAR100
Testing Resset101 CHAR1100
Training Denceset121 CIRAR109
Testine Dessenet121
00
07
060
20'1
2000
1064
2003
6021
2002
6003
2002
101 Cock Time
/cock Time
W cck Tic
Mel Clock Time 11
6.5
12"
0.5
30"1
3014
0.-
1005
2000
W00 Coct 1im
#M/CChch Timm
FC 10T: Tm 03
Madl CEoc Time bl
1942
WNi Cocl Time IK
MilcC Time
Wall Clock Time (0)
WelChac Time
Figure 19: WCT training loss, test error, for CNNs and ViT' CIFARI00 experiments, with Cutout. A batch size of 256 was used, and all networks were tuned using ResNet18 applied on CIFARIO. The final accuracy results are reported Table [3 (b).
37

---

Published as a conference paper at ICLR 2025
Adam
AdamW
AdaFisher
AdaFisherW
AdaHessian
kfac
Shampoo
5G0
Trainin
Testmp srrent Rcanet50 CIPAR10
Training
Teating Basner101 CIFAR10
10
100
Will dltc
Traising Denseneti21 OFARIO
Testing Denseneti21 CIFAR1O
Train Mozilenetre CIFAR100
Testing Nobillnntta CIRAR100
247
33
6x12-1
1993
300
Coct Thmee
U00 Coock True
V6B OAOC TIME 11
1A3 Cook TINO 81
Training Ieet: Receat50 on CIRAR000
Tecting errort Recret5o on CIRA100
Traising F6GS1 Resnotbo CINAR100
Tecoig RoGnECSO CIPARI0O
0.7
309
660
300
1229
305
603
1200
32.00
1000
525
1209
15.00
Waining Moc Desceneti21 CIAR1100
Tecting emeri Dencenet121 on CBAR109
Treining less: ResnetSBL on CHARR10
Tesking enerr: Resnet101 en CIFAR108
0.5
06
022
39
32.5
1200
1502
20.0
1500
2.00
800
Figure 20: WCT training loss and test error for CNNs on CIFAR-10/100 experiments with pretrained weights on ImageNet-Ik. hatch size of 256 was used, and all networks tuned using ResNe150 applied CIFAR1O. final accuracy results are reported Table
38

---

Published as a conference paper at ICLR 2025
D.2.5 COMPARISON WITH OTHER RELEVANT METHODS
In this section, we compare AdaFisher with six baseline optimizers for image classification: SGD, Adam/AdamW, AdaHessian, KFAC, and Shampoo. These baselines were selected because they either represent the current state of the art or utilize second-order gradients, making them suitable comparisons for evaluating second-order optimizers. However, other optimizers, such as AdaFactor Shazeer & Stem (2018) and EVA Zhang et al. (2023), are also relevant in this context. AdaFactor is optimizer that approximates second-orde moments using row and column factorizations, reducing memory consumption for large-scale models. EVA is a second- order optimizer designed to leverage the FIM with efficient matrix inversion techniques. Therefore, we experimentally compare AdaFisher against the optimizer baselines, including Eva and AdaFac- tor. Regarding the HPs for EVA, we used the optimal values reported in its original paper and trained the model for 119 epochs using the WCT technique. For AdaFactor, we fine-tuned the learning rate as described in Appendix p.2.1 identifying 0.001 as the optimal value, and trained the model for two distinct models: ResNet- 18
an enhanced Adam memory-effice
216 epochs. Figure 21 illustrates the performance comparison with CIFAR-100 and MobileNerV3 with CIFAR1O. The same data augmentation techniques were applied across all experiments, detailed in Appendix D.2.2] The best test accuracies achieved are summarized in Table [4] AdaFisher demonstrates superior performance compared to the new optimizer baselines, outperforming both EVA and AdaFactor.
Table 14: Performance comparison of AdaFisher and other optimizers using ResNet-18 (CIFAR 100) and MobileNet-V3 (CIFARIO). Reported using WCT of 200 AdaFisher training epochs the cutoff.
Network Opptmmzzrr
SGD
Adam
AdaFactor
AdaHessian
K-FAC
Eva
Shampoo
AdaFisher
MobileNet-V3 ResNet-18
94.43
93.32 75.74
93.21 69.45
92.86 71.79
94.34
94.41 76.69
93.81 76.78
95. 28
76.56
76.03
77.28
Adam
Adami
AdaFsserr
AdaFishenW
AdaMessian
Shampoo
AdaFactor Testing eneer: Remnet38 CPAR100
Tratning
Mustlenetu3 CIPAR10
Testing ererr Mobilenett3 CIFARLO
Training Resnet18 CIAR110
0.3
10a1
3.1
830
2000
2403
3000
1000
1000
2400
.00
1230
1002
1200
1890
Wal Clack Time- 234
stal Click
Vall Clack Time
Wal Click Time 18
Figure 21: WCT training loss, test error, for ResNet-18 18 on CIFARIOS 100 and MobileNet-Vi V3 CIFARIO. A batch size of 256 was used. The final accuracy and training time results summarized in Table 14
Performance comparison of AdaFisher and other optimizers using (a) ResNet-18 and (b) MobileNet- V3 on CIFAR100 for 200 epochs.
Table 15:
(a)
ResNet-18 18
Optimizer
SGD
Adam
AdaFFactrr
AdaHessian
K-BAC
Shampoo
AdaFisher
Test Ace
76.52
75.71
69.78
76.86
76.96 46.46
77.08 43.18
77.35 216.67
77.28 26.58
Training Time (min)
20.03
23.33
21.67
96.67
(b) MobileNet-V V3
Optimizer
SGD
Adam
AdaFactor
AdaHessian
K-FAC
Eva
Shampoo
AdaFisher
Test
73.42
70.53
71.08 54.22
62.36
75.
75.48 96.78
70.65 487.21
77.56
Training Time (min)
50.03
56.63
206.28
116.86
60.12
(c) ResNet-50
Optimizer
SGD
Adam
Adalfactor
AdaHessian
K-FAC
Eva
Shampoo
AdaHisher
Test Ace
76.12
73.03
70.78 73.32
76.18
77.66
78.01
78.89 583. 11
78. $1.91
Training Time (min)
70.
76.67
502.28
149.36
138.58
83.02
D.2.6 COMPARISON WITH CONSISTENT EPOCH COUNTS
We evaluated AdaFisher and its counterparts, including two prominent optimizers, Eva and Adafac- tor, over 200 epochs R8 ResNet-50 and MobileNet-V3 using the CIFAR101 100 dataset.
39

---

Published as a conference paper at ICLR 2025
Figure 22 illustrates the training loss and test enror trends over epochs, along with the best test error achieved as a function of training time per epoch for all optimizers across both models. Table [5] summarizes the highest test accuracy and total training time for each method on both network chitectures. Notably, while Shampoo achieved marginally better test accuracy than AdaFisher on ResNet- 18, it required approximately eight times longer training time. Conversely, AdaFisher out- performed all baseline optimizers, including Shampoo, in the MobileNet-V3 and ResNet-50 exper- iments, achieving superior test accuracy while maintaining high efficiency comparable first-order optimizers.
Adam
AdamiW
Adafisher
AdafisherW
Adallessian
Kfac
Shampoo
SGD
AdaFactor
Testing
C4
12"4
12-4
Training losa: Beanet50 on CHAAB100
Tasting emeert Remnet50 CIFR110
Restiet-50 NobileMet--3
17.5
Restitt 30
NobilcNat-113
C3
87.5
27.5
22.5
100
Sccands Epacc
Performance comparison of AdaFisher and other well-finetuned optimizers their best perfor- mances using ResNet-18 and MobileNet-V3 on CIFAR- 100 for 200 epochs. A batch size of 256 was used. The final accuracy and training time results are summarized in Table TS
Figure 22:
D.2.7 COMPARISON OF TRAINING SPEED AND MEMORY UTILIZATION
As discussed in Section AdaFisher emerges as a balanced orf between time complexity and performance. Similarly, its memory footprint is comparable to that of Adam, showcasing ef- ficient VRAM utilization. We extend our stability analysis to the CIFAR10 dataset to provide dataset-independent evaluation of performance metrics, as depicted in panel (A) of Figure 23 Ad- ditionally, analyze the memory usage for different batch sizes using the ResNet-50 model on the CIFAR-10/100, presented in panel (B) of Figure 23 The analysis reveals that AdaFisher while maintaining high accuracy levels, uses memory comparably to Adam, especially evident in larger batch sizes. This suggests that AdaFisher can achieve competitive performance without excessive VRAM consumption, making an optimal choice for scenarios with memory constraints.
Best Test Accuurcy vo. BMtck Size
Best Text Accuncy vs. Time per Spech
Bast TosT ACCurRIy NE. Learding Rata
Brch 120
Rosh DE1
CIFAR1O
CIFAR100
Batch Stzee
Time Spoch becondal
Leaming
Figure 23: (A) Performance comparison of AdaFisher and other optimizers acTOSs various batch sizes, epoch times and learning rates (with a batch size of 256), evaluated using the ResNet50 on the CIFAR-10. (B) Per- formance comparison of AdaFisher and other optimizers regarding the memory used, assessed using ResNet50 and CIFAR 10/100 across different batch sizes. This figure highlights how AdaFisher competes closely with Adam in terms of memory efficiency and performance.
40

---

Published as a conference paper at ICLR 2025
Adam
AdaFisher
AdaHessian
kfac
Shampoo
(A)
ResNet18 on CIFAR10
ResNet50 CIFAR10
ResNet101 on CIFAR10
DenseNet121 on CIFAR10
Time (G)
307
203
MobileNetV3 on CIFAR10
Tiny Swin CIFAR10
FocalNet on CIFAR10
CCT-2/3X2 CIFARIO
340
Epocc Time (6
250
(B)
ResNet18 CIFAR100
ResNet50 CIFAR100
ResNet101 CIFAR100
DenseNet121 CIFAR100
Eposh Time (b)
8
MobileNetV3 on CIFAR100
Tiny Swin on CIFAR100
FocaiNet on CIFAR100
CCT-2/3x2 on CIFAR100
100
254
Epoch Time
200
700
100
200
330
300
30
Figure 24: Epoch times for various networks on CIFAR10 (A) and CIFAR100 (B) using Adam, AdaFisher, K-FAC, AdaHessian and Shampoo.
Epoch Times. Continuing our analysis of the time complexity for each optimizer, we present in Figure 24 the epoch times for various network architectures and datasets. Specifically, compare the epoch times of Adam, AdaFisher, K-FAC, AdaHessian, and Shampoo optimizers on CIFARIO and CIFAR100 datasets. As depicted in Figure 24 panel (A), AdaFisher demonstrates comparable training time to Adam across multiple network architectures on the CIFARIO 10 dataset. This indi- cates that AdaFisher achieves efficient optimization without incurring significant additional com- putational costs. Similarly, in Figure 24 panel (B), we observe that the epoch times for AdaFisher remain close those of Adam on the CIFARIOR dataset. While K-FAC and AdaHessian exhibit increased training times. Shampoo shows the highest epoch times across all tested networks. This further highlights the efficiency AdaFisher as an optimizer, combining the advantages of advanced optimization techniques with practical training times.
D.3 LANGUAGE MODELING
D.3.1 DATASET DETAILS
The Wikitext-2 dataset, derived from high-quality Wikipedia articles, contains over two million words and is structured into training, validation, and test sets. widely used for benchmarking language models in natural language processing, especially assessing perplexity to evaluate pre-
41

---

Published as a conference paper at ICLR 2025
dictive performance. This dataset offers balance between computational efficiency and linguistic complexity, making it ideal for practical language model training and evaluation.
D.3.2 NETWORK DETAILS
Network. We utilize streamlined GPT-1 architecture, which incorporates four self-attention lay- ers, reduction from the original twelve. This configuration retains core modeling capabilities while reducing complexity, encompassing total of 08,351,448, learnable parameters. Embeddings & Parameter Sharing. To expedite training, we employ pretrained embeddings from OpenAT's GPT, leveraging the benefits of parameter sharing for enhanced efficiency and faster con- vergence.
D.3.3 HPS
The model underwent training for 50 WCT epochs using AdaFisher on the WikiText-2 and PTB datasets, with the final epoch counts for each optimizer detailed in Table 16 For Adam W, we follow
Table 16: Final selected epoch counts for various optimizers across language modeling task AdamW
AdaHessian
Shampoo
AdaFisherW
55
18
12
50
the learning rate setting in ElNorrahhy et al. (2022). For the other optimizers, select the learning rate by doing a grid search of {0.3, 0.15, 0.1, 0,05, 0.03, 0.015, 0.01,..., 1e- 1e= We tabulate the learning rate that
in Table 17 The batch size was configured to 32, and the weight decay was
established at 0.1. Despite optimizing the configuration of HPs, Shampoo failed to converge, and K-FAC could not be trained at all.
Table 17: Final selected learning rates for each optimizer, tuned using GPTI on WikiText-2 and PTB using batch size of 32. We selected based on final validation PPL.
AdamW
Adallessian
Shampoo
AdaFisherW
Se
0.015
1e-4
D.3.4 RESULTS
Figure 25 displays the training loss and testing curves, clearly showing that AdaFisher passes both Adam and AdaHessian in performance on the WikiText-2 and PTB datasets.
Training
Testing PPL: Small
Training Loas:
on WildTest-2
Teating PPL
Adairi Adaltessian Adolfsheer
Adamii
Adphessian
AdaFstank
Vall
Clock
Cleck Time
Figure 25: Training Loss and Test Perplexity of Small GPT-1 Model on WikiText-2 and PTB Datasets. Ex- periments were conducted using a batch size of 32 and optimal settings for all optimizers.
42


Reviews:
[{"review_id": "3Q9i5XrxSb", "reviewer": "Reviewer_JcNA", "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.", "questions": "Does your shampoo baseline use grafting? \nCASPR is another paper which uses Kronecker sum based approach, what are your thoughts on Kronecker sum based combination to approximate empirical Fisher information matrix.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Kronecker factored preconditioner with diagonal factors hasn't been used with empirical Fisher matrix before. \n2. 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. \n2. The Adafactor already uses Kronecker factored preconditioner with diagonal factors. There is no comparison against Adafactor. \n3. sub-optimal regret bound $O(\\log(T)\\sqrt{T})$, compared to Shampoo - which is optimal $O(\\sqrt{T})$.\n4. There are low rank approaches with similar complexity as proposed methods  such as EVA [1]. There should be comparison against EVA.\n\nIf 2, 3, and 4 are addressed, I am willing to increase my score.\n\n[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.", "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": "", "extra_scores": {}}, {"review_id": "JUmoI0STAk", "reviewer": "Reviewer_5WtX", "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. \n\nMy 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.", "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.\n\n2. 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.", "limitations": "", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "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.", "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": "", "extra_scores": {}}, {"review_id": "4r9hxdBMEJ", "reviewer": "Reviewer_jye4", "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).", "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?\n- 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?\n- It is known that the SGD optimizer performs better than the Adam optimizer in image classification tasks. Why is there no comparison with SGD?\n- 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?\n- 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.\n\nI'm willing to improve my score if you address my concerns.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "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.", "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": "", "extra_scores": {}}, {"review_id": "vrvuHxFH68", "reviewer": "Reviewer_K3HY", "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.", "questions": "1.\t**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?\n\n2.\t**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.\n\n3.\t**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. \n\n4.\t**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?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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. \n\n2. **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.\n\n3. **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.\n\n4. **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. \n\n2. **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.\n\n3. **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).", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)