TLDR: A new optimization method called GUM (GaLore Unbiased with Muon) addresses the convergence issues of existing memory-efficient LLM training techniques like GaLore. By incorporating a layerwise sampling debiasing technique, GUM achieves theoretical convergence guarantees similar to the Muon optimizer while maintaining memory efficiency. Empirical results show GUM consistently outperforms GaLore and often surpasses full-parameter training in LLM fine-tuning and pre-training tasks, attributed to a more uniform knowledge distribution within model layers and improved memorization.
Training large language models (LLMs) demands immense computational resources, particularly GPU memory. To tackle this challenge, researchers have developed memory-efficient optimization techniques, with gradient low-rank projection being a prominent strategy. These methods, such as GaLore, aim to reduce the memory footprint by storing only a projected version of the optimizer states.
However, a significant hurdle for many existing low-rank projection methods is their inherent bias, which can lead to a lack of convergence guarantees and performance gaps compared to full-parameter training. This bias arises because the low-rank projections don’t perfectly preserve the true gradient’s direction and magnitude, especially in high-dimensional spaces, potentially causing slower convergence or reduced model quality.
A new research paper, titled Unbiased Gradient Low-Rank Projection, introduces a novel solution to this problem: GaLore Unbiased with Muon, or GUM. This method investigates a layerwise sampling technique to debias low-rank projection mechanisms, aiming to restore the convergence guarantees while retaining memory efficiency.
How GUM Works
GUM builds upon the GaLore mechanism and integrates concepts from the Muon optimizer. The core idea is to compensate for the biased errors introduced by low-rank projections. It achieves this by allowing some blocks of parameters to be sampled uniformly with a certain probability in each training period to undergo full-rank updates. Other blocks continue with the original low-rank update. By carefully balancing the scaling constants for these two types of updates, the biased low-rank term can be effectively canceled out in expectation, leading to an unbiased estimation of gradients across iterations.
Theoretically, GUM is proven to match the convergence guarantees of the base Muon algorithm, a significant advantage over many biased low-rank methods. This means it can achieve the memory reduction benefits of low-rank techniques without sacrificing the strong theoretical foundation of full-parameter optimizers.
Empirical Performance
The paper presents compelling empirical evidence for GUM’s effectiveness across various LLM training scenarios:
- Synthetic Settings: In a noisy linear regression problem, GaLore-Muon failed to converge, while GUM successfully converged to an accuracy comparable to the full-parameter Muon baseline, demonstrating the clear benefit of the unbiased approach in challenging noisy environments.
- LLM Fine-tuning: GUM consistently outperformed GaLore in instruction-following (IFEval) and mathematical reasoning (GSM8K) tasks. Surprisingly, GUM even achieved better performance than full-parameter training methods like AdamW and Muon in some cases, particularly in enhancing reasoning capabilities. Memory efficiency tests showed GUM could achieve comparable or even better memory consumption than GaLore with optimized configurations.
- LLM Pre-training: On a suite of seven commonsense reasoning tasks, GUM consistently yielded better results than GaLore and, remarkably, often surpassed full-parameter training methods like AdamW and Muon. This improvement is attributed to GUM’s unbiased low-rank update mechanism, which better captures long-tailed gradient updates and enhances model memorization.
Also Read:
- SNOO: Supercharging Large Language Model Training with Nesterov Momentum
- Optimizing Large Language Model Training: A New Approach to Hyperparameter Scaling
Understanding the Improvement
Further analysis revealed that GUM’s empirical gains stem from its inherent high-rank updates. This leads to a higher overall stable rank and a more uniformly distributed set of singular values in the model parameters. This, in turn, results in more long-tailed activation patterns across modules, implying that GUM-trained models utilize their parameter space more efficiently and achieve better memorization of knowledge.
In conclusion, GUM offers a powerful and theoretically sound approach to memory-efficient LLM training. By effectively debiasing low-rank gradient projections, it provides strong convergence guarantees and consistently superior empirical performance, paving the way for more scalable and effective development of large language models.
