Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation

Zhengbo Wang; Jian Liang; Ran He; Zilei Wang; Tieniu Tan

cs.LGFeb 27, 2026

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

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.

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/d_{\text{model}}$	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 (

\downarrow

) 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
Method	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 .

Panel of Figure 2 — 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
$\beta_{1}$ (fix $\beta_{2}=0.95$ )	0.90 ( Default )	32.57	Optimal
	0.95	37.62	Stable
	0.99	1458.92	Diverged
$\beta_{2}$ (fix $\beta_{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

\beta

. We report the validation perplexity on the 60M model. The method exhibits stability around the default settings derived from standard Adam (

\beta_{1}=0.9,\beta_{2}=0.95

). Extreme values (e.g.,

\beta\to 1

) cause the implicit

\gamma

to vanish, leading to training instability.

Memory-Efficient Training

Memory-efficient training reduces the memory needed to optimize neural networks by compressing, recomputing, or relocating training state.

See all 87 papers

PoLoRA: A Preconditioned Orthogonalized LoRA Optimizer
Jul 20, 2026Nikhil Ghosh, Tetiana Parshakova, Robert M. GowerLLM Fine-Tuning Low-Rank Adaptation
Center for Computational Mathematics, Flatiron Institute, New York, NY
TLoRA+: A Low-Rank Parameter-Efficient Fine-Tuning Method for Large Language Models
Apr 15, 2026Yarui Cao, Kai LiuLLM Fine-Tuning Parameter-Efficient Fine-Tuning
Department of Computer Science, Clemson University, Clemson, SC 29634, USA
Clean: Second-order LLM Training at Linear Memory Cost via Nystr{ö}m Sketching
Oct 3, 2026Beheshteh T. Rakhshan, Sahar Rajabi, Maziar Sargordi +3Nyström Method Second-Order Optimization
Mila & DIRO, Université de Montréal · Critical ML, University of Waterloo · Independent Researcher +2

Jul 20, 2026J/K move · Enter open · S save

PoLoRA: A Preconditioned Orthogonalized LoRA Optimizer

Nikhil Ghosh, Tetiana Parshakova, Robert M. Gower

Center for Computational Mathematics, Flatiron Institute, New York, NY

LLM Fine-Tuning Low-Rank Adaptation Parameter-Efficient Fine-Tuning

Open paper AskPDF arXiv

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.

Abstract

Memory-Efficient Training

Explore similar work

PoLoRA: A Preconditioned Orthogonalized LoRA Optimizer

TLoRA+: A Low-Rank Parameter-Efficient Fine-Tuning Method for Large Language Models

Clean: Second-order LLM Training at Linear Memory Cost via Nystr{ö}m Sketching