Sketchy: Memory-efficient Adaptive Regularization with Frequent Directions

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.

Paper

Similar papers

Peer review

Reviewer 6Jc96/10 · confidence 3/52023-07-05

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.

Strengths

I think it's a pretty reasonable and natural idea, and both the theory and experiments seem somewhat promising.

Weaknesses

- The experiments don't get any actual improvement in time or memory relative to existing implementations - 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). Basically 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.

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

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Yes.

Reviewer CbL36/10 · confidence 4/52023-07-05

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?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A

Reviewer fXUQ3/10 · confidence 4/52023-07-06

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.

Strengths

1. The design of the algorithm is well-motived and rigorous theoretical analysis is provided. 2. The experiments on deep nerul network shows the proposed algorithm improved the recent proposed baseline methods Shampoo.

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: 1. To 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. 2. The 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.

Questions

1. Can you provide a detailed discussion to compare the proposed algorithms with RFD, SCFD and related online optimization algorithms? 2. SCFD 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?

Rating

3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

No potential negative societal impact. The detailed discussion for related work is necessary. Please see above for details.

Reviewer rVs56/10 · confidence 4/52023-07-06

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?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

Some discussion is provided in the Discussions section.

Reviewer 6Jc92023-08-12

Thanks for your response. I'm still a little concerned that the theorem seems loose and the experiment isn't very comprehensive, but your point about validating a very concrete proposal is well taken. I'm willing to raise my score to 6.

Reviewer CbL32023-08-13

Thanks!

I read the response and it's satisfactory so I raise my score to 6.

Reviewer rVs52023-08-21

Thank you for the response

Thank you for the clarifications. I did not realize the focus on the Kronecker-factored gradient covariance, although this doesn't change my initial impressions significantly. I will keep my score.

Authorsrebuttal2023-08-21

Please reconsider the focus on Kronecker factors.

We apologize if our earlier reply was too short to emphasize how the focus on Shampoo sketching and Kronecker factors results in dramatically different guarantees from FD-SON (please also note our preconditioner computation is quite different, in the matrix case requiring inverse fourth roots, rather than inverse matrices). A Kronecker-factored covariance results in $O(m^2+n^2)$ memory use for accumulators for an $m\times n$ tensor to optimize. With our sketching approach, this becomes $O(mk+nk)$, for sketching rank $k<\min(m,n)$, which is now sublinear. As our top-level comment emphasizes, this results in a dramatic memory decrease, asymptotically smaller compared to previous sketching approaches and linear-memory algorithms like Adam. Moreover, we provide novel analysis in the non-Kronecker-factored case (S-Adagrad), providing the first $O(\sqrt{T})$ bound for a sketched OCO algorithm in the literature, not present in past works such as FD-SON. Given two strong theoretical differentiations from FD-SON (1 - asymptotically less memory in the matrix case, 2 - square root regret bound for general OCO setting without curvature) and one practical one (our sublinear memory allows application to large deep neural networks, could the reviewer provide reasons for their reasons to reject, which they say outweighs limits to accept? In our top-level comment, we quote the contributions of our work from the introduction, which clarify the novelty in question.

Reviewer rVs52023-08-21

Thank you for the response

I appreciate the authors' efforts in defending their work, and I would like to better understand their points. My understanding of the work, taking the focus on the Kronecker product into account, is that it studies an FD-SON-like algorithm for matrix-shaped parameters (such as those for neural net weights), where instead of preconditioning the mn parameters as a vector, it preconditions the m x n matrix from the left and right. In this setting, it is shown that a low rank approximation of the left and right preconditioners gives a smaller memory algorithm with guarantees that interpolates correctly the full rank left/right matrix preconditioner case. - Applying low rank approximation to asymptotically save on memory is an ubiquitous idea, and I do not think the community particularly benefits from having another paper that instantiates this idea. However, I would like to discuss the regret analysis more to better understand the contributions here. - My initial impression was that analyzing regret bounds for Shampoo (without low rank approximation) would be straightforward and known in the literature. Is the claim that this is already a challenge that is newly contributed by this work? I believe this would be the point that makes the biggest difference in my impression of the work. - Given an analysis of Shampoo, working the analysis of FD into the regret analysis, which I understood to be the main theoretical thrust of this work, seems like a marginal contribution to me. In particular, FD already gives the machinery to reason about how the escaped mass of low rank components evolve through streams of updates, and putting this into a regret analysis seems to me like just applying a theorem in a new context. Please let me know if there are any particular challenges that I should take another look at.

Authorsrebuttal2023-08-21

Thank you for continued interest & questions

We greatly appreciate the reviewer's interest in understanding the work! Thank you for your continued investment. Even in the vector-parameter case (S-Adagrad), low-rank second-moment algorithms have previously required either: 1. curvature (FD-SON and its derivatives do this) to yield an ONS-type algorithms 2. weak bounds (Radagrad is heuristic in low memory form, and see Observation 1 for Ada-FD). So even for analyzing the vector case, our Theorem 3 provides a novel worst-case sqrt-T guarantee when dealing with general convex functions with vector parameters. This is a gap in the literature and may be of interest due to recent work on lower bounds (Efficient Convex Optimization Requires Superlinear Memory by Marsden et al 2022). What we do differently from, say, Ada-FD, is set up a strictly more convex preconditioner, which allows us to establish the approximate isometry in Remark 9 in the Supplement (L569). This two-sided bound around the "ideal" full matrix preconditioner makes the FD preconditioner essentially a modular component in OCO proofs, so long as the algorithm applies the dominating approximation to an exact preconditioner: we rely on the same isometry for Shampoo's Kronecker factors $L_t,R_t$ (implied by Lemma 12). We hope the community can use the same technique in other OCO settings too; surely the community stands to gain from the demonstration of this technique in our work, as the OCO sqrt-T result was missing. Finally, we invite the reviewer to take a look at our empirical results, summarized in Figure 2. We believe that this another "challenge" which has been neglected in this line of work and not yet considered by the reviewer: empirical application of low memory algorithms to real optimization problems of interest to practictioners. As mentioned in Fig.1 and our top-level comment, existing low-memory approaches result in **superlinear** accumulator memory size. Perhaps on the theoretical side, low memory approaches are well known, but in practice algorithms such as RFD, SCFD, Ada-FD, Radagrad, GGT could not be applied to large networks such as the ResNet-50, Conformer, and GNN we have in our work. Storing multiple copies' worth of the parameters for these networks on GPUs is prohibitive. Adafactor may be the only real example of sublinear memory optimizers in practice, but it does not have a robust theory and performs much worse than Adam in terms of quality on all three of our datasets (notice we improve on Adam in all three).

Reviewer rVs52023-08-21

Thank you for the response

I appreciate the authors' in-depth discussion and clarification of their results, and for their patience with my last-minute questions on the last day of the rebuttal period. I still am not 100% convinced about the importance of the low rank-based optimizers in practice since it seems like a "low hanging fruit" idea, given my lack of experience in this area of practical implementations of optimizers and some preliminary searches, I trust that this is indeed a significant contribution. I am convinced that the theoretical results are more significant than I originally perceived, and I would like to raise my contribution score to a 3 (good) and the overall score to a 6 (weak accept), primarily on the basis of the theoretical results. I would prefer for the paper to be written more to focus on the discussion of the theoretical results, although this is perhaps a more personal taste and the other reviewers may disagree. Thank you again to the authors for their patience.

Authorsrebuttal2023-08-21

Thank you for the review!

We really appreciate the reviewer's questions and their engagement on our work. This is a great experience for us as authors and, based on the feedback, helps us modify text to direct attention in the methods section to the construction of the reusable concepts (such as the approx isometry) in our proofs!

Reviewer fXUQ2023-08-21

Thanks for your response. I think the sketching method in proposed method is very similar to existing regularized FD, and the difference between Algorithm 1 and SCFD is still unclear. Hence, I decide to keep my rating.

Authorsrebuttal2023-08-21

Thanks for the feedback

Given the feedback, and the reviewer's note of similarity between SCFD, we wonder if the following presentation would be amenable to the reviewer? First, we can move algorithm 1 (S-Adagrad) to a section in related work, marking its similarity in the past literature to SCFD in the bandit context. The methods section already mostly focuses on our novel regret **analysis** of S-Adagrad / SCFD, which does not appear in prior work. It appears the reviewer is in agreement that S-Adagrad / SCFD has not yet been analyzed in the OCO setting with sqrt-T worst case regret? In this way, we can emphasize the main contribution of the non-Kronecker-factored case to be our analysis, rather than the algorithmic proposal. However, the key new algorithm unlocked by this OCO approach is through integration with Shampoo, which as our top comment demonstrates is critical for sublinear memory and thus application to DNNs.

Reviewer fXUQ2023-08-22

Thanks for your feedback

I believe it is more appropriate to reorganize the paper, emphasizing that the primary contribution lies in the analysis of the non-Kronecker-factored case rather than the algorithmic proposal. This suggests a major revision and an additional round of review are necessary.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC