Low-rank Momentum Factorization for Memory Efficient Training
Authors: Pouria Mahdavinia
Mehrdad Mahdavi
Venue: Accepted by TMLR ()

Abstract: Fine-tuning large foundation models presents significant memory challenges due to stateful optimizers like AdamW, often requiring several times more GPU memory than inference. While memory-efficient methods like parameter-efficient fine-tuning (e.g., LoRA) and optimizer state compression exist, recent approaches like GaLore bridge these by using low-rank gradient projections and subspace moment accumulation. However, such methods may struggle with fixed subspaces or computationally costly offline resampling (e.g., requiring full-matrix SVDs). We propose Momentum Factorized SGD (MoFaSGD), which maintains a dynamically updated low-rank SVD representation of the first-order momentum, closely approximating its full-rank counterpart throughout training. This factorization enables a memory-efficient fine-tuning method that adaptively updates the optimization subspace at each iteration. Crucially, MoFaSGD leverages the computed low-rank momentum factors to perform efficient spectrally normalized updates, offering an alternative to subspace moment accumulation. We establish theoretical convergence guarantees for MoFaSGD, proving it achieves an optimal rate for non-convex stochastic optimization under standard assumptions. Empirically, we demonstrate MoFaSGD's effectiveness on large language model alignment benchmarks, achieving a competitive trade-off between memory reduction (comparable to LoRA) and performance compared to state-of-the-art low-rank optimization methods. Our implementation is available at \url{https://github.com/pmahdavi/MoFaSGD}.

Published in Transactions on Machine Learning Research (07/2025)
Low-rank Momentum Factorization for Memory Efficient Training
Pouria Mahdavinia
pam54260psu.edu
Deportment of Computer Science and Engineering The Pennsyloania State University
Mchrdad Mahdavi
mam616@ppuu edu
Deportment Computer Science and Engineering The Pennsylvania State University
Reviewed on OpenReview: https: // openreview net/ pet//foumm? id= K3D3TVo9a3
Abstract
Fino-tuning large foundation modols presents significant memory challenges due to stateful optimizers like AdamW, often requiring several times more GPU memory than inference. While memory-efficient methods like parameter-effiiiet fine-tuning (e.g., LoRA) and opti- mizer state compression exist, recent approaches like GaLore bridge these by using low-rank gradient projections and subspace moment accumulation. However, such methods may struggle with fixed subspaces or computationally costly offline resampling (e.g,, requiring full-matrix SVDs). We proposc Momentum Factorized SGD (MoFaSGD), which maintains a dynamically updated low-rank SVD represenlation of the first-order momentum, closely approximating its full-rank counterpart throughout training. This factorization enables a memory-efficient fine-tuning method that adaptively updates the optimization subspace at each iteration. Crucially, MoPaSGD leverages the computed low-rank momentum factors to perform efficient spectrally normalized updates, offering alternative to subspace moment accumulation. We establish theoretical convergence guarantees for MoFaSGD, proving it achieves an optimal rate for non-convex stochastic optimization under standard assumptions. Empirically, we demonstrate MoFaSGD's effectiveness on large language model alignment benchmarks, achieving a competitive trade-off between memory reduction (comparable to LoRA) and performance compared to state-of-the-art low-rank optimization methods. Our implementation available
1 Introduction
Advancements in AT have been propelled by initial scaling laws, which boost pre-training performance via ever-larger models and datasets (Kaplan et al. 2020). However, adapting these large foundation models to downstream tasks, such as instruction or preference tuning (Ouyang et al. 2022). often requires several times more GPU memory than inference. This memory burden largely stems from storing optimizer states (c.g.; momentum terms) required by ubiquitous methods like AdamW. Various strategics have emerged alleviate these costs.
Parameter-Efficient Fine- Tine-Tuning (PEFT) methods provide an architectural solution, introducing small set of new, trainable parameters called adapters. For example, the popular LoRA technique (Hu et al. 2021b restricts model updates to these low-rank adaplers, drastically reducing memory overhead by training only a small fraction parameters. Variants such as DoRA [Lu et al. 2024), AdaLoRA (Zang et al, 2023), VeRA (Kopiezko et al., (023) and rsLoRA (Kalajdzievski 2023) have further refined this approach low-rank paraneterization. Other approaches focus on compressing optimizer states directly: methods like Fatoac (Shazeer & Stern 2018), SM3 (Anil et al. 2019 and quantization techniques (Modoranu)

---

Published in Transactions on Machine Learning Research (07))025)
et al. 2024 Feinberg et al. 2024)
factorize reduce the precision of moments. Alternatively, stateless methods like signSGD (Bernstein et al. 2018) and SWAN (Ma et al. 2024 avoid momentum accumulation by carefully normalizing gradients each iteration.
A complementary line research, low-rank subspace optimization, projects gradients onto a smaller, dy- namic subspace to reduce memory. This allows full-parameter updates while operating a lower-dimensional space. Key examples include GaLore (Zhao et al. 2024a). Flora (Hao al.) 2024 and ReLoRA (Lialin 2023 However, applying these dynamic subspace techniques effectively presents difficulties.
et al.
Challenges in Online Subspace Updates. While promising, dynamically updating the optimization subspace online, as in GaLore, faces key obstacles. First, managing optimizer states (particularly momen- tum) during abrupt subspace changes is non-trivial; existing methods may reset, states (Lialin et al. 2023). transform them (Hao et al.] 2024], or leave them unchanged (Hao et al,, 2024). Second, determining the new subspace, often based on gradients, can incur high computational costs (c.g., full SVD on large gradient matrices (Zhao et al.] (2024a)), hindering scalability. These limitations highlight the need for more efficient, iteration- level subspace adaptation strategies.
Contributions
Motivated by the success gradient-based techniques for memory-eficiint fine-tuning, the advantages of dynamic projection subspaces over static ones, and the potential for online subspace updates to better track full-rank optimization dynamics, we address three research questions:
(Q1) Can we mitigate the challenges of online subspace resampling through a computationally efficient update strategy while rigoronsly tracking projection residuals?
(Q2) Can we directly estimate full-rank adaptive optimizer trajectories using low-rank factors, instead of relying on compressed accumulation of subspace moments?
(Q3) Can we match the theoretical per-iteration complexity of standard adaptive optimizers while main- taining LoRA-level memory savings?
We answer these questions affirmatively by introducing MoFaSGD, which maintains dynamically updated low-rank SVD factorization of the first-order momentum. We define this factorization as:
M, = U++1 Diag(o.+1)V R BM.1 + (1-3)). (First-Momentum Factorization)
(1)
Herc, M, denotes the approximate first-order momentum at iteration t, G2 is the gradient at itcration f, and 3 is the momentum decay factor. The low-rank SVD factors U+i1 E Rmxr (left singular vectors), o E R"" (singular values), and V4+1 E Taax (right singular vectors) serve as the optimizer state in our algorithm.
Motivated by our observation that the exponential moving average (EMA) of gradients, defined Et often exhibit low-rank structure, MoFaSGD applies specific low-rank approximation strat- egy: maintaining an online, low-rank SVD factorization of the first-order momentum. This differs from state compression methods that typically target second moments (Shazeer & Stern 2018) and avoids the complex factored approximations of curvature used in methods like Shampoo or KFAC (Gupta et al.) 2018; Martens & Grosse 2015). The factorization is updated efficiently at each iteration using tangent space projections, bypassing costly offline resampling (Zhao et al. 2024a). Crucially, MoFaSGD leverages the same continuously updated low-rank momentum factors (U--1; V.-1) perform spectral normalization (We+1* + Wi- nU VTT).. This integrated approach, inspired by methods like Shampoo (Gupta et al, 2018) and Muon (Jordan et al.] 2024b), provides adaptive step directions without requiring separate, com- putationally intensive matrix operations (e.g., SVD, matrix roots, or Newton-Schulz iterations) at each step.
MoFaSGD achieves LoRA-like memory savings through a unified low-rank momentum representation, while also enabling efficient, adaptive full-paramcter updates via spectral normalization. This positions our method as a low-rank variant of Muon. Unlike MoFaSGD, Muon maintains full-rank momentum buffers (O(m)) memory) and applies full-rank updates after normalization.
2

---

Published in Transactions on Machine Learning Research (07/2025)
2
Related Work
To improve memory efficiency for training large-scale neural networks, various strategies have emerged, including parameter-cffiient fine-tuning, subspace optimization, optimizer state compression, statcless ap- proaches, and partial updale techniques. Moreover, second-order and precouditioning methods such as Shampoo (Gupta et al. 2018) have been gaining popularity due to their faster convergence compared to AdamW, albeit with even higher memory requirements, as studied in detail in Kasimbeg ct al. (2025).
Below, we detail relevant methods like subspace oplimization, optimizer stale compression, stateless ap- proaches, and gradient spectral normalization; other approaches are deferred Appendix A
Subspace Optimization
Rather than introducing additional parameters through PEFT adapters, subspace optimization methods focus on reducing the memory footprint of gradients and optimizer states directly by projecting gradients onto a low-rank subspace and performing moment accumulation in the subspace as well (Zhao et al.] 2024a
Gur-Ari al., 2018 Gressmann al. 2020 Yang et al.) 2023; Vogels et al., 2020). GaLore (Zhao et al. 2024a) projects gradients onto subspace defined by either the left right singular vectors of the gradients, computed via SVD, and accumulates optimizer states within this subspace, thereby reducing the required optimizer memory. Flora (Hao ct al. 2024) periodically resamples the subspace projection matrix using a multivariate Gaussian distribution for low-rank gradient projections.
AdaRankGrad (Refael et al. 2024) dynamically adjusts gradient rank during training. LDAdam (Robert ct al. 2024) and APOLLO (Zhu et al., 2024) maintain optimizer states in low-dimensional representations, adopting similar strategy to subspace moment accumulation in GaLore (Zhao et. al. 2024a). Subspace management is key challenge. Computing it can require costly operations like full SVD (Zhao et al. 2024a), and its update frequency creates trade-off: infrequent updates can lead to stale information. while frequent updates can disrupt optimizer state accumulation. Managing optimizer states across subspace changes also
requires careful strategies (Tialin et al. 2023 Hao et al. 2024)
Optimizer State Compression. This approach directly reduces the memory footprint of existing optimizer states, primarily the second moments used in adaptive methods. AdaFactor (Shazeer & Stern 2018 and SM3 (Anil et al. 2019) achieve this through factorization techniques, reducing memory complexity from O(mn) to O(m + n2) Quantization methods like 8-bit Adam (Dettmers et al. 2021 and 4-bit variants (Li et al. 2023)
reduce the precision of stored moments, achieving near full-precision performance with significant memory savings. CAME (Luo et al. 2023) and Adapprox (Zhao et al. 2024b) refine these compression techniques for potentially better accuracy.
Stateless and Gradient Normalization Methods. Eliminating optimizer states entirely provides maxi- mal memory efficiency. signSGD (Berustein et al., 2018) achieves this by using only the sign of the gradient. Lion (Chen et al. 2024) similarly reduces the memory footprint of AdamW by discarding the sccond-moment estimate. Gradient normalization methods such as SWAN (Ma et al. 2024) and Muon (Jordan et al. 2024b) emulate the behavior of adaptive optimizers without storing second-momen
statistics. Muon, in particular, applies momentum followed by approximate ootooonaalizatinn of the momentum matrix (UM VT, where M = EMEMVTT using efficient Newton-Schulz iterations (Jordan et al., 2024b drawing connections to accumulation-feee variants of Shampoo (Jordan et al., 2024b [eensstein & Newhouse 2024b). However, Muon still retains the Nesterov momentum, and thus is not considered fully stateless.
Positioning MoFaSGD. The landscape of memory-effiiien optimization presents spectrum of trade-offs. Architectural-baseeee solutions such PEFT methods offer high memory savings but restrict fine-tuning up- dates HHu et al.) 2021) to low-rank space. Subspace methods enable full-rank paramcter updates but encounter challenges such as costly or potentially disruptive subspace resampling and moment accunula- tion (Zhao et al. 2024a Hao et al. 2024]; moreover, subspace techniques are mostly applied to AdamW, while their application to non-diagonal based preconditioning methods (such as Shampoo and its variants (Gupta)
et al. 2018 George et al.) 20188 Jordan et al.) (024b)) remains largely unexplored. Optimizer state com-
pression techniques focus on reducing the storage cost of second moments (Shazeer & Stern 2018 Dettmers
3

---

Published in Transactions on Machine Learning Research (07 2025))
a..
2021).
Stateless methods are memory--fficient but discard historical information (Bernstein et al.
2018 8
Ma et al. 2024). This work introduces MoFaSGD to navigate these challenges.
3
Preliminaries
This section introduces notation and provides background on subspace and adaptive optimization methods.
Notation and Conventions. We denote matrices with bold capital letters (e.g., X), vectors with bold lowercase letters (e.g., x), and scalars with non-bold letters (e.g, z). The Frobenius inner product for matrices is denoted as (:-)) (e.g., (X,Y) = Tr(X Y)). Norms are denoted by || |l. with specífic norms indicated using subscripts (e.g., || || for the Frobenius norm). The loss function is represented by f(-). denote the model parameters W, RR=x" and the full-rank gradient aW lw, E R The subscript index denotes the itcration step of any optimizer
Without loss of generality,
W; denoted as G; aC(W) for each weight matrix. Gradients are defined based on the specific batches used in the corresponding iteration. To simplify notation, use the same symbols previously defined where applicable. For any matrix X E 1m the vectorization operator Vec(-) stacks its columns into a single vector, denoted a = Vee(X) E Rm The vectorized version of a matrix is always denoted by the corresponding bold lowercase letter. The Kronecker product is denoted as 69, such that for X E Rm and Y E R the Kronecker product X @ Y E mm Xx1n2, The symbol O denotes the element wise (Hadamard) product between two matrices of the same dimensions. See Appendix D.2 for further notational details.
Adaptive methods with Switched-off Momentums
Without first-order momentum, adaptive methods can be framed as a form online mirror descent:
W++1 arg min L(w,) - (w we, 9t) +
2n 1/ welle. = =0 - -1PP-'a 9:
(2)
w
where the preconditioner P. can be any positive semidefinite matrix, and L(wr) + +(0r+1 - w.,gt) gt) is the local
w.
approximation of the loss function. For instance, if we let Pe.AdaGrad full-matrix Adagrad (Duchi et al. 2011).
T we exactly recover
Similarly, if we define L, = Ei= G.G.T and R. = G. then letting Pt. Pt,Shampoo (R, Ø L:) recovers the Shampoo updates for matrices (Gupta et al. 2018/. Furthermore, if we consider that 9.9, closely approximates the switching off covariance momentum (sccond moment) and considering the following preconditioners basod
components of the true Hessian (Morwani et al. 2024). this motivates
on Adagrad and Shampoo: Pt,1 = [ and Pt. = (G. G, 8 G.G.)+.
A simple derivation shows that using the diagonal version of Pt. in Equation 2 recovers Signed-SGID (Bern- stein et al. 2018). Morcover, if we consider the reduced SVD of the gradient G, UG, Eg, VT, then can write Pt,- = (Vg. VG, & UG. U Using Lemma D.1 in the appendix, we have:
Pt.2 (Va, Ed VJ s Uc. Eg U U2,)= (Vc. & (E. (V.. UC.)
(3)
Furthermore, write 9: the vectorized gradient into Equation 2 yields:
Vec (VG. & Uc.) Vee Substituting the preconditioner Pt, and
w++1 = w. - nVa.. s Ua)) 8 Eo.)(Va. @ Tg.
w. (VG. & UG.) Ee.) Vec EG.)
(4)
w. - (VG. @ UG.) Vec(I,) = w, nUG Vo.
Thus, using Pr2 recovers spectrally normalized updates, sometimes referred to as gradient whitcning (Jordan) et al. 2024b Ma et al.] 2024). For similar derivation, see Bernstein & Newhouse (2024b). Additional background the historical development of adaptive optimization methods provided in Appendix B
4

---

Published in Transactions on Machine Learning Research (07 (225))
Subspace Optimization Methods
Many studies support the conjecture that gradients during the training of deep neural nelworks exhibil low-rank structure, lying in a low-dimensional subspace (Gur-Ari et al. 2018 Gressmann et al.
2020 Yang
ct al. 2023). This property has been studied both thcorctically and empirically and has been leveraged to improve optimization algorithms ranging from commuicatiin--fficenn distributed training (Vogels et al.
2020) to efficient fine-tuning (Hu et al. [202]] Tialin et al., 2023 Hao et al., 2024 ZZhao et al.) 2024a)). Zhhoo iet al. (2024a) aim to explicitly leverage this property to perform subspace training, where gradients are projected and accumulated in low-rank subspace, leading to significant memory footprint reduction.
We briefly clarify subspace methods mathematically, following the formulation of GaLore (Zhao et al. 2024a
Let W, E _ represent the model parameters, and let Qr E Rmxr be the projection matrix defining the subspace at iteration t. Then, GaLore performs the following update:
Ger Q G. E grxa
(Subspace projection)
Mer M + (1-
(First subspace moment)
Vir 32 V211 + (1 (1-Ba) (1-B2)(Gur G.,))
(Second subspace moment)
(5)
W W. n1
M.r
(Project back and update)
Vv...
Moreover, by switching off subspace momentum accumulations, the GaLore update simply be seen as projected gradient descent: W(+1 W. - n Q, G. In addition to subspace accumulation, GaLore performs offline updates of the subspace Q, by updating as U the top-r left singular vectors of Gt. obtained via an SVD operation at predetermined intervals.
4 Momentum Factorized SGD with Spectral Normalization
This section introduces MoFaSGD, our memory-efficient optimization method, and its thcoretical underpin- nings. We begin in Subsection 4.1 by detailing the core algorithmic ideas behind MoFaSGD. This includes the motivation derived from the low-rank structure observed in optimizers, the process of maintaining and efficiently updating low-rank momentum factors using tangent space projection, and the subsequent use of these factors for spectrally normalized parameter updates (Algorithm D Following the algorithmic descrip- tion, Subsection 4.2 presents the convergence properties and theoretical analysis of MoFaSGD, justifying our design choices and establishing formal performance guarantees.
4.1 The MoFaSGD Algorithm
Motivation: Low-Rank Momentum Factors
and Σ B G G maintain low-rank structure and spectral decay properties throughout training, and Zhao et al. (022a) argue that the gradients themselves become low-rank during fine-tuning. Building on these observations, we conjecture that the gradient FMAs exhibit, low-rank properties. We experimentally evaluate this conjecture in Section 5.3 and show that the mass of the gradient EMA largely centered on its top few singular values. Low-rank structure has also been widely leveraged for LLM fino- ano-tuning, c.g., LoRA (Hu et al.) [021] p adapts low-rank matrices during training, while GaLore (Zhao et al., 2024a leverages the low-rank structure of gradients. Thus, leveraging the low-rank property of the gradient EMA connects the implicit assumptions underlying GaLore and LoRA.
Feinberg al. (024)) show that the EMA of gradient covariance
Building on this conjecture, we propose to maintain a low-rank SVD factorized representation of the first momentum. We highlight that this factorization plays two major roles in our algorithm design. First, is leveraged for online subspace sampling, and second, it provides a low-rank estimation of the first momentum (gradient EMA). Formally, we define the low-rank moment factors as:
M, 2 UnI E VII x
Σ g G
(Low-rank Moment Factors)
1=1
5

---

Published in Transactions on Machine Learning Research (07/2025)
Algorithm MoFaSGD: Momentum Factorized Stochastic Gradient Descent
Require: Step size 77, decay 3, rank r
function UMF(G,, (G,, U, En, Vo, 3)
Ensure:
Optimized weights Wi
Compute subspace projections:
Initialize: Wo
2: G.V,, U,G,, Gt, G.G.V.
G.V.
Go CWW()
3: Compute QR factors:
Initialize moment factors: (Un, En, Vo) (VD)
4: (U, Ru.) = QR( G.VI)
t 2+0 0
(Vi, Rv, ) = QR(V, U.
5: repeat
Construct 2r 2r matrix:
6:
G. VwL(WW)
3E U, G.V,
I, Rv.
7:
(U++), E4I V(+1) UMF(G, U, E. Vz:
7: S. = RV,
I
0,
8: W4+1 + W - n
Compute rank-r SVD:
tet+1
(V,"" + SVD,
10: until convergence criterion mel.
10: U t U,U,
11: return W,
11: V41 t V.V."
12: Σ < Σ
13: return (Uz+1, Es+1, Vi+1)
where (G are the observed gradients until iteration t. The left moment factor U4+1 R and the right moment factor V4+1 E IRxxr are orthogonal matrices, while Σ,+1 E RR'r is diagonal.
Our method maps the full-rank gradient G, to the tangent space of the previous moment factor representation (U,, E,, Vi), to ensure a smooth adaptation of the subspace when a new gradient arrives. Formally, we leverage the tangent space of the previous iteration T(.,E..))) as the new subspace for projection, which we denote Te:
T. (UM,,, + UU, + UUV and the projection to this subspace can be derived as follows:
E UP E mm V, E RXT U, = 0, v, =0
(6)
G, - Projj, (G.) U,U,T G. G.V.V - G.V.V.
(Online subspace projection)
where the definition of projection to the tangent space is Proj7, (Gt) arg min || G..lP. GET.
Projecting onto the tangent subspace of previous gradients as shown in Theorem compression error compared to the left, right, or two-sided subspace projections used in GaLore (Zhao et al. 2024a).
4.3 results in lower
Efficient Momentum Factor Updates
The second component of our approach is to efficiently approximate the current moment factor M, given the projected gradient Ge and the previous moment factor M,-1. A naive update, M, = SVD, (G + 3M.-1) where (.D, denotes the rank-r truncated SVD, involves computationally expensive SVD operation, which we aim avoid. Since both M,-1 and G are rank-r, their sum has a rank of most 2r. This observation allows lo approximate M, in O(m + n)r²), which is far more efficient than full-matrix SVD. Let Rv.) QR( G. G.V,)), and (V/, Rv,) QR(VV G. GTU,)). ) where QR stands for the QR decomposition, and U E mm22r V, E Baxx2r and Ru,, Rv, E 222x2 Then we can write:
G G.+ U.U,T G G.V.V. U. G.V.VVV
[EΣ; -.G.V.V I,
V U, G.
= URU. BE, G.V. I, I VT 0,
(7)
= U G.V.
0,
I.
[E, UT,-UTG G.V. I,
I,
0, RT,)) and note that the inner matrix has rank at most r, and U,, V." E 1rrrr We can finally write momentum factor update rule as:
Let U.E"'V'' = SVD, (Rv.
=4 UUU"
V4n ='V'"'V" Σ Σ/

---

Published in Transactions on Machine Learning Research (07/2025)
The computation complexity of our approach includes two QR decompositions, O( +11++11 n)r2), one full SVD a 2r x 2r matrix, ((r)), and hence the total complexity is O((m + n)r² + r
From Momentum Factors to Spectrally Normalized Updates
Inspired by the connection between spectrally normalized gradient updates and effective non-diagonal pre- conditioning methods like Shampoo (Gupta el al. 2018), and motivated by the strong empirical performance of Muon (Jordan et al. 2024b), which applies spectral normalization to gradient momentum, we leverage our low-rank momentum factorization (Equation Low-rank Moment Factors for the main optimizer step. The MoFaSGD update rule is:
(9) Here, Ut+1 E Rmxr and Ve+1 E R represent the left and right singular vectors derived from the effi- ciently computed low-rank approximation of the first-order momentum, M, MoFaSGD (summarized in Algorithm]] 1 contrasts with Muon, which operates on the full-rank momentum M, (requiring ((m)) mem- ory) and uses Newton-Schulz iterations to approximate UM, VM. for its update step. MoFaSGD can thus be viewed as memory-efficient, low-rank variant of Muon.
W+1 W. n/ VIT
Furthermore, we highlight the key distinctions between MoFaSGD and GaLore (Zhao al. 2024a). GaLore employs two-stage process: 1) Projecting gradients onto a low-rank subspace defined by the singular vectors of the gradient itself, updated periodically (referred to as offline subspace resampling), and 2) Accumulating
first and second moments (akin to Adam) within this low-rank subspace.
MoFaSGD adopts different strategics for both stages. Firstly, regarding the subspace dcfinition, MoFaSGD performs gradient low-rank projection, but crucially, onto the tangent space defined by the singular vectors of the gradient momentum (M,). This online, per-iteration subspace adaptation contrasts with GaLore's offline resampling based on single gradients. The choice of the tangent space projection is theoretically motivated by its optimality in minimizing projection residuals (Theorem 4.3). while using the momentum's 's singular vectors aims for a more stable subspace that evolves smoothly, mitigating potential noise in individual gradients.
Secondly, concerning the optimizer update, MoFaSGD deliberately avoids accumulating moments within the subspace, unlike GaLore, thereby avoiding potential error propagation from stale subspaces. MoFaSGD directly uses the computed momentum factors (U(+1, V(+1) to perform the spectrally normalized update in Equation g This design choice aims to circumvent potential crrors arising from subspace moment accumula- tion, particularly when the subspace changes frequently (i.e., near-online updates, or small subspace update intervals in GaLore). As empirically supported in Section 5.3 frequent subspace updates can indeed nega- tively impact GaLore's performance, suggesting that subspace moment accumulation errors might increase with the frequency of subspace changes.
MoFaSGD's novelly lies its unique combination of: 1) Projecting gradients onto the dynamically updated tangent space derived from the low-rank momentum factors, and 2) Utilizing these factors directly for spectrally normalized updates, thereby bypassing the potential pitfalls of subspace moment accumulation inherent in methods like GaLore.
4.2 Convergence and Theoretical Analysis
Our analysis addresses the non-convex optimization problem:
min
L(W) Ee([(W
(10)
WER"X
where we assume access to an unbiasod, barianco-bounded stochastic gradient oracle VL(W,C .)). Below, we first. introduce the necessary definitions and assumptions that underpin our theoretical results.
Definitions and Assumptions
For any optimization itcrate W,, we denote the full-batch gradient by G, = VL(W,) and the stochastic gradient by G; = V.(,W, Formally, we leverage the following standard assumptions throughout our analysis:

---

Published in Transactions on Machine Learning Research (07/2025)
Assumption 4.1. C()) is L-smooth with respect to the nuclear norm ||.] In other words, for any two arbilrary Wi, W2 E TR we have: [VC((W)) V2 & T, - W2ll2
This assumption naturally generalizes the typical smoothness condition from vector optimization and has been previously utilized in the literature (Large et al., 2025 Bernstein & Newhouse 2024b). For further details, please see. Appendix D.1 Additionally, we assume the availability of stochastic gradient oracle satisfying standard properties:
Assumption 4.2. For any model parameter W, we have access to unbiased and variance-bounded stochastic oracle, follows: Eg Eg(VL((.E)) = VC(W), and BeIlVV((..C) V[(W)]ll]
With these definitions and assumptions clarified, we now present an intuitive overview of our main theoretical results, which highlight the strengths and optimality of our proposed MoFaSGD algorithm.
Optimality of Tangent Space Projection (Theorem 4.3 p. We establish that projecting each gradient G. onto the tangent space defined by its singular vectors achieves the minimal projection residual error among broad class of low-rank projection schemes, such as projection onto the left or right singular vector subspaces. Optimal O(1VVT) Convergence Rate (Theorem 4.5). Under Assumptions 4.1 and 4.2 MoFaSGD achieves convergence a stationary point at the optimal O(1/VT) rate. Critically, the factorization of momentum does not degrade the asymptotic convergence rate.
Proof Outline
Our proof structure involves three primary steps. First, we docompose the momentum low-rank approxima- tion error by defining the full-rank momentum M, = - 'G, and breaking down the approximation error into two key components: || BM,-Gii G.Il- < [M, - M,
These terms are individually con- trolled via exponential averaging and tangent-space projections. Second, we rigorously prove the optimality of the tangent-space projection CaceiionnThhoeeeee (Theorem 4.3 demonstrating that G, = U,U, G, + G,VVVV -U,U, G.V.V, minimizos the residual |[G, - GGlle, thus ensuring optimal fitting into the evolving momentum subspace. Lastly, by combining the upper bounds on the aforementioned terms with a standard descent lemma under muclear-nrrm smoothness, we derive our final convergence behavior. By carefully bounding the appraxi- mation terms from the prior steps and leveraging them in our derived descent lemma under muclear-norm smoothness, we arrive at our main convergence result.
The low-rank factorization of gradient momentum as described by Equation Low-rank Moment Factors is the cornerstone of our proposed method. The quality of the momentum approximation plays key role in the effectiveness of our approach. We provide intuitive sketch of the theoretical analysis to bound the factorization residual || - Mell- = Ei= 3"= - U.+ Ex+ V.T l... By recursively bounding this term, we can show that it is sufficient to bound the term |U,++I VIT JU,E.V - -Gele,, which can itself be shown to be bounded by PPorjj, (G.) - GrllE. Thus, the quality of the factored momentum approximation is directly related to the residual of the gradient's 's low-rank projection.
Results
The choice of tangent space projection is optimal in the sense of minimizing the term || Projj, (Gt) - Gelle as demonstrated in the following theorem.
Theorem 4.3. Let L E Rmxr R E Raxr be any arbitrary sketching matrices, and (a1, 022 03) be any arbitrary triple of scalars. Let ( (G) = a, LL G az GRR LagLL GRR Then, the projection residual is minimized when (1) 02, 03) = (1, -1) and R I,. In this case, the residual norm is Proj R) (Gt) Gill = ||(I LL G. RR I|.
Remark 4.4. If we let L = U and R= = Vl"", then we can conclude that the residual error would be upper lamin(m,n) = Emin(mmn) o',, which, considering the low-rank property of the gradient, we bounded by GG. expect to be small, or even zero if the rank of the full gradient is less than r. Note that the subspace projection in GaLore (Zhao et al. (024a) actually equivalent letting either (o1, 0g) = (0,0, 1) or (on, 02 03) (1,0,0).
8

---

Published in Transactions on Machine Learning Research (07 (07/2025)
We now present main convergence bound for Algorithm 1
Theorem 4.5. Let Assumptions 4.1 and 4.2 hold. Moreover, assume rank(Go) S r. By letting 3 and
7) S L, the iterates of Algorithm 1 satisfy the following:
( B[(((W++))
((W))
+ +n +
(11)
T
n/
vT
1=0
Moreover, if we set =
((Wo)-BiCr B[C(Wr
we can derive the following simplified bound as:
TL
(C(W)) B[C(W++
VI++
(12)
VT
Remark 4.6. The convergence metric used for Theorem 4.5 is the stationary point with respect to the nuclear norm, which is averaged over all gradients. When setting o = 0, Algorithm achioves the rate of O vT): which is known to be optimal in the sense of non-convex stochastic optimization under smoothness and an unbiased, bounded stochastic gradient oracle (Arjevani et al. 2023).
5 Experiments
We evaluate MoFaSGD's effectiveness and efficiency across three large language modeling setups: pre- training. natural language undorstanding (NLU) fine-tuning, and instructoon-tuiing These sctups allow us to assess MoFaSGD's performance across different training regimes and task complexities.
5.1 Pre-training setup: NanoGPT Speedrun
We first evaluate MoFaSGD a pre-training context using the Modded NanoGPT benchmark (Jordan et al. 2024a). This benchmark focuses on training a GPT-2 architecture on subset of the FineWeb dataset (Penedo et al., 2025) and measures performance using validation perplexity on a held-out partition of Web. The benchmark's default optimizer is Muon (Jordan et al. 2024b). which holds current training speed records for this task and serves as strong, competitive baseline.
We compare MoFaSGD against full fine-tuning baselines AdamW (Loshchilov et al. 2017 and Muon, as well as the low-rank bascline GaLore (Zhao et al.) 2024a). Following the standard NanoGPT speedrun setup, we use the hyperparameters luned for Muon. We lune the learning rates for AdamW, GaLore, and MoFaSGD, along with GaLore's SVD) frequency and MoFaSGD's momentum decay (3) via grid search,
selecting the best configuration based final validation perplexity (details in Appendix C.2 We evaluate ranks E (16, 32, 128), common choices for low-rank methods in pre-training setup.
Our primary experiment uses budget of 0.73 billion tokens from FineWeb, aligning with the budget used by Muon to reach the target perplexity of 3.27. Figure shows the validation perplexity curves. MoFaSGD consistently outperforms GaLore across all tested ranks, achieving lower final perplexity. The performance advantage particularly noticeable at lower ranks (r = 16). This suggests MoFaSGD's dynamic subspace tracking more effective at capturing important momentum dircctions than GaLore's infrequent updates, especially under strict rank constraints. Both low-rank methods underperform the full-rank AdamW and Muon baselines within this specific token budget. This is likely because full-rank methods have more degrees of freedom, and the 0.73B token budget, optimized for Muon's convergence speed, might be insufficient for low-rank methods to fully match their performance.
To assess longer-term performance, we conduct an extended run for 10, 000 steps (~ 5.3B tokens) using rank r = 32 for both MoFaSGD and GaLore. As shown Figure 3b] MoFaSGD maintains its performance advantage over GaLore, indicating that the benefits of its momentum factorization approach persist during longer training phases.
9

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Ablation: Convergence Efficiency
We conduct. detailed ablation study on the effect of the low-rank parameter {16, 32, 128) during NanoGPT pre-training for both GaLore and MoFaSGD. As illustrated in Figure and Figure 2 higher ranks consistently improve convergence speed and final validation loss. MoFaSGD achieves smoother loss curves and stronger performance across all ranks, particularly under tight memory budgets (e.g., 16), where GaLore exhibits noticeable instability. This supports our claim that MoFaSGD's tangent-space subspace tracking better preserves optimizer continuity under aggressive compression.
In terms of rumtime (Table Db, GaLore shows minimal runtime variation across ranks due to dominant offline SVD costs. In contrast, MoFaSGD's runtime scales more significantly with rank, reflecting its per-step online factorization cost. Nonetheless, MoFaSGD remains faster than GaLore at 32 and achieves superior final loss, highlighting favorable trade-off between expressivity and computational efficiency.
Table 1:
Comparison of MoFaSGD and GaLore across different ranks during NanoGPT pre-training. Best values per row are in bold.
Final Val Loss
Runtime (s)
Throughput
Rank
MoFaSGD
GaLore
MoFaSGD
GaLore
MoFaSGD
GaLore
16
3.8981
4.0773
2156
3755
338,450
194,395
32
3.7208
3.8953
3972
3911
183,770
186,619
128
3.5700
3.6561
4817
3839
151,527
190,131
MoFaSGD
GaLore
MoFaSGD (16)
GaLöre (32)
Validation Loss
MoFaSGD (128)
Validation LOSS
GaLore (128)
200
600
1000
1400
600
1000
Steps
Steps
Figure 1: Validation loss vs. training steps for MoFaSGD (left) and GaLore (right) across ranks E
{16, 32, 128}.
MoFaSGD shows smoother and faster convergence.
MoFaSGD (32)
GaLore (16)
10
McFa5GD (16)
10
Galore (32)
Validation Loss
MoFaSGD (128)
Validation Loss
GaLore
1000
2000
3000
4000
5000
500
1000
1500
2000
2500
3000
3500
4000
Wall-clock Time (seconds)
Wall-clock (seconds)
Figure 2: Validation loss vs. wall-clock time across ranks. runtime efficiency.
MoFaSGD scales better in convergence and
10

---

Published in Transactions on Machine Learning Research (07/2025)
GPT2 Perplexity
GPT2 Perplexity (Longer runs)
12
Galore r=
Galore-332
MoFaSGD r-16
10
MoFaSFD-r32
MoFaSGD r-32
GaLore r=32
Validation Loss
MaFaSGD r=128
Valldation Loss
7
AdamW
Muon
200
400
600
860
1000
1200
1400
2000
4000
6000
8000
10000
Step
Step
(a) GPT-2 Perplexity over 1.39k steps.
(b) GPT-2 Perplexity 10k steps
Figure3 3: Comparison of optimizer performance GPTT2 using validation perplexity loss.
5.2 Post-training Evaluation: Setups and Results
We evaluate MoFaSGD in two post-training scenarios: standard NLU fine-tuning on the GLUE benchmark and large-scale instruction-tuning on the Tulu3 dataset.
GLUE Benchmark. We follow the experimental setup detailed in Zhao et al. (024a for fine-tuning a RoBERTa- Base model (125M parameters) on seven diverse NLU tasks from the GLUE benchmark (Wang) MNLI, QQP, SST-2, MRPC, COLA, QNLI, and RTE. We use the same hyperparameters as Zhao 2024a for comparison baselines to ensure fairness. For MoFaSGD (with ranks r = 4 and r = =), we
2018):
ct al.
tune the learning rate for each task based on validation accuracy, keeping the momentum decaa / fixed at
0.95 (details in Appendix C.3 Final validation accuracy and loss are used for evaluation.
Tulu3 Instruction Tuning. To assess performance more complex alignment tasks, fine-tune the LIaMA-3.1 8B model on the 3-f-----sttmxxteeee-e--- dataset, large (~ 900K samples) and diverse instruction- tuning dataset (Lambert et al. 2024). We adopt most hyperparameters from the Tulu3 setup (Lambert et al. 2024), training for one epoch with an effective batch size of 128; however, we use only subsample of 200K examples to reduce the training budget. We perform a grid search over learning rates for each optimizer. For the low-rank methods (MoFaSGD, LoRA, GaLore), we use rank r 8. Key hyperparameters for MoFaSGD (B) and the baselines (GaLore SVD frequency, LoRA alpha) are as specified in Appendix C.4 We evaluate the final checkpoints using the OLMES evaluation framework (Gu et al.) 2024), benchmarking across MMLU, TruthfulQA, BigBenchHard,, GSMSK, and HumanEval.
Performance Analysis
In these post-training tasks, compare MoFaSGD against LoRA Hu et al. 2021 optimized with AdamW, and GaLore (Zhao et al. 2024a), which are prevalent methods for memory-efficient fine-tuning. We also include results for full-parameter fine-tuning with AdamW as a performance ceiling reference. Table 2 compares the thcoretical complexitics, highlighting MoFaSGD's comparable memory footprint to LoRA alongside efficient online subspace updates, while Figure comprehensively details the memory usage of MoFaSGD compared to other basclines our Tulu3 instruction-tuning setup.
|
To illustrate MoFaSGD's optimization efficiency during instruction tuning the Tulu3 benchmark, we compare its validation loss trajectory against GaLore and LoRA across both training epochs and wall-clock time. As shown in Figure 5a and Figure 5b MoFaSGD consistently achieves lower validation loss over the course of training, indicating superior sample efficiency. Notably, when measured against real-world wall-clock time, MoFaSGD converges faster than both GaLore and LoRA, demonstrating improved practical efficiency in addition to theoretical gains. These trends are further supported by our throughput analysis: MoFaSGD reaches 4206 tokens/sec, outperforming GaLore (3214 tokens/scc) and approaching the high throughput of LoRA (4536 tokens/sec). These findings reinforce our earlier conclusion that MoFaSGD's
11

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Optimizer Memory Complexity
Subspace Resampling
Breakdown Different Optimizer Satups
GaLore
mun wr + 2tur
O(nEn) (offline)
LoRA
+ 3mr + 3nr
MoRaSGD
man + wer nor +r
((Im + n)r") (ouline)
Table 2: Comparison of memory and subspace sampling complexity for low-rank optimizers. Let W E mxxx represent model parameters, and is the rank of low-rank optimizers (w.l.o.g assume m S n). Note that memory complexity includes model param- eters and optimizer states.
Neni (D/8)
Geline Nawdlyed
Figure 4: Empirical memory breakdown (GB) for LJaMA3.1-88 using different. optimizers.
spectrally normalized updates and dynamic momentum factorization enable more effective fine-tuning under memory constraints. We have also included the training loss curves Appendix C.5 showcasing MoFaSGD's convergence behavior on both GLUE and Tulu3 setups.
GaLore
GaLore
1.1
1.1
LoRA
LoRA
Validation Loss
MoFaSGD
Validation Loss
MoFaSGD
1.0
1.0
0.9
0.9
0.8
0.8
0.00
0.25
0.50
0.75
1.00
10000
20000
30000
Epoch
Wall-clock Time (s)
(a) Validation Loss vs. Epoch
(b) Validation Loss vs. Wall-clock Time
Figure MoFaSGD demonstrates superior sample efficiency and faster wall-clock convergence, achieving lower validation loss than both GaLore and LoRA throughout LLaMA3.1-8B instruction tuning on the Tulu3 benchmark.
Table summarizes the final validation accuracies on the GLUE tasks (ranks r=4 and r=8). MoFaSGD achieves performance comparable to, and on average slightly better (+0.5% for r = 4, +0.47% for = 8) than both LoRA and GaLore, while using slightly less estimated memory. This demonstrates MoFaSGD's competitiveness and memory efficiency on standard NLU benchmarks.
Table presents the results the challenging instruction-tuning benchmarks using the Tulu3 setup (rank 8). MoFaSGD outperforms both GaLore (+0.8% avg.) and LoRA (+2.3% avg.) the average across the five benchmarks. This highlights MoFaSGD's potential advantage in complex tasks where accurately capturing the training dynamics over long sequences and diverse instructions is crucial. However, consistent with prior work (Wang el al. 2023), all low-rank methods exhibit a performance gap compared to full fine-tuning with AdamW (MoFaSGD is -42% avg. vs AdamW), This gap underscores the inherent trade-off between memory efficiency and achievable performance, particularly as task complexity increases and may require capturing subtle, high-rank parameter updates that low-rank approximations inherently miss.
12

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Table 3: Comparison of final validation accuracies (%) on seven GLUE tasks (MNLI, QQP, SST-2, MRPC, CoLA, QNLI, RTE) when fine-tuning RoBERTa- base model using different optimizers. For GaLore, LoRA, and MoFaSGD, we report results with rank {4, 8). Memory usage is estimated for each method, and the final column shows the average accuracy. Note that for memory mcasurement, include only the parameters and the optimizer states for fair comparison.
Optimizer
MNLI
QQP
SST-2
MRPC
CoLA
QNLI
RTE
Memory
Avg.
AdamW (Full-Rank)
86.8
91.98
94.48
90.90
62.25
93.15
79.41
747M
85.57
GaLore (r =
85.23
89.62
94.17
90.72
60.33
98.20
77.22
253M
84.36
LoRA (r 4)
84.25
89.73
93.59
90.53
60.42
92.91
78.59
257M
84.29
MoFaSGID (r=
85.12
89.85
94.15
90.78
61.91
93.10
79.08
251M
84.86
GaLore 88=
86.01
89.65
94.04
90.65
59.96
93.16
78.14
257M
84.52
LoRA (r = 8)
85.18
90.15
93.87
90.82
61.11
93.05
78.77
264M
84.71
MoFaSGD (r = 8)
86.32
90.26
94.36
90.75
62.16
93.12
79.28
253M
85.18
Table 4: Final scores of Llama-3.1 8B on the Tulu3-SFT dataset using four different optimizers. The table reports performance on MMLU, TruthfulQA, BigBenchHardl, GSM8K, and HumanEval, along with the average of these five benchmarks (Avg.).
Optimizer
MMLU
TruthfulQA
BigBenchHard
GSM8K
HumanEval
Avg.
AdamW (Full-Rank)
62.8
46.5
66.7
72.7
81.0
65.9
GaLore
58.9
44.2
57.6
68.4
75.2
60.9
LoRA
56.1
42.6
56.5
67.9
73.9
59.4
MoFaSGD
59.4
45.8
58.3
68.4
76.8
61.7
5.3
Ablations
Momentum Spectral Analysis
MoFaSGD motivated by the conjecture that the first. moment (the EMA of gradients) preserves a low-rank structure throughout training. This conjecture stems from the GaLore hypothesis the low-rankness of the gradients themselves (Zhao al. 2024a), and, more importantly, from the observation in Feinberg et al.) (2024), which shows fast spectral decay in the EMA of the gradient covariance, as discussed in more detail in Section 4.1 To investigate our conjecture, we analyze the first-moment buffer M2 from the AdamW
optimizer states generated during the Tulu3 instruction-tuning setup.
For each relevant parameter matrix Mr, we perform SVD and compute the energy ratio captured by the top-r
singular values as Σ d.M. This ratio represents the percentage of the momentum's 's total energy contained within the top-r subspace. We compute the average of this ratio across all 2D weight matrices the model various training steps. Figure (6a] shows the average energy ratio for r 16 and r 32 throughout training. We observe that the top-32 singular values consistently capture around 80% of the momentum's 's energy, while the 16p-16 capture approximately 75%. This persistent and significant concentration of energy in low-rank subspace strongly supports our hypothosis.
M
Impact of GaLore's Subspace Update Frequency
We analyze the impact of subspace update frequency in methods like GaLore compared to MoFaSGD's implicit per-iteration adaptation. MoFaSGD updates its momentum factors at every step, while GaLore performs explicit, costly subspace updates (c.g., full SVD on the gradient) at intervals т. To investigate whether simply increasing GaLore's update frequency (decreasing т) matches the benefits of MoFaSGD's online subspace adaptation, conduct ablation study using the NanoGPT pre-training setup
13

---

Published in Transactions on Machine Learning Research (07 2025)
Mass ratio of top-r singular values in first momentum
GPT2 Perplexity Subspace Update Interval (r)
r-16
T- 300
0.84
r-32
T 75
T-25
Validation Loss
0.82
150
0.80
6-
0.78
0.76
0.0
0.4 Training Completion
1.0
0
200
400
600
800
1000
1200
1400
Step
(a) Momentum Spectral Mass Ratio
(b) GaLore Update Frequency (r) Ablation
Figure 6: (a) Average mass ratio of the AdamW first moment captured by top-r singular vectors during Tulu3 fine-tuning (r = 16, 32). (b) Validation perplexity vs. subspace update interval (r) for GaLore (r = 32) on
NanoGPT.
0.73B token budget with rank r = 32). We vary across {10, 25, 75, 150, 300} steps. Figure 6b shows the validation perplexity curves.
Very frequent updates (r or T = 25) do not yield the best performance and are slightly worse than less frequent updates (e.g., T 150). This aligns with GaLore's findings (Zhao al. 2024a and suggests that overly frequent subspace changes in GaLore can disrupt optimizer state accumulation. This finding underscores the challenge of online subspace changes. This finding suggests that MoFaSGD's approach of avoiding abrupt subspace changes offers more stable and efficient way to leverage low-rank structures.
5.4 Memory Usage Breakdown and Profiling
We assess MoFaSGD's memory efficiency by decomposing its GPU memory usage across five categories: parameters, optimizer states, gradients, activations, and adapters. Figure shoos a comparative breakdown of memory consumption across six optimizer setups on LLaMA3.1-8B. MoFaSGD achieves total memory usage of 29.4 GB, competitive with fused GaLore and LoRA, while enabling full-parameter updates. contrast, AdammW exceeds GB due to high-cost full-rank momentum buffers and persistent gradient accumulation. SWAN and GaLore (non-fused) similarly suffer from gradient buffers, which dominate their memory footprints.
These savings arise from three key design elements: (i) eliminating second-moment buffers entirely, (ii) maintaining low-rank SVD factorization of first-order momentum, and (iii) fusing gradient projection and zeroing operations during backpropagation, which prevents gradient accumulation from persisting across steps.
To further validate these results, Figure 7 shows the memory trace during MoFaSGD training. We observe clean separation between parameter storage, narrow and persistent optimizer state band, and tightly bounded gradient memory. Compared to AdamW (Appendix C.6 which shows 16 GB persistent gradi- eut bullers and 32 GB optimizer states, MoFaSGD significantly reduces runtime memory pressure. A full quantitative table with GB-level memory usage across all optimizers is provided in Appendix C.6)
5.5 Implementation Details
Initialization. To initialize the momentum factors MoFaSGD, we perform a full singular value decom- position (SVD) once at the beginning of training, using the gradient from the first step. Specifically, we set U, Uir Vo Vor and Σ, Ele corresponding the top-r components. We apply MoFaSGD exclusively to the linear layers of transformer blocks, as spectral normalization is particularly effective for
14

---

Published in Transactions on Machine Learning Research (07 (07))
dep3
Actvolorr
Figure 7: GPU memory trace during LLaMA3.1-BB training with MoFaSGD (r 8).
these layers. Notably, both GaLore (haoo et al. 2024a) and Muon (Jordan et al. 2024b) adopt similar approach: applying their custom optimizers only to transformer linear layers while using Adamm for embed- ding weights and 1D layers. We follow this convention and use AdamW (in bf16) for those remaining layers. Consequently, the GPU memory usage for optimizer states including in fused-GaLore and MoFaSGD is approximately 4.2
Gradient Accumulation and Fused Implementation
Gradient accumulation is essential for training large models with limited hardware memory. For methods such GaLore (Zhao et al. 2024a) and MoFaSGD, promptly clearing gradient buffers via backward hooks
is crucial; otherwise, memory savings akin to LoRA are not realized (see Figure 14p. To address this, GaLore performs fused optimizer step, where the gradient is used to update the correspouding parameter immediately during the backward pass, after which the gradient. buffer cleared.
However, this approach is incompatible with gradient accumulation over multiple micro-batches, since the optimizer step must occur only after all gradients accumulated. To resolve this, introduce low-rank gradient buffers into the optimizer state and register backward hook that accumulates low-rank projected gradients instead of performing the optimizer step immediately. For MoFaSGI,, updating the momentum factors only requires the low-rank projections G,V. U, Gú, and U. G.V. We thus implement temporary low-rank gradient buffer, enabling gradient accumulation while avoiding the need for a persistent full-rank buffer (see memory usage in Figure 7
A similar strategy is used for GaLore: we implement gradden--eccuullaioo--reeedy fused version by storing full-rank gradients in low-rank form. Since GaLore only needs the projection Q G. to update its subspace momentum, low-rank accumulation suffices.
Stateless optimizers
Recent stateless optimizers such as SWAN (Ma et al. 2024 and SinkGD (Seetbon et al. 2025) currently do not have open-source implementations, and we encountered challenges in achieving stable convergence with our preliminary implementations. Nevertheless, given the structural similarity between SWAN and Muon (Jordan et al. 2024b) without momentum, we approximate the memory usage of such approaches by profiling Muon with its momentum buffer disabled. These results are reported in Figure under the label 'SWAN' representative proxy for stateless optimizers.
One technical consideration with these methods is their compatibility with gradient accumulation, which is a common practice in low-resource settings. In such scenarios, the fused backward strategy may not applicable, and persistent gradient buffers are still required in memory (see Figure [T]. While stateless optimizers are promising from a conceptual standpoint, this constraint currently poses practical challenges for achieving memory efficiency on par with crrameter-efficiie fino-tuning approaches like LoRA, or subspace- based approaches like fused GaLore and MoFaSGD.
15

---

Published in Transactions on Machine Learning Research (07/ (07/2025)
6 Conclusion
We introduced MoFaSGD, a memory-efficient optimizer that uses memory comparable to LoRA--lke methods while achieving strong convergence. The core idea is to maintain a low-rank factorization of momentum as the optimizer state and leverage this representation to directly update model parameters at each iteration using spectrally normalized updates, building upon the success similar approaches in full-training scenarios. We provide a comprehensive theoretical analysis and establish an upper bound on convergence for stochastie non- convex optimization, matching existing lower bounds. Moreover, we empirically demonstrate the effectiveness of our method compared to standard low-rank optimization approaches in both pre-training and post-training setups.
Limitations and Future Directions. MoFaSGD shows strong empirical performance, but several lim- itations suggest promising directions for future work. First, as a low-rank subepace optimizer, MoFaSGD may underperform compared full-rank methods on more complex tasks. Exploring adaptive or dynamic rank selection strategies, beyond the fixed-rank settings used in our study, could help mitigale this gap; for instance, future work could investigate monitoring projection residuals to allocate more rank to layers with higher approximation error implementing a budget-aware allocation scheme to optimally balance perfor- mance and memory. Second, although our theoretical results establish convergence gnarantees, they rely on assumptions such as muclear-norm smoothness. The practical relevance and limitations of such assumptions in deep learning sctups remain an open area for study.
Acknowledgment
This work was partially supported by NSF CAREER Award #2239374.
References
Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimization. Advances in Neural Information Processing Systems, 32, 2019.
Yossi Arjevani, Yair Carmon, John C Duchi, Dylan J Foster, Nathan Srebro, and Blake Woodworth. Lower bounds for non-convex stochastic optimization. Mathemalicul Programming, 199(1):165-22) 214, 2023.
Jeremy Bernstein and Laker Newhouse. Modular duality in deep learning. arXie preprint orXiv:2410.2126, 2024a.
Jeremy Bernstein and Laker Newhouse. Old optimizer, new norm: An anthology. arx.v.org. 2024b. doi: 10.48550/arm www.iii.409.2003.0
Jeremy Bernstein, Yu-Xiang Wang, Kamyar Avizadeneshelii, and Animashree Anandkunar. signsgd: Com- pressed optimisation for non-convex problems. In International Conference on Machine Learning, pp. 560 569. PMLR, 2018
Jeremy Bernstein, Chris Mingard, Kevin Huang, Navid Azizan, and Yisong Yue. Automatic gradient descent: Deep learning without hyperparameters. arXiv preprint Nii:2234400188 2023.
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, et al. Symbolic discovery of optimization algorithms. Advances in neural information processing systems, 36, 2024.
Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. arxiv preprinl ar Xix:222100000 028661, 2021.
Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized llms. Advances in Neural Information Processing Systems, 36, 2024.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12:2121 2159,, 2011.
16

---

Published in Transactions on Machine Learning Research (07/2025)
Sai Surya Duvvuri, Fnu Devvrit, Rohan Anil, Cho-Jui Hsich, and Inderjit S Dhillon. Combining axes pre- conditioners through kronecker approximation for deep learning. In The Twelfth International Conference on Learning Representations, 2024.
Vladimir Feinberg, Xinyi Chen, Y Jennifer Sun, Rohan Anil, and Elad Hazan. Sketchy: Memory--fficiett adaptive regularization with frequent directions. Advances Neural Information Processing Systems, 36, 2024.
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approximate natural gradient descent in kronecker factored eigenbasis. Advances in neural information processing systems, 31, 2018.
Frithjol Gressmann, Zach Eaton-Rosen, and Carlo Luschi. Improving neural network training in low dimen- sional random bases. Advances in Neural Information Processing Systems, 33:12140-12150, 2020.
Roger Grosse and James Martens. A kronecker-factored approximate fisher matrix for convolution layers. International Conference on Machine Learning, pp. 573-582. PMLR, 2016.
Yuling Gu, Oyvind Tafjord, Bailey Kuehl, Dany Haddad, Jesse Dodge, and Hannaneh Hajishirzi. Olmes: A standard for language model evaluations. arXiv preprint orXiv:2406.08444 2024.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, pp. 1842 1850. PMLR, 2018.
Guy Gur-Ari, Daniel Roberts, and Ethan Dyer. Gradient descent happens in a tiny subspace. arXiv preprinl 1112:1812,.4454, 2018.
Yongchang Hao, Yanshuai Cao, and Lili Mou. Flora: Low-rank adapters are secretly gradient compressors. arXiv preprint ar Xix:22402. 03293, 2024.
Roger A Horn and Charles R Johnson. Topics in matriz analysis. Cambridge university press, 1994.
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Ges- mundo, Mona Atlariyan, and Sylvain Gelly. Parameter-effiiien transfer learning for ulp. In International PMLR, 2019.
conference on machine learning, pp.
Fdward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Lú, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106. 09685, 2021.
Keller Jordan, Jeremy Bernstein, Brendan Rappazzo, &fernbear. Boza Vlado, You Jiacheng, Franz Cesista, Braden Koszarsky, and @Grad62304977. modded-nanogpt: Speedrunning the nanogpt baseline, 2024a. URL http::///gttuub..co///llle//////
nanogpt
Keller Jordan, Yuchon Jin, Vlado Boza, You Jiacheng, Franz Cocista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers neural networks, 2024b. URL https://kellerjordae gi github. io/posts/muon/
Damjan Kalajdzievski. A rank stabilization scaling factor for fine-tuning with lora. arXiv preprint arXiv:2312. 2327:3211203339, 2023.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. Xix preprint arXiv:2001.08300, 2020.
Priya Kasimbeg, Frank Schneider, Runa Eschenhagen, Juhan Bac, Chandramouli Shama Sastry, Mark Saroufim, Boyuan Feng, Less Wright, Edward Z Yang, Zachary Nado, et al. Accelerating neural net- work training: An analysis of the algoperf competition. Nix preprint arxi::2502. 1100022.5011. 15015, 2025.
17

---

Published in Transactions on Machine Learning Research (07/2025)
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Represenlations (ICLR), 2015.
Dawid Kopiczko, Tijmen Blankevoort, and Yuki M Asano. Vera: Vector- random matrix adaptation. arXiv preprint arXiv:2310.1144. 2023.
Nathan Lambert, Jacob Morrison, Valentina Pyatkin, Shengyi Huang, Hamish Ivison, Facze Brahman, Lester James V Miranda, Alisa Liu, Nouha Dziri, Shane Lyu, et al. T\ ulu 3: Pushing frontiers in open language model post-training. arXiv preprint Xrxiv:241. 15124,, 2024.
Tim Large, Yang Liu, Jacob Huh, Hyojin Bahng, Phillip Isola, and Jeremy Bernstein. Scalable optimization in the modular norm. Advances in Newral Information Processing Systems, 37:73501-73355 2025.
Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficien prompt tuning. arXiv preprint Xr14:21004,08611, 2021.
Bingrui Li, Jianfei Chen, and Jun Zhu. Memory efficient optimizers with 4-bit states. Advances in Neural Information Processing Systems, 36:15136-15171, 2023.
Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprinl arXiv:2101.00 90, 2021.
Vladislav Lialin, Sherin Muckatira, Namrata Shivagunde, and Anna Rumshisky, Relora: High-rank training through low-rank updates. In The Twelfth International Conference on Learning Representations, 2023.
Dong C Liu and Jorge Nocedal, On the limited memory bfgs method for large scale optimization. Mathe- matical programming, 45(1):503 528, 1989.
Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for Ilm training, arXiv preprint ar XrXi:22502.1695 2025.
Shih- Liu, Chien-Yi Yi Wang, Hongxu Yin, Pavlo Molchanov, Yu-Chiang Frank Wang, Kwang- Ting Cheng, and Min-Hung Chen, Dora: Weight-decomposed low-rank adaptation. arXiv preprint 2000003302.0933, 2024.
Nyy Loshchilov, Frank Hutter, et al. Fixing weight decay regularization in adam. Xiv preprint arXiv:1711.05101, 5, 2017.
Qijun Luo, Hengxu Yu, and Xiao Li. Badam: A memory efficient full parameter training method for large language models. arXiv preprint XxXx:2444.00227 2024.
Yang Luo, Xiaozhe Ren, Zangwei Zheng, Zhuo Jiang, Xin Jiang, and Yang You. Came: Confidence-guided adaptive memory efficient optimization. arXiv preprint 3rXiv:2307.02047 2023.
Kai Lv, Hang Yan, Qipeng Guo, Haijun Lv, and Xipeng Qiu. Adalomo: Low mowoyy optimization with adaptive learning rate. Xiviv preprint ar orXi:223101100 10195, 2023a.
Kai Lv, Yuqing Yang, Tengxiao Liu, Qinghui Gao, Qipeng Guo, and Xipeng Qiu. Full parameter fine-tuning for large language models with limited resources. arXio preprint arXiv:2306.0974 2023b.
Chao Ma, Wenbo Gong, Meyer Scetbon, and Edward Meeds. Swan: Preprocessing sgd enables adam-level performance on Ilm training with significant memory reduction. arXiv preprint drXim:2412.18148, 2024.
James Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 1114 2020.
James Martens and Roger Grosse. Optimizing neural networks with fronoeker-factored approximate curva- ture. In International conference on machine learning, pp. 2408 2417. PMLR, 2015.
18

---

Published in Transactions on Machine Learning Research (07/2025)
Ionut- Vonut-Vlad Modoranu, Mher Safaryan, Grigory Malinovsky, Eldar Kurtic, Thomas Robert, Peter Richtárik, and Dan Alistarh. Microadam: Accurate adaptive optimization with low space overhead and provable convergence. arXiv.org, 2024, doi: 10.48550/aal 10.48350/arxiv.24000 /ar.:/x222xx00.55553.
Depen Morwani, Itai Shapira, Nikhil Vyas, Eran Malach, Sham Kakade, and Lucas Janson. A new perspective on shampoo's preconditioner. Xrv preprint arXiv:2406. 17748, 2024.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almcida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730-27744, 2022.
Guilherme Penedo, Hynek Kydllček, Anton Lozhkov, Margaret Mitchell, Colin A Raffel, Leandro Von Werra, Thomas Wolf, et al. The fineweb datasets: Decanting the web for the finest text data scale. Advances in Neural Information Processing Systems, 37:30811-30849, 2025.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1 16. LEEE, 2020.
Yehonathan Refael, Jonathan Svirsky, Boris Shustin, Wasim Huleihel, and Ofr Lindenbaum. Adarankgrad: Adaptive gradient-rank and moments for memory-efficient Ilms training and fine-tuning, arXie preprint wxXiv:2410.17880 2024.
Thomas Robert, Mher Salaryan, Ionut-Vlad Modoranu, and Dan Alistarh. Ldadam: Adaptive optimization from low-dimensional gradient statistics. arXiv.org, 2024.
Meyer Scethon, Chao Ma, Wenbo Gong, and Edward Meeds. Gradient multi-normalizztion for stateless and scalable IIm training. arXiv preprint arXiv:2502.061 06742, 2025.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublincar memory cost. In International Conference on Machine Learning, pp. 4596 4604. PMLR, 2018.
T. Ticleman and G. Hinton. Lecture 6.5 - rmsprop: Divide the gradient by a running average of its recent magnitude. Coursera: Neural Networks for Machine Learning. 2012.
Thijs Vogels, Sai Pranceth Karimireddy, and Martin Jaggi. Practical low-rank communication compression in decentralized deep learning. Advances in Neural Information Processing Systems, 33:14171-14181, 2020.
Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. Soap: Improving and stabilizing shampoo using adam. arXiv preprint
w:Xiv:2409.11322 2024.
Alex Wang. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint 1 1ett18044000740 07461, 2018.
Yizhong Wang, Hamish Ivison, Pradeep Dasigi, Jack Hessel, Tushar Khot, Khyathi Chandu, David Wadden, Kelsey MacMillan, Noah A Smith, Iz Beltagy, et al. How far camels go? exploring the state of instruction tuning on open resources. Advances in Neural Information Processing Systems, 36:74764 74786, 2023.
Greg Yang, James Simon, and Jeremy Bernstein. A spectral condition for feature learning. ar arXiv preprinl arXiv:2310. 17813, 2023.
Flad Ben Zaken, Shauli Ravfogel, and Yoav Goldberg. Bitfit: Simple parameter-efficient fine-tuning for transformer- masked languago-models. arXiv preprint 1N1:19906.10199, 2021.
Qingru Zhang, Minshuo Chen, Alexander Bukharin, Nikos Karampatziakis, Pengcheng Hc, Yu Cheng, Weizhu Chen, and Tuo Zhao. Adalora: Adaptive budget allocation for crrameter-efficienn fine-tuning. arXiv preprint or Xic:2303. 10512, 2023.
19

---

Published in Transactions on Machine Learning Research (07/2025)
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-elficien Ilm training by gradient low-rank projection. arXiv preprint or orXim2403.03350, 2024a.
Pengxiang Zhao, Ping Ti, Yingjie Gu, Yi Zheng, Stephan Ludger Kölker, Zhefeng Wang, and Xiaoming Yuan. Adapprox: Adaptive approximation in adam optimization via randomized low-rank matrices. arXiv preprinl ar Kiv:24403. 14958, 2024b.
Hanqing Zhu, Zhenyu Zhang, Wenyan Cong, Xi Liu, Sem Park, Vikas Chandra, Bo Long, David Z Pan, Zhangyang Wang, and Jinwon Lee. Apollo: Sgd-like memory, adamw-level performance. arxie preprinl 24.::2112.002/ 2024.
20

---

Published in Transactions on Machine Learning Research (07/2025)
A
Additional Related Works
Parameter-Efficient Fine-Tuning (PEFT). To reduce the memory associated with full-parameter up- dates, PEFT methods constrain updates to small subset of parameters add minimal trainable compo- nents. Adapters introduce small bottleneck layers into the model (Houlsby et al. 2019), while the popular LoRA technique (Hu et al.] 2021 injects trainable low-rank matrices into existing weight layers, drastically reducing memory by training only these adapters (Hu et al. 2021). Variants include AdaLoRA for adaptive rank allocation (Zhang et al. 2023), VeR. using shared low-rank matrices (Kopiczko et al., [023 and
DoRA decomposing updates into seggnitudd/directtonn (Liu et al., 2024). Other PEFT methods tune only
biases (BitFit. (Zaken el. al. 202T)) or learn input embeddings (prompl tuning (Li & Liang] 2021 Lester
et al. 2021)). QLoRA (Dettmers et al. 2024) further combines LoRA with 4-bit weight quantization. While memory-efficient, PEFT methods generally keep the basc model fixed and rely the capacity of the trained low-rank components.
Second- Orddrrrrrrrrrareee and Preconditioning Methods.
These methods aim primarily to accelerate convergence and improve optimization performance by incorpo- rating curvature information, offering potentially more effective descent directions compared to first-order methods [George et al.,
2018 Duvvuri et al. 2024 Gupta et al.) 2018). Exact second-order information (c.g..
using the full Hessian in Newton's s method or the full gradient covariance in full-matrix Adagrad (Duchi et al.
2011 Duvvuri et al.
2024g) is computationally infeasible for large models due to the prohibitive cost (Old²)
memory, (dd)) computation) of storing and manipulating the required matrices (George et al. 2018 Duvvuri
et al.
2024 Gupta et al. 2018). Therefore, practical methods rely on approximations to make harnessing
second
information tractable. Quasi-Newton methods like LBFGS (Lāu & Nocedall 1989 Duvvuri
et al.
2024) build implicit Hessian approximations but can still be memory-intensive (Duvvuri et al., 2024
A major family of approximations leverages Kronecker product factorizations. KFAC KLronecker-Factoree Approximate Curvature) (Martens & Grosse 2015) approximates the Fisher Information Matrix (an ap-
proximation of the Hessian (Martens 2020]) block-diagonally using Kronecker products specific to network
Grosse & Martens 2016). Shampoo (Gupta et al., 2018), motivated by full-matrix Adagrad (Gupta et al. Anil et al. 2019). uses Kronecker products of gradient statistics (GGT, GTG) as a preconditioner (Gupta ct al., 2018). While powerful, these factored approximations re- quire computing matrix roots or inverses, which are computationally demanding and often amortized over
layer types (Martens & Grosse 2015
multiple steps, potentially using stale curvature estimates (Gupta et al. 2018
George et al. 2018
Refinements seek to improve accuracy or efficiency. EKFAC |Goorge et al. performing cheaper diagonal updates within the coonecker-factoree eigenbasis better Fisher approximation (George et al. 2018). CASPR (Duvvuri et al.
2018
builds on KFAC by yielding a provably
KKFE).
2024 uses Kronecker sums,
aiming for a potentially better Adagrad approximation than Shampoo's 's Kronecker product (Duvvuri et al, 2024). SOAP (Vyas al. 2024) runs AdamW within Shampoo's eigenbasis, aiming for improved stability and fewer hyperparameters, especially when the basis update is infrequent, (Vyas et al. 2024 These methods
highlight the ongoing effort to balance the power of second-order information with computational fcasibility George el al. 2018 Vyas et al. 2024).
On-the-Fly and Partial Updates. These techniques reduce memory by avoiding large buffers or dis-
tributing states. LOMO (Lv et al. 2023b) applics updates immediately, AdaLOMO (Lv et al. 2023a adds
adaptivity, BAdam (Luo el al. 2024 uses weight blocks, and ZeRO (Rajbhandari et al. 2020 shards states in distributed settings.
B Additional Preliminaries
Adaptive Optimization Methods, Adaptive optimization methods have become essential for training deep neural networks. By adjusting learning rates on a per-parameter (or per-dimension) basis using infor- mation from past gradients and the loss curvature, they often lead to significantly faster convergence than vanilla stochastic gradient descent. (SGD) in practice. This line of work began with Adagrad (Duchi et al.
21

---

Published in Transactions on Machine Learning Research (07 (07/2025)
R mm and performing the update rule W++1 w, npp Dt: where in our notation 9: Vec G R and wu Vec(W;) Note that considering only the diagonal terms of P. and using the FMA of gradient covariances (e.g., Pt P==E= "59,) led to an early variant of RMSprop (Tiilemman & Hinton 2012), and further using EMA of the gradients themselves, 9: - "i'ig 9t, instead of the plain 9r, yields the well-known Adam (Kingma & Ba 2015). However, the mentioned methods are limited forms of second-order methods, where the pre- conditioner matrix is restricted to being diagonal. Considering cither the original Adagrad (Duchi et al. 2011 perspective or second-order optimization methods such as Newlon's method, moving beyond diagonal preconditioners a natural step toward achieving even faster convergence. However, storing a non-diagonal proconditioner is often not feasible due to the size of DNNs, requiring O(m²n²) memory. Shampoo (Gupta et al.
2011), which proposed leveraging the gradient preconditioner P. = Ei=1 9:9;
gt,
2018 proposed maintaining Kronecker approximation Pt R, X Li, where L, Rm*" and R2 E Rax thereby reducing the number of parameters required for storing the preconditioning matrix to Om + Morwani et al. (2024) showed that Shampoo (Gupta et al. 2018) closely approximates the full Adagrad preconditioner, unified all aforementioned methods, and argued that the gradient covariance 9:9. closely approximates the Gauss-Newton components of the Hessian at wy, thereby drawing an interesting connection between Adagrad (Duchi et al. 2011 I and sccond-order optimization methods.
C Experimental Details
This section provides detailed configurations, hyperparameters, luning procedures, and dataset information for the experiments presented in the main paper, aiming to ensure reproducibility.
C.1 General Implementation Details
Experiments were implemented using standard libraries for deep learning, including PyTorch, Hugging Face Transformers, and Accelerate. Specific library versions are detailed in the code repository.
Software:
Hardware:
All experiments were conducted on NVIDIA A100 GPUs The number of GPUs may have varied slightly depending on the specific experimental setup (e.g., 4 GPUs for Tulu3 tuning).
Code: git
The implementation for MoFaSGD is available http://gghhb..ooo/oooooccel/////
C.2 Pre-training: NanoGPT Speedrun
Hyperparameters and Tuning. For this benchmark, Muon hyperparameters were kept. at the tuned defaults provided by Jordan et al (2024a). Learning rates for AdamW, GaLore, and MoFaSGD were tuned via grid scarch over [1e- 2. 4. 4,5e - 4, 8e - 4, le 3, 3e - 3, 5e - 3, 8e - 3, le - 2e - 2, 5e - 22.. MoFaSGD's 's momentum decay 8 was tuned over {0.5, 0.85, 0.90, 0.95}. GaLore's SVD frequency was tuned over {10, 25, 75, 300). The best-performing hyperparameters based on final validation perplexity were selected.
Key hyperparameters selected for the NanoGPT pre-training experiments are summarized in Table
C.3 NLU Fine-tuning: GLUE Benchmark
Hyperparameters and Tuning. The experimental setup, including hyperparameters for baseline opti- mizers (AdamW, GaLore, +Adn+,, directly follows the configuration reported in Zhao et al. (2024a) for RoBERTa-Baan Basse on GLUE. For our method, MoFaSGD, learning rales were tuned via grid search over {le - 5, 2e - 5e 5, 1e - 4, - 4] for each task and rank combination (r and r = 8), and the batch size was fixed to 16 for all experiments. The MoFaSGD momentum decay was also kept fixed 0.95. The learning rate yielding the best validation accuracy on each specific task was selected. The final selected hyperparameters for MoFaSGD across the evaluated GLUE tasks presented in Table 0
22

---

Published in Transactions on Machine Learning Research (07/2025)
Table Final Selected Hyperparameters for NanoGPT Pre-training Experiment (0.73B tokens)
Optimizer
Rank
Learning Rate (LR)
SVD Freq.
MoFaSGD B
Muon
5e-2
N/A
N/A
AdamW
2e-3
N/A
N/A
GaLore
r 16
2e-2
150
N/A
r 32
8e-3
75
N/A
r 128
8e-3
75
N/A
MoFaSGD (Ours)
r 16
1e-3
N/A
0.85
r 32
5c-4
N/A
0.85
r 128
3e-4
N/A
0.85
Batch Size (Tokens)
524,288 tokens
LR Schedule
Stable then linear decay with cool-down of 0.4
Table 6: Final Selected MoFaSGD Hyperparameters for GLUE Tasks (RoBERTa-Base)
Hyperparameter
Rank
MNLI
QQP
SST-2
MRPC
COLA
QNLI
RTE
Learning Rate (TR)
r == 4
le-4
5e-5
5e-5
le-4
2e-5
2e-5
5e-5
r = 8
5e-5
5e-5
2e-5
5e-5
1e-5
le-5
5e-5
MoFaSGD 8
4, 8
0.95
Batch Size Epochs
r 4, 8
16
r = 4, 8
13
C.4
Instruction Tuning: Tulu3
Hyperparameters and Tuning. The setup largely follows Lambert al. 2024). Learning rates for all optimizers were selected from the grid {le - 5e le 5e 4, le 3}. The final LR for each method chosen based on final validation loss held-out subset of the tulu-3-sft mixture. We used 5% of the sampled dataset for validation. Final selected hyperparameters for the Tulu3 instruction tuning experiment are summarized in Table
Table 7: Final Selected Hyperparameters for Tulu-3 Instruction Tuning
Hyperparameter
GaLore
LoRA
MoFaSGD (Ours)
Learning Rate (LR)
5e-5
5e-5
1e-4
Rank (r)
8
8
LoRA Alpha
N/A
16
N/ A
GaLore SVD Freq.
200
N/A
N/ A
MoFaSGD
N/A
N/A
0.95
Effective Batch Size
128
Epochs
Momentum Spectral Analysis.
Analysis was performed on the Tulu3 instruction-tuning run using the default optimizer AdamW with learning rate of 1 and batch size of 128 for one epoch. We directly inspected the AdamW state buffer for the first-moment. EMA (M.) (state [' ['exp_avee av] ']) after each
Methodology:
23

---

Published in Transactions on Machine Learning Research (07 (07/2025)
backward pass periodically. SVD performed these buffers, and the energy ratio for rank was computed as Melp
C.5
Training Loss Curves
The training loss curves Figures Sa and 8b illustrate MoFaSGD's convergence behavior. both the simpler GLUE fine-tuning and the complex Tulu3 instuucioon-tuiing MoFaSGD consistently achieves lower training loss compared to LoRA and GaLore. This suggests more effective optimization dynamics, poten- tially stemming from the synergistic elfect of adaptive low-rank momentum factorization and the spectrally normalized updates, which might offer better preconditioning than the standard updates used in LoRA the subspace accumulation in GaLore.
MNLI-Glue on Roberta-base
TuluV3 SFT on LLama3.1 8b
1.1
1.3
Galore
MoFaSGD
LoRA
1.2
LoRA
Validation Loss
0.9
MoFaSGD
Galore
Train Loss
1.1
0.7
1.0
0.5
0.9
0.8
5000
10000
15000
20000
25000
0.4
0.6
0.8
1.0
Step
Epoch
(a) MNLI Training Loss
(b) TuluV3 Training Loss
Figure 8: Training loss curves for post-training setups.
C.6 Memory Profiling Details
We provide full memory trace visualizations and a quantitative breakdown of memory usage across all opti- mizer configurations (LLaMA3.1-8B) BF16, no activation checkpointing, batch size 1, gradient accumulation 8). These results correspond the experiments summarized in Figure and Section 5.4
Profiling Setup. Memory traces collected using PyTorch's native CUDA memory snapshot utility torch. cuda. www.mmmmmmmmmpeee snapshot at each training step. Each trace illustrates the evolution of GPU memory usage over four training steps. Manual annotations highlight key memory components.
Observations. shown in the traces, MoFaSGD maintains a compact and stable memory profile, with minimal transient gradients and compressed optimizer state bands. AdamW and GaLore (non-fusod) show persistent high memory usage due to full-rank gradient and optimizer state buffers. LoRA adds moderate adapter overhead, while fused GaLore shows improvements over its non-fused version.
24

---

Published in Transactions on Machine Learning Research (07 (07))
Skpp
4xp3
Grade
Optinteo
0.000
Figure 9: MoFaSGD (r = 8): Compact optimizer states and minimal gradient spikes. Total memory 29.4 GB.
$8.00
Step
Step
Step
Step
Aptivationa
Grads
Acoppeers
18.00
Parameters
Figure 10: LoRA (r 8): Slightly higher activation and adapter memory, lotal 3336 GB.
10.00
13.00
Optimizer
10.00
Parametors
Figure 11: SWAN: Stateless optimizer, but full-sized gradient buffers lead to higher memory usage (~43.9 GB).
25

---

Published in Transactions on Machine Learning Research (07 (072)
Step
Step
Activatoon
Grads
Figure 12: AdamW (BF16): Full-rank moments and gradient buffers dominate memory (~70.88 GB).
Step
Activations
Optimieer States
Parametera
0.00
Figure 13: GaLore (Fused, r = 8): Low gradient and optimizer state memory; total 330. GB.
Step
Step
Stop
Optimizer Statos
15.00
0.00
Figure 14: GaLore (Non-Fused, r = 8): Gradient accumulation inflates memory cost (~44.55 GB).
Quantitative Breakdown. The table below summarizes the memory footprint by category for all setups.
26

---

Published in Transactions on Machine Learning Research (07/2025)
Table
Memory breakdown (in GB) for LLaMA3.1-8B fine-tuning.
Optimizer
Params
Opt. States
Gradients
Activations
Adapters
MoPaSGD (r=
15.5
4.2
2.1
7.6
LoRA (r y==8
15.5
4.2
2.1
9.8
2.0
SWAN
15.5
4.2
16.0
8.2
0
AdamW (BF16)
15.5
31.8
16.0
7.5
0
GaLore Fused (r = 8)
15.5
4.2
2.1
8.2
GaLore Non- Fused (r 8)
15.5
4.2
16.0
8.8
Proofs and Technical Details
D.1 On the Choice of Nuclear-norm Smoothness
One of the main assumptions behind our convergence bound in Theorem 4.5 is smoothness under the nu- clear norm (Assumption 4.1 Here, we provide detailed argument to justify this assumption. It is primarily molivated by concepts of descent in normed spaces, dualily maps, and modular duality in deep learning advancements detailed in recent works (Bernstein et al. 2023 Large et al. 2025: Bernstein & Newhouse 20244|b). These provide a solid foundation for spectral normalization of the descent update
direction, which has been popularized in strong optimizers such as Muon (Jordan et al., 2024b 2025
Liu et al.
First, we briefly discuss modular optimization theory, inspired by Large et al. 2025 The core idea of gradient descent is to optimize locally linear upper bound of the loss function, characterized by the gradient. Given loss function L: W - R, we ideally want to find tight upper bound on the higher-onder terms, R(w, Aw):
C(w + Aw) S [(w) + (VL(w), Aw) + R(w, Aw)
(13)
For W E Rd, classical second-order theory estimates the upper bound on higher-order terms using the maximum singular value of the Hessian. This leads to employing the &2 norm with the typical upper bound L Awlle, where is the smoothness constant. However, this assumption often does not hold in deep learning I
practice and may yield a very loose upper bound. Modular theory (Large et al. 2025 aims to go beyond this by introducing an architecture-awarr version the upper bound based on the following conjecture:
L(w + Aw) S ((w) + (VL(w), Aw) + 12
(Normed Space Steepest Descent)
Here, W can be any vector space, and ||-|| W E is an arbitrary norm on it. Modular optimization theory aims to find appropriate norms the weight space of all model parameters W Wi x x Wi. The modular norm this parameter product space W is defined as the scaled max over individual norms (e.g.. max (s1 [ ll,,,,... (toz/lw;))) The specific choice of norms for each parameter is based its role and properties. For instance, for a linear layer W, we expect the normalized feature change not to be drastic, which is characterized by the induced operator norm defined ||AW] max [Wale, Empirical observations show that for typical gradient updates, the normalized feature change of linear layer is indeed close to the RMS-RMS operator norm (Large et al. 2025).
Large et al. 2025 show that by assigning norms to atomic modules (e.g., RMS + RMS for Linear layers and & + RMS for embeddings), the modular norm of general architectures like Transformers satisfies the snoothness property. Informally, Proposition of Large et al. (2025) shows that for a general module M
over X, W, and for common loss functions L(w) = EE. [e [ (w, x)) such as cross-entropy, we have:
VL Aw) v((w))l I Tas Awlly
(14)
27

---

Published in Transactions on Machine Learning Research (07 (025))2
where I|-1|rr is the dual norm of the modular norm defined earlier. Since method targets the linear layers of Transformers, and the dual norm of RMS - RMS is the /fu-ouut scaled nuclear norm, we argue that assuming smoothness under the nuclear norm is more natural than under the Frobenius norm. The Frobenius norm assumption essentially treats the parameter space as R ignoring the matrix structure of linear layers and their architectural role.
However, note that our Assumption 4.1 is still stronger than Equation 14 operator norm for all linear layers and assume the remaining parameters are fixed, letting the weight space we would have:
Even if we use the RMS RMS
be Wlinear (W1,... WN:), then by Equation 14
VC(W,+ VC(W)) s L max (1AW,))2)
where ||.|I+ denotes the muclear norm, and ||-1|2 is the spectral norm, i.e., the &o + l operator norm.
D.2 Mathematical Tools and Notations
Kronecker Product and Vectorization. Here, we introduce the background regarding the definition of Kronecker product, Vectorization, and well-known properties of these operators that make it easier to work with them. The Kronecker product is denoted as @, and for any arbitrary X E Rm x and Y E R is defined as follows, where 2i, are the elements on row i and column j of the matrix X:
T1 Y
T12Y Y
T1mm Y
22 Y
F2 Y
F2, Y
X & Y
(15)
Tmz Y
Tmim Y
Moreover, we define the vectorization operator Vec(.) that stacks columns of the matrix as a vector. Partic- ularly, we have:
31,1
22.1
Vec (X) =
Fmi,1
E gminn
(16)
11,2
Ø2,2
[mi.n1
The following lemma summarizes basic properties of & and Vec(.) that we leverage throughout our proofs in the paper. Since this lemma covers basic, well-known properties in Matrix Algebra, we refer the reader to Horn & Johnson (1994) for proofs and details.
Lemma D.1. Let A, B, C and D be arbitrary matrices and assume that operations in each property are well-defined with respect to their dimensions. We have:
Scalor Multiplication: A (a a B) (aA) B
2. Associativity; (A & B) & C
A 8 (B (B&C)
Distributivity over addition:
(A+B)&C (A8C)+(B8C), and Ao(B+C) = (A&B)+(A8C)
4- Not necessarily commudative: (A 2 B) # B A)
5. Transpose: (A B) A
B
28

---

Published in Transactions on Machine Learning Research (07/2025)
6. if A and B positive-semidefnnte have (A @ B)* A @ B for any positive real s, and if A is posilive definile, it holds for any real 8.
7. Mizee-product property: (A 0 () @ D) = AC X BD
Outer product: u we
Mixed Kronecker matria- vector product: (A" @ B) Vec(C) Vec
D.3 Proofs
Proof of Theorem 4.3 To prove Theorem 4.3 we first introduce the following lemma.
Lemma D.2. Consider the reduced SVD decomposition of the left and right sketches as follows L ULE,V, and R UNENVR. Then, the vectorized wnsketched represenlation of any arbitrory matrix
EE Roxn cun be decomposed as follows:
G =UL. Vec(G)
(17)
where UL [eaxr(mntn-t) is semiiorthogonal matrix (i.e., UL = Immmn-r)) and Σn, e g(mto-r)x(mta->) is a diagonal matria, defined us below:
UL, = [U & UL
8
U
U
UL
Σ 2 I
+ I
Σ
0
0
(18)
Σ
0
R
mm-r & Σ
Proof of Lemma D.2. Recall the left singular vectors of the left and right sketch matrices, UL Rm and and
Un Roxr_ We augment these reduced bases with orthonormal bases as follows: LUz, Už] E mm and RUR. U guxa Since each of these expanded bases is square orthonormal matrices, we can write U = Ir
orthogonal vectors to get full
(U)) = Im- and also U = O. Similar properties hold for UR- Moreover, again based on the orthonormality of the augmented bases, we can write:
ULUI + +UZ( Im
URUR + U In
(19)
Leveraging Lemma
D.1
write:
Vec T GRR
Veo Vec LL BOOLLTG)) G) Voc(GRRR)
In LL LL RR Im Vec(G)
(20)
I + U & Im Vec(G)
Using Equation
20 and properties in Lemma
D.1
we can further write:
I, o -OU.UU (UNUk WUUU+UAUUH) @ (ULEAUZ)
URUR 8 U ( @ULNUU
(21)
(Un U.E (Un (U (UE))
+ +
-UUR & UL) 1, Σ (UR & Un) @U & (UZ RUL)
&
U2) (In
Using similar derivation as in Equation 21 we obtain:
UREAUT 2 Im (UR Un) (2) (U) U (UR ( & (m_)) ( T
(22)
29

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Plugging Equation 21 and Equation 22 back into Equation 20
we get:
Vee( LL' + = (UR (
& (UR 8 Uz)
+ &L + U E2 -m-
Veo(G)
(U Un) - + (Un
+ + (
+ (Un 0 Uz) E 2 In- (Un Veo(G)
(23)
To complete the proof, it remains to show that the block matrix
UL, UR @ UL U R mmnxe(m+an-)
(24)
is semi-orthogonal, iee UL. - I r(m+n-r)
First, note that pairs of sub-blocks are mutually orthogonal. For instance,
(Un UL) (U e UL) (U) @ (U, =0
A similar argument applies to all other sub-block pairs.
Moreover, each sub-block is orthonormal in its own right. For example,
(UR @ UL) (UR UL) (UR (UFUL) I, I, I,a.
Hence, ULR is a semi-orthogonal, and the result follows.
We turn to proving Theorem
Let the reconstructed gradient after projection be â (roj(LLR) on LLTG o2GRR
ag LL T.R... We can write:
Veo(G) Vec(a,
+ 0g (RR ? LL Vec(G)
(25)
Replacing L and R with their corresponding reduced SVD decompositions, we have:
(RR LL = UR U.ELU (UR & (z) @ (UR
(26)
Using Equation 23 from Lemma
D.2 and Equation 26 we can rewrite:
Vec(G)
(UR @ Un) (a.I, & Σ + 2 I. E (U)
t @ Un) (1, {U)
(27)
+ai & U E & -m-,) ) (Un Vec(G)
PL.R Vec(G)
where the last equality, we define PL,R E Rmnxmm as the corresponding projection matrix of Proj
We can use Equation 19
write:
Im I, & Im (UR (L)(UR )U)) (U @ UL) @ UL) + (UR (UR COURRUHH T
(28)
+ +U & (U) 8 T
30

---

Published in Transactions on Machine Learning Research (07 (07/2025)
We now propose our Kronecker decomposition of the subspace projection residual.
Vec(G' - G) = (In & Im PL.R) Vec(G)- =
(UR ( @ E +a2 @ 4 + a E -, I
+ ( H 1 - --
U
(29)
+
E 2 Im-, - I, f (UR U
+ (Un U2) 2 U Vec(G))
For simplicity let A (Un @ Un) I, @ Dj + +2 1I 03 & Σ I, 8 (Un & Uz) T
A = RUR @ (2) @ E -n- & (U @ UL) A3 (UR U2) 2 1M- - I, @
Im-r) (UR & U and A₁ = (UR & (UR @ U First note that 42, A3 Rmaxm. For all distinct (i,j) E (1,2,3, 4), their matrix product is zero, A A 0. Without loss of generality, we show it for A2, and A3 here.
In,, A (UR & UL)((znn & Ej & (UR (Un Ut)
1 & UU)) (Un& (U)Ut)) Im-r I, & Im-r)
(U & UL) E In-, - (U Un
U OUTU) (ayEn & Im-r I, 2 m--r)
(UR 8 Ut)
(U 8 Uz) (a2 & Dg - In-r & & On 0n-rr) 0 @ I - I, 2 Im-r)
(Un 8 Ut)
0pn
(30)
Using Equation 29 and above fact yields:
G G = Veo(G G) - || A2gl1 + || + || ||
(31)
where Il-I12 here is l norm of a vector, and g = Veo(G). The equality holds because (A.g) T (A,g) 0 Omn- Based on the above equation, we have the following lower bound on residual || G > A4gll for any given L, and R. The term Ai does not depend on the choice of (a1, 022, az), also Σι and Σz. Thus, the lower bound on this residual would be tight if and only if we have || = | = | la = 0. By
(01,02, setting Σg I,, ΣR = I, and then considering (01, a22 ars) = (1,1,-1), we have A A = Ag = O
Thus, we have shown that under conditions Theorem 4.3 the first three terms in Equation 31 will exactly zero, leaving the following term as the projection residual:
Vee( ( Vec(Grptimmll - G) (U < U (U @ Ut) Vec(G).
(32)
D.3.1 Convergence Proofs
In this section, we aim find an upper bound on the iteration complexity of Algorithm for finding e-stationary point defined by the averaged muclear norm of gradients. First, recall the necessary assumptions and notations.
Update Rule. Let (W,)z be the iterates of Algorithm Then we have the following equivalent update rule:
W41 = W. - U VLL
M, = U--1 E
M, = Σ g gl-iG
M
G
(33)
1=0
31

---

n 2 Published in Transactions on Machine Learning Research (07/2025)
where M, represents the full-rank momentum, and G. VZ(W&. E Rmx" is the stochastic gradient parameter W. Moreover, let M, be the low-rank momentum factorization, where U, E [Rmxxr V, E B and Σ, E 1
Definitions. We recall the necessary definitions. We let the optimization objective be represented as L(W) = Rgg ([,, ,)], and moreover assume we have access to an unbiased, variance-bounded (a) gradient oracle. For any optimization iterate W,, the full-batch gradient is denoted as G. VC(W;), and the stochastic gradient is denoted G V.(W.: Based assumptions the gradient oracle, have G, [ and E[||G, G.ll+] S Moreover, we say a function C()) Rm + is L-smooth with respect to an arbitrary norm ||-|| if for any two paramcters W1, W2, we have VL((W) -CCW)))V((W)))) S L| W2ll2.
Next, we propose our descent lemma.
Lemma D.3 (Descent lemma). Let {W) be the iterates of Algorithm that satisfies Assumption 4.1 Then, we have:
optimized under loss function
((W(+)) s ((W.) - 7 G/ + 2ny1 MM - Gell- +
(34)
First note that C()) is L-smooth with the nuclear norm. Therefore for any Wi, W2 E m have: [WW1) - VL(W2)ll. < L W2 2. where II-Il2 is the spectral norm. Due to this
Proof of Lemma
D.3.
property, can leverage Proposition 5 in Large et al. 2025 to argue that for any W1, W2 we have:
L W1 - Well ((W) < ((W)) + (VZ(W1). W2 Wi) + 2
(35)
Thus, considering the update rule as in Equation 33 for two consecutive iterates W, and W411, we can write:
((W.+1) s L(W)) GG. V + |0++ ViTalaa
(36)
Note that if we take |1-1|= as function, its subgradient is well-known and can be derived follows: the subgradient set at X = UEVT is 0l X (UVT + HH U H 0, V H = 0, W| S 1}. Since the nuclear norm is a convex function, for any X = UEV and Y, we have:
| > |X|| (UV Y -
(37)
Replace ==
= -M,- G, and X M, U1+ Ex+1 ViTii then can rewrite Equation 36
G.ll- + " 2
((W++1) & ((W)
7/ M, + + m
(38)
where we also used the fact that JUV 1|
Using the triangle inequality, we have || 2 ||G completes the proof.
[M, Gell., and plugging this into Equation 38
As indicated in Lemma D.3 the term |,, Gell+ is the main challenge for deriving the upper bound. Note that we can write || G. < M - - + M - Myll.. Thus, we aim to bound cach decomposed term separately. Note that the first decomposed term M, - G:ll- measures the nuclear norm difference bet the full-batched gradient and the first momentum. Moreover, the second decomposed term ||M, Mell+ measures the low-rank compression error of the full-rank momentum.
Lemma D.4. Under Assumptions 4.1 and 4.1
we have:
T
T-1
To
M||M - Gell-] <
Σ
Ell
[[]
(39)
3
(1 B1-B
Proof of Lemma D.4. Let A, G. G. We can write |M: Gell- S | - Gell+ + A M, + A.l..- Moreover, have |M & |G.|l. + 3 M which, upon unrolling, gives
32

---

Published in Transactions on Machine Learning Research (07/2025)
|M E 2i=0 B G
Thus, can write |M, Gill+
B G
Finally, since
EL= 3³ 1-8
-0+1
I 1-88
1
we have:
B
1
T Σ
Σ
M. Grll- S 1- 3 2=
Grll+ I| +
A
||
(40)
1- 3 r=0
Taking the expectation and leveraging Assumption
have:
Ta (1 (1-
E M|I Gell=] -
E B
(41)
3
We now bound the second term M. M,lla We have the following lemun::
Lemma D.5. Let || - M, be the compression error of the low-rank momentum estimation at iteration
L. Then under Assumplion 4.1 and the condition rank (Gn) S r, we have the following bound:
Σ EiIM M.ll] S nLT
20T
(42)
I-B
VB(1- -
Proof of Lemma D.5 First note that we have:
[M M. |13M--1 + G. Mell. / Mr ||G, + 1M,,. - M.ll.
(43)
Considering the momentum factor update rule in we can replace M, with G. + BM where G. = G.V.VVVV - U,U,TG.V.Y, G.V.V" Therefore, we can write:
|M. - Mil. s B1I MA-II- -
(44)
The term |[I- G V.V,") )ll represents the low-rank compression error of gradients happening in our Algorithm [ Thus, we aim to bound this term. We have:
[I - U.U, U, v.V, )ll. - -,, (G. G.-) -- - + ((--VVv.
I|-
I - ( +| U,U Gi-1 V.V, )ll- + |1 A.-Il
< T, 11-00=++100000)))G-
+ |. -A-Il..
$ nL+ + |1(I U,, -.. (t ||A, A,1l|
(45)
where in the last inequality, we used the smoothness property in Assumption 4.1 and also the fact that W. W, U,
|.. that Range((U,_1 G--VV-ll C Range(Ui), and therefore have Range(U,-1) C RRage (U,). We can
Next, bound the term ||[--UU), Gi II. First note that based Equation
write:
|I(I I G. -I-V, )l1, S II[ U
VV,
+1[(1 ((l1--U,,))G-ie VII
(46)
< ||(I (G -VVVVI V.V < UUUI-U,-U, -
where in the last inequality we leveraged the bound from Theorem 4.3 that upper bounds the tangent space
projection error. Particularly, the fact that || G.-1ll- [I - V.-iV
Plugging Equation 46 into Equation 45 we have the following recursive equation:
[I(T - U.U. - (G V.V, Il.. < |I(T U- -U--U-UL,
V.- -I + n || Ar-1ll-
(47)
33

---

Published in Transactions on Machine Learning Research (07 (07/2025)
Unrolling this recursive relation, and considering the SVD initialization at the beginning of Algorithm
|
have:
I(I -I- U,, V.V. )l. < (I - )G. Go VoV. II. nLt A A
(48)
Let eu = -I Go(I -a Vovo |. Note that by our initialization of Un and Vo, eo is expected to be small. For instance, we assume rank(Go) (Go) S r, then we would have co = Therefore, for simplicity we can ignore this term.
Plugging Equation
48 back into Equation
1
have:
21a
E M
M,Il.] S AE [I MA-ill-] 1 +
(49)
/B
Similar to Equation 40
unroll the above equation to get:
E M.ll.] S LLT 20T
(50)
1-3
VB(1- - B)
where we also used the fact that Mo = My.
Proof of Theorem 4.5 Taking the expectation and unrolling the descent lemma, Lemma D.3 we can write:
T
((W)) BBC(WCWTTT T+1)l
2
Gell-]+ 7²
BIG
Σ
EIll
M
1/T
T
t=0
t=0
(51)
((W)) B(c(WT++))
2
2
El| E[JM M. +
EE M. Gell-]
n/T
T
t=
Plugging Equation
41
[Lemma
D.4 and Equation 50
(Lemma
D.5
into Equation yields
23
[|| G.Il ((W)) E[[(W[[TTt w
2
4a
2nL
(1 1- 3
(52)
T/T
2
(1 8)
3
Let E (0, 1), 11 < 1, and B = T; then we can rewrite the above equation as follows:
(((W) BiL(W
Σ E BIIGH
+ ML
(53)
vT
Finally, by letting 1) =
((Wo)- (.Wo)-BCC(rr.. B(C())))))
we have the final upper bound as follows:
7T.
1 T
(C(W)) - E(C(Wr-1)) VL vo++
Σ
(54)
T
VT
1=0
34


Reviews:
[{"review_id": "RnT2zGcUpf", "reviewer": "Reviewer_Gyu5", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper tackles a practical problem, which is optimizer RAM overhead during LLM fine‑tuning and proposes a conceptually neat remedy that sits between PEFT and full‑rank adaptation. The factorised momentum idea is original and well‑motived by an empirical spectral study \n- The tangent‑space projection is theoretically proved optimal and practically easy to compute via two skinny QR decompositions and one 2r\\times2r SVD. \n- The convergence proof is carefully written under nuclear‑norm smoothness assumption, matching lower bounds and improving upon prior low‑rank optimisers that lacked guarantees. \n- Experiments span pre‑training and alignment‑style fine‑tuning, include strong baselines (Muon, GaLore, LoRA) and demonstrate consistent gains. \n- Code availability increases reproducibility.\n\nWeaknesses:\n- First, although Table 1 lists memory costs, actual GPU peak memory during training is not reported. This is important because MoFaSGD still stores full‑precision activations and gradients, and the extra QR‑SVD may raise activation checkpointing overhead. Second, computational cost is discussed only asymptotically\n- Wall‑clock training times and throughput relative to GaLore and LoRA are missing. \n- All experiments adopt fixed ranks, however, I think an adaptive‑rank study would clarify robustness.\n- While the theory assumes unbiased gradient oracles and nuclear‑norm smoothness, modern transformers violate the latter. A discussion of practical mismatch should have been provided.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces Momentum Factorized SGD (MoFaSGD), a memory‑efficient optimization method aimed at fine‑tuning large neural networks under tight GPU‑RAM budgets. The authors observe empirically that the first‑order momentum buffer (the exponential moving average of gradients) is itself approximately low‑rank. Leveraging this, they maintain a rank‑r truncated SVD of the momentum at every iteration and update it efficiently using a tangent–space projection that avoids expensive full‑matrix SVDs. The same low‑rank factors $U_{t+1},\\Sigma_{t+1},V_{t+1}$ are reused to issue spectrally‑normalized updates $W_{t+1}=W_{t}-\\eta\\,U_{t+1}V_{t+1}^{\\top}$, yielding full‑parameter steps with LoRA‑like memory. Empirically, MoFaSGD is benchmarked against AdamW, Muon, GaLore and LoRA on three regimes: (i) a NanoGPT “speed‑run” pre‑training on 0.73 B tokens; (ii) GLUE fine‑tuning of RoBERTa‑Base; (iii) instruction‑tuning Llama‑3‑8 B on the Tulu‑3 mixture. Across ranks $r\\in\\{4,8,16,32,128\\}$ it attains lower validation loss than GaLore, matches or exceeds LoRA on GLUE and improves over both on Tulu‑3 while requiring comparable memory.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Provide actual peak GPU memory and step‑time measurements on at least one configuration (e.g. Llama‑3‑8 B, r = 8) to substantiate the claimed LoRA‑level efficiency. This is critical for acceptance.\n2. Report training throughput (tokens/sec) for MoFaSGD vs AdamW, GaLore and LoRA in the NanoGPT and Tulu‑3 settings. Important but not blocking.\n3. Add an adaptive‑rank experiment (e.g. shrinking or expanding r according to a spectral‑energy threshold) to test robustness; alternatively discuss limitations. Desirable.\n4. Include sensitivity plots for \\beta and learning rate to evidence stability. Desirable.\n5. **Important: clarify whether nuclear‑norm smoothness holds for transformer losses or justify the assumption.**\n6. The proof of Theorem 4.3 currently sketches the optimality argument.", "broader_impact_concerns": "N/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "X3uSQeAU10", "reviewer": "Reviewer_e2hn", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**<Strengths>**\n* The paper shows a balance between rigorous mathematical derivation and supporting empirical observations. Note that I could not check all the mathematical details and theoretical derivations; but in my perspective, the derivation appears reasonable.\n* Thorough comparisons with prior methods such as GaLore, Muon, and LoRA are provided both in the main text and experimentally, clearly demonstrating MoFaSGD’s contributions. In particular, the online updating behavior of MoFaSGD is a notable advantage over GaLore.\n* Evaluation across multiple architectures and tasks strengthens the generalizability.\n\n**<Weaknesses>**\n* Experimental results mainly compare performance at the same number of training steps or epochs (x-axis). Since the reduced computational complexity does not always translate directly into wall-clock speedups, adding experiments that compare baselines at equal real-time budgets would make the evaluation richer.\n* The method relies on the assumption that low-rank gradients yield low-rank momentum. A formal argument of this property would significantly support the paper’s claim.\n* In Figure 3(a), the distribution of singular values seems to be broadened as training proceeds. This suggests that, for large-scale LLM training (more data, longer training time), MoFaSGD might show a wider gap vs. full-rank optimizers.\n* (minor) Table 2 shows only marginal memory efficiency gains over existing approaches. This is understandable given the small rank r; however, this modest improvement makes this paper more focused to performance and convergence speed rather than the memory efficiency.\n* (minor) Technically, LoRA is not an optimizer and introduces additional parameters. Please consider clarifying this distinction in the manuscript.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes MoFaSGD, a novel low-rank, gradient-based optimizer for deep neural network training, focused on the memory efficiency. The core idea is based on the empirical observation that not only the gradients themselves are low-rank, but their associated momentums also exhibit low-rank structure. By applying an SVD on the gradient and updating the resulting $U, \\Sigma, V$ matrices separately (similar to the moving average), MoFaSGD effectively projects gradients into a lower-dimensional “momentum factor” subspace, thereby minimizing the difference in terms of Frobenius-norm. Unlike GaLore, which requires periodic offline updates, MoFaSGD’s fully online updates simplify the algorithm and leads to faster adaptation to training dynamics. Experiments on NanoGPT, RoBERTa, and Llama3.1-8B models with various benchmarks show that MoFaSGD achieves competitive performance compared to existing low-rank based methods.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* Please see the weaknesses above. In particular, consider including the wall-clock time comparisons—plotting Figure 2 with wall-clock time on the x-axis would be acceptable.", "broader_impact_concerns": "I have found no significant broader impact concerns. The paper is proposing a new optimizer which can be used for almost every deep learning models.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "d6TT3BqeRt", "reviewer": "Reviewer_rb63", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\n1. Provides convergence guarantees and justifies tangent-space projection as optimal among low-rank alternatives.\n2. Demonstrates strong results across pretraining and fine-tuning benchmarks, consistently outperforming GaLore and LoRA.\n3. Avoids full SVDs and complex second-moment approximations, making the method scalable and easy to implement.\n\n## Weaknesses\n\n1. While multiple ranks are tested, a deeper analysis of how rank impacts stability and performance is missing.\n2. Methods like signSGD or SWAN are not included in empirical comparisons.\n3. Some implementation choices (e.g., projection normalization, initialization of momentum factors) could be more clearly explained.\n4. Profiling results on the memory consumption of trainable parameters and momentum could be useful.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents MoFaSGD, a novel memory-efficient optimizer that combines the strengths of low-rank subspace methods and adaptive optimization. It introduces a simple yet powerful approach that maintains a low-rank factorization of the momentum and uses it directly for parameter updates, avoiding costly full-rank operations.\n\nKey contributions include:\n\n1. Low-rank momentum representation to reduce memory usage without sacrificing adaptivity.\n\n2. The paper proves theoretical guarantees of optimal convergence rates under standard assumptions.\n\n3. The paper demonstrates strong performance across pretraining and fine-tuning tasks, outperforming other low-rank methods like GaLore and LoRA while retaining similar or better memory efficiency.\n\nIn summary, MoFaSGD offers an alternative solution for efficient training of large models.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Add a more detailed analysis or ablation on how the choice of rank affects training stability, convergence speed, and final performance. \n\n2. Incorporate baselines such as signSGD or SWAN to position MoFaSGD more clearly in the landscape of memory-efficient optimizers.\n\n3. Provide more clarity on implementation choices such as how momentum factors are initialized and how projection normalization is handled.\n\n4. Include memory profiling results that break down the contribution of model parameters, optimizer states, and momentum buffers. This would support claims about memory efficiency with more granularity.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)