A Computationally Efficient Sparsified Online Newton Method
Authors: Fnu Devvrit
Sai Surya Duvvuri
Rohan Anil
Vineet Gupta
Cho-Jui Hsieh
Inderjit S Dhillon
Venue: NeurIPS 2023 poster (2023)
Keywords: Optimization
Second order methods; Deep Learning
TL;DR: We propose a novel computational and memory efficient second order method for deep neural network training which while maintaining computational efficiency.

Abstract: Second-order methods hold significant promise for enhancing the convergence of deep neural network training; however, their large memory and computational demands have limited their practicality. Thus there is a need for scalable second-order methods that can efficiently train large models. In this paper, we introduce the Sparsified Online Newton~(SONew) method, a memory-efficient second-order algorithm that yields a sparsified yet effective preconditioner. The algorithm emerges from a novel use of the LogDet matrix divergence measure; we combine it with sparsity constraints to minimize regret in the online convex optimization framework. Empirically, we test our method on large scale benchmarks of up to 1B parameters. We achieve up to $30\\%$ faster convergence, $3.4\\%$ relative improvement in validation performance, and $80\\%$ relative improvement in training loss, in comparison to memory efficient optimizers including first order methods. Powering the method is a surprising fact -- imposing structured sparsity patterns, like tridiagonal and banded structure, requires little to no overhead, making it as efficient and parallelizable as first-order methods. In wall-clock time, tridiagonal SONew is only about $3\\%$ slower per step than first-order methods but gives overall gains due to much faster convergence. In contrast, one of the state-of-the-art (SOTA) memory-intensive second-order methods, Shampoo, is unable to scale to large benchmarks. Additionally, while Shampoo necessitates significant engineering efforts to scale to large benchmarks, SONew offers a more straightforward implementation, increasing its practical appeal. SONew code is available at: https://github.com/devvrit/SONew

A Computationally Efficient Sparsified Online Newton Method
Devvrit +f
Sai Surya Duvvuri'
Department of Computer Science
Department of Computer Science
The University of Texas Austin devvrit. 03@gmail.com
The University of Texas at Austin mubramanyandvss@ggme com
Rohan Anil
Vineet Gupta
Cho-Jui Hsieh CS Department, UCLA Google chohsieh@cs.c www.uclll.e..... edu
Google Deep rohamanil@google.com
Google
vineet@google.com com
Inderjit Dhillon
Google
isd@google.com
Abstract
Second-order methods hold significant promise for enhancing the convergence of deep neural network training; however, their large memory and computational de- mands have limited their practicality. Thus there need for scalable second-order methods that can efficiently train large models. In this paper, we introduce the Sparsified Online Newton (SONew) method, memory-efficient second-order al- gorithm that yields a sparsified yet effective preconditioner. The algorithm emerges from a novel use of the LogDet matrix divergence measure; we combine it with sparsity constraints to minimize regret in the online convex optimization framework. Empirically, we test our method on large scale benchmarks of up to IB parameters. We achieve up to 30% faster convergence, 3.4% relative improvement in validation performance, and 80% relative improvement in training loss, comparison memory efficient optimizers including first order methods. Powering the method is a surprising fact imposing structured sparsity patterns, like tridiagonal and banded structure, requires little to overhead, making as first-order methods. In wall-clock time, tridiagonal SONew is only about 3% slower per step than first-order methods but gives overall gains due to much faster convergence. In contrast, one of the state-of-the-ee (SOTA) memory-intensine eccond-order methods, Shampoo, is unable scale to large benchmarks. Addition- ally. while Shampoo necessitates significant engineering efforts to scale to large benchmarks, SONew offers practical appeal. SONew code is available at: https:/giihuh.com//vvvtttttSCC
efficient and parallelizable
straightforward implementation, increasing its
Introduction
Stochastic first order methods which use the negative gradient direction to update parameters have become the standard for training deep neural networks (DNNs). Gradient-based preconditioning involves finding update direction, by multiplying the gradient with preconditioner matrix carefully chosen from gradients observed in previous iterations, to improve convergence. (Full-
"quul contribution, Work done while Google
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

matrix) Adagrad [13], online Newton method 125] and natural gradient descent [3] use a full-matrix preconditioner, but computing and storing the full matrix is infeasible when there are millions of parameters. Thus, diagonal versions such as diagonal Adagrad, Adam 33 and RMSprop
(28 are
now widely used to train DNNs due to their scalability.
Several higher-order methods have previously been applied to deep learning ((24) 5 23 38) ). All these methods use Kronecker product factorizations that reduce computational and storage costs to make them feasible for training neural networks. However, to precondition a dj x d2 parameter matrix, these methods require matrix inverse operations, which take O + di) time and Od + dī) space. In comparison, first-order methods use ((11d)) time and memory, which is linear in the number of parameters. For instance, when d1 kdg, the memory used by Shampoo, of + d2 floating point numbers is (K) times the number of parameters, which could be arbitrarily large depending on k. This calls for further research in developing efficient second-order optimization techniques to train DNNs with memory and time complexity linear in the number of parameters.
In this paper, we present a novel Sparsified Online Newton (SONew) method, which only requires linear time and space complexity, to train large-scale DNNs. We derive the algorithm through two steps, classical regret analysis followed by sparsification step. In more detail, regret analysis when using a preconditioner reveals that the error is bounded by two terms, the first depends on the change in the preconditioning matrix, while the second depends on the generalized gradient norm (see Section 3 for more details). We take novel approach of minimizing the second term while regularizing two successive preconditioners to be close the LogDet matrix divergence measure (34 (see Section 3 for the intuition behind choosing LogDet divergence). This analysis naturally yields us an Online Newton method (25). To make it computationally efficient, we further sparsify the preconditioner by finding sparse approximation that is close in LogDet divergence. Thus we are consistent in using the same measure (LogDet divergence) in both the regularization and sparsification steps. This gives us our SONew method, which achieves linear complexity by leveraging structured sparsity patterns, such as tridiagonal and banded, in the preconditioner. This unlike most existing online Newton methods that require quadratic space and cubic time complexity. By making each step linear time, the SONew method can be applied to train modern DNNs as efficiently as first order methods. Further, our method is embarrassingly parallelizable thus making negligible the overhead of computing the preconditioner. We also show that introducing sparsity allows us to reduce the condition number of the problem dynamically to improve numerical stability.
We strengthen the relationship between sparse LogDet divergence minimization and online convex optimization by establishing an optimal (VVT) regret upper bound for tridiagonal sparsity pattern. In our experiments on an MLP Autoencoder and Graph Neural Network (GNN), we found that our method outperformed first-order methods terms of training loss within the same training time, while Shampoo (second-order method) takes significantly longer. In our experiments on Vision Transformers on Imagenet and GNN on OGBG-molpeba, we achieve a target validation performance using 10% and 30% fewer iterations respectively compared to Adam, the SOTA optimizer for both benchmarks. Furthermore, using the same number of iterations as Adam we observe 0.7% and 3.4% relative improvement for ViT and GNN respectively in validation performance. From an optimization point of view, SONew achieves 9% and 80% better relative training loss for VIT and GNN respectively. It is worth noting that Shampoo statistics required ~ 7 x #params for VIT whereas tridiag-SONee uses only 2 x #params for its statistics. We also test another recently proposed memory efficient second order optimizer, rfdSON (37), but found its performance suboptimal the best performing first order method. Owing to SONew's 's scalability. we train a Large Language Model (LLM) with 1 billion parameters and compare it with AdaFactor 45). a popularly used first order optimizer to train LLMs [II. SONew achieves the same performance as AdaFactor using 26% fewer steps, resulting in a 1.35x faster training. When using the same number of steps, SONew obtained a 1.7% relative better train loss. In terms of implementation, SONew is just a few lines of code without complex engineering challenges, rendering it even more useful and practical.
(Equation [3]
2 Background
The inner product between matrices is defined as (A, = Tr(ATB), where Tr(.) denotes the matrix trace. The Frobenius norm of a matrix A is || -VTr(ATA) while its spectral norm is ||.Al|2 = max, ||Ax/|/// / ||س||| We use 7, E R2xxn to denote an identity matrix. We use Sn, S to denote the set of symmetric, and positive definite matrices respectively. The generalized norm of

---

Ar. We use det (A) to vector z E IR" with respect to matrix A E S is defined as || VI Ar.
denote the determinant of matrix A, and (Aag(4) to denote the diagonal matrix with diag( ( is Ag. We use G and G to denote a graph and its sub-graph with vertex set [r|] = {l...., n). Let Eig denote set of edges in graph G. and (t) denote neighbours of vertex i in graph G. A sparse symmetric matrix A E In*x" follows a sparsity structure graph G if As,j = v(i,j) j) & Eg. Note denote the set of positive definite matrices with sparsity structure given by graph G, ie, if X E (g) ++ then Xi = v(i, j) & E(g). (,) ++ is open convex set. Given an index set {i1, d2, .. In use Alr to denote the corresponding principal sub-matrix A.
that set of all such matrices form a linear subspace. We use Sn (g) ++
2.1 LogDet matrix divergence
Let o + - R be a strictly convex, differentiable function. The Bregman matrix divergence between X, E + is defined as 8 34] D Y) x(X) o(Y) Tr(Vo(Y)) (X Y))- Since é is convex, Dx(X,Y) 2 for all X, E 5t For example if ((X) = || the corresponding Bregman divergence (x,)) |X- - -YI
f is the squared Frobenius distance. In this paper, we extensively use the convex function ((X) llg det (X); the corresponding divergence measure ( X, Y) is called the LogDet matrix divergence:
(xa(X, Y) = - llg det (XY-1) Tr(XY n.
(I)
The LogDet divergence is scale invariant to invertible matrices A, i.e. A
Dra(X, Y). LogDet divergence can be written in terms of eigendecooppoitions of X VEVT and Y UOUT 34
(X,)
(2)
These two properties are later used in Section 3 to highlight the significance of LogDet divergence in our algorithm.
3 SONew: Sparsified Online Newton Method
We now present our proposed algorithm SONew.
3.1 Regret minimization via LogDet divergence
We set up our problem under the online convex optimization framework (OCO) round the learner makes a prediction wy and receives a convex loss f((w) and gradient = V (w,) feedback. The goal of the learner reduce regret Rr by predicting w, that a low aggregate loss E/= fe(wr) is achieved compared to the best possible, w = arg min I Se(w). Formally, regret is given by
43 26
where at each
T Ry (r WT) = Σ f((w)) - Σs((m))
T
Using [10, R regret in online setting yields R/T convergence rate in the stochastic setting. To
upper bound this regret, we proceed as in (26) by analyzing the error in the iterates for the update
We+1 = W2 - 77 Ni9t, where X, E Raxn. Then ||0++1 - w*|l w'||x,,- = || - -7 - w ||x =
I|w; w x + g6" X.g: 2n(w; - w w gt- The convexity of f. implies that f((w:) - f (w)) S
(w: - = gt leading to fe(w;) - fi(w)) < 2n (Ilwa - w w1 - ||w--1 - w" 11 + 17² 9t Xege). Summing over all / [T] and rearranging reveals the following upper bound overall regret:
T of X.9. +
T
1 -| |ww RT < 2m
w= 11x
+
1
Σ
(w, w (x, -x,1) w²). (3)
2
2n 2
t=l
3

---

Since w" is unknown, finding X which minimizes is infeasible. So to minimize regret, we attempt to minimize the second term in 3 while regularizing X; to be "close" to x The nearness measure we choose the LogDet matrix divergence, thus leading to the following objective
X = arg min of Xgt, such that Dea (X, X(-1) S C22
xest+
where Dea is as in (T. Why do we use the LogDet divergence? From ( due to the term A. (X,, priontizess matching the smaller eigenvalues of X,-1 with Those of X, i.e. matching the larger eigenvalues of X and X- As consequence, LogDet divergence regularizes X by matching up its large eigenvalues with those of Xx-1. For example if smallest and largest eigenvalue of X,-1 are On and 81. then for an eigenvalue of X, when o > 6.. 01. the penalty from ( for 6,, is higher than for 01. (0/0,, log( (0,) - 1) > (a/01 log(a /01) ) - 1). This intuition leads us to formulate as our objective. We recall that there is precedence of using the LogDet divergence in the optimization literature; indeed the celebrated BFGS algorithm [9] 17 22 can be shown to be the unique solution obtained when the LogDet divergence between successive preconditioners. subject to secant constraint, is minimized (as shown in the 4apa paper by (18).
The optimization problem in (4) is convex in X since the LogDet divergence is convex in its
first argument. The Lagrangian L(X, Ax) = of Xgt + >e(Dea(X, Xt-1) ct) = Tr(Xg:g(
A det (XX;-1) + Tr(Xx - n)) - dect- Setting V2(X,, Xx) = 0, and using the fact that
v log det (X) = X we get the following update rule:
- + 9/95/Ar.
(5)
We emphasize that the update rule (5) arises naturally from our novel use of LogDet divergence to minimize the regret. Moreover, Equation 5 can be seen as general update rule applicable to
numerous existing optimizers. For example, setting G = 0 (equivalently A, = oo) AI E [n] in (4) results in no change to the preconditioner in any round. In this case, with X, In. we get online gradient descent [54]. On the other hand, setting X = 1 gives the update rule of the online Newton method 25 Our update rule differs from (full-matrix) Adagrad 15) which has X, X + 9:9
Maintaining and updating X, as in (5) is possible by using Sherman-Morrison formula but requires O(n²) storage and time complexity. This becomes impractical when n is in the order of millions which is typically the case in DNNs.
3.2 Sparsifying the Preconditioner
To minimize the memory needed for maintaining and updating X, using 5), we adopt the strategy of sparsifying the preconditioner. For existing optimizers such as (full-matrix) Adagrad or the Online Newton method, it is unclear how to sparsify given preconditioner. Specifically, there is no intuitive approach to assessing the quality of a sparse preconditioner compared to a full-matrix preconditioner. However, since our update rule (5) originates from using LogDet divergence in the regret bound analysis, it gives us a natural metric to measure the quality of a sparse preconditioner. Let's 's consider the following problem: find sparse positive definite X with [| I an, a 1, such that the objective Dea X, (X,) + constraint while requiring the sparse preconditioner to remain close to the full-matrix preconditioner in terms of LogDet divergence.
9:96 / / is minimized. Essentially, this problem imposes a sparsity
Due to the normorm constraint, this is a non-convex problem, which makes it difficult to solve exactly. Since L1 norm serves convex relaxation for the Ly norm, we could use it instead, resulting in the following optimization problem also known as graphical lasso estimator [19]
min Dea (X71 + 9:97 />x)=)) + + ||
xes.+
However, the time taken solve the above problem, even with the current best methods 7 29 16
53
can still be too large (as these methods take several minutes for a matrix of size million), making it impractical to embed in DNN training.
In this paper, we take a different direction where we use fixed sparsity pattern constraints, specified by fixed undirected graph G. To sparsify the solution in (5), we formulate the subproblem
X arg min Dea (X, (X + 9:95 /Ae)
(6)
(+)+

---

where (g) denotes the set of positive definite matrices with the fixed sparsity pattern correspond- ing the adjacency matrix of graph G. Note that both steps (4) and (6) use the same LogDet measure.
Owing to the structure of LogDet divergence, ( can be surprisingly solved in (n)) and easily paral- lelizable, for certain sparsity structures G. Algorithm[Tand 2redeeeennttt an instantiation of the proposed SONew method, which solves (6) using (n) time and memory for banded matrices with band size b. In particular tridiagonal matrix, corresponding to chain graph, is banded matrix with bandsize 1.
Algorithm 1 Sparsified Online
Newton
Algorithm (n) flops
in
(SONew) Algorithm
Inputs: = coefficient the update (10).
Inputs:H E (9))), is as (T))
g = sparsity graph (banded/ridiagonaa),
G = the banded graph of band size <
= damping parameter,
Outputs: lower triangular banded L E R
T == total number of iterations/mint haaiiiineemmiiiaabbeess..
and diagonal matrix E R=xa
?) step size/learning rate.
function SPARSIFIED INVERSE(H,
Output: wr-1
0, ;=
1: Ho = cla, W1
3:
L5j == 1, Vjj E [x]
2: for e (l..... T do
4:
for j E (l.....vI) do
parallelizable
3:
compute = V (e(ws)
5:
Let Hyg and HiG5, be defined where 1
Polgig /Xe) (9) with
as in Section + 1.....J + [n].
Pg as in
b (n) time & memory
5: Get L, D SPARSIFIED INVERSE (H,
6:
Solve for L1s5 in the linear system H1, Liy5 Hi D O(b") time.
where X; = LDL solves
Compute descent direction Ne LDD." gb.
7:
D3s = 1/(H,) + 21
7: W1++1 = we - U/
8: end for
8: end for
return L,
9:
return WT+1
10: end function
Maintaining H, E (9) in line 4 Solving the subproblem in 6 naively is impractical since X a dense matrix. However, the structure of the LogDet divergence comes to the rescue; the optimization problem in (6) can be expanded as follows:
arg min -log det (X)+T + 9 /Ae)). (G)++
Let us define the projection onto (,( Pg Ruxn -> R as:
My if (i,j) Eg. to otherwise.
Pa(
(8)
Note that the Tr(.) term in 7 is dependent only on the non-zero elements of X E (g) ++ since Tr(AB) (A, B). for symmetric matrices A and B. Hence, 7 can be written
arg min - log det (X) + (X, Po + 9:9 />+)), (g)+
(9)
Computing the entire matrix X can be avoided by analyzing the optimality condition of 99. Let g(X)- log det (X) + Pa
Pa (X,1 + 9.9. /A:)) denote the objective function in (9), then the
optimality condition of 9 (Vg())) Ps (V( log det()) (X, P& + + 2 /x.)))
since gradients with respective nonzero entries of X should be zero, 8X.,; (gx(g(X)))
0g(X)
v(i.j) E Eg. Using V log = - vx((X, = Y,, and setting X = X, gives:
Po(X, Po g.01 />t) 0,
H, = H:-1 + Po /Ae), where H. = Po (X
(10)
Thus we only need to maintain H1, = Po(X, This matrix is updated as H, = H4-1 Po (g /A.). Since H, E S,, the update can be done in O(EEg) memory and time, while computing the matrix X, would have cost Onn²). In SONew (Algorithm) I. this key observation
used to maintain H, in line
5

---

Computing X, in line in 9 with H, as:
Now that H, is known at every round t, we can replace + /Ne)
X
arg min log det (X) + Tr(XH,). (
(11)
For an arbitrary graph 9, solving 111) might be difficult. Theorems 3.1 and 3.2 show embarrassingly parallelizable explicit solutions to the subproblem (TT) for tridiagonal and banded sparsity patterns.
Theorem 3.1 (Explicit solution of (II for tridiagonal structures/chain graph). Let the sparsity i,j n). Also, let H S, be such that any submatrix of H corresponding to complete subgraph g is positive definite, then the solution of [1 is given by X = LDI", where the unit lower bidiagonal matrix L and diagonal matrix D have the following non-zero entries:
structure be a chain with edges Eg = (i.jj) -i 5 1,1
Haty
1341
Lys - L1+1s H
DjJ²' H5 H
jsn-
D - Hnn (12) nm
Computing this explicit solution involves conducting paralellizable operations on 2 x 2 principle submatrices (highlighted in red) of the tridiagonal matrix H to find the X as shown in the following 3 x 3 example:
H11
H12
H1
H
gi
91g2
H2i
H22
H23
Ha
H22
H23
+
9192
gā
9293
(13)
0
Hsa
Ha3
0
Ns2
NT33
0
92gs
93
0
0
Hn
0
0
1
H21
0
1lz2
71
Raz
1
0
0
Hu2
H Hys
0
0
Haz
Fas
H53
0
Conducting these operations take ((n) time and memory complexity, and similarly the descent direction can be found sequentially by X,g: =19:=/0 L g.)). which can take O(n) time complexity, due to unit lower bidiagonal structure of L. furthermore, these operations can be easily parallelized. We also generalize the explicit solution to banded sparsity structures with band size b.
Theorem 3.2 (Explicit solution of [1] for banded structures). Let the sparsity pattern G be banded matrix of band size b, i.e. E'g {(i,]) ji
< i.j S n}. For every vertex j. let = {j +1.1... 1.....J : + b). Then X, = LDL" is the solution of 11 with nonzero entries of L and D defined as follows:
Lij 1, Lr, - High Djj = -Hg Hig H1,s), 1sism.
(14)
where, H E Sx(9) any submatrix of H corresponding to a complete subgruph of G is positive definite.
Note that Theorem is special case of Theorem 3.2 when is set to and the proof for Theorem 3.2 is given in Appendix [A.I] Computing the above solution requires solving n linear systems of size b (which is small) as shown in Algorithm 2 and takes O((n b 1)b²) flops. Since b & n, the number of flops is (nn).
3.3 Regret bound analysis SONew
The following theorem establishes optimal regret guarantee 26 for SONew in the online convex optimization framework mentioned in Section 3.1
Theorem 3.3. When G = -ridiagonal/chhin graph as defined in Theorem 3.1 then selling =
VG.VT, de = Goc vi and q) P/r in Algorithm I where I| w w0112 I D2, l1ge|loo S Goc incurs a regret RT
The proof sketch involves deriving an explicit expression for entries of X, in Lemma A.2 to upper bound the term (w,- - w²) (XX x,=i) w") in regret upper bound (3.. Upper bounding ET of Xegt involves using the Loewner order X, < || || In = |l. A detailed proof sketch and is given in Appendix A.2 We note here that though the regret bound presented here for convex losses, there are connections to non-convex convergence guarantees by using OCO (online convex optimization) learners, presented in Appendix A.2.5 While our main focus is on deep neural network training, which is typically non-convex, we also conducted convex experiments in Table

---

Optimizer
Time plexity
Memory complexity
Table 1: Consider preconditioning di x d2 param- eter matrix. Time complexity of tridiag and banded SONew inverse scale linearly with number of param- eters, but, Shampoo is cubic in the dimensions of the matrix. Memory used store second-moment of gradients by tridiag-SONee -Swe can be significantly lower than Shampoo, for c-g. dj 4d2, then Shampoo takes 2 times more memory.
Adam
O(dida)
dida
(fdSON(m)
O(m² d1d))
md, d2
Shampoo
Old) + d (d1d22
(dR dg)
Kidiag-SONew
2d1d2
band-4-SONew
(d1122)
5d1d2
3.4 Numerical Stability of SONew
In Theorem 3.1 and Theorem 3.2, as mentioned, any submatrix of H, corresponding to a complete
subgraph of 9 should be positive definite, however, in practice, due to finite precision, each entry of H inherently perturbed with an error proportional to O (Cmachh): where macch machine epsilon 27
We notice in practice that the subtraction operation in Dj SJ Hjj Hf+ HJ+15+1
(line 7 Algorithm|
2), which has condition number Kisub =
|H5sl/|5ph can be high as Sj can
Hii Hi++1 Thus small perturbation
be arbitrarily low due near singular submatrices
Hi+14
H141i+1]
in Il can lead to high perturbations in the preconditioner X. We formalize this notion by deriv- ing an end-to-end componentwise condition number (pg. 135, problem 7.11 in (77) of SPARSI- WIED_INVERS in Theorem A0 Appendix A.33] To reduce this condition number upper bound and be robust perturbations in Hr caused by finite precision, for a tridiagonal graph G, we can remove edges (j.j + 1) which correspond to low Siy < %, where y 2 O denotes tolerance parameter. We show in Theorem Appeendpi [A.3] that this reduces the condition number up- perbound of SPARSIFIED INVERSE. Furthermore, we generalize this to banded sparsity pattern in Algorithm Appendix A.3
4 Related Work
Online Newton method is second order method in online convex optimization framework with properties such as scale invariance [35] and logarithmic regrets in exp-concave and strongly convex functions 23 26]. However, it has time complexity of O(n)), making it infeasible for large However, introduction of LogDet divergence measure SONew allows us to set different sparsity graphs such banded graph with band-size b, for which our preconditioning process is more computationally efficient with a time complexity of (b² b b+ 1)) compared to online-newton method (n²)).
Shampoo 24 5 approximates full gradient statistics matrix using Kronecker factored preconditioners to reduce the memory and time complexity from Onn²) Odd di) and Od de) respectively. Here, n = d1dt d2 denotes number of parameters for a linear layer dimensions d1 x d2. The time complexity of matrix inversion takes heavy toll in Shampoo's o's compute time even with the Kronecker product assumption on the preconditioner, whereas, our method has a time complexity of (bb))12)) quadratic in dimensions of the linear layer (note that = 1 for tridiagonal structure).
KFAC 38 similar to Shampoo, uses Kronecker factored preconditioning, but to approximate the Fisher-information matrix. FishLeg (20) instead approximates the inverse Fisher matrix directly by expressing it in terms of the solution to an optimisation problem. Both these methods have memory and time complexity similar to Shampoo. In this work, we compare with Shampoo among the class of Kronecker factored optimizers due to its widespread testing and adoption within the community [46]. We also point the readers to Eva [52) concurrent work aimed at devising memory efficient optimizer by maintaining rank one matrix approximation to the Kronecker factors of KFAC matrices. For completeness, we include comparison with KFAC, FishLeg, and Eva on Autoencoder benchmark.
There prior work [35 36) in reducing the complexity O(n?) flops of Online Newton Step (ONS) to Onz) flops using sketching. These ONS variants maintain low rank approximation of He (as Algorithm [ and updating it with new gradient d at every iteration requires conducting SVD /Gl/orhonormalizaiii
(35] of a tall and thin matrix in Raxr, where denotes the rank of approximation of H,. In Section we conduct large scale experiments and compare SONew against rfdSON (37] as it's more stable than Oja-SON [35].
7

---

Table 2: float32 experiments on Autoencoder benchmark. We observe that diag-SONew performs the best among all first order methods while taking similar time, tridiag and band-4 SONew perform significantly better than first order methods while requiring similar linear space and time. Shampoo performs best but takes O(d time for computing preconditioner of linear layer of size di xd2, whereas our methods take ( (d)dd) time, as mentioned in Section 3.3
rfdSON takes similar space SONew but performs considerably worse.
Optimizer
First Order Methods
Second Order Methods
Adagrad
RMSProp
Adam
ding-SONew
Shampooi200
rfdSON(1)
rfSSON(4
tridiagsSONew
band-4-SONC
Train lass
54.393
53.330
53.591
53.025
50,702
56.21
55.55
51.723
51,357
Timets)
62
62
62
63
371
85
300
70
260
LogDet problem in equation| 11 is closely related to the Maximum Determinant Matrix Comple- tion (MDMC) 49 The MDMC problem is the dual of LogDet problem (IT). and has explicit
solutions for chordal graphs [4]. Thus the explicit solutions in (14) are the same as the ones proved
in 4. Also, we noticed that the tridiagonal explicit solution has been used previously in KFAC [3
in the context of a gaussian graphical model interpretation of gradients, specifically, KFAC used a brock-tridiagonal preconditioner incorporate correlation within consecutive layers.
Experimental Results
We describe our experiments on standard Autoencoder benchmark 42 trained on MNIST dataset [12. Vision Transformer [13 on Imagenet training, GraphNetwork 6 21 on OGBG- molpeba dataset [30], and a Large Language Model 47. For all second order optimizers, we use grafting [2], a technique used to transfer step size between optimization algorithms. Specifically, given an update e1 of Optimizer-1 and V2 of Optimizer-2, grafting allows to use the direction suggested by Optimizer-2 with step size suggested by Optimizer-1. The final update is given by
[>]] 02 Grafting has been shown to take advantage of tuned optimizer step size and improve performance. For SONew and rfaSON, we use Adam grafting using Adam optimizer step size || ll with SONew/rfdSON direction 12/1|=2]. For Shampoo, we use its default RMSProp grafting. We couldn't 't find rfdSON official implementation, so we use our own implementation using which we reproduced the numbers convex losses (Appendix A.4 reported in their paper [37]
[rll
5.1 Autoencoder benchmark
Setup: We use three sparsity patterns for SONew a) diagonal sparsity, resulting in a diagonal preconditioner similar adaptive first order methods like Adam and Adagrad; b) tridiagonal sparsity, corresponding to a chain graph; and c) banded sparsity, represented by "band-k" in tables and figures for band size of k. We compare SONew against widely used first order methods including SGD (32),
SGD with Momentum [41, Nesterov (40), Adagrad [4], Adam [3 and Rmsprop [48]. We also compare with ridSON [37], a recently proposed memory efficient second order optimizer and with Shampoo [24]. state of the art rocond-order optimizer used in practice, albeit with considerable mem- ory and time requirements. Because of space constraint, we report only the best performing first order methods while include the entire set in the appendix. As previously mentioned, ridSON maintains a low rank approximation of the Online Newton method's 'ss statistics matrix Σ. 9:9!"' We observed rfdSON with adam grafting always performed better than without grafting, hence report the corre- sponding numbers. We evaluate rfdSON with rank m approximation, denoted as rfdSON(m). which requires (m 1) * params space when using grafting. For fair comparison with tridiag-SONew and band-4- we test rfdSON with m and m = 4, respectively. For shampoo, computing preconditioner at every step could be infeasible, instead it is computed every steps referred to as Shampoo(t). Section 3.3 compares time and memory complexities of rfdSÓN, Shampoo, tridiag- SONew, band-4-SONew. Note that df + d2 > 2d, d2 Vdl, d2, thus memory used by tridiag-SONew is never more than Shampoo. We use a 2. 7.72M parameters Autoencoder and each experiment is per- formed using one V100 GPU having 16 GB memory. Further setup details are given in Appendix A.4
Results: In Table 2 we observe that among first order methods, diag-SONew SONew performs the best while taking same amount of time. Increasing the number of edges in the sparsity graph to tridiag or banded sparsity with band size 4 enhances the performance further. Tridiag-SONew runs 5x faster than Shampoo at a marginal cost to the loss even when Shampoo updates preconditioner once every 20 steps. Using same space, rfdSON performs considerably worse than SONew. To test the numerical stability and robustness of SONew, reduce the precision to bfloat 16 and conduct similar
8

---

(a) VIT validation error
(b) Graph validation avg. precision
50:35m
0.299
0.281
KMtSProp
0.2515
30:30m
0.200
0.3433
6605m
Adamm
0.2265
5343m
0.20
5b53us
Shampon
0.280
0.288
Adam
0.279
0.251
vabdaatoo
RMSPhop
Momantaa
0.221
8:39
Mtgg (11.3k
Figure 1: (a) Best validation error for tridiag- triiia-SSOee Momentum, RMSProp. Adam, rfdSON, and Shampoo on (a) VIT benchmark (b) GraphNetwork benchmark. We notice that tridiag- SONew achieves same performance as Adam, the next best baseline using similar space and time, with 10% and 30% less steps/time in ViT and GraphNetwork respectively. While using the same number of steps, SONew achieves relatively 0.7% and 3.4% hetter validation error respectively. Shampoo doesn't 't fit in the 16GB memory of TPU v2 for VIT benchmark, hence couldn't ' perform hyperparameter tuning on On GraphNetwork, compared Shampoo, tridiagSONee gives similar performance while being far more memory efficient (Refer Appendix A.4.2 for more details).
60
dadoyctor 2000 2414
tratigg BoNew 2,547
BMSSPrp
Ahttgoo123)
tridiay
2it lop pyrle
1riin
14
53
2.84
32
51
100k
200k
300%
400k
500k
20
40
opachs
Figure 2: (a) Comparison of SONew with first-order op- Figure 3: Comparison of SONew and Adafactor timizers, ridSON, and Shampoo Autoencoder bench- LLM training. SONew takes 26% less steps to reach mark float32 training. We observe that SONew per- same performance as AdaFactor. Using same forms better than all first order methods, and second number of steps, it achieves ~ 1.7% relative better order methods using the same memory.
log perplexity.
experiments in Appendix [A.4.4 (Table g ) We notice that SONew undergoes the least degradation in performance compared to all other optimizers. We refer the reader to Appendix A.4.4 for thorough comparison and study of bfloat 16 experiments. In Figure 2 we plot the loss curves of all the baselines and SONew for float32 experiments. Moreover, in Appendix A.4.
Table we provide ablation
on performance of SONew with varying batch sizes.
Conparison with other baselines: We further compare SONew with KFAC B8 FishLeg 220. and
Eva (52) for completeness. Since these methods lack a JAX implementation, we adopted the authors' official Pytorch implementations. When we attempted integrate their code with our Autoencoder benchmark, the results were unsatisfactory; for instance, FishLeg recorded a loss of approximately 60.0. This notably unexpected as it underperformed Adam, benchmark that the authors themselves compared against. Given these results and to minimize modifications to the official code, decided test our optimizer, SONew, directly on their provided autoencoder architecture. We present the results in Appendix A.4.4 and notice that SONew outperforms these baselines as well by a large margin.
5.2 VIT and GraphNetwork benchmark
Setup: We compare tridiag-SONew with Momentum, RMSProp, and Adam, on VIT (~22M parameters) and GraphNetwork (~3.5M parameters) benchmark. For each experiment, we search

---

over 200 hyperparameters using 4 16 GB TPUs (v2) for each run. In order to conduct a fair comparison of the running times, we executed the optimal hyperparameter configurations on 4 32GB TPUs (v4) B1 This is because certain operations, including reshapes and transposes, are not optimized on TPUs (v2). Consequently, methods like rfdSON, Shampoo or SONew, which utilize these operations, could potentially be disadvantaged if TPPS (v2) were used, skewing the comparative results. All memory-effice -efoorceefi methods, including rádSON, first-order methods, and SONew, exhibit similar runtimes, with differences of approximately ~ 5%. For VIT, we evaluate their performance based on the same number of steps, as this also effectively compares wall-clock time. However, for GraphNetwork, we train Shampoo for 20% fewer steps to achieve comparable wall-clock time.
Results: We plot the runs that give best validation error rate (for VIT) or validation average precision (for GraphNetwork) in Figure [ Fidiag-SONew requires ~ 10% less steps to reach the same performance as Adam for VIT, and ~ 30% less steps for GraphNerwork benchmark, Training for the same number of steps, we get ~ 0.7% better relative validation error for VIT and 3.4% better relative validation average precision for GraphNerwork. On GraphNetwork benchmark, tridiag- SONew performs 1.3% relatively worse in average precision compared to Shampoo, while being 1.25x faster. On VIT benchmark, Shampoo doesn't 't fit 16 GB TPU v2. Its statistics require 155M entries (~ 7 x #porass) while tridiag-SONew requires only 44M entries (2 x #params). Hence, we could not tune it. rfdSON takes same memory but slightly more time because of its SVD computations. We also notice rfdSON performs worse than Adam on both benchmarks; we leave a thorough investigation of this behavior as a future work.
We show in Appendix A.4 that corresponding to the best validation runs, Tidiag-SONer optimizer's s training loss is also less than that of Adam. Furthermore, from an optimization point of view we also show in Appendix [A.4 that among all the 200 hyperparameter sweeps, the best training loss of is 9% relatively better on ViT and 80% relatively better on GraphNN than that of Adam. We further compare Adam and tridiag-SONew on a 248M parameter Transformer Model in Appendix A.4.4 In next subsection, we present results on a decoder only large scale Language Model.
5.3
Experiments on Language Models
Setup: Owing to SONew's ' scalability, test it on Large Language Model (LLM) 47 with I billion parameters. We compare SONew with AdaFactor (without factoring), a commonly used first order optimizer for training LLMs 51 M AdaFactor is similar to Adam except that in addition it offers benefits like "parameter scaling", which has an effect of layerwise damping of the learning rate. We defer the reader to 45 for more details. We trained the LLM for 5B tokens with batch size of 64k tokens. All experiments were performed on 16 TPU v4s. To support efficient training of large models, we implemented a sharded tridiag-SONew following model parallelism approach.
We report the experiment Figure 3 where we find that SONew beats Adafactor by large margin. Specifically, SONew achieves the same log perplexity using 26% less steps. Morcover, using the same number of tokens, SONew achieves 1.7% relative better performance on train loss, leading to 1.35x speedup. This shows the potential of SONew scalable optimizer that can be used to train large models while using second order information.
Results:
Conclusions and Future Work
In this paper we have introduced a novel Sparsified Online Newtwon (SONew) method that yields computationally efficient sparse preconditioner that can effectively train very large DNNs. The time and memory complexity of SONew is linear in the number of parameters, unlike current Kronecker- factorization based second-order methods for training deep networks. Our experimental results show that SONew uses similar time as first order methods and achieves much better validation and training performance in various benchmarks. In the future, we plan to explore different sparsity graphs for which efficient solutions exist for the LogDet subproblem (IT) and develop corresponding regret bound analyses. Some of the limitations of SONew include: 1) explicit solutions akin to Theorem 3.1 & 3.2 need not exist for all sparsity graphs G: 2) Not all graphs allow for efficient optimizer implementation; 3) Among graphs permitting efficient optimizer implementation like tridiagonal sparsity- ordering of parameters remains unexplored. An alternative ordering might position closely related parameters adjacently, potentially enhancing performance; 4) Comprehensive exploration of methods to scale final updates is needed. While we employ grafting 2 other techniques, such as clipping [45] (38, merit investigation.
10

---

References
[1] N. Agarwal, B. Bullins, X. Chen, E. Hazan, K. Singh, C. Zhang, and Y. Zhang. Efficient full-matrix adaptive regularization. In International Conference on Machine Learning, pages 102-110. PMLR, 2019.
[2] N. Agarwal, R. Anil, F. Hazan, T. Koren, and C. Zhang. Learning rate grafting: Transferability of optimizer tuning, 2022.
[3] S.-I. Amari, Natural gradient works efficiently in learning. Neural computation, 10(2):251-276 1998.
[4] M. S. Andersen, J. Dahl, and L. Vandenberghe. Logarithmic barriers for sparse matrix cones. Optimization Methods and Software, 28(3):396 423, 2013.
[5] R. Anil, V. Gupta, T. Koren, K. Regan, and Y. Singer. Scalable second order optimization for deep learning, arXiv preprint arXiv:2002.00 09018, 2020.
[6] P. W. Battaglia, J. B. Hamrick, V. Bapst, A. Sanchez-Gonzale V. Zambaldi, M. Malinowski,
A. Tacchetti, D. Raposo, A. Santoro, R. Faulkner, C. Gulcehre, F. Song, A. Ballard, J. Gilmer, G. Dahl, A. Vaswani, K. Allen, C. Nash, V. Langston, C. Dyer, N. Heess, D. Wierstra, P. Kohli, M. Botvinick, O. Vinyals, Y. Li, and R. Pascanu. Relational inductive biases, deep learning, and graph networks, 2018. Uttpttps:://// httpp:://axxii:ii org/ abs www.ab/.......... 01261
[7] M. Bollhöfer, A. Eftekhari. Scheidegger, and O. Schenk. Large-scale sparse inverse covari- ance matrix estimation. SIAM Journal on Scientific Computing, 41 41(1):A380-A401, 2019.
[8] L. M. Bregman. The relaxation method of finding the common point of convex sets and its application to the solution of problems in convex programming. USSR computational mathematics and mathematical physics, 7(3):200-217, 1967.
[9] C. G. Broyden. Quasi-Newton methods and their application to function minimisation. Mathe- matics of Computation, 21(99):368-38) 1967.
[10] N. Cesa-Bianchi A. Conconi, and C. Gentile. On the generalization ability of on-line learning algorithms. Advances in neural information processing systems, 14, 2001.
[11] A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung, C. Sutton, S. Gehrmann, P. Schuh, K. Shi, S. Tsvyashchenko, J. Maynez, A. Rao, Barnes, Y. Tay, N. Shazeer,
Prabhakaran, E. Reif, N. Du, B. Hutchinson, R. Pope, J. B.Baddburyy J. Austin, M. Isard, G. Gur-Ari, P. Yin, T. Duke, A. Levskaya, S. Ghemawat, S. Dev, H. Michalewski, X. Garcia, V. Misra, K. Robinson, L. Fedus, D. Zhou, Ippolito, D. Luan, H. Lim, B. Zoph, A. Spiridonov. R. Sepassi, D. Dohan, S. Agrawal, M. Omernick, A. M. Dai, S. Pillai, M. Pellat, A. Lewkowycz, E. Moreira, Child. O. Polozov, K. Lee, Zhou, X. Wang, B. Saeta, M. Diaz, O. Firat, M. Catasta, J. Wei, K. Meier-Hellsterm, D. Eck, J. Dean, Petrov, and Fiedel. Palm: Scaling language modeling with pathways, 2022.
[12] L. Deng. The MNIST database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 191))144 2012.
[13] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale, 2020. https:/::
org/ abs/
2010.11922
[14] J. Duchi, E. Hazan, and Y. Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(61): 201 URL http://jmlr. wtgp/papers/////ducch.
html
[15] J. Duchi, E. Hazan, and Y. Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 201
[16] S. Fattahi and S. Sojoudi. Graphical lasso and thresholding: Equivalence and closed-form solutions. Journal of machine learning research, 2019.

---

[17] R. Fletcher. A new approach to variable metric algorithms. The computer journal, 13(3): 317-322, 1970.
[18] R. Fletcher. A new variational result for quasi Measi-Newtton formulae. SIAM Journal on Optimization, 1(1):18-1,, 1991.
[19] J. Friedman, T. Hastie, and R. Tibshirani. Sparse inverse covariance estimation with the graphical lasso. Biostatistics, 9()):432-441, 2008.
[20] J. R. Garcia, F. Freddi, S. Fotiadis, M. Li, S. Vakili, A. Bernacchia, and G. Hennequin. Fisher- legendre (fishleg) optimization of deep neural networks. In The Eleventh International Con- ference on Learning Representations, 2023. URL htttt::/ net/ www.eeeew/rmr/iee formm?id C91AOPvQHS
[21] J. Godwin T. Keck* P. Battaglia, V. Bapst, T. Kipf, Y. Li, K. Stachenfeld, P. Veličkovié, and A. Sanchez-Gonzalez, Jraph: A library for graph neural networks in jax., 2020. URL
http://github.. http:///gihhuu..cco//eemmmm/// wtt..////deemmdd///neeeee jraph
[22] D. Goldfarb. A family of variable-metric methods derived by variational means. Mathematics of computation, 24(109)233266, 1970.
[23] D. Goldfarb, Y. Ren, and A. Bahamou. Practical quasi-Newton methods for training deep neural networks, Advances Neural Information Processing Systems, 33:2386-2396, 2020.
[24] V. Gupta, Koren, and Y. Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, pages 1842-1850, PMLR, 2018.
[25] E. Hazan, A. Agarwal, and S. Kale. Logarithmic regret algorithms for online convex optimiza- tion. Machine Learning. 69(2):69--192 2007.
[26] E. Hazan et al. Introduction to online convex optimization. Foundations and Trends@ in Optimization, 2(3-4):157-325, 2016.
[27] N. J. Higham. Accuracy and stability of numerical algorithms. SIAM, 2002.
[28] G. Hinton, N. Srivastava, and K. Swersky. Neural networks for machine learning lecture 6a overview mini-batch gradient descent. Cited on, 14(8):2, 2012.
[29] C.-J. Hsieh, M. A. Sustik, I. S. Dhillon, P. K. Ravikumar, and R. Poldrack. BIG & QUIC: Sparse inverse covariance estimation for million variables. Advances in neural information processing systems, 26, 2013.
[30] W. Hu, M. Fey, M. Zitnik, Y. Dong, H. Ren, B. Liu, M. Catasta, and J. Leskovec. Open graph benchmark: Datasets for machine learning graphs, 2020. Utttttss:///arr/ .rg abs/ 00687
[31] N. P. Jouppi, G. Kurian, S. Li, Ma, R. Nagarajan, L. Nai, N. Patil, S. Subramanian, A. Swing, B. Towles, C. Young, X. Zhou, Z. Zhou, and D. Patterson. Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings, 2023.
[32] J. Kiefer and J. Wolfowitz. Stochastic Estimation of the Maximum of a Regression Function. The Annals of Mathematical Statistics, 23(3):462 466, 1952. doi: 10. 2/21144///mm//1/7799222 URL htthhttp::: /doi. org/ 10. 1214/ /tw.////////////////////eeeeeee/ 1177729392
[33] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.6980, 2014.
[34] B. Kulis, M. A. Sustik, and S. Dhillon. Low-rank kernel learning with Bregman matrix divergences. Journal of Machine Learning Research, 10(2), 2009.
[35] H. Luo, A. Agarwal, N. Disa-Bianchi, and J. Langford. Efficient second order online learning by sketching. Advances in Neural Information Processing Systems, 29, 2016.
[36] Luo, C. Chen, 7. Zhang, W.-. Li, and T. Zhang. Robust frequent directions with application in online learning. The Journal of Machine Learning Research, 1077).1677 2019.
12

---

[37] L. Luo, C. Chen, Z. Zhang, W.-J. Li, and T. Zhang. Robust frequent directions with application in online learning. Journal of Machine Learning Research, 20(45):1-41, 2019. URL http
//jmlr www.org/gga.. Siaapess/20/11/
html
[38] J. Martens and R. Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. In International conference on machine learning, pages
PMLR, 2015.
[39] S. Merity, C. Xiong, J. Bradbury, and R. Socher. Pointer sentinel mixture models, 2016.
[40] Y. Nesterov. A method for unconstrained convex minimization problem with the rate of convergence o(1/k2). 1983.
[41] N. Qian. On the momentum term in gradient descent leamning algorithms........ Neural Networks, 12 11::44-151, 1999. ISSN 0893-6080. doi: http:////ooooo//////oooo/oo//0s0
URI. htt.hhttp:////ei http://www.eiieeeeerre....oo http:///wwcccceeeoeooo//eeaaaaa/
[42] Schmidhuber. Deep learning in neural networks: An overview. Neural Networks, 61:85- 117, jan 2015, doi: 10. 10011016j.meumet2114..000000 URL ( neunet. 2014.
[43] S. Shalev-Shwartz et al. Online learning and online convex optimization. Trends in Machine Learning, 4(2):107-195 2012.
Foundations and
[44] D. Shanno. Conditioning of quasi-Newton methods for function minimization. Mathematics of computation, 24(111):6477-5) 1970.
[45] N. Shazeer and M. Stern. Adafactor: Adaptive learning rates with sublinear memory cost, 2018.
[46] H.-J M. Shi, T.-H. Lee, S. Iwasaki, J. Gallego-Posada, Z. Li, K. Rangadurai, D. Mudigere, and M. Rabbat. distributed data-parallel pytorch implementation of the distributed shampoo optimizer for training neural networks at-scale, 2023.
[47] D. R. So, W. Manke, H. Liu, Z. Dai, N. Shazeer, and Q. V. Le. Primer: Searching for efficient transformers for language modeling, 2022.
[48] T. Tieleman and G. Hinton. Lecture 6.5 6.5-rmsp Divide the gradient by a running average of its recent magnitude. coursera: Neural networks for machine learning. 2012.
[49] L. Vandenberghe, M. S. Andersen, et Chordal graphs and semidefinite optimization. Foun- dations and Trends@ in Optimization, 1(4):241-43 2015.
[50] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and
I. Polosukhin. Attention all you need, 2023.
[51] T. Wang, A. Roberts, D. Hesslow, T. L. Scao, H. W. Chung, Beltagy. J. Launay, and C. Raffel. What language model architecture and pretraining objective work best for zero-shot generaliza- tion?, 2022.
[52] L. Zhang, S. Shi, and B. Li. Eva: Practical second-order optimization with kronecker-vectorized approximation. The Eleventh International Conference on Learning Representations, 2023. /ttp::ss:///peeeree// httf://orruurr11......a.066111
[53] R. Zhang, S. Fattahi, and S. Sojoudi. Large-scale sparse inverse covariance estimation via thresholding and max-det matrix completion. In International Conference on Machine Learning, pages 5766-5775. PMLR, 2018
[54] M. Zinkevich. Online convex programming and generalized infinitesimal gradient ascent. In Proceedings of the 20th international conference on machine learning (icml-03), pages 928-936, 2003.
13

---

A
Supplementary material
A.1
Properties of Loglee subproblem
14
Regret bound analysis
15
A.2.1
Regret bound decomposition
15
A.2.2
Properties of tridiagonal preconditioner
16
A.2.3
Upperbounding Regret
17
A.2.4
O(vT) Regret
20
A.2.5
Non-convex guarantees
23
A.3
Numerical stability
23
A.3.I
Condition number analysis
24
A.3.2
Degenerate Hg
25
A.3.3
Numerically Stable SONew proof
26
A.4
Additional Experiments, ablations, and details
26
A.4.I
Ablations
26
A.4.2
Memory Requirements
27
A.4.3
Hyperparaeter search space
27
A.4.4
Additional Experiments
27
A.4.5
Convex experiments
30
A.1
Properties of LogDet subproblem
Proof of Theorem
The optimality condition of 11 is PG(X = Pg(H), X E (G). Let Z = L-1 then PZ(
2T. = -, D-1
= -,
Let J, = I, J.j. where 1,= = (j + 1....... b} defined in the theorem, then select J, indices of vectors on both sides of the second equality above and selecting the J, indices
Zz Ziz
Z11;
11
1/d 0
(15)
ZazJr
L1
Note that L-7 is an upper triangular matrix with ones in the diagonal hence 3 block of L L-Tey will be [1, 1,,,0,0.... Also, since P2(
H
Znj
Zjı,
[g
H5t,
Zt,j
H1
HJ
Substituting this in the linear equation
15
Hg
H31,
1/dg3
His
HJjJ,]
[1
Hjs H3t,
dij
His HJ,J,]
dis L1,
Hygdi + d HT, L1,5
H1g5 + d3s H131, I1, =0
The lemma follows from solving the above equations. Note that here we used that lower triangular halves of matrices L and H have the same sparsity patterns, which follows from the fact that banded
14

---

graph is chordal graph with perfect elimination order f1,2....,nz Furthermore, X, is positive definite, since as (H2) H H1,5) is schur complement of submatrix of H formed by J [5}
Proof of Theorem 3. The proof follows trivially from Theorem 3.1 when b set to
A.2 Regret bound analysis
Proof sketch of Theorem 3.3 We decompose the regret into Ryy - T1 T: + T3 in Lemma A.1
and individually bound the terms. Term T2 2g Σ (We+1 (X("" X, "x,") w depends on closeness of consecutive inverses of preconditioners, (X x, upperbound this we first give explicit expressions of Xx for tridiagonal preconditioner in Lemma A.2 in Appendix A.2.2 This explicit expression is later used to bound each entry of (X X with O(1/vZ) in Appendix A.2.4, this gives a (vVT) upperbound on T2. To show an upperbound on T3 = Ei= '. Xigh, we individually bound of Xt.9t by using a Loewner order X, 3 ||Xx|l2 I, = || IXlllooT,, In and show that X (1/v/T) and consequently T3 O(VT).
A.2.1 Regret bound decomposition
In this subsection state Lemma A.1 which upper bound the regret Rr using three terms Tı, T2. T3.
Lemma A.1 ( 26 In the OCO problem setup, if a prediction E R is made at round t and is updated as Wr41 == Wt 1i (yXt using a preconditioner matrix X, S
1
RT 1
(1ww1
w=1l "1x
-== -
(16)
21)
T-1
1
+
(We+1 (x) xx,1) - w))
(17)
2n
T
+
T X
(18)
2
Proof.
|
w0 10'1,,--
|1ue N9 w0 - x-=
|| w" "1x ni ge X
2n(w:: TT
2n(w: - w -" = ||w, - w "11x,-1 - ||w - w² x1x
n²² g: Xi9t
Using the convexity of ft, fe(w:) fe(w") < (we w²) ( 9t, where 9r Af((w,) and summing over E [T]
T
(Iwm w "1,,, - 1| w 1e=13x,"
Rr <
(19)
1=1 2m
+
T
X
(20)
2
15

---

The first summation can be decomposed as follows
||w, w" 11x- - -11w++1 w0' (1x;')
(lw1 w² 11x7- ||0T+1 w0 11x
T-1 (Wi+1 - w²) (Xx++ ( - w²)
+
Substituting the above identity in the Equation (99 proves the lemma.
Let Ry & Ti + T2 + T3, where
T1
2nj 1 (Ilw1 w² 11,x,' |10T+1 - - 1x,,,,
T-1 Σ (w w² - Xx (2 -w²)
T2 2m
(21)
T3
2t=1
2
gi Xege
A.2.2 Properties of tridiagonal preconditioner
In this subsection, we derive properties of the tridigonal preconditioner obtained from solving the DegDet subproblem [1] with 9 set to a chain graph over ordered set of vertices (I....,n}}
X, arg min - log det (X) + Tr,
(22)
(0)
arg min Dea (X, (X, (9)+
(23)
The second equality holds true only when H₃ is positive definite. Although in Algorithm
maintain sparse H, = H + Po(g.g / />:), Hn = eI,, which further used (22) to find the preconditioner X,, our analysis assumes the full update H, H + 9:92 /Xe, Ho = eT, followed by preconditioner X, computation using (23) Note that the preconditioners X, generated both ways
the same, shown in Section 3.2
The following lemma shows that the inverse of tridiagonal preconditioners used in Algorithm I will restore Hijj when (i.j) fall in the tridiagonal graph, else, the expression is related to product of H1+k,5+*+1 corresponding to the edges in the path from node to j in chain graph. This lemma will be used later in upperbounding Tz.
Lemma A.2 (Inverse of tridiagonal preconditioner). If chain/tridiagonal graph and X arg mgg Dea (X, H-),, then the inverse X-1 has the following expression
Hij Has Hapt 2..
| 1-
(8-1)
(24)
H4118 1.,
Proof.
x-1 =Ej
Where (0) is the 7th column of X. Let denote the right hand side of Equation 24
(8 (8 Xx + x X & 88342
= =Z, + X, 1 H 1 + 8 Hitt =1
16

---

The third equality is by using the following alternative form of Equation
12
0. j -i>1
(8
(H:H; 1,11
-H,,6+7
i j
(25)
R2
1 Senolgol (1)
ifi ifi= =
Bi
where <. Similarly, the offdiagonals of Y X can be evaluated to be follows.
(8 = PigXis + Yij x + 1i+t11+11
Ri + Yú H - Yis 11st1 X+11
15-
Hj
Lemma A.3. Let y R" = max, ma (H,) /V (Hi),) (H4 < 1, then
11+
y² X y < | diag H,)ll2
1
where X₂ is defined as in Lemma A.2
Proof.
Let X- = diag( -1/2
1/2 X, (Htag(H,)--14)
2
y² X < (diag((H))
8--
(26)
Using the identity spectral radius p(X)= S || oc and since X positive definite, ||x, /lo0
|x,
1x-1 < max Σ (8
1 + 2(3 B2 +...) 1+ 8
1-B 8
The second inequality is using Lemma A.2
Substituting this in Equation 26 will give the lemma.
A.2.3 Upperbounding Regret
The following Lemma is used in upperbounding both T, and Tg. In next subsection, we'll upper bound T2 as well.
Lemma A.4. Let B = maX maxig(n-11 i/ | (H))
then
1/(1- - 3) $ 8/62,
where, is constant in parameter e &G VT and consequently used initializing Ho cI, in line 7 in Algorithm
17

---

Proof.
1
1/(1 ((1-3) 3) max max
(27)
1E ||
(ft))
1 +
(Ht)stt1
(Ht)
(where (Ho)):1
max max
1| (He)
V(He)as (Ho)+ti04)
< max max (E (H+) (He)s41e41 (H.)
(since ((He) S
< max max
(28)
iE det
(Hi)s (Hi)i4
(H))a=1 (Ho))+1041 ))
(H
(Ho)+1
[
0
Note that (Hi)+41 (H4)
=
(using line
in Algorithm I thus
0
1
det ([ (H))a
(H))a+1
1
0
(H.))a
(H4)--1441
z det
6
0
62. The numerator last inequality can
be upperbounded by bounding (H)) individually as follows:
(H.)a Σ
Σ
//>
x (Go/a
Go
== vs
< 2G Vi
(29)
Substituting the above in 28
gives
8G2 t
1/(1 B) < max
2G2 T
8
<
Lemma A.5 (Upperbound Tf
16D2Goc VT
Ti S
(30)
En
where D2 maxxe(T) w w/ ll2 and Gx maxellgoloo
18

---

Proof. Since XT is positive definite
| - w0 1xx-1
Ti 5
217
(1()))(X; x1 x
(where y = 11 - w²)
2n
1 10 |diag( (H1)
1 3
(Lemma A.3
1- D 1 +c)
21
1+
(line in Algorithm I
21
1-
GG / A1 +
30
(Lemma A.4
21
Doo + EG (
(Since A Goo Go and &G VT)
En
16
(6 1)
221
Lemma A.6 (O(VT) upperbound on Tz).
T T3
Σ
7
91 Xigi S
4nG VT
2
t=l
where, l1 S Gx and parameters 2G VT, A Goo in Algorithum !
Proof. Using Theorem 3.1
nonzero entries of X₁ be written follows:
H² HnH
(X))a His
1+
Σ
H3
(L.))EF
Hat1
(X))++1 HiHHi4i4 1
where, Eg denote the set of edges of the chain graph ς in Theorem 3.1 Also, for brevity, the subscript is dropped for H,. Let X, = / /diag(H) (H)X, /diag( H), then X, can be written as
(X1)a = 1+ Σ
H?
(F
H
(X1)at1
H2
where, Hy = HH//HH Note that X. 3 [ 3 1eelloode using max)())(X;)),... (A,( < ||:\|| (property of spectral radius). So we upperbound [| maxie(m] (X)) ((X;) + ((,))+11...... ...........,,,. ((x)) + (Xx))nn113
next. Individual terms ((x,) + ((Xc) can be written follows:
19

---

Hi
|Mayl
Σ
((x))=|== =1 Σ
+ - H
1 - H23
(i.))E
=1+
1 |M49l
(
3 2 max iE(n-1)
|1sttl
The last inequality is because Hgy $ 1. Thus, [x maxice 1] I- 1-11
Now
T Xeg S of diag H2) - X, dia (H,) 1/29
96
s || |/agg( (1,) 11,2-112 9t 11
(1|x& S =1xx1l.)
1
< 2 max
of diag H,) - gt.
:Efu Hut+l
Using diag(H:) 2 eln, (step in Algorithm 1), where EG VT as set in Lemma A.8, gives
1
11
ot Xigi S 2 max
Hut1l
2G
VT
nG
< 2 max
1Ein ē(1 [Hs441))
2nGoo
$ ē(1 3VVT
(where max max LE(2) sE[n-1] (He)set2)
Summing up over gives
Σ
11 T' Xege th
- Σ
16nGxo7
(Using Lemma A.4
VT
16nGxo'l VT
A.2.4 O(VT) Regret
In this section derive regret upper bound with O(T1/2) growth. For this, we upper bound T2 as well in this section. In (1). I2 Σtz w" )7 (X,
Xw, w²) can be upper
bounded to a O(T1/2) by upperbounding entries of X, X individually. The following lemmas constructs a telescoping argument to bound (X,- X Xx-1)ii Disl.
Lemma A.7. Let H, H E st such that f H1 gg"" />, where g E R", then
Hg
v Ha
HHaHi
9:93
Hig VH..
(
Hallis Hullis
-1) Oz
+
(V
2v
HulHis
20

---

Proof.
Haj
Hij
v
HigH
(Hi))
H(g)
H Hss
HugHi
Hintt (Enttis)))
v -
9:91
Hij
J
V H3
The following Lemma bounds the change in the inverse of preconditioner Y-1, when there is rank one perturbation H following LogDet problem (11)
Y
arg min log det (X) + TYXH)
(9)1
arg min (xa(X,
(9)+
Lemma A.8 (Rank perturbation of LogDet problem 11
Let H,H S such that
/>, where g E R". Also, Y arg min es, (G)++ Dra(X, Hf) and Y arg (0) (Xa H), where G is chain graph, then
H = H + ////>,
-9-1-- - -Y- Gan(kB + 2)8h
where i, + k S nI, Go = l1 and maxi. H S 3 < 1. Let k(diag(H)) = condition number of the diagonal part of H, then == max(Is(diag(F)) (H)), x(diaag(H))
Proof. Using Lemma A.2
will give the following:
(8-1 1a4k
ilut HAtk-1448
Ha+1 Hitk-194k
Hep ... Hepk-1444 1itk-1
H441441 1i4k-184k-1
V MRNNAII NETE y Mitsi4
/ HRNNAI.... N.. Hitkith
= Haallat Nut1.. Neth-14R Na+1... Ni4 144 / HaHlsatte H HulliratH
where No HH// H/g/ HHaHij < 1 (Since determinants of 2x2 submatrices of H are positive). Expanding Nat1 Na41 0.141 (from Lemma A.7 subsequently Nut2 Nut2 0at2 and so on will give
Halligiin
Not1 ... NetR-1i4k Nut Netk-
V
H Halligii
Biiiiie...N..knnee ............................ 1ith Not1 (Nirletz.. ... Neck litk Not1442... ... NitR-144h -i4k
... Ne4k
Hig
21

---

- 10 = ....Ni44-- 16th + N6-104442NuHS NetR-1446 + +.. No+1 N 1aae--104hhh- 1i+k
Hillitkii
+ NN411 Notk 1-11
Hiilli+ki-
0 (110+4411))8-- +/6--1 1 1 HaHHs4att
H HuHorketk
(8-1- -y-1---1)44 )-xx $ V Hillitttt
Ha Helliran
where maxi, |Na.ll maxis |Ntal S < Expanding 0c41i4441 from Lemma A.7 in the term 0111+1441| Ha Hetketk will give:
141 V HuHH+att
(v HonhntHotnot
-1)
HuHis4t A H He
9i+19it1+1
H HoBHymeNNHuHHH
H (e4nsHH44444s
(1-
MallitRitR AV HitH Hastiiett
9s+19it1+1
+ HalleaaaNNaeen
(1-v
Since H H S HAAHH
I max (1-
H4t0+
1 H4444144441
V His 44444144411
H444
H444414444
git
max
XH AH:44164441
Using the above, H,,/ HJs S K, and 1gil S Go., Vi.j [n]. gives
Gh BG
< G n(1 + B)/A
Thus the following part of 1y-1--11
can be upperbounded:
a+t
V Hillinkink
10
(1= /
Also, / Ha Hi Mi milk
< 3 so
(ÿ (5-1-y-) y-1) äi4k
5G + k +
Lemma A.9 (O(VT) upper bound of T2). Given that s(diag( H.)) S K, l|w; w" 1 < D2.
maxi, (He)) v (H))) 8 < 1, Vt e [T] in Algorithm 7 then T2 in Appendix A.2.1 1.2.1 can
be bounded as follows:
2048
TB <
(Goe D3)
where >, Go and eGoo VT in Algorithm
and S. 1 is constant.
22

---

Proof. Note that T2 = 1 El= (wi+1 wx)7(x(11 - X, (we+1 w²) -
Ei D2 [X x; )12 ((n))
Using ||All2 = =P(A) - || for symmetric matrices A.
we get
|xT+ - Xi 112 = || - -x, -xxxlla '1100
mmx Σ (X - _x, >x,')sl)
Gook
- 16
(Lemma A.8
vi(1 B32
< 1024 Gook
Viei
Now using S. 2/é (using Equation above will give the result.
29 and (He) > é) and summing up terms in T2 using the
Putting together T, T2 and T3 from Lemma A.5 A are defined as in Lemma A.9
Lemma A.9 and Lemma A.6 respectivvly when
16D VT
T1
2n
2048VT
T2
(2)
(31)
1765
4nGoo VT
T3 <
(32)
Setting 7)
D2 tvm
RT Ti - T2 +T3 < S
A.2.5 Non-convex guarantees
Minimizing smooth non-convex functions f is a complex yet interesting problem. In Agarwal al. D. this problem is reduced an online convex optimization, where sequence of objectives fe(w) f(w) + ccll we l are minimized. Using this approach Agarwal et al. [ established convergence guarantees reach stationary point via regret minimization. Thus non-convex guarantees can be obtained from regret guarantees and our main focus in the paper.
A.3 Numerical stability
In this section conduct perturbation analysis to derive end-to-end componentwise condition number (pg. 135, problem 7.11 (27)) upper bound of the tridiagonal explicit solution in Theo- rem 3.1] In addition to this, devise Algorithm to reduce this condition number upper bound for the tridiagonal sparsity structure, and be robust to H, which don't 't follow the non non-degeneracy complete subgraph of G.
condition: any principle submatrix of H, corresponding
Theorem A.10 (Condition number of tridiagonal LogDet subproblem (11). Let H E S
+ be such that Hap = 1 for E [n]. Let ΔH be a symmetric perturbation such that A A = 0 E [n], and
H + AH E S++ Let P)(H) be the input to 11 where G is chain graph, then
id Nog
max xE[n-1]
2/(1 87)
R
(33)
where, B, = Hii4 = componentwise condition number of 11 for perturbation AH.
The tridiagonal LogDet problem with inputs as mentioned in Theorem A.10 has high condition number when 3? Hy / are low and result the preconditioner X, in
23

---

SONew (Algorithm 1 has high componentwise relative errors. We develop Algorithm B to be robust to degenerate inputs H. given that Hü > O. It finds a subgraph of g for which non-degeneracy conditions Theorem 3.2 satisfied and (14) is well-defined. This done by removing edges be singular (H Hi, 1 Hışr) )
which causes inverse H1 theorem we also show that the condition number upper bound in Theorem A. A......rdeee. reduces in tridiagonal case. To test the robustness of this method we conducted an ablation study in Table [5] in an Autoencoder benchmark (from Section 5) in 6íloat16 where we demonstrate noticeable improvement in performance when Igorithm33js used.
to be low. In the following
Theorem A.11 (Numerically stable algorithm). Algorithm finds a subgruph of G. such that explicit solution for 9 in (14) well-defined. Furthermore, when 9 is a tridiagonal/chain graph, the component-wist wise condition number upper bound in (33) is reduced upon using Algorithm 3 Redd Rod
where & KG are defined as in Theorem A.10 for graphs and G respectively.
The proofs for Theorems A.10 and A.11 are given in the following subsections.
Algorithm 3 Numerically stable banded LogDet solution
1: Input: g- tridiagonal banded graph, H/ symmetric matrix in R= with sparsity structure G and Hi > O, 7- tolerance parameter for low schur complements.
2: Output: Finds subgraph 9 of without any degenerate cases from Lemma A.13 and finds preconditioner
X corresponding to the subgraph
3: Let E, = [(i.]) (f.j) E Eg} be edges from vertex to its neighbours in graph G.
4: Let V.+ = {j :1 < J. (1, E }g} and V, = 5. E Eg}, denote positive and negative neighbourhood of vertex i.
5: Let K == {i: Ha IT,, HT,l, H1, is undefined or S {>}
6: Consider new subgraph & with edges Eg Eg (Uiex E. UV, x V,))
7: return X = CERRSIFIEDINNEERE (He, 9), where He Pş(H.)
A.3.1 Condition number analysis
Theorem A.12 (Full version of Theorem A.10 Let H E S such that Ha 1, for i E [n] and a symmetric perturbation AH such that AH = 0, for E [m] and + AH > 0. Let arg min (0) Dra (X, H-1 and X + A arg xes, Dea (X, (H + AH)-1), here G = chain/tridiagonal sparsity graph and S (G)+ denotes positive definite matrices which follows the sparsity pattern G.
AX
Kid lim sup c-+
|AH $ JeHaal, (k. E Eg
{ij]
max 1/(1
{?)
vE[n-1]
where, Ked == condition number of the LogDet subproblem, x²(.) == condition number of a matrix in l2 norm, B. Ha41/ / HHass//HaHaan HaHH41i-1
Proof. Consider the offdiagonals for which (X + A - -Hitt/(11 H
f ((a41) f(z) = -I/(1 x²). Let = fxx, ú = f(x + Ar) and Ax/2l S c then
using Taylor series
(U y)
x (x)
Ar
+ 0((Ar)²
f(x)
(ÿ y)
x (x)
lim c-00
cy
f(x)
24

---

Using the above inequality, with 2 = Hu+1 and
== X
A X N8+1
1 + Hi
lim c00
(34)
eXa+1
Hiit1
2
$ 1 H
Let g(x) = =x2((1 a²), let V1 = g(w1): ). 1/2 g( 01 = g[w] Ar), 02 = g(22 + Ar). Using
Taylor series
(91 - -11)
(11))
An
+ +0((Az))3
y1
f(11)
21
(92 822)
(22)
Az
+ +((Axz))))
1/2
(122)
22
Ay1 + A12
2
2
lim 00
c11 + y1 V2)
max
1
rf
2
1
35
Putting 21 == Ha+1, 32 = - and analyzing yn == -22
1/1+1/111
- H H
and y22 == H? H/-11/11
11 will result in the following
A xx
lim --0
I max
(35)
X
1-
11 144
- 113
Since 8, = + H #++1/ (11
+ H2 11
H² 1) Putting together Equation 35 and
Equation 34 the theorem is proved.
A.3.2 Degenerate H₃
In SONew (1. the H, P=(E:=1969) />e) generated in line could be such that the matrix
E 9=9" />x need not positive definite and so the schur complements Ha - H / H H4+1+1 can be zero, giving an infinite condition number se by Theorem A.10 The following lemma describes such cases in detail for more general banded sparsity structure case.
Lemma A.13 (Degenerate inputs to banded LogDet subproblem). Let H Pg(GG"), when e
(6) in Algorithum where G RaxT and let 91-T be 5th row of G, which is gradients of parameter for (1)
T rounds, then Hij ( 91:T:
Case 1: For tridiagonal sparsity structure G: if oift /2z41/HH+uu = 0.
(++1) then Hjs
Case 2: For b > in 14): If rank(Ha,,,) rank(Hr,,,) b. then (Hzs HT, H1,5) 0 and Djj co. If (ank(HH111)) < b then the inverse H7j,, doesn't exist and Djs is not well-defined.
() Proof. For = 1, if gi gi. then Hjj+1 Hjj H,+1941 | | thus H33
1++1)
HH+1j+1 / 0.
For 1, since H1,1y- using Guttman rank additivity formula, rank(H,; H3 () rank( (,nk(H,,J,) vank(H,,,, = 0, thus Hjj Furthermore, if rank H) < b, then all + x b + principal submatrices of have rank b, thus Vj. HJ,J, have rank b, thus Dij for all j are undefined.
=
If = 9:9: is a singular matrix, then solution to the LogDet problem might not be well- defined as shown in Lemma A.13 For instance, Case can oceur when preconditioning the input layer of an image-based DNN with flattened image inputs, where jth and (j + 1)th pixel can be highly correlated throughout the dataset. Case can occur in the first b iterations in Algorithm when the rank of submatrices rank(HI,1,) and 0.
25

---

Table float32 experiments on Autoencoder benchmark using different band sizes. Band size
corresponds to diag-SONew and corresponds to tridiag- SOoiag We see the training loss getting better as we increase band size
Band size
0 (ddag-SONew)
Stridia-SONeen
4
10
Train CE loss
53.025
51.723
51.357
51.226
A.3.3 Numerically Stable SONew proof
Proof of Theorem A.11
Let I;
(.: Eg} and I
{j:i< (1:5) E Eg] Let K
{i Ha HT. -H,, H H undefined 0, i [n]] denote vertices which are getting removed by the algorithm, then for the new graph G. Da H vi E K since H > 0.
Let K = (i: Hu - H Hi H 0, [n]] Let for some j K, if
arg min j < eiikKlll Kn
denotes the nearest connected vertex higher than j for which Du is undefined zero, then according to the definition Eg in Algorithm 3 I {j + +..... 1.....-111. 1} C 1j. since Dj is well-defined, H1,4, is invertible, which makes it positive definite matrix (since H is PSD), Since H -H Hig,, H1, > using Guttman rank additivity formula HJ > where J, I U.. Since HJJJī is submatrix of 3, it positive definite and hence its schur complement H3j HF , for all j E [m]. the corresponding Djj²s 's are well-defined in the new graph g.
H Hr Hnys > 0. Thus
c Note that Kidd max]e[n-1] 1/(1 () < max 1/(1 - (f) nea for tridiagonal graph, where 3; H5+41, in the case where H;; This is because the arg max;e|n-1] 1/(1 B.) E K.
A.4
Additional Experiments, ablations, and details
A.4.1 Ablations
Effect of band size in banded-SONew Increasing band size in banded-SONew captures more correlation between parameters, hence should expectedly lead better preconditioners. We confirm this through experiments on the Autoencoder benchmark where we take band size 0 (diag-SONew) 1 Stridiag-SNeeva 4, and 10 in Table
Effect of mini-batch size To find the effeet of mini-bate batch size, Table 4 We empirically compare SONew with state of the art first-order methods such Adam and RMSProp, and second-order method Shampoo. We see that SONew performance doesn't deteriorate much when using smaller larger batch size. First order methods the other hand suffer significantly. We also notice that Shampoo doesn't 't perform better than SONew in these regimes.
Table 4: Comparison Autoencoder with different batch-sizes
Baseline Batch size
100
1000
5000
10000
RMSProp
55.61
53.33
58.69
64.91
Adam
55.67
54.39
58.93
65.37
Shampoo(20)
53.91
50.70
53.52
54.90
tds
53.84
51.72
54.24
55.87
bds-4
53.52
51.35
53.03
54.89
Effect of Numerical Stability Algorithm 3 On tridiag-SONew and banded-4-SONew we observe present in Table 5 results where we observed
that using Algorithm g improves training loss. significant performance improvements.
26

---

Table bíloat16 experiments on Autoencoder benchmark with and without Algorithm We
observe improvement in training loss when using Algorithm
Optimizer
Train CE loss without Algorithm
Train CE loss with Algorithm
tridiag-SONew
53.150
51.936
band-4-SONew
51.950
51.84
Table 6: A rough estimate of memory requirement comparisons of different optimizers tested across benchmarks.
model parumeters
K-FAC
Shampoo
FishLeg
Adum
SGD-Momentum
RMSprop
SOsew
Autoencoder
n=1.4M
5,56n
6.56n
4.28n
2n
3n
GraphNetwork
nm3.5M 5M
8.60
10.60
Vision Transformer
n=22M
6.4n
7:2n
3,7n
2n
Language Model
mml 38
5.6m
6.6m
3,30
2n
3n
A.4.2 Memory Requirements
We present list of approximate memory requirements of different optimizers across different benchmarks in T Note that for K-FAC and Shampoo, because preconditioner is updated once only a few steps, they require storing the latest computed preconditioners well along with the statistics, causing even higher memory overhead.
A.4.3 Hyperparaeter search space
We provide the hyperparanter search space for experiments presented in Section 5 We search over 2k hyperparameters for each Autoencoder experiment using Bayesian Optimization package. The search ranges are: first order momentum term 31 E [le 1.9.999), second order momentum term [1e 10, 1e 1]. We give the optimal hyperparameter value for each experiment in Table T2 For VIT and GraphNetwork benchmark, we search $1, B2 E 10.1, 0.999], , Ir E [le- 5, 1e 1]. c [1e 9, le 4), weight decay [le 5, 1.0), learning rate warmup E 12%, 5%, 10%] stotal L___trrinnn__pp_ dropoute [00, 0.1]. label smoothing over 0.1, We cosine learning rate schedule. Batch size kept 1024, and 512 for Vision Transformer, and GraphNetwork respectively. We sweep over 200 hyperparameters in the search space for all the optimizers.
82 E [1e 1, 0.999), learning rate [1e 7, le 1]
For ridSON (37, there's 's no e hyperparameter. In addition to the remaining hyperparameters, we tune {le 5, 1.0} (plays similar role as and #t E [le 5, 0.1]. For LLM 47 benchmark, we only tune the learning rate [le 2, le 3, 1e 4] while keeping the rest of the hyperparams as constant. This is due the high cost of running experiments hence we only tune the most important hyperparameter. For Adafactor [45) we use factored=False, decay method=adam, 31 0.9, weight decay=le 3, decay factor=0.99, and gradient clipping=1.0.
A.4.4 Additional Experiments
VIT and GraphNetwork Benchmarks: In Figure 5 plot the training loss curves of runs corresponding to the best validation runs in Figure Furthermore, from an optimization point of view, we plot the best train loss in Figure got by searching over 200 hyperparameters. We find that tridiag-SONent is 9% and 80% relatively better in ViT and Getwork benchmark respectively (Figure 6p. compared to Adam (the next best memory efficient baseline).
Autoencoder float32 and bfloat16 experiments: We provide curves of all the baselines and SONew in Figure a) and the corresponding numbers in Table 7 for float32 experiments.
To test numerical stability of SONew and compare it with other algorithm in low precision regime, we also conduct bfloat16 experiments on the Autoencoder benchmark (Table We notice that SONew undergoes the least degradation. Fridiigonallsparsii SONew CE loss increases by only 0.21 absolute difference (from 51.72 in float32 (7) to 51.93), whereas Shampoo and Adam incur 0.70 loss increase. It's worthwhile to note that SONew performs better than all first order methods while taking similar time and linear memory, whereas while Shampoo performs marginally better, it is 22x slower
27

---

Table 7: float32 experiments Autoencoder benchmark. We observe that diag-SONew performs the best among all first order methods while taking similar time. tridiag and band-4 perform significantly better than first order methods while requiring similar linear space and time. Shampoo performs best but takes O(d + dğ) time for computing preconditioner of a linear layer of size d1 x d2, whereas our methods take O(d1d2) time, as mentioned in Section 3.3
ifdSON takes similar space as SONew hut performs considerably worse.
Optimizer
First Order Methods
SGD
Nesterov
Adagrad
Momentum
RMSProp
Adam
diag-SONew
Train CE loss
67.654
59.087
54.393
58.651
53,330
53.591
53.025
Time(()
62
102
62
67
62
62
63
Optimizer
Second Order Methods
Shampoo(20)
rfdSON(1)
rfdOON(4)
tridlag-SONew
band-4-SONew
Train CE loss
50.702
53.56
52.97
51.723
51.357
Time(s)
371
85
300
70
260
Table 8: bfloat16 experiments on Autoencoder benchmark test the numerical stability of SONew and robustness of Algorithm We notice that diag-SONew degrades only marginally (0.26 absolute difference) compared to float32 performance. tridiag-SONew and band-4-SONew holds similar observations as well. Shampoo performs the best but has considerable drop (0.70) in performance compared to float32 due to using matrix inverse, and slower due to its cubic time complexity for computing preconditioners. Shampoo implementation uses 16-hit quantization to make work in 16-bit setting, leading to further slowdown. Hence the running time in 6 even higher than in float32.
Optimizer
First Order Methods
SGD
Nesterov
Adagrad
Momentum
RMSProp
Adam
diag-SONew
Train CE
80.454
72.975
68 854
70.053
$3.743
$4.328
53.29
Train time(s)
36
43
37
36
38
Optimizer
Second Order Methods
Shampoo(20)
rfdSON(1)
rldSON(4)
Midiag-SONee
band-4-SONew
Train CF loss
51.401
57.42
55.53
$1.937
51.84
Train timc(s)
1245
80
284
230
60
Nementum
Nesterev
AtaGina
AdoCoo
Adam
Adam
54
deg sonen
dlog 50tmw
53
shamppo020)
Uiclag SON
baddap SONew
52
Band SONOW
bard & 50Xex
riteen!)
52
20
60
80
100
epochs
(a) float32 autoencoder
(b) bfloat16 autoencoder
Figure 4: Training curves of all the basclines for Autoencoder benchmar (a) float32 training (b) bfloat 16 training
than tridiagonal SONew.. The corresponding loss curves are given in Figure b).
Note: In the main paper, our reported numbers for rfdSON on Autoencoder benchmark in Table 2 for float32 experiments are erraneuous. Please consider the numbers provided in Table 7 and the corresponding curve in Figure44a) Note that there's no qualitiative change in the results and none of the claims made the paper are affected. SONew is still significantly better than rIdSON, We also meticulously checked all other experiments, and they do not have any errors.
Autoencoder on KFAC, FishLeg, Eva: For completeness, We compare SONew against KFAC 38), FishLeg (20), and Eva [52) on Autoencoder benchmark as used in their official implementation.
28

---

(a) VIT train CE loss
(b) GraphNetwork train CE loss
0.84
0.70
0
6.221
0.70
eieSON
RMSSPhop
0,545
RMSProp
0.0230
0.67
Momantum
0.711
52.30m
Momsntum
0.0229
Ader
0.753
3143m
8.529
Aeimm
0.0716
Lithay SONeew
0.705
51,53a
TMSON
Stampoo
stypo (1.50)
40
100
120
140
Figure 5: Train loss corresponding the best validation in Figure VTT benchmark (b) GraphNetwork benchmark, We observe that TrdiaagSSNew match or perform better than Adam.
(a) Best VIT train CE loss
(b) Best GraphNetwork train CE loss
2.3-10
RMSProp
Momsentam
0,7511
5a35m
MM8ON
0.5650
6106m
Mar
0.8861
5M3m
Undier---NNee
0.52.9
Momesture
RMSPrup
1c:02
0.48
Adam
1S88ON
St-44
0.2e-0
12/W
triniig NONe
Slammpoo
2.3c-04
1.2e-00
Steps 11).5)
Tiare Cesin
Figure 6: Best train loss achieved during hyperparam tuning. (a) VIT benchmark (b/GraphNetwork benchmark. We observe that tridiag SONew significantly outperforms Adam, while being comparable or better than shampoo.
FishLeg
56.0
55.7
tridlag-SONew
Train 1055
56
56.5
28.5
50
90
time (5)
Figure 7: Autoencoder benchmark run using Pytorch on KFAC, FishLeg, Eva, and tridiag-SONew: notice that tridiag-SONee beats all other baselines by a large margin.
The main difference is their implementation uses ReLU activation compared to Tanh that we used for all our Autoencoder experiments. As these baselines done have JAX implementation, official PyTorch implementation and run tridiag -Sdia in Py PyTorc as well. Hyperparameter search conducted for SONew similar to as reported above, over learning rate, B1, ,2, and For KFAC and Eva, rather than B2, damping factor is tuned over [le - 5, 10] (default value specified is 0.03). kl kl_clip tuned as well over [le 5, 1.0). Preconditioners are updates once every 15 iterations to have same wall clock time as other baselines and SONew. For FishLeg, auxiliary learning rate is tuned E [le - 7, le 1] and damping [le 5, 1.0). All other hyperparameters are tuned similar to SONew. Eva is trained for 100 epochs, and for other methods we change number of epochs such that
use their
29

---

Validation CE loss
b) Train CE loss
3.357 2h:0m
233
tridiag SONew 2h43m 2h10m
trndlag-SONNer 3,310
2.55
0.1
2.2
3.3
135
132
85
xtepec (1g
stags (11
Figure 8: We observe mixed results for 248M parameter language model benchmark. This is possibly because 48 trials were insufficient for optimal tuning. We leave thorough tuning and investigating into the ahove observation future work.
each experiment takes same amount of time. Each optimizer is tuned using 600 hyperparameters. The results are in Figure 7, where notice that Mudiag-SONevs beats all the baselines by large margin.
Adam vs SONew on 248M Language Model: We conduct an additional experiment on a 248M
parameter transformer architecture (50) language model. The model is trained on WikiText103.
introduced in [39]. We train the model for 3 epochs, having 8M tokens per epoch with a batch size of 8k tokens. We search over 48 hyperparameters, tuning learning rate E 2e 2, 1e - 2, 5e 3, le 3. 32 E 0.99. 0.999, weight decay E 0.0, 0.1, and e E le 10, le 8, le - 6, while fixing $1 = 0.9. Validation and training loss is given in Figure 8 Our observations indicate mixed results. While tds-SONew exhibits superior validation performance, Adam outperforms in training metrics. We believe that the 48 trials might have been insufficient for optimal tuning. It's conceivable that with further trials, SONew's 's training loss could surpass that of Adam. We leave this line of investigation for future work.
A.4.5 Convex experiments
As our regret bound applies convex optimization, we compare SONew rfdSON 37 another recent memory-efficient second-order Newton method. We follow (37 for the experiment setup each dataset split randomly in 70%/30% train and test set. Mean squared loss is used. For tridiag-SONew we use a total of 2 d space for d parameters. Hence, for fair comparison we show redsON with m 2. Since the code isn't open sourced, we implemented ourselves. In order to show reproducibility with respect to the reported numbers in (37, we include results with = 5 as well. We see in the Table tridiag-SON SONew consitently matches or outperforms rfdSON across all 3 benchmarks. Each experiment was run for 20 epochs and we report the best model's performance on test set.
Table 9: Comparison of ridSON and SOONe in convex setting on three datasets. We optimize least square loss (v - w" xx,) where is the learnable parameter and (zt, We) is the th training point. Reported numbers is the accuracy on the test set.
Table 10: (a) Dataset stats
Table 11: (b) RFD-SON vs tridiag SONew
Dataset
# total points
dimension
Dataset
RFD-SON, m=2
RFD-SON, m=5
tridiag-SONew
a9a
32,561
123
83,3
83.6
84,6
gisette
6000
5000
gisette
96.1
96.2
96.6
mnist
11791
780
mnist
93.2
94.5
96.5
30

---

Table 12: Optimal hyperparams for Autoencoder Benchmark
Table 13: (a) float32 experiments optimal hy-
Table 14: (b) bíloat experiments optimal
perparamters
hyperparamters
Baseline
B1
B2
Ir
Baseline
81
B2
Ir
SGD
0.99
0.91
8.37e-9
1.17e-2
SGD
0,96
0.98
2.80e-2
1.35e-2
Nesterov
0.914
0.90
3.88e-10 10
5.74e-3
Nesterov
0.914
0.945
9.48e-9
6.19e-3
Adagrad
0.95
0.90
9.96c-7
1.82c-2
Adagrad
0.95
0.93
2.44e-5
2.53e-2
Momentum
0.9
0.99
1c-5
6.89c-3
Momentum
0.9
0.99
0.1
7.77e-3
RMSProp
0.9
0.9
lc-10
4.61e-4
RMSProp
0.9
0.9
2.53e-10 10
4.83e-4
Adam
0.9
0.94
1.65e-6
3.75e-3
Adam
0.9
0.94
3.03e-10 10
3,45e-3
Diag-SONer
0.88
0.95
4.63e-6
1.18e-3
Diag-SONew
0.9
0.95
4.07c-6
8.50e-3
Shampoo
0.95
9.6e-9
3.70e-3
Shampoo
0.85
0.806
6,58c-4
5.03e.3
tridiag
0.96
1.3e-6
8.60c-3
ztridiag
0.83
0.954
1.78c-6
7.83e-3
band-4
0.88
0.95
1.5e-
5.53e-3
bund-4
0.9
0.96
1.52c-6
4.53e-3
31


Reviews:
[{"review_id": "2uoYecrz2d", "reviewer": "Reviewer_XomD", "summary": "This paper propose a memory-efficient second-order algorithm named Sparsified Online Newton (SONew) method. By introducing the LogDet matrix divergence for regret analysis, the authors derives an online Newton update rule. Then the authors introduces structured sparsity patterns on the preconditioner matrix on the preconditioner matrix to reduces storage and computational complexity. They adopt the proposed methods on several benchmarks and show SOTA performances comparing to other optimization methods.", "questions": "1. Could the author provide the actual GPU memory cost and comparing the proposed method with other methods?\n2. Could the author provide more experiments to show the generalization of the proposed methods?", "limitations": "The author does not discuss the limitations and potential negative societal impact of their work. I consider one of the limitations of the proposed method is that it can not be guaranteed that this method can achieve better performance compared to other optimizers on more datasets and network architectures.", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The writing and organization of the paper are very good, the notation representation is accurate.\n2. The method proposed in the paper is reasonable, and the paper applies the proposed method on some benchmarks, achieving better results compared to other optimizers.", "weaknesses": "1. As an optimization method with lower resource usage, the authors only mention the theoretical reduction in space complexity, and the actual decrease in training time compared to other optimizers. However, they do not provide actual comparisons of GPU memory usage.\n2. The authors evaluate the proposed optimization method on deep autoencoders, Vision Transformers, and Graph Neural Networks. However, I think that as a newly proposed optimization method, the authors need to conduct experiments on more benchmarks and more network architectures to verify its general applicability, such as results of more CNN methods on ImageNet, and generative models in computer vision like Diffusion Models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments.\n\n> Could the author provide the actual GPU memory cost and comparing the proposed method with other methods?\n\nWe mention these numbers for ViT on Shampoo (155M optimizer params), SONew (44M optimizer params), and Adam (30M optimizer params). We promise to report similar comparison for all other baselines and benchmark in the final draft\n\n>Could the author provide more experiments to show the generalization of the proposed methods?\n\nWe included additional baseline in the rebuttal pdf to show SONew's performance with respect to other second order methods. We promise to include more large scale experiments in addition to the existing ones in the final draft."}, {"review_id": "i9gVg3NyLe", "reviewer": "Reviewer_KiQx", "summary": "This manuscript proposes SONew, a second-order algorithm that uses sparsity in the preconditioner to distil to linear time and space complexity. Different sparsity patterns (tridiagonal, banded, etc.) allow generalization of the algorithm for training deep neural networks almost as efficiently as first-order methods. SONew achieves up to 30% faster convergence compared to Adam and improves validation performance when trained for the same number of iterations when training with the Adam optimizer. ", "questions": "•\tThe authors compare Shampoo with RMSProp grafting but use Adam grafting for their method. Adam grafting would yield a more accurate step size compared to RMSProp. It would be beneficial to add a baseline of Shampoo with Adam grafting\n•\tFig. 1a mentinos that hyperparameter search could not be performed on Shampoo on 16GB TPU. Could the precision be reduced (bfloat16) to run shampoo?\n", "limitations": "The authors need to expand on the limitation of their method and point out specific cases where their method may be disadvantaged compared to existing methods. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "•\tThe paper provides a strong theoretical foundation for sparsifying the preconditioning matrix, making a clear connection between regret minimization and  LogDet divergence of the preconditioning matrix\n•\tThe paper provides an extensive suite of experiments indicating the computational efficiency and performance gains of their method on large scale model training \n•\tThe authors provide theoretical guarantees of the numerical stability of their method, which is a serious concern for existing second-order approximation OCO algorithms\n•\tThe paper is clearly written\n", "weaknesses": "The paper does a good job providing both theoretical and empirical justification for their method. However, there are a few details about the experimental setup and bottlenecks for baselines that would improve the manuscript. The limitations section should also be expanded to explain under what conditions SONew would not be optimal or theoretical guarantees would fail. Details provided below in the questions section:", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments\n\n> “The authors compare Shampoo with RMSProp…”\n\nWe compare Shampoo using RmsProp grafting since that is the default value of grafting hyperparameter set. Moreover, Adam grafting would require an additional $O(NumParams)$ values to be stored to maintain the exp moving average of the diagonals, leading to extra overhead over already large memory overhead by Shampoo.\n\n> “Fig. 1a mentions that hyperparameter search..”\n\nThanks for the suggestion. Indeed one can use bfloat16 and test Shampoo using that. Though as we move to even larger benchmarks, converting to bfloat16 precision still won’t scale. So the memory bottleneck issue of Shampoo still exists. We leave bfloat16 experiments of Shampoo, as well as SONew as a future work.\n\n> “The authors need to expand on the limitations of their method and ..”\n\nWe’ll surely expand upon the limitations and include specific cases where SONew might not work. For example, when the activations are sparse, then SONew might not be capturing much useful information as two consecutive neurons might not be firing up simultaneously a lot."}, {"review_id": "Fi7RSXtBBM", "reviewer": "Reviewer_UQcK", "summary": "This paper introduces an online Newton method for optimization of convex losses, with the goal of applying it to large neural networks.\nIn particular, the algorithm is derived by optimizing a preconditioning matrix on a regret bound and using the logdet divergence as a regularizer.\nIt introduces a sparsification procedure to alleviate the computational burden of computing and inverting the full pre-conditioning matrix.\nWith a fixed spasification structure (banded) of the preconditioning matrix, the problem regularized by the logdet divergence admits solutions that can be computed efficiently.\nThe method is applied to a variety of (non-convex) benchmarks and compared with other optimizers.", "questions": "See above", "limitations": "See above", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1) Fast optimization of neural networks is very important, to possibly reduce the amount of resources necessary for training neural networks.\n\n2) Theoretical derivation of the algorithm is clear and sound.\n", "weaknesses": "1) The derivation of the algorithm is based on the assumption of a convex loss function.\nMost of previous work discussed in the \"Related work\" section is about other online Newton methods, also designed for convex problems.\nHowever, all numerical experiments are shown on non-convex problems.\nThis is problematic, because if the claim is that the method should be used for second-order optimization of non-convex losses, then it should be compared with the corresponding benchmarks: KFAC, K-BFGS, FishLeg.\nIndeed, the paper acknowledges that KFAC uses a similar sparsification structure (tridiagonal), but no empirical comparison is provided.\nOn the other hand, if the claim is that the method should be used for convex optimization, then it should be clarified what are the (theoretical and empirical) advantages of the algorithm with respect to other online Newton methods.\n\n2) Some of the statements made in the paper are misleading. For example:\ni) The algorithms provided on page 5 are not actually the algorithms used for the numerical experiments. \n\"Adam grafting\" is used on top of those algorithm, but no justification is given, also in the light of the proposed theory and derivations.\nMy assumption is that none of the theory developed holds after including \"Adam grafting\" in the algorithm, and a description of the full algorithm should be given instead of just giving a reference to another paper.\nii) The following statement in section 5 is not true: \"We compare SONew against widely used first order methods including SGD, SGD with Momentum, Nesterov, Adagrad, Adam, and Rmsprop.\"\nThe first four of the six optimizers mentioned in the list are actually not compared with. \niii) The abstract says that the proposed method is slower than first-order methods, but no results are shown in the results about this. \nIn fact, figures show that the proposed method is faster in terms of number of steps, but what is the significance of this results if the method is overall slower? \n\n3) The results of numerical experiments are underwhelming. \nAs I already mentioned, no comparison with other second-order methods is given, besides Shampoo, and comparison with simple and strong benchmarks is also not provided.\nFor example, SGD with momentum is never compared with, and neither Adam nor SGD are compared with in the case of LLM in figure 3.\nFurthermore, it is not clear if the method provides a reliable advantage.\nOnly single training runs are shown, and fluctuations look pretty large in all figures.\nAveraging over multiple training runs would help understanding how reliable is the method.\nFinally, training curves are shown only vs epochs or number of steps, no results in terms of wall clock time are provided, therefore it is hard to tell whether the method provides any real advantage.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments.\n\n> The derivation of the algorithm is based on the … corresponding benchmarks: KFAC, K-BFGS, FishLeg\n\nThere are multiple reasons behind our choice of the baselines for each benchmark. We explain them here:\n\n- Our aim is to present a memory efficient second order optimizer, hence targeted for large models as commonly used in deep learning. While KFAC, FishLeg are popular second order methods, they require $O(d_1^2 + d_2^2)$ memory, which can be significantly larger than $O(d1*d2)$ in the context of deep learning models. Hence, we rather compare with first order methods and a memory efficient second order optimizer rfdSON, as they take similar memory\n\n- We compare optimizers via thorough hyperparameter search comprising several candidate hyperparameter configurations, hence given the limited resources we picked shampoo from the class of kronecker factored optimizers (KFAC, KBFGS), given that shampoo’s empirical performance is similar to that of KFAC and KBFGS as demonstrated in Shampoo paper [4]. \n\n- Furthermore as per your suggestion, we compared against KFAC, FishLeg and Eva (a recently proposed memory efficient second order method) on Autoencoder benchmark. We make the comparison with respect to wall clock time and notice that SONew beats all the baselines by a large margin. The plot is provided in the pdf attached to rebuttal. \n\n> The algorithms ... to another paper\n\nIt is a common practice in second order algorithms to ensure that the preconditioned gradient/update norm is in a trust region [EVA]. For example, i) update clipping in EVA, KFAC ii) grafting in shampoo [4], which ensures that the final update norm is < d, where d is the dimension of the update. In this work we chose grafting and leave exploring other ways to control the update norms as future work.\n\n[4] Anil et al. 2021 Scalable second order optimization for deep learning\n\n[EVA] Eva: Practical Second-order Optimization with Kronecker-vectorized Approximation, ICLR 2023\n\n> The following statement in section 5 is not true: \"We compare SONew against widely used first order methods including SGD, SGD with Momentum, Nesterov, Adagrad, Adam, and Rmsprop.\" The first four of the six optimizers mentioned in the list are actually not compared with.\n\nDue to space constraints we deferred the comparison with the first four to Section A.4.3 as they were least performing among the six (as mentioned in Line 285-286)\n\n> The abstract says that the proposed method is slower ... overall slower?\n\nIn all the experiments, we notice that SONew is indeed slower than first order methods, though by a very small margin (<=3% relative overhead). Though, as you rightly pointed out - it gives significant gains overall (and hence a useful method) as it converges faster. We’ll clarify this in the abstract.\n\n> The results of numerical experiments are underwhelming. As I already mentioned, no comparison with other second-order methods is given, besides Shampoo\n\nAs previously mentioned, given the computation budget, we compared with only the strongest baselines like Adam and Shampoo for large scale benchmarks. Though, as per your request, we have added comparison with FishLeg, KFac, and Eva  on Autoencoder benchmark in the additional rebuttal pdf.\n\n> For example, SGD with momentum is never compared with, “\n\nWe compare Momentum and several other first order methods on Autoencoder benchmark where we noticed they perform significantly suboptimal compared to baseline like Adam. Hence, for larger benchmarks like ViT and OGBG, we fixed Adam as the first order baseline to compare with. Though, as per your request, we’ve now added SGD+Momentum and RmsProp as well to ViT and OGBG benchmark (provided in the additional pdf).\n\n> Adam nor SGD are compared with in the case of LLM in figure 3” \n\nLLM requires a very high computational budget, hence it was not possible to add simple baseline like SGD which will surely perform poorer compared to Adam. We indeed do compare with Adam. The AdaFactor optimizer we use in LLM experiments is without factoring, which resorts to Adam like update with some additional steps like update clippings. The reason to choose AdaFactor (with no factoring) rather than Adam is it has been seen in past that to train LLMs, AdaFactor is an optimal choice (as mentioned in the paper) [10, 46]\n\n[10]Palm: Scaling language modeling with pathways, 2022\n\n[46] T. Wang, A. Roberts, D. Hesslow, T. L. Scao, H. W. Chung, I. Beltagy, J. Launay, and C. Raffel. What language model architecture and pretraining objective work best for zero-shot generalization, 2022.\n\n> Finally, training curves are shown only vs epochs or number of steps, no results in terms of wall clock time are provided, therefore it is hard to tell whether the method provides any real advantage\n\n- Our work presents a memory efficient second optimizer. Hence, a fair comparison would involve optimizers using similar memory, which happens to be first order methods or rfdSON. But for all these methods, the running time are within 5% relative overhead of each other. Hence, comparing against time is as good as comparing against steps. That said, we take your comments and also run Shampoo on OGBG (It couldn’t scale to ViT and LLM) for similar amount of time. We find that SONew performs as good as Shampoo though using significantly less memory! Thanks for your suggestion.\n\n- Additional experiment on Autoencoder benchmark with KFac, FishLeg, and Eva are performed optimizing with respect to a common wall-clock time. The plots can be foud in additional pdf attached to the rebuttal. We find that SONew beats other baselines by a great margin.\n\n\n--------\nWe hope that the rebuttal clarifies questions raised by the reviewer. We would be very happy to discuss any further questions about the work, and would really appreciate an appropriate increase in score if reviewers’ concerns are adequately addressed to facilitate acceptance of the paper.\n"}, {"review_id": "iqar3lp9EP", "reviewer": "Reviewer_3rTc", "summary": "This paper proposes a memory-efficient and scalable second-order optimizer derived from the LogDet matrix divergence measure combined with sparsity constraints. The resulting preconditioner admits a banded/tridiagonal structure found via an efficient algorithm.", "questions": "- How does the method compare with KFAC focusing on generalization on standard tasks explored in the paper and with respect to wall-clock time?\n  \n- What is the practical advantage of the proposed method compared with well-established optimizers such as Adam? Can the proposed method accelerate training speed, in terms of time to achieve SOTA accuracies? If not, how much does generalization suffer compared with SOTA first-order optimizers?\n  \n- Minor comments:\n  \n  - The notation for indexing the matrix in Eq. 12 is confusing.\n    \n  - Personal preference: I wouldn't use subjective terms such as embarassingly parallelizable or beautiful paper.", "limitations": "Some limitations are discussed.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "- The paper is well-written and fairly easy to follow. The steps are clear and logically laid out.\n  \n- I find the derivation of the preconditioner from the LogDet divergence interesting, even though variations of the particular form has been explored before.\n  \n- The theoretical framework around the proposed method looks sound and the analysis of numerical stability is appreciated.", "weaknesses": "- The banded/tridiagonal structure for which an efficient solution exists might not be natural, depending on the model architecture. A block-diagonal pattern such as in KFAC that models inter-layer dependencies would potentially be more suitable.\n  \n- The paper is positioned as a practical, scalable optimizer for large-scale deep learning tasks. As such, I am not completely convinced by the experimental results. See breakdown below.\n  \n- First, comparison with KFAC in terms of generalization and wall-clock performance would be important, as it is one of the very few higher order optimizers in deep learning that has shown adoption in practice.\n  \n- Based on the presented experiments, it is hard to tell when one would pick the proposed optimizer over a simple Adam. Training loss value is not necessarily relevant, as second-order optimizers are known to often quickly converge to local optima with low loss but poor generalization. The plot showing validation error rate does not depict how this translates to classification accuracy. In case of minor/negligible differences one would pick Adam that offers less added complexity and faster wall-clock performance. Overall, there could be more focus on wall-clock performance, as performance with respect to number of steps is not that relevant in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments.\n\n>The banded/tridiagonal structure for which an efficient solution exists might not be natural, depending on the model architecture. A block-diagonal pattern such as in KFAC that models inter-layer dependencies would potentially be more suitable.\n\nOne can indeed capture inter-layer dependencies as well. There are many ways to do this. For example, one can think of all the parameters as one set rather than viewing them as a set of parameters from different layers. This way, there’ll be cross correlations across parameters of different layers which can be captured. We leave testing this as a future work.\n\n> comparison with KFAC in terms of generalization and wall-clock performance would be important, as it is one of the very few higher order optimizers in deep learning that has shown adoption in practice. \n\n- We compare optimizers via thorough hyperparameter search comprising several candidate hyperparameter configurations, hence given the limited resources we picked shampoo from the class of kronecker factored optimizers (KFAC, KBFGS), given that shampoo’s empirical performance is similar to that of KFAC and KBFGS as demonstrated in Shampoo paper [4].\n- Taking your comment into consideration, we also run additional Autoencoder benchmark experiment comparing with KFac, FishLeg, and Eva. The plot is provided in the additional rebuttal pdf document. We notice that SONew beats all other baselines by a significant margin.\n\n> “Based on the presented experiments, it is hard to tell when one would pick the proposed optimizer over a simple Adam. Training loss value is not necessarily relevant, as second-order optimizers are known to often quickly converge to local optima with low loss but poor generalization. The plot showing validation error rate does not depict how this translates to classification accuracy.”\n\nWhen one can afford “num of params” extra memory, SONew finds the best use as demonstrated in our experiemnts. If there’s a strict memory bottleneck, one can prefer Adam.\nJust to clarify validation error rate is = 1 - validation accuracy, and since it is measured in a held out dataset which is not used in training and hence the validation error rate measures the generalization performance of the model and in practice since we don’t have test data, one always uses validation/development/held-out dataset to choose the best generalizing model. \n\n> there could be more focus on wall-clock performance, as performance with respect to number of steps is not that relevant in practice.\n\nThanks for the comment. – (already written above)\n\n> \"How does the method compare with KFAC focusing on generalization on standard tasks explored in the paper and with respect to wall-clock time?” \n\n In the limited time of the rebuttal period, we could set up and run additional baselines including KFac, FishLeg, and Eva only on the Autoencoder benchmark. These comparisons are done wrt wall-clock time. We will further test these baselines on more benchmarks and include the results in our final draft.\n\n> What is the practical advantage of the proposed method compared with well-established optimizers such as Adam? Can the proposed method accelerate training speed, in terms of time to achieve SOTA accuracies? If not, how much does generalization suffer compared with SOTA first-order optimizers?\n\n SONew significantly improves over the training time while taking extra memory only linear in the number of parameters. Specifically, we notice upto 30% faster convergence in OGBG, 10% in ViT, and upto 26% on LLMs. \n\n>notation for indexing the matrix in Eq. 12 is confusing.\n\n Thanks We’ll fix the formatting.\n\n> Personal preference: I wouldn't use subjective terms such as embarassingly parallelizable or beautiful paper.” \n\n Thanks. We’ll keep this in mind and edit such occurrences.\n\n-------\nWe hope that the rebuttal clarifies questions raised by the reviewer. We would be very happy to discuss any further questions about the work, and would really appreciate an appropriate increase in score if reviewers’ concerns are adequately addressed to facilitate acceptance of the paper.\n"}, {"review_id": "xaVHh93Xn5", "reviewer": "Reviewer_jEyd", "summary": "The paper proposes a new second-order method, named SONew, to optimize convex problems theoretically and conducts experiments on non-convex deep learning training tasks, which show SONew trains faster than first-order Adam and second-order fsdSON and Shampoo in some scenarios like the ViT benchmarks. The key idea of SONew is an online Newton method that uses an approximate LogDet divergence with a sparisified preconditioner. The sparse approximation of the LogDet divergence allows the efficient computation of second-order information so that SONew is compute-efficient enjoying exploiting some second-order information to accelerate the training.  Experiments were conducted on various DNNs like Autoencoders, ViT, GNNs, and LLMs, showing SONew’s possible acceleration over some existing first-order methods and second-order methods.", "questions": "1. In the experimental results in Section 5, some results show that SONew is obviously worse than Shampoo, e.g., Fig. 1(b), Fig. 2(b) in terms of iterations to loss (or accuracy). How to conclude that SONew achieves “up to 30% faster convergence, 3.4% relative improvement in validation performance, and 80% relative improvement in training loss” in the abstract? How about comparing to existing compute-efficient and/or memory-efficient second-order methods like M-KFAC and Eva?\n\n2. Given a deep neural network for training, how we can choose tridiag-SONew, band-4-SONew, band-8-SONew, or band-$n$-SONew to achieve good training performance?\n\n3. The paper provides a regret bound for SONew in the online convex optimization framework, but the optimizer is particularly used in DNNs which are known non-convex problems. How is the established bound used in DNN optimization?\n\n[A] M-FAC: Efficient Matrix-Free Approximations of Second-Order Information, NeurIPS 2021.\n[B] Eva: Practical Second-order Optimization with Kronecker-vectorized Approximation, ICLR 2023.", "limitations": "See above.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The proposed sparsified approximation of LogDet divergence is a new direction to reduce the computing cost using second-order information in optimizing DNNs.\n2. A theoretical analysis is provided to give an insight for the sparsified preconditioner.", "weaknesses": "1. The theoretical analysis is particularly for convex problems, so it may not hold as a DNN optimizer under non-convex optimization.\n2. How to choose the sparse graph is unknown. The paper doesn’t conclude what sparse graphs should be used for different DNNs.\n3. Experimental results show that SONew is worse than existing second-order methods like Shampoo (e.g., Table 1). In addition, some recent compute- and memory-efficient second-order methods (e.g., M-FAC [A] and Eva [B]) were not included in the comparison.\n\n[A] M-FAC: Efficient Matrix-Free Approximations of Second-Order Information, NeurIPS 2021.\n[B] Eva: Practical Second-order Optimization with Kronecker-vectorized Approximation, ICLR 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> How to choose the sparse graph is unknown. The paper doesn’t conclude what sparse graphs should be used for different DNNs.\n\n Thanks for the question. Currently we developed SONew primarily for banded diagonal sparsity. That said, exploring other sparse graphs and choosing among them, is a future work.\n> Experimental results show that SONew is worse than existing second-order methods like Shampoo (e.g., Table 1). In addition, some recent compute- and memory-efficient second-order methods (e.g., M-FAC [A] and Eva [B]) were not included in the comparison.\n\n We present a memory efficient optimizer. Hence, a fair comparison would be with other methods using memory linear in the number of parameters. Among all such methods, SONew performs significantly better. Moreover, taking your comments into consideration, we have included comparison with Eva on the Autoencoder benchmark in the additional rebuttal pdf. Please note due to computational and time budget constraints, we preferred comparing to Eva over MFac as authors show in the Eva paper that it performs better than KFac.\n\n> In the experimental results in Section 5, some results show that SONew is obviously worse than Shampoo, e.g., Fig. 1(b), Fig. 2(b) in terms of iterations to loss (or accuracy). How to conclude that SONew achieves “up to 30% faster convergence, 3.4% relative improvement in validation performance, and 80% relative improvement in training loss” in the abstract? How about comparing to existing compute-efficient and/or memory-efficient second-order methods like M-KFAC and Eva?”\n\n By that statement, we meant in comparison to other optimizers which are memory efficient. We’ll clarify this in the final draft Abstract. We also compare with Eva, as mentioned previously.\n\n> The paper provides a regret bound for SONew in the online convex optimization framework, but the optimizer is particularly used in DNNs which are known non-convex problems. How is the established bound used in DNN optimization?”\n\n\nOur analysis is on regret minimization in online convex optimization which has several connections to non-convex optimization a) [1] uses OCO learners for optimizing non-convex objectives b) similarly [2] shows transferability of regret bound analysis to non-convex non-smooth objectives.  We also briefly discuss this in A.2.5\n\n[1] N. Agarwal, R. Anil, E. Hazan, T. Koren, and C. Zhang. Learning rate grafting: Transferability of optimizer tuning, 2022\n\n[2] Ashok Cutkosky, Harsh Mehta, and Francesco Orabona. Optimal stochastic non smooth non-convex optimization through online-to-non-convex conversion. arXiv preprint arXiv:2302.03775, 2023.\n\n------\nWe hope that the rebuttal clarifies questions raised by the reviewer. We would be very happy to discuss any further questions about the work, and would really appreciate an appropriate increase in score if reviewers’ concerns are adequately addressed to facilitate acceptance of the paper.\n"}, {"review_id": "0g1XLbXQsZ", "reviewer": "Reviewer_rRBk", "summary": "This paper strives to present a second-order optimization approach, called Sparsified Online Newton (SONew). SONew solves the regret minimization via introducing the LogDet matrix divergence. Using specific sparsified structures i.e. diagonal or banded sparsity, the proposed method could achieve considerable computational complexity. Finally, the authors show the effectiveness of the proposed method by performing the MNIST task with autoencoder, ImageNet task with ViT, OGBG-molpcba with GraphNetwork and language models.", "questions": "See weakness.", "limitations": "I have not found any discussions about the limitations and potential negative societal impact. But in my opinion, this may not be a problem, since the work only focuses on the learning algorithm. Still, it is highly encouraged to add corresponding discussions.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "***Strengths***\n\n1. The paper is easy to follow and clearly written, although the layout of the paper is arranged too narrow. \n\n2. The proposed method is mathematically sound and reasonable. The paper presents novel insight of online Newton methods in regards to regret minimization with LogDet divergence of the preconditioning matrix, and proposes an elegant way solve it via specific sparsified structures, and gives mathematical analysis in regards to the regret bound, which btw is O(sqrt(T)) the same with many related methods, and numerical stability.\n\n3. Given the large computation cost and memory demand of the second-order approach, the paper shows some promising practical results in regards to the proposed methods.", "weaknesses": "***Weakness***\n\n1. The authors seem to have largely changed the spacing between sections. The current spacing presented in the paper would be much more narrow than that in the latex template. Slight changes in the latex template could be accecptable, but such large changes could not be available. I could understand that the authors try to present as much content as possible. But the writing seems to be quite wordy. For example Abstract contains 24 lines, which btw is near 15 lines for other papers in common. I understand sometimes it is necessary to give a long abstract, but 16 lines in this paper are used to describe the comprehensive results of the propose method. Also, these results are described again in the Introduction. The authors should make the paper more succinct to fit the latex template.\n\n2. For eq(3), the authors use LogDet matrix divergence to force the $X_{t}^{-1}$ and $X^{-1}_{t-1}$ to be \"close\". However, how the such closeness in LogDet matrix divergence could ensure the final term which is a complicated matrix multiplication in Eq3 to be safely omitted? Could the authors provide more details?\n\n3. The authors could compare with more second-order approaches, such as K-FAC [36], K-BFGS [21], SON [33].\n\n4. It would be highly recommended that the authors would release the code. Due to lack of details in the experiment, it may be hard to reproduce the reported results. For example, what is this large language model with 1B parameters used in the experiment?\n\n5. Line 135. The authors seem to omit the item $c_t$ in the presented formula, despite that this item will not make influence regarding this optimization. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> The authors seem to have largely changed the spacing between sections. The current spacing presented in the paper would be much more narrow than that in the latex template. Slight changes in the latex template could be accecptable, but such large changes could not be available. I could understand that the authors try to present as much content as possible. But the writing seems to be quite wordy. For example Abstract contains 24 lines, which btw is near 15 lines for other papers in common. I understand sometimes it is necessary to give a long abstract, but 16 lines in this paper are used to describe the comprehensive results of the propose method. Also, these results are described again in the Introduction. The authors should make the paper more succinct to fit the latex template.”\n\nThanks for bringing to our notice. We promise to fix the formatting in the final draft.\n\n> For eq(3), the authors use LogDet matrix divergence to force the  and to be \"close\". However, how the such closeness in LogDet matrix divergence could ensure the final term which is a complicated matrix multiplication in Eq3 to be safely omitted? Could the authors provide more details?\n\nLogdet Divergence between $X_t^{-1}$ and $X_{t-1}^{-1}$ is used as a regularizer, since the lower the logdet divergence leads to low value of the final term in eq (3). Since $w^*$ is unknown LogDet divergence serves as a proxy for this term.\n\n> The authors could compare with more second-order approaches, such as K-FAC [36], K-BFGS [21], SON [33].”. \n\nThanks for the comment. We compared with rfdSON in the paper. And in addition, we also compare with KFac, FishLeg, and Eva on the autoencoder benchmark. Please refer to the rebuttal pdf attached.\n\n> It would be highly recommended that the authors would release the code. Due to lack of details in the experiment, it may be hard to reproduce the reported results. For example, what is this large language model with 1B parameters used in the experiment?\n\nWe are sharing our code though an anonymized link to the ACs, since sharing with reviewers isn't allowed now. The 1B model used is of Primer architecture [43], as also mentioned in the paper.\n\n>“Line 135. The authors seem to omit the item  in the presented formula, despite that this item will not make influence regarding this optimization\n\nThanks for pointing this out, we’ll correct this typo in the revision. \n\n> I have not found any discussions about the limitations and potential negative societal impact. But in my opinion, this may not be a problem, since the work only focuses on the learning algorithm. Still, it is highly encouraged to add corresponding discussions.”\n\nWe promise to add a discussion, further talking about the limitation of our work.\n\n-----\nWe hope that the rebuttal clarifies questions raised by the reviewer. We would be very happy to discuss any further questions about the work, and would really appreciate an appropriate increase in score if reviewers’ concerns are adequately addressed to facilitate acceptance of the paper."}, {"review_id": "lFAzuld8fQ", "reviewer": "Reviewer_J48q", "summary": "This paper proposes SONew -- a sparsified online Newton method, which uses the LogDet matrix divergence measure to sparsify the preconditioner. The LogDet divergence regularizes by matching the largest eigenvalues of the preconditioner matrix between iterates. It also leads to a convex optimization problem since the LogDet divergence is convex in its first argument. This leads to a simple and cost effective update rule for the preconditioner matrix. By minimizng the LogDet divergence they impose a sparsity constraint while requiring the sparse preconditioner to remain close to the full-matrix preconditioner in terms of LogDet divergence. They demonstrate the capability of their optimizer on autoencoders, vision transformers, graph neural networks, and large language models.", "questions": "How robust is SONew to perturbation in the hyperparameters? Does it require less hyperparameter tuning compared to Adam?", "limitations": "In the end, we need to run many hyperparameter trials to get any optimizer to achieve optimal convergence and generalization. Therefore, an optimizer that converges twice as fast, but requires more than twice the amount of hyperparameter tuning, is unlikely to be adopted by the community.\n\nWhen proposing a novel optimizer, there should be more emphasis on the robustness to change in the hyperparameters, model, and dataset. This is even more critical for LLMs, where we cannot afford to run many hyperparameter trials. The criteria for a good optimizer should be this kind of robustness, and not the convergence.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "SONew shows faster convergence compared to first order optimizers like Adam and AdaFactor.\n\nSONew can be applied to models sizes where Shampoo runs out of memory, while showing competitive convergence and generalization to these second-order optimizers.\n\nWhen using bfloat16, SONew has the least degradation in performance compared to all other optimizers.", "weaknesses": "The experiments can be improved.\n\nTable 1 only shows the final value of the train CE loss and the time to achieve it, but without looking at the training curve it is difficult to see whether SONew is converging faster or simply reaching a lower loss but slowly. What is the criteria they used to stop the training? Is the learning rate schedule tuned for each optimizer?\n\nFor the autoencoder experiments they compare with Adam, RMSProp, Shampoo(20), and diagonal, tridiagonal, and band-4 versions of SONew. For the GraphNetwork experiment they compare with Shampoo, Adam, and rfdSON. For LLMs they only compare with AdaFactor. This inconsistency makes it difficult see the effect of changing the dataset and model on all these different optimizers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> “Table 1 only shows the final value of the train CE loss and the time to achieve it, but without looking at the training curve it is difficult to see whether SONew is converging faster or simply reaching a lower loss but slowly. What is the criteria they used to stop the training? Is the learning rate schedule tuned for each optimizer?”\n\nThe time taken to finish training (for 100 epochs) is almost same for all first order methods and tridiag-SONew. Hence, a comparison wrt wall-clock time will almost be same as wrt steps. The corresponding figure is attached in Fig. 4 (a) in the Appendix and Fig (2) in the main paper. We will also make an explicit wall-clock comparison in our final draft. Moreover, we make explicit wall-clocl time comparison on OGBG benchmark in the submitted rebuttal pdf.\n\n> “For the GraphNetwork experiment they compare with Shampoo, Adam, and rfdSON. For LLMs they only compare with AdaFactor. This inconsistency makes it difficult see the effect of changing the dataset and model on all these different optimizers”\n\nThanks for the comment. As mentioned in the Reviewer jEyd and Reviewer 3rTc, under computational constraints, we make the best calculated judgement of what baselines to compare for each benchmark. Since running Autoencoder benchmark is relatively faster, and is also commonly used to compare optimizers performance for second order methods, we make a thorough comparison on it. Moving to larger benchmarks, we choose the most commonly used first order method baseline Adam. On LLMs as well, as mentioned in the paper we choose AdaFactor, , a commonly used first order optimizer for training LLMs [46, 10].\n\n> How robust is SONew to perturbation in the hyperparameters? Does it require less hyperparameter tuning compared to Adam?”\n\nWe tune all optimizers using same number of hyperparameters. In the final draft, we promise to include an ablation studying robustness with respect to different hyperparameters.\n\n> “In the end, we need to run many hyperparameter trials to get any optimizer to achieve optimal convergence and generalization. Therefore, an optimizer that converges twice as fast, but requires more than twice the amount of hyperparameter tuning, is unlikely to be adopted by the community. When proposing a novel optimizer, there should be more emphasis on the robustness to change in the hyperparameters, model, and dataset. This is even more critical for LLMs, where we cannot afford to run many hyperparameter trials. The criteria for a good optimizer should be this kind of robustness, and not the convergence”\n\n[Graft] describes grafting (which we use in SONew experiments) which the paper proposes to make the optimizers robust w.r.t hyperparameters by using grafts from tuned baselines such as Adam. To obtain the 26% improvement in LLM we use the same hyperparameter as the tuned baseline (Adafactor), however removing grafting in SONew would require us to tune the hyperparameters all over again.\n\n[Graft] Agarwal, Naman, et al. \"Disentangling adaptive gradient methods from learning rates.\" arXiv preprint arXiv:2002.11803 (2020)."}]

Author rebuttal:
Common Rebuttal:
- As Reviewers suggested, we include additional experiment on the Autoencoder benchmark in the pdf attached. We compare SONew with KFac [1], Eva [2], and FishLeg [3]. As these baselines were not available in JAX, we used a PyTorch version as available in their official github repositories. Furthermore, as reviewers mentioned we make wall-clock comparison with these new baselines. Note that since the experiment setup is a bit different, the results of SONew don't match the ones presented in the main paper conducted using JAX. One main reason being we run 600 hyperparam configs for these experiments, and for the Autoencoder experiments presented in the main paper we do 2000 hyperparam configs.

-  Comparison in Wall-clock time: Reviewers mentioned the need to have wall clock time comparison. We want to highlight that SONew performs similar in per step time (<=2% relative overhead) compared to other memory efficient optimizers, which mostly include first order methods or rfdSON. Therefore, in a comparison with other methods using memory only linear in the number of parameters, comparison in wall-clock time will be similar to comparison in steps. Furthermore, we make a wall-clock time comparison with Shampoo in the additional pdf and find SONew performs better than Shampoo while using significantly less memory.

[1] Martens et. al, 2015 Optimizing Neural Networks with Kronecker-factored Approximate Curvature

[2] Eva: Practical Second-order Optimization with Kronecker-vectorized Approximation, ICLR 2023

[3] Gracia et al., 2023 Fisher-Legendre (FishLeg) optimization of deep neural networks.

Meta-review: Motivated by finding scaling matrices that minimize the regret bound in online convex setting, the paper introduces a scaled gradient algorithm that by leveraging the properties of LogDet matrix divergence coupled with sparsity constraints offers a great deal of memory-efficiency. The paper is well written and the theoretical derivations are solid. The numerical experiments can be improved, however, to include more models and alternative methods.
Decision: Accept (poster)