Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation
Authors: Zhengbo Wang, Jian Liang, Ran He, Zilei Wang, Tieniu Tan
Organizations: University of Science and Technology of China · NLPR & MAIS, Institute of Automation, Chinese Academy of Sciences · School of Artificial Intelligence, University of Chinese Academy of Sciences · Nanjing University
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.
Figures & tables
Figure 1: Illustration of our LoRA-Pre method. In this work, we establish a novel connection: the exponential moving average (EMA) update for optimizer momentum is mathematically equivalent to training a linear regressor using online gradient descent . Leveraging this equivalence, we propose compressing the optimizer states (i.e., the momenta) using low-rank matrices to reduce the memory footprint. Finally, the closed-form update rules for these matrices without requiring back-propagation are given by Theorem 3.1 .
Model Size
60M
130M
350M
1B
r/dmodel
128 / 512
256 / 768
256 / 1024
512 / 2048
Training Tokens
1.1B
2.2B
6.4B
13.1B
Adam ( Kingma and Ba, 2015 )
34.09
25.08
18.80
15.56
Muon ( Jordan et al., 2024 )
28.43
21.86
16.17
13.41
GaLore ( Zhao et al., 2024 )
34.88
25.36
18.95
15.64
Low-Rank ( Kamalakara et al., 2022 )
78.18
45.51
37.41
142.53
Table 1: Comparison with low-rank algorithms on pre-training various sizes of Llama models on the C4 dataset. We report the validation perplexity ( ↓ ) on a hold-out C4 test set. The best and second-best performance within the low-rank optimizers are highlighted with bold and underline . ∗ denotes the results are reproduced by ourselves.
Method
Llama-3.1-8B
Llama-2-7B
GSM8K
MATH-500
Average
GSM8K
MATH-500
Average
LoRA ( Hu et al., 2022 )
70.76
17.06
43.91
44.62
7.34
25.98
rsLoRA ( Kalajdzievski, 2023 )
71.06
17.46
44.26
48.79
5.75
27.27
DoRA ( Liu et al., 2024 )
71.06
17.86
44.46
44.39
6.55
25.47
GaLore ( Zhao et al., 2024 )
65.08
18.65
41.87
36.44
8.33
22.39
LoRA-Pre Adam
76.44
17.66
47.05
57.35
6.94
32.15
Table 2: Results of memory-efficient fine-tuning methods. We compare our method with efficient fine-tuning methods including LoRA ( Hu et al., 2022 ) , rsLoRA ( Kalajdzievski, 2023 ) , and DoRA ( Liu et al., 2024 ) , and an efficient optimizer GaLore ( Zhao et al., 2024 ) . The models are fine-tuned on the MetaMath100k ( Yu et al., 2024 ) dataset, and evaluated on GSM8K ( Cobbe et al., 2021 ) and MATH-500 ( Lightman et al., 2024 ) . We highlight the best performance on Adam-like optimizer and Muon-like optimizer with bold .
Figure 2: Rank efficiency comparison across efficient optimization methods. Perplexity versus rank for 60M (left) and 130M (right) models, demonstrating LoRA-Pre’s superior performance at lower ranks compared to baseline methods.
Model Size
60M
130M
350M
Muon ( Jordan et al., 2024 )
28.43
21.86
16.17
Muon w/o momentum
32.15
24.23
17.33
GaLore Muon ( Zhao et al., 2024 )
34.39
25.16
19.24
Fira Muon ( Chen et al., 2024 )
34.45
24.85
17.40
LoRA-Pre Muon
30.76
23.05
16.97
Table 3: Results of pre-training using different efficient Muon optimizers.
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
Parameter
Value
Perplexity
Status
β1 (fix β2=0.95 )
0.90 ( Default )
32.57
Optimal
0.95
37.62
Stable
0.99
1458.92
Diverged
β2 (fix β1=0.90 )
0.90
34.61
Stable
0.95 ( Default )
32.57
Optimal
0.999
1301.58
Diverged
Appendix
Table 4: Sensitivity Analysis of Momentum Parameters β . We report the validation perplexity on the 60M model. The method exhibits stability around the default settings derived from standard Adam ( β1=0.9,β2=0.95 ). Extreme values (e.g., β→1 ) cause the implicit γ to vanish, leading to training instability.
Low-rank adaptation (LoRA) makes finetuning large language models cheaper by adding to each weight matrix a trainable low-rank update parameterized as the product of two matrices. These matrices are usually trained with Adam, which treats them as a single flat vector of parameters and ignores both the matrix and product structure of LoRA. Applying a matrix-aware optimizer such as Muon to each factor does not consistently improve over Adam, and neither do the product-aware Muon variants proposed in concurrent works. To realize consistent gains, we introduce PoLoRA, a Preconditioned Orthogonalized LoRA optimizer built from three ingredients: a product-aware spectral update direction, curvature preconditioning derived from controlling the per-sample loss change, and a magnitude rule that controls the sizes of both the factor and merged updates. We evaluate PoLoRA on instruction-tuning datasets for code and math across models from 1B to 8B parameters, and find that it reaches the final held-out loss achieved by tuned Adam in 1.2-1.7 times fewer steps, while adding at most 3% per-step overhead. Compared to Adam, PoLoRA is also less sensitive to the learning rate, and its optimal learning rate is stable across ranks.
Nikhil Ghosh, Tetiana Parshakova, Robert M. Gower
Center for Computational Mathematics, Flatiron Institute, New York, NY
Fine-tuning large language models (LLMs) aims to adapt pre-trained models to specific tasks using relatively small and domain-specific datasets. Among Parameter-Efficient Fine-Tuning (PEFT) methods, Low-Rank Adaptation (LoRA) stands out by matching the performance of full fine-tuning while avoiding additional inference latency. In this paper, we propose a novel PEFT method that incorporates the TLoRA+ optimizer into the weight matrices of pre-trained models. The proposed approach not only preserves the efficiency of low-rank adaptation but also further enhances performance without significantly increasing computational cost. We conduct experiments on the GLUE benchmark across diverse model architectures. Numerical experiments consistently demonstrate the effectiveness and robustness of our proposed method.
Yarui Cao, Kai Liu
Department of Computer Science, Clemson University, Clemson, SC 29634, USA
Training large language models (LLMs) entails a fundamental trade-off: memory-efficient optimizers such as Adam discard cross-parameter curvature, whereas full-curvature methods such as SOAP can accelerate convergence at prohibitive memory costs. We introduce Clean, a memory-efficient and full-curvature optimizer designed to resolve this bottleneck. Clean leverages the randomized Nystrom method to accurately approximate the left and right preconditioners in SOAP, and to reduce the optimizer's memory complexity from quadratic to linear in terms of model dimensions. We subsequently reintegrate the off-subspace components to capture curvature information beyond the low-rank approximation, preserving rich curvature at minimal memory cost. We further propose Q-Clean, a low-precision variant that aggressively compresses optimizer states. Q-Clean reduces optimizer memory consumption by over 50% compared to Muon when pre-training a LLaMA-1.3B architecture, all while maintaining strong and competitive predictive performance. Notably, Clean operates with a smaller optimizer-state footprint than standard AdamW while reaching AdamW's final performance 26% faster in wall-clock time. Furthermore, our methods uniquely enable the pre-training of a 13B-parameter model on a single 80GB GPU, providing a scalable, efficient, and accessible approach to large-scale model optimization.
Beheshteh T. Rakhshan, Sahar Rajabi, Maziar Sargordi +3
Mila & DIRO, Université de Montréal · Critical ML, University of Waterloo · Independent Researcher +2