Title: Sketchy: Memory-efficient Adaptive Regularization with Frequent Directions
Venue: NeurIPS 2023
Accepted: yes
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

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 secondorder 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.

1 Introduction

DL optimization commonly relies on adaptive gradient methods, namely the Adam optimizer [1]. It differs from stochastic gradient descent in that the learning rate is a 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 is 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-art 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 [5, 9] introduces a factorization of full matrix preconditioning method with adoption in large-scale industrial applications such as training Google's ads click-through-rate model [10]. Furthermore, as hardware evolves, memory efficiency becomes an increasing concern, as 'logic improves much faster than wires and SRAM, so logic is relatively free' [11]: from TPUv2 to TPUv3, per-chip bfloat16 operations per second improved 2 . 67 × but memory bandwidth only improved 1 . 29 × . GPUs exhibit a similar pattern for compute and memory increase, at 5 × and 2 . 2 × , for V100 to A100 [12].

Investigation into the Kronecker-factored gradient covariance matrix reveals a concentrated, but changing, spectrum (Fig. 3), suggesting the majority of the spectral mass can be represented by a lowrank 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 a 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 a novel guarantee without curvature assumptions (Sec. 4.1). Rigorously composing our approach with Shampoo (Sec. 4.2) unlocks a secondorder algorithm which requires sub-linear memory for its accumulators.

By modifying FD for exponential moving averages (Sec. 4.3), we demonstrate a practical algorithm competitive with at-least-linear memory Shampoo and Adam in three modern DL settings (Sec. 5.1). While previous work [8, 14] shows rank-1 preconditioners are effective for trading off quality for memory, these results demonstrate a 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 Kronecker-factored gradient covariance in DL settings (Sec. 5.2).

2 Setting and Definitions

Regret and Optimization. The optimization problem of training a deep neural network has a nonconvex objective loss function f . Since finding the global optimum is computationally intractable in general, theoretical guarantees focus on convergence to an ε -approximate first-order optimum: a point x such that ‖∇ f ( x ) ‖ ≤ ε . A smooth non-convex problem can be reduced to solving a series of offline convex problems [6]. The convex sub-problems have form f t ( x ) = f ( x ) + c ‖ x -x t ‖ 2 , where c is a constant and x t is an iterate in the optimization process. Using online-to-batch 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 B.1. 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 x t ∈ K iteratively, where K ⊆ R d is a convex decision set (take K = R d if unconstrained). After the decision is made, the adversary reveals a convex loss function f t , to which the algorithm suffers cost f t ( x t ) . Upon receiving the cost, the algorithm updates its decision for the next iteration. The regret for an online algorithm is given by

Sketching and the Frequent Directions Method. Given a stream of vectors g t ∈ R d , t ∈ [ T ] , we utilize the FD sketch [13] given in Alg. 1 which maintains a low-rank approximation of the true running covariance G t = ∑ s ≤ t g s g /latticetop s . At each time t , it maintains a matrix B t of size d × /lscript whose last column is 0 and whose square is the sketch B t B /latticetop t = ¯ G t . After seeing g t from the stream, we update the previous matrix using Alg. 1, which updates B t +1 of size d × /lscript whose last column remains 0 ; take B 0 = 0 . B t +1 is obtained by decomposing the sum of ¯ G t and the newly observed matrix, keeping only the top eigendirections, and reducing the eigenvalues uniformly by /lscript -th eigenvalue. At every iteration t , denote ρ t := λ ( t ) /lscript be the removed eigenvalue from the covariance update in Alg. 1.

For convenience, let ρ 1: t def = ∑ t s =1 ρ s be the cumulative escaped mass. For a matrix X , we denote its i -th leading eigenvalue by λ i ( X ) . Let ‖ · ‖ F denote the Frobenius norm of a matrix.

The fundamental property of FD is that applying Alg. 1 over a stream of vectors g t , with B 0 = 0 , the sum of escaped mass ρ t = λ ( t ) /lscript can be bounded by the bottom eigenvalues of G T , formally given by the following lemma:

Algorithm 1 Frequent Directions Update ( FD-update )

Require: Invariant that last column of B t

Ensure: The last column of B t

-1 is 0. is 0.

2: Input: New symmetric PSD matrix M t ∈ R d × d .

1: Input: Previous state ¯ G t -1 = B t -1 B /latticetop t -1 ∈ R d × d

3: Eigendecompose ¯ U t diag λ ( t ) ¯ U /latticetop t = ¯ G t -1 + M t where λ ( t ) contains descending eigenvalues.

5: Update B t = U t diag ( λ ( t ) [1: /lscript ] -λ ( t ) /lscript ) 1 / 2 .

4: Define U t as the matrix whose columns are the first /lscript columns of ¯ U t , and λ ( t ) [1: /lscript ] be its eigenvalues.

output λ ( t ) /lscript , B t B /latticetop t .

Lemma 1 (Liberty [16]) . The cumulative escaped mass ρ 1: T can be upper bounded as

Proof. See Sec. B.2.

3 Related Work

3.1 Spectral Analysis of DL Training

Denote the loss function of the i -th example for weights x as f i ( x ) . The spectrum of the Hessian matrix ∑ i ∇ 2 f i has been the subject of intensive investigation in DL [17, 18, 19, 20] and its properties have been used to devise training methods [4, 21].

Recent papers [22, 23, 24] inspect the covariance matrix, ∑ i ( ∇ f i )( ∇ f i ) /latticetop . In small models, where its computation is feasible, these works identify fast spectral decay.

Agarwal et al. [6] take advantage of this observation by using a low-rank approximation of the whole covariance matrix, based on a limited history of the gradients, ∑ r i = t -r ( ∇ f i )( ∇ f i ) /latticetop . This approach still requires r copies of the model gradients in memory, where typically r should scale with β -1 2 , with β 2 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.

In our work, we validate the decay hypothesis holds across the per-layer factored covariance matrices in several modern neural networks. For a layer's gradient matrix G i at the i -th example and a second moment decay term β 2 , our work inspects spectral decay for L t = ∑ i β t -i 2 G i G /latticetop i and R t = ∑ i β t -i 2 G /latticetop i G i ; the spectral structure for these outer products is not well-documented. Furthermore, as described in Sec. 3.4, approximating the factored covariance L t ⊗ 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.

3.2 Sublinear Memory Methods

Extreme Tensoring [7], AdaFactor [14], and SM3 [8] 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 our 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 to remaining indices.

Crucially, Adam, which uses linear memory for second moment representations, compares favorably in terms of quality to all of these sublinear-memory methods. By increasing rank in factored

√

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 in 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 additional assumptions (Tbl. 1). In this section, we consider the OCO setting over dimension d (Sec. 2).

Random projection (Ada-LR) is most spiritually similar to our work [25]. Although it does not reduce memory usage, it relies on random projections to lower dimension /lscript ≤ d to reduce inverse matrix computation costs. An alternative without formal guarantees, RadaGrad, reduces memory consumption to O ( d /lscript ) ; however, as with all Johnson-Lindenstraus projection methods, it suffers a probabilistic failure rate scaling as O ( /lscript -1 ) (in comparison, our method inherits FD's determinism).

Frequent Directions (FD) [13, 16], provides an alternative matrix sketching approach from the data streaming literature. As an 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 to an adversarially-chosen true covariance G T in operator norm constructed from the corresponding gradients must use O ( d /lscript ) 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 ∑ /lscript -1 i =1 λ i ( G T ) and additive error in ∑ d i = /lscript λ i ( G T ) . However, without the exp-concave assumption, FD-SON falls back to a gradient-descent-like default regret of O ( λ /lscript : d √ T ) , which can be Ω ( T ) without spectral decay. In the context of linear bandits, Chen et al. [28] uses FD for memory reduction. The resulting algorithm, SCFD, is similar to Alg. 2, but SCFD lacks a projection step and does not handle general domains K . We emphasize that our contribution is in the novelty of our regret analysis for general OCO settings, in contrast to linear bandits.

The main prior work exploring FD for general online convex settings, Wan and Zhang [26], extends the FD-SON approach by adding a fixed diagonal perturbation δ I to an FD-based preconditioner, in Ada-FD. However, this approach does not achieve √ T regret even in a non-adversarial setting with stochastic linear cost functions (Observation 2), where learning rate and δ are tuned. Dynamically changing diagonal regularization is essential for worst-case O ( √ T ) performance.

Observation 2. Suppose we receive linear cost functions f t ( x ) = 〈 x, g t 〉 , where g t ∈ R d is a random vector drawn iid from any distribution over r ≤ d orthonormal vectors W . For any sketch size /lscript ≤ r , the bound on the expected regret of Ada-FD is Ω ( T 3 / 4 ) .

Diag. AdaGrad

Adafactor/SM3

Shampoo

GGT

Sketchy (Ours)

m + n

mk + nk

mn

m 2

+ n 2

mnr

sublinear

k{

Lt

tRt

Wan and Zhang [26] remark Ada-FD has √ T regret when G T is low rank with rank below k , so d -k of its eigenvalues are 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 a reduced QR decomposition, rank-1-updated every step), the full matrix AdaGrad algorithm can be perfectly recovered without using more than O ( dk ) memory.

In concrete convex examples, Sketchy compares favorably to these approaches (Appendix A).

3.4 Shampoo

Perhaps our most compelling application is reducing the memory of Shampoo [5, 9, 10]. Shampoo is 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 × n , AdaGrad treats the matrix-shaped parameters as a vector of size mn , and the preconditioner has size m 2 n 2 ; Shampoo instead has left and right preconditioners L, R of size n × n and m × m , respectively, with the preconditioned update L -1 / 4 WR -1 / 4 . Write vec ( W ) as the vectorized weights, then it is equivalent to ( L ⊗ R ) vec ( W ) = vec ( LWR ) , where ⊗ denotes the Kronecker product. Figure 1 illustrates 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 a Kronecker-factored 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-art results in large-scale tasks such as language modeling with BERT-Large and image classification on ImageNet in [9].

Figure 1 elaborates why the composition of FD and Shampoo is essential to avoid memory consumption asymptotically greater than parameter count for approximate full matrix regularization.

However, Shampoo memory costs may still be prohibitive for rectangular weight matrices. In BERTLarge [29], most parameters are in the feed-forward network layers, which consist of 4096 × 1024 dense kernels; other transformers follow similar narrow-to-wide patterns. For large models, occupying even 4 × memory for the left preconditioner can frequently result in OOM in memory-constrained 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 × n as mn/b 2 blocks of size b × b for some block size b < min( m,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, L t ⊗ I or I ⊗ R t . Note, however, that the one-sided approximation doesn't help with vector parameters, such as

those that appear for the bias terms in dense layers or layer norms [30]. For 3D weights, such as those which appear in homogeneous Mixtures of Experts [31], 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.

4 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, in Sec. 4.3, we modify FD to support exponential moving averages.

The main technical novelty in incorporating Alg. 1 to AdaGrad (Alg. 2) and Shampoo (Alg. 3) is 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, as demonstrated in Lemma 10 for AdaGrad and Lemma 14 for Shampoo. The regret guarantee for AdaGrad ([2]) and Shampoo ([5]) 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 -1 / 2 t in Alg. 2 denotes the Moore-Penrose pseudoinverse of the matrix ˜ G 1 / 2 t . Our main algorithm, Sketchy AdaGrad, in this section exploits the FD approach outlined in Alg. 1 as the sketching method in AdaGrad. In particular, at every time step, we pass the newly received subgradient g t into Alg. 1, which updates and maintains a low-rank sketch ¯ G t of the AdaGrad preconditioning matrix G t . We keep track of the cumulative escaped mass ρ 1: t , which we add back to the low-rank sketch to create the Sketchy preconditioner ˜ G t , with which we perform the regular AdaGrad descent and projection.

Theorem 3. Define Ω /lscript = min k< /lscript ( /lscript -k ) -1 ∑ d i = k +1 λ i ( G T ) , then with η = D √ 2 , Alg. 2 guarantees the following additive regret bound:

where D is the diameter of the constraint set K if K is bounded and max t ∈ [ T ] ‖ x t -x ∗ ‖ 2 otherwise.

Proof.

Notably in Thm. 3, Regret T = O ( √ T ) and the lower eigenvalue dependence Ω /lscript is additive.

Corollary 4. We can improve Theorem 3 slightly to

The regret bound above holds under the optimal tuning of the learning rate, which depends on problem quantities that can be unknown a priori. It is possible to design a parameter-free variant of Alg. 2 by using the norm ‖ x ‖ t = ( x /latticetop ( ˜ G t + I ) 1 / 2 x ) 1 / 2 in the projection step of Alg. 2, as seen in [32].

4.2 FD for Shampoo

In this section, we adapt FD-update to Shampoo [5]. For simplicity, we optimize over R m × 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 ¯ L t , ¯ R t for the left and right preconditioning matrices. We keep track of the cumulative escaped masses ρ L 1: t and ρ R 1: t from sketching the left and right preconditioning matrices, respectively, and compensate the cumulative escaped mass to create the left and right Sketchy preconditioning matrices ˜ L t , ˜ R t .

Denote L T def = ∑ T t =1 G t G /latticetop t + ε I and R T def = ∑ T t =1 G /latticetop t G t + ε I .

Theorem 5. Suppose G 1 , . . . G T have rank at most r . Then Alg. 3 run with η = D/ √ 2 r guarantees the following regret bound:

where D = max t ∈ [ T ] ‖ X t -X ∗ ‖ F and Ω L, /lscript , Ω R, /lscript are analogous bounds for ρ L 1: T , ρ R 1: T from Lem. 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 ∑ t G t G /latticetop t and ∑ t G /latticetop t 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 ρ 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 ( β 2 ) = ∑ t β T -t 2 G t G /latticetop t exhibits fast spectral decay (and analogously for R t ). Luckily, thanks to the recursion L t +1 ( β 2 ) = β 2 L t + G t +1 G /latticetop t +1 , the FD sketch may easily be adopted for this setting.

Observation 6. Given a stream g t of vectors for t ∈ [ T ] , sketch size /lscript , updates ( ρ ( β 2 ) t , ¯ G ( β 2 ) t ) = FD-update ( β 2 ¯ G ( β 2 ) t -1 , g t g /latticetop t ) , and G ( β 2 ) T = ∑ T t =1 β T -t 2 g t g /latticetop t , we have

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. 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 /lscript , which we do not tune, but set to /lscript = 256 , which translates to 4 × 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. 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

the second moment, β 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 = ∑ i β t -i 2 G t G /latticetop t and R t = ∑ i β t -i 2 G /latticetop t 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 β 2 = 0 . 999 , we emphasize that the behavior in Fig. 3 is an emergent property of DL training. Though surely a lower β 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 × d matrices with iid N (0 , 1) entries, then numerical experiments show that the average intrinsic dimension of ∑ n -1 i =0 β i 2 x i x /latticetop i 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.

6 Discussion

Up to spectral error, Alg. 2 achieves full-matrix AdaGrad regret despite approximating the smallest part of the spectrum of G -1 / 2 t 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 ˜ G -1 / 2 t 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 [ β 1 / 2 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. [9] 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 th gradient and update its covariance inverse roots (see Appendix 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 /lscript eigenvalues, iterative Lanczos-like routines which are accelerator-friendly, such as LOBPCG [40], may allow incremental updates to ˜ G -1 / 2 t 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.

7 Conclusion

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 /lscript across the many tensors in a network, as the spread in Fig. 3 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].

8 Acknowledgments

Elad Hazan acknowledges funding from the Office of Naval Research grant N000142312156, the NSF award 2134040, and Open Philanthropy.

References

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR (Poster) , 2015.

John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research , 12(7), 2011.

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 pattern recognition , pages 770-778, 2016.

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

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

Naman Agarwal, Brian Bullins, Xinyi Chen, Elad Hazan, Karan Singh, Cyril Zhang, and Yi Zhang. Efficient full-matrix 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.

Xinyi Chen, Naman Agarwal, Elad Hazan, Cyril Zhang, and Yi Zhang. Extreme tensoring for low-memory preconditioning. In International Conference on Learning Representations , 2019.

Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimization. Advances in Neural Information Processing Systems , 32, 2019.

Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09018 , 2020.

Rohan Anil, Sandra Gadanho, Da Huang, Nijith Jacob, Zhuoshu Li, Dong Lin, Todd Phillips, Cristina Pop, Kevin Regan, Gil I. Shamir, Rakesh Shivanna, and Qiqi Yan. On the factory floor: Ml engineering for industrial-scale ads recommendation models, 2022.

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 tpuv4i: Industrial product. In 2021 ACM/IEEE 48th Annual International Symposium on Computer Architecture (ISCA) , pages 1-14. IEEE, 2021.

William J Dally, Stephen W Keckler, and David B Kirk. Evolution of the graphics processing unit (gpu). IEEE Micro , 41(6):42-51, 2021.

Mina Ghashami, Edo Liberty, Jeff M Phillips, and David P Woodruff. Frequent directions: Simple and deterministic matrix sketching. SIAM Journal on Computing , 45(5):1762-1792, 2016.

Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning , pages 4596-4604. PMLR, 2018.

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

Edo Liberty. Even simpler deterministic matrix sketching. arXiv preprint arXiv:2202.01780 , 2022.

Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues of the hessian in deep learning: Singularity and beyond. arXiv preprint arXiv:1611.07476 , 2016.

Levent Sagun, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over-parametrized neural networks. arXiv preprint arXiv:1706.04454 , 2017.

Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net optimization via hessian eigenvalue density. In International Conference on Machine Learning , pages 2232-2241. PMLR, 2019.

Adepu Ravi Sankar, Yash Khasbage, Rahul Vigneswaran, and Vineeth N Balasubramanian. A 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.

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.

Guy Gur-Ari, Daniel A Roberts, and Ethan Dyer. Gradient descent happens in a tiny subspace. arXiv preprint arXiv:1812.04754 , 2018.

Craig Bakker, Michael J Henry, and Nathan O Hodas. Understanding and exploiting the low-rank structure of deep networks. 2018.

Zeke Xie, Qian-Yuan Tang, Zheng He, Mingming Sun, and Ping Li. Rethinking the structure of stochastic gradients: Empirical and statistical evidence. arXiv preprint arXiv:2212.02083 , 2022.

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.

Yuanyu Wan and Lijun Zhang. Efficient adaptive online learning via frequent directions. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2021.

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.

Cheng Chen, Luo Luo, Weinan Zhang, Yong Yu, and Yijiang Lian. Efficient and robust high-dimensional linear contextual bandits. In Proceedings of the Twenty-Ninth International Conference on International Joint Conferences on Artificial Intelligence , pages 4259-4265, 2021.

Jacob Devlin, Ming-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 1 (Long and Short Papers) , pages 4171-4186, 2019.

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016.

Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538 , 2017.

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 in Neural Information Processing Systems , volume 33, pages 8836-8846. Curran Associates, Inc., 2020.

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 (IJCV) , 115(3):211-252, 2015. doi: 10.1007/s11263-015-0816-y.

Anmol Gulati, James Qin, Chung-Cheng Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang, Yonghui Wu, et al. Conformer: Convolution-augmented transformer for speech recognition. arXiv preprint arXiv:2005.08100 , 2020.

Vassil Panayotov, Guoguo Chen, Daniel Povey, and Sanjeev Khudanpur. Librispeech: an 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.

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:1806.01261 , 2018.

Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs.

In Hugo Larochelle, Marc Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and HsuanTien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual , 2020.

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.

Roman Vershynin. High-dimensional probability: An introduction with applications in data science , volume 47. Cambridge university press, 2018.

Andrew V Knyazev. Toward the optimal preconditioned eigensolver: Locally optimal block preconditioned conjugate gradient method. SIAM journal on scientific computing , 23(2):517541, 2001.

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.04836 , 2016.

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 on Learning Representations , 2020.

Luo Luo, Cheng Chen, Zhihua Zhang, Wu-Jun Li, and Tong Zhang. Robust frequent directions with application in online learning. The Journal of Machine Learning Research , 20(1):16971737, 2019.

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.

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.

Adam Kalai and Santosh Vempala. Efficient algorithms for online decision problems. Journal of Computer and System Sciences , 71(3):291-307, 2005.

Koenraad MR Audenaert. A generalisation of mirsky's singular value inequalities. arXiv preprint arXiv:1410.4941 , 2014.

Justin M. Gilmer, George E. Dahl, and Zachary Nado. init2winit: a jax codebase for initialization, optimization, and tuning research, 2021. URL http://github.com/google/init2winit .

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.

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.

Tensorflow datasets, a collection of ready-to-use datasets. https://www.tensorflow.org/ datasets , 2023.

Michael L. Waskom. seaborn: statistical data visualization. Journal of Open Source Software , 6 (60):3021, 2021. doi: 10.21105/joss.03021.

J. D. Hunter. Matplotlib: A 2d graphics environment. Computing in Science & Engineering , 9 (3):90-95, 2007. doi: 10.1109/MCSE.2007.55.

Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, 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.1038/s41586-020-2649-2.

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 J Carey, ˙ Ilhan Polat, Yu Feng, Eric W.

Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio 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.

MLCommons R © open engineering consortium. MLCommons Algorithmic Efficiency. https: //github.com/mlcommons/algorithmic-efficiency , 2023.

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.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 , 2017.

Boris Dayma and Rohan Anil. Evaluation of Distributed Shampoo: Comparison of optimizers: Distributed Shampoo, Adam & Adafactor. Weights & Biases Report, 2022.

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.

Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang. Disentangling adaptive gradient methods from learning rates. arXiv preprint arXiv:2002.11803 , 2020.

Tsuyoshi Ando. Concavity of certain maps on positive definite matrices and applications to hadamard products. Linear algebra and its applications , 26:203-241, 1979.

Rajendra Bhatia. Matrix analysis . Springer, 1997.

Kaare Brandt Petersen, Michael Syskind Pedersen, et al. The matrix cookbook. Technical University of Denmark , 7(15):510, 2008.

Roger W Brockett. Finite dimensional linear systems . SIAM, 2015.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nI think it's a pretty reasonable and natural idea, and both the theory and experiments\nseem somewhat promising.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

comments: 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.

STRENGTHS:
- The observation of the low-rankedness of the KF convariance matrix in Figure 3 is quite intriguing.
- 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.

-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.

- 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.

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
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\nMy 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\nQUESTIONS:\n1.\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\nLIMITATIONS:\nNo potential negative societal impact. The detailed discussion for related work is necessary. Please see above for details.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: 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.

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).

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06