Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation
Authors: Zhengbo Wang
Jian Liang
Ran He
Zilei Wang
Tieniu Tan
Venue: ICLR 2026 Oral (2026)
Keywords: Large Language Models; Efficient Training; Low-Rank; LoRA
Primary area: foundation or frontier models, including LLMs

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.

Published as a conference paper at ICLR 2026
TAMING MOMENTUM:
RETHINKING OPTIMIZER
STATES THROUGH LOW-RANK APPROXIMATION
Zhengbo Wang 1.2 Jian Liang Ran He Zilei Wang'
Tieniu Tan
University of Science and Technology of China
2 NLPR & MAIS, Institute of Automation, Chinese Academy of Sciences
3 School of Artificial Intelligence, University of Chinese Academy of Sciences
Nanjing University
ngggbbowangg@ail www.cstcc.. ustc.edu. cr, cn, 1 liangjian gma com
ABSTRACT
Modern optimizers like Adam and Muon are central to training large language models, but their reliance on first- and second-order momenta introduces sig- nificant memory overhead, which constrains scalability and computational effi- ciency. 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 this equivalence, mizer designed for efficient pre-training. Specifically, LoRA-Pre reduces the optimizer's 's memory footprint by decomposing the full momentum matrix into a compact low-rank subspace within the online linear learner, thereby maintain- ing optimization performance while improving memory efficiency. We empir- ically validate LoRA-Pre's efficacy by pre-training models from the Llama ar- chitecture family, scaling from 60M to 1B parameters. LoRA-Pre achieves the highest performance across all model sizes. Notably, LoRA-Pre demonstrates re- markable rank efficiency, achieving comparable or superior results using only 1/8 the rank of baseline methods. Beyond pre-training, we evaluate LoRA-Pre's ef- fectiveness in fine-tuning scenarios. With the same rank, LoRA-Pre consistently outperforms all efficient fine-tuning baselines. Specifically, compared to standard LoRA, LoRA L-Re-Pre achieves substantial improvements of 3.14 points on Llama- 3.1-8B and 6.17 points Llama-2-7B, validating our approach's effectiveness across both pre-training and fine-tuning paradigms. Our code is publicly available httpp::/ /i huub. com/ http///rffoo.o......
introduce LoRA-Pre, novel low-rank opti-
INTRODUCTION
Large Language Models (LLMs) (Guo et al., 2025; Yang et al., 2025; Grattafiori et al., 2024; Brown et al., 2020; Comanici et al., 2025; Touvron et al., 2023; Jaech et al., 2024) have become the cen- terpiece of modern deep learning. Trained on trillions of tokens from heterogeneous sources and scaled to unprecedented parameter counts, they demonstrate remarkable generalization and transfer capabilities. Beyond this, some LLMs have reasoning abilities and leverage external tools (Guo et al., 2025; Yang et al., 2025). These advances have transformed LLMs from statistical learners into versatile systems, driving breakthroughs across research and real-world applications.
However, the success of LLMs comes with formidable training and adaptation costs (Grattafiori et al., 2024). The vast number of trainable parameters demands enormous memory and compu- tational resources during pre-training and fine-tuning. A key contributor this burden lies in the optimizer states. For instance, Adam (Kingma & Ba, 2015), the de facto optimizer for training LLMs, maintains not only the model weights but also first- and second-order moment estimates of the gradients. This triples memory usage and further exacerbates scalability bottlenecks, underscor- ing the urgent need for more efficient optimization strategies.
"Coreesooningg author.

---

Published as a conference paper at ICLR 2026
To address this, series of low-rank optimization methods has emerged. One prominent line of research achieves optimizer state compression through projected gradient descent (Zhao et al., 2024; Chen al., 2024; Hao et al., 2024; Modoranu al., 2025). These methods initialize projection matrices via SVD) or random mappings, project gradients into smaller subspaces for optimizer state computation, and then project them back to achieve parameter updates, thereby compressing the optimization overhead. Additionally, such methods require periodic subspace updates to enable high-rank parameter updates following W AWT; + AWT, +.... However, due to the inability to update subspaces instantly, error accumulation occurs in optimizer state computation, leading to suboptimal performance. This motivates the need for a more dynamic approach that can rapidly adapt to changing gradient subspaces.
In this paper, we propose LoRA-Pre, a novel low-rank optimizer for L.M pre-training that ad- dresses these limitations through different approach. Our key insight is an interesting mathemati- cal connection between the exponential moving average (EMA) of momentum and linear regression. Specifically, we demonstrate that EMA momentum updates are mathematically equivalent to train- ing an online linear regressor with gradient descent on the online gradient flow:
1 min g) - -1mm 11mmgll-- 12
m + m- 3)
where m E R represents the momentum, is the online gradient, and / is the momentum coefficient. This equivalence reveals that momentum accumulation can viewed as fitting linear model to approximate the gradient history.
Leveraging this theoretical insight, we develop a memory-effiiint optimizer through momentum compression via low-rank factorization. Instead of maintaining the full momentum matrix m, decompose it as the product of two low-rank matrices: m mg mA, where mg E Pg and m. R'xq, with < min(p, q). This factorization reduces memory complexity from p q to (p + 9) x r, yielding substantial memory savings for large-scale models. The low-rank momentum then updated by solving milmm,.. [MMA, m. g) pm glls. with explicit update rules derived in Theorem 3.1.
This theoretical framework enables to compress any momentum-based optimizer. We demon- strate its versatility by developing LoRA-Pre variants for both Adam (Kingma & Ba, 2015) and Muon (Jordan al., 2024) optimizers, with detailed algorithms provided Appendix Extensive experiments across pre-training and fine-tuning tasks validate the effectiveness of our method, while ablation studies demonstrate strong robustness across different rank variations.
Our main contributions are summarized as follows:
We establish a novel theoretical connection showing that exponential moving aver- age (EMA) momentum updates are mathematically equivalent to training linear regressor via online gradient flow.
Based on this insight, we propose LoRA-Pre, memory-efficient low-rank optimizer for pre-training that compresses optimizer states by factorizing the momentum matrix into low- rank components. We construct LoRA-Pre variants for both Adam and Muon optimizers, mathematically derive their low-rank update rules through our regression formulation, and achieve substantial memory reduction while preserving optimization dynamics.
We provide extensive experimental validation across both pre-training and fine-tuning tasks, demonstrating that LoRA-Pre achieves superior performance with remarkable rank efficiency compared to existing basclines, confirming both the efficiency and effectiveness of our approach across diverse model scales and application scenarios.
2 RELATED WORKS
Low-Rank Adaptation. The scaling of Large Language Models (LLMs) has spurred the develop- ment of Parameter- Etrameter-Ffficientt Fine-Tuning (PEFT) methods (Hu et al., 2022; Liu et al., 2024: Wang et al., 2025; Ding et al., 2023; Liu et al., 2022; 2023: Hayou et al., 2024; Wang et al., 2024a; Edalati et al., 2023; Zhang et al., 2023; Tastan et al., 2025; Wang et al., 2023; 2024c;b), which aim to adapt pre-trained models downstream tasks with reduced computational and memory overhead. Among
2

---

Published as a conference paper at ICLR 2026
these PEFT methods, Low-Rank Adaptation (LoRA) (Hu al., 2022) and its variants (Wang al., 2023; 2024a; Hayou et al., 2024; Liu et al., 2024; Yen et al., 2025) have emerged as the predominant methods in the field.
LoRA is grounded in the principle that weight updates during fine-tuning possess an intrinsic low- rank structure (Aghajanyan et al., 2021). By re-parameterizing these updates the product of two low-rank matrices, LoRA substantially reduces the number of trainable parameters while maintain- ing competitive performance, thereby enabling efficient adaptation of LLMs with limited computa- tional resources. The effectiveness of LoRA has inspired a line of research aimed at addressing its shortcomings. For instance, LoRA+ (Hayou et al., 2024) introduces differential learning rates for the two low-rank matrices to improve convergence and final task performance. DoRA (Liu et al., 2024) decomposes pre-trained weights into magnitude and direction components, applying LoRA specifically to the directional component to better capture fine-tuning dynamics. Recent works like LoFT (Tastan et al., 2025) and LoRA-Pro (Wang et al., 2025) establish theoretical connections be- tween LoRA and full fine-tuning via projected gradient equivalents.
While effective for fine-tuning, existing LoRA-based methods face fundamental challenges when applied to pre-training from scratch. Unlike fine-tuning, where small adaptations naturally exhibit low-rank structure, pre-training from random initialization requires full-rank weight updates to learn diverse representations across the entire parameter space (Lialin et al., 2024; Kamalakara et al., 2022). This mismatch between LoRA's low-rank assumption and pre-training's 's full-rank require- ments results in suboptimal performance in the pre-training stage.
Low-Rank Pre- Trainiinn The pre-training cost of LLMs has surged dramatically with the rapid ex- pansion of model scale. A promising direction for mitigating these costs is compressing optimizer states into a low-rank subspace, a strategy that significantly reduces memory footprint and com- munication overhead (Zhao al., 2024: Modoranu al., 2025; Ma et al., 2025; Han et al., 2024; Zmushko et al., 2025; Chen et al., 2024; Hao et al., 2024; Shen et al., 2025; Mahdavinia & Mahdavi, 2025: Zhang al., 2025). For instance, GaLore (Zhao al., 2024) utilizes Singular Value De- composition (SVD) to project gradient information into low-rank subspace for state compression, subsequently projecting the optimized gradients back for parameter updates. To enhance compu- tational efficiency. Flora (Hao et al., 2024) substitutes the expensive SVD operation with random projection, while Fira (Chen et al., 2024) incorporates SGD momentum to leverage gradient infor- mation from the complementary subspace. However, these projection-basee methods typically rely on periodic subspace updates to amortize costs, which often results in optimization discontinuities and error accumulation due to the lag in subspace adaptation.
Recent works have explored online strategies to address these limitations. MLorc (Shen et al., 2025) employs randomized SVD for online momentum compression. MoFaSGD (Mahdavinia & Mahdavi, 2025) utilizes momentum factorization to approximate full-rank momentum online, ensuring non- convex convergence. Similarly, ADAPM (Zhang al., 2025) compresses first-order momentum into a low-rank subspace via linear regression. In contrast, our proposed LoRA-Pre fundamentally reformulates momentum maintenance as an online regression task. By directly evolving low-rank factors via online gradient flow at every step, our approach achieves continuous subspace adaptation, effectively eliminating the instabilities associated with periodic updates or heuristic approximations.
3 METHOD
We begin by revisiting the de facto standard optimizer, Adam (Kingma Ba, 2015), in Section 3.1. Then, we establish connection between the exponential moving average and an online linear re- gressor over past gradients in Section 3.2. Finally, Section 3.3 introduces our efficient optimizer, LoRA-Pre, which compresses optimizer states through low-rank parameterization.
3.1 PRELIMINARY
We first review Adam (Kingma & Ba, 2015), the de facto optimizer in modern deep learning, which combines the benefits of AdaGrad (Duchi et al., 2011) and RMSProp (Hinton et al., 2012) by main- taining estimates of the first and second moments of gradients to achieve adaptive learning rates and robust performance.
3

---

Published as a conference paper at ICLR 2026
Your Momentum Secret Online Linear Learner!
m' + B1 m + (1 - B1)
L(mg, Magi 9c) = mm - -9
Equivalent Form
m2g
1. (m; 9t)
[| Belli
0c+1
St
91-1 01
Low Rank Compression
m
g
Online Gradient Flow update:
Online Gradient Flow
mg mg - -BB)) )
L
at
3mh
3mg
update: m' (1 B.)
m/a - my (1 B1)
0t.
8m
@m
8m
Figure 1: Illustration of LoRA-Pre method. In this work, establish novel connection: the exponential moving average (EMA) update for optimizer momentum is mathematically equiva- lent training linear regressor using online gradienl descent. Leveraging this equivalence, 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.
Consider an optimization problem where X' represents a data point drawn from a distribution Pdotas ((-) : X IR is loss function, and e R2 are the optimized parameters. The Adam optimizer (Kingma & Ba, 2015) updates 9 according to the following steps:
d((x;)
~ Pdota (r),
(Gradient Computation)
(2)
00
m, = B1 - + (1 - 81) lt.
(EMA of the First Moment)
(3)
$2 01-1 + 822 li
(EMA of the Second Moment)
(4)
m4
Ut
m,
0,
(Bias-Correction)
(5)
3
0.11 = 02
(Parameter Update)
(6)
Vol
+
Here, m. and V represent the Exponential Moving Average (EMA) the first- and second-order moments, respectively. The hyper-parameters include the learning rate y, the exponential decay rates 31, 32 E [0, 1) for the moment estimates, and a small constant e for numerical stability.
Similar to the Adam optimization process, momentum also plays critical role in other modern opti- mizers (Shazeer & Stern, 2018; Jordan et al., 2024), enhancing stability and convergence. However, storing momentum states introduces significant memory overhead. Our work directly addresses this by compressing the momentum term to reduce the optimizer's 's memory footprint.
3.2
YOUR MOMENTUM IS A SECRET ONLINE LINEAR REGRESSOR
To begin with, reveal interesting connection: momentum updates in modern optimizers secretly performing online linear regression. Specifically, updating the momentum m via EMA mathematically equivalent optimizing m as the parameter of linear regressor using online gradient flow.
To illustrate this, consider the first-order momentum an example. The standard EMA update for the first-order momentum can be rewritten as follows:
M4+1
3 m. + (1
3)
(7)
my
(1
(m, 9)
(8)
weight
dr gradient
4

---

Published as a conference paper at ICLR 2026
As shown in Equation (8). the EMA update is mathematically equivalent to gradient descent step where the parameter being optimized is the momentum m, the learning rate is 1 1- and the gradient is [L[m.ig] = My g. This reformulation reveals that EMA updates essentially function an online 8m linear regressor that continuously adjusts the momentum weights based on incoming gradients. The underlying objective being minimized is:
1 min [(m; 9) = 2 |lm gll
This insight opens a new avenue for optimizer footprint optimization: since momentum parameters are linear model weights, we can apply standard model compression techniques reduce optimizer memory usage during training.
3.3 LORA-PRE: LOW-RANE ONLINE LINEAR REGRESSION
We now introduce LoRA-Pre, new low-rank optimizer for pre-training. Building on the equiv- alence between exponential moving averages and online linear regression, LoRA-Pre compresses the momentum term m via a low-rank factorization, inspired by the LoRA technique (Hu et al., 2022). This approach can be applied any momentum-based optimizer, such as Adam (Kingma & Ba, 2015) and Muon (Jordan et al., 2024). We detail the compression strategies for both first- and second-order momentum terms below.
First-Order Momentum Compression. Having established that momentum updates are equivalent gradient descent on the objective min.. L(m;; g) 1 I|m-gll in Section 3.2, can now apply low-rank compression to reduce memory usage. Instead of storing and updating the full momentum matrix m R2xq directly, we decompose it into the product of two low-rank matrices mg Rgxr and mA E Rxx,, r < min(p, 9), i.c., m = mg mA. This factorization transforms our original optimization problem into:
1 min L(MB, mAig) - -1 - gll
(10)
m MAMAA
To maintain memory efficiency, we solve this optimization problem using standard gradient descent on the factorized matrices m and m To ensure computational efficiency, derive closed-form update rules for these matrices without requiring back pree--ooppggtion, which is given by Theorem 3.1. We resort to Newton's method for updating since the solution can be expressed in the form of EMA.
Theorem 3.1. Assume both matrices mg E Rexr and mA E 1Rr are full rank. For the objective minman... L(mB, ((mB,mA:g) = |m gll Newton's 's method yields the following closed-form update rules:
mB + (1 - -n) mB + m gm T Am 1)
(11)
mA <- (1 - -1) m. + + (mīms) - mgg.. Here, 7n is the learning rate for the factorized optimization problem.
(12)
Proof. See Appendix
Second-Order Momentum Compression. The compression of second-order momentum v presents additional challenges due to the constraints imposed by Adam's parameter update rule. Since Equa- tion (6) requires the square root of momentum, i.e., Vū, v, the second-order momentum must be clement-wise positive.
A naive approach would parameterize the second momentum as v = Up and optimize using the regression loss L(VR, vAig) = |luDUA - g2 From Theorem 3.1, we derive the corresponding parameter update rule:
UB + (1 72) uB + 72 g² U (
(13)
VA + (1 (2) VA + (DIUB) ) -l 2og
(14)
Unfortunately, this approach cannot guarantee that Uū.j > 0, Vi,j. j. making the computation of /i VenuA problematic.
5

---

Published as a conference paper at ICLR 2026
To address this issue, we re- re-parameterize the second-order momentum (VB vA)²² where denotes the Hadamard product. This pe-aammtterritiionn ensures clement-wise positivity while maintaining the low-rank structure. We then formulate the optimization of low-rank matrices un and 1 as:
1
min L(UB, vAig) = lI -
(15)
B ad
Its update rule can thus be directly derived from Theorem 3.1.
v - (1 - 72) - un + 72 (v 3 - (1 - 72) UA + 7 (ugun) Egll
(16)
(17)
Low-Rank Optimizer Algorithms. As shown before, our method can be applied any optimizer with momentum to compress its optimizer state during pre-training and fine-tuning stages. The detailed pseudo-codes of LoRA-Pre optimizer for AdamW (Kingma & Ba, 2015) and Muon (Jordan et al., 2024) are provided in Appendix B.
4 EXPERIMENTAL RESULTS
In this section, we present extensive experiments to evaluate the effectiveness of our proposed method, LoRA-Pre. Our evaluation encompasses both memory-effiiiet pre-training and memory- efficient fine-tuning on downstream tasks.
We begin by assessing LORA-Pre's pre-training capabilities in Section 4.1. Following the exper- imental setup of GaLore (Zhao et al., 2024), we train Llama (Touvron et al., 2023) models from scratch with varying model sizes of 60M, 130M, 350M, and IB parameters. All models are trained on the Colossal Clean Crawled Corpus (C4) dataset (Raffel et al., 2020), a large-scale cleaned dataset specifically designed for language model pre-training. To simulate realistic pre-training conditions, the models are trained sufficiently large volumes of data without repetition.
Subsequently, we evaluate LoRA-Pre's fine-tuning performance in Section ??. We fine-tune both Llama-3.1 3.1--BB (Grattafiori al., 2024) and Llama-2-7B (Touvron et al., 2023) models 100k subset sampled from the MetaMathQA dataset (Yu et al., 2024). The fine-tuned models are then evaluated on the GSM8K (Cobbe et al., 2021) and MATH-500 (Lightman et al., 2024) datasets. Finally, we present an ablation study of LoRA-Pre in Appendix 4.3.
Implementation Details. To ensure fair comparison, we align the experimental setup with that of GaLore (Zhao et al., 2024). By default, LoRA-Pre is applied to all parameters in the attention and MLP layers, while other parameters are optimized using the standard Adam (Kingma Ba, 2015) optimizer. We set the default ranks for the 60M, 130M, 350M, and 1B parameter models to 128, 256, 256, and 512, respectively. The optimal learning rate is selected from the set (0.01, 0.005, 0,001, 0.0005, 0.0001} based on validation perplexity. To maintain strict fairness in comparison, we retain the same scale factor of 0.25 as used GaLore (Zhao et al., 2024). For memory-efficient fine-tuning tasks, we set the default rank as 8 and set the learning rate as 2e 5 by default.
4.1 MEMORY- EFFICIENT PRE- TRAINING
In this section, we evaluate the pre-training performance of our proposed method, LoRA-Pre. Our experimental setup strictly follows that of GaLore (Zhao et al., 2024). We compare LoRA- Pre against several baseline methods, including both full optimizers and low-rank optimizers: 1) Adam (Kingma & Ba, 2015): The de facto optimizer in modern deep learning that utilizes first- and second-order momentum statistics to dynamically adjust lcarning rates and stabilize training. 2) Muon (Jordan al., 2024): A novel preconditioned optimizer that updates parameters by or- thogonalizing the first-order momentum. 3) GaLore (Zhao et al., 2024): A low-rank optimizer that projects gradients using SVD and computes optimizer states a reduced subspace. 4) Low- Rank (Kamalakara et al., 2022): A traditional low-rank approach that directly represents weights through learnable low-rank factorization W = BA. 5) LoRA (Hu et al., 2022): The most widely adopted low-rank method for fine-tuning that factorizes weights as W W + B... For pre-training scenarios, we maintain Wo as the full-rant initialization matrix. 6) ReLoRA (Lialin et al., 2024): A LoRA variant designed for pre-training that periodically merges BAA into W and initializes BA

---

Published as a conference paper at ICLR 2026
Table Comparison with low-rank algorithms on pre-training various sizes of Llama models on the C4 dataset. We report the validation perplexity (4) on hold-out C4 test set. The best and second- best performance within the low-rank optimizers highlighted with bold and underline. * denotes the results are reproduced by ourselves.
Model Size
60M
130M
350M
1B 512 / 2048
r mmddel
128 512
256
768
256
1024
Training Tokens
1.1B
2.2B
6.4B
13.1B
Adam (Kingma & 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
LoRA (Hu ct al., 2022)
34.99
33.92
25.58
19.21
ReLoRA (Lialin et al., 2024)
37.04
29.37
29.08
18.33
SLTrain (Han et al., 2024)
34.15
26.04
19.42
16.14
LORO (Mo et al., 2025)
33.96
24.59
18.84
15.19
Fira* (Chen et al., 2024)
31.19*
24.51*
17.22
14.31
LoRA-Pre Adam
32.57
23.78
16.36
13.53
LoRA-Pre Muon
30.76
23.05
16.97
13.92
with optimizer state resets. 7) SLTrain (Han al., 2024): A sparse plus low-rank approach that pa-
rameterizes weights as W S BA, where both components are jointly optimized. 8) LORO (Mo al., 2025): A method that optimizes LoRA parameters by strictly constraining updates within the low-rank manifold. 9) Fira (Chen et al., 2024): A method that improves upon GaLore with the Norm-Based Scaling and the Norm-Growth Limiter.
We pre-trained Llama series models of different sizes to evaluate LoRA-Pre against these baseline methods. By default, all low-rank optimizers are built upon the Adam (Kingma Ba, 2015) opti- mizer foundation. To demonstrate the generalizability of our approach, we also evaluate LoRA-Pa with Muon (Algorithm 2), as our method is compatible with any momentum-
optimizer.
The results, presented in Table 1, demonstrate that method achieves superior performance across multiple model scales. Specifically, LoRA-Pre Adam and LoRA-Pre Muon attain either the highest second- second-highest performance across almost all four model sizes (60M, 130M, 350M, and IB), val- idating the effectiveness of our approach. While Fira yields competitive results on the 60M model, LoRA-Pre consistently outperforms it the larger scales (130M, 350M, and 1B), likely because our method avoids the етгоr accumulation associated with Fira's projected gradients. Furthermore, LoRA-Pre Adam outperforms the previous best efficient baselines by substantial margins of 0.81, 2.45, and 1.6 perplexity points for the 1 130M, 350M, and 1B models, respectively. Furthermore, when integrated with the Muon (Jordan al., 2024) optimizer, LoRA-Pre Muon achieves addi- tional improvements on both 60M and 130M scale models, demonstrating our method's 's ability to generalize across different optimizers.
4.2 MEMORY-EITICIENT FINE-TURE TUNING
In this section, we evaluate the fine-tuning performance of LoRA-Pre on mathematical tasks. We fine-tune Llama-2-78 and Llama-3.1 3.BBBB models on the MetaMath 100k dataset and assess their per- formance on GSM8K (Cobbe et al., 2021) and MATH-500 (Lightman et al., 2024). To ensure fair comparison, we maintain identical typeer-aaaamtter and training configurations across all methods.
We select several memory-efficient fine-tuning baselines for comparison, including: 1) LoRA (Hu et al., 2022); the standard low-rank fine-tuning method. 2) rsLoRA (Kalajdzievski, 2023): an im- proved LoRA variant that optimizes the scaling factor through rank-stabilized normalization. 3) DoRA (Liu et al., 2024): LoRA extension that decomposes weight updates into magnitude and directional components for more effective optimization. 4) GaLore (Zhao et al., 2024): memory- efficient optimizer that projects gradients into low-rank subspaces using SVD decomposition. To demonstrate cross-optimizer compatibility, evaluate Muon-hased versions, including: 1) Ga-

---

Published as a conference paper at ICLR 2026
Table 2: Results of memory-efficien fine-tuning methods. We compare our method with effi- cient fine-tuning methods including LoRA (Hu et al., 2022), rsLoRA (Kalajdzievski, 2023), and DoRA (Liu et al., 2024), and an efficient optimizer GaLore (Zhao al., 2024). The models are fine-tuned on the MetaMath (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.
Llama-3.1-8B
Llama-2-7D
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
L (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
GaLore Muon (Zhao et al., 2024)
63.41
18.06
40.74
33.11
4.37
18.74
LoRA Muon (Hu et al., 2022)
70.30
19.25
44.78
35.15
6.15
20.65
LoRA-Pre Muon
72.65
20.83
46.74
47.20
6.15
26.68
Lore Muon: which apply the GaLore (Zhao al., 2024) algorithm the Muon optimizer, and 2) LoRA Muon: which optimizes LoRA with the Muon optimizer.
The results are presented in Table LoRA-Pre consistently achieves the highest scores across all experimental configurations, demonstrating superior performance regardless of the base model or optimizer used. The improvements are particularly notable across different settings: when training Llama-3. 3--8B with Adam, LoRA-Pre shows an average improvement of 2.59 points over the second- best method, while with Llama-2-7B and Adam, this improvement increases to 4.88 points. When using the Muon optimizer, LoRA-Pre maintains its advantage with improvements of 1.96 and 6.03 points for the respective models. These results confirm LoRA-Pre's 's effectiveness across diverse experimental conditions and its robust compatibility with different optimizers.
4.3
ABLATION STUDY
Perplexity vs Rank (60M)
Perplexity vs Rank (130M)
Galore
Galore
70
Adam
Pre Much
LoftA-Pre Muon
Adam
Adam
Perplexity
60
Perplexity
50
40
30
4
64
128
64
128
256
Rank
Rank
(a)
(b)
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.
Ablation Different Rank. To systematically evaluate how rank selection affects the performance of LoRA-Pre compared to other efficient optimization methods, conduct comprehensive experi- ments across different rank configurations. We evaluate LoRA-Pre (both Adam and Muon variants) against GaLore (Zhao et al., 2024) on 60M and 130M parameter models. We test ranks of {4, 16, 64, 128} for the 60M model and {16, 64, 128, 256} for the 130M model to observe performance trends across different memory budgets.
8

---

Published as a conference paper at ICLR 2026
Figure shows that all methods improve with increasing rank, but exhibit different rank efficiency. LoRA-Pre consistently achieves better perplexity at lower ranks compared to GaLore. First, all methods show improved performance with increasing rank, but they differ significantly in their rank efficiency. When comparing specific configurations, the efficiency differences become clear. On the 60M model, A-Re Adam rank=16 achieves comparable performance Galoree rank= representing an 8x reduction in rank requirement. Similarly, on the 130M model, LoRA- Pre Adam rank=16 matches GaLore's 's performance rank=256, representing 16x efficiency improvement. LoRA-Pre Muon shows higher rank tolerance than LoRA-Pre Adam. We attribute LoRA-Pre's rank efficiency its continuous subspace adaptation mechanism. GaLore performs periodic subspace updates, creating intervals where the subspace becomes misaligned with the gradient structure. To compensate for this error accumulation, GaL requires larger subspaces. contrast, LoRA-Pre adjusts its subspace at each step, maintaining better alignment and thus achieving effective optimiza- tion with smaller subspaces.
To gain deeper insights into this rank efficiency, we examine the training dynamics of LoRA-Pre Muon across different rank configurations. Figure 3 visualizes the perplexity trajectories for the 130M model with ranks of 256, 128, 64, and 16.
The results reveal an intriguing convergence pat- tern: while smaller ranks initially exhibit higher perplexity values, this performance gap dimin- ishes rapidly as training progresses. This behav- ior demonstrates that LoRA-Pre's dynamic sub- space update mechanism can efficiently capture the evolving momentum structure during training, even when operating with constrained ranks. This rapid adaptation capability explains why LoRA- Pre maintains competitive performance across a wide range of rank settings, making it both bust to rank selection and practically appealing for memory-constrained training scenarios.
Evaluation Perplexity Iterations 130M model
rank-16 16
mank-66
ma
(56
Perplexity
8
30k
Iteration
Figure 3: Test perplexity for LoRA-Pre Muon with different ranks during training.
Table Results of pre-fraining using different efficient Muon optimizers.
Model Size
60M
130M
350M
Muon (Jordan et al., 2024) Muon w/o momentum
28.43
21.86
16.17
32.15
24.23
17.33
Gal ore 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
Ablation of Low-Rank Muon Optimizers. In this section, we evaluate the effectiveness of current efficient optimizers by extending them to the recently proposed Muon optimizer (Jordan et al., 2024). Since existing efficient optimizers were originally designed for Adam (Kingma & Ba, 2015), their compatibility and performance with other optimizers remain unexplored. We conduct experiments on 60M, 130M, and 350M parameter models, comparing LoRA-Pre against GaLore (Zhao el al., 2024) and Fira (Chen et al., 2024) by adapting their implementations to use Muon. Standard Muon serves as the upper bound, while Muon without momentum provides the lower bound. The Muon- based algorithm for LoRA-Pre is presented in Algorithm 2.
The results in Table reveal two significant findings. First, LoRA-Pre Muon consistently outper- forms all other efficient optimizers, achieving improvements of 3.54, 1.80, and 0.43 points over the second-best method at 60M, 130M, and 350M parameters, respectively. Second, projection-based methods surprisingly perform worse than basic Muon without momentum, despite incorporating mo- mentum computation. This counterintuitive result exposes fundamental generalization limitations of projection- based gradient descent methods when applied to different optimizers. We attribute this phenomenon to the periodic subspace updates in projection based methods, which introduce mo- mentum computation errors that subsequently affect Muon's orthogonal update calculations. In

---

Published as a conference paper at ICLR 2026
contrast, LoRA-Pre continuously updates its subspace, enabling better capture of the orthogonal space during Muon's update process and achieving superior performance.
5 CONCLUSION
In this paper, we present LoRA-Pre, novel low-rank efficient optimizer. We establish that EMA momentum updates are mathematically equivalent to training online linear regressor with gra- dient descent on the online gradient flow. Building on this insight, we propose compressing the momentum component through low-rank factorization, deriving update rules that maintain the EMA form while operating in a compressed parameter space. We provide two variants: LoRA-Pre Adam and LoRA-Pre Muon. Extensive experiments on pre-training and fine-tuning tasks demonstrate that LoRA-Pre achieves competitive or superior performance across all evaluated tasks and model sizes. Notably, our method exhibits excellent rank robustness, requiring only 1/8 or fewer ranks compared to previous methods while achieving comparable results. The approach generalizes effectively to various optimizers, making it versatile solution for memory- efficceet optimization.
ACKNOWLEDGEMENT
This work funded by the National Natural Science Foundation of China under Grants (62276256, U2441251, 62550062, 62425606, 32341009) and the Young Elite Scientists Sponsorship Program by CAST (2023QNRC001), We gratefully acknowledge Jie Cheng and Zhen Yang for providing computational resources and for their valuable discussions, which were critical this work.
REFERENCES
Armen Aghajanyan, Sonal Gupta, and Luke Zettlemoer. Intrinsic dimensionality explains the ef- fectiveness of language model fine-tuning, In ACL-IJCNLP, 2021.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In NeurlPS, 2020.
Xi Chen, Kaituo Feng, Changsheng Li, Xunhao Lai, Xiangyu Yue. Ye Yuan, and Guoren Wang. Fira: Can we achieve full-rank training of 1lms under low-rank constraint? arXiv preprint arXiv:2410.01623, 2024.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reicchrro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10. 14168, 2021.
Gheorghe Comanici, Eric Bieber, Mike Schaekermann, Ice Pasupat, Noveen Sachdeva, Inderjit Dhillon, Marcel Blistein, Ori Ram, Dan Zhang, Evan Rosen, al. Gemini 2.5: Pushing the frontier with advanced reasoning, multimodality, long context, and next generation agentic capa- bilities. arXiv preprint urXiv:2507.06261 2025.
Ning Ding, Yujia Qin, Guang Yang, Fuchao Wei, Zonghan Yang, Yusheng Su, Shengding Hu, Yulin Chen, Chi-Min Chan, Weize Chen, et al. Parameter-efficient fine-tuning of large-scale pre-trained language models. Nature Machine Intelligence, 5(3):220-25 235, 2023.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. JMLR, 12(7), 2011.
Ali Edalati, Marzich Tahaci, Ivan Kobyzev, Vahid Partovi Nia, James J Clark, and Mchdi Reza- gholizadch. Krona: Parameter efficient tuning with kronecker adapter. In NeurIPS Workshop, 2023.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv: 1783, 2024.
10

---

Published as a conference paper at ICLR 2026
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song. Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501. 2025.
Andi Han, Jiaxiang Li, Wei Huang, Mingyi Hong, Akiko Takeda, Pratik Kumar Jawanpuria, and Bamdev Mishra. Sltrain: a sparse plus low rank approach for parameter and memory efficient pretraining. In NeurIPS, 2024.
Yongchang Hao, Yanshuai Cao, and Lili Mou, compressors. In ICML, 2024.
Flora: Low-rank adapters are secretly gradient
Soufiane Hayou, Nikhil Ghosh, and Bin Yu. Lorat: Efficient low rank adaptation of large models. In ICML, 2024.
Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning lecture 6a overview of mini-batch gradient descent, 2012.
Edward J Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. In ICLR, 2022.
Aaron Jaech, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Carney, et al. Openai ol system card. arXiv preprint arXiv:2412.16720, 2024.
Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https: //kellerjordan. github... itoiio//oaae ttoppsstss.... /muon/.
Damjan Kalajdzievski. rank stabilization scaling factor for fine-tuning with lora. arXiv preprint arXiv:2312.03772,. 2023.
Siddhartha Rao Kamalakara, Acyr Locatelli, Bharat Venkitesh, Jimmy Ba, Yarin Gal, and Aidan N Gomez. Exploring low rank training of deep neural networks. arXiv preprint arXiv:2209. 13569, 2022.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
Vladislav Lialin, Sherin Muckatira, Namrata Shivagunde, and Anna Rumshisky. Relora: High-rank training through low-rank updates. In ICLR, 2024.
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. In ICLR, 2024.
Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig, Pre- train, prompt, and predict: A systematic survey of prompting methods in natural language pro- cessing. ACM Computing Surveys, 55(9):1-33) 2023.
Shih- Shih-Yang Liu, Chien-Yi Wang, Hongxu Yin, Pavlo Molchanov, Yu-Chiang Frank Wang, Kwang- Ting Cheng, and Min-Hung Chen. Dora: Weight-decomposed low-rank adaptation. In JCML, 2024.
Xiao Liu, Kaixuan Ji, Yicheng Fu, Weng Tam, Zhengxiao Du, Zhilin Yang, and Tang. P-tuning: Prompt tuning can be comparable to fine-tuning across scales and tasks. In ACL, 2022.
Chao Ma. Wenbo Gong, Meyer Scetbon, and Edward Meeds. Swan: Sgd with normalization and whitening enables stateless Ilm training. In ICML, 2025.
Pouria Mahdavinia and Mehrdad Mahdavi. Low-rank momentum factorization for memory efficient training. TMLR, 2025. URL https: /ttp:r//i////e/eeee/e/// net/forum? wtt.rrrrmnnddw333ooo333333
Zhanfeng Mo, Long-Kai Huang, and Sinno Jialin Pan. Parameter and memory efficient pretraining via low-rank riemannian optimization. In ICLR, 2025.
Ionut- Vlad Modoramu, Mher Safaryan, Erik Schultheis, and Dan Alistarh. Svd-free low-rank adap- tive gradient optimization for large language models. arXiv preprint arXiv:2505. 17967, 2025.
11

---

Published as a conference paper at ICLR 2026
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. JMLR, 21(140):1-61 2020.
Noam Shazeer and Mitchell Stern. Adaptive learning rates with sublinear memory cost. In ICML, 2018.
Wei Shen, Yaxiang Zhang, Minhui Huang, Mengfan Xu, Jiawei Zhang, and Cong Shen. Mlore: Momentum low-rank compression for large language model adaptation. arXiv preprint arXiv:2506.01897, 2025.
Nurbek Tastan, Stefanos Laskaridis, Martin Takáč, Karthik Nandakumar, and Samuel Horváth. LoFT: Low-rank adaptation that behaves like full fine-tuning, In ICML Workshop, 2025. URL
https: httpp/ppppeeeeeeeeeeeeeeeeeeeeep net www.nneee//toineeeeeeeeeent forum?id-AiggssDD:...
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Niko- lay Bashlykow, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint vrXiv:2307.09288. 2023.
Shaowen Wang. Linxi Yu, and Jian Li. Lora-ga: Low-rank adaptation with gradient approximation. In NeurIPS, 2024a.
Zhengbo Wang, Jian Liang, Ran He, Nan Xu, Zilei Wang, and Tieniu Tan. Improving zero-shot generalization for clip with synthesized prompts. In ICCV, 2023.
Zhengbo Wang, Jian Liang, Ran He, Zilei Wang, and Tieniu Tan. Connecting the dots: Collaborative fine-tuning for black-box vision-language models. In ICML, 2024b.
Zhengbo Wang. Jian Liang. Lijun Sheng, Ran He, Zilei Wang, and Tieniu Tan. A hard-to-beat baseline for training-free clip-based adaptation. In ICLR, 2024c.
Zhengho Wang, Jian Liang. Ran He, Zilei Wang, and Tieniu Tan. Lora-pro: Are low-rank adapters properly optimized? In ICLR, 2025.
An Yang, Anfeng Li, Baosong Yang. Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao. Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint 0085082505.09388 2025.
Jui-Nan Yen, Si Si, Zhao Meng, Felix Yu, Sai Surya Duvvuri, Inderjit S Dhillon, Cho-Jui Hsieh, and Sanjiv Kumar. Lora done rite: Robust invariant transformation equilibration for lora optimization. In ICLR, 2025.
Longhui Yu, Weisen Jiang, Han Shi, YU Jincheng, Zhengying Liu, Yu Zhang, James Kwok, Zhen- guo Li, Adrian Weller, and Weiyang Liu. Metamath: Bootstrap your own mathematical questions for large language models. In ICLR, 2024.
Qingru Zhang, Minshuo Chen, Alexander Bukharin, Pengeheng He, Yu Cheng, Weizhu Chen, and Tuo Zhao. Adaptive budget allocation for parameter-efficient fine-tuming. In ICLR, 2023.
Yimu Zhang, Yuanshi Liu, and Cong Fang. Adapm: partial momentum algorithm for Ilm training. arXiv preprint arXiv:2510.0910: 2025.
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong
Tian. Galore: Memory-efficient Ilm training by gradient low-rank projection. In ICMI,, 2024.
Philip Zmushko, Aleksandr Beznosikov, Martin Takáč, and Samuel Horváth. Frugal: Memory- efficient optimization by reducing state overhead for scalable training. In ICML, 2025.
12

---

Published as a conference paper at ICLR 2026
Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation
Appendix
The structure of the Appendix is follows.
Appendix A contains the proofs of the theorems in the main manuscript.
Appendix B details the optimization algorithms of the proposed method.
Appendix C provides theoretical analysis of approximation error and convergence of the proposed method.
Appendix provides additional experiments of our method.
Appendix details the LLM usage in this paper.
A PROOF OF THEORETICAL RESULTS
Theorem. Assume matrices mg E Rpx and mA E R are both full rank. For the objective minm p.mm m. L[mB, ma;g) | gllp, Newton's method yields the following closed-form update rules:
mg - - (1-71) mB + 711 gmn Am )-1 mA + (1 -1) mA + m (mpmn)) " m.... Here, "/ is the learning rate for the factorized optimization problem.
(18)
(19)
Proof. We aim to derive Newton's 's method update rules for the optimization problem mm L(m mm:g) 2/2 || gll Our approach begins with computing the first- order gradients, then proceeds to the Hessian computation, and finally establishes the connection to exponential moving average (EMA) updates. To start, we compute the first-order partial derivatives:
aL
(mpma -9)m
(20)
dmg
aL mmmmmAA- - g))-
(21)
0m
While standard gradient descent would directly use these gradients to update the parameters, we instead pursue Newton's method because it yields a more elegant form that naturally resembles EMA updates. For Newton's method, we need the second-order derivatives (Hessian matrices). Computing these second-order partial derivatives gives
g a [(mgmA- - m
HBB =
= mAm' & Ip:
(22)
3m
3MB
0
[m
(m - 9) I, & m
HA^
(23)
0m3
0ma
Using these Hessian matrices, we can now compute the Newton directions by solving the linear systems H dd VI, which yields:
aL
vec(dma) = H vec
dmn mm -gm't (mam²))
(24)
0mg
vec(dma)) = H vec
dma m. (mbmmB) 11 -m
(25)
Öma
13

---

Published as a conference paper at ICLR 2026
The key insight emerges when apply these Newton directions with learning rate "/1. Substituting the Newton directions into the update formula +- x 7/1 dgs we obtain:
mB + mB dmg
mB 71 [m 9m (m) "1
(1 7) mn + 71 gm (mam²²)
(26)
MA + MA - dms
Ma [MA mm)) -
(1 7) m. + 71 (mbma))''n
(27)
These final expressions reveal the remarkable property that Newton's method naturally produces up- date rules in the form of exponential moving averages, where each new parameter value is weighted combination of the previous value and a target value derived from the optimization objective.
To further illustrate this connection, we note that in the uncompressed case where we optimize minm (mm; g) = I - glli. Newton's method similarly yields the classic EMA update:
m H QL Ōm
7) 7.
(28)
This consistency across problem formulations demonstrates the fundamental nature of this EMA- like structure in Newton's method and justifies our preference for this approach over standard gradi- descent.
B DETAILED ALGORIHTS OF LORA- FOR ADAM AND MUON OPTIMIZERS
This section presents the LoRA-Pre algorithms for both Adam (Kingma & Ba, 2015) and Muon (Jor- dan et al., 2024) optimizers.
B.1 ALGORITHM OF LORA-PRE FOR ADAM
The Adam optimizer update rules under LoRA-Pre have been established in Section
First-order momentum updates: For the first-order momentum term with parameterization MBMA, the update rules are:
m'a +- (1 (1)?)) mm + 71- gm (m)
(29)
m'a + (1 -1) mA + n1 (m
(30)
By default, set 1 - 71 B1, which ensures that after the update, m' m'gm'a mgmA + ..., making the EMA coefficient consistent with standard Adam.
B1
Second-order momentum updates: For the second-order momentum term with parameterization v (2 the update rules are:
8'B t (1 (1) UB + 72 Iglu
(31)
V - (1 -1-72) 72) UA 7 (W
(32)
Analogously, set 1 - 7/ p by default, which ensures that o' (1/g0%))2 (2
83
Complete algorithm: Based these update formulas, Algorithm presents the complete LoRA- Pre implementation for the Adam optimizer, demonstrating how these factorized momentum updates integrate seamlessly into the standard Adam framework.
14

---

Published as a conference paper at ICLR 2026
Algorithm Comparison of Adam and
Adam with LoRA-Pre
Require:
Initial learning rate 7, weight decay A, ,1, 82 [0, 71, 72
[0,
0
1:
Initialize parameters fg. time step t +- 0,
first moment mo + 0, second moment vo +- 0,
first low-rank moment m,, < 0, m.Ao0 ~ NO 0.02),
second low-rank moment v8,0 +- 0, v4.0 (.0,0.0.).
2:
repeat
3:
++t+ ++1
4:
+ (0-1(0)-1)
5:
# Update first moment
6:
m. + 31 m + (1 (1)
7:
m. -1, 1 M, (1 31) gt
8:
mB, (1 m) mB,r1 7/1 9 - 1
9:
mA, (1 71) m.Ac-1 71 mB,.1- -mm
# Update second moment
11:
+ B2 1 + (1 - 32) g
02 82 (vB.t-1UA,t-)) + (1 - 32) g2²
UB + (1 72) UB,--1 + 7/2 | u --1 (2A,
14:
AAA < (1 72) WAtt-1 7/- (Vb
1 1g:
15:
me + mz/( (1 3
δ 01/(1= (2)
17:
02 0-1 Y
+ 1 0 vocte
dr
18: until stopping criterion is met
19: return Optimized parameters 8
15

---

1 gm Published as a conference paper at ICLR 2026
B.2 ALGORITHM OF LORA- FOR MUON
In this section, we present the LoRA-Pre algorithm for the Muon (Jordan et al., 2024) optimizer.
First-order momentum updates: For the Muon optimizer, derive the LoRA-Pre algorithm by first reformulating the momentum update. The Muon momentum term can be equivalently written as:
m' μ m +
m (1 - [1) +
m (1 1-H1 [1) (m - 9) 1.
1 m (1 μi) (m 9)
(33)
1 p
By treating the Muon update as the solution to an optimization problem, we can derive the equivalent objective function:
L(m; g)
1
||m -gll
# (m, g)
(34)
2
μ After applying low-rank factorization m mB mA. the objective becomes:
1 L(mB, mAig) = || Elle
(mBMA.9)).
(35)
1-
We aim to derive Newton's method update rules for this modified objective. Computing the first-
order gradients:
ØL (mgma - g)m't
(36)
dmg
1-
OL m 9)
1 mBg-
(37)
1- r
The Hessian matrices have the same structure as before since the additional linear term doesn't 't affect the second derivatives:
a [mMmA- - g)
HBB 3mg B
-mAmAA & Ip:
(38)
2
ŌmB
8 L
2'' (m - g)l
HA^
I m
(39)
0m
Oma
Using these Hessian matrices, we can now compute the Newton directions by solving the linear systems - d= VI, which yields:
0L
vec(dmmn) H BB vec
dmn
mg gma" 2 (mAmm))
(40)
3mg
r 1
aL
vec(dm,)) H vec
dm mA 1 P
(mfmp) - mm
(41)
3m
The key insight emerges when we apply these Newton directions with learning rate Vı. Substituting the Newton directions into the update formula x + x dz. we obtain:
mg + mg - 71 dmn
(mamm))
= -71
mB
1 4
= (1 - - mm +
(m)(
gm
(42)
1
MA +- MA - ?1 dms
1 (mfmm)
-
= mA - - ma
1 1
= (1 - (1-7) mA
?1
(mFmB) -1
(43)
16

---

Published as a conference paper at ICLR 2026
Similarly, set 71 Vp.
Complete algorithm: Based on these update formulas, Algorithm presents the complete LoRA- Pre implementation for the Muon optimizer, demonstrating how these factorized momentum updates integrate seamlessly into the standard Muon framework.
Algorithm Comparison of Muon and Muon with LoRA-Pre
Require: Initial learning rate 7, weight decay X, momentum μ E (0, E [0, Initialize parameters oy. time step 4-0.
first moment mi f-
first low-rank moment mB, < O, mA,0 ~ (O,
2:
repeat
3:
< +++
4:
< (f,1()
5:
# Update first moment
6:
M1 n. 14-11 9t
7:
my + mBB.-177A,.11 +
8:
MB, + - (1-7) mB,L-1 + I-p
71
1
9:
MA, (1 (1- 71) mA.--1 - 1 m 'm 1. 19c
10:
O2 - New-oonSShhhul)(mm
0. 0111 - ( -
12: until stopping criterion is met return Optimized parameters 0,
C THEORETICAL ANALYSIS OF APPROXIMATION ERROR AND CONVERGENCE
In this appendix, provide rigorous theoretical analysis of the LoRA-Pre Adam optimizer. We explicitly analyze the approximation error introduced by the low-rank factorization of the optimizer states, and prove the convergence fidelity of the algorithm in non-convex settings.
C.I PROBLEM SETUP AND ALGORITHM DYNAMICS
Consider the unconstrained optimization problem min f()). Let gt (f(0,-1) be the stochastic gradient at step f. We denote the states of Standard Adam as my, and the effective states of LoRA-Pre Adam as m,, 0t.
Standard Adam Dynamics The standard optimizer updates its moments using exponential moving averages (EMA) with decay rates B1, .3 E [0, 1):
m: - 11-1 + (1) -9.. B2 01-1 + (1 32) g.
(44)
(45)
2. LoRA-Pre Dynamics LoRA-Pre maintains low-rank factors (m m.A.. approximate the gradient history. Let 71 be the update rate. The exact simultaneous update rules (Online Least Squares) be compactly expressed using the Moore-Penrose pseudo-inverse (-)1:
mn.: (1 m1) MB.-- + 9.m 1,t-1
(46)
M. (1 (1) MA + m BBt119h
(47)
17

---

Published as a conference paper at ICLR 2026
where the Tikhonov regularized pseudo-inverses for the full-rank factors are defined m (mamm A- -)- and m|3 mm + - with damping factor À >
We define the canonical projection operators associated with the factors at step t 1:
PA _ m A,1- A,1-1 - (Projection onto Row Space of m.).
PB mg (Projection onto Column Space of mg).
Let m m. A, be the low-rank history reconstruction. Analogous updates apply the second moment factors using the gradient magnitude [gc], producing the reconstruction h4 = U,,.
3. Effective Moments for Update Crucially, LoRA-Pre computes the effective moments for the parameter update by combining the low-rank history with the exact current gradient:
m 31 m + (1 11--11 31) 9t.
(48)
U2 = 32 12- + 82) 9t
(49)
Note that for the second moment, LoRA-Pre approximates the history of magnitudes h and then squares it.
C.2 ASSUMPTIONS
Assumption (Regularity and Boundedness). The objective function and stochastic gradients sat- isfy the following conditions:
1. L-Smoothness: The objective function f is L-smooth: [V - S L||x y
2. Bounded Gradients: The stochastic gradients are uniformly bounded in both Frobenius and infinity norms. There exist constants G and Gos such that for all t, |[ || S G and
l1g. ol S G
3. Lipschitz Smooth Update: The optimizer uses a smoothed damping term e O. The update
mapping is defined as @(m, v) Vute" This function is Lipschitz continuous with respect to both arguments, with constants Lmm = J/2 and Lp
G
Ze1.5
Assumption 2 (Subspace Approximation Capability). The gradient dynamics admit a low-rank structure. Crucially, we assume this structure holds for both the gradient direction and its element- wise magnitude. Let PB,,, PA,t denote the projections onto the subspaces maintained by the opti- mizer at step t. We assume there exists bound 5 2 0 such that:
l1gt (PB,1 + AF, S 6,
(50)
|||ge| (PB.l99l + 1o Al I &.
(51)
The second inequality ensures that the second-moment estimator (based on (gr)) also admits bounded reconstruction error.
Assumption (Reference Optimizer Descent). Let u, mm/(/ c) be the update direction of the standard full-rank Adam optimizer. We assume that in expectation, ue is valid descent direction aligned with the true gradient:
E ((O.), u.)1 CE(II
(52)
for some constant c O. This assumption anchors the convergence of LoRA-Pre Adam to the theoretical behavior of standard Adam.
C.3 BOUNDEDNESS OF FACTOR RECONSTRUCTION ERROR
We first prove that the error of the stored low-rank history fy uniformly bounded. We strictly enforce the time-scale alignment condition: 81 (1 - =
Lemma C.1. Let C.HH = ||, mellr- Under Assumptions and 2, EYn is uniformly bounded by constant Ebound-
18

---

Published as a conference paper at ICLR 2026
Proof. Step 1: Exact Expansion of LoRA Dynamies Substitute the update rules (46) and (47) into
my mB,.m.
mi = (11-1) (1-2))mB.-1 7 - B
= 1 m-) + (1 - (P + 9.P: + 7tt
(53)
where Qr gem
A, A, 1m p2t- 19t is the quadratic interaction term. Using the condition B1
n)² ,, this implies "/1 = vB1. The expansion becomes:
m m,m+1 (1 B1) (P + 9:Pa) (1 Q
(54)
Step Constructing the Recursive Error We form the difference with the standard Adam update 81m--1 + 31)9:
my - Th = (i - -m+-1) + Re:
(55)
where the residual driving term Rs is:
R2 (1 Bi)ge (1
- (1 -11-VBi)'
(56)
Step 3: Bounding the Residual Magnitude
We explicitly bound the norm of the driving residual R₃- Applying the triangle inequality to the residual definition (56), we decompose the total bound into linear and quadratic contributions:
RellF S
B1)
gr + BBl Png. - (.PPa) +
(57)
Quadratic Term
Linear
Using the homogeneity of the norm and the triangle inequality, followed by the substitution the gradient bound |[9e|1 S (Assumption and the subspace residual bound 5 (Assumption 2), we derive:
1. The Linear Term:
||Linear||r-
= = v /1) 99 + VBi PBgt - -9:PN)
$ (1 B (1191lle + B1ll9 PB99 -.PAlle)
$ (1 VBNG + VBr(1 (1-v11))
(58)
2. The Quadratic Term: Bounding ||Q/] F requires the spectral norm of the regularized pseudo- inverses. As established in the setup, m mz ( and m (m- Am Let the Singular Value Decomposition (SVD) of the factor be UEVT. Substituting this into either inverse form reveals that the singular values are governed by l((a) = - Analyzing the function VX, yielding the uniform spectral bound
f(x)
-+ for x 2 0, we find its maximum
|| 112 < 2VX
Using the sub-multiplicativity of the Frobenius norm (specifically ||ABC|| S || we strictly bound Q;
| = [Igem) [ 1"
< 1lgell
|l 1,
1|2
|1 B,
112
119.11
11
1|1
11-m 1212
2
6
< G2
Co.
(59)
2v
4A
Substituting both bounds back into the residual equation:
|| r S
VBiG VBi(I
v Bi)5 (1
C8 $ Ares-
(60)
19

---

Published as a conference paper at ICLR 2026
Step 4: Convergence to Steady State
The error dynamics follow the linear recursion inequality E S 3, + Ares. Since the mo- mentum parameter satisfies 0 S $1 < 1, this geometric series converges finite steady state as t 00:
Ares - Ebound- lim Etth - 31 (-00
(61)
To derive the explicit bound, we substitute the expression for Ares derived in Equation (60) and factor the denominator using 1 - 31 = (B + vB1):
VBi) vBi( vBi + (1 VB1)
Ebound
(1 (1 Bi)
vBi) (B+)[ VB,8 + (1 VBi)col
(62)
VB vB1)
Canceling the common factor (1 vB1) from the numerator and denominator, we obtain the final uniform bound:
G 1 (1-
Ebound
(63)
1 VBI
This confirms that the reconstruction error is uniformly bounded by constant determined by the gradient magnitude G, the subspace residual ò, and the quadratic interaction Cq.
C.4 JOINT EFFECTIVE MOMENT ERROR
We now derive the error bounds for the effective moments me and ün used in the parameter update, explicitly accounting for the non-linear square term in Üy.
Lemma C.2. Let Am I|m, Teell F and A, 1| ūzllF denote the effective moment errors. Let be the total number of parameters. The errors are bounded by:
Am 31Ebounds
(64)
A B2 <2 + oinrat).
(65)
where oottl
4 L2
represents the intrinsic variance bound derived from Popoviciu's 's inequality.
Proof. 1. First Moment Error Analysis
We compare the effective first moment ih with the standard Adam moment my, Subtracting their definitions:
m: m = (m - m4-1) + (1 (1) ge) = (m m
(66)
The current gradient terms cancel out. Taking the Frobenius norm and applying Lemma C.I, we obtain:
A S B1 Eboundd.
(67)
2. Second Moment Error Analysis
We analyze the error in the second moment estimate. Standard Adam updates the second moment as
0t = 1-4411 (1 - "2 LoRA-Pre updates effective second moment as ū = 3y(hh-1)) + (1 -2)99
Subtracting the two equations, the term (1 B2)gi² involving the current gradient cancels out ex-
actly. The error is strictly propagated from the history terms:
1 6 B2 ut-1 (h (h)-1)o2]
(68)
To rigorously bound this, we introduce an auxiliary variable h defined the exact exponential moving average the gradient magnitudes [gl (i.e., hg = 32hx-1 + (1 32)194)). We apply the triangle inequality to decompose the error into two structurally distinct terms:
A, < 32
1 (he-1) + 1 (fc-1)²2113
(69)
Term Intrinsic Variance
Term II: Approximation Error
20

---

Published as a conference paper at ICLR 2026
Term I: Intrinsic Variance (Popoviciu's Bound). This term captures the discrepancy between "the mean of squares" and "the square of the mean". Let us focus on a single parameter index j. The EMA formulation can be interpreted as weighted expectation E,, over the history of gradients, where the weights sum to 1. Thus:
1.5 = Ew B and (Iv-1,)) (B (
(70)
The difference Uz-1j (1 is precisely the variance of the random variable [gr under the EMA probability measure. Since the gradient magnitude is bounded, ie., @riil E [0, Goc]. we invoke Popoviciu's Inequality, which states that the variance of any bounded random variable on interval [a, b] is bounded by (b a)². Here, [a, bj = [0, Goo]. Thus, for every element j:
1 GG 002 G 4
M1-1 (m1-1))
(71)
We generalize this the full parameter matrix by computing the Frobenius norm over all elements:
11-1. (he-1g) S d. (1
( 6
Vd
||:-1 (h4-1) 2
G
(72)
4
We define this constant bound Piotal
vd C2
Term II: Approximation Error This term represents the propagation the low-rank reconstruction error through the squaring operation. Consider the function f(x) x². For inputs restricted to the
domain [0, Gool, the derivative bounded by [x)) = =1ar| I 2G By the Mean Value Theorem, the function is Lipschitz continuous with constant Log 2G Applying this element-wise to the
vectors h and h;-1;
(hi-1g)² (he-1,)²² S 2G h h2-15
(73)
Squaring both sides and summing over all j to recover the Frobenius norm:
d
d (he-15)² (hr (hr-15))²² < 4G2 Σ
Σ
he-15 - Fu-1o
(74)
1=1
Taking the square root yields:
11 02 (fr-1) S 2G 1111-1111111 Pre
(75)
Substituting the bound from Lemma C.1 (11 - -u-illl S Ebound). we have:
Term Il < 2G Ehoundd-
(76)
Combining the bounds for Term and Term II, we obtain the final bound for the second moment сπrоr:
D S B2 (o + Ehoundd)
(77)
C.5 CONVERGENCE ANALYSIS
Theorem C.3. Let the step size be v/Vi. Under Assumptions 1, 2, and 3, LoRA-Pre Adam (with update rule Φ(m, v) = m/ Vv E) converges to a neighborhood of a stationary point:
Cinst VT
E ( ISST
min
+ Cnoose (Ebound + otoaa)
(78)
where Cinit depends on the initial function gap and Cnoise depends on the Lipschitz constants of the update rule.
21

---

Published as a conference paper at ICLR 2026
Proof. 1. Bounding the Update Direction Error
Let di = o(m, 0z) and ü, = (,,, Ut) be the ideal and actual update directions, respectively. We analyze the theoretical update function @(m,v) ve+t (using the smoothed denominator for Lipschitz continuity). The partial derivatives are bounded follows:
Wir m:
V ml - V L.n
Wi
[ e0 I sup
-m m 2i++jji
S
G 2
-
Lag-
Using the moment error bounds from Lemma C.2, the total direction error ||u, [lu--ūell -elle satisfies:
& S Lm Am + Lu
1
G
< (3)Ebound) + 2e (0 + oiotal)
B1 30
GB
Ebourdd +
Etotal
(79)
el.5
1.5
total
R1
R₂
Here, Stotal is uniform deterministic upper bound the direction error.
One-Step Descent Analysis in Expectation
Since the objective function is L-smooth, apply the Descent Lemma conditioned the current state 0;:
Lnk
E, [f s f(0)) (.E. a,)l +
(80)
2
Let Gstep = G/ Ve be the bound on the update norm || We decompose the inner product term:
-E [V ū.)l -.. u))l + E, un uz)l.
(81)
Using Assumption 3 (EE, [ [.. uz)l 2 /|| and applying Cauchy-Schward followed by Young's
Inequality (ab Şa² + b²) to the error term:
E, [ uc)l VV Elottl
1
+
2c Stotel
22
(82)
Substituting this back into the descent inequality, the descent term the error:
-clV is partially offset by
[cl
c 1IVSS
1
Stotal
+
Laf
G2
2
2c
-(0))
CT 2
|
/||
+
2c Stotal +
c2
Lnd steep
(83)
3. Global Convergence
Taking the total expectation and summing from 1 to T:
LG2
ct E[I f(1,) - ff Stotal 2c
step
nt
+
m
(84)
2
Dividing by SS (where ST nt), we isolate the minimum gradient norm:
E|l / f(| 2(f(0,) fr) + Shotal cST
1 12 LG Σnz
(85)
ST
Given π). n//t, have ST > VT and Σ R In Thus, T 0o, the first and third terms vanish (O(1/VT) and O(IT/vT)) leaving the constant error floor. Defining Cnoise e max(K1, K2)² (noting that e max K1, (Ehound + t
we conclude:
Cina VT
miim EinE vf
+ Cnoise Ebound + 2.
(86)
22

---

Published as a conference paper at ICLR 2026
This confirms that the algorithm converges to neighborhood determined by the approximation quality Ebound and intrinsic variance of dooal-
D ADDITIONAL EXPERIMENTAL RESULTS
D.1 HYPERPARAMETER SENSITIVITY ANALYSIS: THE COUPLING OF 8 AND Y
In this section, we evaluate the sensitivity of LoRA-Pre Adam to variations in its momentum hyper- parameters. key design principle of LoRA-Pre that the low-rank update coefficients (% 72) are independent hyperparameters requiring separate tuning. Instead, they are analytically coupled with the standard Adam momentum coefficients (81, 82) to ensure consistent decay dynamics.
Derivation of the Coupling Strategy. Formally, the update rules for the low-rank momentum components m and a are defined as:
m'g (1 71) mB 7i .m m'a - (1 - 7) mA + 71 (m 1 mm...
(87)
(88)
When reconstructing the effective full-rank momentum m ~ m the decay behavior can be approximated by expanding the product of the updated factors. Ignoring higher-order cross terms, the effective decay rate is the product of the individual decay rates:
m' m's (1-m)² mgmA + O6+)) (1 (1 m.
(89)
To align this trajectory with standard Adam optimization (where momentum decays by B1), we enforce the coupling constraint (1 - 7/2 )2 3r. Similarly, for the second moment, since LoRA-Pre approximates the gradient magnitude h == [g| and subsequently squares it (ū = fr²),, the decay rate compounds twice, leading the condition (1 - 72 2 B2. Consequently, the learning dynamics are strictly governed by 8, eliminating the need to grid-search 7,.
Empirical Verification. We conducted sensitivity studies the 60M parameter model by varying 3, and 32 around their standard default values (3) = 0.9, B2 0.95). The results are summarized in Table 4.
Table 4: Sensitivity Analysis of Momentum Parameters 8. We report the validation perplexity on the 60M model. The method exhibits stability around the default settings derived from standard Adam (3: 0.9, 32 0.95). Extreme values (e.g., 1) cause the implicit y to vanish, leading to training instability.
Parameter
Value
Perplexity
Status
0.90 (Default)
32.57
Optimal
B1 (fix B2 0.95)
0.95
37.62
Stable
0.99
1458.92
Diverged
0.90
34.61
Stable
32 (fix 81 = 0.90)
0.95 (Default)
32.57
Optimal
0.999
1301.58
Diverged
As shown in Table 4, LoRA-Pre achieves the best performance at the configuration inherited from standard Adam. While the optimizer is robust within reasonable range, extreme values (e.g., 31 = 0.99 or 32 0.999) lead to numerical instability. This is expected: update rate y approaches zero (e.g., for B₂ = 0.999, 7/2 x 2.5 x 10-4), causing the low-rank factors to adapt too slowly to track the shifting gradient subspace. This confirms that coupling strategy effective, but the "sweet spot" for low-rank adaptation dynamics aligns with, but is more sensitive than, full-rank optimization.
3 + 1, the derived
23

---

Published as a conference paper at ICLR 2026
E STATEMENT OF THE USE OF LARGE LANGUAGE MODELS
The use of LLMs in this work was restricted to paper writing assistance. They were not used to generate results, derive proofs, or conduct analysis without human verification. The disclosure here, as well as in the submission form, fulfills the ICLR requirement that all contributions of LLMs be acknowledged transparently.
24


Reviews:
[{"review_id": "VP2940RljQ", "reviewer": "Reviewer_QWoo", "summary": "This paper introduces LoRA-Pre, a new memory-efficient optimizer designed to reduce the memory footprint of optimizer states during large language model (LLM) pre-training and fine-tuning. The central insight reframes the exponential moving average (EMA) momentum update as an instance of online linear regression, then leverages this equivalence to propose a low-rank factorization of the optimizer's momentum matrices. Detailed update rules are derived for Adam and Muon optimizers, with closed-form solutions shown via Newton's method. Extensive experiments on Llama-family models (from 60M to 1B parameters) demonstrate that LoRA-Pre achieves competitive or superior performance with much smaller rank than baseline low-rank methods, both for pre-training and memory-efficient fine-tuning.", "questions": "1) How does the method handle ill-conditioning when computing matrix inverses for the low-rank updates, especially as $r$ increases? Is there any regularization, and how is numerical stability ensured in practice?\n\n2) Under the settings of Table 1, which has the larger peak memory: LoRA-Pre (Adam) or plain Muon?\n\n3) In Table 1, classical Muon consistently outperforms LoRA-Pre(Muon) (and all other methods), while LoRA-Pre(Adam) always outperforms Adam. Also as training tokens increase, LoRA-Pre(Adam) overtakes LoRA-Pre(Muon), whereas among “full” methods Muon remains superior to Adam. What interactions between Muon’s orthogonalized updates and factorized momentum explain this gap, and can they be mitigated?\n\n4) Your Theorem 3.1 derives closed-form factor updates via an online Newton step on $L(m_B, m_A, g)$. How does this update relate to the closed-form updates used in LoRA-Pro [3] and LoFT [4], which did not explicitly employ Newton’s method?\n\n[1] Shen W. et al. MLorc: Momentum Low-rank Compression for Large Language Model Adaptation //arXiv preprint arXiv:2506.01897. – 2025.\n\n[2] Shivagunde N. et al. Approximations may be all you need: Towards pre-training LLMs with low-rank decomposition and optimizers. – 2024.\n\n[3] Wang Z. et al. Lora-pro: Are low-rank adapters properly optimized? //arXiv preprint arXiv:2407.18242. – 2024.\n\n[4] Tastan N. et al. LoFT: Low-Rank Adaptation That Behaves Like Full Fine-Tuning //arXiv preprint arXiv:2505.21289. – 2025.\n\n[5] Bogachev V. et al. LoRA meets Riemannion: Muon Optimizer for Parametrization-independent Low-Rank Adapters //arXiv preprint arXiv:2507.12142. – 2025.\n\n[6] Zhang Y., Liu Y., Fang C. AdaPM: a Partial Momentum Algorithm for LLM Training //arXiv preprint arXiv:2510.09103. – 2025.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "__1) Theoretical Insight:__ The paper establishes a non-trivial mathematical equivalence between EMA momentum updates and online linear regression (Section 3.2), which is both elegant and enables a new way to think about optimizer state compression beyond ad-hoc engineering.\n\n__2) Methodological Contribution:__ Closed-form Newton-based update rules are derived for the low-rank factors, with careful exposition (Section 3.3, Theorem 3.1, Appendix A). This enables stable and efficient momentum updates without full-rank storage.\n\n__3) Empirical Validation:__ Experiments are thorough, covering both pre-training (Section 4.1, Table 1) and fine-tuning (Section 4.2, Table 2) across a range of Llama models and tasks (GSM8k, MATH500). \n\n__4) Competitive baselines:__ Adam, Muon, LoRA, Galore, and others are compared, and LoRA-Pre outperforms or matches the best across most benchmarks.", "weaknesses": "__1) Incomplete related work coverage.__ The literature review appears incomplete. The paper omits several relevant lines of work where optimizer momenta or gradient statistics are decomposed or constrained in low-rank form. For example, MLorc [1] and MoFaSGD [2] address momentum/gradient factorization. Within PEFT, papers like LoRA-Pro [3] and LoFT [4] are highly relevant (the latter yielding updates that resemble those proposed here). Another closely related direction is Riemannian low-rank pretraining/optimization [5], which explicitly performs muon-like steps on a low-rank manifold. While the paper does cite LoRA-Right and LORO, this seems insufficient to position the contribution within the broader space of low-rank momentum and manifold-constrained optimization. In addition, AdaPM [6] also decomposes momentum into a low-rank representation, conceptually aligning with the spirit of the present work. Although the specific algorithmic designs differ, ADAPM appears to be a highly relevant work. I recognize that ADAPM was posted after the ICLR submission deadline, but it should be cited and discussed in the camera-ready to properly contextualize the contribution.\n\n__2) Memory accounting is unclear and potentially unfavorable in some regimes.__ LoRA-Pre reduces the persistent state, but each step still materializes full momentum via $m = m_B m_A$ and momentum with bias correction to apply the update. This allocation increases peak memory, unlike methods that avoid full momentum materialization (e.g., standard LoRA or GaLore’s projected states). For small models (for example with 1 layer), LoRA-Pre can yield no benefit and may increase both memory and runtime. A complete peak-memory breakdown (parameters, activations, gradients, optimizer states, and transient products) is missing.\n\n__3) Computational overhead and numerical stability are underexplored.__ Per-step inversions of $r \\times r$ matrices add nontrivial overhead and can be ill-conditioned as $r$ grows. The paper does not specify stabilization strategies nor quantify runtime overhead relative to Adam/Muon/GaLore. Empirical stability analysis is absent.\n\n__4) Fairness of comparisons and the $r/d{model}$ control variable.__ Using $r/d_{model}$ as the primary control for fairness is inadequate: peak memory and step time depend on whether full momentum is materialized, number of layers, matrix operations performed at each step, etc. Comparisons should be matched on peak memory and/or throughput wall-clock time.\n\n__5) Typo:__ in Algorithm 1 line 12 it should be $(1 - \\beta_2)$, not $(1 - \\beta_1)$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jM293GTVYU", "reviewer": "Reviewer_dPBR", "summary": "This paper introduces LoRA-Pre, a low-rank optimizer that reinterprets the exponential moving average (EMA) in momentum computation as an online linear regression problem. By exploiting this equivalence, the authors compress the optimizer’s momentum matrices via low-rank factorization (m_B,m_A) and derive closed-form Newton-style EMA updates without requiring back-propagation through the optimizer state. This leads to  LoRA-Pre achieves similar or better performance using only 1/8 of the rank, and it generalizes effectively to other optimizers such as Muon.", "questions": "see weaknesses", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tThe relation between EMA and online linear regression is novel and conceptually clean, providing a unified foundation for memory-efficient optimizer design.\n2.\tThe method applies seamlessly to multiple modern optimizers (Adam, Muon), offering a general solution for optimizer-state compression in large-scale model training.\n3.\tLoRA-Pre consistently outperforms both projection-based and fine-tuning low-rank optimizers across different scales and tasks. Rank-efficiency studies reinforce the robustness under constrained memory.", "weaknesses": "1.\tMissing Key Baseline in main table: Although Fira [1] is mentioned in the related work and ablation (Table 3), it is missing from the main results table (Table 1). This omission is non-trivial, since Fira operates under **exactly the same setting** for pre-training and it also deals with the optimizer state itself under a low-rank constraint, making it a direct and essential baseline.\nThe paper should include a comparison with Fira-Adam in the main results.\nEven if LoRA-Pre does not necessarily outperform Fira, the paper should explicitly analyze their differences and trade-offs.\n\n2.\tThis paper does not analyze approximation error or convergence fidelity under low-rank constraints.\n\n3.\tThe ablation focuses solely on rank. Adding sensitivity analysis to other hyperparameters (like $\\beta$ and $\\gamma$) would offer a more comprehensive understanding.\n\n[1] Fira: Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "V5Ngd6X2ca", "reviewer": "Reviewer_Mkay", "summary": "The paper proposes LoRA-Pre, a new optimizer framework that reinterprets momentum as an online linear regression process and replaces traditional exponential-moving-average (EMA) updates with a low-rank closed-form solution. By decomposing momentum into low-rank factors and updating them analytically, LoRA-Pre significantly reduces optimizer-state memory while preserving training stability. The method is applied to both Adam and Muon, yielding LoRA-Pre-Adam and LoRA-Pre-Muon variants. Experiments on large-scale pretraining (LLaMA-1B and smaller models) and instruction fine-tuning (LLaMA-8B, LLaMA-7B) show that LoRA-Pre matches or surpasses strong baselines such as Galore, DoRA, and ReLoRA, while offering notable memory savings and rank efficiency.", "questions": "Can the authors clarify how the proposed low-rank updates scale computationally in distributed or large-batch training? An analysis of communication and synchronization cost would help assess real-world deployment.\n\nHow sensitive is LoRA-Pre to the chosen rank or initialization of the low-rank factors? A more detailed study could clarify whether the observed gains primarily come from rank efficiency or improved optimization dynamics.\n\nHave the authors considered extending comparisons beyond projection-based methods to include modern adaptive or second-order optimizers, to better quantify the relative benefits of the proposed approach?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper introduces a clear and technically sound rethinking of momentum optimization by framing EMA updates as an online regression problem and deriving a low-rank closed-form solution. This perspective is both conceptually original and practically valuable, bridging optimization theory with efficient model training. The proposed LoRA-Pre method is well-integrated with existing optimizers like Adam and Muon, achieving strong empirical results with lower memory costs. The experiments are comprehensive, covering both large-scale pretraining and fine-tuning, and the presentation is mathematically rigorous and well-organized.", "weaknesses": "The empirical evaluation primarily compares with projection-based low-rank optimizers; incorporating a wider range of modern baselines (e.g., Sophia, Shampoo) would provide a more complete assessment of practical advantages.\n\nThe ablation studies emphasize rank efficiency but do not clearly disentangle the effect of the low-rank momentum representation from other implementation factors such as learning-rate scaling or normalization.\n\nThe scalability and computational trade-offs of the proposed updates, particularly under distributed or large-batch training, are not fully analyzed, leaving some uncertainty about deployment efficiency in large-scale settings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)