Research map / Optimization / Optimizers

Memory-Efficient LLM Training with Online Subspace Descent

NeurIPS 2024 · Optimization › Optimizers · approach: Memory-efficient LLM optimizers · 3 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Online Subspace Descent, which updates the GaLore-style projection matrix with online PCA instead of SVD, and proves a convergence guarantee for subspace descent with arbitrary projection update rules under a Hamiltonian-descent analysis.

Problem

Memory-efficient LLM training via low-rank gradient projection (GaLore-style) relies on SVD to compute projection matrices, which is expensive and lacks convergence guarantees for arbitrary projection update rules on non-convex objectives.

Approach

Replace periodic SVD of the gradient with continuous online PCA updates of the projection matrix (one optimizer step on a PCA-style reconstruction loss per iteration, run in parallel with weight updates), and show via Hamiltonian-descent Lyapunov analysis that the optimizer converges to stationary points for any smooth projection update rule satisfying a mild condition.

Key results

• On LLaMA 7B pretrained on C4 for 10K steps, the method reaches perplexity 43.72 versus 51.21 for GaLore, with wall-clock time 7.14 h versus 9.74 h.
• On LLaMA 60M (seq len 1024), the method with Adafactor+AdamW8bit reaches perplexity 34.15 versus 46.90 for GaLore with Lion (the best GaLore variant reported in the optimizer table is 34.32 with Adafactor).
• Online PCA single-step updates are reported as up to 142x faster than a PyTorch SVD on representative LLaMA weight tensors.

Measured effects reported in the paper

Surprising findings

• Updating the projection matrix with plain SGD gives a poor result for AdamW8bit weights (53.53 perplexity), yet the authors note SGD can match Adam-based projection updates in some settings.
• Without any SVD, the method achieves lower perplexity than GaLore on the 7B model, and the paper notes its advantage grows with model size and sequence length.

Limitations

• Convergence analysis is for continuous-time dynamics and omits weight decay, which the authors leave to future work.
• The learning rate of the projection update must be kept small (alpha = 5) to avoid instability, and larger learning rates cause loss spikes.
• Experiments are mostly on small LLaMA models pretrained for only 10K steps on C4; the 7B run is short.

Experimental setup

Datasets: C4

Benchmarks: C4 validation perplexity GLUE

Metrics: perplexity wall-clock time GPU memory

Baselines: GaLore Full-rank training LoRA ReLoRA Adam8bit

Builds on: GaLore Stochastic Subspace Descent Adam Lion ReLoRA

Closest papers by meaning

• Compress to Impress: Efficient LLM Adaptation Using a Single Gradient Step on 100 Samples (NeurIPS 2025 · similarity 0.89)
• SUMO: Subspace-Aware Moment-Orthogonalization for Accelerating Memory-Efficient LLM Training (NeurIPS 2025 · similarity 0.88)
• From Low Rank Gradient Subspace Stabilization to Low-Rank Weights: Observations, Theories, and Applications (ICML 2025 · similarity 0.87)
• A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models (ICML 2025 · similarity 0.87)
• Fira: Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint? (NeurIPS 2025 · similarity 0.87)
• SLoPe: Double-Pruned Sparse Plus Lazy Low-Rank Adapter Pretraining of LLMs (ICLR 2025 · similarity 0.87)