Research map / Large language models / Instruction and fine-tuning
SubTrack++ : Gradient Subspace Tracking for Scalable LLM Training
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 SubTrack++, a memory-efficient LLM training method that tracks gradient low-rank subspaces on the Grassmannian via rank-1 geodesic updates, combined with a projection-aware Adam and recovery scaling, to cut wall-time while matching GaLore's memory.
Problem
Low-rank gradient projection methods for LLM training (e.g., GaLore) trade off memory, wall-time, and performance; periodic SVD is expensive and noise-sensitive, and discarded gradient components hurt accuracy.
Approach
Replace periodic SVD subspace recomputation with Grassmannian geodesic tracking using rank-1 updates driven by the estimation error, make Adam's moment statistics projection-aware under subspace changes, and add recovery scaling that reintegrates the gradient component discarded by projection.
Key results
- Reduces pre-training wall-time by up to 65% versus existing SOTA methods while keeping GaLore-level memory (abstract).
- On Llama-1B pre-training, SubTrack++ reaches evaluation loss 4.52 at 1B versus LDAdam 4.70 and GaLore 6.53 (Table 1).
- SubTrack++ achieves 36% lower wall-time than GaLore in Llama-2-7B-chat SFT on Alpaca.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add the projection-aware optimizer to Grassmannian tracking (ablation) | SubTrack++ on Llama 1B · C4 pre-training | evaluation loss: 6.53 to 5.43 | Helped |
| Add recovery scaling to Grassmannian tracking (ablation) | SubTrack++ on Llama 1B · C4 pre-training | evaluation loss: 6.53 to 5.28 | Helped |
| Add both the projection-aware optimizer and recovery scaling (ablation) | Grassmannian tracking on Llama 1B · C4 pre-training | evaluation loss: 6.53 to 4.51 | Helped |
| Use rank-1 subspace updates (ablation) | SubTrack++ · C4 pre-training, Llama 1B | evaluation loss | Helped |
| Increase subspace update frequency (ablation) | SubTrack++ · C4 pre-training, Llama 1B | evaluation loss | Mixed |
| Use Grassmannian tracking instead of frequent SVD updates (ablation) | SubTrack++ · C4 pre-training | wall-time | Helped |
| Use SubTrack++ rather than GaLore | Llama 1B pre-training · C4 | wall-time: 43% faster | Helped |
Surprising findings
- SubTrack++ occasionally outperforms full-rank training in evaluation loss, attributed to implicit regularization from low-rank projection.
Limitations
- Not evaluated on models larger than 7B due to limited resources.
- Projection-aware optimization is only derived for the Adam family; extension to other optimizers needs further work.
- Convergence theory assumes fixed projection matrices, which does not fully cover slowly changing subspaces.
Experimental setup
Datasets: C4 GLUE SuperGLUE Alpaca
Benchmarks: C4 GLUE SuperGLUE Alpaca
Metrics: evaluation loss wall-time peak memory consumption accuracy
Baselines: Full-Rank GaLore BAdam Online Subspace Descent LDAdam Fira APOLLO
Builds on: GaLore Fira LDAdam APOLLO Online Subspace Descent
Closest papers by meaning
- A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models (ICML 2025 · similarity 0.89)
- Breaking the Frozen Subspace: Importance Sampling for Low-Rank Optimization in LLM Pretraining (NeurIPS 2025 · similarity 0.88)
- AdaRankGrad: Adaptive Gradient Rank and Moments for Memory-Efficient LLMs Training and Fine-Tuning (ICLR 2025 · similarity 0.88)
- Fira: Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint? (NeurIPS 2025 · similarity 0.87)
- GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection (ICML 2024 · similarity 0.87)
- From Low Rank Gradient Subspace Stabilization to Low-Rank Weights: Observations, Theories, and Applications (ICML 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.