Title: Memory-Efficient LLM Training with Online Subspace Descent
Venue: NeurIPS 2024
Accepted: yes
Keywords: Large Language Model Pretraining; Optimizer; Memory-Efficient LLM Training
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

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 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 the projection matrix with eigenvectors, Online Subspace Descent updates the projection matrix with online PCA. Online Subspace Descent is 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-ofthe-art low-rank training methods across different settings and narrows the gap with full-rank baselines. 1

1 Introduction

The continual advancement in training large language models (LLMs) presents a compelling challenge in balancing computational efficiency with model performance. As the scope and complexity of these models grow, so does the necessity for innovative strategies that optimize memory usage without compromising the learning capabilities of the model. Recent approaches in low-rank adaptation strategies, including Stochastic Subspace Descent [13], LoRA [11], ReLoRA [15], Gradient LowRank Projection (GaLore) [25] and Sketchy [9] , have paved the way for memory-efficient training by utilizing a 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 as projection matrices and provides convergence analysis on convex functions and objectives. However, to the best of our knowledge, no one has offered any guarantee of convergence for this class of methods on 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 is broadly applicable to a wide range of optimizers that can be analyzed within the Hamiltonian descent framework [18]. By establishing this convergence guarantee, we demonstrate that our approach is 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

Algorithm 1 Online Subspace Descent

1: Required: Optimizer OptimizerW , learning rate ϵ t W , weight decay λ W for model weights W t ; and { OptimizerP , ϵ P t , λ P } for the projection matrix P t . Proper initialization.

2: for iteration t do

3: Calculate gradient G t = ∇ L ( W t ) ; Update model weights W t by

4: Calculate G P t = ∇ L G t ( P t ) for L G t ( · ) in Eq (6); Update the projection P t by

5: end for

6: Remark: We added weight decay as a common heuristic. We recommend using Adam for both optimizers, and set ϵ P t = αϵ t W with a constant α (e.g., α = 5 ), and λ W = λ P . 2

includes popular algorithms such as LION [4] and Adam [12], which are widely used in various machine learning and optimization tasks. Our results therefore offer a robust theoretical foundation for understanding and analyzing the behavior of these optimizers, ensuring their effectiveness and reliability in diverse applications.

Inspired by our theoretical understanding, we introduce a novel family of memory-efficient optimizers 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 ability to 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 analysis, we demonstrate that our approach presents lower perplexity in pretraining LLaMA models (ranging from 60M to 1B parameters) on the C4 dataset compared to state-of-the-art 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

where W is the set of weight matrices of the model. For simplicity, we assume W ∈ R n × m is a 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 ∥ A ∥ 2 = tr( A ⊤ A ) the Frobenius norm. We use A ⊙ B to denote the elementwise product, and A ⊙ 2 = A ⊙ A .

Example 2.1. Update rules of common optimizers:

where ϵ t are step sizes, and M t , V t are the first and second order momentum, and β, β 1 , β 2 are momentum coefficients in (0 , 1) , with β it = β i -β t +1 i 1 -β t +1 i , i = 1 , 2 for β 1 , β 2 ∈ (0 , 1) in Adam, and K is any convex function with ∇ K ( 0 ) = 0 for LionK [4], and Lion [5] uses K ( X ) = ∥ X ∥ 1 , 1 and ∇K ( X ) = sign( X ) .

These optimizers can be unifiedly viewed as updating W t together with an optimizer state S t :

with some mapping ϕ t , ψ t . We have S t = M t for momentum and S t = { M t , V t } for Adam. Note that both M t , V t are of the same size as the model weights W t , resulting in 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 M t and V t . One key challenge is to retain the high performance of Adam while enhancing its memory efficiency.

Hamiltonian+Descent One powerful approach to studying the dynamic properties of optimizers is to examine their continuous-time ODE forms in the limit of infinitesimal step size. The continuoustime 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 a "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,

where H ( W , S ) is a Hamiltonian (or Lyapunov) function that satisfies

so that minimizing L ( W ) reduces to minimizing H ( W , S ) ; and Φ( · ) , Ψ( · ) are two monotonic mappings satisfying

With Φ( X ) = Ψ( X ) = 0 , the system in (2) reduces to the standard Hamiltonian system that keeps H ( W t , S t ) = const along the trajectory. When adding the descending components with Φ and Ψ , the system then keeps H ( W , S ) monotonically non-decreasing:

where we write ∂ W H t = ∂ W H ( W t , S t ) and similarly for ∂ S H t . The main idea is that the cross terms ⟨ ∂ W H t , ∂ S H t ⟩ are canceled, leaving only the negative terms.

Example 2.2. The momentum method yields following continuous-time form and Hamiltonian:

2

Example 2.3. Adam [12] yields the following continuous-time form and Hamiltonian,

for which we can show that d d t H ( W t , M t , V t ) ≤ 0 when a ≥ b/ 4 .

Example 2.4. The LionK optimizer [5, 4] (without weight decay) can be written into

where a ≥ 0 , b ∈ [0 , 1] and K ( X ) is any convex function that attains the minimum at X = 0 . One of its Hamiltonians that yields the Hamiltonian+descent structure (Eq (13) in Chen et al. [4]) is

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 in the full space while only using subspace descent (Section 3.2). In 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 in 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 4.

3.1 Static Subspace Descent

One popular approach to improving memory efficiency is to confine the optimization to a lowdimensional space. To do this, we impose a low rank structure of W = P ˆ W , where P ∈ R n × k is a projection matrix to be determined later and ˆ W ∈ R k × m is a dimension-reduced parameter. When k ≪ n , P and ˆ W are much smaller in size compared to W . Now consider

Applying the optimizer from (1) to update ˆ W along with an optimizer state ˆ S , and mapping the update rule ˆ W t +1 = ˆ W t + ϕ t ( ˆ S t ) to that of W = P ˆ W t , we get

where we used the fact that ∇ ˆ W L ( P ˆ W ) = P ⊤ ∇ W L ( W ) . This yields a more memory-efficient optimizer, as the size of ˆ S t is proportional to that of ˆ W t , much smaller than S t in (1) when k ≪ 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 [11], it is 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:

where χ t is a update rule of P t that will be determined in the sequel. The intuition is to open up different projection directions at different iterations, so that optimization can be conducted in different subspaces across different iterations. This is similar to the update of coordinate descent, except in a continuous fashion. Note that the update of P t can be done in parallel with that of ( W t , ˆ S t ) , and incurs no slowdown once it is fast enough to not cause a speed bottleneck.

Example 3.1. Examples of common optimizers equipped with online subspace descent:

How Should P t be Updated? It is useful to draw intuition from the projected gradient descent rule

in which P t P ⊤ t can be viewed as a low rank preconditioning of G t . To make it follow the exact gradient descent, we hope to make P t P ⊤ t G t approximate G t as much as possible. In Galore, this is achieved by performing singular value decomposition (SVD) on G t periodically every T iterations:

where T ⌊ t/T ⌋ is the largest multiple of T less than or equal to t . However, numerical SVD incurs a large computational cost for very large models. Also, since P t is fully determined by G T ⌊ t/T ⌋ calculated from a single mini-batch at the last periodic point, it does not incorporate the gradient information from all data in a timely fashion.

In this work, we propose to update P t in a continuous online fashion that incorporates the most recent gradient information in a timely fashion, without calling torch.linalg.decompositions routines. We view the update of P t as conducting an online principal component analysis (PCA) based on the streaming of { G t } . In particular, we propose to update P t at time t by minimizing the following PCA objective:

where ∥ A ∥ = tr( A ⊤ A ) 1 / 2 and I k × k is the k × k identity matrix; we introduced an auxiliary loss to encourage the columns of P to be orthonormal and normalizes G t to increase stability.

The key property of L G t ( P ) in (6) is that all its stable local minimum is a global minimum, and P is a global minimum iff PP ⊤ ˜ G t forms the optimal rankk approximation of ˜ G t [e.g., 3]; moreover, we have P ⊤ P = I k × k at optima when λ > 0 .

Instead of minimizing L G t ( P ) exactly, to retain computational efficiency, we propose to update P t by only performing one step of optimization on L G t ( P ) :

where OptimizerP.step can be a favorite optimizer, such as gradient descent or Adam. Note that when using Adam, we introduce a copy of optimizer state S P t for P t . See Algorithm 1. Compared to the exact SVD, each online update of P t here is fast and can be executed in parallel with the ( W t , ˆ S t ) updates to avoid slowdown.

3.3 Difficulty in Theoretical Understanding

The idea above of projecting an arbitrary optimizer with a dynamically changing P t is heuristically motivated and lacks an immediate rigorous theoretical justification. The main challenge lies in the complex interaction between the update of U t and the optimization state S t , 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 t impacts the optimizer's convergence behavior and provide guidance for the design of the update rules of P t .

To gain understanding, it is useful to first exam the simple case of projected gradient descent in (5) which does not have an optimizer state ( S t = ∅ ). In this case, since P t P ⊤ t is positive semi-finite, the update P t P ⊤ t G t is always non-increasing direction of L ( W ) for any P t . The algorithm is essentially a variant of coordinate or subspace descent, where P t defines the subspace on which one step of gradient descent is conduced at iteration t . To ensure that (5) finds a local optimum, we mainly need to ensure that P t G t = 0 only if G t = 0 to prevent the optimizer from stopping prematurely; this is a mild condition that can be satisfied e.g. when P t is updated by (online) PCA on G t .

Unfortunately, this coordinate-descent-like interpretation does not apply to more advanced optimizers that track a momentum state S t . This is because S t accumulates the information from the projected gradient P τ G τ at all earlier iterations τ ≤ t . As P t changes across time, it is unclear whether the gradient projected to different subspaces P τ would be coherent with each other, and useful for future updates that are conducted in different subspaces P t for t > τ . The difficulty is the inertia effect of S t 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 4 via the Hamiltonian descent framework.

4 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 is not a problem for optimizers that yields the Hamiltonian+descent structure in (2). Our result is two-fold:

Section 4.1: When applying Online Subspace Descent on systems in (2), the Hamiltonian+descent structure is preserved once the update rule of P t has a smooth continuous-time limit. Hence, under very mild conditions, Online Subspace Descent equipped with common optimizers like Adam and Lion automatically yield a Lyapunov function and hence benign continuous-time convergence. Moreover, P t can, in fact, be generalized to an arbitrary linear operator as shown in Section 4.3.

Section 4.2: For any smooth P t update rules that eliminates the degenerate case of P ⊤ t G t = 0 while G t = 0 at convergence, the online subspace optimizer guarantees to converge in continuous time to a stationary point of the loss L ( W ) . This mild condition is satisfied, for example, when P t is updated by a typical optimizer on the online PCA objective L G t ( P ) .

4.1 Online Subspace Descent Preserves the Hamiltonian+Descent Structure

Applying dynamic projection to Hamiltonian descent in (2), we obtain the following systems:

where Γ specifies the update rule of P t . Following essentially the same derivation as (3), one can show that H ( W , S ) remains a Lyapunov function of (7), regardless of the choice of Γ :

where the key is to use the adjoint property of P and P ⊤ that ⟨ P t X , Y ⟩ = ⟨ X , P ⊤ t Y ⟩ , which cancels the crossing terms, independent of the values of P t . There is no requirement on Γ here, besides that the derivative in (8) should exist. As shown in Section 4.3, we can generalize (8) by replacing P t and P ⊤ t with a general linear operator P t and its adjoint P ∗ t .

Please refer to Appendix A for continuous-time Momentum, LionK and Adam with subspace descent and their Hamiltonian functions.

4.2 Convergence to Local Optima

In addition to the Lyapunov structure, we need an additional mild condition on the update rule of P t 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 zero gradient G t = 0 by excluding the degenerate case of P t G t = 0 while G t = 0 in the invariant set of the system.

̸

Assumption 4.1. Assume the functions in system (7) are continuously differentiable and

̸

This is a mild condition. Assumption i) says that the optimizer should stop at a point with ˆ G t = 0 , which is easy to verify for the common optimizers like momentum, Adam, LionK . Assumption ii) ensures ˆ G t = 0 would imply G t = 0 in invariance sets, which is satisfied when for example, P t is updated by a reasonable optimizer of the online PCA loss that converges to a stable local minimum.

Theorem 4.2. Assume Assumption 4.1 holds. Let ( W t , S t , P t ) t be a bounded solution of (7) , then all the accumulation points { W t } as t → + ∞ are stationary points of L ( W ) .

Proof. By LaSalle's invariance principle, the positive limit set of ( W t , S t , P t ) t must be contained in I , where I = { the union of complete trajectories satisfying d d t H ( W t , ˆ S t ) = 0 , ∀ t } .

From the Assumption i), the trajectories contained in I must satisfy d d t W t = 0 , which implies d d t G t = d d t ∇ L ( W t ) = 0 and ˆ G t = 0 and hence G t ≡ G is a constant with P ⊤ t G = 0 . Moreover, from Assumption ii), we must have ∇ L ( W t ) = G t ≡ 0 , since otherwise the trajectory is not invariant. As a result, all trajectories in the limit set I must have ∇ L ( W t ) = 0 . Because d d t W t = 0 , these trajectories are static points of W t .

4.3 Online Subspace Descent with General Linear Projection Operators

We can generalize the online subspace descent with general linear operators:

where we generalize P t to be any linear operator P t with an adjoint operator P ∗ t , satisfying

The derivation of Lyapunov follows a similar way:

where the crossing terms are again canceled due to the adjoint property.

As an example of the general framework, consider P t ( X ) = P t XQ t , where Q t is another projection matrix applied on the different dimension of X (see also [25]). The adjoint operator of P t is P ∗ t ( X ) = P ⊤ t XQ ⊤ t . This can be verified by

The subspace descent system of this operator is

where P t , Q t can be updated jointly via an online SVD on G t .

Another linear operator that involves two matrices is P t ( X ) = P t X + XQ t , which yields P ∗ t ( X ) = P ⊤ t X + XQ ⊤ t .

5

Experiment

We answer a 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 a 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.

First, Online Subspace Descent closes the gap between the state-of-the-art low-rank method and full rank baseline uniformly across different model sizes, as shown in figure 1. A highlight amongst these results is LLaMA 1B (SS 256). As shown in table 1, Online Subspace Descent attains significant improvement over GaLore in perplexity, while consuming a 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 on the convergence of the online update of P t becomes more obvious. See more details in Appendix B.

Adamw(8bit)

4.0

Loss(LogPerplexity)

Ours(AdamW8bit)

3.9

GaLore(AdamW8bit)

3.8

3.7

3.6

3.5

3.4

00060000000000090000000000000070000

Another favorable characteristic of Online Subspace Descent is its minimum overhead. In figure 2, we measure and analyze the execution time of SVD and online PCA on a popular data center GPU (A100) and a consumer GPU (RTX 3090). 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. Online PCA is fast because it is implemented as a single optimization step with respect to a 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.

RTX3090-SVD

100

RTX3090-OnlinePCA

A100-SVD

Execution Time (s)

A100-OnlinePCA

58.01x

10-1

142.43

49.75x

72.88

10-16.38x

20.71x

51.00x

57.93K

10-3

512 × 1376

1024 × 2736

2048 x5461

4096×11008

Figure 2: The execution time of torch.svd and that a single-step backward() call for online PCA in PyTorch, on matrices of typical shapes in linear layers in the LLaMA 60M to 7B. Thanks to the high speed of single-step online PCA, P t updates can be executed in parallel with weight updates, adding no overhead to the training process. In contrast, SVD incurs significant overhead as the model and weight tensor sizes increase.

5.2 What Rank Should we Pick for Online Subspace Descent ?

We conduct an ablation study on the rank of Online Subspace Descent. Figure 3 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 as the rank increases, eventually reaching a 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 lowerfrequency tokens requires higher ranks. Once these lower-frequency 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?

α and λ : The parameter α controls the update speed of P t , while λ determines the regularization strength on the optimization objective of P t . Empirically, we find that the result is not sensitive to λ

for small models (60M). and set λ = 0 . 1 for all subsequent experiments. We find that α must be kept small to avoid instability (Figure 3), and we set α = 5 for all experiments.

Learning rate : For the small model (60M), learning rate choices are more flexible, producing similar results. However, for larger models (350M, 1B), we recommend using a learning rate that is 10 times smaller, specifically 0.001. Larger learning rates cause unrecoverable spikes and instability, a general characteristic observed across all methods. See additional hyperparameter choices in Appendix B.

8.0

6.0

Rank = 1376

Rank = 512

7.5

α=500

6.0

入=100.0

入 =10.0

Loss (Log Perplexity)

5.5

Rank = 128

Rank = 32

7.0

α = 5

α = 0.5

5.5

入 =1.0

入 = 0.1

6.5

入=0.01

5.0

6.0

5.0

入=0.001

入=0.0

5.5

4.5

5.0

4.5

4.0

4.5

4.0

4.0

3.5

0

2000

4000

6000

8000

10000

0

2000

4000

6000

8000

10000

0

2000

4000

6000

8000

10000

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 AdamW8bit. We conduct ablation studies on LION [6] and Adafactor [21], finding that Online Subspace Descent behaves similarly to how it does with AdamW8bit. Despite the initial observation that updating P t with AdamW8bit consistently yields better results, we discover that updating P t with simple SGD can achieve similar performance.

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 P t 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 on the GLUE benchmark in Table 4. Our method consistently outperforms Galore when the model size scales up.

6 Related Works

We discuss related works on memory-efficient optimization and low-rank adaptation techniques.

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 et 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-Efficient Optimization Several approaches aim to reduce memory costs associated with gradient statistics in adaptive optimization algorithms [21, 2, 7]. In particular, Adafactor [21] factorizes the second-order statistics by a row-column outer product and update the factorized bases on the fly, hence achieving a sub-linear memory cost. K-Fac [19] presents a factorized approximation of the Fisher information matrix which leads to a sublinear natural gradient method. More recently, Feinberg et al. [9] observes that the spectra of the Kronecker-factored gradient covariance matrix in deep learning (DL) training tasks are concentrated on a 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-order 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 a result, our method not only provides a 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 1B parameters) on the C4 dataset compared to state-of-the-art 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 a strong foundation for exploring these questions.

8 Acknowledgment

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

References

Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. Qsgd: Communication-efficient sgd via gradient quantization and encoding. Advances in neural information processing systems , 30, 2017.

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

Pierre Baldi and Kurt Hornik. Neural networks and principal component analysis: Learning from examples without local minima. Neural networks , 2(1):53-58, 1989.

Lizhang Chen, Bo Liu, Kaizhao Liang, and Qiang Liu. Lion secretly solves constrained optimization: As lyapunov predicts. arXiv preprint arXiv:2310.05898 , 2023.

Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Yao Liu, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, et al. Symbolic discovery of optimization algorithms. arXiv preprint arXiv:2302.06675 , 2023.

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.

Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. arXiv preprint arXiv:2110.02861 , 2021.

Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized llms. Advances in Neural Information Processing Systems , 36, 2024.

Vladimir Feinberg, Xinyi Chen, Y Jennifer Sun, Rohan Anil, and Elad Hazan. Sketchy: Memoryefficient adaptive regularization with frequent directions. Advances in Neural Information Processing Systems , 36, 2024.

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

Edward J 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:2106.09685 , 2021.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

David Kozak, Stephen Becker, Alireza Doostan, and Luis Tenorio. Stochastic subspace descent. arXiv preprint arXiv:1904.01145 , 2019.

Brett W Larsen, Stanislav Fort, Nic Becker, and Surya Ganguli. How many degrees of freedom do we need to train deep networks: a loss landscape perspective. arXiv preprint arXiv:2107.05802 , 2021.

Vladislav Lialin, Sherin Muckatira, Namrata Shivagunde, and Anna Rumshisky. Relora: Highrank training through low-rank updates. In Workshop on Advancing Neural Network Training: Computational Efficiency, Scalability, and Resource Optimization (WANT@ NeurIPS 2023) , 2023.

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 arXiv:2404.00438 , 2024.

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.

Chris J Maddison, Daniel Paulin, Yee Whye Teh, Brendan O'Donoghue, and Arnaud Doucet. Hamiltonian descent methods. arXiv preprint arXiv:1809.05042 , 2018.

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

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.

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

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

Yiming Wang, Yu Lin, Xiaodong Zeng, and Guannan Zhang. Multilora: Democratizing lora for better multi-task learning. arXiv preprint arXiv:2311.11501 , 2023.

Wei Wen, Cong Xu, Feng Yan, Chunpeng Wu, Yandan Wang, Yiran Chen, and Hai Li. Terngrad: Ternary gradients to reduce communication in distributed deep learning. Advances in neural information processing systems , 30, 2017.

Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient llm training by gradient low-rank projection. arXiv preprint arXiv:2403.03507 , 2024.

A Hamiltonian Examples

Example A.1. Momentum + Online Subspace Descent is

with Lyapunov function H ( W , ˆ M ) = L ( W ) + ∥ ˆ M ∥ 2 2 a , for which

with Lyapunov function H ( W , M , V ) = L ( W ) + 1 2 a 〈 ˆ M √ ˆ V + e , ˆ M 〉 , for which

where we assume a ≥ b/ 4 .

Example A.3. The LionK + Online Subspace Descent is

Consider the Hamiltonian function in Eq (13) of [4]:

where we defined [ X ; Y ] ∇K = ⟨ Y , ∇K ( X + Y ) -∇K ( X ) ⟩ and used the fact that [ X ; Y ] ∇K ≥ 0 by the convexity of K ; we used ⟨ G t , P t X t ⟩ = ⟨ P ⊤ t G t , X t ⟩ = ⟨ ˆ G t , X t ⟩ .

In all examples above, although the form of Hamiltonian H ( W , ˆ S ) is independent of the update rule of P t , the decreasing rate d d t H ( W t , ˆ S t ) depends on P t in a complicate way through ˆ M t , ˆ V t , ˆ G t . An interesting direction for future investigation is to find optimal rules of P t to maximize the decreasing rate as an optimal control problem.

B Experiments

B.1 Hyperparameters

We sweep learning rate from [0.01, 0.005, 0.001]. For GaLore as well as Adam8bit, we follow the recommended hyperparameters as much as possible. For instance, GaLore update gap is set to recommended default, 200 . Warmup is set to 10% of the total training steps. Batch size is set to 512 and gradient clipping is set to 1 . 0 .

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 as the rank increases, both Dynamic Subspace Descent and GaLore improve.

B.3 Optimizer Sweep

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.

It is fine to include aspirational goals as motivation as long as it 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]

Justification: We discussed the limitation and difficulty in analyzing the class of method theoretically.

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 in 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 on 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 a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

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

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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 or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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)?

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 a 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 a 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 it 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 to 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 a 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 submissions 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 to reproduce that algorithm.

(b) If the contribution is primarily a new 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 a way to access this model for reproducing the results or a 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 to 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., to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: 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 ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we 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 a 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: //nips.cc/public/guides/CodeSubmissionPolicy ) 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 a subset of experiments are reproducible, they 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 in 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, hyperparameters, 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 not include experiments.

The experimental setting should be presented in the core of the paper to a 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, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer:[No]

Justification: Training LLM is 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.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence 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.)

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.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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 error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: GPU type used in experiment is specified.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or 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 as 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 it 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 https://neurips.cc/public/EthicsGuidelines ?

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 a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: Our paper aims to promote deeper theoretical understanding of a class of method. The algorithm itself is not practical enough to 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.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: same reasoning as above

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 to 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 a best faith effort.

12. Licenses 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 not use 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 a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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.

If this information is not available online, the authors are encouraged to reach out to the asset'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/model 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 as 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 crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, 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 nor 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 NeurIPS 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe paper has discussed limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-06-29"}

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

The authors tested their method against GaLore using perplexity as a metric while training on the C4 dataset.

STRENGTHS:
- the paper is well written and easy to follow
- having a dynamically changing projection matrix based on the gradient landscape makes total sense
- the theoretical explanation shows an interesting perspective on why this proposed method of using online pca for matrix projection could work
- 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.
- while the theory is interesting, it seems to include many strong assumptions that do not reflect real-world datasets
- 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.

QUESTIONS:
please answer how you could address the weaknesses above

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

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

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.

QUESTIONS:
Could you provide more high-level intuitions on why the methods could converge under the complications?

LIMITATIONS:
Yes, the authors discuss the limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nMemory-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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\nIt 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\"\n\nLIMITATIONS:\nEvaluations for production-size LLMs are missing.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}