Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation
| Authors | Zhengbo Wang et al. |
| Year | 2026 |
| Field | Machine Learning |
| arXiv | 2602.24283 |
| Download | |
| Categories | cs.LG, cs.AI, cs.CL |
Abstract
Modern optimizers like Adam and Muon are central to training large language models, but their reliance on first- and second-order momenta introduces significant memory overhead, which constrains scalability and computational efficiency. In this work, we reframe the exponential moving average (EMA) used in these momenta as the training of a linear regressor via online gradient flow. Building on this equivalence, we introduce LoRA-Pre, a novel low-rank optimizer designed for efficient pre-training. Specifically, LoRA-Pre reduces the optimizer's memory footprint by decomposing the full momentum matrix into a compact low-rank subspace within the online linear learner, thereby maintaining optimization performance while improving memory efficiency. We empirically validate LoRA-Pre's efficacy by pre-training models from the Llama architecture family, scaling from 60M to 1B parameters. LoRA-Pre achieves the highest performance across all model sizes. Notably, LoRA-Pre demonstrates remarkable rank efficiency, achieving comparable or superior results using only 1/8 the rank of baseline methods. Beyond pre-training, we evaluate LoRA-Pre's effectiveness in fine-tuning scenarios. With the same rank, LoRA-Pre consistently outperforms all efficient fine-tuning baselines. Specifically, compared to standard LoRA, LoRA-Pre achieves substantial improvements of 3.14 points on Llama-3.1-8B and 6.17 points on Llama-2-7B, validating our approach's effectiveness across both pre-training and fine-tuning paradigms. Our code is publicly available at https://github.com/mrflogs/LoRA-Pre.
Engineering Breakdown
Plain English
This paper proposes LoRA-Pre, a new optimizer that reduces memory overhead in large language model training by decomposing momentum matrices into low-rank approximations. Instead of storing full first- and second-order momentum tensors like Adam does, the authors reframe momentum as an online linear regression problem and compress it, achieving better memory efficiency without sacrificing training performance.
Key Engineering Insight
The critical insight is treating exponential moving averages (EMA) in optimizers as equivalent to online gradient flow in a linear regressor—this mathematical reframing enables you to apply low-rank decomposition techniques to optimizer states, which weren't previously understood as compressible in this way.
Why It Matters for Engineers
For teams training large models, optimizer memory is a real constraint today. Adam requires storing two full copies of model parameters (for first and second moments), doubling effective memory usage. LoRA-Pre directly addresses this bottleneck by shrinking optimizer state size, which means you can either train larger models on the same hardware, use smaller GPUs, or fit longer context windows—all with the same training dynamics.
Research Context
Prior work (Adam, Muon) focused on improving convergence speed but accepted high memory costs as a trade-off. This paper shifts the problem frame by asking whether optimizer state itself can be compressed without sacrificing convergence. It advances the emerging pattern of applying structural compression (low-rank, sparsity) to parts of the training pipeline beyond just model weights, enabling more efficient scaling of LLM pre-training.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::