Lean and Mean Adaptive Optimization via Subset-Norm and Subspace-Momentum with Convergence Guarantees
Authors: Thien Hang Nguyen
Huy Nguyen
Venue: ICML 2025 poster (2025)
Keywords: Memory-efficient optimization
adaptive optimization
parameter efficient fine-tuning
PEFT
LLMs
high-probability convergence
stochastic optimization
Primary area: optimization
TL;DR: We introduce two complementary techniques for efficient optimization that reduce memory requirements while accelerating training of large-scale neural networks.

Abstract: We introduce two complementary techniques for efficient optimization that reduce memory requirements while accelerating training of
large-scale neural networks. The first technique, Subset-Norm step size, generalizes AdaGrad-Norm and AdaGrad(-Coordinate) through step-size sharing. Subset-Norm (SN) reduces AdaGrad's memory footprint from $O(d)$ to $O(\sqrt{d})$, where $d$ is the model size. For non-convex smooth objectives under coordinate-wise sub-gaussian noise, we show a noise-adapted high-probability convergence guarantee with improved dimensional dependence of SN over existing methods. Our second technique, Subspace-Momentum, reduces the momentum state's memory footprint by restricting momentum to a low-dimensional subspace while performing SGD in the orthogonal complement. We prove high-probability convergence rates for Subspace-Momentum under standard assumptions. Empirical evaluation on pre-training and fine-tuning LLMs  demonstrates the effectiveness of our methods. For instance, combining Subset-Norm with Subspace-Momentum achieves Adam's validation perplexity for LLaMA 1B in approximately half the training
tokens (6.8B vs 13.1B) while reducing Adam's optimizer-states memory footprint by more than 80\% with minimal additional hyperparameter tuning.

Lean and Mean Adaptive Optimization via Subset-Norm and Subspace-Momentumn with Convergence Guarantees
Thien Hang Nguyen 12
Huy Le Nguyen
Abstract
Optimizers comparison LLaMA 1B pro-training for 13.18 tokens
We introduce two complementary techniques for efficient optimization that reduce memory require- ments while accelerating training of large-scale neural networks. The first technique, Subset- Norm step size, generalizes AdaGrad-Narr -orm and AdaGrad(-Coonn -Coordinate) through step-size sharing. Subset-Norm (SN) reduces AdaGrad's memory footprint from (dd) to O(vd), where is the model size. For non-convex smooth objectives under coordinate-wise sub-gaussian noise, show noise-adapted, adapted high-probabiilty conver- gence guarantee with improved dimensional de- pendence of SN over existing methods. Our sec- ond technique, Sebssacce-Mmentung reduces the momentum state's memory footprint by restrict- ing momentum low-dimensional subspace while performing SGD in the orthogonal comple- ment. We prove high-probability convergence result for Subspace-Momentuu under standard as- sumptions. Empirical evaluation on pre- pre-training and fine-tuning LLMs demonstrates the effective- ness of our methods. For instance, combining Subset-Norm with Subspace-Momentun achieves Adam's validation perplexity for LLaMA IB in approximately half the training tokens (6.8B 13. I3.IB) while reducing Adam's optimizer-states memory footprint by more than 80% with mini- mal additional hyperparameter tuning.
Adam (1.99 GB)
GaLore (1.61 GB)
Valitation Perplesity
AdamSN (ours) (2.62 GB)
AdamSNSM (ours) (0.84 GB)
strpo Perpesity:
SN+SSM
Training Steps
Figure 1. Validation perplexity for Adam, GaL.ore orr (Zhao et al., 2024), AdamSN, and AdamSNSM (ours) during L.LMMA 1B model training for 13.1B tokens (100K steps). Optimizer memory footprint is shown in parentheses. Adam achieves perplexity of 16.00 at 100,000 steps, while AdamSN and AdamSNSM exhibit lower perplexity earlier in training 58,000 and 48,000 steps.
2012) are de facto methods for training large-scale deep neu- ral networks. However, the optimizer states for the momen- tum and second moment (or adaptive step size) terms are memory intensive, consuming as much as twice the size of the model. As deep neural networks continue grow in the of large-language models (LLMs), concerns that were previously overlooked, such as the memory consumption of optimizer states, have become an active area of research. Indeed, numerous methods have recently emerged to reduce the memory footprint of optimizer states (e.g. Adam's 's mo- mentum and second moment terms) with approaches rang- ing from quantization (Li et al., 2024a; Dettmers al., 2021; 2024), low-rank decomposition (Hu et al., 2021; Lialin et al., 2023; Zhao et al., 2024; Shazeer & Stern, 201 8), sketching- based dimensionality reduction (Muhamed et al., 2024; Hao et al., 2024), etc. Existing methods either lacks theoretical guarantees, requires strong assumptions, trades too much performance, or requires expensive additional tuning for the memory saving, especially in pretraining tasks.
1. Introduction
Adaptive optimizers like Adam (Kingma & Ba, 2014), Ada- Grad (Duchi et al., 2011), and RMSProp (Ticleman et al., Khoury
'Work done at Northeastern University. College of Computer Sciences, Northeastern University, Boston. *Contextual AI. Correspondence to: Thien Nguyen <nguyen.thiens northcastem. edu; thien. wonpyetucnnotttttaa.ii..a...a Huy Nguyen <hu.nguyen northeastern.ddu:
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
We aim to reduce memory consump- tion while maintaining strong performance and theoreti- cal guarantees. To this end, we introduce two memory-
Our contributions.

---

Subset-Norm and Subspace-Momentum
efficient optimization algorithms for large-scale DNN train- ing: Subset-Norm (SN) for adaptive step-size memory reduction (Section 3) and Subspace-Momentum (SM) for momentum compression (Section 4). While existing approaches trade performance for memory savings, our theoretically-grounee methods achieve both a reduced memory footprint and faster training:
Algorithm Generic Template for Stochastic Adaptive
Optimizers with Momentum
Input: Initial point E R base step size n > 0, constant > 0.
for t = 1 to T do
Obtain stochastic gradient (1s))
m/ = update ( (xz);; me-1)
Subset-Norm (SN): memory-efficient adaptive step- size algorithm with high-probability convergence guaran- tees for non-convex objectives under coordinate-wiss wise sub- gaussian noise. By unifying
of =
V (zz)): of
T++1 Zt 77 4
m.
{Update step}
end for
- 's and AdaGrad-Norm's 's analysis, we show that the SN adaptive step size (Algorithm 2) achieves improved dimensional dependence, while reducing the memory footprint from (dd) to roughly O(vd). On LLaMA models' pretrain- ing tasks, SN step sizes achieves better perplexity than coordinate-wise step size across a range of optimizers and model sizes, while using significantly less memory and introducing minimal additional hyperparameters. !
Table I. Update rules for common optimizers in the framework of Algorithm 1. We omit bias correction terms and numerical stahilizer for simplicity. Memory for optimizer state shown for model of size (and memory footprint) d.
Optimizer
Memory
Update Rules
mi 1 + B1) f(z))
Adam
my 2d
" Hanf + (1 (1)) 8 FZze))²
Bm +1 (2-3)
Lubspace-Momentuun (SM): A momentum compression method that applies momentum in a chosen subspace and SGD the orthogonal complement with high-probability convergence guarantees under sub-gaussian noise for non- convex smooth objectives. When combined with SN, for some selected dimensionk less than d,² our method (SNSM) reduces the memory footprint of Adam and Ada- Grad+momentum from 2d to k + vd (see Table 4) while delivers improved training speed and performance.
SGDm
"? ID
mi ((x))
AdaGrad-Coord
d
v( vi-1 + [[(m)))
m. ff(m))
RMSProp
d
v Bi (1- f(x))
f(()))
AdaGrad Norm
" li_ + /
m, {s(())
SGD
Empirical evaluations on LLaMA models from 60M to 1B parameters demonstrate that algorithms scale ef- fectively and attain better performances than existing op- timizers. Our proposed methods are simple to implement, require minimal additional hyperparameter tuning, and are compatible with modern distributed training frame- works like FSDP (Zhao et al., 2023; Rajbhandari et al., 2020). We provide an implementation in PyTorch at https: https://githul: com/ iimyttngg///a snn///nnntt sm.
more memory-efficient alternatives like SGD and RMSProp. We aim to design principled algorithms that achieve the best of both worlds: strong performance and memory-effii -efooy-efficiennt.
2.2. Assumptions and Notations
For our theoretical analysis, we consider the unconstrained non-convex stochastic optimization problem minger< f(x) where f: Rd - IR is the objective function. We assume ac- cess to an history independent, non-biased gradient estima- tor fzz) for any x E X, that is E Furthermore, we assume that f is an L-soooth function:
2. Preliminaries
fZ(z)
f(z).
2.1. Common Optimizers and Memory Footprint
Consider the generic template in Algorithm 1, which cap- tures a broad range of first-order optimizers that leverage either momentum or adaptive step sizes. Many standard optimizers can be represented within this framework by varying choices of momentum and adaptive step-size terms, as shown in Table 1. Generally, optimizers with higher memory requirements, such as Adam, tend to outperform
[ ff(ill < L || - all, for all v, E Rdd
Smoothness implies the following quadratic upperbound that will extensively utilize: for all x, y E IRd we have f(y) - f()) < (Vf()) x) + 4 lly 2||
Notations. Let us denote the i-th coordinate of a vector v E Rd. If a vector 2; is already indexed as part of a sequence of vectors (where I2 denotes the f-th update) then we use Zt, to denote z, 's i-th coordinate and 21, E RR to denote the indexing with respect to an ordered subset vp S
Although the subset size can be tuned (Section 5.3), we pro- vide heuristic in Section 3.2 that works effectively across model sizes, climinating the need for additional tuning.
Typically, & is chosen to be around d/4
2

---

Subset-Norm and Subspace-Momentum
[d] of size & where (1 g0,0 with Ty(k) denoting the k-th element of V. For gradients, we let Vaf(x) == of oæ, denote the partial derivative with respect to the i-th coordinate. Similarly, for stochastic gradients (ar), we let (ar) denotes its i-th coordinate. If a, b IR: then ab and a/ b denotes coordinate-wise multiplication and division, respectively i.e (ab); a&b; and (a/b); 0z/b,.
We compress the second moment adaptive step size by par- titioning parameters into subsets for which they share the same adaptive step size AdaGrad-Norm (McMahan & Streeter, 2010; Ward et al., 2019). Formally, we need to specify a partition function 1/- [d] - [c] that splits the d coordinates into c non-empty subsets W, = + - C [d]. where LI V, = [d]. For example, one can (G) = (j/c) mod & to get consecutive equipartitioned subsets V {ik: ik 1,..., ik 1)) for some subset-size A E N so that ke = d.3
Coordinate-wise sub-gaussian noise assumption. A ran- dom variable X is o-sub-Gaussiae (Vershynin, 2018) if
Given a stochastic gradient (1(z)) E Rd at time t for parameter 2t, denote ff.) E R to be the subset of the coordinates of the stochastic gradient with respect to the subset v, For example, given eb = (i/c) mod & as above, we have (Vv.J(ze)), = - Similarly, we can define Vv, E RI to be .es;
1
E [exp (A9x9)] s exp (A2x2) for all À such that |>] S
If we denote the stochastic gradient noise as & == (zx)) V((a)) and f as the i-th coordinate of ξr, then we assume the noise is per-coordinate subgaussian i.c. there exists a; > O for i E [d] such that &, satisfies
of(z.)
Now, we define the subset-norm (SN) adaptive step size b for subset V, and the update rule for #++1 (see Figure 2):
E exp (A²G2.)] S exp (APo7) _ Vi [d].
(1)
Df. = bi-1, + | f(x)) i
(2)
Note that ||& being o-subgaussian implies that each εt, is also o-subgaussian, so coordinate-wise sub-gaussian is more general than standard scalar sub-gaussian noise as- sumption. Furthermore, when ||-I| used without explic- itly specifying the norm, we assume it is the lz norm I1-112- We also use the O-indexing convention i.e. [n] == {0. 1..1......111 for integer n E N.
7 fz((z)), for E [d].
Te+1. = It.j bl,(()
Note that choosing c c= d and 1 recovers AdaGrad- Coordinate and AdaGrad-Norm, respectively. We now show a convergence guarantee on arbitrary partitions that will inform us on how to select good partition strategy.
-
3. Subset-Norm (SN) Adaptive Step Size
3.1. High-Probability Convergence of Subset-Norm
Coordinate
Subset-Norm (ours)
Norm
We show the following high-probabilit convergence result for the subset-norm adaptive step size:
Grales
Starc
Coddert
Stae
Grdden
Sax
Theorem 3.1. Suppose that f Rdd R is L-smooth and lower bounded by fa. Given unbiased stochastic gradients (xx)) with stochastic gradient noise & == fxx)
24
10
K2-2
65-4
26-2
fxxx) that is a--per-coordinate subgaussian for i [d]. For partitions of the parameters into E N+ disjoint subsets
Figure 2. AdaGrad variants: Coordinate, Subset-Norm, and Norm
[d] = = V with V,nV = 0., for i & j. the iterates In given by Algorithm 2 satisfies the following with probability at least 1 - (c0) failure probability 6 > 0)
Subset-Norm generalizes Coordinate (k: 1) and Norm (k = d).
Algorithm 2 SGD with Subset-Norm Adaptive Step Size
T T where
Σ 0
lov 1la
N(0) T
Input: Initial point æ1 E IRd, base step size >> 0, partition function v [d] + [c] that splits the coordinates into e subsets V,-
[V =C())
G(O)
(i) C [d]. where Ii v. [d]. and bo,: > 0 for
[c
for = to T do
Obtain stochastic gradient (Zv)
G(Ò) = Ô
Σ lov, 11 +|olll (11011++ +213) + cl.).
, = 1 + v f(xe) for i E (c)
(Update
accumulated gradient norms]
N(6) = 1101/2
Σ
|1 + Lc.
Pt+1k = Ft.k b()) (af(E)), for & E [d]
end for
Wee use this strategy all our implementations for simplicity.
3

---

Subset-Norm and Subspace-Momentum
Polylog terms are hidden in Theorem 3.1 for simplicity. The full result, Theorem E. and proofs are presented in Appendix E. Theorem 3.1 provides guarantee for all partitions of the parameters into arbitrary disjoint subsets and generalizes AdaGrad-Norm (c = 1) and AdaGrad- Coordinate (c = d) results. The result is noise-adapted: Ei= ||ov,11 ||2 2 is small enough, the rate becomes the opti- mal deterministic rate of O(}) regardless of the base step size 17. The next section explores implications of Theorem 3.1 and strategies for selecting subsets.
important for modern neural network, since the number of parameters d is typically greater than or on the same order as the total number of iterations T.
0000
12-00 7
0.30
-20-0
(1.00
40-00
0.20
13-01
0.00
22:0
0.30
16 06
Layer Indo:
3.2. Coordinate-Noise Density and Subset-Norm's Improved Dimensional Dependency
Figure 3. Noise density per parameter across layers for LLaMA 60M on pre-training task after 100 steps.
Theorem 3.1 presents trade-off's between the partition strate- gies and the stochastic gradient noise, where we need to bal- ance between the number of subsets с and noise-reduction benefits of paramiters-grouping e.g., ||=|||2 I ||~||11-
Subset-size heuristics to avoid additional hyperparame- ters. Providing a useful and robust default setting for an algorithm is important to justify claims of reduced costs. We provide a simple partitioning scheme for SN for 2D parameters p E Rmxx; we simply group along the smaller dimension. For example, if p is of shape (2048, 1024), we group by rows to get state of size 2048. This is a natural grouping scheme that groups the latent dimension together and aims for the rough ddt.45 subset size discussed in the pre- vious paragraph. Another simplification is that subset-norm is applied only on linear modules, since 2D linear modules makes up the vast majority of parameters in transformers. This means we compress all the attention, MLP, and final LM head weights. This implementation is presented in more details in Appendix A.2. Section 5.3 shows that this heuristic grouping, while simple, is not optimal and can be improved by tuning the subset size, but we opt for simplicity over performance- -taaannnntttttnnooeeeeae in our experiments.
Coordinate-noise density d=. To make the intuition above concrete, consider a scenario with various coordinate- noise density rate: fix a rate 3 E [0, 1). some d coordinates have noise a > while the rest are O. The rate 3 controls the density of coordinate noise. When 3 O, only 1 coordi- nate have noise. When B = 1, all coordinates have noise. To get a feel for 8's relationship to the fraction of coordinates containing noise, half the coordinates contain noise when 3 = 0.96 for d = 60M and / = 0.97 for d = 10B and 3 R 0.98 for d = 15 (See also Figure 10). See Figure 3 for noise density of LLaMA 60M (details in Appendix C.1). Furthermore, a upper bounds all coordinate noise, i.e. ||a|| S a, which is common in coordinate- analysis (Défossez et al., 2022).
Derivation of convergence rate given coordinate noise density d8. Given 8 E [0, 1]. we can obtain a concrete expression for the convergence rates of various methods (different subset sizes) from Theorem 3.1. For SGD with Subset-Norm, we consider an equipartition strategy, where we divide the coordinates into e = d subsets of size d /k each with the d noisy coordinates into just is subsets so that the rest of the c c- k subsets have no noisy coordi- nate. We defer the derivation details to Appendix C.2 and summarize the results in the first row of Table 2.
Generic Implementation. We provide pseudocode for the generic equipartition strategy of Algorithm 2 in Section A.3 that we use for the subset sizes ablations in Section 5.3.
Furthermore, in contrast to methods like AdaFactor or Ga- Lore that are limited to 2D parameters, the generic subset- norm algorithm is coordinate-wisg and admits easy im- plementation to FSDP (Zhao et al., 2023; Rajbhandari et al., 2020), where parameters are flattened to 1D tensors for efficient communication.
Subset Selection. In Table 2, the equal subset-size par-
4. Subspace-Momentumm
tition strategy for Subset-Norm has better dependency on the dimension d when the noise not completely sparse i.e. 3 = O. Hence, if we expect the actual noise density 3 to be around 0.75 to 0.90, then compressing with subset size of around o.5 to 0066 is optimal. The dependency on d is
Existing algorithmic compression approaches like GaLore (Zhao et al., 2024), GRASS (Muhamed et al., 2024), and FLORA (Hao et al., 2024) project the gradient to lower dimensional space & for updating the optimizer state via some bounded linear operator P : Rd - I² such that P'P Rd + Rd is a projection, i.c. (P* = P'P, P, where P R - Rd is the adjoint operator of P. More
"Figure 1I shows that overall noise is quite sparse but varies more when limited to particular layer as in Figure 3. See Section 5.3 for more experiments on subset size selection.
4

---

Subsct-Norm and Subspace-Momentum
Table Algorithms comparison between dimensional dependencies and convergence rates under different coordinate-nois density settings. Given a density rate 3. convergence rates' dimensional dependency are highlighted in red and green to denote the worst and best dependency on the dimension. Note that memory usage of CooodiaaCe ((d) while SGD with Subset-Norm (with the partition strategy presented here) is O(d/k). where & d' 22-45-00 0.6 AdaGrad-Coordinng
is chosen as optimal noise dependent subset size.
Density rate
AdaGrad-Nooe Norm
Subset-Norm (equipartition subsets)
(8.3+1.87// //T + d=*1 /r) if 3 [0, 2/3] 8.341.90 /VT + d 6040.0 /r 3 13/4 1]
B E [0, 1]
(
/VT /VT +
/T)
Ô (1 /VT (VT + /2)
0
d- IVT + 1T)
(1/T + 1/T)
8 //T/ 4/T)
3 0.5
Ô
( (&/VT + 8 (T)
(1 /TT VT + /
1 //T a sT//T))
8 = 0.9
Ô
(8 /VT + +8/11
8.20 VT + 8.
d VT + (804/))
B =
Ô
8
/TT + +
B /TT + (T)
B /VT + /T)
Algorithm SGD with Subspace Momentum (SM)
U
Projection P'P
Input: Bounded Lincar Operator P PR + RR such that P is a projection, where P is P's adjoint. for f = 1, .,2,.... ; T do
Eyto
somspant ())
SGD in t
Ohtain stochastic gradient f(.))
Monsntum
m. - + (1 B1 31)P f2)(())
(Momentum
subspace U = rowspan P))
Ti = (xx) p PPP ((x)
[Component in U SSeep in both spaces)
21+1 = 3t (p me + r1)
mm
end for
Figure Subspace Momentum Illustration.
U == rowspan(P) while performs SGD in U-, Unlike joint compression techniques like GaLore (4), SM only affects the momentum term. Hence, SM is modular and fits into the framework of Algorithm there, we can combine it with different adaptive step sizes such as subset-norm.
concretely, given a stochastic gradient A (wx)) E Rd at time t, low-dimensional version C₁- := PV (zz) E g* is
computed that is used to update the states before projecting back to Rd for update:
= 81m=- + (1 31
4.1. High-Probabilii Convergence of Subspace- Momentuum
of 820,-1 - 32) ot
(4)
T+41 = P* (me/v:).
We show that Piospacc-Moomntunne Algorithm 3, converges with high- brobbprrbbbility under the standard assumptions of smoothness and o-subgaussian gradient noise:
This update performs adaptive optimization in the row span U C Red of P when viewed as linear operator, with dUm = k. For example, Galore (Zhao et al., 2024) utilizes the top k singular vectors of stochastic gradients, and FLORA (Hao et al., 2024) simply projects to a random subspace using dense Gaussian matrices. Due to the opti- mization operating in a low rank subspace, convergence is not guaranteed unless stronger conditions are assumed.
Theorem 4.1. Suppose that f RRd + R is L-smooth and lower bounded by f.. Assume unbiased stochastic gradients (xe) with o-subgaussian stochastic gradient noise. Then, the iterutes vt given by SGD with Subspuce- Momeentun
(Algorithm 3) with step size 1 min
a=oT for
20
(3-8)L a = 21133 satisfies the following with probability at least
We propose Subspace Momentum (SM) presented in Al- gorithm 3 and illustrated in Figure 4 where SM guar- antees convergence by incorporating the orthogonal com- plement of Pa (xz)) E U that lives in the orthogo- nal complement U of U (with U @ U- RR). We can compute the orthogonal complement of fz( via E
T
8A1a
ToVaD
480 log (1/8)
Σ
+
T
T
VT
T
where A = f(x1) - fo is the initial function gap.
We observe that Theorem 4.1 has similar rate vanilla SGD. Unlike adaptive algorithms, we need to know the
Subspace Momentum maintains the same memory foot- print, Okk), as existing low-rank optimizers. However, SM's update step full rank: it uses momentum only in
Section B.9 contains detailed ablation different momen- tum and adaptive step sizes combinations.
5

---

Subset-Norm and Subspace-Momentum
problem parameters get an adaptive convergence rate. The proof is presented in Appendix D.2, where we also provide some intuition for the algorithm.
results and Table shows the memory footprint of different optimizers across a range of model sizes.
Additional Baselines. We provide additional comparisons with FLORA (Hao et al.. 2024), LoRA (Hu et al., 2021 1), and ReLoRA (Lialin et al., 2023) in Table 8. Note that these memory-efficient methods sacrifice performance (over Adam) to save memory while our method, AdamSNSM, achieves the best of both worlds.
4.2. Subspace Selection, Subspace Switching, and Projection Updates
In our experiments, we use the top-k singular vectors of a stochastic gradient snapshot as our main subspace, similarly to GaLore (Zhao et al., 2024).
5.1.1. DISCUSSIONS
Algorithm 3 and the accompanying theory in Section 4.1 are only for a fixed projection. However, from our experiments, we find that performing subspace switching every G steps (as in GaLore)) can be beneficial, especially for smaller ranks. Section B.7 contains an ablation studies on this. We incorporate projection updates in our main algorithms by picking a projection update gap and then fully resetting the momentum term to zero when we switch (in contrast to GaLore's accumulated statistics when switching subspace).
Subset-Norm (SN) improves upon all existing adaptive methods while reducing memory. Modifying Adam, Ada- Gradm, AdaGrad, and RMSProp with the SN adaptive step size not only reduces memory footprint but improves their performance across different scales. Notably, AdaGrad and AdaGradm benefit the most from the SN step size, pro- viding empirical support for the theoretical benefits of SN presented in Section 3.
(SM) with SN further improves performance while saving additional memory. Perhaps surprisingly. limiting the use of momentum to a subspace improves performance in SN-adaptive step sizes rather than degrading it. Our experiments show that SNSM, combining SN and SM, gives the best performance for the least amount of memory across model sizes. While adding SM introduces additional hyperparameters, Section B.7 sug- gests that these parameters are not too sensitive.
Combining
5. Experiments
We evaluate Subset-Norm (SN) and Subspace-Momentie (SM) on LLM pretraining and supervised fine-tuning tasks, where memory is often a bottleneck. We compare against several baselines, with memory estimates given for parame- ters of size m x n, where we assume WLOG m 2 n.
Baselines. We consider AdaGrad (Duchi et al., 2011), AdaGradm where we incorporate momentum 0.9 to Ada- Grad, Adam (Kingma& & Ba, 2014), and RMSProp (Tiele- man et al., 2012) as standard optimizers. We also consider GaLore (Zhao et al., 2024) as a recent memory-efficien method that projects the optimizer states into a low-rank subspace (typically rank n/4), using 2mmm44) memory but requiring 6 hyperparameters including subspace rank, pro- jection update frequency, and scaling parameters.
Furthermore, Section B.8 shows that the choice of the sub- space matters i.c. the subspace spanned by a top-k singular vectors of snapshot of a stochastic gradient seems to be the most beneficial for momentum as opposed to simpler choices like a random subspace. Our current guarantee for SM, presented in Section D, does not yet explain why or when subspace momentum is useful, and theoretical under- standing of (EMA style) momentum in stochastic optimiza- tion is still limited (Kidambi et al., 2018). We believe this could be related to how momentum is beneficial when noise is low (and harmful when noise is high) and the choice of the subspace could correlate to the amount of gradient noise or optimization landscape that harm or benefit momentum (Wang et al., 2024; Gitman et al., 2019).
Our methods. We incorporate SN and SM to AdaGrad, AdaGradm, Adam and RMSProp. SN reduces the adaptive step size (e.g. Adam's s second moment term) memory from mnn to m for parameter of size m X n. SNSM further compresses the momentum term of momentum methods like Adam and AdaGradm by adding SM with SVD at the cost of additional hyperparameters (See Algorithm 6 for the full implementation used in our experiments). RMSPropSN and AdaGradSN achieves minimal memory footprint of just m while requiring only hyperparameters.
Hyperparameter robustness. In Table 3, the best learning rate (LR) found via grid search is displayed and is high- lighted in red as the best LR changes across scales. This indicates potential sensitivity to tuning for each respective algorithm. We see that Adam requires smaller LR for larger models, but using SN and SNSM does not. AdaGradm seems less sensitive to the base LR overall.
5.1. LLM Pre-Training Experiments
We test our method on the task of pre-training LLaMA models (Dubey et al., 2024; Touvron et al., 2023) on the C4 dataset (Raffel et al., 2023) with a standard setup- - details in Appendix A.1. Table presents the main pre-training
'The memory footprint is the total parameters in the optimizer states multiplied by 16 bits. See Listing for more details.
6

---

Subset-Norm and Subspace-Momentum
Table 3. Final perplexity ("Perpl.") along with the number of tokens in parentheses of different optimizers on pretraining L.L.MA models task. Bolded methods are ours. Columns LR and #TP denote the learning rate and the number of tunable parameters of the corresponding method, respectively. We only tune for the base learning and set other parameters as in previous implementations. The memory column
shows the optimizer's ' states memory consumption given parameter of shape m. x with m 2 n. Red LR highlights instability.
Memory (for m x n)
60M
(1.38B)
130M (2.62B)
350M (7.86B)
1B (13.1B)
Methods
#TP
Perpl.
LR
Perpl.
LR
Perpl.
LR
Perpl.
LR
Adam
2mm mn + m
30.46
0.005
24.60
0005
18.67
0.001
16.00
0.0005
AdamSN
29.75
0.05
22.90
0.05
17.49
0.05
14.96
0.05
AdamSNSM
rn + m
5
29.74
0,05
22.43
0.05
16.91
0.05
14.05
0.05
AdaGradm
2mn
2
30.40
0.10
24.86
0.10
18,30
0.10
17.42
0.10
AdaGradmSN
mn + m
2
29.73
2.00
22.58
2,00
17. 17.14
2,00
14.48
2.00
AdaGradSNSM
+ m
4
29.81
1.00
22.43
1.00
16.99
13.96
AdaGrad
mn
1
37.12
0.05
25.76
0.05
18.14
0.05
15,25
0.01
AdaGradSN
29.85
2.00
24.19
1.00
17.72
14.82
RMSProp
mn
2
35.51
0.001
25.94
0.01
20.01
0.001
17.03
0.001
RMSPropSN
2
34.57
0.01
25.67
0.01
18.72
0.01
15.97
0.001
Gal (Zhao et al., 2024)
2rn
6
34.73
0.01
25.31
0.01
18.95
0.01
16.76
0.001
Rank r Dimension m
128/512
256/768
256/1024
512/2048
Table 4. Optimizer states memory footprint (in GB for BF16 dtype) for different LLaMA models. Our methods, AdamSN, AdamSNSM, and RMSPropSN (RMSPSN), modifications of Adam and RMSProp (RMSP) to utilize Subset-Norm (SN) and Subspace- MMoeett (SM). For GaLore and AdamSNSM, the subspace is of dimension d/r, where the memory accounts for additional space for storing the projection matrices.
Table 5. Last and minimum validation perplexity for SFT of LLaMA 7B on the UltraFeedback dataset between Adam, LoRA, and AdamSNSM for different ranks. We also show the wall- clock time and peak memory for batchsize I for these optimizers.
Adam
LoRA (1-64)
AdamSNSM (r-64)
SNSM (r=32)
Opt. AdamW
AdamSN
RMSP
GaLore
AdamSNSM
RMSPSN
Last
2.622
2.632
2,581
2,580
Mem. 2d
d vi
d
4d/r
24/r va
v2
2
2.410
2.892
2.390
60M
130M
0.22
0.14
0.11
0.25
0.15
0.08
0,03
Time (min.)
Memory
77.11
266
20.75
249
42.89
503
42.89
301
350M
1.37
0.75
0.69
0.53
0.28
0.06
IB
4.99
2.62
2.49
161
0.84
0,12
10.01
5,16
5.00
2.96
0.15
25.10
7.01
Discussion. We observe similar improvement over Adam as pre-training tasks. Surprisingly, the smaller rank (for momentum) is more beneficial than the larger rank. In contrast to LoRA, since report peak-memory here, due to the full parameter training of SNSM, the primary memory bottlenecks are gradients and activations. Furthermore, we note that the primary contributor to SNSM's slower wall clock time is the SVD computation on large dimension. We try larger projection update gaps in Table 6 which reduce this cost while maintaining good performance for our methods. Furthermore, we discuss potential more efficient alternatives in Section 5.3 and leave further exploration to future works.
Closing the theory-practice gap. While there is a non- trivial performance gap between Adam and AdaGrad(m) for larger models, using the SN step size closes this gap across scales. This shows that AdaGrad style algorithms can be competitive to Adam when using the SN step size. Interestingly, vanilla AdaGrad seems to perform well as model size increases. This is important because AdaGrad enjoys stronger theoretical understanding than Adam and has one fewer parameter- 32 - to tune.
5.2. LLMs Supervised Fine-Tuning (SFT) Experiments
GLUE Fine-tuning. Additional results on fine-tuning on GLUE tasks with BERT models are in Appendix B.4.
We further evaluate on a supervised firriieedd-----tinnggge task, where we fine-tune a pre-trained LLaMA 7B model on the U1- traFeedback dataset (Cui et al., 2024) using the chosen re- sponses with max sequence length of 1024. We train for epoch with linear decay and gradient clipping of 1. Table 5 contains the result with the time and memory of one training epoch on a single A100-80GB GPU. Note SNSM's r de- notes the dimension of SM but the optimization is full-rank.
5.3. Ablation Studies
In this section, we present ablation studies on various pa- rameters of SN and SM.
Subset-Norm's 's subset size ablation. While we use a simple scheme to compress the adaptive step size of linear
7

---

Subsct-Norm and Subspace-Momentum
AdunSN Subset Size va, Perplesity
Table 6. Effects of less frequent subspace update schedule (gap). Compared to Table 3 where the gap is fixed to 200 across all scales.
31.10
Aamm pporpaandiyy $1.001
Model Size
60M
130M
350M
1B
Pesplcatty
26.50
CapiSSpps (5%)
200V10K
1K/20K
3K/60K
5K/100K
32
AdamsN Lauuiia meladaacatatnal
AdamSNSM
29.84
22.71
18:43
15,28
20.30
AdalGralSNSM
30.28
22.76
17.02
13.90
GaL.ore
36.69
29.37
21.27
19.14
20:75
512
1024
Fixed Subspace
10K/I0K
20K/20K
60K/6UK
100K//00K
123
23
AdamSNSM
30.65
23:65
18.94
15.16
AdsGradSNNSM
31.43
24.85
14.62
Subsot Size (Jag soale)
Gal.ore
37.95
26.63
21.49
27.11
Figure 5. Subset size ablation for AdamSN on LLaMA 60M trained for 1.38B tokens (batch size of 512 of max length 256 for 10,000 steps). The higher the subset size, the smaller the memory footprint of the second moment optimizer state.
every 200 steps can be arbitrary. In Table 6, we examine more structured schedules: (1) updating every 5% of the total training steps (correspondiig to 200/10K steps for the 60M model) and (2) only using a fixed subspace at the start. Compared to Table, 3 where a fixed gap of 200 is used across scales, we see SNSM's performance stay relatively similar when we increase the update gap to 5% of the total training steps, whereas GaLore's 's performance suffers more.
modules in the previous experiments, Table 2 suggests that there is an optimal subset size that depends on the noise. Figure 5 shows performance for various subset-size selec- tion. Since the step size scales with the subset size, the optimal base LR should be decreased as we decrease the subset size closer towards Adam. We include additional results for 130M model in Figure 8.
While one can use the heuristics discussed on models where linear modules make up the vast majority, for arbitrary mod- els with weights of d elements, we found that subset size of v 2 is probably a reasonable choice. If more resources are available, the subset size can also be tuned.
Table 7. Fixed Subspace Choices on LLaMA 60M. We examine GaLore and SNSM with top-k singular vectors projections (SVD) and random subspaces (Random) using dense gaussian projections.
GaloreSVD
GaloreRandom
SNSM-SYD
SNSM-
Perpexitt
37.95
38.23
30.65
40.15
Subspace selection. While the top-k singular vectors of stochastic gradients gives subspace with strong perfor- mance, performing SVD can be expensive for larger models and storing the dense projection consumes non neglliiie memory for large ranks. Gradient-independemt projections like random gaussian as in FLORA (Hao et al., 2024) avoids SVD and can save memory by storing the pseudorandom seed (at the cost of recomputating the projection at every step). One can further speed up the random projection by using fast subspace embedding like the Subsampled- Randomized Hadamard Transform (SRHT) used in the Fast- JL transform (Ailon & Chazelle, 2009). Random projections like SRHT can also be used to approximate SVD (Appx- SVD) computation (Halko et al., 201 that can be much faster than full SVD. Finally, the cheapest projection is a subspace of random standard bases. Recently, GRASS (Muhamed et al., 2024) explores this idea and tests sam- pling random bases with large gradient norms. We examine different choices for the subspace and compare their time, space, and performance in Appendix B.8 (Table 13).
Interestingly, for fixed subspace (100% gap), GaLore still achieves decent performance even though the optimization only happens in a small subspace up until the 1B model, where the training stops improving after 50K/100K steps. In Table 7, we see that a random subspace seems to work decently well too. This suggests that a majority of progress can be made in small subspace in smaller models. In contrast, this is not the same for restricting momentum to subspace. Furthermore, we notice that there training loss spikes at the times when we switch subspace for Ga- Lore that impacts training with 5% gap, most likely due to incompatible optimizers' statistics between subspaces. This could explain why GaLore's 100% gap performs similarly or even better than 5% gap for certain run. Finally, we note that AdaGradSNSM performs the best here with the larger gaps as the dimension increases.
Additional experiments and ablations. We provide addi- tional experiments and detailed ablations on wall-clock time speedup. peak-memory savings, the effect of clipping, batch sizes, random seeds, combinations between adaptive step sizes and momentum, and more in Appendix B.
Larger projection update gaps. Frequently updating the projection map using SVD can be expensive. especially for larger models. Furthermore, updating the projection
8

---

Subset-Norm and Subspace-Momentum
6. Related Works
have highlighted fundamental limits in non-convex opti- mization. Here, obtaining convergence results for EMA updates (Adam style) for subset-norm and under further re- laxed assumptions like affine smoothness (Wang et al., 2023; Attia & Koren, 2023), affine noise (Hong & Lin, 2024; Faw et al., 2022), heavy-tailed noise (Zhang et al., 2019; 2020; Nguyen et al., 2023a;b) are of great interest.
As model sizes grow, memory-efficient training techniques have become crucial. Following up on AdaFactor (Shazeer & Stern, 2018), low-rank methods like Galore (Zhao et al., 2024), LoRA (Hao et al., 2024), and ReLORA (Lialin et al., 2023) reduce memory usage by approximating large weight matrices with low-rank representations. Projection-basee approaches, such as GRASS (Muhamed et al., 2024) and FLORA (Hao et al., 2024), compress gradients or combine low-rank ideas with projections to reduce memory require- ments. Recently, AdaMeM (Vyas et al., 2024a) proposes to incorporate the orthogonal subspace to the AdaFactor opti- mizer; this is related to but different from our simpler SM algorithms, where we use subspace decompositions to de- couple the momentum and SGD. BAdam (Luo et al., 2024), a block coordinate descent method that utilizes Adam as an inner solver, has been proposed for fine-tuning large lan- guage models. In contrast to our proposed methods, these methods are largely heuristic-driven and often lack conver- gence guarantees under standard assumptions. On the other hand, methods like SM3 (Anil et al., 2019), which uses sub- set (cover) statistics to show convergence in online learning, and MicroAdam (Modoranu et al., 2024), which provides convergence guarantees for a gradient compression scheme with error correction, offer theoretical guarantees.
Comparison with Adam-mini. Very recently, Adam-mini (Zhang al., 2024) also uses shared step sizes as Subset- Norm; however, the partition strategy is quite different from ours. While Adam -ini also employs a grouping strategy for the adaptive step size, it is primarily motivated empir- ically and lacks a general grouping strategy for general parameters. In contrast, our theory results show that group- ing by noise magnitude leads to improvement. In experi- ments, our AdamSNSM uses less memory than Adam-mini, due to the fact that Adam-mini uses full momentum while we use momentum only in a subspace (which outperforms full momentum in many cases given a good choice of sub- space). Furthermore, in terms of perplexity, Adam-mini performs very closely to Adam W while our methods out- perform Adam (which performs similarly to AdamW) on a range of language tasks and model sizes.
7. Conclusion and Future Works
Additional approaches reducing memory during training include optimizer quantization (Li et al., 2024a; Dettmers et al., 2021; 2024), attention computation compression/op- timization (Wu et al., 2022; Dao et al., 2022; Dao, 2023; Shah et al., 2024), activation checkpointing (Chen et al., 2016), and distributed training (Rajbhandari et al., 2020).
In this paper, we introduce two principled optimizer states' memory reduction methods -Subset-Norm (SN) and subspace-Mooennun
to address the high memory costs associated with adaptive optimizers in large-scale deep learning. SN and SM achieve memory savings without compromising performance and admit high- probability convergence guarantees under relaxed assump- tions. Extensive experiments pre-training and fine-tuning LLMs validate our methods' effectiveness and efficiency.
For inference, compression techniques are also actively be- ing explored (Sakr& & Khailany, 2024; Dettmers al., 2022;
Xiao et al., 2024; Lin et al., 2024; Frantar al., 2023). These are orthogonal directions to our work and can be combined. Another orthogonal direction is approximated second-order optimization, where one aims to approximate the Hessian preconditioner using only first-order informa- tion in order to achieve faster convergence. Some works in this area include (Gupta et al., 2018; Liu et al., 2023a; Vyas et al., 2024b). These methods typically demonstrate faster training but at the cost of super-linear memory and additional computational overhead.
Future works. Promising directions include exploring SN and SM on additional domains like Reinforcement Learning, where high memory and high noise are also bottlenecks when scaling up to large models. Fully generalizing our methods involves developing general projections (beyond matrix decomposition) for SM on higher order tensors for use on additional architectures like CNNs. A more in-depth investigation to more optimal subset partition strategies for SN is also an interesting open question, since our analysis in Section 3.2 only applies to equipartition subsets. Further theoretical understanding for SM's subspace dependency for improving the subspace selection for SM is desirable. Furthermore, the convergence of SNSM is still unknown. Finally, the benefits of momentum in stochastic optimization in general is still a mystery; using SM to study effects of momentum on particular subspaces could open doors to obtain more fine-grained understanding for why or when momentum helps.
Convergence analysis of non-convex optimization methods has seen significant progress, with recent works providing convergence proofs for adaptive algorithms like Adam (Li al., 2024b; Défossez et al., 2022). Numerous studies have explored convergence properties of various adaptive and stochastic gradient methods (Chen et al., 2018; Défossez
et al., 2022; Ene & Nguyen, 2021; Liu et al., 2023c;b; Ward et al., 2019; Zou et al., 2019; Reddi et al., 2018; Nesterov, 1983), while lower bound analyses (Arjevani et al., 2023)

---

Subset-Norm and Subspace-Momentun
Acknowledgement
Cutkosky, A. and Mehta, H. High-probabilig bounds for non-convex stochastic optimization with heavy tails. Ad- vances in Neural Information Processing Systems, 34: 4883-4895, 2021.
This work was supported by NSF CCF 231 1649, We thank Alina Ene, Themistoklis Haris, and Duy Nguyen for insight- ful discussions and Hieu Nguyen for assistance with the experiments. This work was completed in part using the Discovery cluster, supported by Northeastern University's Research Computing team.
Dao, Flashattention-2: Faster attention with bet- ter parallelism and work partitioning. arXiv preprint arXiv:2307.0800 2023.
Dao, T., Fu, D., Ermon, S., Rudra, A., and Ré, C. Flashat- tention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Process- ing Systems, 35:16344-16359, 2022.
Impact Statement
This paper presents work whose goal advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
DeepSeek-AN AI, Liu, A., Feng. B., Wang, B., Wang, B., Liu, B., Zhao, C., Dengr, C., Ruan, C., Dai, D., Guo. D., Yang, D., Chen, D., Ji, D., Li, E., Lin, F., Luo, F., Hao, G., Chen, G., Li, G., Zhang, H., Xu, H., Yang, H., Zhang, H., Ding, H., Xin, H., Gao, H., Li, H., Qu, H., Cai, J. L., Liang. J., Guo, J., Ni, J., Li, J., Chen, J., Yuan, J., Qiu, J., Song, J., Dong, K., Gao, K., Guan, K., Wang. L.., Zhang, L., Xu, L., Xia, L., Zhao, L., Zhang, L.,, Li, M., Wang, M., Zhang, M., Zhang, M., Tang, M., Li, M., Tian, N., Huang, P., Wang, P., Zhang, P., Zhu, Q., Chen, Q., Du, Q., Chen, R. J., Jin, R. L., Ge. R., Pan, R., Xu, R., Chen, R., Li, S. S., Lu, S., Zhou, S., Chen, S., Wu, S., Ye, S., Ma, S., Wang, S., Zhou, S., Yu, S., Zhou, S., Zheng, S., Wang, T., Pei, T., Yuan, T. Sun, T., Xiao, W. L., Zeng, W., An, W., Liu, W., Liang, W., Gao, W., Zhang, W., Li, X. Q., Jin, X., Wang, X., Bi, X., Liu, X., Wang, X., Shen, X., Chen, X., Chen, X., Nie, X., Sun, X., Wang, X., Liu, X., Xie, X., Yu, X., Song, X., Zhou, X., Yang, X., Lu, X., Su, X., Wu, Y., Li, Y. K., Wei, Y. X., Zhu, Y. X., Xu, Y., Huang. Y., Li, Y., Zhao, Y., Sun, Y., Li, Y., Wang, Y., Zheng, Y., Zhang, Y., Xiong, Y., Zhao, Y., He, Y., Tang, Y., Piao, Y., Dong. Y.. Tan. Y.. Liu, Y.. Wang. Y., Guo, Y., Zhu, Y., Wang, Y., Zou, Y., Zha, Y., Ma, Y., Yan, Y., You, Y., Liu, Y., Ren, Z. Z., Ren, Z., Sha, Z., Fu, Z., Huang, Z., Zhang, Z., Xie, Z.. Hao, Z.. Shao, Z., Wen, Z., Xu. Z., Zhang, Z., Li, Z., Wang. Z., Gu, Z.. Li, Z., and Xie, Z. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model, 2024. URL https://arxis: org abcp://aa/// 04434.
References
Ailon, N. and Chazelle, B. The fast johnson-lindensrrane transform and approximate nearest neighbors. SIAM Jour- nal on computing, 391)3302-332, 2009.
Anil, R., Gupta, V., Koren, T., and Singer, Y. Memory efficient adaptive optimization. Advances in Neural In- formation Processing Systems, 32, 2019.
Arjevani, Y., Carmon, Y., Duchi, J. C., Foster, D. J., Srebro, N., and Woodworth, B. Lower bounds for non-convex stochastic optimization. Mathematical Programming, 199 (1):165-214, 2023.
Attia, A. and Koren, T. Sgd with adagrad stepsizes: Full adaptivity with high probability to unknown parameters, unbounded gradients and affine variance. In Interna- fional Conference on Machine Learning, pp. 1147-1171. PMLR, 2023.
Chen, T., Xu, B., Zhang, C., and Guestrin, C. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604 06174,111111 2016.
Chen, X.,
Liu, S., Sun, R., and Hong, M. On the conver- gence of a class of adam-type algorithms for non-convex optimization. In International Conference on Learning Representations, 2018
Défossez, A., Bottou, L., Bach, F., and Usunier, N. A simple convergence proof of adam and adagrad. Transactions on Machine Learning Research, 2022.
Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-115 2023.
Dettmers, T., Lewis, M., Shleifer, S., and Zettlemoyer, L. bit optimizers via block-wise quantization. arXiv preprint arXiv:2 2021.
Cui, G., Yuan, L., Ding, N., Yao, G., He, B., Zhu, W., Ni, Y., Xie, G., Xie, R., Lin, Y., Liu, Z., and Sun, M. Ultrafeedback: Boosting language models with sealed ai feedback, 2024. URL https:/ // /arxiv...../ org/ abs/ 2310. 01377.
Dettmers, T., Lewis, M., Belkada, Y., and Zettlemoyer, L. Gpt3.int80: 8-bit matrix multiplication for transform- ers at scale. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in
10

---

Subset-Norm and Subspace-Momentum
Neural Information Processing Systems, volume 35, pp. 11818-30332. Curran Associates, Inc., 2022.
Halko, N., Martinsson, P.-G,, and Tropp, A. Finding
structure with randomness: Probabilisticc algorithms for constructing approximate matrix decompositions, SIAM review, 53(2):217-288 2011.
Dettmers, T., Pagnoni, A., Holtzman, A., and Zeulemoyer, L. Qlora: Efficient finetuning of quantized llms. Advances in Neural Information Processing Systems, 36, 2024.
Hao, Y., Cao, Y., and Mou, L. Flora: Low-rank adapters are secretly gradient compressors. arXiv preprint wxXi:2442.03229. 2024.
Ding, J., Ma, S., Dong, L.,, Zhang, X., Huang, S., Wang, W., Zheng, N., and Wei, F. Longnet: Scaling transformers to 1,000,000,000 tokens. arXiv preprint arXiv:2307.02486, 2023.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E.,, Casas, D. d. L.,, Hendricks, L.. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 2022.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenbom,
D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M.,, Heigold, G., Gelly, S., et al. An image is worth 16x 16 words: Transformers for image recognition at scale. International Conference on Learning Representations, 2020.
Hong, Y. and Lin, J. Revisiting convergence of adagrad with relaxed assumptions. arXiv preprint arXiv:2402.13799, 2024.
Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., and Chen, W. Lora: Low-rank adaptation of large language models. arXiv preprint urXiv:2106.0966 2021.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The Ilama 3 herd of models. arXiv preprint arXiv:2407.21788 2024.
Huang, D., Niles Weed, J., and Ward, R. Streaming k-pea: Efficient guarantees for oja's s algorithm, beyond rank-one updates. In Conference on Learning Theory, pp. 2463- 2498. PMLR, 2021.
Duchi, J., Hazan, E.,, and Singer, Y. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7). 2011.
Kidambi, R., Netrapalli, P., Jain, P., and Kakade, M. On the insufficiency of existing momentum schemes for stochastic optimization, 2018. URL https::/arxtt
Ene, A. and Nguyen, H. L. Adaptive and universal algo- rithms for variational inequalities with optimal conver- gence s. arXiv preprint arXiv:2010.07799, 2021.
org/ abs 1olbs/1800 11//111000.0 05591.
Faw, M., Tziotis, I., Caramanis, C., Mokhtari, A., Shakkot- tai, S., and Ward, R. The power of adaptivity in sgd: Self- tuning step sizes with unbounded gradients and affine variance. arXiv preprint arXiv:2202.05791, 2022.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Krizhevsky, A. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009. URL https: /wwww/wwww.///./////////.// cs toronto edu/ -rrii striz/learnigg features- 2000--200---000 TR. pdf.
Frantar, E., Ashkboos, S., Hoefler, T., and Alistarh, D. Gptq: Accurate post-training quantization for generative pre- trained transformers, 2023. URL https: ptttps://areee /aaxiv.. org/ aal//bb//22111 17323.
Li, B., Chen, J., and Zhu, J. Memory efficient optimiz- ers with 4-bit states. Advances in Neural Information Processing Systems, 36, 2024a.
Gitman, I., Lang, H.. Zhang, P., and Xiao, L. Understanding the role of momentum in stochastic gradient methods. Advances in Neural Information Processing Systems, 32, 2019.
Li, H., Rakhlin, A., and Jadbabaie, A. Convergence of adam under relaxed assumptions. Advances in Neural Information Processing Systems, 36, 2024b.
Gorbunov, E., Danilova, M., and Gasnikov, A. Stochas- tic optimization with heavy-tailed noise via accelerated gradient clipping. Advances in Neural Information Pro- cessing Systems, 33: 13:15042-55053, 2020.
Li, S. and Liu, Y. High probability guarantees for noncon- vex stochastic gradient descent with heavy tails. In In- ternational Conference on Machine Learning, pp. 12931 12963. PMLR, 2022.
Gupta, V. Koren, T., and Singer, Y. Shampoo: Pre- conditioned stochastic tensor optimization. Interna- tional Conference on Machine Learning, pp. 1842-1850. PMLR, 2018
Lialin, V., Muckatira, S., Shivagunde, N., and Rumshisky, A. Relora: High-rank training through low-rank updates. The Tweifth International Conference on Learning Representations, 2023.

---

Subset-Norm and Subspace-Momentum
Liang, K., Liu, B., Chen, L.,, and Liu, Q. Memory-efficient Ilm training with online subspace descent. arXiv preprint arXiv:2408. 12857, 2024.
Nguyen, T. D., Nguyen, T. H., Ene, A., and Nguyen, H. L. High probability convergence of clipped-sgd under heavy- tailed noise. arXiv preprint arXiv: 05437, 2023b.
Lin, J., Tang, J., Tang, H., Yang, S., Chen, W.-M., Wang, W.-C., Xiao, G., Dang, X., Gan, C., and Han, S. Awq: Activation-aware weight quantization for IIm compres- sion and acceleration, 2024, URL https://aare org/ abs/ 2306. 00978.
Raffel, C.,, Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Explor- ing the limits of transfer learning with a unified text-to- text transformer, 2023. URL https: hatpps///axiiii //tpxxxxx.or//////////// abs /1t.0190...20 106833........
Liu, H., Li, Z., Hall, D., Liang. P., and Ma, T. Sophia: A scalable stochastic second-order optimizer for language model pre-training. arXiv preprint arXiv:2305. 2023a.
Rajbhandari, S., Rasley, J., Ruwase, O., and He, Y. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Per- formance Computing, Networking, Storage and Analysis, pp. 1-16. IEEE, 2020.
Liu, Z., Nguyen, T. D., Ene, A., and Nguyen, H. On the con- vergence of adagrad(norm) on I mathhb (IR"dd 'S: Beyond convexity, non-asymptotic rate and acceleration. In The Eleventh International Conference on Learning Represen- tations, 2023b. URL hitps: // iopenreviewwnet.. forum? d:0.ddd=2.HH.........
Reddi, S. J., Kale, S., and Kumar, S. On the convergence of adam and beyond. In International Conference on Learning Representations, 2018.
Sakr, C. and Khailany, B. Espace: Dimensionality reduction of activations for model compression. arXiv preprint arXiv:2410.05433 2024.
Liu, Z.,, Nguyen, T. D., Nguyen, T. H., Ene, A., and Nguyen, H. High probability convergence of stochastic gradient methods. In International Conference on Machine Learn- ing, pp. 21884-21914 PMLR, 2023c.
Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., and Dao, T. Flashattention-3: Fast and accurate atten- tion with asynchrony and low-precision. arXiv preprint arXiv:2407.0880 2024.
Luo, Q. Yu, H., and Li, X. Badam: memory efficient full parameter optimization method for large language mod- els, 2024. URL https: httx//aivvvv//// org/ abs 2404. 02827.
Shazeer, N. and Stern, M. Adafactor: Adaptive learning rates with sublinear memory cost. International Con- ference on Machine Learning, pp. 4596-4604. PMLR, 2018.
McMahan. H. B. and Streeter, M. J. Adaptive bound opti- mization for online convex optimization. In Conference on Learning Theory (COLT), pp. 244-256. Omnipress, 2010.
Tieleman, T., Hinton, G., et al. Lecture 6.5-rmsprop: Divide the gradient by rumning average of its recent magnitude. COURSERA: Neural networks for machine learning, 4 (2):26-31,200 2012.
Modoranu, L.-V., Safaryan, M. Malinovsky, G., Kurtic, E., Robert, T., Richtarik, P., and Alistarh, D. Microadam: Ac- curate adaptive optimization with low space overhead and provable convergence. arXiv preprint arXiv:2405. 15593, 2024.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation lan- guage models. arXiv preprint arXiv:2302. /397/, 2023.
Vershynin, R. digh-ddinn probability: An introduc- tion with applications data science, volume 47. Cam- bridge university press, 2018.
Muhamed, A., Li, O., Woodruff, D., Diab, M., and Smith, V. Grass: Compute efficient low-memory Ilm train- ing with structured sparse gradients. arXiv preprint arXiv:2406.176660 2024.
Vyas, N., Morwani, D., and Kakade, S. M. Adamem: Mem- ory efficient momentum for adafactor. In 2nd Work- shop on Advancing Neural Network Training: Com- putational Efficiency, Scalability, and Resource Opti- mization (WANT@ICML. 2024), 2024a. URL https: // //openreview..ommmmmmee net wwtetee/oruueeee forum? fetpummmidffffffiii f=
Nesterov, Y. A method for unconstrained convex minimiza- tion problem with the rate of convergence o (1/k* 2). In Doklady an ussr, volume 269, pp. 543-547, 1983.
Nguyen, T. D., Nguyen, T. H., Ene, A., and Nguyen, H. Improved convergence in high probability of clipped gra- dient methods with heavy tailed noise. Advances in Neu- ral Information Processing Systems, 36:24191-24222, 2023a.
Vyas, N., Morwani, D.. Zhao, R., Shapira, I., Brandfon- brener, D., Janson, L., and Kakade, S. Soap: Improving and stabilizing shampoo using adam. arXiv preprint arXiv:2409. /32/, 2024b.
12

---

Subsct-Norm and Subspace-Momentum
Wang, B., Zhang, H., Ma, Z.,, and Chen, W. Convergence of adagrad for non-convex objectives: Simple proofs and relaxed assumptions. In The Thirty Sixth Annual Conference on Learning Theory, pp. 161-190. PMLR, 2023.
Pytorch fsdp: experiences scaling fully sharded data
parallel. arXiv preprins Xixx:3304.11277 2023.
Zou, F., Shen, L., Jie, Z., Zhang, W., and Liu, W. A sufficient condition for convergences of adam and rmsprop. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11127-11135, 2019.
Wang, R., Malladi, S., Wang, T., Lyu, K., and Li, Z. The marginal value of momentum for small learning rate SGD. In The Twelfth International Conference on Learning Representations, 2024. https: httpttttst////eeeeeeeeeee //ppereeeee......//// net/ wtr/m//ruum///// i (i3
Ward, R., Wu, X., and Bottou, L. Adagrad stepsizes: Sharp convergence over nonconvex landscapes. In Interna- tional Conference on Machine Learning, pp. 6677-6686. PMLR, 2019.
Workshop, B., Scao, T. L., Fan, A., Akiki, C., Pavlick, E., Ilié, S., Hesslow, D., Castagné, R., Luccioni, A. S., Yvon, F., et al. Bloom: A 176b-parameter open-access multilin- gual language model. arXiv preprint arXiv:2211.05111 2022.
Wu, Y., Rabe, M. N., Hutchins, D., and Szegedy, C. Mem- orizing transformers. arXiv preprint arXiv:2203.08911 2022.
Xiao, G., Lin, J., Seznee, M., Wu, H., Demouth, J., and Han, S. Smoothquant: Accurate and efficient post-training quantization for large language models, 2024. URL https:// /wtp//aaavvvv/vaa///a//eaa/a.aaa wweeiii.rgg//eo orgg/bs//////eooooooooooo a//aabbb/2211 10438.
Zhang, J., He, T., Sra, S., and Jadbabaie, A. Why gradient clipping accelerates training: A theoretical justification for adaptivity. arXiv preprint arXiv: /905. 1881, 2019.
Zhang, J., Karimireddy, S. P., Veit, A., Kim, S., Reddi, S., Kumar, S., and Sra, S. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems (NeurIPS), 33:15383-1533 2020.
Zhang, S., Roller, S., Goyal, N., Artetxe, M., Chen, M., Chen, S. Dewan, C., Diab, M., Li, X., Lin, X. V., al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01060 2022.
Zhang, Y., Chen, C., Li, Z., Ding, T., Wu, C., Kingma, D. P., Ye, Y., Luo, Z.-Q., and Sun, R. Adam-mini: Use fewer learning rates to gain more, 2024. URL https: //arxiv. org/ abs/ /2////46..11 16793.
Zhao, J., Zhang, 7.., Chen, B., Wang, 7.,, Anandkumar, A., and Tian, Y. Galore: Memory-effiiien Ilm train- ing by gradient low-rank projection. arXiv preprint 2003002403.03507, 2024.
Zhao, Y. Gu, A., Varma, R., Luo, L., Huang, C.-C,, Xu, M., Wright, L., Shojanazeri, H., Ott, M., Shleifer, S., et al.
13

---

Subset-Norm and Subspace-Momentum
Algorithm 4 Adam-Subset-Nomm with a Simple Partitioning Scheme
Input: Learning rate '/. EMA parameters 31 and 32, e > O, optional weight decay wd 2
for each p E IR mB= in params do
grad + p.grad
< if m
else 1
k 4 p.shape(r]
{Set k = m if 0, else k: - n
gradN + grad.norm(dim-1 - r) E R
{ {Subset norm]
m - f1m + (1 - B1) grad E 1Rmx
v - 320 + +1 (1- 11-Ba 32) gradN" E IR*
{Omitting bias correction terms]
p - p o
{Broadcast division]
p - p 17- wd
(Weight decay)
end for
A. Experimental Details
In this section, we provide hyperparameters details, implementation details (pseudocode), and other practical considerations.
A.1. Experimental details
All of our pre-training experiments are conducted on NVIDIA RTX4090/3090 GPUs. Unless specified otherwise, we run all experiments on BF16 format, weight decay of O, gradient clipping of 1.0, cosine learning rate decay to 10% of the max learning rate with 10% linear warmup steps, and batch size of 512 (similarly to (Zhao et al., 2024) and (Touvron et al., 2023: Dubey et al., 2024)).? For all our experiments, we use the default (61, [2) = (0.9, 0.999) and only tune for the base learning rate within grid of {0.5, 0.1, 0-05, 0.01, 0.005, 0.001). 8 We train for 1.38B, 2.62B, 7.86B, and 13.1B tokens for models of sizes 60M, 130M, 350M, and 1B parameters, respectively, following (Zhao al., 2024) and matches roughly the scaling laws in (Hoffmann et al., 2022).
For GaLore, we use the same hyperparameters as in (Zhao et al., 2024), where we use rank 128/512, 256/768, 256/1024, and 512/2048 for the 60M, 130M, 350M, and 1B models, respectively (Table 2 of (Zhao et al., 2024)). ? For AdamSNSM, we use the same ranks and projection update gap (of 200) as GaLore for all models." 10 However, we do not tune for an additional scaling parameter unlike GaLore, and we compresses the LM head (final linear layer) with SN and SM also."
A.2. Maam-Subet-Norm Implementation
Algorithm 4 presents the pseudocode for Adaam-Suset-Noomd as mentioned in Section 3.2 where we partition the coordinates (for each parameter) into subsets of equal sizes.
A.3. Generic Subset-Norm Adaptive Step Size Implementation
The heuristic implementation in Section A.2 is simple and does not require any tuning. However, to modify existing algorithms to work with arbitrary subsets, one could utilize reshape as in Algorithm 5 as an example.
A.4. AdamSNSM Implementation Details
Algorithm 6 provides the pseudocode and implementation details for the version of AdamSNSM with SVD subspace momentum and heuristics subset-norm (as described in Section 3.2) used our experiments.
7 Note that these addition improve the performance for baselines. See Appendix B.
Except AdaGradSNm where we find higher learning rates in (0.5. 1, 2, 5} to he better. We tune the Ir on the 60M model and use the same learning rate for the larger model, where the hase learning rate is only reduced if the method fails converge.
"Note that our reproduced results for GaLore and baselines are similar to (Zhao et al., 2024).
Note that smaller gap is more expensive than larger gap. Our experiments below show that we can increase the projection update gap without much performance loss. If data not limited, one could use a larger gap to speed up training. However, if data is limited, then a smaller gap to converge in fewer tokens is potentially more desirable.
Il Existing methods typically do not compress the emhedding layer and final LM head, while our methods seem robust to this choice. Compressing these layers save additional memory.
14

---

Subsct-Norm and Subspace-Momentum
Algorithm 5 Generic Subset-Norm Adaptive Step Size Update Rule (PyTorch-y Notation)
Input: Parameter P E R2 step size > 0, 3, > 0, and partition size k such that & divides d
R +
d/k: k)
{Reshape gradient into shape & x k} {Update state V via subset norm reduction dim 1 [Broadcast addition and division for update step) {Reshape U back to R and update P}
V + @V + (1 ( Rd/k
U E R E Rd/kxk
V++c
P P U.iie((d
Algorithm AdamSNSM with Subspace Momentum via Top-k Singular Vectors from SVD
Input: Learning rate 1), rank k, update gap G, momentum parameters $1.32 E (0, 1), and stability parameter for č = 1,..., I do
Obtain stochastic gradient E Rm*x then
(WLOG, assume m 2 n)
if mod
U.S, V (VD(g))
[Compute singular value decomposition) {Extract top k singular vectors)
P U[; k] Rmnk
end if
3.m + (1 - PT) It E R
{Update subspace momentum,
91 -PPp
{Compute orthogonal SGD component {Sum all columns for subset-norm heuristic {EMA of subset-norm,
sum(gr, dim = 1) E IR
G = 822 + - 32)s² E R"
21 .t-1 + +7 Pm-r lB+c
{Update with subspace momentum and subset-norm step size)
end for
A.5. Measuring Memory Footprint of Optimizers
In PyTorch, obtain the number of parameters in optimizer states using the code in Listing 1.
A.6. Peak memory measurement during training for different optimizers
We measure peak memory consumption directly via running nvidia-smi in Figure 6 while training as oppose controlled measurement as in Table 4. Note that these peak measurements incur additional memory from gradient computation and algorithms' overhead.
def seteiiiiirr_sttt__stteeee (optimizer) -> Mmplelint, Dict (str, total_state_ stieettttteeeee_____e___e_eeeeet state_ breakdown () for group in for p in groupl'parane
int]]:
state optimizer..tateelp
for state_key, state_valus in state.items ():
if torch.is_tensen (state_value) value) if state_value.numell F we do not count singleton continuc lotal_state_izz if state_kes not state_sized sire_size_bbte wiw._breaad..o.n
== 1:
10
12
state_value numel()
13
14
saate_size_bbrekddnn [state_key]
state_si be_akkkooown [state_ keyl +- atate_ leellve.mee nume"
return total_state state size breakdown
Listing 1. PyTorch function to calculate optimizer state size
15

---

Subsct-Norm and Subspace-Momentum
GPU Memory Usage Across Model Sizes Optimizers
25 GPU Memory Limil (24 GB RTX-4090
17.4
Memory (GB)
15.8
16
RMSPropSN
AdamSNSM
Galore
11.8
SGDm
9,5
AdamSN
Adam
6.7
2.8
3.1
3.3
2.0
150m
350m
Model Size
Figure Peak Memory Usage (Gb) various model sizes, obtained with batch size and activation checkpointing to measure the optimizer state footprint.
B. Additional Experiments and Ablation Studies
B.1. Additional Comparisons with Memory Efficient Optimizers for Pre-Training LLaMA Models
Table 8 extends Table 3 to compare against other recent methods.
Table 8. Additional comparisons with other memory efficient optimizers
Method
LLaMA 60M
LLaMA 130M
LLaMA 350M
LLaMA 1B
AdamW
30.46
24.60
18.67
16.00
AdamSNSM (ours)
29.74
22.43
16.91
13.96
GaLore (Zhao et al., 2024)
34.73
25.31
20.51
16.76
FLORA (Hao et al., 2024)
32.52
23.69
LoRA (Hu et al., 2021)
34.99
33.92
25.58
19.21
ReLoRA (Lialin et al., 2023)
37.04
29.37
29.08
18.33
B.2. Wall-clock speedup and peak memory
We provide the per iteration time, peak memory (via nvidia-smi), and time Adam's val perplexity after 100K steps for the 1B model for each method on a 2x4090 machine with the same setup as in Table 3 (seq length 256, total batch size 512, miero batchsize 16) in Table 9.
Table Per iteration time, peak memory (via nvidia-smi), and time to Adam's val perplexity after 100K steps for the 1B model for each method on a 2x4090 machine with the same setup as in Table 3 (seq length 256, total batch size 512, miero batchsize 16) in Table 9
Adam
AdamSNSM (Gap=500n)
AdamSNSM (Gap=200)
AdamSN
GaLare (Gap=200)
Time for IK iters
7426
7465 s
7624 s
7399
7827
Time per iteration
7.43 xii
7.47
7.62
7.39 sñi
7.83 xfi
Time perplexity 16
-0066 4 hrs (100K
11366 hrs (<66K itens)
101 (<48K iters)
~- 18.9 (<558K iters)
>-17 brs 100K iters)
Peak mom
21.554 GB/GPU
16.642 GBIGPU
16.642 GBIGPU
193193 GB/GPU
18.187 GB/GPU
16

---

Subset-Norm and Subspace-Momentum
B.3. Vision Tasks
Diffusion Transtormers. While our main focus is on large models that are more typical language models where memory is often bottleneck, vision models also increasing in size. Hence, conduct further evaluations using the DiT-L/2 model (458M) on setup with batch size 2048, image size 64, and 8xA6000 GPUs. We compared our method (SNSM) with Adam. As shown in Table 10, SNSM outperforms Adam in FID similarly to LLM tasks.
Table 10. FID scores over training iterations for the DiT-L/2 model (458M parameters) on 64 x 64 images.
FID
Iter
200k
300k
400k
500k
700k
Adam
56.69
56.63
40.69
41.15
39.61
AdamSNSM
66.76
66.31
34.05
32.31
32.26
CIFARIO and CIFAR100. We further evaluate Adam, AdamSN, and AdamSNSM (rank 64 and no update gap) by training vit basee -pttccheee www.bbaeseeaachhheeeeee 6_224 13 (Dosovitskiy et al., 2020) (around 85M params) from the t timn library on the CIFARIO and CIFAR100 (Krizhevsky, 2009) datasets for 10 epochs with a batch size 64 and weight decay on 2x4090 machine. tune the across (1e-3, 5e-3, le-4, 5e-4, 5e-5, le-5) grid for all methods. The results are shown in Table 11.
Table II. Performance Comparison of Optimizers on Vision Transformers CIFARIO and CIFARI00
Best val accuracy (10 epochs)
Adam
AdamSN
AdamSNSM (r=64, g=1000)
CIFARIO0 100
43.30%
45.20%
45.60%
CIFAR10
69.02%
69.18%
71.21%
Peak Mem (bs 64)
9.288GB
8.886GB
8.878GB
These preliminary experiments show promising results for the application of our methods to vision tasks where models are becoming larger.
B.4. Fine-tuning on GLUE Tasks
Table 12 presents results for fine-tuning on GLUE dataset for various methods. The SN step size maintains good performance while reducing the memory footprint.
B.5. AdaGrad, AdaGrad-Norm, and AdaGrad-Suhset-Nome
We examine the subset-norm step size for AdaGrad in Figure 7. We again see that subset- norm is slightly better than the full coordinate version while using lot less memory. This is consistent with our observations for Adam and RMSProp when we replace the standard coordinate step size with the subset-norm adaptive step size.
phttps://huugiiggtaa http/timm/vvee http:///iit_aaeppttttttee 224. aeg.aaugee__ia. in21k ft inik
"nttps://timm. fast.
Table Performance metrics across GLUE tasks. QQP, RTE, SST-2, MRPC, STSB, QNLI, and MNLI use accuracy as the metric, while CoLA uses the Matthews correlation coefficient. The best and runner-up results for each task and the average score highlighted.
Method
QQP
RTE
SST2
MRPC
STSB
QNLI
MNLI
COLA
Avg
Adam
92.0
77.9
94.9
89.2
90.5
93.0
87.6
65.4
86.3
GaLore (r 4)
90.9
79.4
95.2
88.7
90.8
92.4
86.9
61.9
85.8
RMSProp
91.9
79.4
95.2
91.4
90.3
92.8
87.6
65.1
86.7
RMSPropSN
91.9
80.1
95.1
90.0
90.7
93.1
87.5
63.8
86.5
AdamSN
91.2
74.4
94.5
89.5
90.4
92.0
86.7
64.4
85.4
17

---

Subset-Norm and Subspace-Momentum
Training (Smoolhed)
Evaluation Perplexity
800
Add (am=
Adatirad Susset-Nom Man-
Acstimesooodrwie Mem=O(d)
Acstind-Comiine Mem=O(tt
AdaSrod Norm
MemrO11)
800
AdaGraa Nom:
Menm011)
400
Perpledly
800
200
100
0
2000
4000
6000
3000
10000
2000
4000
5000
8000
10000
Updale Step
Updalee Step
Figure 7. Pretraining LLaMA 60M on the C4 dataset for AdaGrad variants. Memory consumption estimate as a function of parameter count is shown in legend.
AdamSNSM Projection Cap Ablation
AdamSN Subset Size Perplexity for LLaMA 130M
AdemNNum
33.0
Qperpienity: 30,30
512
24
Perglesty
51.5
Adas
256
X244
31.0
64
1.26
24
38.5
AdamSN heuristic
100
2000
7010
Dlog
Update Projection Gap
Figure 8. Subset size ablation for AdamSN on LLaMA 130M trained for 2.62B tokens (batch size of 512 of max length 256 for 20,000 steps). The higher the subset size, the smaller the memory fontprint of the second moment optimizer state.
Figure 9. Rank and gap ablation for AdamSNSM on LLaMA 60M for 10,000 steps. The lower the rank, the less memory consumption used by the momentum state. The higher the projection gap, the less SVD computation is performed which is cheaper.
B.6. Additional Subset-Size Experiments for 130M model
We provide additional subset-size experiments similar to the ones in Section 5.3 for LLaMA 130M in Figure 8.
B.7. Subspace--Momentum Rank and Gap Ablations
Rank and gap ablations. We examine the impact of varying rank and update gap of subspace momentum, similarly to (Zhao et al., 2024), in Figure There, we see that the higher the rank, the better the results. For the update gap, seems like there is an optimal choice. However, due to the SVD computation, larger gap will be cheaper than more frequent gap.
B.8. Subspace-Momentum Projection Choice Ablations
Projection types. Table 13 tests different choices for projection in SM discussed Section 5.3. Note that for memory storage, SVD, Random Projection via dense Gaussian projection (Gaussian). and Approximated-SWV SVD (Appx-SVD) need store the projection matrix (unless recompute every step). The remaining methods only need to store the indices for sampling and/or the random seed to regenerate any random choices.
Note that the choice of the projection is important as some projections are more computationally and memory expensive than other, although trading other qualities for given the cost. Simple projections like selecting a subset of coordinates for momentum (Subset-Mone Momeentum) are not only faster but enables simple distributed training like FSDP unlike more complex subspace selection mechanism that requires additional priors about the parameters (shape, low-rank, etc.) that might not
18

---

Subsct-Norm and Subspace-Momentum
Table 13. Different projections selection for Suhspace-Momentum and validation perplexity. All methods are evaluated L.L.MA 60M with rank 128/5 128/512 and a projection update gap of 200. Time and space rows denote time and space to compute and store the projection.
AdamSNSM's s projection
SVD
Approx-VVD
Gaussian
SRHT
Top-k
Random-rows
OPCA
Oja
Time m Space (for rank k)
O(mn²) (lkm)
(mmn log kn?)
O(kn)
((max (m. n))
O(mn)
O(k)
O(ka)
O(kn)
(kk)
((km)
O(k)
O(k)
((k)
((kmn
((km)
Validation Perplexity
29.74
31.51
42.48
33.33
31.42
33.17
29.63
30.69
always satisfied.
Online k-PCA and Streaming k-PCA for Up-to-date Subspace. Computing subspace from stochastic gradient snapshots can be noisy. Recently, (Liang et al., 2024) proposes a formulation of online-PCA to handle the problem of staled top-k components as the stochastic gradients evolve. We test this algorithm in the OPCA column. Another natural algorithm to ensure the top-k components stay up-to-date is Oja's algorithm for streaming k-PCA (Huang et al., 2021). We also test this algorithm in Table 13. While we can maintain up to date projection using these schemes, more frequent updates suffer from the same issue of transferring optimization statisties from one subspace to another. We only test for not resetting the statisties in this setting and leave additional investigation for future works. Furthermore, these schemes are more expensive computationally due to additional computation requirement at every step. OPCA further uses Adam for inner optimization which incurs additional memory.
B.9. Step Sizes and Momentum Choices Full Ablations
We investigate various combinations of momentum and adaptive step size approaches in Table 14. For adaptive methods, we compare EMA, which uses exponential moving average to accumulate the second moment (u² = . + (1 - -B)g?), ) with AdaGrad's 's cumulative accumulation approach (b? = 6b_ + g Methods with the SN suffix utilize subset norm for parameter grouping, contrasting with per-coordinate approaches that are standard, While EMA momentum follows the standard momentum implementation, subspace momentum employs reduced rank approximation with rank 128 for this model size.
Table 14, Different combinations of momentum (columns) and adaptive step-size (rows) and the effect of the learning rate schedule on each combination (cosine learning rate decay schedule with warmup "coslr" or constant learning rate "Ir"). Memory footprint for each adaptive step size and/or momentum are shown. Green and red highlight runs with perplexity below 30 and above 50 respectively.
Final eval perplexity (Ir)
No momentum
EMA momentum
Subspace momentum
LLaMA 60M for I.3IB tokens
Mcm=0
Mem = m n
Mem = max max(m, n) rank
SGD
SGDm
SGD+SM
No Adaptive Step-size
86.60 (cosIr=le-3)
55.76 (cosIr=le-3)
89.97 (cosIr=le-3)
Mem = 0
100.04 (Ir=1.0)
56.07 (Ir=1.0)
213.21 (lr=5e-4)
RMSProp
Adam
AdamSM
EMA Coordinate
35.01 (cosIr=le-3)
30.46 (cosIr=5e-3)
32.34 (coslr=le-3)
Mem
36.46 (Ir=5e-4)
33.47 (Ir=1e-2)
32.25 (lr=5e-4)
RMSPropSN
AdamSN
AdamSNSM
EMA Subset-Norm
34.86 (cosIr=le-2) le-2)
29.75 (coslr=5e-2)
29.74 (coslr=5e-2)
Mem max(m,n)
34.57 (Ir=le-2)
33.69 (Ir=le-2)
32.49 (Ir=le-2)
AdaGrad
AdaGradm
AdaGradSM
AdaGrad Coordinate
37.12 (coslr=5e-3)
31.48 (cosIr=5e-2)
30.99 (coslr=5e-2)
Mem = m n
46.47 (l=55e-4)
43.99 (Ir=le-2)
41.32 (Ir=5e-4)
AdaGradSN
AdaGradSNm
AdaGradSNSM
AdaGrad Subset-Norm
33.19 (cosIr=5e-3)
29.73 (cosIr=5e-3)
29.81 (coslr=5e-3)
Mem = max(m, n)
41.23 (I==0.1)
44.98 (lr=0.1)
40.11 (lr=0.1)
19

---

Subset-Norm and Subspace-Momentum
Method
60M (no clipping)
60M (with clipping)
130M (no clipping)
130M (with clipping)
Adam
30.58
30.46
25.07
25.07
AdamSN
30.06
29.75
22.89 23.54
Gal.ore
34.91
34.73
25.43
25.31
Table 15. Pre-training LLMs ablation experiments for gradient clipping. We compare validation perplexity between LLaMA 60M and 130M with and without clipping. We use the same hyperparameters Section A.1 but just add clipping..
Table 16. Batch size ablation for various optimizers along with optimal learning rate.
Adam
GaLore
AdamSN
AdamSNSM
Batch size
Perpl.
LR
Perpl.
LR
Perpl.
LR
Perpl.
LR
1024
27.94
0.005
32.75
0.01
27.68
0.05
28.02
0.05
512
30.46
0.005
34.73
0.01
29.75
0.05
29.74
0.05
256
36.65
0.001
44.71
0.001
37.03
0.001
32.82
0.05
128
41.72
0.001
49.75
0.001
42.04
0.001
36.82
0.05
Discussions. From Table 14, Subset norm (SN) step sizes consistently outperform coordinate- wiste implementations while requiring less memory. Adaptivity proves crucial for optimization effectiveness, where the first row without adaptivity perform consistently poorly. The addition of momentum is beneficial in all configurations while SM is more beneficial for adaptive step sizes. The impact of learning rate scheduling also evident across configurations, with cosine decay consistently outperforming constant learning rates. Notably, we observe varying degrees of learning rate sensitivity: adaptive methods demonstrate greater robustness learning rate selection, while non-adaptive methods require more precise tuning.
B.10. Gradient Clipping
Gradient clipping standard in training LLMs for many open source models like LLaMA, DeepSeek, OPT, etc. (DeepSeek- Al et al., 2024; Touvron al., 2023; Workshop et al., 2022: Zhang al., 2022: Chowdhery et al., 2023; Ding al., 2023). Clipping has strong connection to stochastic gradient noise being heavy-tailed (Zhang et al., 2019) and many theoretical results have been shown to suggest some form of clipping is beneficial when the noise could follow heavy-tail distribution(Cutkeoe & Mehta, 2021; Gorbunov et al., 2020; Li & Liu, 2022; Nguyen et al., 2023b;a). We present the results with clipping equal to 1.0 for each method in Table 15.
In Table 15, we see that gradient clipping indeed helps most of the methods achieve slightly better perplexity. In our experiments, we notice that adding some form of gradient clipping produces more stable training.
B.11. Batch Sizes and Random Seeds
Fixed number of steps. We measure the impact of different batch sizes on pre-training LLaMA 60M for 10,000 steps in Table 16. 15 We use the same configuration in other experiments. Typically, smaller batch sizes require smaller learning rates, but curiously, AdamSNSM seems to be stable with the choice of learning rates. Even more interestingly, AdamSNSM's final performance seems to be affected less by the smaller batch size opposed other methods, especially GaLore.
Fixed data quantity. number of steps. In this section, we fix the amount of data to 1.3B tokens for pre-training LLaMA 60M. Hence, adjusting the batch size would also adjust the number of steps. Table 17 contains the result where SNSM shows consistently better performance than Adam across different batch sizes.
the previous section, we compare the performances on different batch sizes fixing the same
Random seeds. Throughout our experiments, we fix the random seed for all runs within a same table. Table 17, we investigate the effects of random seeds by running each batch size 3 random seeds and report the mean and standard deviation. We see that SNSM has better variance than Adam for many batch sizes overall. We also examine the random
This reduces the amount of total tokens trained. However, we only compare optimizers against another. To compare the aptimizer against different batch sizes, one should train for the same amount of tokens.
20

---

Subsct-Norm and Subspace-Momentum
Fraction of coordinates with noise
d 60M
0.8
d = 1B
0.6
0.4
0.2
0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1
Noise density rate B
Figure 10. Fraction of coordinates with noise over noise density rate.
variation the 130M model in Table 18.
Table 17. Mean and standard deviation (in parentheses) evaluation perplexities of Adam and AdamSNSM optimizers when pretraining
LLaMA 60M for 1.3B tokens over 3 random seeds. SNSM rank 128 and gap 200. Learning rates tuned over grid for each
batch size.
Batch size
1024
512
256
128
64
32
16
Adam
31.80 (1.87)
30.46 (0.29)
32.11 (1.53)
34.57 (0.16)
36.34 (0.18)
38.91 (0.12)
43.12 (0.26)
48, 48.88 (0.17)
57.28 (0.80)
AdamwSN
30.11 (0.15)
29.81 (0.12)
30.32 (0.07)
31.30 (0.02)
32.72 (..11)
35.38 (..11)
40.46 (0.97)
45.81 (0.11)
51.01 (0.25)
AdamSNSM
31.39 (0.17)
29.93 (..m)
30.08 (0.19)
30.57 (0.08)
32.35 (0.14)
34.51 (0.14)
37.05 (0.20)
39.39 (0.02)
44.27 (0.00)
Adam
AdamSN
Adagrad
AdaGradSN
Mean Stdev
24.69
22.98 0.07
25.95
24.57
0.07
0.16
0.37
Table 18. Mean and standard deviation across runs for different optimizers on pretraining LLaMA 130M task.
C. Coordinate-Noise Density
This section further examine the coordinate-noiss density model by providing additional empirical results across the train progress. We also provide the full derivation for the convergence of AdaGrad algorithms under various noise density rate.
C.1. Empirical Validation
Coordinate-noise density experiments. To validate the coordinate-noiss density model, sample stochastic gradients repeatedly (via different mini batches) to obtain a sample variance estimate for the true sub- -a-ussainin parameter a, for each coordinate: if .i.... In E R are independent stochastic gradient samples, we can calculate the sample variance 52 as an estimator for o² as 52 = n-1 Σn (g: g)² where ğ = Σn Ji is the sample mean. We pick n = 200 samples (with batch size equals 128) for estimating coordinate-noise on LLaMA 60M across various steps during training process. Figure shows the aggregated noise distribution across all parameters for LLaMA 60M after 100 training steps. There, the noise is quite low for the vast majority of coordinates except for some outliers. While the noise seems sparse in aggragate, more fine-grate fine-grained analysis, presented in Figure 3, shows that noises are dense per parameter, except for the Q and K attention projections in the deeper layers. Figures 12 to 17 in Appendix C. present more noise density rates across various parameters throughout different points of the training progress.
Figure 12 to show the normalized noise density ratio for different parameters LLaMA 60M as described in Section 3.2. The noise patterns show a clear layer-dependent structure, where early layers (like layer 0) maintain consistently high density
21

---

Subset-Norm and Subspace-Momentum
Coordinate-Nole Distribution Layers 60M
10
Density Statistics:
Total Elements: 58,07 18,073, 600
107
Near Zero Elements (<1.0e: 0e-00000:: 58 060
Non-zero Elements: 12,957
109
Non-zero Fraction: 0.02%
Count (log scale)
109
104
10a
10
10
109
0.5
100
1.5
2.0
3.0
3.5
Noise Values
1e-5
Figure 11. Aggregated noise distribution across all parameters after 100 steps of training.
(Step 0) Gradient Noise Density 3 (Num Non-zero- d") for LLaMA 60M
Q proj (d-2.1M)
proj (d-2.1M)
Noise Density 3 (Num Noe-zero-
Embedding:
head:
proj (
proj (d-2.1M)
Gate (d-5.6M)
0.5
pooj (did
proj (4=5
Up pij (5-5.580) 000
Gate (d-3.0M) 1,000
Down proj (d==5.6M)
0.4
(0-222M) 0.489
Embadd (d-16.4M)
K anoj (4-2.1M) 0.472
Head (d-16.4M)
0.2
0.0
0
Layer
Figure 12. Noise density for different parameters of LLaMA 60M at Step 0.
(close to 1.0) throughout training, while deeper layers start very sparse and gradually become denser training progresses. Notably, the embedding layer shows opposite trend, starting relatively dense and becoming increasingly sparse by step 5000, suggesting different dynamics for embedding updates compared to attention layers. The middle layers show an interesting transition pattern, starting sparse but rapidly becoming dense after about 1000 steps, indicating potential critical phase in training where these layers become more actively involved learning.
C.2. Convergence Rate Derivation
We derive the dimensional dependency of convergence rates for different AdaGrad variants below.
GaGGradCCoordinate For = d AddaGra-CCorddnaaa we get Σi= ||wvl1l = add, 11a112 a2d0, and 1
add so the bound from Theorem 3.1 becomes
1
(a'a + od + dI. + oz) O add +
add + ad + Ld
V.ll
T
T
f=1
The dependency on d for the slow term O(1/VT) is dl dl.5+3 The dependency on for the fast term O(1/T) is d1-5d= d2-5, Note that there an inherent d' dependency for the slow term that does not reduce as the coordinate-no noise density decrease.
22

---

Subset-Norm and Subspace-Momentum
(Step 10) Gradient Noise Density (Num Non-zerom d²) LLaMA BOM
Q proj (d-2.1M)
1.0
K pnj (d=2.1M)
Nosse Density 3 (Num Non-zero-
Embedding: 0.971
LM head: 0.973
V proj (e=2 (M)
0.8
Layer
O proj (d-2.1M)
V proj (d=2.1M): 1.000
Gate proj (d-5 GM)
0.6
O proj (d=2.1M): 1.000
Up proj (d=5 6M)
Gate proj (d=5.6M)) 1.000
Up-prei_(d5S.66)) 1.000
Down proj (d=5.5M)
0.4
Down proy (d==:660))11000)
Embed (d-16.4M)
Q proj (d-2.1M): 0.581
Head (d-16.4M)
0.2
K proj (dm) 0.548
0.0
0
Layer
Figure 13. Noise density for different parameters of LLaMA 60M at Step 10.
(Step 100) Gradient Noise Density 3 (Num Non-zero= d') for LLaMA 60M
Noise Dersity 8 (Num NNo--000 69)
Q proj (d-2.1M)
1.0
K proj (d-2.1M)
IM head 0 967
V proj (el2 1M)
0.8
Embedding
Layer
O proj (d-2.1M)
proj (d- (d-2211): 1.000
proj (d-5 (d-5.0M)
0.6
V (d-2. 1.000
proj (d-5.6M)
Down (d=5.6M)..
Gate proj [d=5.081) 0:95
proj 6M)
0.4
Up proj 0.945
Embed (d-16.4M)
proj (d=2.1M): 0.608
LM Head (d-16.4M)
0.2
K proj (d-2 0.577
Layer Index
Figure 14. Noise density for different parameters of LLaMA 60M at Step 100.
(Stop 1000) Gradient Noise Density (Num Non-zero- dd)) for LLaMA 60M
Q proj (d-2.1M)
1.0
K proj (d-2.1M)
Noise Densiry 8 (Num NNon-zero=
IM head 976
V proj (c-2.1M)
Layer
proj (4=2.1M)
V proj (d)2.1M): 1.000
Gate (d-5
0.6
Embedding: 0.737
proj (d=2.1M):
Up proj (d-5.6M)
Down proj (d-5.6M) 1.000
Down proj (d-5.6M)
Gate proj (d=5.6M): 1.000
0.4
Up proj (d=5.6M)::1000
Embed (d:16 4M)
K praj (d) 0.999
LM Head (d-16.414)
0.2
proj (d=2.1M): 0.999
0.0
0
Layer Index
Figure 15. Noise density for different parameters of LLaMA 60M at Step 1000.
23

---

Subsct-Norm and Subspace-Momentum
(Step 5000) Gradient Noise Density (Num Non-zerom d') for LLaMA 60M
Q proj (d=2 MM)
1.0
K proj (d-2.1M)
Noèse Deniityy 3 (Num Non-zero-
LM head 0.967
V proj (e-2.1M)
0.8
Layer
O pro (d-2.1M)
proj (d=2.1M): 1.000
Gate proj (d=55 6M)
0.55
Embedding: 0.501
O (d=2.1M): 1.000
Down proj (d-5.6M): 1.000
proj [d-5 (6-).MM
Down proj (d-5
proj (d=5.660)
Up proj (.-55. 1000):))::00000
Embed (d-15.4M)
proj (d=2:1M)
LM Head (d=16.4M)
0.2
Q proj (d-2.1M): 0.9% 0.995
Layer Index
Figure 16. Noise density for different parameters of LLaMA 60M at Step 5000.
(Step 9999) Gradient Noise Density (Num Non-zerom dd)) for LLaMA 60M
Q proj (d-2.1M)
1.0
K proj (d-2.1M)
Noise Density 3 (Nom Now-2000
LM head: 0.968
V proj (d-2.3M)
0.8
Layer
0 proj (d-2.1M)
proj (d-2.1M) 1.000
Gate proj (6-5 6M)
0.5
Embadding: 0.501
0 proj (d-2.1M): 1.000
Down proj (d=5.6M) 1.000
proj (d=5.6M)
proj [d [d-5.6My)
proj 6M)
Up proj (d=5_66)):1100 1.000
Embed (d-15.4M)
proj (d-2.1M): 0.997
LM Head (d-16.4M)
Q proj (d-2.1M) 0.996
Layer Index
Figure 17. Noise density for different parameters of LLaMA 60M Step 9999.
24

---

Subset-Norm and Subspace-Momennum
For c (AdaGrad-Norm)) we get |0112 = E% |lo,11 a2de, 1|01112 = a2d and ||o||* a a8. This means that our bound from Theorem 3.1 becomes
AdaGrad-Norm
T
d²d + /2 + L add VT
V s 1=1
Ō (a + a + L + a) 0
+
T
The dependency on d for the slow term O(1/VT) is d23 d = d2.5,, The dependency on d for the fast term O(1/T) is d23 d = d Note that when 3 0, or when all the noise is on single coordinate, we recover the dimension-fir free results of previous works.
Subsetor-S Now, consider the following partition strategy, where we divide the coordinates into c = d' =*k subsets of size d²/ /k each with the d noisy coordinates into just subsets so that the rest of the c k subsets do not contain any noisy coordinate. This is reasonable choice due to the empirical validation from Section C.1: The noisy parameters seem to cluster in groups corresponding to the architecture.
With this strategy, we have Ilou; 112 a²dd/ kk =
=1lowallz 3/2 /0//2/k055 /ko.5 if j is noisy subset. We can compute
2 low, || 20.55 | = Ci 11o,, 112= ====aa a²dd. and Di=o lov, 1I = /k. From Theorem 3.1, we get
c--
bound of
vell2 S Ō (a-t88 /k 2d + d + (k- 3/2 a) a²d + //2.5 + -dd-
T
odd
+
VT
T
Set k /73/553/55 so that (d- d28 /k 83//5+3/5
Then we can simplify
- Σ t=1
V 2 S δ
a (0
+ + ((9+1)/5 I (3(++1)/5
+ and +
LL
δ
VT
T
The dependency on d for the slow term O(1/VT) is 1(5+1)/5 d (29-33/110 d for the fast term O11TT) is a bit more complicated: For B E [0, }. we have the dependency on is d3(+1)/55 1218+1)/5
d (1+6)) = 0.3+ The dependency on
ds+1, For 8 [3 1]. we have the dependency on is (B d (( 6--63-0.6 Note that this is only possible partition strategy where the subset sizes of equal size (which is probably the most natural and easiest to implement). There, the optimal subset size is k = 1..6, for which if we plug in 3 E [0, 1] we get range from 1 to d...
C.3. From Theory to Practice
Our theory provides an optimal grouping strategy that depends the noise density. However, in practice, we must trade off the cost to figure out good grouping and the performance gain from it. The key from the theory improvement is to group the coordinates with similar noise magnitudes together. However, any expensive method to figure out these groups (c.g. the Hessian in Adam-mini) would have detrimental effects on the wall clock time and memory. Instead, our heuristic as in Section 3.2 is meant to be a simple method to capture most of these groups
Intuitively, coordinates in the same row/column either act on the same input or are used to compute the same output. The noise and normalization on each input and output would affect coordinates in the same row/columns in correlated way. To provide some evidence, we perform the experiments in Section C.I again in Figure 18, but with the noise grouped by the corresponding dimension according to the heuristics. There we see most groups have very low noise (very close to 0, namely less than 10 12) while a small number of groups (top I percentile in the annotation) have much larger noises. Overall, our heuristics aim to capture the similar noise coming from the same inputs and outputs. Our experiments suggest that this is a major part of the gain. There might be other simple sources of correlation in the noise magnitudes, which we leave for future work.
25

---

Subset-Norm and Subspace-Momentumm
Distribution (s 99th Percentile)
(Values x30-)
Coser
120000
Ran 1.30-00
40000
20000
0.09
Figure 18. Coordinate noise grouped by the corresponding dimension according to the heuristics
D. Subspace-Momentum Convergence Proof
In this section, we provide high-probability convergence proof for SGD with Subspace-Momentun for non-convex smooth objective under sub-gaussian gradient noise.
D.1. Setup and intuition
Given linear operator P Rd - R we have P" : IR + Rd is P's adjoint 16 and we consider P" RR + Rd is projection operator i.c. P P is bounded linear operator such that (P* pP* P.P. Given space V C Rd we denote its orthogonal subspace by V == {u E IRd (v. u) = 0, Vu E V).
Notations.
Let U row (P) S Rdd be the row span of P. Let v Rd + (x) p P* and V Rd - U be
4-- (x) p Px. Then for any vector 12d, have the orthogonal decomposition
((xx V (z).
Let = fz, denotes the stochastic gradient at time t. Let è Pgt, of gg, = P" Pg. E U, and 9r 9: g E U- Let V, ff(x,) be short hand for the gradient at time t and let V = v (Vf(z)) E U and V = (Vf(z))) e U be the orthogonal decomposition of ff(r)) with respect to U and U-L, so that V = Vi + V. Note that the superseript of a variable tries to suggest the space that it lives (either U or U.) We have the following update rule for subspace momentum:
SGD with Subspace Momentum.
mi Bin + (1- P)
9r = 9t -P*PPtt
m. pmmnt
Jt+1 = It - -1 (me + gt
Note that
m4 MP m + (1 3) P* Pgt
= -m,-1 + (1 (1-3)9gr
1' Rdd the adjoint Pa of linear operator is the linear operator given by the transpose of the matrix representation of P. We also generalize Subspace-Moment Momentum to general Hilbert spaces.
26

---

Subset-Norm and Subspace-Momentum
Expanding the terms, that this is just momentum in U
= - + (1 -1- (1-8)PPPM
3P + B)gg
32 pm m + (1 3) (1 B)g
= (1- Σ B gi
(5)
Hence, we can think of the update of SGD-SM as performing two separate algorithms in the orthogonal subspaces:
momentum in the subspace U and SGD in the subspace U (see also Figure 4) i.e. if we decompose into its orthogonal components 24 = 2t + 2t. then
F(+1 xz 177
= xe - .7(11 3)
Σ gi-
2741 17
199
For our analysis, let & := 9t- V, denote the stochastic gradient error time We can further decompose the error into its subspace components:
&t = c + &t (o! v.) + (9r ( V/).
Basic facts. We establish some facts for subspace momentum.
1. Pythagorean: || || 1 2| and V = || 12 + | || and on for these decompositions.
2. Subspace smoothness: If f is smooth [v] ( < T. || y/l- then due to contraction property of the projection operator, we have that the projected gradients of are also L-Lipschitz:
P'' - P [Vf(() - vf(vi)
(6)
3. Subspace non-bias:
E
9
VII = E [
= E [P* P (ge V.)]
0,
and similarly for the orthogonal subspace
[i]] [$(t]
IE |.. 2
4. Subspace bounded variance: if the stochastic gradient's variance is bounded, then its subspace components also bounded TE [|]5
E
11 U
v = E ||5. [12]
E
||5, 11
||5. [1]
- 0² IE [||- 11
and similarly,
E
||s |1
S <
IE [1
27

---

Subset-Norm and Subspace-Momentum
D.2. Subspace-Momentum convergence proof
Suppose that f 12 IR is L-smooth and stochastic gradients (xx) gr is unbiased, i.e. E8 [g.] Vz(z)), and has o-sub-gaussian noise, i.e. Elexp(X² |19. - v[(+))]l()) S exp(X9a²) for all À s.t. [^] 1/o. First, we will show an error bound that is starting point for the high-probability convergence results.
Lemma D.I. If f is L-smooth, then SGD with Subspace-Momentum (Algorithm 3) yields
(3 Laf
((oT+))- -
Σ V
Σ
Σ 11
2(1 (3)
Remurk D.2. Lemma D.I shows that the optimization error of SGD-SM quite similar SGD-M.
Proof. Note that m, E U and E U-_ -Starting with smoothness, we have
L f(x+11) S f(x)) + (1))) xz)
2 || zell?
f(()) 17 (V f(re) mi + + L
m + g 11
2
= -((x)) - (V/ m.) (VI( gt) + r/2L
|| + gt 12
2
We have
1
((441)-f(az S -n / ma) 1 (V ggtt +
| + 11'
2
L
1121
-n V mi) n (+
||me|l 22 +
||
(Pythagorean)
2
Summing it up, we get
T
L
f(TT+1) - f(x1)
< m
V m.)
[mel
V
19
|1
(7)
2
2
t=1
SGD with momentumn error
vanilla SGD in U-
We analyze -17 (V m.) + 2 || and -n subspace is controlled by their respective algorithm. Investigating the momentum term, we have
+ 22 |
separately. Intuitively, the error within each
2
(VV, - (VV. fm.Bmmm + (1 (1-)) =- (V m--1)) - (1
(V( 9 )
BB -VV vi-1,m--1) (V/1,. mm-1) (1--)) (1 (V.
We examine (AV V m+-1):
-(V ( = (V. V.-1, 1) + (V+ V m-11)
(V. V.-1,
I || V, V [m
< L || -,-1 |m-1l|
1 m - |m-
< 1L
m gt
28

---

Subset-Norm and Subspace-Momentum
Now we have
(V) S nT. || m- + ---111 - 3 (VV (1- -(1-)) (V. .
1-1
n
Σ
8
m: + 2. - -3 -3 Σ
(1
gr (V.
==1
1=1
Summing over I, we have
Vi mz) S L
8
m: + g5 1/2 - (1 B)n
3 (V. )
t=
tml
T
= Lnj² 2
8-
m. + g, 11 - (1 - f)n
3
A (
(swap the sum)
i=1
T
T
T
T
< Lij²
mi + 9:
2
Σ
8 - (1 - 3)n
V
9:
3
t=1
Ln²
mm + 9/ 112
-1 Σ
V
1 !
9:
1 8
i=1
2
(Imell 2 + ||g 12) 17 >
T
-
17
1 B
Now, we look E |m; 11:
|mpl 2
3mm-1 + (1 - "1
(1-33) 3 3|mme-1ll 11 + (1- 3) 19 11
(convexity of 11-113
T
B 1llmell + +(1 33) Σ
|| 1|
t=1
[=1
mmm
2
||
1=1
Examining the momentum error terms, we get
1 L T
U m,)
+
||mel1
2
1=1
Ln²² 1 B
+
|m:
11
Ln²
99
17
T
Σ | 2
2
1
1=1
1=1
(3 - L)Ln²
2(1 - B)
Σ
T
|1
112
+
Ln²
3
|
||2
17
Σ
|
(8)
We consider the SGD error terms in the orthogonal subspace:
n V ) + 2
L
n
(V
- - |+ 2 +
7L
|
2
1| L
=
=
| 2
(9)
29

---

Subset-Norm and Subspace-Momentum
Now we are ready to combine (8) and (9). First note the common terms | in both equations combine to a sum similarly
9 II"
Tap
T
17
T
(3 - Lnp (1 5)
T
1
2
==1
Combining both terms, we that the terms are combined from both subspaces (red from (8) and blue from (9))
T
T
Σ
Σ
1
=
1
T
T
= -17 Σ
(Ve. Et)
t=1
t=1
(3 Ln
+
Lnj²
L
96 1 112 =
(3 Ln
Σ
T
|1
2(1 (3)
B
2
2 - 33)
t=1
t=1
Plugging everything back into (7), we have
(3- Lm
f(TT-)) S
V
2
2 Ee)
Σ |1g. 1
(10)
2(1 3)
D.2.1. PROOF OF THEOREM 4.1.
Our proof uses the technical tools from (Liu et al., 2023c), although the strategy here has been simplified. We use Lemma A.1 from (Liu et al., 2023c):
Lemma D.3 (Lemma A.1 from (Liu et al., 2023c)). For any 2 0.00
I and a nonnegative o-subgaussian random 2n
variable X, we have
20
E
+
b²
x2
(ax + +6xx3) s <xxp (30² (a² (a²+b%) +b²)). zl
We present some useful tools:
Corollary D.4. Suppose that is a mean zero random vector in Rd, where ||X|| is o-subgaussian For 0 S a S 462
1 and
B E Rd then
E exp a
|||] + (B.x)] s exp (30 (1B1)+
Proof. We have
(a (XXI++(B.x) x
E exp
(a [1][x]/ + (B,x))]
=E E
+ a 1 + (B,X)
Σ kd
E
+ a X 2 +
Σ Kt!
(2 (B,x))
k=2
< E
+ a x + Σ _
Sexp
(302 + [B13))
30

---

Subset-Norm and Subspace-Momentum
We control martingale via:
Lemma D.5. If we have a sequence of random variable X, with Fi = X2.,,, Xor t= 1,2,........ T. we can
bound exp (X,) Fi] I exp(Yi), where Y, Fe-measurable, then
T
Σ E1
X1 S Σ
Y2
log (1/6)
t=1
holds with probability at least 6.
Proof. Define the Z. = Xx Y, and S, = Ei= Zi. Then
E (exp (Zt) Fi] E lexp (X, - Yt) Fil
= exp (-Y;) E (exp (X,) Fil
(Y, is F--measurable)
< exp( ( exp(Yr) = exp(0) 1.
Then we show E (exp (S:)) < via an induction: we have E (exp (ST) Frl = E exp (ZT) [F]] S 1. Suppose that
E eexp (8(+1) Fit1l
E (e))((.)) Fi]
E jexp (z) exp(S++1) [F]
E eep (Z,) E (exp (S++1) Fi+1] Fi]
< eepp 2,) Fi] 1
Hence, this implies that [exp(S:)] S By Markov's inequality, this means that S1 log(?) with probability least 1 8:
S
Σ Z. Σ
X. Y Y,s log (1/8)
t=1
T
Σ X, S Σ
Y, log (1/8).
t=1
We now prove Theorem 4.1:
Proof of Theorem 4.1.
Starting from Lemma D.I and letting a 2[-B) and A1 = f(x1) - S- for simplicity, we have
(3-8)L
ATt1 A1
T
T
Vel2
(., + +an2
119.112
i=1
Σ
| 11 1 Σ
V Ex) + anj 2 Σ 115. V./
"
1=1
=1 (ar) - 1) Σ
V. + (an) - 1) Σ (V., + an) 2 Σ |
Rearranging and defining some weight w 0, we have
T
T
w (Ar+1 A1) + mpw (1 am)
| < ma (am- 1) Σ
(ee Et) + oown Σ
||eal1
31

---

Subset-Norm and Subspace-Momentum
Let Fi == a (fi....,, 1 denote the natural filtration. Now, since E E (Ve, E2)
0 and & is o-sub-gaussian we
have that V E Fr and so if 0 I wour2² I
1 4
Corollary D.4 implies
wn (an - 1) (V Ex) + wanj 115:11) Ft I exp (30² [won] + = (am - 1) 1) (Vv.)))))
E
exp
Then Lemma D.5 implies that with probability least 1 - 8, we have
T
wiy (an - 1) Ev, + wann"
2
Ilelll 2 S 30
wan + + (an - 1)² [| + log (1/6)
[=1
1=1
T
30 an aT + 30 w2 (am - 1)²
|| |.ll + log (1/0).
Then with probability at least 1 8, we have
T
T
T
w (AT+1 A1) + my (1 - on)
| S my (az) 1)
Σ
&v + awnp
2
||Ell1
1=1
i=1
< 30 aor'wnaal + 3o 21 (an) 1)²
Σ
| 2 + log (1/0)
t=
mw (1 am)) Σ
|vll2 < ww
w²n (am 1)² Σ
V 2 + +oo (1/0).
Combining the Ve terms, we get
(nw (1- an) - -302 (am 1)²)
Σ
11 < wAi + + log (1/8).
(11)
Setting 12 then
my (1 -an) - w30 (am-
1)² 1/
1 aT)
-30² un (an 1)²)
1 (an 1)2
1 am
2 m7 4
if - oz) 2 1/2. Furthermore, we have that wonn) = then both requirements are satisfied. Consider the LHS of 11, we can bound
on < :
if < 3. as required for Corr D.4. Hence, if S 1
12
Ta
T
(mu (1 an) - 302 w2 (an 1)²)
Σ
V
T/C
Σ
V.
Finally, we have
T
m7
|.11² Swi + + 2
aT + log (1/8)
Σ
|v,12
4 - + 2 + 480 log (1/8).
"
32

---

Subset-Norm and Subspace-Momentum
Setting 17 min
we have that with probability at least - d
o
4 I
Vell
A
30
+ 480 log (1/8)
1
4
4 oaaTT
A + 30 min
aT + 48o² log (1/8)
2a
min
2xr
a-oo
aPaT
A
<4 4
2a
V
A
A 30
o²uT
of' + 48o log (1/6)
8A1a + 40 AT 1 + 30 A + 480 log (1/6)
8A1a + Ta OTA
+ 480² log (1/8)
1
T
8Ai T
48o log (1/8) TovaAi
V
T
VT
T
We are done.
E. Subset-Norm adaptive step size full theorem and proof
We show the full result in Theorem F. with all the polylog terms omitted from Theorem 3.1.
Theorem E.1. Suppose that f Rdd R is L-smooth and lower bounded by fa. Given unbiased stochastic gradients
f(x) with stochastic gradient noise & == f(,)) f(x)) being o--per-coordinate subgaussian for i [d]. For partitions of the parameters into disjoint subsets [d] = - V; with V n v = 0, if + j. the iterates æt given by (3)
satisfies the following inequality with probability at least 6eö (for failure probability 8 > 0):
42
c- lov.
I(6) T
V. 12 G(8).
where G(o) and I(0) are polylog terms:
VT
A G(D)
4v
i 1 ll + 1())
H())+ + T/ ||0112 2 onL 3/ max log
log
"
b0,min
2A,
8 8log
Tn by,min
1(6) 1b0ll +
||0
log
loo SnLelog
bp,min
8 8810v.
log
1 + Ilov T 1lo 11 log } ))
H(6) = Σ In (T/6) 12 (+2a)
2 log
1
where
- of |oov 1== ==== Ejev, og Fmax maxxeld de A -((1)= So, bo. minsejd bo, 0.
E.1. Proof of Theorem E.1
For simplicity, in our analysis, we will use V., := v.f(z.) and V., == v.f(x,) to denote the i-th coordinate of the stochastic gradients and gradients at iterate respectively. The proof utilizes techniques and follows the strategies (Liu et al., 2023c), where the main effort is adapt the techniques for handling subsets from the AdaGrad-Norm and AlaGrad-Coordingg proofs in (Liu et al., 2023c).
@ Proof. We write b, be crucial in the main analysis). We start with the smoothness of f and A, = f(xx) fa.
G
(ef(e)) for ls E V, (we will use this notation briefly to show some steps and will bis
denote
33

---

Subset-Norm and Subspace-Momentun
L
A441- - S (Vf(2)) I2+1 x1) +
| - zell
2
4
+
(12)
1t
2
bt
V. V LJ
2T
6-1
+
Σ
2
by
1= JEW,
i=0
jEV;
c. (Er.)
V1.j)
I.
Σ
(G,, Vrs Ve,i)
b1,5
bg,t
JEV;
jEV;
V2
VeySt.
T
V
fj
11
+
bt.i
2
1=0 jEV.
1=0 JEV,
6-1
Veskeed
6-1
2 C-1
+ 17
Σ
Veshes +
Σ
6 (13)
Ue
be,i
2
JEW
1= JEV,
1=0
jeV;
1=0
jeV;
Now, we analyze
for
i
0, 1, 1
c 1:
bii
1
1
b atii
04.5
be,i
arpbbrii
b9,1 on..
ae,bbr, (br,s + (4,5)
1
18-1,5 +
12 12-1, - Vv.
dee (be,s + 01.))
be,abe, (be,s + ax,i)
- ((x))l
+ ([V,
tee (be + az,s)
Since by. v 19-1, + Pc. ||
2 ( and ar, = 1 /Vv. 2 f(vll we have
flor (1 + ((ve)(()))l)
<
b
+ [(.))l arybee (1 (x.))
F - ((x))l
ae,bbei
[(re) - (xx))
ar,
34

---

Subset-Norm and Subspace-Momentum
Hence, we have
1
1
C. ||
br,5
ar,
Then from 13, taking the absolute value of E
Ejev, oes
(1
1
1 bes
Ve.rSt.j we can bound:
VIJ
Veslis Ge,s
1
1
2T
c-1
A411
1)
+1
V
agss
b
2
1=0 JEV;
JEV:
i=0 jEV,
JEV;
Verfees Or
||st. |I ar,sbe,s JEV:
<1 Σ
+17
Σ Vewtrul +
Σ
bt,5
2
1=0 jEV.
1= JEV,
1=0
1=
JEV
(1)
6-1
Ve,seea Or
|1.. || aeabi
+ n/L
c-1 1
+
bE.5
2
1=0 JEV.
1=0 JEW.
1 JEV.
G-1
Sーク Σ
be,i
de,s
1=0 iEV,
jEV:
||E+. 11
1
+ 17
v.,1ll ||
+
2b₁,,
2a
2
i=
jEV,
where (1) is due to Ejew. when we apply to vectors. The last incquality is due to 2ab S a² + b². Now, we can sum both sides for t = 1,..., T to telescope the L.HS:
|e = (V v.vl.
|
v, ll 1EC.v. 11l and |-| denotes coordinate-wise absolute value
Ve,See. atis
AT41-
-
A
17
b
1=1
1=U
JEV;
1=0 jEV;
11r. 111,,
2
Vv,. 112
n
+ 1
2b,
20²,
2
1=0
jEW,
Rearranging gives
ΣΣΣ =
T Tc
Vi,
A AT+1
AT+1
T c-1
Veslie
bb,5
atii
JEW;
T
1. 11
Vv.3.11 2a
n
ΣΣ Pi
T c-1
v.
+
262
2
t=1 im
jEV;
On the LHS,, we note that
T
V
T
Vv,v, 112
ΣΣΣ
ΣΣ
by.
b
JEV.
1=
35

---

Subset-Norm and Subspace-Momentum
We now bound each term separately. It's easiest to bound C: Ei Li
Zjev,
Tel ΣΣΣ by
Vi
T
b
br b,,*
1
2 log
jEV,
JEV,
6/ 12-1
67-15
log
T
b
log
1,
br.
2
log
bo.
We now have a useful inequality
v. |
brx vi= = 0,...,c- 1. bo,""
< 2 log
(14)
Next, we deal with EL
Vealen via martingale argument. Let J3 = (6I,...,((I)) denote the natural Gest filtration. Note that v, is F--measurable. For any w > 0, we have for each E [c]:
Ei=o Ejev
V
E
exp
2w
F
at,i
of
JEV.
jeW:
22 1
VerSee
exp
2w
E
exp
Fr
Ot,i
jEV.
JEW.
1
Then simple inductive argument and using Markov's inequality gives with probability at least 6-
V aZ
1
< 2w²
log
8
Or
1=1 je
By a union bound across all e subsets, we have w.p. at least - co:
ΣΣΣ
T
Vesles
c-1 ΣΣΣ
wo? tj
log 1
(15)
Ot,i
w 8
jEW
iE
36

---

Subset-Norm and Subspace-Momentum
Let's call the event that (15) happens Ej. Now, consider E
Ei=o
TimoEjen Ejev,
We have
V2
V V. 11
V. 11
6, b/-1,5
VV,
JEV;
||
2 v,
+ v1.
V v. 1|
12-1,5 +
2 11E,,v.ll
1 |
V
|15. v. 11?
2
+2
b7
02
For (+) the fact that c+x
is increasing function and V v, 1 = |r,,, v, E,u;
1| 2 v..0
1|
+ 1.,,
Let max == maxie(d] then under event Ej, we have with probability at least 1 - cō:
T' c-1
Vi.ift.j
T
WO ?
+
log
Ctps
0 JEW;
1=1
i=
JEW;
T c-1
wa max
2
log
w
JEV.
V.
Ev. 11?
S wo max
+
+ log
1=l 1=0
[.,. 12
T
11Ee 1|2
(set = Velog?
max
clog
8
2
02
+ m V /clog 5
Gana
1=0
mia
12
11 112
2a
+ a.
D2
1=1
1=0
Volog 1
where the second last equality is due to choosing readability.
and the last equality is letting m e log 1/2 for
Omax
Let Mr; maxIsT KG,ll Using our notation, we can define Mr.e, = maxesT v. Under event Fg (and our new
37

---

Subset-Norm and Subspace-Momentum
bound for C), we have that with probability least 1 - co:
1.
C
A
T
VK 11 2b,
V v.
12
by bot
u.
1
log
2a7
im0 jEV;
1
VenGeni
(16)
Ot,i
JE
1...v.11
1r.0.11
log bT
Mr
+ 1
(def of Mr,w,)
20/
202
bo,
0
2m
||
(B))
1
T
11 by.
20
+ +a+
b
7
1=1 2=0
bound with (C)
c-1
|kt,v 11?
Vv,, 112
log bT.
Σ
MT..
nT.
(17)
262
2a
bo.i
9
T
1K,, 11,
20
+ a+
7
t=1 mm
|Er,,,
Vvv, 112
br. log bo.
MT.P
+ (nL 4a)
(18)
2
20
t= 1=0
1=
A
VEE,, 112
20
+ a+
(19)
e-1
EE 11?
T c-1
V.,v. 11?
C-
by
MT. MT.W.
MT.v.
+ (nL + 4a)
log
(20)
26
20
bo.4
1=
==
Let us turn our attention to Mr.s, := maxier V,, Note that
m
IK,, 11?
Pr
|1 v, 2 = Pr [E[T]
max
exp
2 exp
(for w >> 0)
maxeerr]
1112
S
exp
E
exp
(Markov)
W
I1Er. 11
exp
E
max exp 6E|2]
W
| Wr 113
exp
exp
107
1E (T)
38

---

Subsct-Norm and Subspace-Momentum
We have
E2 ZjeV:
E
exp
E
exp
St.j
Σ
jEV; St.J
E
St.j
exp
JEV,
E2
E
exp
(independence)
JEW.
Since sub- sabb-aaussiantty give us
1 exp (X2a7), , l oi
E
exp (A2
vi E [d]
we have E
exp
]
exp
if
We pick = lov, 11? Ejev, of 2 oj. Vj v, Hence, we have
See, Vill 112 ) ||o
oZ ||o0.11 11
E
exp
II exp
jEV.
llov, 11?
exp
1.
(21)
||ow,11" 11?
We have actually shown that ε.,,, is a || -ubbauussian random variable in IR* (see Proposition 2.5.2 in (Vershynin, 2018)). This fact will come in handy later. Now, have
A
[ 11? |low, 11?
]
Pr max [E(T]
1,, 2 A
S exp
E
exp
l ||
1E(2]]
A
exp
T.
|lov, ||?
Setting exp (TT- T 8 gives A = ||00,11 In In T/8. Hence, we have with probability at least 1 - 8,
Mr.v max & w. :E[2]
Iov, 1| 2 T/8.
(22)
Union bounding across all 0, 0.......c c- 1, we have that with probability least 1 co,
Mr.w. |low, 2 T/5, Vi = 10.
(23)
Let us denote the event in (23) by Fg.. Combining it with event En and starting from 19, we have that with probability
39

---

Subset-Norm and Subspace-Momentum
1 co:
T
V,,, 11
A
V.V. 1|
VV 11
+
2a
+
Σ
Σ
Mr, v.
b,.i
1
26
Aml i-0
T
c-1
V 11
b booi
Mr.V
+ (nL + 4ax)
log
207
1=0
A
+
2a
[L 11?
+
1/8
Σ
T
c-
loo 11
115.. 112
+ a+
26
11
1=1 1=
V 11
br.
1/
|1op 11
+ +nn
4a)
log bo.
20h,
t=1 1=0
1=0
A
llo
T
11?
+
In T/8
+
+2a
+a+
2
1=1
c-1
||a
11
T
V.V
11
c-1
br,
In T
+ (nI + 4aa
log
2
bu.s
0=
1=
|vo; 4
|0,,0
11
Recall that
< 2
+ 2
we then have
br,
| r. ||
ov:ll 1|
T
Vr,,11
100:
T
115,. 112
In T/
< 1/1
2
2
2
0f
2
0=
1=0
|?
T
e
1
T
IEE.. 11
= Ln
10 112
2
In
1/8
lo 2
62.
62.
1=1
1=0
[=1
bT
T
11Er, 11
$ In T/6
1o 1|0
+ In T/8
1ow 2
(from 14)
bot
b2
1=
Hence, we have with probability at least 1 2c5:
T
Ve,,
A
+
[n 1/..//| +
(00o.ll+2a///)))) +22x)
Σ
15..v. 11?
(24)
br,
17
62
c-1
br.s
c-1
brg bo.
+ a
+
T/8 llov. 11 11
Σ
(nL + 4ox) log
bo,s
1=0
1=0
A
T
112
+
+ 1n 1000.1' 2 +2xx)
Σ
(25)
b4,5
1=0
1=1
br bo.
+ ta+ +
[mT/6 100, + nL + 4a)
log
(26)
Now, we bound ΣT.
and log bru bo.s
We need to first lower bound Σ | || We proceed by noting that
V _| = VV,, + &..u, 1|2
= -V,,, 112 + 2(E, V.N.)) + 1E,,,,
//Vr, -1[ 1/ + 16.. 2- -2(E V.N..).
40

---

Subset-Norm and Subspace-Momentumg
Define for L {0,1, ",,, T
and some constant U, be specified later:
1
U411= = exp Σ w/ (IV s. 11 v.v. u11s. Us v..1V
Ur. - exp w (IVV,,0))) - 11 112 ||E,, IP) well v..lVa.v.11?)
Uc U4xexp (we (2(5,, V...)) - 2/ V...,1
First, note that U, E Fi. We show that Ug is supermartingale
E (U++1
[F]] = E [U, exp (wr (2(5r V.,o,)) 1/ V.,, v,1133) Fi]
Ug exp (- V., [1 E jexp (E V.v. Fi)
S Un exp (- V.. |1 E exp (ww²" 1oo, I1 V...lll' 112 Fi
Us.
(wr=4w; l1ov 113)
where (x) is due to Lemma 2.2 of (Liu et al., 2023c) and the fact that E,v, is llov, 112 -ubgaasusni from (21). Hence, by Ville's supermartingale inequality, we have
Pr
max rE[++]]
U/ 2 8-1
S SE = 6.
This implies w.p. 2 1 - 5, SIST.
( 01.2.112 2 +11,,,v,1" 11) -uallV.,e 2 log 3
ws
W.
Awu 1ov. V 122 112 WellEo, 1.,1 11 5
Wg || 1 V... 11
|++ +log
1
1
Aw, |lo,, 11ov, V.v.
I1 v. 11
4 8. w: 11
+
log
8
00
s=1
$=1
Set w.
to get
1|/0011
1
Σ
I v.
1.
Σ
1 w..1 + 1|lw,,11
log 8
11 T.
(27)
We are now ready bound
T Σ
|,,2 bi.
Starting by applying (27), we have that with probability at least 8
T
|| v.
||
b2,
[=1
t=
b/
+
Σ1
V.
||Er 11
+ E
1,,, 112
Allov,11 11 log

---

Subset-Norm and Subspace-Momentum
where (x)* = max (z, Let 7 = max (10) u { Ngr Exm 11E, 11 s 20/)
for some C' 2 0. We have
1. 1|
|| v, 11
T
1 1|
Σ
Di,
b,, + s= V
t=1
t=
1
10,0,11 11
+ Σ
11,,,111 11
00, t=1
00, + Σ 25=1 115,, 112 4 1000 1 log
1=7+1
2C'
T
|1E,. 115,,v111
Σ
+
0,
+ Σx
4 |10,11" log
100.1110g 4
(=++1
Σ |.
Now, since
2 C for t > T, we have b3. + Ex=1 [6,, 12 4 1l 11 log 2 b0 - -411ovo,, log
log + C+
2
122 If b - 4 [||ov,1 2 log : 2 0, then we pick C = 0 and b/. - 4 loo, 1| log 1 + + _ 2 |1E, 12
1 Σ [,,,,
122 1f. -41|ov,1 1| log : O, pick C 4 llov, 12 log 1/ bg. 0. which gives 1g - 4 | 1og +
C' E 112 2 2 E. I1E., 11 In either case, we have bg. - 4 "|o,,l1" 1 ! + C 2 E' 1, 2
Ei= V,,,,, 112 Hence, letting C = max (0, llov, 1| log _/ bb.) 4 ,lov, |log|] 1.- we have w.p. at least 6:
1 M; 11
20
1 v; 1
2
Σ
1,
b/
s=1 11
1=++1
20
T
|| w. 112
+2
b
E
t=1
8 ||ow,,1"
log
|1,,, 1|
+2
by
15.. 112
s=1
Let X, = =1 Ei 1. = 2 -x,- + I1Ee, 21 where No- = 1. Then,
T 11
T
X1
X
T
X1-1
Σ
E 1I5
X1
Xx
1=1
1=1
t=1
T
X
log
X₂-1
tml
Xt
log
X
1=l
Xx
T
= log
= log
1+
Σ t=1
|1 ve, 112)
Xo
Hence, with probability at least --:: 8:
T Er, w. 12
8 $1|ow,1 2 logg 1/
T ++>un
<
+
2log
1+
(28)
t=l
42

---

Subset-Norm and Subspace-Momentum
remains bound Et=
vi 11
Note that
T
Vx 112
Pr
Σ
a
=
Pr
exp
|| 0
11?
2 exp
u l 11?
11E
exp
Σ 1=
Ioo.
exp
|
exp(T)
v,, is 11.00 -
cxp
|Tou.|7
Choosing llo,, 11 |lov; 112 log 1
z gives that with probability at least 1 - 8, we have
Σ 1 S 11ov,11 2 l1o 11 log
(29)
Having high probability bound the sum of the stochastic error of the subset-norm, combine both events from (28) and (29) to get that with probability at least 1 - 20:
w. 11?
81100, 1...,112 log
1 |lo 11,0,111 11 T + 1 1l log 1
+ 2
(30)
0,4
Then we can also condition on the event that (30) happens and combine with the event in (26) to get that with probability least 1 2có (assuming e 2 2), we have
Σ 1e 11
In 1/0 10ov 1 2ax
(31)
bii
17
by,
br bo,i
+ +a+
In T/8 10o,,11 + nL + 4ax
log
(32)
1=0
( 8 (llos,
1| 1Plog] J
log 1
1n |oo, 11 # 2a)
+ log 1 + 1lov, 112 T + 11ov,
1
Vo,
1=0
=:H(6)
(33)
br
T/8 |lov, 1 12+nL. nL + 4a) log
+a+
bo,
A
br.
H()) + In T/8 1oo, 112 1/2 11 + 4ex) log bo.
(34)
First, note that by.s I |brll Di-d by. Letting bu min, bu.s, we then have
[n76 1ov.11 + nL + 4a) log br bo.:
|
s log
[InT//] 1|o0 2 nL + 4a)
bo,min
1=0
|
log bp,min
[n 1
+enL +caa)
43

---

Subsct-Norm and Subspace-Momentum
v.l bea
Now, note the LHS term
>
of (32):
V Ve
c--
2
1
Σ
b
Ve,ull2
2
V2 1 \
b
i=0
1=0
Ve
G-1
Vv.v
lb
Eizo
c-1
be,s)
be,i
1=0
Now, Di= bii = ===
1bra1 = ||bell so with probability 1 - 2cd:
1
T
Vv.v. 11
|brll
l
by.5
1=1
1=0
T
T
Vv,, 12
|Vell2
< ||br|||
bi,
1=1
(=1 1=0
A1
| ll1
s ||brll1
+ cH( + (In - enL
+ + 4ca)
log
(35)
77
bo,min
- 1bbllla
AA1
+ (8) + (nn (/n - en + 400) log
|| 11
(36)
77
bo,win
It remains to bound|
|brl11- We start again from smoothness of f:
L
A++1 A S (V, - 2z)
+
|0
xzll?
2
4
2
V1, 6
+
2
b
VE
1
62 L.V,
E
by
2
i=U
JEV;
4
V
L c-1
( 0,0
2
+
E1
+
b:
2
62
1=
2
2
G-1
St, V..
2
C-1
Σ
Σ
+ 11
+
b
br.i
2
1=0 JEM
i=0 JEV,
1=0
12
1.
1
Ee, Vi.
+1
(37)
be,b
2
b,
bet
JEV;
Note that
Eej V
ΣΣ
S1.j
br,b
br.
2
b
1=0 jEV,
jEV;
V.V.
ト
ト
b
b,,
0= jE
2=0
Plugging back in, we have
44

---

Subset-Norm and Subspace-Momentum
V
V V.V:
1 11? b,,5
A++1 A S
Σ
+
1
L
f=0
1=
Summing over T and rearranging, we get
V 1,
2A1
E.w.ll 1??
T c-1
V 1,0.
|
De,i
1
be,i
D2.
(=1 f=0
12
4A1
T
c-1
T
4nL
1
||
2
b
7
b
6
be,6
We can bound Σ 21=0
LnL
b,,s
11)
v
as follows. Consider i E [c]. Let T2 max {1 ST bes S 4nL} so
that t 2 Ti implics b > 17L
4nl.
< bz.""
1.2
T
4nL
1
2 = Σ
T2
1
1
4n
be,s
D2
v:
1
b
1=1
(=+++1
<0
T:
1nL
1
|e,u,
1
D2
bi,
2
$ 4
62,5
(=1
brad < 1 log
4nL.
s 8 log
bo.
bot
Hence, we have
4A1
T
c-1
4nI
+
2
+ 8
Σ
log
b1,5
bes
bo.
(=1 10
[=1
1=0
Consider the LHS
T
V V.
T
c-1
62 1 1,1
T
c-1
62 1,1 b
b
b,s
b
t=1
==
1=1 ==0
T
b3- 1,1
T
bi,
b 1
1=1 ==
1=1 ==
ΣΣ
bi.
1
Σ
br, booi
0=0
|
|1bo
Hence, we have
2
1 T
4nL
|| 11, < |oll, +
+
SnLolog
1
b
bo min
0
| 2| bet
It remains bound ΣT
T
for each i E [c]. Recall from (30), with probability at least 1
45

---

Subset-Norm and Subspace-Momentum
1
Σ
Σ
/
Iloo
10 14 S T.
We have with probability at least
1 2co,
r
1|5,,, 1|
1 10.5 Ex=1 v.,,
T
1=1
V b9, +
(Eezi 115,, 11?
llo
log
8 11o0
log
T
+ +222
Σ
1112
bo.i
(2)
8 $|00, log 1/
Av lov, 112T T + 11o log 1 8¹
b
where (1) is due to (27) and (2) is due to Lemma (29). Hence, we have that with probability at least 1 - 2cd,
2
I1o
11
log
4nL
|bbll1 - |60l11
+
11ov,112 1 + 1100,111 2| log + &s,Lclogg
bo,i
bo,min
1=0
=
241
8
log
4nL
|Poll1 <
+
+
loog 1| + AVT
|lo, || + log
1 + 8n log bo,min 1=0
"
1=0
1=0
2A
8log
+ Ln 4nL
4
lov.
|1bol1
+
1|01
log
Ioo
"
ba.
bo
Hence, we can combine (36) with the bound for ||bا||, to get that with probability 6cδ:
T
A + H())+ +
|brll
V Vill2
< |bbll11
In 1/ + cnL + mcx V clog 8)
log
bo.min
4VT
lo + +(8)
1=0
A
4VT - (=0
lla + )
(5) +
In || +nL + 4/2 o max V log log
17
bo,
Dividing both sides by T, we get the theorem that with probability Ger.:
T
4 low.
ll
I(0) T
V < G(6)
where G(d) and 1(6) are polylog terms:
VT
cnT. + 3/2 0 max /og 5
4VT
D Iloo, || 1(6)
G(6)
H8)) +
In T/5 ||al1/
log
do,min
241
8 log 1
I(6) = |bull1 +
||0| + log
lov.l1 ll + LnnLlog bo,min
bo.min
1=0
(8 [||T 1|² llo
2og (1 Ilov, 112 T + 1lov, 11 log ( ))
H(O)= =
In (TT/)) 1222a)
46


Reviews:
[{"review_id": "rF1Ar7Wavf", "reviewer": "Reviewer_6d8c", "summary": "This paper proposes two techniques, Subset-Norm (SN) and Subspace-Momentum (SM), to reduce the memory footprint of adaptive optimization methods (like Adam and AdaGrad) when training large deep learning models, particularly large language models (LLMs). Subset-Norm reduces the memory of the adaptive step-size term by sharing it across subsets of parameters. Subspace-Momentum reduces momentum memory by restricting momentum updates to a low-dimensional subspace and using SGD in the orthogonal complement. The paper provides theoretical convergence guarantees for both methods and demonstrates empirically that combining them (SNSM) improves both memory efficiency and performance (lower perplexity) compared to standard adaptive optimizers on LLM pre-training and fine-tuning tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s careful consideration and constructive feedback. Below, we clarify points raised and address specific concerns.\n___\n\n> the \"why\" behind the performance improvement of SM is not fully explained.\n\nThis is a limitation that is general to the class of algorithms that utilize momentum and it is an important open challenge for the community at large because the theoretical bound for momentum and without momentum are the same. We believe that any advance for SGD with momentum will lead to a better understanding of SM in our setting. \n\n> in addition to Adam and GaLore, a wider range of baselines would strengthen the empirical evaluation.\n\nWe will add comparison with the following methods to Table 3 of our revision for the pretraining task. Some comparisons are shown below: \n| Method          | LLaMA 60M (1.3B tokens) | LLaMA 130M (2.62B tokens) | LLaMA 350M (7.86B tokens) | LLaMA 1B (13.1B tokens) |\n|-----------------|-------------------------|----------------------------|----------------------------|--------------------------|\n| AdamSNSM (ours) | 29.74                  | 22.43                      | 16.91                      | 13.96                   |\n| AdamW           | 30.46                  | 24.60                      | 18.67                      | 16.00                   |\n| GaLore          | 34.73                  | 25.31                      | 20.51                      | 16.76                   |\n| FLORA           | 32.52                  | N/A                        | 23.69                      | N/A                     |\n| LoRA            | 34.99                  | 33.92                      | 25.58                      | 19.21                   |\n| ReLoRA          | 37.04                  | 29.37                      | 29.08                      | 18.33                   |\n\n>  some important ones are not included in the comparison such as Adam-mini and FLORA.\n\nWe currently cite and compare with these works in the related works in the related works, Section A of the Appendix. However, we will include a more detailed comparison to Adam-mini in our revision, where we include a draft below:\n\n\"While Adam-mini also employs a grouping strategy for the adaptive step size, it is primarily motivated empirically and lacks a general grouping strategy for general parameters. Adam-mini gives a heuristic approach to form the group for transformers’ parameters without showing theoretically what this grouping is trying to achieve and if this goal is achieved, why the convergence is improved. In contrast, our theory results show that grouping by noise magnitude leads to improvement. Our heuristic is an efficient method toward this goal and the experiment validates this by showing that our groups successfully isolate the noise into a small number of groups and perhaps as a consequence, have improved performance (also see our additional experiments below).  We further demonstrate in Table 2 that there are scenarios where subset-norm attains better convergence than existing methods. \nIn experiments, our AdamSNSM uses less memory than Adam-mini, due to the fact that Adam-mini uses full momentum while we use momentum only in a subspace (which outperforms full momentum in many cases given a good choice of subspace). Furthermore, in terms of perplexity, Adam-mini performs very closely to AdamW while our methods outperform Adam (which performs similarly to AdamW) on a range of language tasks and model sizes.\"\n\n> Compared to standard adaptive optimizers, the proposed methods are fairly complex\n\nThe theories we developed for SN and SM are general and allow for many options, including future works. However, as reviewer t4t5 and KhrK pointed out, we also developed specific implementations that are simple and effective.\n\n> The reliance on SVD for subspace selection in SM is a computational bottleneck.\n\nAs demonstrated in Table 6 on page 8, our methods do not depend on the SVD update gap to be as frequent as GaLore. In contrast to GaLore, our algorithms (e.g. AdaGradSNSM) could even benefit from a larger update gap as in Table 6, further reducing the cost of SVD, while achieving the best performance out of all baselines. Section C.4 of the Appendix contains additional ablation studies on the projection gap and rank which control the additional overhead cost and benefit from the projection computation.  \nFinally, please refer to our response to reviewer t4t5 which contains additional results showing the wallclock time and peak memory for different optimizers. \n____ \nWe sincerely appreciate your detailed reviews and constructive suggestions. We believe that the discussion and additional experiments improve the quality of our submission and hope this provides sufficient reasons to raise the score.", "claims_and_evidence": "Most claims are supported by evidence, but some require further strengthening. For example, the main claim is that SN and SM improve performance and memory efficiency. While the empirical evidence is generally strong, the gains are sometimes modest, and the \"why\" behind the performance improvement of SM is not fully explained.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria make sense for the problem.", "theoretical_claims": "I didn't thoroughly check the correctness of the proofs.", "experimental_designs_or_analyses": "The experimental setup is generally well-described, following standard practices for LLM pre-training and fine-tuning. The ablation studies on subset size and subspace selection are useful. However, in addition to Adam and GaLore, a wider range of baselines would strengthen the empirical evaluation.", "relation_to_broader_scientific_literature": "The paper's contributions are related to memory-efficient adaptive optimization and LLM training.", "essential_references_not_discussed": "The paper cites many relevant works, but some important ones are not included in the comparison such as Adam-mini and FLORA.", "other_strengths_and_weaknesses": "Strengths: The problem of optimizer memory is clearly motivated, and the proposed solutions are intuitive.\nWeaknesses: Compared to standard adaptive optimizers, the proposed methods are fairly complex, and the reliance on SVD for subspace selection in SM is a computational bottleneck.", "other_comments_or_suggestions": "None.", "questions_for_authors": "See above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "44Rdq3sOrI", "reviewer": "Reviewer_KhrK", "summary": "This work studies two modifications to widespread adaptive optimization meta-algorithm procedures, with the goal of reducing the memory footprint of adaptive optimization while simultaneously improving performance.\n\nThe first modification is referred to as adaptive subset-norm (SN) stepsizes, which is similar to existing methods such as Adam-Mini. Rather than maintaining an adaptive learning rate for each parameter (as in traditional Adagrad) or maintaining only a single adaptive learning rate for all parameters (Adagrad-Norm), they propose partitioning parameters into groups and maintaining an adaptive coordinate for each group. They state that their analytical results (e.g. Thm. 3.1) show that by choosing an intermediate number of groups, they can obtain better convergence guarantees. This reduces the memory footprint of maintaining adaptive stepsizes from $O(d)$ for $d$ the number of parameters to $O(c)$ for $c$ the number of groups.\n\nTheir second proposed method, subspace momentum, computes the momentum update in some $k$-dimensional subspace instead of the full-dimensional space, similar to existing methods like GaLore. However the authors propose simultaneously performing a step of (stochastic) gradient descent in the orthogonal complement of the selected subspace on each iteration. This allows them to obtain convergence guarantees (Thm. 4.1) in contrast with other methods which restrict momentum to some low-dimensional subspace, which cannot guarantee convergence without stronger assumptions about the span of this low-dimensional subspace relative to the objective function. This reduces the memory cost of maintaining historical momentum terms from $O(d)$ to $O(k)$. They suggest using the leading-$k$ singular subspace of a stochastic gradient sample to set the subspace, which is an intuitive rule and closely related to GaLore.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s careful consideration and constructive feedback. Below, we clarify points raised and address specific concerns.\n___\n> Question 1 and 2\n\nWe answer question 1 and 2 of the reviewer together as they are related:\n- The reviewer is correct in that arbitrarily increasing the number of subsets $c$ towards $d$ (moving from AdaGrad-Norm towards AdaGrad-Coordinate) would yield a suboptimal dependency and each n-th order term of the noise requires a more careful of balancing the subset size. Our derivation in Appendix B.2 shows the effect of the number of subsets c on all n-th order noise terms that appear in the bound. \nIn particular, Table 2 shows that the grouping of $c=d^{2/5 (1+\\beta)}$ groups of size $d^{3/5-2\\beta/5}$ each (which depends on the fraction of noisy coordinates) gives better dependency than $c=1$ (AdaGradNorm) or $c=d$ (AdaGrad-Coordinate) on most noise-density $\\beta$ setting. \n- And as the reviewer points out in question 2, indeed that in the low noise regime, AdaGrad-Norm would attain the best guarantee (second row of Table 2) and our grouping strategy of $c=d^{2/5 (1+\\beta)}$ groups would be suboptimal. There, a more careful case-work analysis to handle edge cases could help attain a more optimal strategy. \n- Regarding the reviewer’s point on $\\delta$, the probability grows with $c$ but since our guarantee is high-probability, the worst case is $log(d/ \\delta)$ so it is not the dominant term in the final bound. \n- Finally, we thank the reviewer for the attention to details and constructive feedback. We will incorporate the discussion above and revise the writing in the section after Theorem 3.1 to improve clarity and intuition.\n\n> Question 3\n\n- Our theory provides an optimal grouping strategy that depends on the noise density. However, in practice, we must trade off the cost to figure out a good grouping and the performance gain from it.\n- The key from the theory improvement is to group the coordinates with similar noise magnitudes together. However, any expensive method to figure out these groups (e.g. the Hessian in Adam-mini) would have detrimental effects on the wall clock time and memory (which are pointed out as crucial concerns by reviewer t4t5). Instead, our heuristic is meant to be a simple method to capture most of these groups.\n- Intuitively, coordinates in the same row/column either act on the same input or are used to compute the same output. The noise and normalization on each input and output would affect coordinates in the same row/columns in a correlated way. \nTo provide some evidence, we perform the experiments in Section D.1 again but with the noise grouped by the corresponding dimension according to the heuristics. We show it in the plot at https://imgur.com/a/uKfHJ1B. \n- There we see most groups have very low noise (very close to 0, namely less than 10^-12) while a small number of groups (top 1 percentile in the annotation) have much larger noises. \n- Overall, our heuristics aim to capture the similar noise coming from the same inputs and outputs. Our experiments suggest that this is a major part of the gain. There might be other simple sources of correlation in the noise magnitudes, which we leave for future work.\n\n> Question 4\n- We agree with the reviewer. If there is an effective adaptive strategy to picking the subspace dimension, then one should not count that as a hyperparameter. Unfortunately, at this time, we are not aware of an effective strategy. \n- We do not tune for the number of groups. With tuning, as in Figure 5, the performance can be improved. We will rename that column to more precisely reflect the number of tunable parameters as opposed to hyperparameters. The number of tuning runs is a more accurate measure of tuning cost but is not easily obtainable. \n___\nOther points\n> modifications also improve performance on smaller models/other architectures\n- Please see our response to Reviewer oW9K for additional experiments on vision tasks. \n\n> discuss Adam-mini in the main body so that readers are aware of the relevant work\n- We want to note that Adam-mini is concurrent work to ours, to be published in ICLR this year, but we agree with the reviewer that Adam-mini is highly relevant and including a detailed comparison in the main body improves the quality of our paper. We will include a more detailed comparison in the main body of our revision with a draft as in our response to reviewer 6d8c.\n\n> clarify what is meant by k would be an appropriate addition. \n> move norm notation to first paragraph.\n\nWe thank the reviewer for the suggestions and attention to details. We will clarify the meaning of $k$ in our revision and include all the norm subscripts in theorems for clarity in our revision. \n___\nWe sincerely appreciate your detailed reviews and constructive suggestions. We believe that the discussion and additional experiments improve the quality of our submission and hope this provides sufficient reasons to raise the score.", "claims_and_evidence": "Claim 1: Subset-norm (SN) can improve performance over either Adagrad or Adagrad-Norm.\n\nEvidence:\n- Thm. 3.1, which the authors describe as capturing a tradeoff between the number of subsets into which parameters are partitioned and the coordinate-wise stochastic gradient noise.\n- Fig. 1, which shows that Adam+SN (e.g. Adagrad+SN with full-dim. momentum) minimizes validation perplexity during model pre-training compared with Adam and also has a lower memory footprint (2.6 Gb vs 5 Gb).\n\nClaim 2: Subspace-momentum (SM) can reduce the memory footprint of momentum while still maintaining convergence guarantees.\n\nEvidence:\n- Thm. 4.1, which shows that SM is still guaranteed to converge in appropriate settings.\n- Fig. 1, which shows that Adam+SNSM minimizes validation perplexity in fewest iterations and has the lowest memory footprint out of Adam, GaLore, and Adam+SN.", "methods_and_evaluation_criteria": "They consider pre-training of LLaMa models on the C4 dataset. This seems reasonable given their motivating interest in reducing memory footprints for large-scale models. It would be interesting to see if their modifications also improve performance on smaller models/other architectures, as their claims suggest that such improvements should hold in generality, but this may be out-of-scope.", "theoretical_claims": "I looked at the proof of Thm. 3.1. The lines I checked seemed correct, and the techniques seem to be reasonable extensions of traditional convergence bounds for Adagrad under L-smoothness assumptions.", "experimental_designs_or_analyses": "I did not check the soundness beyond what is presented in the main body.", "relation_to_broader_scientific_literature": "Both of the proposed modifications, SN and SM, are closely related to existing methods. SN is closely related to Adam-mini, which also partitions parameters into blocks and maintains adaptive learning rates per-block. However the partitioning strategy in Adam-mini is related to the block-diagonal structure of the Hessian, while the theoretically-optimal partitioning strategy for SN is related to coordinate-wise stochastic gradient noise levels. However, the practical procedures for parameter partitioning proposed in the appendix are much more simplistic and do not analyze coordinate-wise noise; instead they merely group matrix-valued parameters into blocks of rows and/or columns depending on which dimension is larger.\n\nSM is closely related to low-rank learning strategies, particularly GaLore, which also restricts updates to a low-dimensional subspace chosen based on estimating leading singular subspaces. Unlike GaLore, SM involves also performing an update in the orthogonal complement of the identified subspace, allowing the authors to still derive convergence guarantees.\n\nThe authors claim that related methods (e.g. Adam-mini, factorization like Shampoo/SOAP, low-rank parameterization like LoRA and GaLore) either lack theoretical guarantees, sacrifice too much in performance, or greatly increase the cost of parameter tuning.\n\nZhang, Yushun, et al. \"Adam-mini: Use fewer learning rates to gain more.\" arXiv preprint arXiv:2406.16793 (2024).\n\nZhao, Jiawei, et al. \"Galore: Memory-efficient llm training by gradient low-rank projection.\" arXiv preprint arXiv:2403.03507 (2024).\n\nVyas, Nikhil, et al. \"Soap: Improving and stabilizing shampoo using adam.\" arXiv preprint arXiv:2409.11321 (2024).\n\nGupta, Vineet, Tomer Koren, and Yoram Singer. \"Shampoo: Preconditioned stochastic tensor optimization.\" International Conference on Machine Learning. PMLR, 2018.", "essential_references_not_discussed": "Adam-mini only mentioned briefly in the appendix, but I think the main idea behind parameter partitioning is as closely-related to SN as GaLore is to SM. I think given the strength of the connection to Adam-mini and the level of discussion that the authors devote to GaLore in the main body, it would be appropriate to discuss Adam-mini in the main body so that readers are aware of the relevant work.", "other_strengths_and_weaknesses": "Strengths:\n- The problem studied (“How can adaptive optimization be improved and made more memory-efficient for large models?”) is very well-motivated.\n- The authors focus on adding theoretical analysis to heuristics that have begun to be explored in practice.\n- The numerical results show that Adam+SMSN yields good training performance while also having a drastically smaller memory footprint.\n\nWeaknesses:\n\nMy overall impression of the paper is that its main contribution is to combine two pre-existing heuristics that already existed in literature. The authors argue that their analysis makes their approach more theoretically-grounded, but the most novel analysis about stochastic noise (Thm. 3.1) is not utilized at all in their implementation, which instead chooses a simple heuristic partitioning of parameters based on matrix rows/columns. If I am mis-understanding the implications of their theory, or not seeing how it is informing their choice of parameter grouping, then I would be happy to increase my score (see Questions).\n\n- Both SN and SM are closely related to existing methods, and the authors argue that one of the central contributions of this work is new theory that better explains these methods. Given the fact that the authors hold up the analysis as a significant portion of the contribution, I think the intended take-aways from the analysis need to be better-explained in text. For example, I found the explanation of the tradeoff between parameter partitioning after Thm. 3.1 to be confusing. \n- Moreover, the practical implementations used in the paper seem very divorced from the theory; it is my understanding that when choosing parameter partitions, the authors merely group matrix-valued parameters into either rows or columns (depending on which is larger), whereas the guarantees in Thm. 3.1. The fact that this is the proposed partition also weakens the argument made in the appendix that the SN partitioning scheme is supposedly more theoretically-motivated than the scheme in Adam-mini.", "other_comments_or_suggestions": "Parameter $k$ is referenced in “Our Contributions” before it is defined, even informally (line 90, left column). A reasonable reader can infer that $k$ must be related to the dimension of the momentum subspace, but a few words to clarify what is meant by $k$ would be an appropriate addition.\n\nThe note about norm notation (norm without subscript defaults to ell-2) is currently located at 126 left-hand column. I would suggest moving it up to the first paragraph on notation: when I was reading the statement of Thm. 3.1, I was very confused about whether the fact that some norms had subscripts and others did not meant that I should assume these norms are different. I could not find the note about norm notation because I looked in the notation paragraph, rather in the sub-Gaussian noise assumption paragraph where it is currently located. I might suggest switching to consistent norm notation within the same equation; in the statement of Thm. 3.1, all global-vector quantities have subscript-norms, while all subset quantities do not have subscripts, potentially adding to readers’ confusion.", "questions_for_authors": "As discussed in Strengths/Weaknesses, if the authors can clarify the insights provided by their theory and whether these insights are leveraged in their implementation, then I am open to increasing my score. \n\n1. Can you please explain how the result in Thm. 3.1 expresses a trade-off between partition strategies and stochastic gradient noise? It almost seems to me as if apart from the 4-th order terms in $G(\\delta)$, every other part of the guarantee gets worse as c increases: as $c\\rightarrow d$, $\\sum\\_i ||\\sigma\\_{\\Psi\\_i} ||\\_2 \\rightarrow ||\\sigma||\\_1$ whereas for $c\\rightarrow 1$, $\\sum_i ||\\sigma_{\\Psi_i} ||_2 \\rightarrow ||\\sigma||_2$ where $||\\sigma||_2 \\leq ||\\sigma||_1$. On top of that, for $\\delta$ fixed, the probability of success $1-O(c\\delta)$ gets smaller as $c$ grows, and of course all terms with explicit dependence on $c$ grow. Is my understanding correct?\n\n2. If this understanding is correct, then it seems like for in the small-noise regime (e.g. $\\sigma_i \\leq 1$), the 4-th order terms are not dominant terms in the expression, and thus all dominant terms get worse as $c$ increases, suggesting that Adagrad-Norm has the best guarantee. Is this correct?\n\n3. As I've written elsewhere, it seems to me that the parameter-partitioning strategy used in practice (group by rows/columns) is not particularly informed by the theoretical results, other than the fact that the theoretical results suggest that in some regimes an intermediate number of groups may outperform both Adagrad and Adagrad-Norm. Is this correct?\n\n4. In Table 3, it is my understanding that the column HP counts the number of hyperparameters per method. Based on this, it seems to me that the authors treat SN as not introducing any additional hyperparameters; this seems slightly misleading; the number of groups in the parameter partition/choice of parameter partitioning creates new hyperparameters. It is true that using a particular fixed strategy (i.e. the row/column grouping, applied only to 2D linear modules) fixes the hyperparameters associated with this procedure, but this is true of any algorithm. For example, the number of hyperparameters associated with GaLore can be reduced if one commits ahead of time to using a particular procedure for picking subspace dimension.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lypkxx6RX3", "reviewer": "Reviewer_t4t5", "summary": "This paper proposes two modifications to AdaGrad: 1) instead of coordinate wise adaptive learning rate, subset norm adaptative learning rate can provide adaptive learning rate scaling for different subsets of parameters 2) similar to GaLore, it keeps momentum in low rank and recovers momentum by up-projection, whose projection matrix is periodically updated via SVD.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s careful consideration and constructive feedback. Below, we clarify points raised and address specific concerns.\n____\n> Theorem 3.1 is highly dependent on the variance of the stochastic noise\n\n- The reviewer is right that the bounds depend on the noise and this is standard for any bounds on stochastic gradient descent. The novelty is that we show grouping the coordinates based on noise magnitudes can lead to a better performance. Our experiments show that our heuristic successfully isolates the noise into a small number of groups and perhaps as a result, delivers better training results.\n\n> Theorem 4.1 has no dependency on k the rank of the subspace\n\n- This is a limitation that we admit in the discussion section of the paper: the worst case convergence of SM is simply the SGD with momentum bound. We believe that a better analysis (that probably depends on a better subspace choice rather than the worst case one) could potentially show improved rates over naive SGDm. Our experiments show promising results for the top-k singular vector subspace so we believe that this is an interesting point for future works to show theoretical improvement.\n- And similarly to our response to Reviewer 6d8c, this is a limitation that is general to the class of algorithms that utilize momentum. \nThis is an important open challenge for the community at large because the theoretical bound for momentum and without momentum are the same. We believe that any advance for SGD with momentum will lead to a better understanding of SM in our setting. \n\n> Besides Galore, a few of its follow-ups are highly relevant to the discussion too.\n\n- Please refer to our response to Reviewer 6d8c. We will add additional comparisons to Table 3 of our revision and a more detailed comparison with our Adam-mini in our main body.\n\n> I would like to see wall clock speed up on pretraining task as well as peak memory for different implementations.\n\nWe provide the per iteration time, peak memory (via nvidia-smi), and time to Adam’s val perplexity after 100K steps for the 1B model for each method on a 2x4090 machine with the same setup as the paper (seq length 256, total batch size 512, micro batchsize 16) below:\n|                     \t| Adam    \t| AdamSNSM  (Gap=5000) | AdamSNSM  (Gap=200)| AdamSN                 \t| GaLore   (Gap=200)|\n|-------------------------|-------------|----------------------------|----------------------------|----------------------------|----------------------------|\n| Time for 1K iters  \t| 7426s   | 7465 s            |  7624 s\t| **7399s**              \t| 7827s              \t|\n| Time per iteration  \t| 7.43 s/it   | 7.47 s/it     |         7.62 s/it\t| **7.39 s/it**              \t| 7.83 s/it              \t|\n| Time to perplexity < 16 | ~206.4 hrs  (100K iters)| ~136.9 hrs (<66K iters) | **~101.6 hrs (<48K iters)**  | ≈118.9 hrs (<58K iters)   | >217 hrs (>100K iters)\t|\n|Peak mem\t| 21.554 GB/GPU | **16.642 GB/GPU**    | **16.642 GB/GPU**     \t| 19.193 GB/GPU          \t| 18.187 GB/GPU          \t|\n\n\n- The results above show that the dimensionality reduction can achieve some speedup and the improved performance can offset some overhead of SVD. Furthermore, as demonstrated in Table 6 on page 8, our methods do not require the SVD update gap to be as frequent as GaLore. \n- Peak memory for batch size 1 is also shown in Figure 6 on page 16. \n\n> Another setting that is important but not really discussed in the paper is how friendly the proposed method is to distributed training.\n\n- Thank you for bringing up this great point. Our non-heuristic version of the SN algorithm (Algorithm 5) is FSDP-friendly due to being shape agnostic and can group coordinates locally. Our supplementary material contains the code for that version (`adamw_sng.py`). Figure 5 shows that local grouping with group size of around $\\sqrt(d)$ gives good results. \n- A brief discussion regarding SM’s distributed training is on line 932 where it is more subtle and depends on the subspace choice where subspace that aligns with the standard bases would work on distributed setup but there are other tradeoffs.\n- We will include additional discussion regarding this point in our revision.\n\n> This method also introduce extra hyperparameter tuning... \n\n- We do not tune for the block size as it is simply the row/column of each matrix. While we show that tuning can bring additional benefits as shown in Figure 5, we can already attain a large part of that gain with no tuning via our heuristics on line 213. Finally, due to the modularity of our approaches, we can even abandon the use of momentum and still achieve strong performance relative to vanilla Adam with RMSPropSN as shown in Table 3 and 8 (versus vanilla RMSProp). \n___\nWe sincerely appreciate your detailed reviews and constructive suggestions. We believe that the discussion and additional experiments improve the quality of our submission and hope this provides sufficient reasons to raise the score.", "claims_and_evidence": "The experiments and ablations are very solid. I have nothing to complain about.", "methods_and_evaluation_criteria": "The method makes sense and evaluation all follows current common practices. I don't see anything suspicious in particular.", "theoretical_claims": "I checked Theorem 3.1 and 4,1, they seem correct. Though I have not carefully checked the proofs in appendix, the conclusion fits my expectation to the convergence rate, which is similar to that of the SGD's convergence under smooth conditions.\n\nOne complaint here is Theorem 3.1 is highly dependent on the variance of the stochastic noise, which might render the bound vacuous.\n\nTheorem 4.1 has no dependency on k the rank of the subspace, I would imagine it should because when k is full rank it should exactly recover the SGD baseline case.\n\nThere is also no analysis on how the projection matrix would impact the convergence.", "experimental_designs_or_analyses": "The experiment designs follow the common practice and ablation studies are abundant, nothing much can be complaint about.", "relation_to_broader_scientific_literature": "this work is related to a series of optimizers proposed to pretrain large language models faster and more efficiently, the subspace momentum can be seen as a form of preconditioning, which fits in the broader optimizer literatures, such Muon, MARS and Lion.", "essential_references_not_discussed": "Besides Galore, a few of its follow-ups are highly relevant to the discussion too. I believe it will benefit the readers if the authors can discuss the difference and commonality compared to those.", "other_strengths_and_weaknesses": "The proposed SN and S momentum are simple and straightforward, though they are not really new or surprising. especially subspace momentum, this is an incremental on top of GaLore.\n\nFor this method to be truly convincing, I would like to see wall clock speed up on pretraining task as well as peak memory for different implementations. Despite of momentum being saved in low rank, it might still have to materialize larger tensor due to intermediate computation of residual r, which cancels out the benefit of lower memory consumption.\n\nAnother setting that is important but not really discussed in the paper is how friendly the proposed method is to distributed training.  \n\nThis method also introduce extra hyperparameter tuning, default hyperparameters of AdaGrad can't be reused and non-trivial effort to find the optimal block size and etc. It also inherits all the problems of GaLore in practice: SVD can become a real bottleneck when weight matrices become large, randomized/approximated SVD has no guarantee that it wouldn't degrade the perplexity as shown in appendix.\n\nNonetheless, I am still leaning towards acceptance due to the comprehensive ablations and detailed guidelines in hyperparameter choices and etc. Theoretical propositions also seem solid and interesting, though I am not completely familiar with the mathematical tools/framework used in this paper.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7nbO1F2PVl", "reviewer": "Reviewer_oW9K", "summary": "This paper introduces two complementary adaptive optimization techniques that reduce the memory footprint of optimizer states while accelerating LLM large-scale neural network training.\n\n- Subset-Norm (SN): A generalization of AdaGrad-Norm and AdaGrad-Coordinate that shares step sizes across subsets of parameters. SN reduces AdaGrad’s memory footprint from $O(d)$ to $O(\\sqrt{d})$ while maintaining strong theoretical convergence guarantees under sub-Gaussian noise assumptions.\n\n- Subspace-Momentum (SM): A technique that applies momentum in a low-dimensional subspace while using SGD in the orthogonal complement. The authors prove high-probability convergence guarantees for SM under standard assumptions.\n\nEmpirical results demonstrate that combining SN and SM achieves Adam’s validation perplexity on LLaMA 1B with half the training tokens (6.8B vs. 13.1B) while reducing Adam’s optimizer state memory footprint by more than 80%, with minimal additional hyperparameter tuning.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s careful consideration and constructive feedback. Below, we clarify points raised and address specific concerns.\n___\n> No discussion of quantization-based approaches like 8-bit Adam\n\nDue to space limit, we cite and discuss quantization approaches (and more) in the related works section on page 13 of the Appendix. This is an important orthogonal direction that can be combined with our method for further future improvement.\n\n> Question 1: MicroAdam and Adam-mini comparisons\n\nWe provide brief comparisons to MicroAdam and Adam-mini in the related works section on page 13, but will provide more detailed comparison for Adam-mini due to its relevance as also pointed out by other reviewers in our revision. \n- For MicroAdam, while it delivers good performance for fine-tuning, it has a limitation for the pre-training LLM setting (see page 10 of their paper) which is a setting we want to focus on due to the relative lack of effective memory efficient optimizers.\n- For Adam-mini, we want to note that it is concurrent work to ours, to be published in ICLR this year, but we agree with the reviewer that Adam-mini is highly relevant and including a detailed comparison in the main body improves the quality of our paper. We will include a more detailed comparison in the main body of our revision with a draft as in our response to reviewer 6d8c.\n\n> Question 2: vision tasks\n\n- While our main focus is on large models that are more typical to language models where memory is often a bottleneck, to address the reviewer’s request for additional experiments, we conducted further evaluations using the DiT-L/2 model (458M) from [facebookresearch/DiT](https://github.com/facebookresearch/DiT) on a setup with batch size 2048, image size 64, and 8×A6000 GPUs. \nWe compared our method (SNSM) with Adam using the same training configuration as in the paper. As shown in the table below, SNSM outperforms Adam in FID similarly to LLM tasks:\n\n| FID/Iter | 200k  | 300k  | 400k  | 500k  | 700k  |\n|-----------|-------|-------|-------|-------|-------|\n| Adam | 56.69 | 56.63 | 40.69 | 41.15 | 39.61 |\n| AdamSNSM | 66.76 | 66.31 | 34.05 | 32.31 | **32.26** |\n\n- We further evaluate Adam, AdamSN, and AdamSNSM (rank 64 and no update gap) by training “vit_base_patch16_224” (~85M params) from the `timm` library on the `CIFAR100` dataset for 10 epochs with bs64 and wd0 on a 2x4090 machine. We tune the lr across {1e-3, 5e-3, 1e-4, 5e-4, 5e-5, 1e-5} grid for all methods.\n\n| | Adam  | AdamSN  | AdamSNSM  |\n|-----------|-------|-------|-------|\n|Best Val Acc| 43.3%| 45.2% |  **45.6%**| \n\n- Please see the loss curve here: https://imgur.com/a/4rtgxOx. These preliminary experiments show promising results for the application of our methods to vision tasks.\n\n> Question 3: numerical precision\n\n- We use BF16 as mentioned in Appendix B.2 and ensure that all our experiments across different optimizers use consistent numerical precision for fair comparisons. Due to space constraint, a lot of the experimental setup and implementation details are moved to the Appendix. We will make this clearer in the main body of our revision.\n\n> Question 4: tune hyperparameters\n\n- Tuning effort is provided in Appendix B.2. We use a similar setup (rank, projection update gap, etc.) to Zhao et al 2024 and try to minimize tuning (in the spirit of this resource constraint setting) and only tune for the learning rate. However, we perform extensive ablation studies (Section 5.4 and Appendix C) to show that performance can further be improved with additional tuning. We are grateful that reviewers t4t5 and 6d8c looked at these ablations and considered them \"very solid\" and \"quite extensive and contains important details.\"\n\n> Question 5: red LR\n\n- The red LR simply highlights the results where the best LR found for the larger model differs from the smaller model and hence tuned similarly to the other LR. It is desirable for hyperparameters to transfer from smaller models to larger models. However, this is a complex topic in general (e.g. https://arxiv.org/abs/2203.03466).\n\n> Question 6: SN heuristics\n\n- On line 213 of the paper, we provide a simple heuristic of grouping the latent dimension together.  In Figure 5, we examine different subset sizes where one can tune to get a better result. However, our heuristic seems to be able to capture much of the performance. In general, we find that choosing a subset size around sqrt(mn) for each block seems to be a good heuristic and balance that is supported by our analysis. \n\n___\nWe sincerely appreciate your detailed reviews and constructive suggestions. We believe the discussion, analysis, explanations and additional experiments clarify your concerns and improve the quality of our submission. We hope this provides sufficient reasons to raise the score.", "claims_and_evidence": "The claims in the paper are generally well-supported by:\n\n- Theoretical analysis: The authors provide strong mathematical guarantees for both SN and SM. The convergence proofs are well-presented and adhere to standard assumptions in stochastic optimization.\n- Empirical validation: The experiments on pre-training and fine-tuning large-scale LLMs demonstrate clear improvements in memory efficiency while maintaining or improving convergence speed.\n- Comparison to baselines: SN and SM are compared against AdaGrad, Adam, and GaLore, showing consistent improvements in both perplexity and memory efficiency.\n\nHowever, some additional validations would strengthen the claims:\n\n- CV benchmarks : The proposed methods should be tested on ImageNet/CIFAR-10 with ViTs or ResNets to evaluate performance in a broader range of architectures and compare with well-established optimizers in vision tasks.\n- Robustness: The paper lacks hyperparameter robustness analysis, which would help practitioners understand how sensitive SN and SM are to tuning.\n- Comparison with recent memory-efficient optimizers: The paper does not compare against newer low-memory optimizers such as Adam-mini or MicroAdam, which have been proposed to reduce the memory footprint of optimizers.", "methods_and_evaluation_criteria": "The proposed SN and SM methods are well-formulated and grounded in existing literature on AdaGrad-style optimizers and gradient compression techniques. \n\nThe authors evaluate SN and SM across multiple model sizes (from 60M to 7B parameters) and compare optimizer memory footprints, training perplexity, and convergence rates. The evaluation is relevant, but additional benchmarks (e.g., ImageNet, CIFAR-10) would improve generalization.", "theoretical_claims": "The convergence proofs for SN and SM seem correct.", "experimental_designs_or_analyses": "- The experimental design is strong, with multiple scales of LLaMA models tested.\n- The results clearly demonstrate the memory efficiency of SN and SM, particularly on LLM pre-training tasks.\n- However, there are no details on the numerical precision used during training (e.g., FP32, BF16, FP16). Clarifying this would be important since precision affects memory consumption. \n- Hyperparameter tuning details are unclear:\n\t- How were the “red LR” values tuned? Do you adjust them based on instability?\n\t- How should practitioners tune SN and SM hyperparameters? A guideline for selecting subset sizes and subspaces would be helpful.", "relation_to_broader_scientific_literature": "The paper is well-positioned within the literature on adaptive optimizers (Adam, AdaGrad) and memory-efficient optimizers (GaLore, FLORA, GRASS). However, no comparisons are made to newer memory-efficient optimizers like MicroAdam or Adam-mini, which should be included.", "essential_references_not_discussed": "Missing citations to recent memory-efficient optimizers such as Adam-mini. \n\nNo discussion of quantization-based approaches like 8-bit Adam, which also aim to reduce optimizer memory.", "other_strengths_and_weaknesses": "**Strengths:**\n- Well-motivated problem with significant practical relevance for training LLMs.\n- Strong theoretical backing with convergence guarantees.\n- Empirical results demonstrate state-of-the-art memory savings.\n\n**Weaknesses**:\n- No evaluation on computer vision benchmarks (ViT, ResNet).\n- No comparison with MicroAdam or Adam-mini.\n- Lack of clear recommendations for tuning SN and SM hyperparameters.\n- Inconsistent notation between tables (Adam vs. AdamW).\n- No discussion of numerical precision used in training.", "other_comments_or_suggestions": "The notations in Table 3 and Table 4 are inconsistent (Adam vs. AdamW). This should be clarified.", "questions_for_authors": "(1)\tHow does SN and SM compare to recent memory-efficient optimizers like MicroAdam and Adam-mini?\n\n(2) Can SN and SM be effectively applied to ViTs and ResNets for computer vision tasks?\n\n(3) What numerical precision was used during training (BF16, FP16, FP32)?\n\t\n(4) How do you tune hyperparameters for SNSM?\n\t\n(5) How was the “red LR instability” tuned?\n\t\n(6)\tWhat heuristics can be used to select the optimal subset size in SN?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This submission considers two modifications to adaptive algorithms for training neural networks, e.g., the adagrad/adam family of algorithms, which can reduce the memory footprints of these algorithms. There is both theoretical analysis in the form of convergence guarantees for their modified algorithms and there is also experimental validation in the form of pre-traing/fine-tuning language models and CIFAR100 classification; in both settings comparisons are made with baselines like adam, GaLore, or adagrad. The empirical benefits in terms of memory are clear and the performance as an optimization algorithm seems reasonable as well, with many ablation experiments in different settings (language, vision).

In the rebuttal period, the authors addressed the majority of the concerns of the reviewers. The main sticking point seemed to be comparison with adam-mini, recently proposed optimizer for the same tasks that also claims a reduced memory footprint. To this, the authors agreed to cite and discuss adam-mini in their paper.

The other sticking point was the fact that the analysis presented is not completely connected to what is done in practice; indeed the analysis relies on certain knowledge about the stochastic noise in the gradient computations that is typically not available nor easy to estimate and a different heuristic is used in the actual implementation of these methods. Although I agree that it would be a much stronger contribution if the theoretical results could characterize what is done in practice, I do not see this as a strong negative for the paper because of the accompanying experiments, which are extensive.

For these reasons, I recommend to accept this paper.
Decision: Accept (poster)