Memory-Efficient LLM Training with Online Subspace Descent
Authors: Kaizhao Liang
Bo Liu
Lizhang Chen
qiang liu
Venue: NeurIPS 2024 poster (2024)
Keywords: Large Language Model Pretraining; Optimizer; Memory-Efficient LLM Training
Primary area: optimization_for_deep_networks
TL;DR: a new family of optimization algorithms adaptively descending on projected subspaces found by online PCA

Abstract: Recently, a wide range of memory-efficient LLM training algorithms have gained substantial popularity. These methods leverage the low-rank structure of gradients to project optimizer states into a subspace using projection matrix found by singular value decomposition (SVD). However, convergence of these algorithms is highly dependent on the update rules of their projection matrix. In this work, we provide the \emph{first} convergence guarantee for arbitrary update rules of projection matrix. This guarantee is generally applicable to optimizers that can be analyzed with Hamiltonian Descent, including most common ones, such as LION, Adam. Inspired by our theoretical understanding, we propose Online Subspace Descent, a new family of subspace descent optimizer without SVD. Instead of updating projection matrix with eigenvectors, Online Subspace Descent updates projection matrix wtih online PCA. Online Subspace Descent is flexible and introduces only minimum overhead to training. We demonstrate that, for the task of pretraining LLaMA models ranging from 60M to 1B parameters on the C4 dataset, Online Subspace Descent achieves lower perplexity than state-of-the-art low-rank training methods across different settings and narrows the gap with full-rank baselines.

Memory-Efficient LLM Training with Online Subspace Descent
Kaizhao Liang', Bo Liu", Lizhang Chen', Qiang Liu'
*The University of Texas at Austin (kaizhaol, lzchen, ,qiinn) @utexas. (qtang/uutexxs.edd
Abstract
Recently, a wide range of memory-efficient LLM training algorithms have gained substantial popularity. These methods leverage the low-rank structure of gradients to project optimizer states into subspace using projection matrix found by singular value decomposition (SVD). However, convergence of these algorithms is highly dependent on the update rules their projection matrix. In this work, we provide the first convergence guarantee for arbitrary update rules of projection matrix. This guarantee is generally applicable to optimizers that can be analyzed with Hamilto- nian Descent, including most common ones, such LION, Adam. Inspired by our theoretical understanding, we propose Online Subspace Descent, a new family of subspace descent optimizer without SVD. Instead of updating the projection matrix with cigenvectors, Online Subspace Descent updates the projection matrix with online PCA. Online Subspace Descent flexible and introduces only minimum overhead to training. We show that for the task of pretraining LLaMA models ranging from 60M to 7B parameters on the C4 dataset, Online Subspace Descent achieves lower perplexity and better downstream tasks performance than state-of- the-art low-rank training methods across different settings and narrows the gap with full-rank baselines.
Introduction
The continual advancement in training large language models (LLMs) presents a compelling challenge in balaneing computational efficiency with model performance. As the scope and complexity of these models grow, does the necessity for innovative strategies that optimize memory usage without compromising the learning capabilities the model. Recent approaches
low-rank adaptation strategies, including Stochastic Subspace Descent [13], LoRA [11], ReLoRA [15], Gradient Low- Rank Projection (GaLore) [25] and Sketchy [9], have paved the way for memory-efficien training by utilizing periodically updated low-rank projection matrix to manage parameter updates. In particular, GaLore and Sketchy both utilize expensive singular value decomposition to determine the projection matrix, whereas stochastic subspace descent suggests using random matrices projection matrices and provides convergence analysis on convex functions and objectives. However, to the best of our knowledge, one has offered any guarantee of convergence for this class of methods non-convex functions and objectives.
In this work, we provide the first convergence guarantee for arbitrary update rules of the projection matrix. This guarantee is significant because it broadly applicable to wide range of optimizers that can be analyzed within the Hamiltonian descent framework [18]. By establishing this convergence guarantee, we demonstrate that our approach not limited to specific or narrowly defined update rules, but can be extended to include many commonly used optimizers in the field. In particular, this
Code is available at https: com. wtt.ghtwww..oo//gglllaagggg///// www.iinnn-uaaeeeeeeeeeeee Descent.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Algorithm 1 Online Subspace Descent
1: Required: Optimizer OptimizerW, learning rate e!", weight decay \W for model weights W and |OptimizerP, <f X²) for the projection matrix Pt. Proper initialization.
2: for iteration t do
3: Calculate gradient G, = ()) Update model weights W, by
(Ât; : 5.) Optimizemeiiiee.. P. G.,, S.-),, Wi- W, + e (P, -
4: Calculate G. = VLG, for LG. (-) in Eg (6); Update the projection Pi by
, S..)) OptimizerP(Gi, 1 P4+1 P P (A XP)
5: end for
6: Remark: We added weight decay as common heuristic. We recommend using Adam for both optimizers, and set e," =ae," W with constant o (e.g., 5), and W XP.
includes popular algorithms such LION [4] and Adam [12], which widely used in various machine learning and optimization tasks. Our results therefore offer robust theoretical foundation for understanding and analyzing the behavior of these optimizers, ensuring their effectiveness and reliability in diverse applications.
Inspired by theoretical understanding, we introduce a novel family of memory-efficien opti- mizers named Online Subspace Descent, which incorporates a dynamically changing projection matrix, replacing the conventional periodic updating approach (SVD) with online PCA. By allowing the projection matrix to evolve in response to the changing gradient landscape, Online Subspace Descent enhances the model's 's ability navigate the parameter space more effectively. This dynamic adaptation aligns more closely with the natural progression of learning in deep neural networks, which is characterized by changes in the importance of different characteristics and interactions as training progresses. Through extensive experiments and comparative amalysis, we demonstrate that our approach presents lower perplexity in pretraining LLaMA models (ranging from 60M IB parameters) on the C4 dataset compared to state-of-the ofathe-r-te--rt low-rank training methods, closing the perplexity gap with full-rank baselines on language model pretraining.
2 Optimization Background
The training of deep learning models reduces to an optimization problem
min L (W),
W
where W is the set of weight matrices of the model. For simplicity, we assume W Rexm is single matrix of size (n, m) without loss of generality. For notation, we write (A, B) = tr(A B)
for inner products of matrices, and ||AA]]2 tr(A] T A) the Frobenius norm. We use A O B to denote the elementwise product, and A
A.
Example 2.1. Update rules of common optimizers:
Gradient Descent
W W. - v.,,,,
Momentum
W4+1 W. E,M,,
M. (1- +
Lion-K [4]:
W W. - -VeevVVKNNN VK(NN),
N; (1 (W))
M (1 32)VL( 32Mr-1,
M.
Adam [12]: W W. E VV.+e
M. (1- (I.)VL(W,)) + M,,M,
V, = - 02 + B2e V--1 i-1,
where e, step sizes, and M., Ve the first and second order momentum, and B, Br, B2 are 8. g
momentum coefficients in (0, 1), with 3
1- 1 r, i for 31, E (0, 1) in Adam, and K.
is any convex function with (K()) = 0 for Lion-K [4], and Lion [5] uses K(X) = X 1111 and
VC((X sign(X:

---

These optimizers can be unifiedly viewed as updating W together with optimizer state Sy: S. ( VL(W 1)),
W++1 t+1 Wt (5(e)
with some mapping de, th. We have S, = M, for momentum and S. = (M,, V.} for Adam. Note that both M, V are of the same size the model weights Wt c. resulting high memory consumption for large models. This issue is particularly pronounced for Adam, which typically yields the best performance for large language models (LLMs) but incurs the highest memory cost due to the need to maintain both Mr and V.. One key challenge is to retain the high performance of Adam while enhancing its memory efficiency.
Hamiltonian+Deseeee One powerful approach to studying the dynamic properties of optimizers is examine their continuous-time ODE forms in the limit of infinitesimal step size. The continuous- time forms provide clearer insights into the asymptotic convergence of the algorithm, abstracting away the choices of step size, discretization, and stochastic errors. The underlying logic is that "sound" optimizer should be guaranteed to converge to local optima of the loss, at least when using sufficiently small step sizes.
Inspired by [4, 18], we observe that the continuous- time form of many common optimizers yields a Hamiltian+ Descent structure,
d
W. OsH(WW S.) (Ow H(W:: S.))
dt d
S. -ö H (Wa 6. S.) (0 (W S,)),
dt
where HWW,S) .S is Hamiltonian (or Lyapumov) function that satisfies
min HWW .) = L(W), VW,
S
so that minimizing L(W) reduces to minimizing H H(W S); and (--), v-() are two monotonic mappings satisfying
| 11, (X)
| - (x, V (x) 0, VX.
With x(X) = (X)) O, the system in (2) reduces to the standard Hamiltonian system that keeps H (W., S)) const along the trajectory. When adding the descending components with & and v. the system then keeps H(W,S) S) monotonically non-decreasing:
d H(W.,S.) dt
d
HsH,, S. dt (Öw He, 0g H, 45 (Bw H,)) (0s Ow H, ((gHH,))
@ow He. dt W.
(3)
- =-1@aw H - H, 1| 00
where we write ow H, = Hw ((W.,S.) and similarly for Hg The main idea is that the cross terms (Ow Hi, Hi) canceled, leaving only the negative terms.
Example 2.2. The momentum method yields following continuous-tin -time form and Hamiltonian:
d W. = M, dt M. = (LVL((WL().) -M,), with H(W, MM = L(W) + ||
dt
2a
Example 2.3. Adam yields the following continuous-tin time form and Hamiltonian,
M.
d V. = b(V V.).
W.
M, a(VL(W.)) -M.),
di
W.+e*
di
di
M
with
H(W, M, V) = L(W)= + 2a VV +e
M
for which we can show that d HWW, i, M, V.) 0 when 2 b/4. di Example 2.4. The Lion-K optimizer [5, 4] (without weight decay) be written into
d MM = A.M.=a(VL(W>)) .) + M,),
w. = vK((1 - b)M, - (V ))) dt
di
where a 2 0, E 1] and K(X) is any convex function that attains the minimum at X 0. One of its Hamiltonians that yields the Hamiltonian+ structure (Eq (13) in Chen et al. [4]) is
H(W, M) aL(W)
1 xC(11 - -b)
3

---

3 Memory-Efficient Optimizers via Online Subspace Descent
We introduce the idea of constructing memory efficient optimzers by descending in the subspaces that dynamically changes across iterations as motivated by GaLore [25] and Sketchy [9]. We first derive static subspace descent by restricting the whole optimization on a subspace (Section 3.1). and then propose to dynamically change the subspace across iterations as a heuristic to attain the optimization the full space while only using subspace descent (Section 3.2). particular, we propose to update the subspaces via continuous online PCA like updates to avoids the need of exact SVD like in GaLore and Sketchy (Section 3.2). Finally, we remark Section 3.3 the heuristic nature of the derivation of the method and highlight the difficulty in theoretical understanding, which motivates our analysis based on Hamiltonian dynamics in Section
3.1 Static Subspace Descent
One popular approach to improving memory efficiency is to confine the optimization to a low- dimensional space. To do this, we impose a low rank structure W = PW, where P E BRn is projection matrix to determined later and W E RR is a dimension-reduced parameter. When k < n, P and W are much smaller in size compared to W. Now consider
min L(PW).
W
Applying the optimizer from (1) update w along with an optimizer state S, and mapping the update rule W 1+1 W, (5.) that of W PW. c, we get
W x+1 W. + (S.).
S. v. VLWW.)))
(4)
where we used the fact that V (PW) P V LWW.. This yields more memory-efficient optimizer, the size of S, is proportional to that of W, much smaller than S in (1) when < n.
3.2 Online Subspace Descent
With a static P. regardless of its values, the parameter W is restricted to have a low rank structure. Although low rank assumption is proved to be useful for fine-tuning with LoRA-like methods [1]], it often too limited for pre-training or when the desirable model weights are not inherently low-rank.
To address this problem. Zhao et al. [25] suggested to keep the projected updated in (4), but use different P across the iterations:
W 6+1 W. + (S));
b V.W.(W.). P = P (P,,, W. S.),
where Xt is a update rule of P, that will be determined in the sequel. The intuition is to open up different projection directions at different iterations, that optimization can be conducted in different subspaces across different iterations. This is similar to the update of coordinate descent, except in continuous fashion. Note that the update of P, can be done in parallel with that of (W,, S,), and incurs no slowdown once it fast enough
cause a speed bottleneck.
Example 3.1. Examples of common optimizers equipped with online subspace descent:
Gradient Descent
G., W t+ W. .P..P ..P.G.
G. VL(W.),
Momentum:
W41 = W. - -GP,M
M, = (1 3)P. G. + +3M.-1
Lion-K:
W W. -EP..K((N.),
G. P..
N. = (1 -B:)ar + S,M,,
M, = (1 1-32) + 32Mi-1
M.
Adam
W W. - -.P VV+t
= G.PP..G.
M, (1 G. + BI0 M-1
V. (1 f2e c.
22
Bzg V 1-1-
How Should P, be Updated? It is useful to draw intuition from the projected gradient descent rule
W = W. G.P.T.G.
G. = VLW.))
55)

---

in which P,P, can be viewed as a low rank preconditioning of Gg. To make it follow the exact gradient descent, we hope to make P; P,P.GG G, approximate G, much as possible. In Galore, this is achieved by performing singular value decomposition (SVD) on G, periodically every T iterations:
P - = torch. linalg. svd GT(//)))
where TLt/TJ is the largest multiple of T less than or equal to L. However, numerical SVD incurs large computational cost for very large models. Also, since P. is fully determined by GTU/TI calculated from a single mini-batch the last periodic point, it does not incorporate the gradient information from all data in timely fashion.
In this work, propose to update P. in a continuous online fashion that incorporates the most recent gradient information in a timely fashion, without calling torch.linalg decompossoossssss routines. We view the update of P. as conducting an online principal component analysis (PCA) based on the streaming of [Gr}. In particular, we propose to update P. at time t by minimizing the following PCA objective:
LG, (P) PP G.l +1 PAPP Toxa| G.
G |G
(6)
where |||| = tr(AA A)1/2 and Tack is the k x l identity matrix; we introduced an auxiliary loss to encourage the columns of P to be orthonormal and normalizes G, increase stability.
The key property of LG, (P) in (6) that all its stable local minimum is global minimum, and P is a global minimum iff PP G. forms the optimal rank-k approximation of G. [e.g., 3]: moreover, we have P P = Lxck at optima when 0.
Instead of minimizing LG, (P) exactly, to retain computational efficiency, we propose to update Pt by only performing one step of optimization on LG. (P):
P =
VpLG,
where OptimizerP. sttep can be a favorite optimizer, such as gradient descent or Adam. Note that when using Adam, we introduce copy of optimizer state S," for P.. See Algorithm 1. Compared to the exact SVD, each online update of P, here fast and can be executed parallel with the (W, .. 5:) updates to avoid slowdown.
3.3 Difficulty in Theoretical Understanding
The idea above of projecting an arbitrary optimizer with a dynamically changing P, is heuristically motivated and lacks immediate rigorous theoretical justification. The main challenge lies in the complex interaction between the update of U, and the optimization state S., which could potentially degrade the convergence and other theoretical properties of the original optimizer. A key question is whether we can develop a theoretical framework to understand how P. impacts the optimizer's convergence behavior and provide guidance for the design of the update rules of Pt-
To gain understanding, it is useful to first exam the simple case projected gradient descent in (5) which does not have an optimizer state (S, = 0). In this case, since P.P. i is positive semi-finite, the update P.P. G3 is always non-increasing direction of L(W) for any P,. The algorithm is essentially variant of coordinate subspace descent, where P. defines the subspace step of gradient descent is conduced at iteration f. To ensure that (5) finds a local optimum, we mainly need to ensure that P,G, 0 only if G; 0 to prevent the optimizer from stopping prematurely; this is mild condition that can be satisfied e.g. when P: is updated by (online) PCA on G,.
which one
Unfortunately, this coordinate-Oescente that track a momentum state S. This because S; accumulates the information from the projected gradient P,G, all earlier iterations < t. As P, changes across time, it is unclear whether the gradient projected to different subspaces P- would coherent with each other, and useful for future updates that are conducted in different subspaces P, for > The difficulty is the inertia effect of S, that entangles the different subspaces, making the dynamic behavior fundamentally more complicated than naive coordinate descent where the descent in different subspaces is uncoupled. This is what we address in Section via the Hamiltonian descent framework.
interpretation does not apply to more advanced optimizers
5

---

Hamiltonian Descent Meets Subspace Descent: A Lyapunov Analysis
In this section, we show a surprising result that the complication outlined above in Section 3.3 not a problem for optimizers that yields the Hamiltonian+dessent structure in (2). Our result is two-fold:
When applying Online Subspace Descent on systems in (2), the Hamiltonian-descent structure is preserved once the update rule of Pt has a smooth continuous-time limit. Hence, under very mild conditions, Online Subspace Descent equipped with common optimizers like Adam and Lion automatically yield Lyapunov function and hence benign continuous convergence. Moreover, P. can, fact, be generalized to an arbitrary linear operator as shown in Section 4.3.
Section 4.1:
Section 4.2: For any smooth P. update rules that eliminates the degenerate case of P G, 0 while G; = at convergence, the online subspace optimizer guarantees to converge in continuous time stationary point of the loss L(W). This mild condition is satisfied, for example, when Pa is updated by a typical optimizer on the online PCA objective Lg, (P).
4.1 Online Subspace Descent Preserves the Hamiltonian+ Desceen Structure
Applying dynamic projection to Hamiltonian descent in (2), we obtain the following systems:
d
W. P,0g S.) (8w HW,,, (.))
dt
d
5, p Hw H( Si) (W. S.)
dt
d P. I(Pi, (W)..))
di
where I specifies the update rule of P.- Following essentially the same derivation (3), one can show that H (W, S) remains Lyapunov function of (7), regardless of the choice of I':
d H(W, 2, S.) - 11ów H1 3/2 |8 & + (Ōw H,, P,OgH)) - (0gHr. P. Ow Hi) dl
(8)
1|1 11,13 11 S O,
where the key is to use the adjoins property of P and P that (PiX,Y) = P..P. cancels the crossing terms, independent of the values Pt. There is no requirement on I here, besides that the derivative in (8) should exist. As shown in Section 4.3, we can generalize (8) by replacing P. and P. with general linear operator P, and its adjoint Pt
which
Please refer to Appendix A for continuous-tine Momentum, Lion-K and Adam with subspace descent and their Hamiltonian functions.
4.2 Convergence to Local Optima
In addition to the Lyapumov structure, we need additional mild condition the update rule of P. to ensure the system converges to the local optimum of the loss L(W). The main idea is to prevent the system from stopping prematurely before reaching the degenerate case of P,G, = O while G, + O in the invariant set of the system. Assumption 4.1. Assume the functions in system (7) are continuously differentiable and
gradient G, = 0 by excluding
de (W,,, ,, 3.) = 0 implies G. P. L(W.)) x) 0 and गt W.=0. 00
ti) When G: G + the set (P: P G 0} is not a positive invariant set of àt. P. PT=T(PP,, Gt).
This is a mild condition. Assumption i) says that the optimizer should stop at point with G; which is easy to verify for the common optimizers like momentum, Adam, Lion-K. Assumption ii) ensures G; = 0 would imply G, = 0 in invariance sets, which is satisfied when for example, P; is updated by a reasonable optimizer of the online PCA loss that converges to stable local minimum.
Theorem 4.2. Assume Assumption 4.1 holds. Let (W,, P.). be bounded solution of (7), then all the accumulation points (W,) as t - +00 are stationary points of (())..
Proof. By LaSalle's invariance principle, the positive limit set of (W,. (, S, Pt), must be contained in I, where I {the union of complete trajectories satisfying 4 H (W, c. S.) 0, Vt

---

From the Assumption i), the trajectories contained in I must satisfy dr dt GG di WVW and G, 0 and hence G; G is a constant with P G O. Moreover, from Assumption ii), we must have (L( invariant. As result, all trajectories in the limit set I must have ( = O. Because di these trajectories are static points of Wt.
d W. = 0. which implies
G2 = 0, since otherwise the trajectory is not
4 W 0.
4.3 Online Subspace Descent with General Linear Projection Operators
We can generalize the online subspace descent with general linear operators:
W. P. ..)) @www (W, S.)
dt S (Ow H(W (.)) V(Og S.)
P. I(P,, ()).))
dt
where generalize P, to be any linear operator P. with adjoint operator P;, satisfying
(x, (XP)( (X), Y), Y.
The derivation of Lyapunov follows a similar way:
d H(W,, s, S.) 1@0 H, 1 | 1/ + (Öw H, P.(Og H., Pt (Ow H.))
dt
- 1 1 2 (0sHH S 0.
where the crossing terms are again canceled due to the adjoint property.
As an example of the general framework, consider Px( P,XQ,, where Q, is another projection matrix applied on the different dimension (X) = P.P XQ. This can be verified by
(see also [25]). The adjoint operator of Pa is
(P Y = tr(::XQ,Y)) = tr )
The subspace descent system of this operator is
d
W. P. (,,, (WW,))WWWW.... (.))
dt
d
8, P. dw ((W,, S.))Q v(0g ((W S.))
dt
d
P, (P(P,, Q. VL(W.))
dt
d
Q: (P Q .)),.)),
di
where Pt, Q, can be updated jointly via online SVD on G., Another linear operator that involves two matrices Px(X) P,X+XQ, which yields P;(X)- P. x XQ.
Experiment
We answer number of key questions with pretraining experiments of LLaMA [22] on the C4 dataset [20]. All experiments except for large 7B experiments are conducted on single NVIDIA A100 GPU.
5.1 Why do we Need Online Subspace Descent?
Overall, Online Subspace Descent offers two major advantages over previous methods that rely on SVD, better convergence and lower overhead. In this section, we discuss both in detail.
7

---

First, Online Subspace Descent closes the gap between the date-of-the Ei--hhe-nn art low-rank method and full rank baseline uniformly across different model sizes, as shown figure 1. A highlight amongst these results is ILaMA 1B (SS 256). As shown in table 1, Online Subspace Descent attains significant improvement over Lore in perplexity, while consuming similar amount of GPU memory (8.64 GB v.s 9.01 GB). One additional observation in 1 shows as model size and sequence length grow, Online Subspace Descent becomes more effective. We hypothesize that this is due to the higher intrinsic rank of the underlying optimization problem in larger models. Hence, the positive impact the convergence of the online update P. becomes more obvious. See more details Appendix B.
Perplexity()
Method
4.0
AdamW (8bit)
60M
350M
Loos (Log Berrllein
3.9
GaLore (AdamW Sbit)
(AdamW Sbit)
8bit-Adam (Full Rank)
32.75
30.43
29.40
3.8
GaLore (Rank = 512)
57.03
44.34
35.52
37
Ours (Rank 512)
56.12
43.67
3.6
35
Table Pretraining LLaMA with a sequence length of 256 and for 10K steps, perplexity was re- ported as the training average of the last 10 steps. AdamWSbit serves the base optimizer for both.
34
1000 2000 3000 4000 5000 6000 7000 8000 9000
Figure 1: Validation perplexity LLaMA IB with sequence length 256, rank 512 for 10K steps.
Another favorable characteristic of Online Subspace Descent its minimum overhead. In figure 2, measure and analyze the execution time of SVD and online PCA popular data center GPU (A 100) and a consumer GPU (RTX 3090). The typical Pytorch implementation of SVD can be up 142 times slower than runming single-step online PCA on representative weight tensors from LLaMA architectures. Online PCA is fast because it is implemented as single optimization step with respect simple loss function. Hence, each step of online PCA can be cleverly scheduled and hidden in the weight optimization step when executed in parallel, whereas SVD is too expensive to be hidden.
The execution time of torch.svd and that single-step backward0 call for on- line PCA in PyTorch, shapes linear layers in the LLaMA 60M 7B. Thanks to the high speed single-step online PCA, P. updates can be executed in parallel with weight updates, adding no over- head to the training process. In contrast, SVD incurs significant overhead as the model and weight tensor sizes increase.
Figure 2:
RTX 3090 SVD
RTX 3090 Orline PCA
A100 5VD
matrices of typical
A100 Online PCA
58.01x
142.43k
72,88k
20.7 1x 51,00k
57.93K
5.2 What Rank Should we Pick for Online Subspace Descent?
We conduct ablation study the rank of Online Subspace Descent. Figure shows that the final perplexity is inversely correlated with rank: higher ranks result in lower convergent perplexity. However, the rate of reduction of perplexity decreases the rank increases, eventually reaching saturation point. We propose an intuitive explanation for this phenomenon. In language modeling, high-frequency tokens can be effectively learned with low-rank training. However, learning lower- frequency tokens requires higher ranks. Once these lower-free freeerrfreuencc tokens are adequately learned, further increasing the rank does not significantly decrease perplexity. In conclusion, given sufficient time and resources, higher ranks yield better performance for Online Subspace Descent. It is recommended that the highest rank be selected until the perplexity reduction saturates.
5.3 What are the Best Hyperparameters?
The parameter a controls the update speed of Pt, while determines the regularization strength on the optimization objective of Pt. Empirically, we find that the result is not sensitive to À
and A:

---

for small models (60M). and set À = 0.1 for all subsequent experiments. We find that must be kept small avoid instability (Figure 3), and we a 5 for all experiments.
For the small model (60M), learning rate choices are more flexible, producing similar results. However, for larger models (350M, 1B), recommend using learning rate that is 10 times smaller, specifically 0,001. Larger learning rates cause unrecoverable spikes and instability, general characteristic observed across all methods. See additional hyperparameter choices in Appendix B.
Learning rate:
8.0
6.0
Rank 1375
4 560
A- 100
Rami 312
50
10.0
Uog Sepe sopeato
5.5
Ranit 120
Rami
5
0.5
55-
-10
0.1
6.5
0.03
6.0
5.0
45-
4.0-
4.5
4.0
55
2000
4020
6000
9000
10000
2000
4200
6000
8000
10300
2000
4000
6002
8000
10000
Figure 3: From left right are loss curves of 10K steps LLaMA 60M: leftmost is the sweep of rank, middle the sweep and rightmost the sweep of X.
5.4 Can Online Subspace Descent be Applied to Different Optimizers?
One straightforward extension of Online Subspace Descent is to apply it to other base optimizers beyond AdamWSbit. We conduct ablation studies on LION [6] and Adafactor [21], finding that Online Subspace Descent behaves similarly to how it does with Adam Despite the initial observation that updating P. with AdamWSbit consistently yields better results, we discover that updating Pt with simple SGD can achieve similar performance.
Method
GaLore
Ours
Lion
Adaf.
Lion+Lion
Adaf.+Adaf.
Ldam+
Addmm+
Perplexity
46.90
34.32
57.97
47.61
44.76
34.15
Table I.IaMA 60M C4 with sequence length 1024, with optimizers on P. and W., .. denote "Ours (W,
optimizer) (P, optimizer)". Adaf., and Adam refer to Adafactor and Shit-Adam W, respectively.
5.5
Can Online Subspace Descent Scale to Larger Model?
We pretrain from scratch a 7B LLaMA model on the C4 dataset for 10K steps, where the Pa matrix is updated by SGD. The perplexity is the lower the better. The final perplexity and training wall-clock time are provided in Table 3. We further provide the downstream evaluation of the pretrained checkpoints using Galore and our method the GLUE benchmark Table 4. Our method consistently outperforms Galore when the model size scales up.
Method
Perplexity
Wall Clock Time (hours)
Galore 51.21
9.7439
Ours
43.72
7.1428
Table 3: Perplexity and Clock Time models pretrained on C4 10K steps. Lower perplexity better. Online Subspace Descent can be upto 1.3x faster than GaLore.
6
Related Works
We discuss related works memory-efficee optimization and low-rank adaptation techniques.

---

Method
MRPC
RTE
SST2
MNLI
QNLI
QQP
AVG
Galore
0.6838
0.5018
0.5183 0.5233
0.3506 0.3654
0.4946
0.3682
0.4862 0.4951
Ours
0.6982
0.4901
0.5142
0.3795
Table 4:
Standardized GLUE evaluation for 7B model checkpoints using eval-harness. Results are reported for various
downstream tasks.
Low-Rank Adaptation Low-Rank Adaptation (LoRA) [11] adds a low-rank adaptor to speficic linear layers in a model, and finetune only the low-rank adaptor. As the adaptors are small, LoRA is widely applied for finetuning large models. Many variants have been proposed since LoRA, including support for multi-task learning Wang et al. [23] and further memory reductions Dettmers al. [8]. Notably, Lialin et al. [15] proposed ReLoRA for pretraining, requiring a full-rank training warmup to match standard performance levels. It's important to note that LoRA is fundamentally distinct from subspace descent. While subspace descent optimizes within the original model parameter space, LoRA focuses its optimization efforts within the space of the adaptors.
Memory-Efficien Optimization Several approaches aim reduce memory costs associated with gradient statistics in adaptive optimization algorithms [21, 2, 7J. In particular, Adafactor [21] factorizes the second-order statisties by row-column outer product and update the factorized bases on the fly, hence achieving sub-linear memory cost. K-Fac [19] presents a factorized approximation of the Fisher information matrix which leads a sublinear natural gradient method. More recently, Feinberg et al. [9] observes that the spectra of the Cronecker-factorg gradient covariance matrix in deep learning (DL) training tasks are concentrated on small leading eigenspace and propose to maintain a matrix preconditioner using the frequent directions sketch. However, their method requires conducting the eigendecomposition at every step, which can be costly for large models. Other than factorization methods, quantization techniques [7, 1, 24, 16] are also widely used, where the gradient (or the momentum and the preconditioner) are directly quantized to tradeoff performance for memory. Fused gradient computation method [17] have also been used to minimize memory costs during training. GaLore [25] is the most relevant work to ours. GaLore focuses on low-rank gradient structures, reducing memory costs for both first and second-ordes statistics. Our method can be viewed as a general extension to GaLore where we replace the infrequent SVD by a continuous subspace descent [14, 10]. As result, our method not only provides more general framework to study memory-efficient subspace descent, but is also more performant than GaLore in practice.
7 Conclusion
In conclusion, we provide the first convergence guarantee for arbitrary update rules of projection matrix, applicable to a range of optimizers that can be analyzed using Hamiltonian Descent, including common ones like LION, AdamW, and Adafactor. Inspired by this theoretical foundation, we introduce Dynamic Subspace Descent, a novel family of subspace descent optimizers that eschews SVD in favor of online PCA for updating projection matrix. Dynamic Subspace Descent is both flexible and minimally intrusive, and our experiments show that it achieves lower perplexity in pretraining LLaMA models (ranging from 60M to IB parameters) the C4 dataset compared to state-of-the-as low-rank training methods, while also closing the perplexity gap with full-rank baselines.
For future research, we propose several open and intriguing questions: (1) Are there alternative methods for updating projection matrix that could accelerate convergence? (2) What is the impact of weight decay on convergence in Dynamic Subspace Descent? (3) Can low-rank gradients and updates be combined with dynamic low-rank weights (e.g.. Mixture of Experts) to further enhance training efficiency? (4) Can this method be applied to problems beyond language modeling? We hope that our work provides strong foundation for exploring these questions.
10

---

8
Acknowledgment
The research is conducted in Statistics & AI group UT Austin, which receives supports in part from NSF CAREER INFER18466 1846421, SenSE2037267, Office of Navy Research, and NSF Al Institute for Foundations of Machine Learning (IFML).

---

References
[1] Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. Qsgd: Communication-effiiienn sgd via gradient quantization and encoding. Advances in neural information processing systems, 30, 2017.
[2] Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimization. Advances in Neural Information Processing Systems, 32, 2019.
[3] Pierre Baldi and Kurt Horiik. Neural networks and principal component analysis: Learning from examples without local minima. Neural networks, 2 2(1):53 58, 1989.
[4] Lizhang Chen, Bo Liu, Kaizhao Liang, and Qiang Liu. Lion secretly solves constrained optimization: As іyapunov predicts. arXiv preprint arXiv:2310.05898, 2023.
[5] Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Yao Liu, Hieu Pham, Xuamyi Dong, Thang Luong, Cho-Jui Hsich, et al. Symbolic discovery of optimization algo- rithms. arXiv preprint arXiv:23302.666 06675, 2023.
[6] Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, et al. Symbolic discovery of optimization algorithms. Advances in Neural Information Processing Systems, 36, 2024.
[7] Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. arXiv preprint arXiv: 10.02861, 2021.
[8] Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized lims. Advances in Neural Information Processing Systems, 36, 2024.
[9] Vladimir Feinberg, Xinyi Chen, Y Jennifer Sun, Rohan Anil, and Elad Hazan. Sketchy: Memory- efficient adaptive regularization with frequent directions. Advances in Neural Information Processing Systems, 36, 2024.
[10] Guy Gur-Ari, Daniel A Roberts, and Ethan Dyer. Gradient descent happens in tiny subspace. arXiv preprini arXiv: 18/2.04754, 2018.
[11] Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv: 106.v:2106.09685, 2021.
[12] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[13] David Kozak. Stephen Becker, Alireza Doostan, and Luis Tenorio. Stochastic subspace descent. arXiv preprint 1rXiv:1904.01 2019.
[14] Brett W Larsen, Stanislav Fort, Nic Becker, and Surya Ganguli. How many degrees of freedom do we need train deep networks: loss landscape perspective. arXiv preprint arXiv:2107.05802. 2021.
[15] Vladislav Lialin, Sherin Muckatira, Namrata Shivagunde, and Anna Rumshisky. Relora: High- rank training through low-rank updates. In Workshop on Advancing Neural Network Training: Computational Efficiency, Scalability, and Resource Optimization (WANT@ NeurIPS 2023). 2023.
[16] Bo Liu, Lemeng Wu, Lizhang Chen, Kaizhao Liang, Jiaxu Zhu, Chen Liang, Raghuraman Krishnamoorthi, and Qiang Liu. Communication efficient distributed training with distributed lion. arXiv preprint Sit::2444.00434 2024.
[17] Kai Lv, Yuqing Yang, Tengxiao Liu, Qinghui Gao, Qipeng Guo, and Xipeng Qiu. Full parameter fine-tuning for large language models with limited resources. arXiv preprint arXiv:2306.09782. 2023.
[18) Chris J Maddison, Daniel Paulin, Yee Whye Teh, Brendan O' D'Donoguue, and Arnaud Doucet. Hamiltonian descent methods. arXiv preprint ww..:vv8009.004440 2018.
12

---

[19] James Martens and Roger Grosse. Optimizing neural networks with frottcder-factored approx- imate curvature. In International conference on machine learning. pages 2408-2417. PMLR. 2015.
[20] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv e-prints, 2019.
[21] Noam Shazeer and Mitchell Stemn Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pages 4596-4604. PMLR, 2018.
[22] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 2023.
[23] Yiming Wang, Yu Lin, Xiaodong Zeng, and Guannan Zhang. Multilora: Democratizing lora for better multi-task learning. arXiv preprint arXi::231.I1500,, 2023.
[24] Wei Wen, Cong Xu, Feng Yan, Chunpeng Wu, Yandan Wang. Yiran Chen, and Hai Li. Temgrad: Ternary gradients to reduce communication in distributed deep learning. Advances in neural information processing systems, 30, 2017.
[25] Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficien IIm training by gradient low-rank projection. arXiv preprini arXiv:2403.03507, 2024.
13

---

A Hamiltonian Examples
Example A.1. Momentum + Online Subspace Descent
d M, - alG: - -M.),
W. -P,M.
=. G.W.)
dt
dt
|
2
with Lyapunov function
H(W (WWM)= = L(W) +
for which
2a
dt H(W M,) = VI(W.) + M (P. - M,) = - | 2 5.
Example A.2. Adam + Online Subspace Descent is
d W. P.
M.
G. = V.W.),
d M, a(G. M.), dt
dv. - b(G - -V.)
dt
VV,+e*
dt
M
M²)
M² )) for which
with Lyapumov function H(W M, V) L(W) +
2a
VV
HWW,, M,, Vi)
dt
P.
Mr
M,
a(P, G.
M.)
)
b
M 22
GG, 02 V.)
VV,+C
VV. +e
4a
VVOOVV
e)
b VV.
M,
b
M
t
82
4a
VV
te
VV,+e
4a
VV.
+
te)02
b)
M.
M,G.
-
1
4a
4a
VV.VV
+e)
50,
where we assume b/4
Example A.3. The Lion-K Online Subspace Descent is
d W. - B) b&.).
d
M. MM==o(GG-- +.GG+MM M,), G P.VL(W.)
dt
dt
Consider the Hamiltonian function in Eq (13) of [4]:
H(W M) = aL(W) + - KC(11 (()1)))
d
H( (W,,,M). a(G,, P.VK((1 P.VK. b)M, (G.)) of aVVX((1 6)M, G. M,)
dt
- = aG VK((1 - b)M, - (G.) - VK((1 - -b) - (VV((((1-6)))))))))))) b),
-11 b) -bGoloK (1-b) b) 10; (1 bb 50
where we defined (X;Y (Y. +C( -X(X++)-VK(X)) and used the fact that [X;Y]vk 2 by the convexity of KG; we used (G., P.X.) = (P, Gr,GG, G1, X.) = (G., X.).
In all examples above, although the form of Hamiltonian H (W,S) S) is independent of the update d H (W,, S.) depends on P. in complicate way through M., V., rule of Pt. the decreasing rate πr
G An interesting direction for future investigation is to find optimal rules of P; to maximize the decreasing rate as optimal control problem.
14

---

B Experiments
B.1 Hyperparameters
We sweep learning rate from [0.01, 0.005, 0.001J. For GaLore as well as Adamšbit, we follow the recommended hyperparamcters much as possible. For instance, GaLore update gap set to recommended default, 200. Warmup is set to 10% of the total training steps. Batch size is set to 512 1.0.
and gradient clipping
B.2 Rank Sweep
In the following table, is a sweep on different ranks and their final perplexity of LLaMA 60M (SS 1024) on C4. All other hyperparameters are fixed and using recommended default. Notice that the rank increases, both Dynamic Subspace Descent and GaLore improve.
Rank
Perplexity (Ours)
Perplexity (GaLore)
32
85.90
86.16
128
49.01
48.05
512 Full
37.41
36.93
37.18
36.51
Table 5: On I.LaMA 60M SS 1024, we sweep across different ranks, the trend is clear and intuitive that higher rank is preferred when it's feasible.
B.3
Optimizer Sweep
Method
Perplexity
Lion
52.65
Adafactor
33.45
Adam&bit
29.77 77
SGD
3469.14 14
Galore LION
46.90
Galore Adafactor
34.32
Galore Adam W8bit
48.05
Ours LION + LION
57.97
Ours Adafactor + Adafactor
47.61
Ours Adam W Sbit + AdamWShh
49.01
Ours LION Adam8bit
44.76
Ours Adafactor + AdamWSbir
34.15
Ours AdamWSbit + SGD
53.53
Table this experiment, we train LLaMA 60M on C4 with sequence length of 1024. We combine different base optimizers to update both Pt and Wt, denote as "Ours weight optimizer + Pt optimizer".
15

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Our theoretical and empirical results align with our main claims.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
We discussed the limitation and difficulty in analyzing the class of method theoretically.
Justification:
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
16

---

Justification:
All assumptions are clearly stated in the body of the paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibiliyy
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Hyperparameters choices are provided and later we will make code publicly available anonymously.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
17

---

Answer:
[Yes]
Justification:
reproducible, they All dataset and codebase used in the experiments are opensourced.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (uttps://nips.cm co/
mublic/guides Saliidee/CoddSummiis wwwsssssssooooooooooo for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: wwwpc//ggg////eeeeeeeeeeeeeeeeee
/ //nips.comm
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All hyperparameters, including learning rate, rank of projectors, sequence length are specified in each experiments.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: No]
Training LLM expensive, we can afford running multiple runs on large models. However, previous empirical studies have shown that variation across seed is generally small in language model pretraining.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by emror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
18

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative eтror rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: GPU type used in experiment is specified.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
We comply with the Code of Ethnics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our paper aims to promote deeper theoretical understanding of class of method. The algorithm itself is not practical enough produce powerful enough LLM that could potentially have any negative societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
19

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: same reasoning above
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Everything is opensourced.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/dataseta has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
20

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: No new asset has been introduced.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: No human subject has been invovled.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We did not perform any studies involving human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
21


Reviews:
[{"review_id": "ABiMZtOJZ0", "reviewer": "Reviewer_NLbB", "summary": "The paper presents Online Subspace Descent, a memory-efficient modification applicable to a wide class of gradient-descent based algorithms where low-rank projections of gradients can be employed to reduce the memory overhead. Contrary to recent techniques such as GaLore, which require infrequent but costly updates of the projection matrix, the proposed method updates the projection matrix in an on-line manner, and proves the convergence of the resulting algorithm.", "questions": "It has been observed that scaling up LLMs is notrivial also from a numerical stability point of view, e.g., in quantization, starting with 7B models, outliers become much more of a problem than they are for smaller models. As such, it would be great to see this method applied to larger models, at least to 7B, ideally even larger. If compute is a problem, maybe scaling the amount of training steps less than would be optimal could help, to at least have some proof-of-concept for the method at larger model scales.\n\n\nl.145: `this is a mild condition that can be satisfied e.g., when P_t is updated by (online) PCA on G_t` Is this really true, in the generality that is suggested here? If so, I'd like to see a formal argument. The paper only fixes the objective eq. (6), but leaves the choice of optimizes to the user, and my intuition says it should be possible to construct an example where minimizing (6) with a momentum-based optimizer could lead to $<G_t, P_t>$ being zero without $G_t$ being zero.\n\nl 158, I'm not sure what \"once\" is supposed to mean here. Should it just be \"if\"?\n\nRegarding the convergence guarantee for arbitrary $P_t$, I think it would be illuminating to consider what happens in case of an adversarial choice; for example, for SGD with momentum, choose $P_t = min <G_t, P_t^T P M_t>$, i.e., choose P so that it projects the momentum buffer _against_ the direction of the current gradient. Why does this choice of P not break convergence?\n\nl. 235\n> In conclusion, given sufficient time and resources, higher ranks yield better performance for Online Subspace Descent. It is recommended to select the highest rank until the perplexity reduction saturates.\n\nCan this reliably be selected a priori, i.e., does picking a rank optional for running optimization for, e.g., 10k steps, also result in optimal rank for training until \"convergence\"", "limitations": "Evaluations for production-size LLMs are missing.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Memory-efficient optimizers are crucial to enable LLM training/fine-tuning outside of huge datacenters; as such, this is a potential high-significance paper. While I think the idea of gradually updating the projection matrix is pretty much \"the obvious thing to do\", I really like that the authors didn't just propose \"Adam[Add-some-random-letter-here]\", but instead took a step back and showed that the same principle applies to a wide class of optimization algorithms.", "weaknesses": "* There are some grammatical errors and typos, I think (not a native speaker), but not to the degree that they hamper the understanding of the text (e.g., I think it should be: update rules of _the_ projection matrix; in the subspaces that dynamically _change_, How Should P be _updated_, is not a problem for optimizers that _yield_)\n\nl.1ß6: `Note that the update of P t can be done in parallel with that of 106 (W t, ˆ St), and incurs no slow down once it is fast enough to not cause a speed bottleneck` Generally, Adam updates are memory-bound operations, so I'm not sure if you can schedule a second operation concurrently without slowing down the update. Also, the same could be said for GaLore, you run you SVD in parallel to the main network; if it takes multiple steps, you could just wait till it is done, and update only afterwards, so GaLore can be made as parallel as the algorithm that is proposed here, I think.\n\nl.119: `In this work, we propose to update P t in a continuous online fashion that incorporates the most 119 recent gradient information in a timely fashion, without calling linear algebra routines.`\nMatrix multiplications are linear algebra routines, so I think most P updates will involve some linear algebra.\n\nI think the description around l. 119-128 could do with some rewriting. The sentence `we propose to update P t at time t by minimizing the following PCA objective` is misleading, because P_t is not actually updated like that; instead P_t is one step along some optimizer that decreases this objective. Second, `Instead of minimizing [...] exactly, to retain computational efficiency, we propose` seems to be misleading, too; in an earlier paragraph, it was GaLore was criticized for using only a single gradient information for its update. This sentence suggests that if you had enough compute, you would do the same here, only due to computational constraints do you end up with something that integrates multiple gradients, which I think goes against the message you want to make in the paper.\n\nl. 170:  `There is no requirement on Γ here, besides that the derivative in (8) should exist` Γ does not appear in (8), so it isn't clear which derivative should exist.\n\nl. 223: `The typical Pytorch implementation of SVD can be up to 142 times slower than running a single step online PCA on representative weight tensors from LLaMA architectures.` Given that GaLore runs SVD only every 300 steps, this would make GaLore comutationally more efficient than the proposed method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## l.1ß6 “Generally, Adam updates are memory-bound operations”\n- This is precisely the bottleneck that this family of online subspace descent optimizers aims to tackle. By projecting into subspace, the memory footprint of Adam will be greatly reduced, allowing us to schedule another operation concurrently. \n- In addition, since we are running tensor-wise operation, the incremental memory consumption over running adamW in subspace is simply backward of P for a single tensor. This is reflected in our experiment for Llama 1B, where the peak memory increase compared with Galore AdamW is only 0.4 GB, that is counting fragmentations, redundant intermediate saving in practice.\n\n## l.1ß6 “the same could be said for GaLore, you run you SVD in parallel to the main network”\n- It can definitely run SVD in parallel synchronously. But it won’t benefit from doing that. Instead, it will slow down that specific training step SVD is paralleled with if SVD on the latest gradient is computed. SVD uses Jacobi iterations until the principal components converge, in some cases it fails and has to restart. Overall on weight matrix sizes we care about, we observe that  it’s 20x - 150x slower than running a single backward step + update which makes it impossible to be hidden by the adamW operations. Waiting for SVD to be done slows down training, whereas we essentially propose to distribute and amortize that workload into every step of training instead of doing it all at once and wait at a single step. In our latest experiment of the 7B model, we also show that our method is 1.37x faster than baseline GaLore in wall clock time.\n- In another scenario, it’s also possible to run SVD completely asynchronously and use $P_t$ computed from some arbitrarily old gradients to perform projections and updates. However, this becomes a very different class of algorithms from baseline GaLore. It overly complicates implementation in practice, since it needs to manage schedules for every tensor in the network. SVD finish time can also vary and it is unpredictable. Hence, it's possible to have part of the network using $P_t$ calculated from a different timestep, if it's running in a non-blocking fashion. The implication of this is unknown and it's a very interesting direction of exploration. Yet, given the limited span of this work, we would like to defer this discussion to future works.\n\n## l.119 \"Matrix multiplications are linear algebra routines\"\nWe meant to refer to torch.linalg.decomposition routines. We will make it more specific in revision.\n\n## Re: \"119-128 could do with some rewriting\"\n- We agree that we can rewrite it to highlight the goal of doing online PCA, not just for saving computational cost.\n\n## l. 170 \"There is no requirement on Γ here, besides that the derivative in (8) should exist Γ does not appear in (8), so it isn't clear which derivative should exist.\"\n- The derivatives in (8) refers to the partial derivatives of $H_t$ that appear in Eq (8). We will make it specific in the revision.\n\n## l. 223 \"Given that GaLore runs SVD only every 300 steps, this would make GaLore computationally more efficient than the proposed method.\"\nPlease refer to the above explanation on system efficiency/parallel implementations. Even if running Galore every 300 steps, it’s still an unmitigable slowdown where no parallelization tricks can help. And as tensor shape grows, as figure 2 shows, the gap of latency becomes higher. Meanwhile the cost of online subspace descent can be hidden with a smart parallelization trick.\nOur latest experiment of 7B model shows that our method is 1.37x faster than baseline GaLore in terms of wall clock time with the exactly the same hardware setup.\n\n## Re: \"Larger model\"\nWe pretrained from scratch 7B llamas on C4 for 10K steps. Perplexity is the lower the better. The results are consistent with our claim in production model scale.\n| Method | Perplexity | Wall Clock Time (hours) |\n|--------|------------|-------------------------|\n| Galore | 51.21      | 9.7439                  |\n| **Ours**   | **43.72**      | **7.1428**                  |\n\n## l.145 \" Adversarial projections\"\nThe loss in $L_{G_t}(P)$ attains global optima only if P spans the linear space of the top k eigenvectors of $G_t$ in Loss (the other eigenvectors are saddle points). Hence, if an optimizer minimizes $L_{G_t}(P)$ correctly, it should find the top k eigenvectors of G, which is zero iff G ==0. Regardless of the choice $P_t$, the system yields a given Hamiltonian function and converges to some invariant set. But the invariant set coincides with the local optima of the objective function f only when Assumption 4.4 satisfies. The choice you provide may not satisfy Assumption 4.4.\n\nMore concretely, solution of $P_t=min<G_t,P_t^TPM_t>$ gives the opposite direction of $G_t$. Considering the negative direction won’t harm it, because it’s still the same subspace even the gradient has an opposite sign and while it’s projected back, it’s still “gradient descent”, the two negative signs simply cancel out. The real adversarial scenario is $P_t=min|<G_t,P_t^TPM_t>|$. It is projecting into complementary orthogonal subspace, aka, $<G_t,P_t^TPM_t> = 0$. This is against ii) in our Assumption 4.4, which states we can not pick $P$ such that $P^TG$, the projected gradient is always 0. Otherwise, it just doesn’t update the model weights at all and it defeats the purpose.\n\n## l 158 \"I'm not sure what \"once\" is supposed to mean here. Should it just be \"if\"?\"\nYes, it should be “if”. We will correct that in revision.\n\n## l. 235 \"choices of rank\"\nJust in terms of speed of convergence, in our ablation study, within 10K steps, the higher rank it’s the faster it goes. Eventually, if it’s run long enough, our theory indicates that it will always converge. However, in practice, it might be difficult to observe that due to constraint in resources. So as a safe recommendation, if it can be afforded, choose a higher rank.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "SdPqmAGMox", "reviewer": "Reviewer_6zai", "summary": "Utilizing low-rank structure has become a popular way for memory-efficient LLM training. The authors are the first to provide convergence guarantees for general low-rank updating rules. Furthermore, based on their theoretical result, they propose a family of optimizers called online subspace descent. The empirical result shows that their performance is better than existing low-rank training methods.", "questions": "Could you provide more high-level intuitions on why the methods could converge under the complications?", "limitations": "Yes, the authors discuss the limitations.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This work provides a general theoretical guarantee for a range of low-rank updating rules, which can guide and help future research in this direction.", "weaknesses": "While the authors provide mathematical proofs, it could be more beneficial to the readers if they can provide more high-level intuitions why these methods could converge. As stated by the authors, the convergence result is surprising under the complications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Re: “Intuitions”\nThe rough intuition is that $P_t$ serves as a kind of preconditioning matrix in the Hamiltonian systems. But to arrive at the precise mathematical conclusion, we find that the best and quickest way to understand it is through the derivation in Eq (8), together with physical understandings of Hamiltonian systems (that, e.g., the hamiltonian consists of potential energy and kinetic energy, and the kinetic energy term is different for different optimizers)\n\nMore specifically, as long as the projected subspace defined by $P_t$ keeps changing and not always be orthogonal to the gradient as we stated in Assumption (4.4), it will always make some progress towards the fixed point. Given enough time, it will eventually reach there. Hence the consequence of convergence does not depend on $P_t$.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "ivyqFsevFF", "reviewer": "Reviewer_RBsr", "summary": "The authors propose a variation of GaLore that doesn't require a fixed SVD for achieving a memory efficient LLM training. Instead, the authors project the gradients dynamically into a small sub-space using online PCA which depends on the evolving trace of the gradient landscape. \n\nThe authors tested their method against GaLore using perplexity as a metric while training on the C4 dataset.", "questions": "please answer how you could address the weaknesses above", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- the paper is well written and easy to follow\n- having a dynamically changing projection matrix based on the gradient landscape makes total sense\n- the theoretical explanation shows an interesting perspective on why this proposed method of using online pca for matrix projection could work\n- the results is promising in C4 as the memory seems more or less similar in efficiency as GaLore with similar perplexity.", "weaknesses": "- the perplexity results is not enough to show that the method works well on standard NLP tasks like GLUE which consists of the datasets MRPC, COLA, STS-B, RTE, SST2, MNLI, QNLI, QQP to see if the actual downstream task performance is still retained with the proposed projection method.\n- while the theory is interesting, it seems to include many strong assumptions that do not reflect real-world datasets\n- the memory savings and the perplexity is very similar to GaLore and it is not clear that it is significant. The authors need to do multiple runs and compute the standard deviation to make the results are not due to noise.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Re: \"Downstream tasks\"\nWe did standardized GLUE evaluation for the above two 7B checkpoints with eval-harness.\n\n| Method | MRPC  | RTE   | SST2  | MNLI  | QNLI  | QQP   | AVG    |\n|--------|-------|-------|-------|-------|-------|-------|--------|\n| GaLore | 0.6838| **0.5018**| 0.5183| 0.3506| 0.4946| 0.3682| 0.4862 |\n| **Ours**   | **0.6982**| 0.4901| **0.5233**| **0.3654**| **0.5142**| **0.3795**| **0.4951**|\n\nBoth Cola and STSB require further fine-tuning, which brings in more confounding factors (even Llama-2 7B trained by Meta doesn't get meaningful results), so we exclude them from the comparison. \n\n## Re: “Many strong assumptions that do not reflect real world datasets”\nWe respectfully argue that our assumptions (Assumption 4.4) are in fact, very much minimum. \n1) The model needs to be continuously differentiable. (pretty much all modern neural nets)\n2) The optimization stops when the projected update hits 0. (this is rather a fact than a assumption)\n3) It says we should pick $P$ such that $P^TG$, the projected gradient is not always 0. Otherwise, it just doesn’t update the model weights at all and it defeats the purpose.\n\nEvery theoretical work has to make some assumptions. We want to assure our reviewers that it will be hard to find papers with same level of theoretical rigorousness that is making fewer assumptions than we do: we did not assume convex losses; we did not assume specific forms of the model (linear model or MLPs); we did not even assume specific optimizers (Adam, momentum, lion and etc)...\n\n## Re: “The memory savings and the perplexity is very similar to GaLore and it is not clear that it is significant”\nWe agree that this does have very similar memory saving as GaLore. However, we argue that our contributions lie majorly in our analysis that introduces a new family of optimizers that are guaranteed to converge. Among them, online subspace descent not only possesses the same memory saving properties, but also can be implemented more efficiently in a parallel manner, whereas Galore doesn’t share the same flexibility.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "gm27PHhKzP", "reviewer": "Reviewer_FtMv", "summary": "The paper considers memory-efficient optimization for large language model. In particular, the focus is on optimizers leveraging low-rank projection. The paper first derives asymptotic convergence for such methods with arbitrary choice of projection matrix. Then the paper identifies the inefficiency of using SVD to construct the projection matrix and proposes online subspace descent method by updating the projection matrix with online PCA. Experiments show such method improves the efficiency of training LLM while achieving comparable performance compared to low-rank baseline with SVD projection.", "questions": "1. In Line 233, what does it mean to be high-frequency versus low-frequency tokens? And is there any verification or justification that high-frequency tokens are learned with low-rank training while learning low-frequency tokens requires higher ranks?\n\n2. In Line 252, can you comment further why AdamW is more suitable for online PCA?", "limitations": "The paper has discussed limitations.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper derives the first convergence result for optimizers with low-rank projection.\n\n2. The paper addresses the inefficiency of using SVD to construct the projection matrix by online subspace descent.\n\n3. Experiments show the benefits of the online subspace descent over the SVD baseline.", "weaknesses": "1. The main downside of the proposed method is that it could potentially increase the memory overhead. It seems to me that the idea of online updating $P_t$ trades off computation with memory. Although it reduces the computational cost of not computing SVD, it nevertheless increases the memory especially using Adam type of optimizer for $P_t$. In Line 216, the paper states the memory overhead is 9.01GB versus 8.64GB. However, it is not clear what model this numbers refer to. I suspect with the increase in rank and model size, the memory overhead could become larger. \n\n2. The use of Hamiltonian descent and Lyapunov analysis for showing convergence is not properly motivated. Because the main result (Theorem 4.5) only shows asymptotic convergence to stationary point of $L(W)$, it is not convinced why to derive such results under continuous time limit. Can such result also be derived in the discrete case and can non-asymptotic result be derived? \n\n3. Figure 2 is not a fair comparison because in practice, SVD is not used every update.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Re: “Memory and efficiency”\nUpdating $P_t$ with AdamW will inadvertently increase the memory consumption. However, in practice, we observe that the increase in peak memory is minuscule compared to the overall size of the model. Meanwhile, we can enjoy the nice properties of online subspace descent for its flexibility and system efficiency when implemented in parallel (as we discussed with reviewer NLbB and Rbsr). In addition, our subsequent experiment shows that SGD can perform similarly to AdamW when updating P, in which case, there is no additional memory cost over GaLore, since no momentum needs to be saved.\n\nHere is a comparison between GaLore and our method on 7B llamas trained for 10K steps on C4 (updating $P_t$ with SGD)\n| Method | Perplexity | Wall Clock Time (hours) |\n|--------|------------|-------------------------|\n| Galore | 51.21      | 9.7439                  |\n| Ours   | 43.72      | 7.1428                  |\n\nThe hardware setups are completely identical, both using a single node of H800s with 8 cards. It's clear that our method is not only faster in wall clock time, but also reaches lower perplexity when seeing same number of tokens.\n\n## Re: \"Discrete case\"\nThe derivation of discrete case is highly similar to the continuous example we've shown in the paper. \n\n**Example 4.1.** Momentum + Online Subspace Descent is\n\n$$\n\\frac{d}{dt} W_t = -P_t \\hat{M}_t, \\quad \\hat{G}_t = P_t^\\top \\nabla L(W_t), \\quad \\frac{d}{dt} \\hat{M}_t = a(\\hat{G}_t - \\hat{M}_t),\n$$\n\nwith Lyapunov function $H(W, \\hat{M}) = L(W) + \\frac{\\|\\hat{M}\\|^2}{2a}$, for which\n\n$$\n\\frac{d}{dt} H(W_t, \\hat{M}_t) = - \\nabla L(W_t)^\\top P_t \\hat{M}_t + \\hat{M}_t^\\top (P_t^\\top \\nabla L(W_t) - \\hat{M}_t) = -\\|\\hat{M}_t\\|_2^2 \\leq 0.\n$$\n\nFollowup this example as a demo, we have the following discrete derivation:\n\n$$W_{t+1} = W_t - \\epsilon P_t \\hat{M}_t$$\n\n$$\\hat{M}_{t+1} = \\hat{M}_t + a(\\hat{G}_t - \\hat{M}_t)$$\n$$ \\hat{G}_t = P_t^\\top \\nabla L(W_t)$$\n\nwhere $\\hat{M}_{t+1} = \\hat{M}_t + a(\\hat{G}_t - \\hat{M}_t) = (1-\\alpha \\epsilon)\\hat{M}_t + \\alpha \\epsilon \\hat{G}_t = \\beta \\hat{M}_t + (1-\\beta)\\hat{G}_t$.\nAssuming $L(\\cdot)$ is $L$-smooth, we have\n\n$$\nH(W_{t+1}, \\hat{M}_{t+1}) - H(W_t, \\hat{M}_t) \\leq -\\epsilon \\|\\hat{M}_t\\|_2^2 + \\alpha \\epsilon^2 \\|\\hat{G}_t - \\hat{M}_t\\|_2^2 + \\frac{L}{2}\\epsilon^2\\|P_t\\hat{M}_t\\|_2^2.\n$$\n\nHence, a telescoping sum yields\n\n$$\n\\frac{1}{T}\\sum_{t=0}^{T-1} \\|\\hat{M}_t\\|_2^2 \\leq \\frac{H(W_0, \\hat{M}_0) - H(W_T, \\hat{M}_T)}{\\epsilon T} + \\epsilon B_t,\n$$\n\nwhere $B_t = \\frac{1}{T} \\sum_{t=0}^{T-1} \\alpha \\|\\hat{G}_t - \\hat{M}_t\\|_2^2 + \\frac{L}{2}\\|\\hat{M}_t\\|_2^2$.\n\n## Re: “high-frequency versus low-frequency tokens”\nThis is our attempted explanation for the phenomenon we observe in our ablation study of different ranks. Loss of higher rank runs tend to drop faster and converge to a lower loss after 10K steps. We suspect that it’s due to the unique property of language modeling tasks, where intuitively rarer tokens (tokens that are not common) require high gradient rank to be learnt. However, we would like to delay the justification of this hypothesis to future study, since it’s not the main focus of our paper. \n\n## Re: \"AdamW is more suitable for online PCA\"\nThis is what we find empirically, when trying different optimizers such as LION and Ada Factor. However, it’s not unfathomable that SGD in principle can perform similarly, which we show above. It’s just for the default learning rates and schedules, adamW seems to work more robustly and reliably and it’s the least sensitive one.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:
We really appreciate our reviewers for their constructive reviews and suggestions. Here we summarize and highlight our response to a few main points of concerns. Hope that will help address the majority of those questions.

## Re: “System efficiency compared to Galore” (NLbB, FtMV, RBsr)
In terms of speed, online subspace descent’s additional cost can be easily hidden by simple parallelization, whereas GaLore can not. As model size increases, the gap in latency between our method and GaLore becomes larger. Additional experiment shows on 7B model on same hardware setup, our method can be up to **1.3x** faster than GaLore.

In terms of memory, we have flexibility in choosing different optimizers for updates of $P_t$. When choosing AdamW, inadvertently, it increases memory consumption slightly over GaLore. We chose to start with AdamW due to its robustness to hyperparameters. However, additional experiment on the 7B model shows that SGD also shares similar properties and can indeed be used for $P_t$ updates, incurring no extra memory cost over GaLore.

## Re: “Larger scale experiment” (NLbB)
We pretrained from scratch 7B llamas on C4 for 10K steps. Perplexity the lower the better. The results are still consistent with our claim in production model scale. $P_t$ is updated by SGD.
| Method | Perplexity | Wall Clock Time (hours) |
|--------|------------|-------------------------|
| Galore | 51.21      | 9.7439                  |
| **Ours**   | **43.72**      | **7.1428**                  |

## Re: “Downstream tasks” (RBsr)
We did standardized glue evaluation for the above two 7B checkpoints with eval-harness.
| Method | MRPC  | RTE   | SST2  | MNLI  | QNLI  | QQP   | AVG    |
|--------|-------|-------|-------|-------|-------|-------|--------|
| GaLore | 0.6838| **0.5018**| 0.5183| 0.3506| 0.4946| 0.3682| 0.4862 |
| **Ours**   | **0.6982**| 0.4901| **0.5233**| **0.3654**| **0.5142**| **0.3795**| **0.4951**|

Both Cola and STSB require further fine-tuning (even Meta Llama-2 7B gets unmeaningful results), so we exclude them from the comparison. 

## Re: “Motivation for the Lyapunov Analysis” (FtMV, 6zai)
The rough intuition is that Pt serves as a kind of preconditioning matrix in the Hamiltonian systems. But to arrive the precise mathematical conclusion, we find that the best and quickest way to understand it is through the derivation in Eq (8), together with physical understandings of Hamiltonian systems (that, e.g., the hamiltonian consists of potential energy and kinetic energy, and the kinetic energy term is different for different optimizers)

Meta-review: After the author response, all reviewers ultimately recommended acceptance of the paper.
Decision: Accept (poster)