Lotus: Efficient LLM Training by Randomized Low-Rank Gradient Projection with Adaptive Subspace Switching

TL;DR

Lotus is a novel training method for large-scale language models that reduces training time and memory usage by adaptively switching gradient subspaces, outperforming existing low-rank gradient projection techniques.

Contribution

It introduces an adaptive subspace switching criterion that improves training efficiency and performance in large-scale models compared to prior low-rank gradient methods.

Findings

30% reduction in training time

40% decrease in memory consumption

Outperforms baseline in pre-training and fine-tuning

Abstract

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.