SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
L _ {j j} = 1, L _ {j + 1 j} = - \frac {H _ {j + 1 j}}{H _ {j + 1 j + 1}}, D _ {j j} ^ {- 1} = H _ {j j} - \frac {H _ {j + 1 j} ^ {2}}{H _ {j + 1 j + 1}}, j \leq n - 1 \& D _ {n n} ^ {- 1} = H _ {n n} \tag {12}
$$
Computing this explicit solution involves conducting parallellizable operations on  $2 \times 2$  principle submatrices (highlighted in red) of the tridiagonal matrix  $H$  to find the  $\hat{X}$  as shown in the following  $3 \times 3$  example:
$$
\begin{array}{l} H = \left( \begin{array}{c c c} \framebox {H _ {1 1}} & H _ {1 2} \\ \framebox {H _ {2 1}} & H _ {2 2} \\ 0 & H _ {3 2} \end{array} \right) = \left( \begin{array}{c c c} \tilde {H} _ {1 1} & \tilde {H} _ {1 2} & 0 \\ \tilde {H} _ {2 1} & \tilde {H} _ {2 2} & \tilde {H} _ {2 3} \\ 0 & \tilde {H} _ {3 2} & \tilde {H} _ {3 3} \end{array} \right) + \left( \begin{array}{c c c} g _ {1} ^ {2} & g _ {1} g _ {2} & 0 \\ g _ {1} g _ {2} & g _ {2} ^ {2} & g _ {2} g _ {3} \\ 0 & g _ {2} g _ {3} & g _ {3} ^ {2} \end{array} \right) \tag {13} \\ \rightarrow \hat {X} = \left(\begin{array}{c c c}\framebox {1}&0&0\\- \frac {H _ {2 1}}{H _ {2 2}}&1&0\\0&- \frac {H _ {3 2}}{H _ {3 3}}&1\end{array}\right)\left(\begin{array}{c c c}\framebox {H _ {1 1} - \frac {H _ {2 1} ^ {2}}{H _ {2 2}}}&0&0\\0&H _ {2 2} - \frac {H _ {2 3} ^ {2}}{H _ {3 3}}&0\\0&0&H _ {3 3}\end{array}\right)\left(\begin{array}{c c c}\framebox {1}&- \frac {H _ {2 1}}{H _ {2 2}}&0\\0&1&- \frac {H _ {3 2}}{H _ {3 3}}\\0&0&1\end{array}\right) \\ \end{array}
$$
Conducting these operations take  $\mathcal{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  $\mathcal{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  $\mathcal{G}$  be a banded matrix of band size  $b$ , i.e.  $E_{\mathcal{G}} = \{(i,j):|i - j|\leq b,1\leq i,j\leq n\}$ . For every vertex  $j$ , let  $I_{j} = \{j + 1,\dots ,j + b\}$ . Then  $X_{t} = LDL^{T}$  is the solution of (11) with nonzero entries of  $L$  and  $D$  defined as follows:
$$
L _ {j j} = 1, L _ {I _ {j} j} = - H _ {I _ {j} I _ {j}} ^ {- 1} H _ {I _ {j} j}, \quad D _ {j j} ^ {- 1} = \left(H _ {j j} - H _ {I _ {j} j} ^ {T} H _ {I _ {j} I _ {j}} ^ {- 1} H _ {I _ {j} j}\right), 1 \leq j \leq n. \tag {14}
$$
where,  $H \in S_n(\mathcal{G})$  any submatrix of  $H$  corresponding to a complete subgraph of  $\mathcal{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  $\mathcal{O}((n - b + 1)b^3)$  flops. Since  $b \ll n$ , the number of flops is  $\mathcal{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  $\mathcal{G} =$  tridiagonal/chain graph as defined in Theorem 3.1, then setting  $\epsilon = \hat{\epsilon} G_{\infty}\sqrt{T}$ ,  $\lambda_t = G_{\infty}\sqrt{t}$  and  $\eta_t = \frac{D_2}{\hat{\epsilon}\sqrt{n}}$  in Algorithm 1, where  $\| w_t - w^*\| _2\leq D_2$ ,  $\| g_t\|_{\infty}\leq G_{\infty}$  incurs a regret  $R_{T} = \mathcal{O}(\sqrt{n} G_{\infty}D_{2}\sqrt{T})$ .
The proof sketch involves deriving an explicit expression for entries of  $X_{t}^{-1}$  in Lemma A.2, to upper bound the term  $(w_{t} - w^{*})^{T}(X_{t}^{-1} - X_{t - 1}^{-1})(w_{t} - w^{*})$  in regret upper bound (3). Upper bounding  $\frac{\eta}{2}\sum_{t = 1}^{T}g_{t}^{T}X_{t}g_{t}$  involves using the Loewner order  $X_{t}\preceq \| X_{t}\|_{2}I_{n}\preceq \| X_{t}\|_{\infty}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.
Optimizer	Time complexity	Memory complexity
Adam	O(d1d2)	d1d2
rfdSON(m)	O(m2d1d2)	md1d2
Shampoo	O(d13+d23)	(d12+d22)
tridiag-SONew	O(d1d2)	2d1d2
band-4-SONew	O(d1d2)	5d1d2
Table 1: Consider preconditioning a  $d_{1} \times d_{2}$  parameter matrix. Time complexity of tridiag and banded SONew inverse scale linearly with number of parameters, but, Shampoo is cubic in the dimensions of the matrix. Memory used to store second-moments of gradients by tridiag-SONew can be significantly lower than Shampoo, for e.g. if  $d_{1} = 4d_{2}$ , then Shampoo takes  $>2$  times more memory.
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  $\mathcal{G}$  should be positive definite, however, in practice, due to finite precision, each entry of  $H$  is inherently perturbed with an error proportional to  $\mathcal{O}(\epsilon_{mach})$ , where  $\epsilon_{mach}$  is machine epsilon [27]. We notice in practice that the subtraction operation in  $D_{jj}^{-1} = S_{jj} = H_{jj} - H_{j+1j}^{2} / H_{j+1j+1}$  (line 7 Algorithm 2), which has a condition number  $\kappa_{sub} = |H_{jj}| / |S_{jj}|$ , can be high as  $S_{jj}$  can be arbitrarily low due to near singular submatrices  $\begin{bmatrix} H_{ii} & H_{ii+1} \\ H_{i+1i} & H_{i+1i+1} \end{bmatrix}$ . Thus small perturbation in  $H$  can lead to high perturbations in the preconditioner  $\hat{X}$ . We formalize this notion by deriving an end-to-end componentwise condition number (pg. 135, problem 7.11 in [27]) of SPARSIFIED_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  $\mathcal{G}$ , we can remove edges  $(j,j+1)$  which correspond to low  $S_{jj} < \gamma$ , where  $\gamma \geq 0$  denotes a tolerance parameter. We show in Theorem A.11, Appendix A.3 that this reduces the condition number upper bound of SPARSIFIED_INVERSE. Furthermore, we generalize this to banded sparsity pattern in Algorithm 3, Appendix A.3.
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  $\mathcal{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  $\mathcal{G}$  such as banded graph with band-size  $b$ , for which our preconditioning process is more computationally efficient with a time complexity of  $\mathcal{O}(b^3 (n - b + 1))$  compared to online-newton method  $\mathcal{O}(n^2)$ .
Shampoo [24, 5] approximates full gradient statistics matrix using Kronecker factored preconditioners to reduce the memory and time complexity from  $\mathcal{O}(n^2)$  to  $\mathcal{O}(d_1^2 + d_2^2)$  and  $\mathcal{O}(d_1^3 + d_2^3)$  respectively. Here,  $n = d_1d_2$  denotes number of parameters for a linear layer of dimensions  $d_1 \times 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  $\mathcal{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 -  $\mathcal{O}(n^2)$  flops of Online Newton Step (ONS) to  $\mathcal{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  $\mathbb{R}^{n\times 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].
Table 2: float32 experiments on Autoencoder benchmark. We observe that diag-SONew performs the best among all first order methods while taking similar time. tridiag and band-4 SONew perform significantly better than first order methods while requiring similar linear space and time. Shampoo performs best but takes  $\mathcal{O}(d_1^3 + d_2^3)$  time for computing preconditioner of a linear layer of size  $d_1 \times d_2$ , whereas our methods take  $\mathcal{O}(d_1 d_2)$  time, as mentioned in Section 3.3. rfdSON takes similar space as SONew but performs considerably worse.
Optimizer	First Order Methods	Second Order Methods
Adagrad	RMSProp	Adam	diag-SONew	Shampoo(20)	rfdSON(1)	rfdSON(4)	tridiag-SONew	band-4-SONew
Train CE loss	54.393	53.330	53.591	53.025	50.702	56.21	55.55	51.723	51.357
Time(s)	62	62	62	63	371	85	300	70	260
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 OGBG-molpcba 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 us to use the direction suggested by Optimizer-2 with step size suggested by Optimizer-1. The final update is given by  $\frac{\|v_1\|}{\|v_2\|} \cdot 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  $\left\| v_{1}\right\|$  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 "band- $k$ " in tables and figures for band size of  $k$ . We compare SONew against widely used first order methods including SGD [32]), SGD with Momentum [41], Nesterov [40], Adagrad [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  $\sum_{i} g_{i} g_{i}^{T}$ . 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) * \# \text{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, tridiag-SONew, band-4-SONew. Note that  $d_{1}^{2} + d_{2}^{2} \geq 2d_{1}d_{2} \forall d_{1}, d_{2}$ , thus memory used by tridiag-SONew is never more than Shampoo. We use a 2.72M parameters 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 \times$  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
(b) GraphNetwork validation avg. precision
Figure 2: (a) Comparison of SONew with first-order optimizers, rfdSON, and Shampoo on Autoencoder benchmark in float32 training. We observe that SONew performs better than all first order methods, and second order methods using the same memory.
Figure 3: Comparison of SONew and Adafactor on LLM training. SONew takes  $26\%$  less steps to reach the same performance as AdaFactor. Using the same number of steps, it achieves  $\sim 1.7\%$  relative better log perplexity.
Figure 1: (a) Best validation error runs for tridiag-SONew vs Momentum, RMSProp, Adam, rfdSON, and Shampoo on (a) VIT benchmark (b) GraphNetwork benchmark. We notice that tridiag-SONew achieves same performance as Adam, the next best baseline using similar space and time, with  $10\%$  and  $30\%$  less steps/time in ViT and GraphNetwork respectively. While using the same number of steps, SONew achieves relatively  $0.7\%$  and  $\sim 3.4\%$  better validation error respectively. Shampoo doesn't fit in the 16GB memory of TPU v2 for ViT benchmark, hence we couldn't perform hyperparameter tuning on it. On GraphNetwork, compared to Shampoo, tridiag-SONew gives similar performance while being far more memory efficient (Refer Appendix A.4.2 for more details).
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.
Comparison 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  $\sim 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 ( $\sim 22\mathrm{M}$  parameters) and GraphNetwork ( $\sim 3.5\mathrm{M}$  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  $\sim 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  $\sim 10\%$  less steps to reach the same performance as Adam for VIT, and  $\sim 30\%$  less steps for GraphNetwork benchmark. Training for the same number of steps, we get  $\sim 0.7\%$  better relative validation error for VIT and  $\sim 3.4\%$  better relative validation average precision for GraphNetwork. On GraphNetwork benchmark, tridiag-SONew performs  $1.3\%$  relatively worse in average precision compared to Shampoo, while being  $1.25\times$  faster. On VIT benchmark, Shampoo doesn't fit in a 16 GB TPU v2. Its statistics require 155M entries ( $\sim 7\times$  #params) while tridiag-SONew requires only 44M entries ( $2\times$  #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  $64k$  tokens. All experiments were performed on 16 TPU v4s. To support efficient training of large models, we implemented a sharded tridiag-SONew following model parallelism approach.
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 \times$  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 Newton (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 Kronecker-factorization based second-order methods for training deep networks. Our experimental results show that SONew uses similar time as first order methods and achieves much better validation and training performance in various benchmarks. In the future, we plan to explore different sparsity graphs for which efficient solutions exist for the LogDet subproblem (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  $\mathcal{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.