Unbiased Gradient Low-Rank Projection
Rui Pan, Yang Luo, Yuxing Liu, Yang You, Tong Zhang
TL;DR
This paper tackles bias and convergence gaps in memory-efficient low-rank gradient methods for training large language models. It introduces GaLore Unbiased with Muon (GUM), a debiasing approach that uses layerwise sampling to mix full-rank compensated updates with low-rank Muon updates, preserving memory efficiency while achieving convergence guarantees comparable to full-parameter Muon. The authors prove a non-convex convergence bound for GUM that matches Muon's rate under constant sampling, and demonstrate, across synthetic tasks and LLM fine-tuning/pretraining, that GUM outperforms GaLore and can even surpass full-parameter training on several benchmarks. They further show that unbiased, high-rank updates yield a higher stable rank and more uniform singular-value distributions, which translate into better memorization and long-tail activation patterns across layers. Collectively, these results establish GUM as a scalable, theoretically solid alternative for memory-constrained training of large models.
Abstract
Memory-efficient optimization is critical for training increasingly large language models (LLMs). A popular strategy involves gradient low-rank projection, storing only the projected optimizer states, with GaLore being a representative example. However, a significant drawback of many such methods is their lack of convergence guarantees, as various low-rank projection approaches introduce inherent biases relative to the original optimization algorithms, which contribute to performance gaps compared to full-parameter training. Aiming to tackle this problem, this paper investigates the layerwise sampling technique for debiasing low-rank projection mechanisms. In particular, an instantiation of the paradigm gives rise to a novel and unbiased low-rank optimization method built upon GaLore's mechanism and the Muon algorithm, named GaLore Unbiased with Muon (GUM). We theoretically prove our method matches the convergence guarantees of the base Muon algorithm while preserving the memory efficiency of low-rank techniques. Empirical experiments on LLM fine-tuning and pretraining also demonstrate non-trivial improvements over GaLore and even better performance than full-parameter training. Further investigation shows that the improvement of this technique comes from a more uniform distribution of knowledge inside layers, leading to more efficient utilization of the model parameter space and better memorization.