SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In this section, we adapt FD-update to Shampoo [5]. For simplicity, we optimize over  $\mathbb{R}^{m\times n}$  in Alg.3 projection may be handled as in Alg.2 Similar to Sketchy AdaGrad, Sketchy Shampoo uses the FD approach outlined in Alg.1 to sketch the left and right preconditioning matrices for Shampoo. In particular, we maintain two parallel sketching streams using Alg.1 to produce sketches  $\bar{L}_t,\bar{R}_t$  for the left and right preconditioning matrices. We keep track of the cumulative escaped masses  $\rho_{1:t}^L$  and  $\rho_{1:t}^R$  from sketching the left and right preconditioning matrices, respectively, and compensate the cumulative escaped mass to create the left and right Sketchy preconditioning matrices  $\tilde{L}_t,\tilde{R}_t$ .
Algorithm 3 Sketchy Shampoo (S-Shampoo)
1: Input: step size  $\eta$ , time horizon  $T$ .
2: Initialize  $X_0 = 0_{m\times n}$ $\tilde{L}_0 = \varepsilon I_m$ $\bar{R}_0 = \varepsilon I_n$ $\bar{L}_0 = 0_m$ $\bar{R}_0 = 0_n$
3: for  $t = 1, \dots, T$  do
4: Play  $X_{t}$ , suffer  $f_{t}(X_{t})$ , receive  $G_{t} \in \partial f_{t}(X_{t})$ .
5: Sketch  $(\rho_t^L,\bar{L}_t) = \mathsf{FD - update}(\bar{L}_{t - 1},G_tG_t^\top)$ $(\rho_t^R,\bar{R}_t) = \mathsf{FD - update}(\bar{R}_{t - 1},G_t^\top G_t)$
6: Update  $\tilde{L}_t = \bar{L}_t + \rho_{1:t}^L I_m$ ,  $\tilde{R}_t = \bar{R}_t + \rho_{1:t}^R I_n$  and  $X_{t+1} = X_t - \eta \tilde{L}_t^{-1/4} G_t \tilde{R}_t^{-1/4}$ .
7: end for
Denote  $L_{T}\stackrel {\mathrm{def}}{=}\sum_{t = 1}^{T}G_{t}G_{t}^{\top} + \varepsilon I$  and  $R_{T}\stackrel {\mathrm{def}}{=}\sum_{t = 1}^{T}G_{t}^{\top}G_{t} + \varepsilon I.$
Theorem 5. Suppose  $G_{1}, \ldots, G_{T}$  have rank at most  $r$ . Then Alg.3 run with  $\eta = D / \sqrt{2r}$  guarantees the following regret bound:
$$
R e g r e t _ {T} (S - S h a m p o o) \leq \sqrt {2 r} D \left(\operatorname {t r} \left(L _ {T} ^ {1 / 4}\right) + m \Omega_ {L, \ell} ^ {1 / 4}\right) \left(\operatorname {t r} \left(R _ {T} ^ {1 / 4}\right) + n \Omega_ {R, \ell} ^ {1 / 4}\right),
$$
where  $D = \max_{t\in [T]}\| X_t - X^*\| _F$  and  $\Omega_{L,\ell},\Omega_{R,\ell}$  are analogous bounds for  $\rho_{1:T}^{L},\rho_{1:T}^{R}$  from Lem.
Proof. See Sec.B.5.1
Bounds may be improved analogous to Cor.4 for Alg.3, but we omit the similar statement due to space.
4.3 Exponentially Weighted FD
This section discusses the modification of Alg.1 to support exponential moving averages. Early in algorithm development, we noticed that attempting to approximate the unweighted sum of factored gradient covariances  $\sum_{t}G_{t}G_{t}^{\top}$  and  $\sum_{t}G_{t}^{\top}G_{t}$  with FD tended to an estimate of covariance that was roughly 0, creating numerical instabilities. Note that FD guarantee (Lem.1) still holds—but the error term  $\rho_{1:T}$  becomes greater than  $\| G_T\|$ , resulting in a vacuous bound due to lack of spectral decay.
Indeed, Fig.3 motivating this work only confirmed that the exponential moving average  $L_{t}(\beta_{2}) = \sum_{t}\beta_{2}^{T - t}G_{t}G_{t}^{\top}$  exhibits fast spectral decay (and analogously for  $R_{t}$ ). Luckily, thanks to the recursion  $L_{t + 1}(\beta_2) = \beta_2L_t + G_{t + 1}G_{t + 1}^{\top}$ , the FD sketch may easily be adopted for this setting.
Observation 6. Given a stream  $g_t$  of vectors for  $t \in [T]$ , sketch size  $\ell$ , updates  $(\rho_t^{(\beta_2)}, \bar{G}_t^{(\beta_2)}) = FD\text{-update}(\beta_2 \bar{G}_{t-1}^{(\beta_2)}, g_t g_t^\top)$ , and  $G_T^{(\beta_2)} = \sum_{t=1}^T \beta_2^{T-t} g_t g_t^\top$ , we have
$$
\left\| \bar {G} _ {T} ^ {(\beta_ {2})} - G _ {T} ^ {(\beta_ {2})} \right\| \leq \rho_ {1: T} ^ {(\beta_ {2})} \leq \min  _ {k <   \ell} \frac {\sum_ {i = k + 1} ^ {d} \lambda_ {i} \left(G _ {T} ^ {(\beta_ {2})}\right)}{\ell - k}.
$$
Proof. See Sec.B.4.2
Figure 2: Test metrics are classification error rate for top-1 accuracy for ImageNet v2, word error rate for Librispeech, and one minus average precision for OGBG-MolPCBA. We plot the mean of 5 random seeds, with 1.96 times the standard error as error bars. For readers familiar with ImageNet v1, final validation accuracy for Shampoo was  $77.69\%$ $(0.03\%)$ , S-Shampoo having  $77.18\%$ $(0.04\%)$ , and Adam having  $76.76\%$ $(0.03\%)$ , but we emphasize that due to tuning, the test set performance pictured above should be of primary concern.
We investigate how much of Shampoo's quality our low-memory approach can recover (Sec. 5.1) and whether the factored covariance exhibits spectral decay amenable to sketching (Sec. 5.2). We relegate convex examples to Appendix A due to space constraints, which check that Sketchy compares favorably to related work in a setting directly related to the theory.
5.1 Deep Neural Networks
We evaluate the effectiveness of S-Shampoo as a practical second-order algorithm for training networks, including
- ResNet-50 [3] for ImageNet image classification task of ImageNet [33] with random cropping and flipping augmentations.
- A 16-layer Conformer model [34] for the audio transcription task, Librispeech [35].
- A GNN with 5 message-passing steps [36] on ogbg-molpcba [37], which classifies structural properties of graphically encoded molecule inputs.
Our FD variant of Shampoo introduces only one new hyperparameter, the rank  $\ell$ , which we do not tune, but set to  $\ell = 256$ , which translates to  $4 \times$  memory savings for Shampoo blocks of size 1024 for the accumulators. Shampoo, Adam, and the underlying architectures introduce their own hyperparameters. S-Shampoo inherits those of Shampoo. We tune only common parameters between the three optimizers with the same budgets, selecting based on validation set accuracy. In the ImageNet case, we evaluate final test set performance using ImageNet v2 [38], as the ImageNet test set is unavailable. Additinoal training information is available in Appendix C
As Fig.2 demonstrates, the second-order information leveraged by Shampoo results in improvements over Adam, a first-order method. Our method performs at least as well as Adam in all cases, despite using asymptotically less memory to represent covariance (as Fig.1 shows, Adam uses  $O(mn)$  for a rectangular weight matrix's diagonal accumulators, whereas S-Shampoo uses  $O(mk + nk)$ ). In the GNN case (OBGB-MolPCBA), Shampoo does not perform as well; its training curves indicate overfitting, but we note that S-Shampoo was less susceptible to overfit like Adam.
5.2 Spectral Analysis
To explain Sketchy's strong performance in Sec. 5.1 we inspect the exponential moving average of Kronecker-factored gradient covariance for fast spectral decay. We find that this is indeed the case in practice, so Sketchy's low-rank plus diagonal covariance is representative of true training statistics.
For all our architectures, we tune Shampoo and extract the intermediate gradient covariances over the course of training. To make our curves comparable across architectures, we fix the parameter for
Figure 3: Two measures of spectral decay. As described in Sec. 5.2 we demonstrate spectral decay of  $L_{t}$  and  $R_{t}$  covariance factors of shape  $1024 \times 1024$ . On the left, we take the factors  $C$  from the first layer of each network and plot the proportion of spectral mass captured by its top 256 eigenvalues throughout training, i.e.,  $\sum_{i=1}^{256} \lambda_{i}(C) / \sum_{i=1}^{1024} \lambda_{i}(C)$ . On the right, we plot a continuous measure of eigenvalue concentration, the intrinsic dimension  $\operatorname{tr} C / \lambda_{\max}(C)$ , typically  $10 \times$  smaller than nominal dimension. We plot the average across all covariances for each network's weight's dimensions, with the shaded regions capturing the interquartile range.
the second moment,  $\beta_{2} = 0.999$  for these runs. Furthermore, ResNet-50 has a few parameters with dimension 2048, but the largest dimension for any parameter from the other two architectures is 1024, so we use the Blocked Shampoo variant discussed in Sec.3.4 with block size 1024. In other words, weights containing a dimension 2048 are split into two. We tune other Shampoo parameters for each architecture, and plot statistics of Kronecker factors  $L_{t} = \sum_{i}\beta_{2}^{t - i}G_{t}G_{t}^{\top}$  and  $R_{t} = \sum_{i}\beta_{2}^{t - i}G_{t}^{\top}G_{t}$ .
In Fig. 3 we plot the intrinsic dimension of Kronecker covariance factors over training for our three settings. The intrinsic dimension determines the rate at which empirical covariance estimates concentrate to their expectation, rather than a random vector's actual dimension, up to logarithmic factors (Vershynin [39], Remark 5.6.3). Despite actual dimensionality being over 1024, intrinsic dimension across all architectures stays below 105. A conspicuous phase shift  $10\%$  of the way through training may be the result of a change from linear learning rate warmup to a learning rate decay, starting at roughly  $5\%$  of the way into training.
Given  $\beta_{2} = 0.999$ , we emphasize that the behavior in Fig.3 is an emergent property of DL training. Though surely a lower  $\beta_{2}$  would naturally result in lower intrinsic dimension (which can still be taken advantage of by Alg.2 and 3), we would still expect higher intrinsic dimension if covariances were near-isometries. If we observe some large number  $n = 10000$  draws  $x_{i}$  of  $1024\times d$  matrices with iid  $N(0,1)$  entries, then numerical experiments show that the average intrinsic dimension of  $\sum_{i = 0}^{n - 1}\beta_2^i x_ix_i^\top$  is 324.63 (0.52) and 862.13 (0.25) for  $d = 1,64$ , respectively, with parenthesized numbers denoting standard error across 20 trials. Values generated this way are larger than the average intrinsic dimension of roughly 10, 30, 50 observed in Fig.3
Up to spectral error, Alg.  $\boxed{2}$  achieves full-matrix AdaGrad regret despite approximating the smallest part of the spectrum of  $G_{t}^{-1 / 2}$  at each step. Remarkably, these eigenvectors correspond to the most easily discernible signals of the covariance for the stream  $g_{t}$ . This apparent (and fortuitous) coincidence is resolved by considering the covariance of  $\tilde{G}_t^{-1 / 2}g_t$ : whitening the gradient to facilitate optimization best reflects on regret; as a result, approximating top eigenvectors of  $G_{T}$  helps more than the bottom ones.
Our initial implementation focused on correctness rather than physical speed or memory reduction. Engineering optimizers competitive with existing industrial-strength implementations of Adam and Shampoo was out of scope. In implementing FD, we performed updates via the factored SVD of  $[\beta_2^{1/2}B_t;G_t]$  rather than the eigendecomposition depicted in Alg.1 this avoids squaring, which is unavoidable in Shampoo. For speed, Shampoo subsamples gradients for its covariance estimation and updates its inverse matrix roots intermittently, every fixed number of steps. A tuning script provided by Anil et al. included gradients from every step, but updated roots every 10 steps. Since FD does
not separate sampling from its computation of estimated covariance eigendecomposition, we took the more difficult setting for S-Shampoo, only allowing it to simultaneously observe every  $10^{\mathrm{th}}$  gradient and update its covariance inverse roots (see Appendix  $\mathbb{G}$  for a theoretical justification).
Though step-skipping makes Shampoo and S-Shampoo tractable, future work may explore further speedups: since FD only requires the top  $\ell$  eigenvalues, iterative Lanczos-like routines which are accelerator-friendly, such as LOBPCG [40], may allow incremental updates to  $\tilde{G}_t^{-1/2}$  in factored form with only a few matrix multiplies, S-Shampoo may be able to update more frequently than its non-sketched counterpart, further improving quality.
In this work, we address a gap in the OCO literature for low-memory optimization with the novel Alg. 2 and demonstrate its relevance to practical non-convex problems such as neural net training (Sec. 5.1) by leveraging a new observation about gradient covariance (Sec. 5.2).
The growing disparity between compute capability and memory bandwidth [11] underscores the need for further research in this direction. Further, large-batch settings reduce the performance gap between first and Shampoo-based second order methods, since the batch-size independent runtime of the optimizer is amortized per example used for the gradient calculation. Even in performing experiments for this work, we would frequently find that faster accelerators were unavailable, but many previous-generation ones were, encouraging us to leverage data-parallel training. For datasets such as Imagenet, we notice the advantage of second order methods in dealing with large batches even at relatively modest sizes, such as 1024; many works on explore several larger multiples of this [41].
Potential for future work includes numerical methods outlined in the previous section as well optimizing the rank  $\ell$  across the many tensors in a network, as the spread in Fig. highlights the large variance in covariance intrinsic dimension. Furthermore, the inductive biases conferred by the minima which different-rank representations of curvature reach may have problem-dependent generalization implications, a question which we leave for future work. For a comparison of full rank preconditioning's effect versus first-order minima, see Amari et al. [42].
Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antonio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-272, 2020. doi: 10.1038/s41592-019-0686-2.
[56] MLCommons® open engineering consortium. MLCommons Algorithmic Efficiency. https://github.com/mlcommons/alGORITHMic-efficiency, 2023.
[57] Vijay Janapa Reddi, Christine Cheng, David Kanter, Peter Mattson, Guenther Schmuelling, Carole-Jean Wu, Brian Anderson, Maximilien Breughe, Mark Charlebois, William Chou, Ramesh Chukka, Cody Coleman, Sam Davis, Pan Deng, Greg Diamos, Jared Duke, Dave Fick, J. Scott Gardner, Itay Hubara, Sachin Idgunji, Thomas B. Jablin, Jeff Jiao, Tom St. John, Pankaj Kanwar, David Lee, Jeffery Liao, Anton Lokhmotov, Francisco Massa, Peng Meng, Paulius Micikevicius, Colin Osborne, Gennady Pekhimenko, Arun Tejusve Raghunath Rajan, Dilip Sequeira, Ashish Sirasao, Fei Sun, Hanlin Tang, Michael Thomson, Frank Wei, Ephrem Wu, Lingjie Xu, Koichi Yamada, Bing Yu, George Yuan, Aaron Zhong, Peizhao Zhang, and Yuchen Zhou. Mlperf inference benchmark, 2019.
[58] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017.
[59] Boris Dayma and Rohan Anil. Evaluation of Distributed Shampoo: Comparison of optimizers: Distributed Shampoo, Adam & Adafactor. Weights & Biases Report, 2022.
[60] Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning lecture 6a overview of mini-batch gradient descent. Cited on, 14(8):2, 2012.
[61] Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang. Disentangling adaptive gradient methods from learning rates. arXiv preprint arXiv:2002.11803, 2020.
[62] Tsuyoshi Ando. Concavity of certain maps on positive definite matrices and applications to Hadamard products. Linear algebra and its applications, 26:203-241, 1979.
[63] Rajendra Bhatia. Matrix analysis. Springer, 1997.
[64] Kaare Brandt Petersen, Michael Syskind Pedersen, et al. The matrix cookbook. Technical University of Denmark, 7(15):510, 2008.
[65] Roger W Brockett. Finite dimensional linear systems. SIAM, 2015.