Reproduce “Lotus: Efficient LLM Training by Randomized Low-Rank Gradient Projection with Adaptive Subspace Switching” on OpenResearch

About this paper

Tianhao Miao, Zhongyuan Bao, Lejun Zhang

LOTUS, a method for efficient Large Language Model (LLM) training, integrates randomized low-rank gradient projection with an adaptive subspace switching mechanism. It demonstrated a 30% reduction in training time and a 40% decrease in memory consumption for gradient and optimizer states compared to GaLore, while also outperforming baseline methods in both pre-training perplexity and GLUE benchmark scores.

Key insights

• Randomized SVD can efficiently approximate gradient decomposition for low-rank projection without significant loss in model quality, mitigating the computational overhead of exact SVD.
• The effectiveness of a low-rank gradient subspace can be quantified and adaptively managed by monitoring the consistency of unit-norm gradient directions within it.
• Dynamically adapting the gradient subspace based on optimization progress, rather than fixed schedules, leads to more effective gradient capture and faster convergence.

Summary by alphaXiv.