SubTrack++ : Gradient Subspace Tracking for Scalable LLM Training

2025/02/03 by Sahar Rajabi, Rajabi, Sahar, Nayeema Nonta +3 · 2 voices · 3 citations
Computer Science · #FOS: Computer and information sciences #Machine Learning (cs.LG) #cs.LG

paper · pdf · doi:10.48550/arxiv.2502.01586

arxiv published 2025/02/03 · arxiv updated 2025/10/25

Abstract

Training large language models (LLMs) is highly resource-intensive due to their massive number of parameters and the overhead of optimizer states. While recent work has aimed to reduce memory consumption, such efforts often entail trade-offs among memory efficiency, training time, and model performance. Yet, true democratization of LLMs requires simultaneous progress across all three dimensions. To this end, we propose SubTrack++ that leverages Grassmannian gradient subspace tracking combined with projection-aware optimizers, enabling Adam's internal statistics to adapt to subspace changes. Additionally, employing recovery scaling, a technique that restores information lost through low-rank projections, further enhances model performance. Our method demonstrates SOTA convergence by exploiting Grassmannian geometry, reducing pre-training wall-time by up to 65% and fine-tuning time by 36% compared to existing SOTA methods, while maintaining the same memory footprint.

Cited by

• FOAM: Blocked State Folding for Memory-Efficient LLM Training
• Parameter-Efficient Subspace Optimization for LLM Fine-Tuning
• Randomized Gradient Subspaces for Efficient Large Language Model Training

Discussions

• 💡Okay, I am so excited to announce the preprint of my lab’s work on memory and time-efficient training of #LLMs, especially with #DeepSeek’s success bringing this to the centerstage! arxiv.org/abs/25 [bsky, 3 points, 1 comments]
• Some training methods find a low-rank subspace to update the gradient in. This improves upon that by using something called a Grassmanian. arxiv.org/abs/2502.01586 [bsky, 0 points, 0 comments]

Related

• Randomized Gradient Subspaces for Efficient Large Language Model Training
• Fira: Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint?
• Breaking the Frozen Subspace: Importance Sampling for Low-Rank Optimization in LLM Pretraining
• LDAdam: Adaptive Optimization from Low-Dimensional Gradient Statistics
• Gradient Multi-Normalization for Stateless and Scalable LLM Training