Title: A Computationally Efficient Sparsified Online Newton Method
Venue: NeurIPS 2023
Accepted: yes
Keywords: Optimization

Second order methods; Deep Learning
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

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

1 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 an update direction, by multiplying the gradient with a preconditioner matrix carefully chosen from gradients observed in previous iterations, to improve convergence. (Full-

matrix) Adagrad [15], online Newton method [25] 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 d 1 × d 2 parameter matrix, these methods require matrix inverse operations, which take O ( d 3 1 + d 3 2 ) time and O ( d 2 1 + d 2 2 ) space. In comparison, first-order methods use O ( d 1 d 2 ) time and memory, which is linear in the number of parameters. For instance, when d 1 = kd 2 , the memory used by Shampoo, d 2 1 + d 2 2 floating point numbers is O ( 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 a 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 a novel approach of minimizing the second term while regularizing two successive preconditioners to be close in 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 a 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 is 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 O ( √ T ) 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 in 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-molpcba, 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 × # params for ViT whereas tridiag-SONew uses only 2 × # params for its statistics. We also test another recently proposed memory efficient second order optimizer, rfdSON [37], but found its performance suboptimal to the best performing first order method. Owing to SONew'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 [11]. SONew achieves the same performance as AdaFactor using 26% fewer steps, resulting in a 1 . 35 × 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 (Equation (13)) without complex engineering challenges, rendering it even more useful and practical.

2 Background

The inner product between matrices is defined as ⟨ A,B ⟩ = Tr( A T B ) , where Tr( . ) denotes the matrix trace. The Frobenius norm of a matrix A is ∥ A ∥ F = √ Tr( A T A ) , while its spectral norm is ∥ A ∥ 2 = max x ∥ Ax ∥ 2 / ∥ x ∥ 2 . We use I n ∈ R n × n to denote an identity matrix. We use S n , S ++ n to denote the set of symmetric, and positive definite matrices respectively. The generalized norm of a

vector x ∈ R n with respect to matrix A ∈ S ++ n is defined as ∥ x ∥ A = √ x T Ax . We use det ( A ) to denote the determinant of matrix A , and diag( A ) to denote the diagonal matrix with diag( A ) ii = A ii . We use G and ˜ G to denote a graph and its sub-graph with a vertex set [ n ] = { 1 , . . . , n } . Let E G denote set of edges in graph G , and neig G ( i ) denote neighbours of vertex i in graph G . A sparse symmetric matrix A ∈ R n × n follows a sparsity structure graph G if A i,j = 0 ∀ ( i, j ) / ∈ E G , . Note that set of all such matrices form a linear subspace. We use S n ( G ) ++ to denote the set of positive definite matrices with sparsity structure given by graph G , i.e, if X ∈ S n ( G ) ++ , then X ij = 0 ∀ ( i, j ) / ∈ E ( G ) . S n ( G ) ++ is an open convex set. Given an index set I = { i 1 , i 2 , .., i n } , we use A II to denote the corresponding principal sub-matrix of A .

2.1 LogDet matrix divergence

Let ϕ : S ++ n → R be a strictly convex, differentiable function. The Bregman matrix divergence between X,Y ∈ S ++ n is defined as [8, 34]: D ϕ ( X,Y ) = ϕ ( X ) -ϕ ( Y ) -Tr( ∇ ϕ ( Y ) T ( X -Y )) . Since ϕ is convex, D ϕ ( X,Y ) ≥ 0 for all X,Y ∈ S ++ n . For example if ϕ ( X ) = ∥ X ∥ 2 F , the corresponding Bregman divergence D ϕ ( X,Y ) = ∥ X -Y ∥ 2 F is the squared Frobenius distance. In this paper, we extensively use the convex function ϕ ( X ) = -log det ( X ) ; the corresponding divergence measure D ℓ d ( X,Y ) is called the LogDet matrix divergence :

The LogDet divergence is scale invariant to invertible matrices A , i.e. D ℓ d ( A T XA,A T Y A ) = D ℓ d ( X,Y ) . LogDet divergence can be written in terms of eigendecompositions of X = V Σ V T and Y = U Θ U T [34]:

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

Weset up our problem under the online convex optimization framework (OCO) [43, 26], where at each round the learner makes a prediction w t and receives a convex loss f t ( w t ) and gradient g t = ∇ f t ( w t ) as feedback. The goal of the learner is to reduce regret R T by predicting w t so that a low aggregate loss ∑ T t =1 f t ( w t ) is achieved compared to the best possible, w ∗ = arg min w ∑ T t =1 f t ( w ) . Formally, regret is given by

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 w t +1 := w t -ηX t g t , where X t ∈ R n × n . Then ∥ w t +1 -w ∗ ∥ 2 X -1 t = ∥ w t -ηX t g t -w ∗ ∥ 2 X -1 t = ∥ w t -w ∗ ∥ 2 X -1 t + η 2 g T t X t g t -2 η ( w t -w ∗ ) T g t . The convexity of f t implies that f t ( w t ) -f t ( w ∗ ) ≤ ( w t -w ∗ ) T g t leading to f t ( w t ) -f t ( w ∗ ) ≤ 1 2 η ( ∥ w t -w ∗ ∥ 2 X -1 t -∥ w t +1 -w ∗ ∥ 2 X -1 t + η 2 g T t X t g t ) . Summing over all t ∈ [ T ] and rearranging reveals the following upper bound on overall regret:

Since w ∗ is unknown, finding X t which minimizes (3) is infeasible. So to minimize regret, we attempt to minimize the second term in (3) while regularizing X -1 t to be 'close' to X -1 t -1 . The nearness measure we choose is the LogDet matrix divergence, thus leading to the following objective

where D ℓ d is as in (1). Why do we use the LogDet divergence? From (2), due to the term λ i /θ j , D ℓ d ( X,X t -1 ) prioritizes matching the smaller eigenvalues of X t -1 with those of X , i.e., matching the larger eigenvalues of X -1 t -1 and X -1 . As a consequence, LogDet divergence regularizes X by matching up its large eigenvalues with those of X t -1 . For example if smallest and largest eigenvalue of X t -1 are θ n and θ 1 , then for an eigenvalue σ of X , when σ > θ n , θ 1 , the penalty from (2) for θ n is higher than for θ 1 , ( σ/θ n -log( σ/θ n ) -1) > ( σ/θ 1 -log( σ/θ 1 ) -1) . This intuition leads us to formulate (4) 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, 44] can be shown to be the unique solution obtained when the LogDet divergence between successive preconditioners, subject to a secant constraint, is minimized (as shown in the 4-page 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,λ t ) = g T t Xg t + λ t (D ℓ d ( X,X t -1 ) -c t ) = Tr( Xg t g T t ) + λ t ( -log det ( XX -1 t -1 ) +Tr( XX -1 t -1 ) -n )) -λ t c t . Setting ∇L ( X,λ t ) = 0 , and using the fact that ∇ log det ( X ) = X -1 we get the following update rule:

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 a general update rule applicable to numerous existing optimizers. For example, setting c t = 0 (equivalently λ t = ∞ ) ∀ t ∈ [ n ] in (4) results in no change to the preconditioner in any round. In this case, with X 0 = I n , we get online gradient descent [54]. On the other hand, setting λ t = 1 gives the update rule of the online Newton method [25]. Our update rule differs from (full-matrix) Adagrad [15] which has X -2 t = X -2 t -1 + g t g T t .

Maintaining and updating X t as in (5) is possible by using Sherman-Morrison formula but requires O ( n 2 ) 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 t 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 a 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 consider the following problem: find a sparse positive definite X with ∥ X ∥ 0 ≤ αn , α > 1 , such that the objective D ℓ d ( X, ( X -1 t -1 + g t g T t /λ t ) -1 ) is minimized. Essentially, this problem imposes a sparsity constraint while requiring the sparse preconditioner to remain close to the full-matrix preconditioner in terms of LogDet divergence.

Due to the L 0 -norm constraint, this is a non-convex problem, which makes it difficult to solve exactly. Since L 1 -norm serves as a convex relaxation for the L 0 norm, we could use it instead, resulting in the following optimization problem also known as graphical lasso estimator [19]:

However, the time taken to 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 a fixed undirected graph G . To sparsify the solution in (5), we formulate the subproblem

where S n ( G ) ++ denotes the set of positive definite matrices with the fixed sparsity pattern corresponding to 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, (6) can be surprisingly solved in O ( n ) and easily parallelizable, for certain sparsity structures G . Algorithm 1 and 2 presents an instantiation of the proposed SONew method, which solves (6) using O ( n ) time and memory for banded matrices with band size b . In particular a tridiagonal matrix, corresponding to a chain graph, is a banded matrix with bandsize 1.

Algorithm 1 Sparsified Online Newton (SONew) Algorithm

Inputs : λ t := coefficient in the update (10), G := sparsity graph (banded/tridiagonal), ϵ := damping parameter, T := total number of iterations/mini-batches, η t := step size/learning rate. Output : w T +1

1:

2:

H

0

=

ϵI

d

,

w

1

= 0

for

t

∈ {

1

, . . . , T

}

do

3: compute g t = ∇ f t ( w t )

4: H t := H t -1 + P G ( g t g T t /λ t ) ∈ S n ( G ) with P G as in (8). ▷ O ( n ) time & memory

5: Get L, D = SPARSIFIED_INVERSE ( H t , G ) , where X t = LDL T solves (11).

6: Compute descent direction u t = LDL T g t , 7: w t +1 = w t -η t u t

8:

end for

9: return w T +1

Algorithm 2 SPARSIFIED_INVERSE( H, G ) in O ( n ) flops

Inputs : H ∈ S n ( G ) , is as (10). G := the banded graph of band size b ≪ n Outputs : lower triangular banded L ∈ R n × n and diagonal matrix D ∈ R n × n 1: function SPARSIFIED_INVERSE( H , G ) 2: L := 0 , D := 0 3: L jj := 1 , ∀ j ∈ [ n ] 4: for j ∈ { 1 , . . . , n } do ▷ parallelizable 5: Let H jI j and H I j I j be defined as in Section 2, where I j = { j +1 , . . . , j + b } ∩ [ n ] , 6: Solve for L I j j in the linear system H I j I j L I j j = -H I j j ▷ O ( b 3 ) time. 7: D jj := 1 / ( H jj + H T I j j L I j j ) 8: end for 9: return L, D 10: end function

Maintaining H t ∈ S n ( G ) in line 4 . Solving the subproblem in (6) naively is impractical since X -1 t -1 is a dense matrix. However, the structure of the LogDet divergence comes to the rescue; the optimization problem in (6) can be expanded as follows:

Let us define the projection onto S n ( G ) , P G : R n × n → R n × n as:

Note that the Tr( . ) term in (7) is dependent only on the non-zero elements of X ∈ S n ( G ) ++ , since Tr( AB ) = ⟨ A,B ⟩ , for symmetric matrices A and B . Hence, (7) can be written as

Computing the entire matrix X -1 t -1 can be avoided by analyzing the optimality condition of (9). Let g ( X ) = -log det ( X ) + ⟨ X,P G ( X -1 t -1 + g t g T t /λ t ) ⟩ denote the objective function in (9), then the optimality condition of (9) is P G ( ∇ g ( X )) = P G ( ∇ ( -log det( X )+ ⟨ X,P G ( X -1 t -1 + g t g T t /λ t )) ⟩ = 0 , since gradients with respective nonzero entries of X should be zero, ∂g ( X ) ∂X i,j = ( ∇ X ( g ( X ))) i,j = 0 , ∀ ( i, j ) ∈ E G . Using ∇ ( -log det( X )) = -X -1 , ∇ X ( ⟨ X,Y ⟩ ) = Y , and setting X = X t gives:

Thus we only need to maintain H t = P G ( X -1 t ) . This matrix is updated as H t = H t -1 + P G ( g t g T t /λ t ) . Since H t ∈ S n ( G ) , the update can be done in O ( | E G | ) memory and time, while computing the matrix X -1 t would have cost O ( n 2 ) . In SONew (Algorithm 1), this key observation is used to maintain H t in line 4.

Computing X t in line 5 . Nowthat H t is known at every round t , we can replace P G ( X -1 t -1 + g t g T t /λ t ) in (9) with H t as:

For an arbitrary graph G , solving (11) might be difficult. Theorems 3.1 and 3.2 show embarrassingly parallelizable explicit solutions to the subproblem (11) for tridiagonal and banded sparsity patterns.

Theorem 3.1 (Explicit solution of (11) for tridiagonal structures/chain graph) . Let the sparsity structure G be a chain with edges E G = { ( i, j ) : | i -j | ≤ 1 , 1 ≤ i, j ≤ n } . Also, let H ∈ S n ( G ) be such that any submatrix of H corresponding to a complete subgraph of G is positive definite, then the solution of (11) is given by ˆ X = LDL T , where the unit lower bidiagonal matrix L and diagonal matrix D have the following non-zero entries:

Computing this explicit solution involves conducting paralellizable operations on 2 × 2 principle submatrices (highlighted in red) of the tridiagonal matrix H to find the ˆ X as shown in the following 3 × 3 example:

Conducting these operations take O ( n ) time and memory complexity, and similarly the descent direction can be found sequentially by X t g t = L ( D ( L T g t )) , 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 (11) for banded structures) . Let the sparsity pattern G be a banded matrix of band size b, i.e. E G = { ( i, j ) : | i -j | ≤ b, 1 ≤ i, j ≤ n } . For every vertex j , let I j = { j +1 , . . . , j + b } . Then X t = LDL T is the solution of (11) with nonzero entries of L and D defined as follows :

where, H ∈ S n ( G ) any submatrix of H corresponding to a complete subgraph of G is positive definite.

Note that Theorem 3.1 is a special case of Theorem 3.2 when b is set to 1 , and the proof for Theorem 3.2 is given in Appendix A.1. 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 3 ) flops. Since b ≪ n , the number of flops is O ( n ) .

3.3 Regret bound analysis of 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 = tridiagonal/chain graph as defined in Theorem 3.1, then setting ϵ = ˆ ϵG ∞ √ T , λ t = G ∞ √ t and η t = D 2 ˆ ϵ √ n in Algorithm 1, where ∥ w t -w ∗ ∥ 2 ≤ D 2 , ∥ g t ∥ ∞ ≤ G ∞ incurs a regret R T = O ( √ nG ∞ D 2 √ T ) .

The proof sketch involves deriving an explicit expression for entries of X -1 t in Lemma A.2, to upper bound the term ( w t -w ∗ ) T ( X -1 t -X -1 t -1 )( w t -w ∗ ) in regret upper bound (3). Upper bounding η 2 ∑ T t =1 g T t X t g t involves using the Loewner order X t ⪯ ∥ X t ∥ 2 I n ⪯ ∥ X t ∥ ∞ I n . A detailed proof sketch and proof is given in Appendix A.2. We note here that though the regret bound presented here is 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 9.

3.4 Numerical Stability of SONew

In Theorem 3.1 and Theorem 3.2, as mentioned, any submatrix of H t corresponding to a complete subgraph of G should be positive definite, however, in practice, due to finite precision, each entry of H is inherently perturbed with an error proportional to O ( ϵ mach ) , where ϵ mach is machine epsilon [27]. We notice in practice that the subtraction operation in D -1 jj = S jj = H jj -H 2 j +1 j /H j +1 j +1 (line 7 Algorithm 2), which has a condition number κ sub = | H jj | / | S jj | , can be high as S jj can be arbitrarily low due to near singular submatrices

in H can lead to high perturbations in the preconditioner ˆ X . We formalize this notion by deriving an end-to-end componentwise condition number (pg. 135, problem 7.11 in [27]) of SPARSI-FIED_INVERSE in Theorem A.10,Appendix A.3. To reduce this condition number upper bound and be robust to perturbations in H t caused by finite precision, for a tridiagonal graph G , we can remove edges ( j, j + 1) which correspond to low S jj < γ , where γ ≥ 0 denotes a tolerance parameter. We show in Theorem A.11,Appendix A.3 that this reduces the condition number upperbound of SPARSIFIED_INVERSE. Furthermore, we generalize this to banded sparsity pattern in Algorithm 3,Appendix A.3.

4 Related Work

Online Newton method is a 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 [25, 26]. However, it has a time complexity of O ( n 2 ) , making it infeasible for large n . However, introduction of LogDet divergence measure in SONew allows us to set different sparsity graphs as G such as banded graph with band-size b , for which our preconditioning process is more computationally efficient with a time complexity of O ( b 3 ( n -b +1)) compared to online-newton method O ( n 2 ) .

Shampoo [24, 5] approximates full gradient statistics matrix using Kronecker factored preconditioners to reduce the memory and time complexity from O ( n 2 ) to O ( d 2 1 + d 2 2 ) and O ( d 3 1 + d 3 2 ) respectively. Here, n = d 1 d 2 denotes number of parameters for a linear layer of dimensions d 1 × d 2 . The time complexity of matrix inversion takes a heavy toll in Shampoo's compute time even with the Kronecker product assumption on the preconditioner, whereas, our method has a time complexity of O ( b 3 d 1 d 2 ) quadratic in dimensions of the linear layer (note that b = 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], a 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 is prior work [35, 36] in reducing the complexity O ( n 2 ) flops of Online Newton Step (ONS) to O ( n ) flops using sketching. These ONS variants maintain a low rank approximation of H t (as in Algorithm 1) and updating it with a new gradient g t at every iteration requires conducting SVD [36]/orthonormalization [35] of a tall and thin matrix in R n × r , where r denotes the rank of approximation of H t . In Section 5, we conduct large scale experiments and compare SONew against rfdSON [37] as it's more stable than Oja-SON [35].

LogDet problem in equation 11 is closely related to the Maximum Determinant Matrix Completion (MDMC) [4, 49]. The MDMC problem is the dual of LogDet problem (11), 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 [38] in the context of a gaussian graphical model interpretation of gradients, specifically, KFAC used a block-tridiagonal preconditioner to incorporate correlation within consecutive layers.

5 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 OGBGmolpcba 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 v 1 of Optimizer-1 and v 2 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 ∥ v 1 ∥ ∥ v 2 ∥ · v 2 . Grafting has been shown to take advantage of a tuned optimizer step size and improve performance. For SONew and rfdSON, we use Adam grafting - using Adam optimizer step size ∥ v 1 ∥ with SONew/rfdSON direction v 2 / ∥ v 2 ∥ . For Shampoo, we use its default RMSProp grafting. We couldn't find rfdSON official implementation, so we use our own implementation using which we reproduced the numbers on convex losses (Appendix A.4) reported in their paper [37].

5.1 Autoencoder benchmark

Setup: We use three sparsity patterns for SONew - a) diagonal sparsity, resulting in a diagonal preconditioner similar to adaptive first order methods like Adam and Adagrad; b) tridiagonal sparsity, corresponding to a chain graph; and c) banded sparsity, represented by "bandk " 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 [14], Adam [33], and Rmsprop [48]. We also compare with rfdSON [37], a recently proposed memory efficient second order optimizer and with Shampoo [24], a state of the art second-order optimizer used in practice, albeit with considerable memory 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, rfdSON maintains a low rank approximation of the Online Newton method's statistics matrix ∑ i g i g T i . We observed rfdSON with adam grafting always performed better than without grafting, hence report the corresponding numbers. We evaluate rfdSON with rank m approximation, denoted as rfdSON( m ), which requires ( m +1) ∗ # params space when using grafting. For a fair comparison with tridiag-SONew and band-4-SONew, we test rfdSON with m = 1 and m = 4 , respectively. For shampoo, computing preconditioner at every step could be infeasible, instead it is computed every t steps - referred to as Shampoo( t ). Section 3.3 compares time and memory complexities of rfdSON, Shampoo, tridiagSONew, band-4-SONew. Note that d 2 1 + d 2 2 ≥ 2 d 1 d 2 ∀ d 1 , d 2 , thus memory used by tridiag-SONew is never more than Shampoo. We use a 2 . 72 Mparameters Autoencoder and each experiment is performed 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 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 5 × 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, we reduce the precision to bfloat16 and conduct similar

(a) VIT validation error

Momentum

0.2811

5h35m

0.281

RMSProp

0.2515

5h36m

rfdSON

0.2432

6h05m

Adam

0.2285

5h43m

validation error rate

tridiag-SONew 0.2268

5h53m

0.251

0.243

8:239

10%less steps

55

60

65

70

75

steps (*1.8k)

(b) GraphNetwork validation avg. precision

0.289

30%speedup

0.280

0.270

0.26

Shampoo

0.290

tridiag-SONew0.288

Adam

0.279

rfdSON

0.277

RMSProp

0.267

Momentum

0.221

0.22

0

20

40

60

80

100

120

140

Time (min)

60

Adam

RMSProp

diag_sonew

shampoo(20)

tridiag-SONew

band-4-SONew

train loss

54

53

52

51

0

20

40

60

80

100

epochs

adafactor

2.897

2d14h30m

tridiag-SONew 2.847

2d16h25m

26%less steps

2.84

100k

200k

300k

400k

500k

steps

experiments in Appendix A.4.4 (Table 8 ). 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 a thorough comparison and study of bfloat16 experiments. In Figure 2 we plot the loss curves of all the baselines and SONew for float32 experiments. Moreover, in Appendix A.4.1 Table 4 we provide ablation on performance of SONew with varying batch sizes.

Conparison with other baselines: We further compare SONew with KFAC [38], FishLeg [20], and Eva [52] for completeness. Since these methods lack a JAX implementation, we adopted the authors' official Pytorch implementations. When we attempted to integrate their code with our Autoencoder benchmark, the results were unsatisfactory; for instance, FishLeg recorded a loss of approximately ∼ 60 . 0 . This was notably unexpected as it underperformed Adam, a benchmark that the authors themselves compared against. Given these results and to minimize modifications to the official code, we decided to 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) [31]. 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 TPUs (v2) were used, skewing the comparative results. All memory-efficient methods, including rfdSON, 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 a 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 1. tridiag-SONew requires ∼ 10% less steps to reach the same performance as Adam for VIT, and ∼ 30% less steps for GraphNetwork 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 GraphNetwork. On GraphNetwork benchmark, tridiagSONew performs 1 . 3% relatively worse in average precision compared to Shampoo, while being 1 . 25 × faster. On VIT benchmark, Shampoo doesn't fit in a 16 GB TPU v2. Its statistics require 155M entries ( ∼ 7 × # params ) while tridiag-SONew requires only 44M entries ( 2 × # 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, tridiag-SONew optimizer'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 tridiag-SONew 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, we test it on a Large Language Model (LLM) [47] with 1 billion parameters. We compare SONew with AdaFactor (without factoring), a commonly used first order optimizer for training LLMs [51, 11]. 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 a batch size of 64 k 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.

Results: We report the experiment in Figure 3 where we find that SONew beats Adafactor by a large margin. Specifically, SONew achieves the same log perplexity using 26% less steps. Moreover, using the same number of tokens, SONew achieves 1 . 7% relative better performance on train loss, leading to 1 . 35 × speedup. This shows the potential of SONew as a scalable optimizer that can be used to train large models while using second order information.

6 Conclusions and Future Work

In this paper we have introduced a novel Sparsified Online Newtwon (SONew) method that yields a 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 Kroneckerfactorization 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 (11) 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-the 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.

References

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.

N. Agarwal, R. Anil, E. Hazan, T. Koren, and C. Zhang. Learning rate grafting: Transferability of optimizer tuning, 2022.

S.-I. Amari. Natural gradient works efficiently in learning. Neural computation , 10(2):251-276, 1998.

M. S. Andersen, J. Dahl, and L. Vandenberghe. Logarithmic barriers for sparse matrix cones. Optimization Methods and Software , 28(3):396-423, 2013.

R. Anil, V. Gupta, T. Koren, K. Regan, and Y. Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09018 , 2020.

P. W. Battaglia, J. B. Hamrick, V. Bapst, A. Sanchez-Gonzalez, 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. URL https://arxiv.org/abs/1806.01261 .

M. Bollhöfer, A. Eftekhari, S. Scheidegger, and O. Schenk. Large-scale sparse inverse covariance matrix estimation. SIAM Journal on Scientific Computing , 41(1):A380-A401, 2019.

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.

C. G. Broyden. Quasi-Newton methods and their application to function minimisation. Mathematics of Computation , 21(99):368-381, 1967.

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.

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, P. Barnes, Y. Tay, N. Shazeer, V. Prabhakaran, E. Reif, N. Du, B. Hutchinson, R. Pope, J. Bradbury, 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, D. Ippolito, D. Luan, H. Lim, B. Zoph, A. Spiridonov, R. Sepassi, D. Dohan, S. Agrawal, M. Omernick, A. M. Dai, T. S. Pillai, M. Pellat, A. Lewkowycz, E. Moreira, R. Child, O. Polozov, K. Lee, Z. Zhou, X. Wang, B. Saeta, M. Diaz, O. Firat, M. Catasta, J. Wei, K. Meier-Hellstern, D. Eck, J. Dean, S. Petrov, and N. Fiedel. Palm: Scaling language modeling with pathways, 2022.

L. Deng. The MNIST database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine , 29(6):141-142, 2012.

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. URL https://arxiv.org/abs/ 2010.11929 .

J. Duchi, E. Hazan, and Y. Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research , 12(61):2121-2159, 2011. URL http://jmlr.org/papers/v12/duchi11a.html .

J. Duchi, E. Hazan, and Y. Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research , 12(7), 2011.

S. Fattahi and S. Sojoudi. Graphical lasso and thresholding: Equivalence and closed-form solutions. Journal of machine learning research , 2019.

R. Fletcher. A new approach to variable metric algorithms. The computer journal , 13(3): 317-322, 1970.

R. Fletcher. A new variational result for quasi-Newton formulae. SIAM Journal on Optimization , 1(1):18-21, 1991.

J. Friedman, T. Hastie, and R. Tibshirani. Sparse inverse covariance estimation with the graphical lasso. Biostatistics , 9(3):432-441, 2008.

J. R. Garcia, F. Freddi, S. Fotiadis, M. Li, S. Vakili, A. Bernacchia, and G. Hennequin. Fisherlegendre (fishleg) optimization of deep neural networks. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id= c9lAOPvQHS .

J. Godwin*, T. Keck*, P. Battaglia, V. Bapst, T. Kipf, Y. Li, K. Stachenfeld, P. Veliˇ ckovi´ c, and A. Sanchez-Gonzalez. Jraph: A library for graph neural networks in jax., 2020. URL http://github.com/deepmind/jraph .

D. Goldfarb. A family of variable-metric methods derived by variational means. Mathematics of computation , 24(109):23-26, 1970.

D. Goldfarb, Y. Ren, and A. Bahamou. Practical quasi-Newton methods for training deep neural networks. Advances in Neural Information Processing Systems , 33:2386-2396, 2020.

V. Gupta, T. Koren, and Y. Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning , pages 1842-1850. PMLR, 2018.

E. Hazan, A. Agarwal, and S. Kale. Logarithmic regret algorithms for online convex optimization. Machine Learning , 69(2):169-192, 2007.

E. Hazan et al. Introduction to online convex optimization. Foundations and Trends® in Optimization , 2(3-4):157-325, 2016.

N. J. Higham. Accuracy and stability of numerical algorithms . SIAM, 2002.

G. Hinton, N. Srivastava, and K. Swersky. Neural networks for machine learning lecture 6a overview of mini-batch gradient descent. Cited on , 14(8):2, 2012.

C.-J. Hsieh, M. A. Sustik, I. S. Dhillon, P. K. Ravikumar, and R. Poldrack. BIG & QUIC: Sparse inverse covariance estimation for a million variables. Advances in neural information processing systems , 26, 2013.

W. Hu, M. Fey, M. Zitnik, Y. Dong, H. Ren, B. Liu, M. Catasta, and J. Leskovec. Open graph benchmark: Datasets for machine learning on graphs, 2020. URL https://arxiv.org/abs/ 2005.00687 .

N. P. Jouppi, G. Kurian, S. Li, P. 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.

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.1214/aoms/1177729392. URL https://doi.org/10.1214/aoms/1177729392 .

D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

B. Kulis, M. A. Sustik, and I. S. Dhillon. Low-rank kernel learning with Bregman matrix divergences. Journal of Machine Learning Research , 10(2), 2009.

H. Luo, A. Agarwal, N. Cesa-Bianchi, and J. Langford. Efficient second order online learning by sketching. Advances in Neural Information Processing Systems , 29, 2016.

L. Luo, C. Chen, Z. Zhang, W.-J. Li, and T. Zhang. Robust frequent directions with application in online learning. The Journal of Machine Learning Research , 20(1):1697-1737, 2019.

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.org/papers/v20/17-773.html .

J. Martens and R. Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. In International conference on machine learning , pages 2408-2417. PMLR, 2015.

S. Merity, C. Xiong, J. Bradbury, and R. Socher. Pointer sentinel mixture models, 2016.

Y. Nesterov. A method for unconstrained convex minimization problem with the rate of convergence o (1 /k 2 ) . 1983.

N. Qian. On the momentum term in gradient descent learning algorithms. Neural Networks , 12 (1):145-151, 1999. ISSN 0893-6080. doi: https://doi.org/10.1016/S0893-6080(98)00116-6. URL https://www.sciencedirect.com/science/article/pii/S0893608098001166 .

J. Schmidhuber. Deep learning in neural networks: An overview. Neural Networks , 61:85117, jan 2015. doi: 10.1016/j.neunet.2014.09.003. URL https://doi.org/10.1016%2Fj. neunet.2014.09.003 .

S. Shalev-Shwartz et al. Online learning and online convex optimization. Foundations and Trends® in Machine Learning , 4(2):107-194, 2012.

D. F. Shanno. Conditioning of quasi-Newton methods for function minimization. Mathematics of computation , 24(111):647-656, 1970.

N. Shazeer and M. Stern. Adafactor: Adaptive learning rates with sublinear memory cost, 2018.

H.-J. M. Shi, T.-H. Lee, S. Iwasaki, J. Gallego-Posada, Z. Li, K. Rangadurai, D. Mudigere, and M. Rabbat. A distributed data-parallel pytorch implementation of the distributed shampoo optimizer for training neural networks at-scale, 2023.

D. R. So, W. Ma´ nke, H. Liu, Z. Dai, N. Shazeer, and Q. V. Le. Primer: Searching for efficient transformers for language modeling, 2022.

T. Tieleman and G. Hinton. Lecture 6.5-rmsprop: Divide the gradient by a running average of its recent magnitude. coursera: Neural networks for machine learning. 2012.

L. Vandenberghe, M. S. Andersen, et al. Chordal graphs and semidefinite optimization. Foundations and Trends® in Optimization , 1(4):241-433, 2015.

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need, 2023.

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.

L. Zhang, S. Shi, and B. Li. Eva: Practical second-order optimization with kronecker-vectorized approximation. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=_Mic8V96Voy .

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.

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.

A Supplementary material

A.1 Properties of LogDet subproblem

Proof of Theorem 3.2

The optimality condition of (11) is P G ( X -1 ) = P G ( H ) , X ∈ S ++ n ( G ) . Let Z = L -T D -1 L -1 , then P G ( Z ) = H

Let J j = I j ∪ j , where I j = { j +1 , . . . , j + b } as defined in the theorem, then select J j indices of vectors on both sides of the second equality above and selecting the J j indices :

Note that L -T is an upper triangular matrix with ones in the diagonal hence J th j block of L -T e j will be [1 , 0 , 0 , . . . ] . Also, since P G ( Z ) = H

Substituting this in the linear equation 15

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

graph is a chordal graph with perfect elimination order { 1 , 2 , . . . , n } . Furthermore, X t is positive definite, since as ( H jj -H T I j j H -1 I j I j H I j j ) is a schur complement of submatrix of H formed by J j = I j ∪ { j } .

Proof of Theorem 3.1 The proof follows trivially from Theorem 3.1, when b is set to 1 .

A.2 Regret bound analysis

Proof sketch of Theorem 3.3. We decompose the regret into R T ≤ T 1 + T 2 + T 3 in Lemma A.1 and individually bound the terms. Term T 2 = 1 2 η · ∑ T -1 t =1 ( w t +1 -w ∗ ) T ( X -1 t +1 -X -1 t )( w t +1 -w ∗ ) depends on closeness of consecutive inverses of preconditioners, ( X -1 t +1 -X -1 t ) , to upperbound this we first give explicit expressions of X -1 t for tridiagonal preconditioner in Lemma A.2 in Appendix A.2.2. This explicit expression is later used to bound each entry of ( X -1 t +1 -X -1 t ) with O (1 / √ t ) in Appendix A.2.4, this gives a O ( √ T ) upperbound on T 2 . To show an upperbound on T 3 = ∑ T t =1 η 2 · g T t X t g t , we individually bound g T t X t g t by using a Loewner order X t ⪯ ∥ X t ∥ 2 I n ⪯ ∥ X t ∥ ∞ I n and show that ∥ X t ∥ ∞ = O (1 / √ T ) and consequently T 3 = O ( √ T ) .

A.2.1 Regret bound decomposition

In this subsection we state Lemma A.1 which upper bound the regret R T using three terms T 1 , T 2 , T 3 .

Lemma A.1 ( [26] ) . In the OCO problem setup, if a prediction w t ∈ R n is made at round t and is updated as w t +1 := w t -ηX t g t using a preconditioner matrix X t ∈ S ++ n

Proof.

Using the convexity of f t , f t ( w t ) -f t ( w ∗ ) ≤ ( w t -w ∗ ) T g t , where g t = ∆ f t ( w t ) and summing over t ∈ [ T ]

The first summation can be decomposed as follows

Substituting the above identity in the Equation (19) proves the lemma.

Let R T ≤ T 1 + T 2 + T 3 , where

A.2.2 Properties of tridiagonal preconditioner

In this subsection, we derive properties of the tridigonal preconditioner obtained from solving the LogDet subproblem (11) with G set to a chain graph over ordered set of vertices { 1 , . . . , n } :

The second equality holds true only when H t is positive definite. Although in Algorithm 1 we maintain a sparse H t = H t -1 + P G ( g t g T t /λ t ) , H 0 = ϵI n which is further used in (22) to find the preconditioner X t , our analysis assumes the full update H t = H t -1 + g t g T t /λ t , H 0 = ϵI n followed by preconditioner X t computation using (23). Note that the preconditioners X t generated both ways are the same, as shown in Section 3.2.

The following lemma shows that the inverse of tridiagonal preconditioners used in Algorithm 1, will restore H i,j , when ( i, j ) fall in the tridiagonal graph, else, the expression is related to product of H i + k,i + k +1 corresponding to the edges in the path from node i to j in chain graph. This lemma will be used later in upperbounding T 2 .

Lemma A.2 ( Inverse of tridiagonal preconditioner ) . If G = chain/tridiagonal graph and ˆ X = arg min X ∈ S n ( G ) ++ D ℓ d ( X,H -1 ) , then the inverse ˆ X -1 has the following expression

Where ˆ X ( j ) is the j th column of ˆ X . Let ˆ Y denote the right hand side of Equation (24).

Proof.

The third equality is by using the following alternative form of Equation (12):

where i < j . Similarly, the offdiagonals of ˆ Y ˆ X can be evaluated to be zero as follows.

n

where X t is defined as in Lemma A.2.

Using the identity of spectral radius ρ ( X ) ≤ ∥ X ∥ ∞ and since ˜ X is positive definite, ∥ ∥ ∥ ˜ X -1 t ∥ ∥ ∥ 2 ≤ ∥ ˜ X -1 t ∥ ∞

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 1 and T 3 . In next subsection, we'll upper bound T 2 as well.

where, ˆ ϵ is a constant in parameter ϵ = ˆ ϵG ∞ √ T and consequently used in initializing H 0 = ϵI n in line 1 in Algorithm 1,

Proof.

Note that [ ( H t ) ii ( H t ) ii +1 ( H t ) i +1 i ( H t ) i +1 i +1 ] ⪰ ϵ [ 1 0 0 1 ] (using line 1 in Algorithm 1), thus det ([ ( H t ) ii ( H t ) ii +1 ( H t ) i +1 i ( H t ) i +1 i +1 ]) ≥ det ( ϵ [ 1 0 0 1 ]) = ϵ 2 . The numerator last inequality can be upperbounded by bounding ( H t ) ii individually as follows:

Substituting the above in (28) gives

Lemma A.5 ( Upperbound of T 1 ) .

where D 2 = max t ∈ [ T ] ∥ w t -w ∗ ∥ 2 and G ∞ = max t ∥ g t ∥ ∞

Proof. Since X T is positive definite

Lemma A.6 ( O ( √ T ) upperbound on T 3 ) .

where, ∥ g t ∥ ∞ ≤ G ∞ and parameters ϵ = ˆ ϵG ∞ √ T , λ t = G ∞ √ t in Algorithm 1.

Proof. Using Theorem 3.1, nonzero entries of X t can be written as follows:

where, E G denote the set of edges of the chain graph G in Theorem 3.1. Also, for brevity, the subscript is dropped for H t . Let ˆ X t = √ diag( H ) X t √ diag( H ) , then ˆ X t can be written as

where, ˆ H ij = H ij / √ H ii H jj . Note that ˆ X t ⪯ ∥ ˆ X t ∥ 2 I n ⪯ ∥ ˆ X t ∥ ∞ I n , using max {| λ 1 ( ˆ X t ) | , . . . , | λ n ( ˆ X t ) |} ≤ ∥ ˆ X t ∥ ∞ (property of spectral radius). So we upperbound ∥ ˆ X t ∥ ∞ = max i ∈ [ n ] {| ( ˆ X t ) 11 | + | ( ˆ X t ) 12 | , . . . , | ( ˆ X t ) ii -1 | + | ( ˆ X t ) ii | + | ( ˆ X t ) ii +1 | , . . . , | ( ˆ X t ) nn | + | ( ˆ X t ) nn -1 |} next. Individual terms | ( ˆ X t ) ii -1 | + | ( ˆ X t ) ii | + | ( ˆ X t ) ii +1 | can be written as follows:

The last inequality is because | ˆ H ij | ≤ 1 . Thus, ∥ ˆ X t ∥ ∞ ≤ 2 max i ∈ [ n -1] 1 1 -| ˆ H ii +1 | . Now

Using diag( H t ) ⪰ ϵI n (step 1 in Algorithm 1), where ϵ = ˆ ϵG ∞ √ T as set in Lemma A.8, gives

Summing up over t gives

A.2.4 O ( √ T ) Regret

In this section we derive a regret upper bound with a O ( T 1 / 2 ) growth. For this, we upper bound T 2 as well in this section. In (21), T 2 = ∑ T t =2 ( w t -w ∗ ) T ( X -1 t -X -1 t -1 )( w t -w ∗ ) can be upper bounded to a O ( T 1 / 2 ) by upperbounding entries of X -1 t -X -1 t -1 individually. The following lemmas constructs a telescoping argument to bound ∣ ∣ ( X -1 t -X -1 t -1 ) i,j ∣ ∣ .

Lemma A.7. Let H, ˜ H ∈ S ++ n , such that ˜ H = H + gg T /λ , where g ∈ R n , then

Proof.

The following Lemma bounds the change in the inverse of preconditioner Y -1 , when there is a rank one perturbation to H ≻ 0 in following LogDet problem (11) :

Lemma A.8 ( Rank 1 perturbation of LogDet problem (11) ) . Let H, ˜ H ∈ S ++ n , such that ˜ H = H + gg T /λ , where g ∈ R n . Also, ˜ Y = argmin X ∈ S n ( G ) ++ D ℓ d ( X, ˜ H ) and Y = arg min X ∈ S n ( G ) ++ D ℓ d ( X,H ) , where G is a chain graph, then

where i, i + k ≤ n , G ∞ = ∥ g ∥ ∞ and max i,j | H ij | / √ H ii H jj ≤ β < 1 . Let κ (diag( H )) := condition number of the diagonal part of H , then κ := max( κ (diag( H )) , κ (diag( ˜ H ))) .

Proof. Using Lemma A.2 will give the following:

where N ij = H ij / √ H ii H jj < 1 (Since determinants of 2x2 submatrices of H are positive). Expanding ˜ N ii +1 = N ii +1 + θ ii +1 (from Lemma A.7), subsequently ˜ N ii +2 = N ii +2 + θ ii +2 and so on will give

where max i,j | N i,j | , max i,j | ˜ N i,j | ≤ β < 1 . Expanding θ i + li + l +1 from Lemma A.7 in the term | θ i + li + l +1 | √ ˜ H ii ˜ H i + ki + k will give:

Since H i + li + l H i + l +1 i + l +1 ≤ ˜ H i + li + l ˜ H i + l +1 i + l +1 ,

Using the above, H i,i /H j,j ≤ κ , and | g i | ≤ G ∞ , ∀ i, j ∈ [ n ] , gives

Thus the following part of ∣ ∣ ∣ ∣ ( ˜ Y -1 -Y -1 ) ii + k ∣ ∣ ∣ ∣ can be upperbounded:

Lemma A.9 ( O ( √ T ) upper bound of T 2 ) . Given that κ (diag( H t )) ≤ κ , ∥ w t -w ∗ ∥ 2 ≤ D 2 , max i,j | ( H t ) ij | / √ ( H t ) ii ( H t ) jj ≤ β < 1 , ∀ t ∈ [ T ] in Algorithm 1, then T 2 in Appendix A.2.1 can be bounded as follows:

where λ t = G ∞ √ t , and ϵ = ˆ ϵG ∞ √ T in Algorithm 1, and ˆ ϵ ≤ 1 is a constant.

Proof. Note that T 2 = 1 2 η · ∑ T -1 t =1 ( w t +1 -w ∗ ) T ( X -1 t +1 -X -1 t )( w t +1 -w ∗ ) ≤ ∑ T -1 t =1 D 2 2 ∥ ∥ ( X -1 t +1 -X -1 t ) ∥ ∥ 2 / (2 η ) . Using ∥ A ∥ 2 = ρ ( A ) ≤ ∥ A ∥ ∞ for symmetric matrices A , we get

Now using κ ≤ 2 / ˆ ϵ (using Equation (29) and ( H t ) ii > ˆ ϵ ) and summing up terms in T 2 using the above will give the result.

Putting together T 1 , T 2 and T 3 from Lemma A.5, Lemma A.9 and Lemma A.6 respectively, when ϵ , λ t are defined as in Lemma A.9:

A.2.5 Non-convex guarantees

Minimizing smooth non-convex functions f is a complex yet interesting problem. In Agarwal et al. [1], this problem is reduced to an online convex optimization, where a sequence of objectives f t ( w ) = f ( w ) + c ∥ w -w t ∥ 2 2 are minimized. Using this approach Agarwal et al. [1] established convergence guarantees to reach a stationary point via regret minimization. Thus non-convex guarantees can be obtained from regret guarantees and is our main focus in the paper.

A.3 Numerical stability

In this section we conduct perturbation analysis to derive an end-to-end componentwise condition number (pg. 135, problem 7.11 in [27]) upper bound of the tridiagonal explicit solution in Theorem 3.1. In addition to this, we devise Algorithm 3 to reduce this condition number upper bound for the tridiagonal sparsity structure, and be robust to H t which don't follow the non-degeneracy condition: any principle submatrix of H t corresponding to a complete subgraph of G .

Theorem A.10 (Condition number of tridiagonal LogDet subproblem (11)) . Let H ∈ S ++ n be such that H ii = 1 for i ∈ [ n ] . Let ∆ H be a symmetric perturbation such that ∆ H ii = 0 for i ∈ [ n ] , and H +∆ H ∈ S ++ n . Let P G ( H ) be the input to 11, where G is a chain graph, then

where, β i = H ii +1 , κ ℓd ∞ := componentwise condition number of (11) for perturbation ∆ H .

The tridiagonal LogDet problem with inputs H as mentioned in Theorem A.10, has high condition number when 1 -β 2 i = H ii -H 2 ii +1 /H i +1 i +1 are low and as a result the preconditioner X t in

Setting η = D 2 ˆ ϵ √ n

SONew (Algorithm 1) has high componentwise relative errors. We develop Algorithm 3 to be robust to degenerate inputs H , given that H ii > 0 . It finds a subgraph ˜ G of G for which non-degeneracy conditions in Theorem 3.2 is satisfied and (14) is well-defined. This is done by removing edges which causes inverse H -1 I j I j to be singular or ( H jj -H T I j j H -1 I j I j H I j j ) to be low. In the following theorem we also show that the condition number upper bound in Theorem A.10 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 bfloat16 where we demonstrate noticeable improvement in performance when Algorithm 3 is used.

Theorem A.11 (Numerically stable algorithm) . Algorithm 3 finds a subgraph ˜ G of G , such that explicit solution for ˜ G in (14) is well-defined. Furthermore, when G is a tridiagonal/chain graph, the component-wise condition number upper bound in (33) is reduced upon using Algorithm 3, ˆ κ ˜ G ℓd < ˆ κ G ℓd , where ˆ κ ˜ G ℓd , ˆ κ G ℓd are defined as in Theorem A.10 for graphs ˜ G 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: Gtridiagonal or banded graph, H -symmetric matrix in R n × n with sparsity structure G and H ii > 0 , γ -tolerance parameter for low schur complements.

2: Output: Finds subgraph ˜ G of G without any degenerate cases from Lemma A.13 and finds preconditioner ˆ X corresponding to the subgraph

3: Let E i = { ( i, j ) : ( i, j ) ∈ E G } be edges from vertex i to its neighbours in graph G .

4: Let V + i = { j : i < j, ( i, j ) ∈ E G } and V -i = { j : i > j, ( i, j ) ∈ E G } , denote positive and negative neighbourhood of vertex i .

5: Let K = { i : H ii -H T I i i H -1 I i I i H I i i is undefined or ≤ γ }

6: Consider a new subgraph ˜ G with edges E ˜ G = E G \ ( ⋃ i ∈ K E i ∪ ( V + i × V -i ))

7: return ˆ X := SPARSIFIED_INVERSE ( ˜ H t , ˜ G ) , where ˜ H t = P ˜ G ( H t )

A.3.1 Condition number analysis

Theorem A.12 ( Full version of Theorem A.10 ) . Let H ∈ S ++ n such that H ii = 1 , for i ∈ [ n ] and a symmetric perturbation ∆ H such that ∆ H ii = 0 , for i ∈ [ n ] and H +∆ H ≻ 0 . Let ˆ X = arg min X ∈ S n ( G ) ++ D ℓ d ( X,H -1 ) and ˆ X + ∆ ˆ X = arg min X ∈ S n ( G ) ++ D ℓ d ( X, ( H +∆ H ) -1 ) , here G := chain/tridiagonal sparsity graph and S n ( G ) ++ denotes positive definite matrices which follows the sparsity pattern G .

where, κ ℓd := condition number of the LogDet subproblem, κ 2 ( . ) := condition number of a matrix in ℓ 2 norm, β i = H ii +1 / √ H ii H i +1 i +1

Proof. Consider the offdiagonals for which ( ˆ X + ∆ ˆ X ) ii +1 = -H ii +1 / (1 -H 2 ii +1 ) = f ( H ii +1 ) ,where f ( x ) = -x/ (1 -x 2 ) . Let y = f ( x ) , ˆ y = f ( x + ∆ x ) and | ∆ x/x | ≤ ϵ then using Taylor series

Using the above inequality, with x := H and y := ˆ X

Let g ( x ) = x 2 / (1 -x 2 ) , let y 1 = g ( w 1 ) , y 2 = g ( x 2 ) , ˆ y 1 = g ( w 1 +∆ x ) , ˆ y 2 = g ( x 2 +∆ x ) . Using Taylor series

Putting x 1 := H ii +1 , x 2 := H ii -1 and analyzing y 1 := H 2 ii +1 / (1 -H 2 ii +1 ) and y 2 := H 2 ii -1 / (1 -H 2 ii -1 ) will result in the following

Since ˆ X ii = 1 + H 2 ii +1 / (1 -H 2 ii +1 ) + H 2 ii -1 / (1 -H 2 ii -1 ) . Putting together Equation (35) and Equation (34), the theorem is proved.

A.3.2 Degenerate H t

In SONew (1), the H t = P G ( ∑ t s =1 g s g T s /λ t ) generated in line 4 could be such that the matrix ∑ t s =1 g s g T s /λ t need not be positive definite and so the schur complements H ii -H 2 ii +1 /H i +1 i +1 can be zero, giving an infinite condition number κ ℓd ∞ by Theorem A.10. The following lemma describes such cases in detail for a more general banded sparsity structure case.

Lemma A.13 (Degenerate inputs to banded LogDet subproblem) . Let H = P G ( GG T ) , when ϵ = 0 in Algorithm 1, where G ∈ R n × T and let g ( i ) 1: T be i th row of G , which is gradients of parameter i for T rounds, then H ij = 〈 g ( i ) 1: T , g ( j ) 1: T 〉 .

Case 1: For tridiagonal sparsity structure G : if g ( j ) 1: T = g ( j +1) 1: T , then H jj -H 2 jj +1 /H j +1 j +1 = 0 .

Case 2: For b > 1 in (14): If rank( H J j J j ) = rank( H I j I j ) = b , then ( H jj -H T I j j H -1 I j I j H I j j ) = 0 and D jj = ∞ . If rank( H I j I j ) < b then the inverse H -1 I j I j doesn't exist and D jj is not well-defined.

For b > 1 , since H I j I j , using Guttman rank additivity formula, rank( H jj -H 2 jj +1 /H j +1 j +1 ) = rank( H J j J j ) -rank ( H I j I j ) = 0 , thus H jj -H 2 jj +1 /H j +1 j +1 = 0 .

Furthermore, if rank( H ) ≤ b , then all b +1 × b +1 principal submatrices of H have rank b , thus ∀ j , H J j J j have a rank b , thus D jj for all j are undefined.

If GG T = ∑ T i =1 g i g i is a singular matrix, then solution to the LogDet problem might not be welldefined as shown in Lemma A.13. For instance, Case 1 can occur when preconditioning the input layer of an image-based DNN with flattened image inputs, where j th and ( j +1) th pixel can be highly correlated throughout the dataset. Case 2 can occur in the first b iterations in Algorithm 1 when the rank of submatrices rank( H I j I j ) < b and ϵ = 0 .

A.3.3 Numerically Stable SONew proof

Proof of Theorem A.11

Let I i = { j : i < j, ( i, j ) ∈ E G } and I ′ i = { j : i < j, ( i, j ) ∈ E ˜ G } Let K = { i : H ii -H T I i i H -1 I i I i H I i i is undefined or 0 , i ∈ [ n ] } denote vertices which are getting removed by the algorithm, then for the new graph ˜ G , D ii = 1 /H ii , ∀ i ∈ K since H ii > 0 . Let ¯ K = { i : H ii -H T I i i H -1 I i I i H I i i > 0 , i ∈ [ n ] } . Let for some j ∈ ¯ K , if

denotes the nearest connected vertex higher than j for which D ll is undefined or zero, then according to the definition E ˜ G in Algorithm 3, I ′ j = { j +1 , . . . l -1 } ⊂ I j , since D jj is well-defined, H I j I j is invertible, which makes it a positive definite matrix (since H is PSD). Since H jj -H T I j j H -1 I j I j H I j j > 0 , using Guttman rank additivity formula H J j J j ≻ 0 , where J j = I j ∪ j . Since H J ′ j J ′ j is a submatrix of H J j J j , it is positive definite and hence its schur complement H jj -H T I ′ j j H -1 I ′ j I ′ j H I ′ j j > 0 . Thus for all j ∈ [ n ] , the corresponding D jj 's are well-defined in the new graph ˜ G .

Note that κ ˜ G ℓd = max i ∈ [ n -1] 1 / (1 -β 2 i ) < max i ∈ ¯ K 1 / (1 -β 2 i ) = κ G ℓd , for tridiagonal graph, where β i = H ii +1 , in the case where H ii = 1 . This is because the arg max i ∈ [ n -1] 1 / (1 -β 2 i ) ∈ 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 to better preconditioners. We confirm this through experiments on the Autoencoder benchmark where we take band size = 0 (diag-SONew), 1 (tridiag-SONew), 4, and 10 in Table 3.

Effect of mini-batch size To find the effect of mini-batch size, in Table 4, We empirically compare SONew with state of the art first-order methods such as Adam and RMSProp, and second-order method Shampoo. We see that SONew performance doesn't deteriorate much when using smaller or larger batch size. First order methods on the other hand suffer significantly. We also notice that Shampoo doesn't perform better than SONew in these regimes.

Effect of Numerical Stability Algorithm 3 On tridiag-SONew and banded-4-SONew, we observe that using Algorithm 3 improves training loss. We present in Table 5 results where we observed significant performance improvements.

A.4.2 Memory Requirements

We present a list of approximate memory requirements of different optimizers across different benchmarks in Table 6. Note that for K-FAC and Shampoo, because preconditioner is updated once only a few steps, they require storing the latest computed preconditioners as well along with the statistics, causing even higher memory overhead.

A.4.3 Hyperparaeter search space

We provide the hyperparamter search space for experiments presented in Section 5. We search over 2 k hyperparameters for each Autoencoder experiment using a Bayesian Optimization package. The search ranges are: first order momentum term β 1 ∈ [1 e -1 , 0 . 999] , second order momentum term β 2 ∈ [1 e -1 , 0 . 999] , learning rate ∈ [1 e -7 , 1 e -1] , ϵ ∈ [1 e -10 , 1 e -1] . We give the optimal hyperparameter value for each experiment in Table 12. For VIT and GraphNetwork benchmark, we search β 1 , β 2 ∈ [0 . 1 , 0 . 999] , lr ∈ [1 e -5 , 1 e -1] , ϵ ∈ [1 e -9 , 1 e -4] , weight decay ∈ [1 e -5 , 1 . 0] , learning rate warmup ∈ [2% , 5% , 10%] ∗ total_train_steps, dropout ∈ [00 , 0 . 1] , label smoothing over { 0 . 0 , 0 . 1 , 0 . 2 } . We use cosine learning rate schedule. Batch size was kept = 1024, and 512 for Vision Transformer, and GraphNetwork respectively. We sweep over 200 hyperparameters in the search space for all the optimizers.

For rfdSON [37], there's no ϵ hyperparameter. In addition to the remaining hyperparameters, we tune α ∈ { 1 e -5 , 1 . 0 } (plays similar role as ϵ ) and µ t ∈ [1 e -5 , 0 . 1] .

For LLM [47] benchmark, we only tune the learning rate ∈ [1 e -2 , 1 e -3 , 1 e -4] while keeping the rest of the hyperparams as constant. This is due to the high cost of running experiments hence we only tune the most important hyperparameter. For Adafactor [45], we use factored=False, decay method=adam, β 1 = 0 . 9 , weight decay= 1 e -3 , decay factor= 0 . 99 , and gradient clipping=1.0.

A.4.4 Additional Experiments

VIT and GraphNetwork Benchmarks : In Figure 5 we plot the training loss curves of runs corresponding to the best validation runs in Figure 1. Furthermore, from an optimization point of view, we plot the best train loss runs in Figure 6 got by searching over 200 hyperparameters. We find that tridiag-SONew is 9% and 80% relatively better in ViT and GraphNetwork benchmark respectively (Figure 6), 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 4(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 8). We notice that SONew undergoes the least degradation. Tridiagonal-sparsity 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 22 × slower

100

(a) float32 - autoencoder

(b) bfloat16 - autoencoder

60

SGD

Momentum

SSOI

Nesterov

train

AdaGrad

Adam

54

RMSProp

diag_sonew

53

shampoo(20)

tridiag-SONew

52

band-4-SONew

rfdson(1)

51

rfdson(4)

0

20

40

60

80

100

epochs

70

SGD

train loss

Momentum

Nesterov

AdaGrad

60

Adam

RMSProp

diag_sonew

shampoo(20)

tridiag-SONew

54

band-4-SONew

52

51

53

rfdson(1)

rfdson(4)

0

20

40

60

80

epochs

than tridiagonal-SONew. The corresponding loss curves are given in Figure 4(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 Figure 4(a). Note that there's no qualitiative change in the results and none of the claims made in the paper are affected. SONew is still significantly better than rfdSON. 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.

(a) VIT train CE loss

(b) GraphNetwork train CE loss

loss

0.025

0.024

train

0.023

0.022

rfdSON

0.0238

RMSProp

0.0230

0.021

Momentum

0.0229

0.020

Adam

0.0210

tridiag-SONew0.0208

Shampoo

0.0200

20

40

60

80

100

120

140

train CE loss of the best valid error rate run(s)

0.84

0.78

0.75

0.70

RMSProp

0.846

5h36m

0.67

Momentum

0.711

5h35m

Adam

0.752

5h43m

tridiag-SONew0.706

5h53m

0.62

rfdSON

0.625

6h05m

55

60

65

70

75

steps (*1.8k)

Time (min)

(a) Best VIT train CE loss

(b) Best GraphNetwork train CE loss

0.84

0.71

train CE loss

RMSProp

0.8462

5h36m

Momentum

0.7511

5h35m

rfdSON

0.5650

6h05m

best

0.56

Adam

tridiag-SONew 0.4239

0.4661

5h43m

5h53m

0.48

0.46

0.44

0.42

55

60

65

70

75

steps (*1.8k)

2.3e-02

1e-02

best train CE loss

9e-04

Momentum

2.3e-02

RMSProp

1e-02

rfdSON

9.2e-03

2.3e-04

Adam

9e-04

Shampoo

2.3e-04

1.2e-04

tridiag-SONew1.2e-04

20

40

60

80

100

120

140

Time (min)

FishLeg

56.0

KFac

55.7

Eva

55.5

tridiag-SONew53.5

Train loss

56

55.5

54.5

53.5

30

40

50

60

70

80

90

time (s)

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, we use their official PyTorch implementation and run tridiag-SONew in PyTorch as well. Hyperparameter search is conducted for SONew similar to as reported above, over learning rate, β 1 , β 2 , and ϵ . For KFAC and Eva, rather than β 2 , damping factor is tuned over [1 e -5 , 10] (default value specified is 0.03). kl_clip is tuned as well over [1 e -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 ∈ [1 e -7 , 1 e -1] and damping ∈ [1 e -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

(a) Validation CE loss

adam

3.357 2h40m

3.48

tridiag-SONew

3.3102h43m

loss

validation1

3.4

3.36

3.34

3.32

3.31

80

85

90

95

100

steps (*1k)

(b) Train CE loss

2.59

tridiag-SONew2.51 2h43m

adam

2.482h40m

2.55

SSO

train

2.52

2.5

2.48

80

85

90

95

100

steps (*1k)

each experiment takes same amount of time. Each optimizer is tuned using 600 hyperparameters. The results are in Figure 7, where notice that tridiag-SONew beats all the baselines by a 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 ∈ 2 e -2 , 1 e -2 , 5 e -3 , 1 e -3 , β 2 ∈ 0 . 99 , 0 . 999 , weight decay ∈ 0 . 0 , 0 . 1 , and ϵ ∈ 1 e -10 , 1 e -8 , 1 e -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 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 to convex optimization, we compare SONew to rfdSON [37], another recent memory-efficient second-order Newton method. We follow [37] for the experiment setup - each dataset is 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 rfdSON with m = 2 . Since the code isn't open sourced, we implemented it ourselves. In order to show reproducibility with respect to the reported numbers in [37], we include results with m = 5 as well. We see in the Table 9 that tridiag-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 10: (a) Dataset stats

Table 14: (b) bfloat16 experiments optimal hyperparamters

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nSONew 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nHow robust is SONew to perturbation in the hyperparameters? Does it require less hyperparameter tuning compared to Adam?\n\nLIMITATIONS:\nIn 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n***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.\n\nWEAKNESSES:\n***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.\n\nQUESTIONS:\nSee weakness.\n\nLIMITATIONS:\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nSee above.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\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  \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.\n\nQUESTIONS:\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- 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.\n\nLIMITATIONS:\nSome limitations are discussed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1) 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\nWEAKNESSES:\n1) 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\nQUESTIONS:\nSee above\n\nLIMITATIONS:\nSee above", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n•\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\nWEAKNESSES:\nThe 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:\n\nQUESTIONS:\n•\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\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 1, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-29"}