Sketchy: Memory-efficient Adaptive Regularization with Frequent Directions
Authors: Vladimir Feinberg
Xinyi Chen
Y. Jennifer Sun
Rohan Anil
Elad Hazan
Venue: NeurIPS 2023 poster (2023)
Keywords: online convex optimization
deep learning
matrix sketching
frequent directions

Abstract: Adaptive regularization methods that exploit more than the diagonal entries exhibit state of the art performance for many tasks, but can be prohibitive in terms of memory and running time. We find the spectra of the Kronecker-factored gradient covariance matrix in deep learning (DL) training tasks are concentrated on a small leading eigenspace that changes throughout training, motivating a low-rank sketching approach. We describe a generic method for reducing memory and compute requirements of maintaining a matrix preconditioner using the Frequent Directions (FD) sketch. While previous approaches have explored applying FD for second-order optimization, we present a novel analysis which allows efficient interpolation between resource requirements and the degradation in regret guarantees with rank $k$: in the online convex optimization (OCO) setting over dimension $d$, we match full-matrix $d^2$ memory regret using only $dk$ memory up to additive error in the bottom $d-k$ eigenvalues of the gradient covariance. Further, we show extensions of our work to Shampoo, resulting in a method competitive in quality with Shampoo and Adam, yet requiring only sub-linear memory for tracking second moments.

Sketchy: Memory-efficient Adaptive Regularization with Frequent Directions
Vladimir Feinberg Google Deep vladf@google.com com
Xinyi Chen
Y. Jennifer Sun Princeton University Google DeepMind
Rohan Anil Google DeepMind
Princeton University
Google DeepMind
Elad Hazan
Princeton University
Google DeepMind
Abstract
Adaptive regularization methods that exploit more than the diagonal entries exhibit state of the art performance for many tasks, but be prohibitive in terms of memory and running time. We find the spectra of the Kronecker factored gradient covariance matrix in deep learning (DL) training tasks are concentrated leading eigenspace that changes throughout training, motivating a low-rank sketch- ing approach. We describe generic method for reducing memory and compute requirements of maintaining matrix preconditioner using the Frequent Directions (FD) sketch. While previous approaches have explored applying FD for second- order optimization, we present a novel analysis which allows efficient interpolation between resource requirements and the degradation in regret guarantees with rank k: in the online convex optimization (OCO) setting over dimension d, we match mall-matrix d2 memory regret using only dk memory up additive error the bot- tom d k eigenvalues of the gradient covariance. Further, we show extensions of our work Shampoo, resulting method competitive in quality with Shampoo and Adam, yet requiring only sub-linear memory for tracking second moments.
small
1 Introduction
DL optimization commonly relies on adaptive gradient methods, namely the Adam optimizer [T. It differs from stochastic gradient descent in that the learning rate structured diagonal matrix built from previous gradients rather than a scalar. In full matrix AdaGrad [2], the inverse matrix square root of the sum of outer products of previous gradients the learning rate.
Full matrix preconditioning is impractical for modern deep learning architectures: for instance, the ResNet-50 architecture (3) has over 23 million parameters, requiring more than 2 petabytes to represent its gradient covariance. Thus, diagonal preconditioning methods remain popular. However, previous work has demonstrated state-of-the-an results in some settings, such as large-batch data parallel training, for nondiagonal forms of preconditioning [4] [5 6
[7] 8 [9] In particular, Shampoo [ introduces factorization of full matrix preconditioning method with adoption in large-scale industrial applications such as training Google's 's ads dlck-throughh
model [10]. Furthermore, as hardware evolves, memory efficiency becomes an increasing concern, "logic improves much faster than wires and SRAM, so logic is relatively free" [11]: from TPUv2 to TPUv3, per-chip bfloa1 operations per second improved 2.67x x but memory bandwidth only improved 1. x.29x. GPUs exhibit a similar pattern for compute and memory increase, at 5x and 2.2x, for V100 to A100 [1].
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Investigation into the ponnecker-factoree gradient covariance matrix reveals a concentrated, but changing. spectrum (Fig.B) suggesting the majority of the spectral mass can be represented by a low- rank matrix, albeit rotating over time. The Frequent Directions (FD) sketch provides a mechanism to track the top eigenvectors without materializing the full covariance matrix, as proposed in [13]. Is large portion of the spectral mass sufficient to retain the performance of adaptive regularization in theory and practice? In this work, we investigate this hypothesis.
In the setting of online convex optimization, by applying a dynamic diagonal regularization to the FD sketch, we can recover full-matrix AdaGrad regret up to additive spectral terms under a memory constraint, providing novel guarantee without curvature assumptions (Sec. 4.1. Rigorously composing our approach with Shampoo (Sec. 4.2) unlocks a second- order algorithm which requires sub-linear memory for its accumulators.
By modifying FD for exponential moving averages (Sec. 4.3p, we demonstrate practical algorithm competitive with at-lcast-linear memory Shampoo and Adam in three modern DL settings (Sec. 5.1]. While previous work 8 14 shows rank- I preconditioners are effective for trading off quality for memory, these results demonstrate Pareto improvement by using higher-rank approximations.
We explain the competitive performance of Sketchy with observations of fast spectral decay in the moving average of -rotecker-factoree gradient covariance in DL settings (Sec. (.2).
2 Setting and Definitions
Regret and Optimization. The optimization problem of training a deep neural network has non- convex objective loss function f. Since finding the global optimum computationally intractable in general, theoretical guarantees focus on convergence to an e-approximate first-order optimum: a point x such that [z)] S E. A smooth non-conved problem can be reduced to solving a series of offline convex problems (6). The convex sub- uub-problems have form ft (x) = f(z) + cllx 2z11?, where e a constant and x, an iterate in the optimization process. Using online-to-baie
conversion, we can translate the regret bound of an online convex optimization (OCO) [15] algorithm to convergence guarantees for offline optimization. For more details of this reduction, see Appendix [..] Therefore, non-convex optimization guarantees can be obtained from regret bounds, and we focus on the latter in this paper.
In this setting, an OCO algorithm chooses a point E K. iteratively, where K. C Rd is convex decision set (take K. IRd if unconstrained). After the decision is made, the adversary reveals a convex loss function ft. to which the algorithm suffers cost f((z)). Upon receiving the cost, the algorithm updates its decision for the next iteration. The regret for an online algorithm given by
Regrety Σ f(.))
min Σ f()). zEK
Sketching and the Frequent Directions Method. Given a stream of vectors 9. E RRd, E [T), we utilize the FD sketch [13 given in Alg- which maintains a low-rank approximation of the true running covariance G2 = E.st 9e9. each time t, it maintains a matrix B, of size d x whose last column is 0 and whose square is the sketch B, Gr. After seeing 9: from the stream, we update the previous matrix using Alg. [ which updates B of size d l whose last column remains O; take Bo 0. B++1 1 is obtained by decomposing the sum of G, and the newly observed matrix, keeping only the top eigendirections, and reducing the eigenvalues uniformly by ε-th eigenvalue. At be the removed eigenvalue from the covariance update in Alg.
every iteration t, denote Pe : X
For convenience, let P1:t i-th leading eigenvalue by (x)).. Let || ||r denote the Frobenius norm of matrix.
def
D 11=1 Ps be the cumulative escaped mass. For a matrix X, we denote its
The fundamental property of FD that applying Alg. [over stream of vectors gr, with Bg = 0, the (t)
sum of escaped mass Pt = x can be bounded by the bottom eigenvalues of Gr, formally given by the following lemma:

---

Algorithm Frequent Directions Update (FD-update)
Require:
Invariant that last column of B:-1 is O.
Ensure: The last column of B, is O.
1:
Input: Previous state G = B E Rdxd
2: Input: New symmetric PSD matrix M2 E Raxd
3: Eigendecompose U, diag AT = G + Mj where ( contains descending eigenvalues.
4: Define U, the matrix whose columns the first l columns of U2, and \ioo) be its eigenvalues.
5: Update B, = U2 diag (c) (1u x(
1/2
output 1 B,B.
Lemma (Liberty [16). The cumulative escaped mass Pir can be upper bounded as
E M(GTT)
def
Pı.r S min
Xei
k=
k-k
1=8
Proof. See Sec. B.2
Related Work
3.1 Spectral Analysis of DL Training
Denote the loss function of the i-th example for weights x as fi().. The spectrum of the Hessian
matrix Σ. V2 f. has been the subject of intensive investigation in DL [7 19 20 and its properties
have been used to devise training methods 21.
inspect the covariance matrix, E, (V.)(V In small models, where its computation feasible, these works identify fast spectral decay.
Recent papers 22 23 24
Agarwal et al. 6 take advantage of this observation by using low-rank approximation of the whole covariance matrix, based on a limited history of the gradients, -f approach still requires copies of the model gradients in memory, where typically r should scale with 32 with 32 the exponential moving average for second order statistics (the authors set r = 200). Fundamentally, approximating the whole covariance matrix constrains Agarwal et al. 6 application to small models.
This
In our work, validate the decay hypothesis holds across the per-layer factored covariance matrices in several modern neural networks. For a layer's gradient matrix G; the i-th example and second moment decay term 32, our work inspects spectral decay for L. = Σ. 8 G and Rg = E Gş; the spectral structure for these outer products is not well-documented. Furthermore, described in Sec. 3.4 approximating the factored covariance L, & R, requires less memory than the full covariance and explains why our method can scale to large modern architectures whereas Agarwal al. 66 cannot.
3.2 Sublinear Memory Methods
Extreme Tensoring 77), AdaFactor [4]. and SM3 [ are methods that require sublinear memory relative to the number of parameters, at the other end of the memory-quality tradeoff beyond methods that rely on the diagonal of the gradient covariance such as Adam. Owing to different structural assumptions on the set of feasible preconditioners, comparison with these methods is out of scope. However, these methods may compose with our approach. One may apply Extreme Tensoring first, then sketch the resulting reshaped tensor covariances with method to further reduce memory consumption. Similarly, an SM3-like approach which reduces the indices for each dimension to be preconditioned can be applied before Sketchy is applied remaining indices.
Crucially, Adam, which uses linear memory for second moment representations, compares favorably in terms of quality to all of these sublincar-men memory methods. By increasing rank in factored

---

Memory Table I: Memory-efficient adaptive gradient methods, in the OCO setting with dimension d (Sec. We describe the worst-case regret bounds without exp-concavity assumptions, asymptotically. hiding logarithmic factors, treating the decision set diameter constant, and assume optimally- tuned hyper-
parameters. e refers to the controllable preconditioner rank. Note tr G /mmnnee E, V the optimal preconditioner's regret among the class of positive semi-definite, unit-trace matri- ces, H. and GT is the of gradient outer products. We let eigenvalues X; ((Gr) with Am-
Auj >
Reference
Regret (general convex)
Full Matrix AdaGrad
tr G
d²
Ada-LR [25
tr Gr 1/ + X 1/2 d3/4
d²
T
(+1
Ada-FD 26
Ω (T I
de
SON 227
Td
d?
FD-SON 27
v VEARNT
de
This paper
tr tr() 1/2 + d(d E)7ed
de
covariance representation, Sketchy is competitive with Adam, despite sublinear memory for second moment representation. Thus, for simplicity, we compare to only Adam, which dominates the alternative sublinear approaches terms of quality.
3.3
Sketching-based Approaches
Several works have explored sketching like approximations to the gradient covariance matrix, but none provide an adaptive bound exploiting fast spectral decay in gradient covariance without addi- tional assumptions (Tbl. D. In this section, we consider the OCO setting over dimension d (Sec. 2
Although it does not reduce memory usage, it relies on random projections to lower dimension E < to reduce inverse matrix computation costs. An alternative without formal guarantees, RadaGrad, reduces memory consumption to (dé); however, as with all Joonson--Lindnstreie projection methods, it suffers probabilistic failure rate scaling as O(E-!) (in comparison, our method inherits FD's determinism).
Random projection (Ada-LR) is most spiritually similar our work 25
Frequent Directions (FD) B 16,, provides an alternative matrix sketching approach from the data streaming literature. As adaptive sketch, it dominates random projection in terms of matrix recovery, and lower bounds show its memory usage is optimal in the sense that any equally-accurate approximation
daverssrialll-choseee true covariance Gr in operator norm constructed from the corresponding gradients must use (dE) bits.
In the context of exp-concave cost functions, Luo et al. 27 provide an FD sketched version of Online Newton Step (ONS), FD-SON, In this setting, their approach nearly recovers classical ONS regret, up to logarithmic error in E 1 Xe(GT) and additive error in Σi= A=(G)). However, without the
exp-concave assumption, FD-SON falls back to gradient-dessent-like default regret of O (DeaVT)
which can be Q(T) without spectral decay. In the context of linear bandits, Chen et al. (8 uses FD for memory reduction. The resulting algorithm, SCFD, is similar to Alg. 2 but SCFD lacks projection step and does not handle general domains K. We emphasize that our contribution is in the novelty of regret analysis for general OCO settings, in contrast linear bandits.
The main prior work exploring FD for general online convex settings, Wan and Zhang 26 extends the FD-SON approach by adding fixed diagonal perturbation 5T to an FD-based preconditioner, in Ada-FD. However, this approach does not achieve VT regret even in non-adversarial setting with stochastic linear cost functions (Ohservation 2). where learning rate and & are tuned. Dynamically
changing diagonal regularization is essential for worst-case O (VT) performance.
Observation 2. Suppose receive linear cost functions [x(x) (a,, g.). where g, E Rdd is random vector drawn iid from any distribution over r S. d orthonormal vectors W. For any sketch size & S. the bound on the expected regret of Ada-FD FD is 2/13/11)

---

AdafactorSM3
Sketchy (Ours)
Diag.AdaGres AdaGrad
Shampoo
GGT
k
L R
m
mk nk
mn
+
mnr
cublinear
Figure I: Asymptotic memory consumption for representing gradient covariance in adaptive reg- ularization approaches for a single matrix parameter of size n x m. Here, r refers to the GGT history buffer size and k to the approximation rank of FD (both typically set to hundreds). Past sketching approaches like Ada-FD and Radagrad take memory similar to GGT, with r being sketch size; these are all asymptotically superlinear. This figure demonstrates optimizer memory usage in the theoretical OCO setting; in practice for deep learning workloads there are additive O(mn) factors for momentum, the parameters themselves, and grafting parameters for Shampoo and Sketchy.
Proof. See Sec. B.3
Wan and Zhang 26 remark Ada-FD has VT regret when Gr is low rank with rank below k, so d- of its eigenvalues precisely zero. However, this setting does not require any sketching in the first place. By tracking the column space of observed gradients (e.g., with reduced QR decomposition, rank-l-updated every step), the full matrix AdaGrad algorithm can be perfectly recovered without using more than ((dk) memory.
In concrete convex examples, Sketchy compares favorably to these approaches (Appendix
3.4 Shampoo
Perhaps our most compelling application reducing the memory of Shampoo [ 9 10 Shampoo
an adaptive preconditioning method that takes into account the structure of the parameter space, and thus is more efficient than full matrix AdaGrad. For example, if the parameter is a weight matrix W of size m x n, AdaGrad treats the matrix-shaped parameters as vector of size mn, and the preconditioner has size n Shampoo instead has left and right preconditioners L, R of size x n and m X m, respectively, with the preconditioned update L 1.1//WR R 1/4 Write ved (W) as the vectorized weights, then it is equivalent to (L R)vee (W) vee (LW R), where denotes the Kronecker product. Figure [ fillustrates the updates of AdaGrad and Shampoo, where AdaGrad update uses the entire matrix instead of the diagonal, which is shown in the figure. In other words, Shampoo uses camneckrr-fcctoee preconditioner, and the factorization preserves the matrix structure of the parameters. Since in DL optimization parameters often have matrix structure, Shampoo has strong empirical performance, and has improved upon state-of-the-an results in large-scale tasks such as language modeling with BERT-Large and image classification on ImageNet in 9
Figure eelaborates why the composition of FD and Shampoo is essential to avoid memory consump- tion asymptotically greater than parameter count for approximate full matrix regularization.
However, Shampoo memory costs may still be prohibitive for rectangular weight matrices. In BERT- I Large [29] most parameters are in the feed-forward network layers, which consist of 4096 x 1024 dense kernels; other transformers follow similar narrow-to-wide patterns. For large models, occupying 4x memory for the left preconditioner can frequently result in OOM in memory-constraine settings; this was in fact one of the practical motivations for our proposed approach.
Anil et al. 9 introduces two workarounds for the problem of rectangular matrices based on limiting covariance modelling. Furthermore, both approximations can be applied to our method, so we do not compare against them. First, the authors propose Blocked Shampoo, which views each weight matrix W of shape m X n as mm/b² blocks size b x for some block size b < min n) (in the limit b = 1, this recovers diagonal AdaGrad). This approach is dependent on the ordering of neurons in hidden layers. Another approximation relies on only one-sided covariance upper bounds, I,, 60 T or I & Rg Note, however, that the one-sided approximation doesn't 't help with vector parameters, such as
5

---

those that appear for the bias terms in dense layers or layer norms 33 For 3D weights, such as those which appear in homogeneous Mixtures of Experts BT. blocking increases memory consumption. These approaches do not take into account the fast decay of the preconditioner's spectrum, which is the focus of our work.
Algorithms and Main Theorems
In this section, we introduce the adaptation of Frequent Directions (FD) to AdaGrad (Sec. 4.1 and Shampoo (Sec. 4.2]. the corresponding algorithms and regret guarantees. Additionally, Sec. 4.3
modify FD to support exponential moving averages.
The main technical novelty in incorporating Alg. to AdaGrad (Alg. 2 and Shampoo (Alg. 3 the construction of preconditioning matrices with FD-sketched matrices compensated by the cumulative escaped masses. The insight of such construction lies in the observation that while the FD sketch lower bounds the full preconditioning matrix, the FD sketch compensated with the cumulative escaped masses upper bounds the full preconditioning matrix, demonstrated Lemma I0 for AdaGrad and Lemma 14 for Shampoo. The regret guarantee for AdaGrad ([2)) and Shampoo (55) directly depends on the trace of the preconditioning matrices, therefore obtaining upper and lower bounds on the preconditioning matrices allows explicit additive dependence on the cumulative escaped mass. We expect this approach to be reusable for alternative approximation schemes.
4.1 FD for AdaGrad
Our main algorithm in this section is Alg. 2 run with FD (Alg. 1 as the sketching method. G. in
Alg. denotes the Moore-Penrose pseudoinverse of the matrix G1/2 Our main algorithm. Sketchy
AdaGrad, in this section exploits the FD approach outlined in Alg. D as the sketching method in
AdaGrad. In particular, every time step. we pass the newly received subgradient Ir into Alg T which updates and maintains low-rank sketch G2 of the AdaGrad preconditioning matrix G,. We keep track of the cumulative escaped mass Pi:s, which we add back the low-rank sketch create
the Sketchy preconditioner G,, with which we perform the regular AdaGrad descent and projection.
Algorithm 2 Sketchy AdaGrad (S Ad-Arddra
1: Input: constraint set K., step size time horizon T.
2: Initialize X, Go Go = 0.
3: for = I,..., do
4: Play receive E (xt),, suffer cost fe(e)).
5: Sketch (pr, G.) = FD-updateeGt-........ 9t
6: Update G. G. + PIrT, M++1 n/, gt. and 2441 argmin
1/2
l| x/1 2//22//2-
aEK:
7:
end for
Theorem 3. Define Ω4 ming - Ei=kt >r(GT), then with 11
'2' Alg.
guarantees
the following additive regret bound:
Dt Regret, (s AdAdGGGad) S D (2 + +dy
2
where D is the diameter the constraint set K. K. is bounded and max || x²122 otherwise.
Proof. See Sec. B.4.1
Notably in Thm.
Regret, = 0 and the lower eigenvalue dependence Ot additive.
Corollary 4. We improve Theorem 3 slightly to
+ d(d f)see
Regret, (S-AdaGradd s D
12 G2/2
2

---

Proof. See Sec. Sec.B..42 B.4.2
The regret bound above holds under the optimal tuning of the learning rate, which depends problem quantities that can be unknown a priori. It is possible to design a parameter-free free variant Alg by
using the norm || = (x T (Ga + T)1/2x))1/2 1/2 in the projection step of Alg. as seen in 322
4.2 FD for Shampoo
In this section, we adapt FD-update to Shampoo [5]. For simplicity, we optimize over R mx in Alg. [3] projection may be handled as in Alg. Similar to Sketchy AdaGrad, Sketchy Shampoo uses the FD approach outlined in Alg. to sketch the left and right preconditioning matrices for Shampoo. In particular, we maintain two parallel sketching streams using Alg. [ to produce sketches L,, R for the left and right preconditioning matrices. We keep track of the cumulative escaped masses pi: and pf: from sketching the left and right preconditioning matrices, respectively, and compensate the cumulative escaped mass to create the left and right Sketchy preconditioning matrices Li, R.-
Algorithm 3 Sketchy Shampoo (S-Shampoo)
I: Input: step size time horizon T.
2: Initialize Xo 0mxn. Lo eIm, Ro = ETn, T.o = Om, Ro = On-
3: for = 1, ....., T do
4: Play X,, suffer (x,), receive G, E @f((X.).
5: Sketch (pt L.) = FD- FD-update(L1-1) G,G, (ol Ra) = FD-update R4-1 Gi).
6: Update L, Le + platm. RS Rg + plaly and X141 X₁ nLi G,R, 1114
7: end for
Denote LT del Et, G,O.T eel and RT del Et G. el.
Theorem 5. Suppose Tt....Gr have rank at most Then Alg. the following regret bound:
run with D/ guarantees
( (S-Shampoo) & 2rD (tr(L//*) + mL///) (IT(R4/1) + 12,
where D maxxelr [| 1-=1= and Lr,ee SRRe are analogous bounds for plava R from Lem,
Proof. See Sec. B.5.1
Bounds may be improved analogous to Cor, for Alg. 3 but we omit the similar statement due to space.
4.3 Exponentially Weighted FD
This section discusses the modification of Alg. to support exponential moving averages. Early in algorithm development, we noticed that attempting to approximate the unweighted sum of factored gradient covariances Σ. G.G. and E, G, G, with FD tended to estimate of covariance that was roughly 0, creating numerical instabilities. Note that FD guarantee (Lem. [ still holds-but the error term P1.7 becomes greater than || resulting in vacuous bound due to lack of spectral decay.
Indeed, Fig. motivating this work only confirmed that the exponential moving average Le(3)) E. - G,G. exhibits fast spectral decay (and analogously for R2). Luckily, thanks to the recursion L441 (32) B2L, G+41GTT11 the FD sketch may easily be adopted for this setting.
Observation Given a stream gt of vectors for E [T]. sketch size updates (82) G((t)) FD weduuddaeeBattt ((9TT) ). and (3))
ΣT
- we have
Ei d G
c c((52
(B2)
<
PP..T
min
k

---

Imageteet (Eror Rate)
Librapeach (WITR)
OGBG: MCIPAIA (1:4P)
0.44
0.09
0.95
0.42
0.00
0.50
0.05
0.35
0.34
0.33
90
0.05
30
0.70
10 20
Figure 2: Test metrics classification error rate for top-1 accuracy for ImageNet v2, word error rate for Librispeech, and one minus average precision for OGBG-MoIPCBA. We plot the mean of random seeds, with 1.96 times the standard error as 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 emphasize that due tuning, the test set performance pictured above should be of primary concern.
5 Experiments
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. (.22. We relegate convex examples to Appendix due to space constraints, which check that Sketchy compares favorably related work in 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 for ImageNet image classification task of ImageNet 33 with random crop-
ping and flipping augmentations.
A 16-layer Conformer model 34 for the audio transcription task, Librispeech 135
A GNN with 5 message-passing steps [36] on ogbg-molpcba 37, which classifies struc- tural properties of graphically encoded molecule inputs.
Our FD variant of Shampoo introduces only one new hyperparameter, the rank E, which we do not tune, but set to l 256, which translates 4x 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 asympotically less memory to represent covariance (as Fig. [ shows, Adam uses ((m)) for rectangular weight matrix's s diagonal accumulators. whereas S-Shampoo uses O(mk + nk)). ). In the GNN case overfitting, but we note that S S-Shampoo was less susceptible to overfit like Adam.
Shampoo does not perform well; training curves indicate
5.2 Spectral Analysis
To explain Sketchy's strong performance in Sec. 5.1
inspect the exponential moving average of Kronecker- facttrre gradient covariance for fast spectral decay. We find that this is indeed the case in practice, Sketchy's low-rank plus diagonal covariance is representative of true training statistics.
For our architectures, we tune Shampoo and extract the intermediate gradient covariances over the course of training. To make curves comparable across architectures, fix the parameter for

---

30
buiting COmpReTIn
we demonstrate spectral decay of Th and Rg covariance factors of shape 1024 1024. On the left, take the factors C from the first layer of each network and plot the proportion of spectral mass captured by its top 256 eigenvalues 1024 )C().. On the right, we plot a continuous measure of
Figure 3: Two measures of spectral decay. As described in Sec. 5.2
throughout training. i.e., E 256 X/CC)/ E eigenvalue concentration, the intrinsic dimension tr C/ (C) typically 10x smaller than nominal dimension. We plot the average across all covariances for each network's weight's 's dimensions, with the shaded regions capturing the interquartile range.
the second moment, 32 0.999 for these runs. Furthermore, ResNet-50 has 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 = Σ, GGG and R. = E, - G.-
In Fig. 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 random vector's '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 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 32 0.999, we emphasize that the behavior in Fig. an emergent property of DL training. Though surely a lower B3 would naturally result in lower intrinsic dimension (which can still be taken advantage by Alg- 2 and 3). we would still expect higher intrinsic dimension covariances were near-isometries. If we observe some large number = 10000 draws of 1024 x d matrices with iid (0, 1) entries, then numerical experiments show that the average intrinsic dimension of m- 3i is 324.63 (0.52) and 862.13 (0.25) for = 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
Discussion
Up to spectral error, Alg. achieves full-matrix AdaGrad regret despite approximating the smallest part of the spectrum of G, 1/2 at each step. Remarkably, these eigenvectors correspond to the most easily discernible signals of the covariance for the stream gt. This apparent (and fortuitous) coincidence is resolved by considering the covariance of 1/2 g.: g:: whitening the gradient to facilitate optimization best reflects regret; result, approximating top eigenvectors of GT helps more than the bottom ones.
Our initial implementation focused on correctness rather than physical speed or memory reduction. Engineering optimizers competitive with existing andustrial-strengtn implementations of Adam and Shampoo was out of scope. In implementing FD, performed updates via the factored SVD of /8 B,; G.] rather than the eigendecomposition depicted in Alg. this avoids squaring, which unavoidable in Shampoo. For speed, Shampoo subsamples gradients for its covariance estimation and updates inverse matrix roots intermittently, every fixed number of steps. A tuning script provided by Anil et al. [9] included gradients from every step, but updated roots every 10 steps. Since FD does

---

not separate sampling from its computation of estimated covariance cigendecomposition, we took the more difficult setting for S-Shampo,, only allowing simultaneously observe every 10th gradient and update its covariance inverse roots (see Appendix G for theoretical justification).
Though step-skipping makes Shampoo and S-Shampoo tractable, future work may explore further speedups: since FD only requires the top e eigenvalues, iterative Lanczos-like routines which are soeelcrator-friendly. such as LOBPCG [40], may allow incremental updates G in factored form with only few matrix multiplies, S-Shampoo may be able to update more frequently than its non-sketched counterpart, further improving quality.
7 Conclusion
In this work, we address a gap in the OCO literature for low-memory optimization with the novel Alg. 2 Zladd demonstrate its relevance practical non-convex problems such as neural training (Sec. 5.1 DIIb by leveraging new observation about gradient covariance (Sec. 5.2).
underscores the The growing disparity between compute capability and memory bandwidth need for further research 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 -errouuggenerttion ones were, encouraging to leverage data- data-parallel training. For datasets such 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 [1
Potential for future work includes numerical methods outlined in the previous section as well optimizing the rank l across the many tensors in network, the spread in Fig. highllghhs 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 's effect versus first-order minima, see Amari et 44
8 Acknowledgments
Elad Hazan acknowledges funding from the Office of Naval Research grant N000142312156, the NSF award 2134040, and Open Philanthropy.
10

---

References
[1] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR (Poster), 2015.
[2] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
[3] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and paltern recognition, pages 770-778, 2016.
[4] James Martens and Roger Grosse. Optimizing neural networks with cronecken-factored approx- imate curvature. In International conference machine learning, pages 2408 2417. PMLR, 2015.
[5] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. International Conference on Machine Learning, pages 1842-1850. PMLR, 2018.
[6] Naman Agarwal, Brian Bullins, Xinyi Chen, Elad Hazan, Karan Singh, Cyril Zhang, and Yi Zhang. Efficient full-matris adaptive regularization. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning. volume 97 of Proceedings of Machine Learning Research, pages 102-110. PMLR, 09-15 Jun 2019.
[7] Xinyi Chen, Naman Agarwal, Elad Hazan, Cyril Zhang, and Yi Zhang. Extreme tensoring for low-memory preconditioning. In International Conference on Learning Representations, 2019.
[8] Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimization. Advances in Neural Information Processing Systems, 32, 2019.
[9] Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002. 19 2020.
[10] Rohan Anil, Sandra Gadanho, Da Huang, Nijith Jacob, Zhuoshu Li, Dong Lin, Todd Phillips, Cristina Pop, Kevin Regan, Gil Shamir, Rakesh Shivanna, and Qiqi Yan. On the factory floor: MI engineering for industrial-scale ads recommendation models, 2022.
[11] Norman P Jouppi, Doe Hyun Yoon, Matthew Ashcraft, Mark Gottscho, Thomas B Jablin, George Kurian, James Laudon, Sheng Li, Peter Ma, Xiaoyu Ma, et al. Ten lessons from three generations shaped google's 's tpuv4i: Industrial product. In 202/ ACM/IEEE 48th Annual International Symposium on Computer Architecture (ISCA), pages 1-14. IEEE, 2021.
[12] William J Dally, Stephen W Keckler, and David B Kirk. Evolution of the graphics processing unit (gpu). IEEE Micro, (1(6):22-51, 2021.
[13] Mina Ghashami, Edo Liberty, Jeff M Phillips, and David P Woodruff. Frequent directions: Simple and deterministic matrix sketching. SIAM Journal on Computing, 452-:1172 2016.
[14] Noam Shazeer and Mitchell Stem. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pages 6696-4604. PMLR, 2018.
[15] Elad Hazan et al. Introduction to online convex optimization. Foundations and Trends in Optimization, 2(3-4):157-325, 2016.
[16] Edo Liberty. Even simpler deterministic matrix sketching, arXiv preprint arXiv:2202.0/ 2022.
[17] Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues of the hessian in deep learning: Singularity and beyond. arXiv preprint 16X 2016.
[18] Levent Sagun, Utku Evei, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over-parametriveed neural networks. arXiv preprint arXiv: 04454, 2017.
[19] Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net opti- mization via hessian eigenvalue density. In International Conference on Machine Learning, pages 2232-2241 PMLR, 2019.

---

[20] Adepu Ravi Sankar, Yash Khasbage. Rahul Vigneswaran, and Vineeth N Balasubramanian. deeper look at the hessian eigenspectrum of deep neural networks and its applications to regularization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 9481 -9488, 2021.
[21] Naman Agarwal, Zeyuan Allen-Zhu, Brian Bullins, Elad Hazan, and Tengyu Ma. Finding approximate local minima faster than gradient descent. In Proceedings of the 49th Annual ACM SIGACT Symposium on Theory of Computing, pages 1195- 1199, 2017.
[22] Guy Gur-Ari, Daniel A Roberts, and Ethan Dyer. Gradient descent happens in tiny subspace. arXiv preprint 1812.044,, 2018.
[23] Craig Bakker, Michael J Henry, and Nathan O Hodas. Understanding and exploiting the low-rank structure of deep networks. 2018.
[24] Zeke Xie, Qian- Yuan Tang, Zheng He, Mingming Sun, and Ping Li. Rethinking the structure of stochastic gradients: Empirical and statistical evidence. arXiv preprint (xXiv:2212.02083 2022.
[25] Gabriel Krummenacher, Brian McWilliams, Yannic Kilcher, Joachim M Buhmann, and Nicolai Meinshausen. Scalable adaptive stochastic optimization using random projections. Advances in Neural Information Processing Systems, 29, 2016.
[26] Yuanyu Wan and Lijun Zhang. Efficient adaptive online learning via frequent directions. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2021.
[27] Haipeng Luo, Alekh Agarwal, Nicolo Cesa-Bianchi, and John Langford. Efficient second order online learning by sketching. Advances in Neural Information Processing Systems, 29, 2016.
[28] Cheng Chen, Luo Luo, Weinan Zhang, Yong Yu, and Yijiang Lian. Efficient and robust high-dimennional linear contextual bandits. In Proceedings of the Twenty-Ninth International Conference on International Joint Conferences on Artificial Intelligence, pages 4259-4265, 2021.
[29] Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Pupers), pages 4171-4186, 2019.
[30] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint 1xX.v:1677.06450. 2016.
[31] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gas mixture-of-expprts layer. arXiv preprint /701 7065388.. 2017.
[32] Ashok Cutkosky. Better full-matrix regret via parameter-free online learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances Neural Information Processing Systems, volume 33, pages 1836-8846 Curran Associates, Inc., 2020.
[33] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision doi: 10. 1007/s1 223--11--0816-7
(ICCV), 115(3):211 252, 2015.
[34] Anmol Gulati, James Qin, Chung- Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang, Yonghui Wu, et al. Conformer: Convolution-augmennte transformer for speech recognition. arXiv preprint arXiv:2005. 08000, 2020.
[35] Vassil Panayotov, Guoguo Chen, Daniel Povey, and Sanjeev Khudanpur. Librispeech: asr corpus based on public domain audio books. In Acoustics, Speech and Signal Processing (ICASSP), 2015 IEEE International Conference on, pages 5206-5210. IEEE, 2015.
[36] Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: 01261, 2018
[37] Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen I Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs.
12

---

In Hugo Larochelle, Mare Aurelio Ranzato. Raia Hadsell, Maria- Maria-Florina Balcan, and Hsuan- Tien Lin, editors, Advances in Neural Information Processing Systems 33: Anmual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6- 12, 2020, virtual, 2020.
[38] Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize to imagenet? In International Conference on Machine Learning, pages 5389-5400, 2019.
[39] Roman Vershynin. High-dimensional probabiiity: An introduction with applications in data science, volume 47. Cambridge university press, 2018.
[40] Andrew V Knyazev. Toward the optimal preconditioned eigensolver: Locally optimal block preconditioned conjugate gradient method. SIAM journal on scientific computing, 23(2):517- 541, 2001.
[41]
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv: 1609.
1609.04836 04836, 2016.
[42] Shun-ichi Amari, Jimmy Ba, Roger Baker Grosse, Xuechen Li, Atsushi Nitanda, Taiji Suzuki, Denny Wu, and Ji Xu. When does preconditioning help or hurt generalization? In International Conference Learning Representations, 2020.
[43] Luo Luo, Cheng Chen, Zhihua Zhang, Wu-Jun Li, and Tong Zhang. Robust frequent directions with application online learning. The Journal of Machine Learning Research, 20(1):1697-
1737, 2019.
[44] Chih-Chung Chang and Chih-Jen Lin. Libsvm: a library for support vector machines. ACM transactions on intelligent systems and technology (TIST), 2(3):1-27, 2011.
[45] Elad Hazan, Tomer Koren, and Kfir Y Levy. Logistic regression: Tight bounds for stochastic and online optimization. In Conference on Learning Theory, pages 197-209. PMLR, 2014.
[46] Adam Kalai and Santosh Vempala. Efficient algorithms for online decision problems. Journal of Computer and System Sciences, 71(3):291-34) 2005.
[47] Koenraad MR Audenaert. A generalisation of mirsky's singular value inequalities. arXiv preprint arXiv: 1,00. 2014.
[48] Justin M. Gilmer, George E. Dahl, and Zachary Nado. init2winit: a jax codebase for initialization. optimization, and tuning research, 2021. httt:tttttttthut http:itttuu.mmoooooooeeeee haoaa//gool///aae ini t2winit
[49] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Handerman-Milne and Qiao Zhang. JAX: composable transformations of Python+NumPP programs, 2018.
[50] Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee Flax: A neural network library and ecosystem for JAX, 2020.
[51] Tensorflow datasets, a collection of ready-to-use datasets. https: http:::///ww.tesoofllooo. www.srennsoooooooo datasets 2023.
org/
[52] Michael L. Waskom. seaborn: statistical data visualization. Journal of Open Source Software, 6 (60):3021, 2021. doi: 10.21105/joss.0000 03002..
[53] D. Hunter. Matplotlib: A 2d graphics environment. Computing in Science & Engineering, (35 2007. doi: 10.1 109/MCSE.2007.55
[54] Charles R. Harris, K. Jarrod Millman, Stéfan van der Walt, Ralf Gommers, Pauli Vir- tanen, David Cournaeeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Rio, Mark Wiebe, Pearu Peterson, Pierre Gérard-Marchant, Kevin Sheppard, Tyler Reddy. Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature, 585(7825):357 362, September 2020. doi: 10. 10.1038//41586 050.4441-6--000-9--2222 264--2.
[55] Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C Carey, ilhan Polat, Yu Feng, Eric W.
13

---

Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E, A. Quintero, Charles Harris, Anne M. Archibald, António 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.
[56]
MI.Common( open engineering consortium. MI.Commons Algorithmic Efficiency. httpp //github. com wtt.o//mooomoo///////mmmmm hal.aporithmi----ffiieeeeeee 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. MIperf 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.11880 2020.
[62] Tsuyoshi Ando. Concavity of certain maps on positive definite matrices and applications to hadamard products. Linear algebra and its upplications, 26:203-241, 1979.
[63] Rajendra Bhatia. Matrix analysis. Springer, 1997.
[64] Kaare Brandt Petersen, Michael Syskind Pedersen, al. The matrix cookbook. Technical University of Denmark, 7(15):510, 2008
[65] Roger W Brockett. Finite dimensional linear systems. SIAM, 2015.
14


Reviews:
[{"review_id": "6p1qIv06XX", "reviewer": "Reviewer_rVs5", "summary": "It is known that dense preconditioning can be highly beneficial for optimization, but this has high memory requirements and is often prohibitive. The classic solution to this problem is the use of a diagonal preconditioner. This work proposes the use of a low rank preconditioner by using a streaming approach to low rank approximation known as Frequent Directions, which maintains small memory requirements throughout training. The authors then give a regret bound on this algorithm, by combining the analysis of Frequent Directions with the regret analysis of AdaGrad, which results in a trade-off between the regret and the memory used for the preconditioner. While the idea of using Frequent Directions for compressing preconditioners in a stream has been used before, the authors obtain a bound that has advantages over prior regret bounds. Experimental results are nice.", "questions": "I still do not understand the difference between the current work and FD-SON, could you clarify?", "limitations": "Some discussion is provided in the Discussions section.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This work gives a highly practical idea for efficient training of neural networks, with a nice theoretical analysis that combines AdaGrad and Frequent Directions as well as nice empirical results.", "weaknesses": "I am confused about the primary contribution of this work, and perhaps the writing needs to be improved to clarify this. The abstract and early introduction seem to mostly focus on the contribution of this work as the idea of applying Frequent Directions during optimization. This is confusing, as the idea itself has been suggested and investigated in prior work (which the authors thoroughly discuss).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to read and review our work. Here are our responses to your questions and concerns, with which we hope to boost your confidence in our work:\n\nWeaknesses:\n\n**rVs5**: I am confused about the primary contribution of this work, and perhaps the writing needs to be improved to clarify this. The abstract and early introduction seem to mostly focus on the contribution of this work as the idea of applying Frequent Directions during optimization. This is confusing, as the idea itself has been suggested and investigated in prior work (which the authors thoroughly discuss). \n\n**Response**: Please see our top-level comment. In short, in L31, we emphasize the focus on the Kronecker-factored gradient covariance, not the full gradient covariance.\n\nQuestions:\n\n1. I still do not understand the difference between the current work and FD-SON, could you clarify?\n\n**Response**: We highlight Table 1, page 3 in our current work, where FD-SON regret simplifies to $\\sqrt{\\ell \\lambda_{\\ell:d} T}$ and our work has $\\sqrt{d(d-\\ell)\\lambda_{\\ell:d}}+tr(G_T)^{1/2}$. In the case that there is no spectral decay, so $\\lambda_{\\ell:d}=\\Theta(T)$, the S-Adagrad analysis degrades gracefully, with still $\\sqrt{T}$ regret. The FD-SON bound becomes linear in the OCO case with no spectral decay.\n\nWe’ve updated the text to describe this case.\n"}, {"review_id": "dqsn6EIOee", "reviewer": "Reviewer_fXUQ", "summary": "This paper proposed a sketched online optimization method. The main idea is using frequent directions (FD) with adaptive regularization to maintain a preconditioner in the training process, which leads to $O(dk)$ memory and additive error associated with the small eigenvalues of the gradient covariance. ", "questions": "1.\tCan you provide a detailed discussion to compare the proposed algorithms with RFD, SCFD and related online optimization algorithms? \n2.\tSCFD is designed for linear contextual bandits and the analysis shows its positive definite monotonicity is important to establish the regret bound of this problem. In the scenario of convex online optimization, is there any theoretical result to show using SCFD is better than FD and RFD? \n", "limitations": "No potential negative societal impact. The detailed discussion for related work is necessary. Please see above for details.", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The design of the algorithm is well-motived and rigorous theoretical analysis is provided.\n2. The experiments on deep nerul network shows the proposed algorithm improved the recent proposed baseline methods Shampoo.\n", "weaknesses": "My main consideration is the novelty of this paper. The relationship between proposed Sketchy AdaGrad (S-AdaGrad) and existing work [27, 43] is unclear. Here are the detailed comments:\n\n1.\tTo my understanding, the main idea of S-AdaGrad (Algorithm 2) is using adaptive regularized frequent directions to estimate the gradient covariance. However, this idea is not novel. Luo et al. [43] proposed robust frequent directions (RFD) and applied it to online optimization. The main idea of RFD is also introducing an adaptive regularization, which is very similar to FD-update (Algorithm 1) in this paper. RFD uses $(\\rho_{1:t} I)/2$ as the regularization term (rather than $\\rho_{1:t} I$ in this paper). It is shown that this strategy has tighter error bound and is more well-conditioned than standard FD, while the results of Observation 6 looks no better than FD. \n2.\tThe FD-update (Algorithm 1) looks the same as spectral compensation frequent directions (SCFD) [27]. Since sketched convex online optimization has been studied by Wan and Zhang [24], Luo et al. [25] and Luo et al. [43], the design of S-AdaGrad seems only applying a different sketching method SCFD to convex online optimization.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time and effort to review our paper. Here are our responses to your concerns:\n\nWeaknesses:\n\nRFD has the same regret bound (though different constants) as FD-SON, which as per Table 1, page 3 can be linear in T when there is no spectral decay or curvature (see our response to rVs5). More importantly, RFD gives up the approximate isometry between its preconditioner approximation and the true gradient covariance (see Remark 9 in the Appendix). This is precisely because of their choice to tighten the Frobenius norm by giving up the property that their preconditioner dominates the true covariance. However, this isometry is required in our analysis which demonstrates additive spectral error. The RFD paper essentially re-analyzes the FD-SON proof, and as such inherits its linear-regret bound in the no-spectral-decay case. We updated our related work to emphasize this.\n\n\nQuestions:\n\n1. Please see our top-level official comment. To summarize, all of the approaches you have mentioned have two fundamental deficiencies which our work addresses: (a) They require superlinear $O(mnr)$ memory, versus Sketchy’s $O(mk + nk)$, and (b) even for the S-Adagrad case, prior work does not provide a worst-case $O(\\sqrt{T})$ regret bound when no spectral decay is present $(\\Omega_k = \\Theta(T))$. We present the first analysis which achieves this and demonstrate how this guarantee can rigorously compose with Shampoo.\n\n2. The entire tree of citation descendants for SCFD has 9 papers according to Google Scholar. None focus on the general OCO case (in fact, all focus on linear bandits or bandits in general). Perhaps the reviewer has an analysis of SCFD for the general OCO case in mind and could provide a citation?\n"}, {"review_id": "qHT4Q3Zztr", "reviewer": "Reviewer_CbL3", "summary": "Motivated by empirical observation that Kronecker-factored gradient covariance matrix is (approximately), this work proposed a sketching scheme for memory-efficient preconditioner for gradient-based optimizer.", "questions": "- Although it's standard, it would be helpful for readers to summarize how the regret guarantee can be translated into the stationary point guarantee for nonconvex problems. Maybe, include a corollary/proposition so that readers can understand how the guarantees in this paper translate into the nonconvex land?", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The observation of the low-rankedness of the KF convariance matrix in Figure 3 is quite intriguing.\n- The regret analysis of the sketchy for OCO looks sound.", "weaknesses": "- From Figure 1, it seems that the main application of sketchy is S-Shampoo that achieves the sublinear memory. However, I wonder for shampoo (where the factorization has a certain structure) doing a sketching according to sketchy would actually lead to a low-rank KF covariance matrix? For the adagrad case, it makes a complete sense that Figure 3 provides an empirical justification for S-Adagrad, but I'm quite not sure whether this actually justifies S-Shampoo.\n\n-Also, given the factorization structure of Shampoo, I'm not sure how doing a SVD and sketching individual factor leads to capturing the major component of the entire KF covariance matrix.\n\n- Although the authors said they didn't focus on the implementation of S-Shampoo, it would be helpful for readers how the optimizers compare in terms of wall-time to see how sketchy fare in practice.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time and effort to review our paper. Here are our responses to your concerns:\n\nWeaknesses:\n\n1. From Figure 1, it seems that the main application of sketchy is S-Shampoo that achieves the sublinear memory. However, I wonder for shampoo (where the factorization has a certain structure) doing a sketching according to sketchy would actually lead to a low-rank KF covariance matrix? For the adagrad case, it makes a complete sense that Figure 3 provides an empirical justification for S-Adagrad, but I'm quite not sure whether this actually justifies S-Shampoo.\n\n**Response**: Figure 3 is plotting the intrinsic trace of the Kronecker factors (left and right) from Shamoo, individually. I.e., both factors simultaneously have fast spectral decay. Due to the nature of the way the Kronecker product composes eigensystems, this does also imply that the resulting overall matrix has a certain spectral decay as well, but because the total matrix dimension prohibitively large, S-Adagrad wouldn’t be practical to implement for DNNs. \n\n2. Also, given the factorization structure of Shampoo, I'm not sure how doing a SVD and sketching individual factor leads to capturing the major component of the entire KF covariance matrix.\n\n**Response**: This is a great question, and gets at the harmony of why sketching composes well with Shampoo. It’s not required to approximate the entire KF covariance matrix well to have good preconditioner – thanks to Thm. 5, the regret from S-Shampoo scales with the individual spectral terms of each kronecker factor, not the overall one. This is due to the fact that the regret is bounded by the product of the trace of individual KF terms, rather than the entire matrix, in the Thm. 5 proof, ultimately from the inequality on L638 in the supplement.\n\n3. Although the authors said they didn't focus on the implementation of S-Shampoo, it would be helpful for readers how the optimizers compare in terms of wall-time to see how sketchy fare in practice.\n\n**Response**: Please see the top-level comment, where we add wall-time numbers.\n\nQuestions:\n\nWe are happy to include a corollary in the revision to translate our regret bound to non-convex optimization. In summary, given a non-convex optimization problem $\\min_x f(x)$, where $f$ is $L$-smooth and bounded, one can reduce it to solving strongly convex subproblems. Consider the following algorithm: \n\nFor $t=1, ..., T$, \n\n  $x_{t+1} = \\text{argmin}_x f(x) + 2L ||x - x_t||^2$\n\nOutput $x_{t^*} = \\text{argmin}_x ||\\nabla f(x_t)||$. \n\nThis algorithm is guaranteed to output an $\\epsilon$-stationary point after $O(1/\\epsilon^2)$ iterations. Note that each iteration is a convex optimization problem that can be solved with our algorithm, and hence it will translate into a nonconvex optimization guarantee. For stochastic non-convex optimization, the full algorithm requires $O(\\mu^2\\sigma^2/\\epsilon^4)$ calls to the gradient oracle, where $\\sigma^2$ is the variance of the stochastic gradient oracle, and $\\mu$ is the ratio of adaptivity defined in [1].\n\n[1] Efficient Full-Matrix Adaptive Regularization. Naman Agarwal, Brian Bullins, Xinyi Chen, Elad Hazan, Karan Singh, Cyril Zhang, Yi Zhang Proceedings of the 36th International Conference on Machine Learning, PMLR 97:102-110, 2019."}, {"review_id": "bZi9S3ClJ5", "reviewer": "Reviewer_6Jc9", "summary": "This paper looks at using Frequent-Directions to do (non-diagonal) adagrad more efficiently.  It suggests doing this for both full-matrix adagrad and Shampoo, a version that takes into account the matrix structure of the parameters.", "questions": "- Line 93 says that FD is \"optimal up to universal multiplicative   constants,\" but IMO that's a bit of an overstatement -- it's minimax   optimal for its particular guarantee, but you could either (1) want  a different guarantee, or (2) want a guarantee for a distribution that isn't that particular hard instance.  In particular, it might well not be optimal up to constants for the instances that appear in this paper.\n\n- Is Section 5.1 using beta_2 = 0.999?  It never says that it is, but the rest of the paper implies that it should.\n", "limitations": "Yes.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "I think it's a pretty reasonable and natural idea, and both the theory and experiments\nseem somewhat promising.", "weaknesses": "- The experiments don't get any actual improvement in time or memory relative to existing implementations\n\n- The theorems seem a bit weak.  Take Theorem 3; in the case of no or limited spectral decay, the second error term (with Omega_L) is about d/sqrt(l) times bigger than the first term -- so it doesn't reduce to the first term, even for l=d/2 when it seems like it should be possible.\n\n  If l is (d/10), as suggested by the experiments, isn't that still a sqrt(d) factor increase in the loss?  This seems poor... so poor relative to the experiments that I suspect that the true bound shouldn't have that extra sqrt(d).\n\n\nBasically my concern with this paper is: the idea of using FD for this problem doesn't seem that innovative, which would be fine if this paper really demonstrated that FD is a good idea... but having read this paper, I'm still not confident that FD *is* a good approach here.  It's a reasonable thing to try, certainly, but that was already true before reading the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time and effort to read and review our paper! Here are the responses to your concerns:\n\nWeaknesses:\n\n1. The experiments don't get any actual improvement in time or memory relative to existing implementations.\n\n**Response**: Please see the top comment regarding performance numbers.\n\n2. The theorems seem a bit weak. Take Theorem 3; in the case of no or limited spectral decay, the second error term (with Omega_L) is about d/sqrt(l) times bigger than the first term -- so it doesn't reduce to the first term, even for l=d/2 when it seems like it should be possible. If l is (d/10), as suggested by the experiments, isn't that still a sqrt(d) factor increase in the loss? This seems poor... so poor relative to the experiments that I suspect that the true bound shouldn't have that extra sqrt(d).\n\n**Response**: We have a tighter bound in Corollary 4. We agree that this bound is not tight, but we want to point out that previous works have shown for second-order methods superlinear memory is required (see the paper \"Efficient Convex Optimization Requires Superlinear Memory\" by Marsden et. al.). In our case, the guarantee is established for our proposed algorithm which requires near linear memory $O(mk+nk)$. More importantly, the performance of our algorithms depends on the adaptivity of the regret bound from the term $\\Omega_\\ell$, which we showed empirical evidence for in Figure 2. \n\nBasically my concern with this paper is: the idea of using FD for this problem doesn't seem that innovative, which would be fine if this paper really demonstrated that FD is a good idea... but having read this paper, I'm still not confident that FD is a good approach here. It's a reasonable thing to try, certainly, but that was already true before reading the paper.\n\nThe set of “ideas to try'” may have included FD-for-optimization before this paper. But our work shows that a much more specified, concrete variant, the composition of FD with Shampoo is rigorously-defined belongs to a much narrower set of optimization algorithms, those which are theoretically and practically validated to provide low generalization error (for NNs) or regret (for convex online learners).\n\nThis required a novel integration between FD, Shampoo, and their regret analysis. Previous sketching work was never applied to large neural networks (even empirically focussed ones, like GGT, limit themselves to small resnets). We hope that the shift to approximate second order methods on the KF covariance factors, rather than full per-layer covariance, spur further interest in this algorithmic direction, including further development of efficient implementations.  those which require sublinear memory for accumulators, which were shown on large neural benchmarks such as ResNet-50 on imagenet to actually work better than Adam (unlike all previous low-memory works), and which have a path to an efficient implementation (with iterative top-k eigenvalue routines instead of SVD).\n\nQuestions:\n\n1. Line 93 says that FD is \"optimal up to universal multiplicative constants,\" but IMO that's a bit of an overstatement -- it's minimax optimal for its particular guarantee, but you could either (1) want a different guarantee, or (2) want a guarantee for a distribution that isn't that particular hard instance. In particular, it might well not be optimal up to constants for the instances that appear in this paper.\n\n**Response**: This is a fair catch. We’ve updated our language.\n\n2. Is Section 5.1 using beta_2 = 0.999? It never says that it is, but the rest of the paper implies that it should.\n\n**Response**: $1 - 10^{-4}$ is the largest $\\beta_2$ in the hyperparameter grid (for all optimizers & all neural nets); full grids in the appendices. Using lower $\\beta_2$ values (which in practice are selected by the hyperparameter search) results in even more dramatic spectral decay. For instance, for resnet shampoo, the optimal selected $\\beta_2$ was actually 0.988. Repeating the synthetic data experiment on L272 to L274, the covariance factor’s intrinsic rank for random normal gradients falls from 862.13 (0.25) to 585.30 (0.45) as $\\beta_2$ goes from 0.999 to 0.988.\n"}]

Author rebuttal:
Thank you for all the reviewers for taking the time to read and review our work. Many reviewer comments focus on the differences between S-Adagrad and prior work which inspects FD for optimization.

We invite the reviewers to look at the following sections of our paper. As mentioned in L31, the Kroneker-factored gradient is of particular interest because FD applied to each factor results in sublinear usage. We discuss this in:

The first of our contributions L37-L41 “we recover full-matrix AdaGrad regret up to additive spectral terms under a memory constraint, providing a novel guarantee without curvature assumptions” (not present in the literature) and “[r]igorously composing our approach with Shampoo (Sec. 4.2), unlocks a second order algorithm which requires sub-linear memory for its accumulators.”
Related work L62-68: “approximating the factored covariance $L_t\otimes R_t$ requires less memory than the full covariance and explains why our method can scale to large modern architectures whereas Agarwal et al. [6] cannot”. Note this same explanation applies to RFD,  SCFD, Ada-FD, Radagrad, etc.
Figure 1 on page 4, which diagrams how historical sketch-based methods require superlinear memory $O(mnr)$ for a matrix size $m\times n$ and Sketchy requires $O(mk + nk)$, which is sublinear.

We attempted to emphasize the two differences (novel guarantee, sub-linear memory) from prior work by bolding them in the introduction in the contributions portion at the end of the introduction, but we can include additional discussion in the first page as well.

Many of these historical low-memory approaches have failed to gain traction in the community. We believe that because they did not have $\sqrt{T}$ regret in the OCO case, composition with Shampoo was never a natural option for these methods. Since they all require essentially multiple copies of the neural network weights to be stored to operate, they were never of much interest to practitioners (whose expectations are based on Adam, which requires just one copy for accumulators).

Several reviewers have asked about wall-clock time. We emphasize that we did not perform extensive engineering optimizations. Both Sketchy and Shampoo have iterative variants, but SVD/eigh were used, respectively, for implementation simplicity. For all datasets, numbers below represent the 95% bootstrap confidence interval constructed from taking the median of steps/sec measurements across 5 runs per dataset/algorithm pair. Each run had about 50 such measurements taken uniformly over the course of training. Librispeech data was unfortunately unusable as the cluster scheduler we used interrupted execution frequently, invalidating measurements.

Steps per second (higher is better), for setup described in appendix / Figure 2

**Imagenet**

   Sketchy [6.173, 6.348]

   Adam [14.909, 18.904]

   Shampoo [2.956, 3.032]

**OGBG-MolPCBA**

   Sketchy [4.193, 4.238]

   Adam [4.426, 4.526]

   Shampoo [3.522, 3.766]

Even without heavy engineering, Sketchy improved in runtime over Shampoo in this replicated data-parallel setting (16 TPUv4 per run) because even without an iterative routine, it performs a reduced-size SVD (per L287 in the paper) instead of a full decomposition for gradient KF factors.


Meta-review: Dear Authors,

Thank you for your valuable contribution to Neurips and the ML community. Your submitted paper has undergone a rigorous review process, and I have carefully read feedback provided by the reviewers and considered the author rebuttal in detail.

This paper proposes the use of low rank preconditioners in optimization via the Frequent Directions (FD) method, which is a streaming low-rank approximation algorithm maintaining a low memory budget. The authors provide a theoretical regret bound and present numerical simulations.

While all the reviewers agree that the work is interesting and potentially impactful, the paper received some critical reviews. In particular, Reviewer fXUQ was initially highly critical regarding the novelty of the method. The reviewer pointed out connections with existing works that apply FD in optimization. The author rebuttal addressed some of these comments and clarified that the main contribution is in the analysis of the non-Kronecker-factored case. The reviewer argues that a reorganization is needed to clarify the contributions and connections to existing proposals (e.g., Spectral Compensation Frequent Directions). Although the theoretical analysis appears to be a definite improvement, in the reviewer discussion phase, it was not entirely clear how much of the algorithmic proposal is novel. In addition, Reviewer 6Jc9 was critical about the theory and experiments. The reviewer claimed that the theoretical bound is loose and the experiments are not conclusive. The author rebuttal led the reviewer to increase their score. We congratulate the authors for the effective rebuttal, which clarified some concerns.

I am pleased to recommend the acceptance of your paper for publication at our conference. This is an excellent achievement, and I congratulate you on your success. As we move forward with the publication process, I would like to remind you to carefully review the feedback and comments provided by the reviewers. While the overall assessment is positive, the reviewers have offered valuable suggestions that can further strengthen the quality of the paper. 

In particular, as Reviewer fXUQ suggested, I strongly urge the authors to discuss connections with Spectral Compensation Frequent Directions (SCFD) and make sure the contributions are not overstated.

I encourage you to use this feedback to make any necessary improvements and refinements to the manuscript before submitting the final version for publication.

Once again, thank you for submitting your work to our conference, and congratulations on this well-deserved recognition. I look forward to seeing your paper published in our conference proceedings and witnessing the impact of your research on the scientific community.

Best,
AC
Decision: Accept (poster)