Lotus: Efficient LLM Training by Randomized Low-Rank Gradient Projection with Adaptive Subspace Switching
This addresses the efficiency challenge in training large language models, offering improvements in both time and memory usage, though it appears incremental as it builds on existing low-rank gradient methods like GaLore.
The paper tackles the trade-off between memory consumption, training time, and model performance in large-scale model training by proposing Lotus, a method that modifies the projection process to enable efficient transitions between low-rank gradient subspaces, resulting in a 30% reduction in training time and a 40% decrease in memory consumption while outperforming baselines in pre-training and fine-tuning tasks.
Training efficiency in large-scale models is typically assessed through memory consumption, training time, and model performance. Current methods often exhibit trade-offs among these metrics, as optimizing one generally degrades at least one of the others. Addressing this trade-off remains a central challenge in algorithm design. While GaLore enables memory-efficient training by updating gradients in a low-rank subspace, it incurs a comparable extra training time cost due to the Singular Value Decomposition(SVD) process on gradients. In this paper, we propose Lotus, a method that resolves this trade-off by simply modifying the projection process. We propose a criterion that quantifies the displacement of the unit gradient to enable efficient transitions between low-rank gradient subspaces. Experimental results indicate that Lotus is the most efficient method, achieving a 30% reduction in training time and a 40% decrease in memory consumption for gradient and optimizer states. Additionally, it outperforms the baseline method in both pre-training and fine-tuning tasks.