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

TL;DR

Lotus addresses the memory-time-performance bottleneck in large-language-model training by replacing fixed, SVD-based subspace switching with an adaptive, randomized low-rank gradient projection strategy. It introduces a path-efficiency metric $\rho_t$ to quantify subspace drift and triggers switching when alignment degrades, using a randomized SVD to approximate gradient decomposition. Theoretical results show faster convergence under adaptive switching, and experiments demonstrate roughly a 30% reduction in training time and a 40% reduction in memory usage compared to GaLore, with improvements on pre-training and GLUE benchmarks. This provides a practical, scalable approach for efficient large-model optimization with reduced hardware demands.

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.

Figures (2)

• Figure 1: The comparison between previous method(e.g. GaLore) with fixed switching frequency and our greedy search strategy that updates the subspace adaptively. $G^*$ is the displacement of the unit gradient in a subspace. When the average displacement of unit gradient vector $G_{unit}$ is lower than $\gamma$, the subspace will be switched.
• Figure 2: We visualize training time consumption comparsion between Galore, Adarankgrad, Apollo and Lotus on both pre-training and fine-tuning tasks. (a) shows the ETA of pre-training LLaMA-type 3B models in C4. (b) shows the average time cost on 8 GLUE tasks. Lotus is the most effective method in terms of computational time efficiency. (Less is better)