SUMO: Subspace-Aware Moment-Orthogonalization for Accelerating Memory-Efficient LLM Training
Authors: Yehonathan Refael
Guy Smorodinsky
Tom Tirer
Ofir Lindenbaum
Venue: NeurIPS 2025 poster ()
Keywords: Low-rank optimization
Steepest descent
Moment orthogonalization
Newton–Schulz method
Large language model
Memory efficient finetuning
Pretraining
Primary area: deep_learning

Abstract: Low-rank gradient-based optimization methods have significantly improved memory efficiency during the training of large language models (LLMs), enabling operations within constrained hardware without sacrificing performance. However, these methods primarily emphasize memory savings, often overlooking potential acceleration in convergence due to their reliance on standard isotropic steepest descent techniques, which can perform suboptimally in the highly anisotropic landscapes typical of deep networks, particularly LLMs. In this paper, we propose SUMO (Subspace-Aware Moment-Orthogonalization), an optimizer that employs exact singular value decomposition (SVD) for moment orthogonalization within a dynamically adapted low-dimensional subspace, enabling norm-inducing steepest descent optimization steps. By explicitly aligning optimization steps with the spectral characteristics of the loss landscape, SUMO effectively mitigates approximation errors associated with commonly used methods like Newton-Schulz orthogonalization approximation. We theoretically establish an upper bound on these approximation errors, proving their dependence on the condition numbers of moments, conditions we analytically demonstrate are encountered during LLM training. Furthermore, we both theoretically and empirically illustrate that exact orthogonalization via SVD substantially improves convergence rates while reducing overall complexity. Empirical evaluations confirm that SUMO accelerates convergence, enhances stability, improves performance, and reduces memory requirements by up to 20\% compared to state-of-the-art methods.

SUMO: Subspace-Aware NomeentOrttnooonalzzttoo for Accelerating Memory-Effcient LLM Training
Yehonathan Refael Faculty of Engineering Tel Aviv University
Guy Smorodinsky Department of Computer science
Tom Tirer Faculty of Engineering Bar-Ilan University tirer. tom@biu ac il
Ben Gurion University smorodin4post. bgu. agu.ac.11 ac.
yehonathanGtau.cc ac.
Ofir Lindenbaum
Faculty of Engineering
Bar-Ilan University
ofir. wwwdenbaum@bb...emmm ac.
Abstract
Low-rank gradient-hasee optimization methods have significantly improved mem- ory efficiency during the training of large language models (LLLMs), enabling operations within constrained hardware without sacrificing performance. However, these methods primarily emphasize memory savings, often overlooking potential acceleration in convergence due to their reliance on standard isotropic steepest descent techniques, which can perform suboptimally in the highly anisotropic landscapes typical of deep networks, particularly LLMs. In this paper, we pro- pose SUMO (Subspace-Ares Awaree Moment-Oriee Orthogonalization), an optimizer that employs exact singular value decomposition (SVD) for moment orthogonalization within dynamically adapted low-dimensional subspace, enabling norm-ind norm-indueing steepest descent optimization steps. By explicitly aligning optimization steps with the spectral characteristics of the loss landscape, SUMO effectively miti- gates approximation enrors associated with commonly used methods, such the Newton-Schul orthogonalizatton approximation. We theoretically establish an upper bound on these approximation error, proving their dependence on the condi- tion numbers of moments, conditions we analytically demonstrate are encountered during LLM training. Furthermore, both theoretically and empirically illustrate that exact orthogonalization via SVD substantially improves convergence rates while reducing overall complexity. Empirical evaluations confirm that SUMO accelerates convergence, enhances stability, improves performance, and reduces memory requirements by up 20% compared state-of-the-an methods. Code: https://gitthu: com com/guy120494/SUN
Introduction
Low-rank gradient-based optimization methods have become powerful tools for reducing memory consumption during the pre-training and fine-tuning of large language models (LLMs), often without sacrificing performance, and sometimes even improving it. For instance, while pre-training Llama typically requires around 58GB of memory, far exceeding the 24GB available consumer GPUs like RTX 4090, recent advances, such as those discussed [133,, have demonstrated that Llama 7B can now be trained from scratch on single 24GB GPU without the need for costly memory offloading. The theoretical analysis [ attributes this efficiency to the inherent low-rank structure of the gradients, which enables optimization in significantly reduced latent space. Furthermore, 22 found
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

a consistent decrease in gradient rank throughout training, suggesting that low-rank optimization not only reduces memory usage but also converges toward increasingly compact subspaces.
However, despite these advancements, existing methods primarily focus on memory savings and often overlook the potential to accelerate convergence. Current approaches typically rely standard steep- est descent methods and assume isotropic geometry, which can hinder efficiency in ill-conditioned settings. This observation motivates our primary objective: to develop a subspace-aware optimizer that leverages low-rank structure while adapting to the geometry of the loss landscape. By reevaluat- ing the choice of norm and its influence on gradient descent dynamics, we aim to design an algorithm that improves generalization, accelerates convergence, while preserving the memory advantages of low-rank methods.
Classical gradient descent, including SGD performs steepest descent under the Euclidean norm, which reflects isotropic curvature. However, deep networks exhibit highly anisotropic loss landscapes, making this assumption suboptimal. Recent work shows that adaptive optimizers, such as Shampoo [5). SOAP [6), and Muon 77]. can be interpreted as steepest descent under non-Euclidean norms tailored to the network architecture and data structure. As shown in [8) these methods implicitly adapt spectral or operator norms, which better capture local curvature and improve convergence. This motivates the design of subspace-an aware optimizers that exploit both low-rank structure and appropriate geometry to accelerate training.
To formalize the role of geometry in optimization, consider a neural network with a differentiable loss function & W R defined on weight space W R... The local behavior around a point
can be approximated by the Taylor expansion, L(w + Aw) se L(w) + g Aw + À A|| where g = ((w) À > captures the sharpness or curvature of the loss surface and - || is chosen norm reflecting the geometry of the optimization landscape. Minimizing this approximation corresponds precisely to performing steepest descent under the given norm constraint. According to [9]. the solution to this minimization explicitly takes the form,
l|gll- A
Aw-
argmax g t:|||||=1
where ||||+ denotes the dual norm of -|- Adaptive optimizers differ primarily in their norm choices. Adam utilizes a dynamic Max-of-Max norm constraint. Recent optimizers consider matrix norms while applying steepest descent at the layer level. Muon imposes a fixed Schatten-p norm constraint for large p. effectively using the spectral norm on weight matrices [IO] 77. Shampoo [5] dynamically learns the optimal approximate Schatten-p norm for steepest descent, with its variants like SOAP [ apply momentum to efficiently navigate the space of possible norms. Muon, by contrast, operates within a relatively fixed but large Schatten-p norm, balancing the dynamic adaptability of Shampoo with static spectral-norm constraints. Since neural network weights locally act as linear operators on Euclidean spaces, the induced operator (spectral) norm provides a natural constraint that aligns with the curvature of the loss surface. This perspective motivates gradient orthogonalization, which ensures that optimization updates respect the spectral norm, thereby controlling perturbation magnitude and enhancing optimization stability and efficiency [9]
While norm-induced optimization methods offer a principled way to align updates with the geometry of the loss landscape, their practical deployment often incurs substantial computational overhead. For instance, Shampoo requires computing matrix inverses or root operations at every optimization step, which can be computationally expensive for large-scale neural networks. Similarly, Muon's though effective, involves a costly approximation to spectral decompositions, computed by applying five iterations of -ewton-Schulz [12] (referred to as Newton- Schulz5). Therefore, there is an inherent trade-off between the theoretical optimality provided by these norm-induced optimization approaches and their practical computational demands.
first-order moments-orthooonalizaa
To bridge the gap between the geometric ad advantages of norm- indu-induced methods and their computational costs, we first analyze the limitations of existing approximations. We derive an upper bound on the error introduced by the Newton-Schula orthogonalization, demonstrating that this error increases with the condition number of the moment matrix. This finding explains the increasing instability of the Newton-Schulz method in ill-conditioned scenarios, which we subsequently demonstrate to occur in the first-order moment matrices during the training of large language models (LLMs). Building on this analysis, we establish a convergence rate for Muon optimization and compare to an alternative method that replaces the Newton-Schule approach with exact Singular Value Decomposition (SVD).
2

---

Remarkably, we find that the SVD-baser approach converges faster, with improvements directly proportional to the accumulated errors from the moments' orthogonalization by the Newton-Schulzz method. Motivated by the empirical observation that gradients in LLMs often exhibit a low-rank structure, especially during early training, we propose subspace-aware optimization scheme. This scheme performs exact SVD-based moment orthogonalization within a now-dimensional adaptive subspace. This approach benefits from the relatively low computational cost of SVD calculations for low-rank input matrices and enhances convergence stability. Also, our approach attains an even greater reduction in memory usage than all previous low-rank training methods by relying solely on the first-order moment, as detailed below in Table I We support our method with a theoretical convergence guarantee and validate its empirical benefits through experiments, demonstrating faster training and better model performance compared existing methods.
Table I: Comparison of properties between SUMO, GaLore, Adam, Shampoo, and SOAP. Assume W E R with m 2 nI, a constant projection rank and subspace update rate K.
SUMO
Adam
Shampoo
SOAP
GaLore
Computation
O(mar mmn/
O(mm)
O(m + +n²)
O(m² +
O(mer + mn? /K)
Optim. states memory
pr + mr
2ma
m² + +
2mn + 2m² + 2n
2nr + mr
Subspace-aware
x
Orthogonalization
x
x
2 Related Work
Low-rank gradient optimization. Low-rank gradients naturally emerge during neural network training, as shown in both theoretical and empirical studies (13H15). This structure has been leveraged to reduce memory and computational costs during training [16118]. Recent work [2] showed that gradients in reversible layers [19] tend to collapse to rank one over time and used this to adaptively adjust gradient rank in Adam. In this paper, we demonstrate that the same low-rank trend is present in the first-order moment, which we utilize to efficiently apply exact anthogonalization-avviii avoiding the accumulation of approximation errors, such as those encountered in Newton-Schulzz, during optimization.
Memory efficient optimizers. Reducing the memory demands of training large language models (LLMs) has driven extensive algorithmic research. One research direction, initiated by LoRA (20), reduces the number of trainable parameters via low-rank adaptation. Yet, such methods often fall short of fully parameterized models, especially during pre- pre-training Another direction does not restrict the set of trainable parameters but instead optimizes the training methods, with notable examples including AdaRankGrad, GaLore, Fira, Flora, Adam-mini, GaLooreminii LDAdam, GoL.ore, LoQT, and Apollo [1- 25), integrating low-rank gradient projections in optimization. In this work, reduce memory usage even further by relying solely on a first-order momentum, as shown in Table []
Gradient preconditioning. Preconditioning the Gradient method is critical in enhancing the ef- ficiency and effectiveness of optimizers. Several notable approaches for using preconditioner have emerged, including methods based on signed gradients [26 29 gradient clipping (30]. nor- malization (30 31, and gradient whitening 7 B2 37 Recent studies 7 38) explored gradient- orthogonalization strategies, thereby speeding up training. Orthogonalizing gradients effectively constrains updates to lie on directions of uniform magnitude (spectral radius 1), preventing up- dates from exaggerating specific gradient directions over others. This procedure ensures a form of normalization that mitigates potential instabilities from ill-conditioned gradients. Unlike these methods, which apply preconditioning or approximate orthogonalization in the high-dimensional parameter space, our approach performs exact SVD-based orthogonalization within an adaptively selected low-rank subspace, offering improved stability and lower computational overhead.
Orthogonal Stochastic Gradient Descent with Momentum (OSGDM). OSGDM 38 is a re-
cently introduced first-order optimization method that speeds up neural network training by orthogo- nalizing gradients before the optimization step. Specifically, for data batch c(t), OSGDM applies SVD to the gradient matrix G = Ww,C (& (g()).) 0)) of each neural network layer to generate an orthonormal gradient approximation O, This ensures diversity among learned representations
3

---

and reduces redundancy. The update rule for OSGDM with momentum term y and learning rate is defined as,
0 orth (Gi"), M + , + n w + w - ((++))
where orth G) = (GG ) is the ortogonalization operator, and M; is the first order moment of layer 1. Despite the additional computational overhead of SVD, OSGDM empirically converges faster and achieves higher accuracy than common methods such as Adam.
Muon optimizer.
At iteration t, given weight W momentum p., learning rate T/, and objective L1, Muon, introduced by (7), constructs the update rule,
M -M + G 0 (t) Newton-SShulls((MM))
W(+1) W - "))
Here, ( is the momentum at iteration t, initialized as zero matrix when t 0. The Newton- Schulz5 method [35] approximates (M T)-1/2M() (M
orthogonalizing M and thus ensuring uniform update directions, avoiding dominance by few directions. Muon explicitly controls the norm of gradient updates-pann particulalyy the spectral norm (or Schatten-p norm with large p), which limits updates to smaller, well-conditioned steps in parameter space. By constraining the spectral norm, moment orthogonalization implicitly prevents overly large or ill-conditioned parameter updates. Such updates often lead to poor generalization due to instability or overfitting. Shortly after the introduction of Muon, the study in [39] proposed a framework to scale Muon for larger LLMs, mainly adding weight decay, and carefully adjusting the per-parameter update scale.
3 Method and Main Results
3.1 Theoretical Motivation: convergence
Exact moments orthogonalization leads to significantly faster
Previous work on pre- wre-training and fine-tuning large language models (LLMs) has primarily focused on reducing memory usage for constrained hardware lowering computational cost (c.g.. I 40 3).
In this paper, we take a step toward accelerating LLM optimization by showing that applying exact orthogonalization (e-g., via SVD) to the first-order moment offers practical advantage, even over the most accurate approximations, such as the commonly used Newton-Schulz5 method. Specifically, we find that SVD converges faster and incurs lower computational overhead. To support this, we first present a new observation: the moment matrix in LLM training tends to decrease in rank over time. Building on this, we derive an upper bound on the approximation error of Newton- Schulz5, showing that it depends on both the number of iterations and the matrix condition number, thereby highlighting its limitations ill-conditioned or low-rank settings (precisely the case LLM optimization moments). This motivates the need for more accurate orthogonalization of moment matrices during LLM training. Applying SVD directly to full-sized layers generally impractical. The surprising result, however, is that when integrated into low-rank optimization scheme, the use of SVD becomes not only feasible but preferable. We conclude with convergence analysis of Muon optimization, which, under these conditions, converges significantly more slowly than the SVD-based alternative. To the best of our knowledge, our convergence analysis of Muon optimization is the first to avoid neglecting the error in the Newton-Schulte approximation [41] The proofs of all lemmas and theorems of this section are relegated to the Appendix A
Lemma 3.1 (Moment Becomes Low-Rank During Training). Let (()) E Raxm denote the first mo- ment of reversible layerinn jin a moment- moment-based optimization algorithm, updated according to M(t) 3, ((t11) + G(t), where G(t) is the gradient matrix at iteration I. Let ( = U(!) (() v() be the singular value decomposition (SVD) of M(O, and define the rank-r orthogonal projection matrix
as (r) = U.) 1:r) U( 1:r]T. Then the relative error of the best rank-one approximation,
(l( |M (1) 11
(1)
[M
satisfies (t)(t) S O(C-1) ) for some constant C >> 1.
Reversibble networks formally defined in Appendix B.1

---

The above result, in 1b. implies that M approaches its rank-one approximation (1) (1) the iteration number increases, namely, M becomes rank-one. The following Lemma 3.2 characterizes the impact of the moments' low-rank structure on the approximation error of the Newton- -Shtul5 orthogonalization.
Lemma 3.2 (Orthogonalization error E,). For matrix A Rm let 01 be the largest singular value of AA and om be the smallest (without the loss of generality, assume m S n.. Let S m be the largest index where o, > gr+1
on 2 O. Let K d by the condition number of AA Denote& & the error of Newton-Schullz after i iterations. Then we have
am
2'
| - vr (1 1
(2)
K
According to the lemma, the approximation enror grows exponentially with the condition number. Given the low-rank structure of the first-order moments, dimensinaionl optimization can mit- igate this error. Specifically, projecting the moment estimates (1() onto their dominant (small) r-dimensional subspace ensures that the squared moment M (1) is constructed using only the top squared eigenvalues. These dominant components significantly larger and exclude near-zero values, resulting in a substantially lower condition number compared to that of the full-rank squared moment matrix. This observation motivates the use of the Muon optimizer within low-rank optimization framework for ILMs, including 2D) reversible layers. Such an approach not only preserves the inherent memory efficiency of low-rank methods but also reduces approximation error during optimization, potentially leading to faster convergence and improved performance compared to full-dimensional training. However, we also empirically observe that the eigenvalues of the moment matrix decay gradually. As shown in Figure ! even when projecting onto the dominant subspace, the resulting matrix (+ composed of the top 16 squared eigenvalues, can still exhibit large condition number, thereby introducing negliiiig approximation error.
Condition Kumber Steps
Moment's Singular Values
800
0:00
600
0.0-
000
Condron
400
0.04
200
0
200
-00
500
800
1000
20
32
24
15
Sxop
(a) Condition number of the first-order moment vs.
(b) Illustration of the moment's 's singular value decay,
training step. The red line marks value 10.
taken arbitrarily at step 100.
Evidence of anisotropy and ill-conditioning in the first-order moment matrix as a function of the Galore steps of the Roberta-batt model [42 on the GLUE dataset RTE task 43 (a) condition number growth, (b) spectral decay of moment.
Figure I:
To comprehend the cumulative error of Newton-Schulz. orthogonalization each optimization step, we proceed to derive the convergence rate of the Moun optimization. To that end, we now provide some notations. Consider a neural network denoted as o(:: 0), which consists of L layers and is
parameterized by 0 A wih ..... W 1 w xydxxd Here, W; represents the weights tensor parameters associated with the i-th layer, for i E [T.). We denote the differential loss L. where, with slight abuse of notation, we write the training problem by minw L(W) Eg Eg [C((@ {))], the context refers to the weights of a certain layer. We use the Frobenius norm, denoted ||F, which is induced by the inner product (X. Y) = tr(X Y). Assume that the stochastic gradient v VL(W, E) is an unbiased estimator of the full gradient V VLWW. with variance bounded by o², i.e.,
B g) V2(W)] < o². Let E(!) ())) Newton Newton-Schulz(MMie Schuuz(((()))) denote the approximation error of the Newton-Schule (with i iteration) at time 4, where M(9 denotes the moment at iteration t.
5

---

Lemma 3.3 (Exact convergence rate of Muon). Consider the Muon optimizer update w.k. layer W E x defined by
( 3M + +(1 (1- 33)
O00 < V( + (i iterationn Newtoon-Schulttzppppxxiim
w 00
where M = U() v(0 denotes the singular value decomposition of M and E represents
the Schull- approximation ervor after iterations. Suppose the following:
The gradient V((W) is L-Lipschitz continuous.
*There exists 8 such that ||E = 81 | 1 5V for all t.
If we take 3=11 a with a min(
RT. ovT
1), T/E =1 V(10/11 (07011-8444 -)+ 2m82)TT.
VAR
and
(batch free convergence) than T
E/= E [ is bounded by
VRLmu22 + 48 282)
(RL) Vo(LL)
VT
VRLT
T1/4
1mm8
where R = ((W(O)) - L' If we take as an arbitrary constant, and B T, we have,
BIV s o
([ |RLmml
+ 45 + 26t)
VRL
1
VT
VT
73/2
VT - 4
Remark 3.4 (Comparison: slower convergence vs exact orthogonalization) When 5 = 0, indicating an absence of error, the convergence rate is aligned with the one derived in 41), Theorem 2.1, that is
Σ T. [IVCCW()]]] O
VnRL VT
T4/2
VT
This result overlooks the error associated with the Newton-Schulz5 approximation because it is based on theoretically exact method of orthogonalization.
Remark 3.5 (The impact of 8 the convergence rate). A reduction in is associated with an improvement in the convergence rate. Furthermore, it should be noted that influences the step size m,: larger 8 yields a smaller 7). providing an additional explanation for the convergence rate.
Remark 3.6 (The size of 5). We acknowledge that the findings of our analysis applicable only under the conditions specified in 1 - 4vnă > 0 < ' In scenarios where d > I the algorithm may fail to converge. To ensure that á remains sufficiently small, the Newton-Schulz3 method necessitates a substantial number iterations, consequently slowing down the convergence. Remark 3.7 (Speed-up by SVD vs Newton-Schulzs approximation). According to Lemma 3.2 these low-rank moments, which inherently possess exceptionally high k, result in an error expressed by (1 - e)²² concerning a remarkably small This situation necessitates numerous iterations for the Newton-Schultz method to converge. For example, if (1 - E) = 0.99 considered and Newton- Schultz5 is utilized with 5 iterations, the error would be 0.9932 0.725, relative to the norm of the moment, namely M.
4yn applies,
Recall that in the low-rank setting, accurately computing the pseudoinverse using singular value decomposition (SVD) is numerically advantageous and reasonably computationally affordable com- pared to iterative methods such as Newton-Schula For general matrix A E Rexr, the SVD provides a decomposition A = UEVT, with U E Raxn Σ E R and V E Rrrr The Moore- Penrose pscudoinverse is then calculated as A = VE requiring approximately 4nr + 8 floating-poin -pint operations (FLOPs) for the initial decomposition, and an additional rn? + r2n FLOPs for subsequent multiplications, totaling roughly 4nr - 8 + rn² + r"n FLOPs.
Alternatively, approximating the inverse of A A E RR" using Newton-Sedee Schulz iterations involves 2r FLOPs form the matrix A T A, approximately 20 10r FLOPs for five iterations, and an additional nm FLOPs to multiply by A resulting in a total of about nr +r +10,2 FLOPs. For example, when the rank is r and n 1024, the SVD approach requires approximately twice as many operations as and inherent optimality of the SVD-based method, the increase computational effort is negligible relative to the other full-matrix operations required in each update step.
Nonetheless, given the superior numerical stability

---

3.2 Method
We are now ready to present our main algorithm designed accelerate the low-rank optimization scheme outlined in Algorithm ! A detailed mathematical formulation of the weight update rule proposed in this paper given in Appendix [ The algorithm consists of four primary blocks, all contained within an outer loop that continues until convergence is achieved or predefined number of epochs is reached. Each block serves a specific purpose, which will be explained in detail below.
Block 1 (Adaptive Subspace Selection): We select the subspace along the directions of the largest eigenvectors, but since computing full SVD for large matrices is computationally intensive and seemory-demanding for producing a "good" proxy for the optimal low-rank approximation. It solves the optimization problem arg min
we leverage the Randomized-SVD by [4]] which is an efficient technique
Q G(t) and approximates the matrix G as G s3
Q G(t), that requires Ommrr + m²?) operations, instead (min(mm²² , m²n)) applied by SVD.
Block 1.1 (Moment Subspace Transformation): We transform the first-order moments evaluated during the low-rank optimization steps, which occur in Block 2, between the preceding and the newly updated subspace. This transformation is required because, as will be demonstrated later, the first moments of the gradients in Block are aligned with the previously projected subspace. Consequently, transformation is necessary to translate them from the former subspace to the current one.
Block 2 (Low-Rank Steepest Descent Optimization): Here we calculate the (steepest) optimiza-
tion step. SVD operation is adopted to solve exactly (M())() ) (() Let UEV = M()) be the singular value decomposition (SVD) of M we will have (M( - UVT, which orthogonalizes (()), Formally, the
Orthogonalization, { arg min ||O Alle either O²0 = Ior 00 = I].
O
Block 3: Rather than using standard gradient clipping, we adopt the Norm-growth Limiter (NL) introduced in (21]. which has been shown to slightly outperform traditional clipping techniques by better constraining the progression of gradient magnitudes. Specifically, the gradient update is
modified as follows,
|0(*)
y then (())
0.0
> ||0((=>)||, where the scalar 7 serves
|
as a growth threshold to regulate abrupt increases in gradient norm from one iteration to the next. We use 7 1.1, which empirically yields the best results.
Block 4 (Update Step in the Original Space). To better exploit already computed gradient information, we suggest to use the orthogonal term of the gradient that lies outside the low-rank
subspace spanned (()), namely ( (()) Q(!) Q(1) G Importantly, (!)) + does
Git)
interfere with the orthogonalized moment estimation (t), meaning it can be incorporated without compromising stability. Since G( is already computed and stored in each iteration, no extra memory is required. Furthermore, because Q(1 is of a low rank (typically rank 4, 8, or 16), the additional computationattins overhead is negligible. For efficient memory usage, instead of explicitly forming the perpendicular part G (1) we work only in the rank-r subspace spanned by Q (f) Practically, we use fact that
G (+)
Q (1)
(
+ Q ( ( = G( Q (G -0()) -0()))
GOOD
which utilizes the already calculated projected gradient G(0) Q(0) G(O, Accordingly, the pre-trained model parameters are updated along with weight decay,
w ((t-1) n (G Q ( (G o())) 17 W(-1)
To ensure stable training across parameter matrices of different shapes, we interpret the root mean square (RMS) magnitude of updates implicit layer-wise learning rate adaptation, following
7

---

the approach in [3)] By scaling updates by /ma n),, our method compensates for shape- induced magnitude differences, achieving consistent effective learning rates across layers, similar to adaptive optimizers like AdamW.
Algorithm 1 SUMO: Subspace-Aware Monent-Orthooggnalzzttooon Optimization
Input: weight matrix W R=** with m 2 n. Step size 7/, decay rate μ., weight decay rank subspace update frequency K, step clipping ratio 7/ Initialize: 4-+00 repeat
Block 1: Calculate low rank gradient projection.
Sample mini-batch B = ...... }1B1
Compute G <-
|6 - ow (0( 0),
if't mod K 0 then
Q(1) + Truncated. Randomized
BerrmniieddSD((CCCC SVD(G # Alternatively Liincated_SVDCGt
# Block 1.1: Moment subspaces transformation
Rexr +- Q ((1-1) if 2 1, else 0rxr
M RM if 2 1, else 0r*n
[1 order moment]
end if Alternatively criteria || S
G Q(0 (t))
# Block 2: Low-rank steepest-decent step (moment ortogonalization) ( - (M + (())
(
# Block (Optional):
if - then 0 o
|
||((=))]|
o
# Block 4: Update weight in original space.
W + ((+-1)
( Q (t) ( - (G()
17 W(+-1)
+ +++1 until convergence criteria met (e.g. epoch number, gradient norm return W7T)
|G S 5
Note that, for clarity, can assume, without loss of generality, that 2 In the opposite scenario, the projection matrix would multiply the gradient from the right-hand side.
Theorem 3.8 (Convergence SUMO). For a loss function C, and given architecture 4 suppose that the compositions of f = C (o(-) is B-smooth non-convex function that is bounded by some M E R.. Let (()) denote the gradient matrix w.rl. the j-th reversible layer ((()) at time N, for all
j j E [T.,, and Te, E N times are set by a convergence criterion (that is, [|G S Se). Then, there
10) | S. Namely,
exist C CRR and such that for all TN
T and Ts Σi Σ
Algorithm I achieves an e-critical point iee. (0 < e, for some & N, and any j E [L].
We emphasize that the convergence proof for Galore [T] addresses only optimization within fixed subspaces, ignoring dynamic updates. AdaRankGrad's proof [2] first established guarantees for the complete dynamic-subspace updates, yet both prior works simplified the inner steps as standard SGD. In contrast, SUMO's convergence proof explicitly considers the exact optimization steps without simplifications.
The proof of Theorem I can be found in Appendix A
To reduce memory consumption, Algorithm applies per-layer weight updates during backprop- agation, following recent works such as 45 This contrasts with conventional optimizers, which store full gradients and update weights afterward, potentially leading to inefficiencies. Details for post-hoc adapter extraction are discussed in Appendix

---

Accuracy over Optimization Steps
0.91
PApRedup 1.A
0.83
0.87
SUMDISVDI
Gelore
SUNDONenton Schulz2)
1000
2000
3000
4000
5000
600D
7000
Step
Figure 2: SUMO with SVD demonstrates superior convergence speed 1.6x faster), attaining comparable or higher accuracy than GaLore and SUMO with Newton-Schultz with significantly fewer optimization steps on QNLI.
Experiments
Fine-tuning on GLUE benchmark. Our model was evaluated using the GLUE benchmark 4 through the fine-tuning of the pre-trained Roberta-base model 42 across eight tasks. The comparative analysis includes full fine-tuning, LoRA, and GaLore methodologies, with the results enumerated Table 2 The reported metrics are the overall accuracy (matched and mismatched) for MNLI, Matthew correlation for CoLA, Pearson correlation for STS-B, TTS, Fl-score for MRPC, and accuracy for the remaining tasks. Evidently, our approach improves fine-tuning accuracy while requiring less training memory, using only single moment compared to GaLore. The experiments were carried out using the NVIDIA A100 GPU.
Table 2: Comparison of SUMO against state-of-the-art memory-efficient fine-tuning methods on the GLUE benchmark using the pre-trained RoBERTa-Base model. For comparison, we provide detailed
results for SUMO using both SVD and Newton-Schulz5 orthogonalizations (ablation study).
Model
Memory
CoLA
STS-B
MRPC
RTE
SST2
MNLI
QNLI
QQP
Full Fine-Tuning
747M
62.24
90.92
91.30
79.42
94.57
87.18
92.33
92.28
LoRA (rank=4)
257M
61.38
90.57
91.07
78.70
92.89
86.82
92.18
91.29
GaLore (rank=4)
253M
60.35
90.73
92.25
79.42
94.0
87.0
92.24
91.06
SUMO (Newton-Schulz5, rank=4)
90.82
92.43
79.36
94.17
86.92
92.26
91.27
SUMO (SVD, rank=4)
197M
62.3
91.04
93.5
81.07
87.34
93.26
91.68
LoRA (rank=8)
264M
61.83
90.80
91.90
79.06
93.46
86.94
92.25
91.22
GaLore (rank=8)
257M
60.06
90.82
92.0
79.78
94,38
87.17
92.2
91.11
SUMO (Newton-Schule5, rank=4)
198M
61.74
90.79
91.94
79.69
94.17
87.21
92.24
91.38
SUMO (rank=8)
61.7
93.7
81.37
94.82
87.58
93.67
91.72
Pre-training Llama C4 Dataset. To highlight the effectiveness of our method in pre-training, we pre-train Llama models following the evaluation protocol of [T] and compare performance with the state-of-the-ar method, in terms of perplexity and memory usage. Specifically, we compare SUMO's 'sss performance with sotte-of-theasee methods in terms of perplexity and memory efficiency. For this evaluation, we trained large Llama-based models on the C4 dataset, curated and extensive version of the Common Crawl web corpus [46]. This dataset is widely used for pre-training language models and developing word representations. To better reflect real-world pre-training scenarios, we trained non-repeating, large-scale dataset and scaled model sizes up to billion parameters. The results of these experiments are shown Table Experiments were conducted using an NVIDIA H200 GPU.
Few/Zero-shot reasoning and long-context generalization. To evaluate the performance of our method a complex reasoning task, we utilize the GSM8K dataset [47] to test systematic generalization. For these experiments, we used a batch size 32 and 10 epochs for fine-tuning. We present the performance result in Table training Phi-2 (2.7B) model [48] and in Table training
Also known e-stationary, see, e.g..

---

Table 3: Comparison of sfate-of-the-ar low-rank algorithms for pre-training Llama models of varying
sizes on the C4 dataset. The results are reported in terms of validation perplexity. As shown,
SUMO leads to improved performance with substantial memory reduction compared to leading
parameter-effiiient fine-tuning schemes.
Method
60M
130M
350M
IB
Full-Rank
34.06 (0.36G) 25.08 (0.76G) 18.80 (2.06G)
15.56(7.80G)
GaLore
34.88 (0.24G) 25.36 (0.52G) 18.95 (1.22G)
15.64(4.38G)
Low-Rank
78.18 (0.26G) 45.51 (0.54G) 37.41 (1.08G)
(32.53
LoRA
34.99 (0.36G) 33.92 (0.80G) 25.58 (1.76G) (9.21
ReLoRA
37.04 (0.36G) 29.37 (0.80G) 29.08 (1.76G) 18.33(6.17G)
SUMO
34.26 (0. (3)
24.87
(0.51G) 18.69
(1.16G)
14.68
(3.84G)
Training Tokens
1.IB
2.2B
6.4B
13.1B
T mmddel
128/256
766/768
256/1022
512/2048
Llama (3B) model 44
The results demonstrate that the proposed method significantly improves generalization out-of-distribution data. We used NVIDIA H200 GPU.
Table 4: Zero-shot evaluation on GSM8K dataset
Table 8-shot evaluation on GSM8K dataset (LLaMA, 3B).
(Phi-2, 2.7B).
Phi-2 (2.7B)
Rank
Accuracy (0-shot)
LLaMA (3B)
Rank
Accuracy (8-shot)
Base Model
64
15.16%
Base Model
64
17.93%
Galore
64
52.24%
Galore
64
74.9%
LoRA
64
42.8%
LoRA
64
68.3%
SUMO
64
54.13%
SUMO
64
76.7%
Additional experiments and ablation studies are presented in the Appendix
D
Zero-shot generalization. Following this, we extend our evaluation to diverse commonsense and reasoning tasks using zero-shot methods, and then conduct the exact evaluation protocol described in Table 4 of (3] (which suggests the Apollo optimizer). The task details are presented in Appendix [D. To ensure a fair comparison, we repeat the exact experimental setup used for AdamW, APOLLO, and APOLLO-Mini. The results are summarized in Table beblo..
Table 6: Zero-shot evaluation of Llama-350M models pretrained with sequence length 1024 across reasoning tasks (lower is better for perplexity; higher better otherwise).
Memory
Perplexity
Bool()
RTE
OBQA
ARC-E
ARC-C
PIQA
SclQ
MathQA
Avg.
AdanWW
1.37G
16.30
0.4911
0.4693
0.3688
0.5233
0.352
0.3729
0.2449
0.6534
0.609
0.2064
0.4272
APOLLO
0.34G
15.64
0.5373
0.4198
0.3850
0.4925
0.322
0.3788
0.2483
0.6681
0.624
0.2127
0.4406
APOLLO-ME Mini
0.15G
16.12
0.5376
0.1562
0.3707
0.5217
0.324
0.3758
0.2312
0.6638
0.619
0.2224
0.4374
SUMO
0.5313
0.321
0.3832
0.6709
0.2246
As shown, the SUMO-pretrained Llama-350M model achieves lower perplexity on average and consistently outperforms on downstream benchmarks.
5 Discussion
Our results highlight that exact moment orthogonalization within low-dimensional adaptive subspace significantly improves both convergence and stability in memory- LLM training. By avoiding the approximation errors of Newton-Schull. the proposed SUMO leverages the low-rank structure of gradients to enable accurate, spectral- aalinndd updates with minimal overhead.
Empirically, SUMO outperforms prior low-rank methods on both fine-tuning and pre-training tasks, achieving greater memory reduction than memory-effiiien benchmarks such as Galore. Our theoreti- cal analysis further confirms its superior convergence properties under practical conditions. These findings position SUMO as a simple yet effective alternative to approximate geometric optimizers. Future work may investigate parallel computations for orthogonalization, integrate quantization tech- niques, and assess the effectiveness of the method in knowledge editing [50] or domain generalization
51 52..
10

---

Acknowledgment
TT was supported by the Israel Science Foundation (No. 1940/23) and MOST (No. 0007091) grants.
OL was supported by the MOST grant No. 0007341.
References
[1] Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient IIm training by gradient low-rank projection, 2024.
[2] Ychonathan Refacl, Jonathan Svirsky, Boris Shustin, Wasim Hulcihel, and Ofir Lindenbaum. efffccient LLMs training and fine-tuning. In The Thirteenth International Conference on Learning Representations, 2025.
Adarankgrad: Adaptive gradient rank and moments for memory-efficient
[3] Hanqing Zhu, Zhenyu Zhang, Wenyan Cong, Xi Liu, Sem Park, Vikas Chandra, Bo Long, David Pan, Zhangyang Wang, and Jinwon Lee. Apollo: Sgd-like memory, adamw-level performance. arXiv preprint arXiv:24/2.0522 2024.
[4] Barak Battash, Lior Wolf, and Ofir Lindenbaum. Revisiting the noise model of stochastic gradient descent. In International Conference on Artificial Intelligence and Statistics, pages 4780-4788. PMLR, 2024.
[5] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization, 2018.
[6] Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. Soap: Improving and stabilizing shampoo using adam, 2025.
[7] Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cecista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024.
[8] Louis Cesista Franz. The Case for Muon, October 2024.
[9] Jeremy Bernstein and I Laker Newhouse. Old optimizer, new norm: An anthology. arXiv preprint arXiv:2409.203: 20322,, 2024.
[10] Dmitry Kovalev. Understanding gradient orthogonalization for deep learning via non-euclidean trust-region optimization, 2025.
[11] Hao Sun, Li Shen, Qihuang Zhong, Liang Ding, Shixiang Chen, Jingwei Sun, Jing Li, Guangzhong Sun, and Dacheng Tao. Adasam: Boosting sharpness-aware minimization with adaptive learning rate and momentum for training deep neural networks, 2023.
[12] Nicholas J Higham. Newton's method for the matrix square root. Mathematics of computation, 46(174):537-544 1986.
[13] Jingzhao Zhao, Frederik T. Schaefer, and Anima Anandkumar. Zero initialization: Initializing neural networks with only zeros and ones. Transactions on Machine Learning Research, 2022.
[14] Romain Cosson, Ali Jadbabaie, Anuran Makur, Armin Reisizadeh, and Devavrat Shah. Low- Rank Gradient Descent. IEEE Open Journal of Control Systems, 2023.
[15] Greg Yang, Jacob B. Simon, and Jeremy Bernstein. spectral condition for feature learning. arXiv preprint arXiv:2310. /78/3, 2023. arXii:2310.17813.
[16] M. Gooneratne, K. C. Sim, P. Zadrazil, A. Kabel, F. Beaufays, and G. Motta. Low-rank gradient approximation for memory-efficient on-device training of deep neural network. In 2020 IEEE International Conference Acoustics, Speech and Signal Processing, ICASSP 2020, Barcelona. Spain, May 4-8, 2020. IEEE, 2020.
[17] Shuang Huang. Brian D. Hoskins, Michael W. Daniels, Matthew D. Stiles, and George C. Adam. Low-Rank Gradient Descent for Memory-Ffficient Training of Deep In-Memory Arrays. ACM Journal Emerging Technologies in Computing Systems, 2023.

---

[18] Ionut- Vlad Modoranu, Alexander Kalinov, Ermin Kurtic, Erwin Frantar, and Dan Alistarh. Error Feedback Can Accurately Compress Preconditioners. ArXiv preprint arXiv:2306.060% 2023. arX.v:2306.06098.
[19] Yuandong Tian, Lantao Yu, Xinlei Chen, and Surya Ganguli. Understanding self-supervised learning with dual deep networks, 2021.
[20] Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models, 2021.
[21] Xi Chen, Kaituo Feng, Changsheng Li, Xunhao Lai, Xiangyu Yue, Ye Yuan, and Guoren Wang. Fira: Can we achieve full-rank training of Ilms under low-rank constraint? arXiv preprint arXiv:2410001623, 2024.
[22] Yongchang Hao, Yanshuai Cao, and Lili Mou. Flora: Low-rank adapters are secretly gradient compressors. ArXiv, abs/2402.03293, 2024.
[23] Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Yinyu Ye, Zhi- Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer leamiig rates to gain more. arXiv preprint arXiv::406. 16793, 2024.
[24] Thomas Robert, Mher Safaryan, Ionut-Vlat Modoranu, and Dan Alistarh. Ldadam: Adaptive optimization from low-dimensional gradient statistics, 2025.
[25] Sebastian Loeschcke, Mads Toftrup, Michael J. Kastoryano, Serge Belongie, and Vésteinn Snæebarmarson. Loqt: Low-rank adapters for quantized pretraining, 2024.
[26] Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Animashree Anandkumar. signsgd: Compressed optimisation for non-convex problems. In International Conference Machine Learning, pages 560-569. PMLR, 2018.
[27] Michael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, and Zhenxun Zhuang. Robustness to unbounded smoothness of generalized signsgd. Advances in neural information processing systems, 35:9955-9968 2022.
[28] Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, and Quoe V. Le. Symbolic discovery of optimization algorithms. NeurlPS, 2023.
[29] Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be. arXiv preprint arXiv:2304. 13600, 2023.
[30] Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33:15383-15393, 2020.
[31] Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song. James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arXiv:1904.00962. 2019.
[32] Zhirong Yang and Jorma Laaksonen. Principal whitened gradient for information geometry. Neural Networks, 21(2-3):232-24 2008
[33] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR (Poster), 2015.
[34] Dongseong Hwang. Fadam: Adam is a natural gradient optimizer using diagonal empirical fisher information. arXiv preprint arXiv:2405.12807 2024.
[35] Jeremy Bernstein and Laker Newhouse. Old optimizer, new norm: An anthology, 2024.
[36] Jeremy Bernstein and Laker Newhouse. Modular duality in deep learning, 2024. 2024.
12

---

[37] David E Carlson, Edo Collins, Ya-Ping Hsieh, Lawrence Carin, and Volkan Cevher. Precondi- tioned spectral descent for deep learning. Advances in neural information processing systems, 28, 2015.
[38] Mark Tuddenham, Adam Prügel- Prügel-Bennett, and Jonathan Hare. Orthogonalising gradients to speed up neural network optimisation. arXiv preprint arXiv:2202.07052, 2022.
[39] Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, Yanru Chen, Huabin Zheng, Yibo Liu, Shaowei Liu, Bohong Yin. Weiran He, Han Zhu, Yuzhi Wang, Jianzhou Wang. Mengnan Dong, Zheng Zhang, Yongsheng Kang, Hao Zhang, Xinran Xu, Yutao Zhang. Yuxin Wu, Xinyu Zhou, and Zhilin Yang. Muon is scalable for IIm training. 2025.
[40] Ychonathan Refael, Adam Hakim, Lev Greenberg, Tal Aviv, Satya Lokam, Ben Fishman, and Shachar Seidman. Slip: Securing Iims ip using weights decomposition, 2024.
[41] Jiaxiang Li and Mingyi Hong. A note on the convergence of muon and further. arXiv preprint arXiv:250..02900, 2025.
[42] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint 1907.11 2019.
[43] Alex Wang. Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Superglue: A stickier benchmark for general-purpose language understanding systems. In Advances in Neural Information Processing Systems, volume 32, 2019.
[44] Nathan Halko, Per-Gunnar Martinsson, and Joel A. Tropp. Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions, 2010.
[45] Kai Lv, Hang Yan, Qipeng Guo, Haijun Lv, and Xipeng Qiu. Adalomo: Low-memory optimiza- tion with adaptive learning rate, 2024.
[46] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21(140):1677 2020.
[47] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems, 2021. Arxiv, 2021.
[48] Mojan Javaheripi, Sébastien Bubeck, Marah Abdin, Jyoti Aneja, Sebastien Bubeck, Caio César Teodoro Mendes, Weizhu Chen, Allie Del Giorno, Ronen Eldan, Sivakanth Gopi, et al. Phi-2: The surprising power of small language models. Microsoft Research Blog, 1(3):3, 2023.
[49] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint (tt::2307.092883 2023.
[50] Amit Rozner, Barak Battash, Lior Wolf, and Ofir Lindenbaum. Knowledge editing in lan- guage models via adapted direct preference optimization. In Findings of the Association for Computational Linguistics: EMNLP 2024, pages 4761-4774, 2024.
[51] Yehonathan Refael, Iftach Arbel, Ofir Lindenbaum, and Tom Tirer. Lorenza: Enhancing generalization in low-rank gradient Ilm training via efficient zeroth-order adaptive sam, 2025.
[52] Amit Rozner, Barak Battash, Lior Wolf, and Ofir Lindenbaum. Domain samain-generalizable multiple- domain clustering. Transactions Machine Learning Research, 2023.
[53] Ashok Cutkosky and Harsh Mehta. Momentum improves normalized sgd, 2020.
[54] Kenji Kawaguchi. Deep learning without poor local minima. In D. Lee, M. Sugiyama, U. Luxburg, Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29, Curran Associates, Inc., 2016.
13

---

[55] Rik Koncel- -enzi--Kedziirski,, Subhro Roy, Aida Amini, Nate Kushman, and Hannanch Hajishirzi. MAWPS: A math word problem repository. In Proceedings of the 2015 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1152-1157, 2015.
[56] C. Clark, K. Lee, M.W. Chang. T. Kwiatkowski, M. Collins, and Toutanova. Boolq: Exploring the surpissing difficulty of natural yes/no questions. arXiv preprint
2019.
[57] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. In International Conference on Learning Reprexentations, 2018.
[58] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguisties, pages 4791 1790, 2019.
[59] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: adversarial winograd schema challenge scale. Communications of the ACM, 64(9):99-106, 2021.
[60] Todor Mihaylov, Peter Clark, Oyvind Tafjord, and Tushar Khot. Can a suit of armor conduct electricity? new dataset for open book question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 2381-2391, 2018.
[61] Peter Clark and Oren Etzioni. Think you have solved question answering? try are, the ai2 reasoning challenge. In arXiv preprint arXiv:1803.055453 2018.
[62] Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 7432-7439, 2020.
[63] Matt Gardner Johannes Welbl, Nelson Liu, Crowdsourcing multiple choice science questions, 2017.
[64] Aida Amini, Saadia Gabriel, Peter Lin, Rik Concel-Kedzionske Yejin Choi, and Hannaneh Hajishirzi. Mathqa: Towards interpretable math word problem solving with operation- based formalisms, 2019.
14

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: we propose SUMO (Subspace-Awar Moment- oramenn-Ortooonlizzaiionnn an opti- mizer that employs exact singular value decomposition (SVD) for moment orthogonalization within dynamically adapted low-dimensional subspace, enabling norm-inducing steepest descent optimization steps. By explicitly aligning optimization steps with the spectral characteristics of the loss landscape, SUMO effectively mitigates approximation errors associated with commonly used methods like Newton-Schulz orthogonalization approxima- tion (claim. 3.3). We theoretically establish an upper bound on these approximation errors (claim. 3.2 proving their dependence the condition numbers of moments, conditions we analy amalytically demonstrate are encountered during LLM training. Furthermore, we both theoretically and empirically illustrate that exact orthogonalization via SVD substantially im- proves convergence rates while reducing overall complexity. Empirical evaluations confirm that SUMO accelerates convergence, enhances stability, improves performance, and reduces memory requirements by up to 20% compared to state-of-the oaf-the-ad art methods. (Figure. Table.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Our methods work when the layers are reversible, and improves other methods especially when the moments are ill-conditioned.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
15

---

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: all assumptions of theoretical results are presented in the paper, and all proof are in the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: Everything needed for reproducibility is presented in Section
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
16

---

(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e-g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: We provide experimental details in the main text and appendix. A GitHub link well be published soon.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips..om
public/guides oowes//ooooouuussisooooooooooeeo mubblc//enieeeeoooeettbisssss..s
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
//nips.cc http/uullcgggguugee
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permirted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer:[Yee [es]
Justification: All training and test details are presented in Section
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
17

---

Answer: [Yes]
Justification:
All detailed in Section 4
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by emoor bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
All detailed in Section 4
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
The research conforms with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special cireumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
18

---

Question:
Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer[[NA]
The paper talks about optimizers. does not have societal impact. Guidelines:
Justification:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (c.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deploymeett of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer[[NA]
Justification:
The paper poses such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We followed and gave proper credits for our use of data and models in Section
Guidelines:
19

---

The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/code/mooel as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsoureing nor research with human objects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing research with human objects.
20

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does involve LLMs as any important, original, or non standaadd components.
Please refer to our LLM policy for what should or should not be described.
21

---

A Proofs of Section 3
In this section, we prove all the theorems and results of Section 3
We aim to show that the gradient G becomes approximately rank-one ex- ponentially fast, then the exponentially weighted moving average of the gradients (i.e., the momentum ( also exhibits exponential decay of higher-rank components.
Proof of Lemma 3.7
Consider the singular value decomposition of the gradient G U T at iteration t. For all natural numbers r < m., we define H (X)** (r) = U[: 1 r To enhance notation clarity, denote (r) = ((((r) (() (r),, where (r) represents an orthogonal projection matrix, sat- isfying the conditions P(") (r) (r). and (() = P(t) (r). Without compromising generality, it is assumed that at f 0, the rank of ( is characterized by rank (G r. For reversible networks, it has been established in [heoorem 3.2] that the gradients assume the form GO = & Σ (A, B, C(). characterized by constant matrices {A,): and positive semi- definite (PSD) matrices (B,, C. for i 2 to. where E N holds. It is pertinent to recall that the vanilla weight update can be represented as W(t) = ( + (/ Let S - + C, & B; and λ < 22 denote its two smallest distinct eigenvalues. To substantiate our findings, we utilize several results and arguments presented in the proof of Lemma 3.3 in (.. Specifically, consider (t)) as the projection of (too onto the minimal eigenspace V of S corresponding to X1 According to our assumption, the rank of (tto is L, and its singular value decomposition (SVD) is given by (tt) = Ek= qæzy! where { and {yr)/= are orthonormal unit vectors, and {a) are the corresponding singular values. Therefore, as per Lemma 3.3 in [1. the gradient can be decomposed
into,
G
| S (1 (1) |1 1 + - -
9n
il where 9o is the projection of ( onto the minimal eigenspace V of S= n Σ C. 00 B, and o6 is orthogonal to V1. Here, a x2 are the smallest distinct eigenvalues of S.
We now unroll the momentum update, ( = E._ G Substitute the decomposition of G
Σ gt-s (1 m)r1)'al + (1 -1=)))''1
= Σ gi (1 m/11)"
)' 90 $1-=(111))))))
Let us define a A E'= g1= -71)" b, A E'= gr - n1)))² nt so that |M 12
a.gl + b. Now, compute the squared Frobenius norm:
I = ||a.9l + bell a?1l - + 20 (06- bb) + [1bel1
Since 90 T 90 and be lies in the span of gb. we have (gb. by) = O. thus.
| = at llad 13: + 11b.112-
Likewise, the spectral norm || af @1l.ab112 Hence, the ratio
(M (1M
M P..
M 12
a + 1|b 112 afllob 12
Kmm (1)
|M
M
a?llo!! P.2 + 116/112.
Using that 11980 13 of and the decay bound Ω ((max nd1 )2) we conclude:
|bb, O s722))29), while of
22

---

m - 17221 1 max 1//11
) O(C-9)
Rm (t) 50
for some constant C
Before proving Lemma 3.3 we shortly present the following two preliminary lemmas. To that end,
we present the following notations,
M The moment in iteration f. Its dimensions
m, where
|| || The Frobenius norm: ||||| || VAAT
L A stationary point to which the loss L converges.
B Batch size.
For A E IRmx and B E Rexn we denote A A B by AB for convenience.
Additionally, note that our proof is based equivalent but slightly modified formulation of moment's update. Specifically, instead of using the standard formulation of the moment's update
( = MM + (
we consider the convex combination,
( BM + (1 B)
This alternative formulation simplifies the analysis, but equivalent. To show that, we point out that can choose modefied learning step 1j" 1 0 get the same weight's updating step.
1) Orth
(MM +1 ()
Orth
1/
(M
+ -1-11)) 3)
3
Orth
" -I-B"
M +7)((())
1/rth
3
M(1) + G(
where Orth is the SVD orthogonalization step, formaly solving
arg min min [ -A || either O = or 00" =I].
Obviously, 3 0 could be chosen in way that M would result in any required positive real number.
assume following assumptions throughout our proofs:
(A1) The gradient VC(W) is L-Lipschitz continuous.
(A2) VC(W, <) is an unbiased estimator of VL(W) where L(W, is the gradient of L(W) when taking single training sample &.
(A3) E VL { VL(W) < o².
(A4) There exists 8 > 0 such that ||E|51| o| (() = 6.m for all t.
Lemma A.1 (Descent Lemma with Newton-Schull Approximation Error), Consider the Muon optimizer update defined by
9 + (--1) + - "))
0 + vl) +
(Newton-Schulz, iteration approximation),
w(+1) - w(t) ((0)
23

---

where M ( is the singular value decomposition of M and 3. represents the Newton-Schute (5 iterations) approximation error. Additionally, assume (A1) (A4). Then the following holds:
((Wat))) 8
((Wio)
(% ml [V + + (Vc(w()) MO
DELL
mL + 1/ Lm8*
Proof. Since L is L-lipschitz function, the descent lemma holds. Thus we have
((w(tt))) & ((w(n) + (VL(W"))) with witt) + - [/ -
- UN E Life | + 1
- (VCC( VN -
+ LuE (n+ ( T E + 112)
2 ((w())- Th (VL( II nh /V(((())- 11 nLmL
2
+ # 11 LoT 15 11+
Lni
|1
2
S((W)) _ VL( 7h 3/5 [V
M + ne mL
2
Lai
mmvmm VL Lid mm
8²n
2
= =(((()) - 2 mp |VL(W()) + th
5
/V((())) -M
nimL
2
1/ ni Linz
2
Where in (+) we used
53
cquation 2.8.
Lemma A.2. For constant 7/- = n > O, the following holds
71 41l md
mL
Tn² Lm82
((W(1)) - C + m
MOI |
+ Trp mLd +
2
2
Proof.
Using Lemma A.1
isolating
|V /V((())) and summing over all steps
T
41l
md
VL(W())]
4
1=1
5
mLL
Trj Lmdd
((W(1)) L + M
(C(((0)-MM')
T mL5
Proof of Lemma
3.3
The proof follows 53 Using the notations as
53
denote 5,()
M(t) VLWW (),
700 ((t) VL(W() and S(X, Y) VL(X) V L L(Y). Note that
have the following
(6 0 from A(2).
(0 < e from A(3).
24

---

E((>/9, = 0, Vi + j since 710) and 700 are independent.
|S(X, Y)II S L X Y|| from A(I).
Now following the update in (2), we get
g(++1) B9(0) + (1 8)7 (9) + S(X(0, ,(x((xx((t)))
3'8(0 + (1- - 3) Σ B gTy((-1)
Σ
gr S(x(1)) (0(-),,(((+11)))) x((+1-1))
r=0
therefore
t-1
t-1
||t+1211 S 8111; + (1 - 3)
Σ
B (1-T)
+ t Σ g.
0=0
Taking expectation get (using the fact that 81 = 81):
g2h 0²
E g
(1 - B)
Σ
+ 1L
Σ gT
B
1=0
r=0
a
1- B
+ n 1- 8
m
m
VI - g2
0
V
n
1- 3
m
m
All all, we get
(
+
3
TTL
(1-3) 3)B
m
1 - 3
Using Lemma A.2
we get
1/: 4L / má
T
VL((()))] S
4
((W(1)) - 1/
5
|L(W(())
M
m2'LT
+ Tn 2L6 +
Tnp Lmn²
2
2
Dividing both sides by 2-4nLymb we get
/LC(W())
A(C(W()) -2)
10 T
(V(((() -M 2Tm + 47 + Lm
7
/mo
[4R
+10 TTI 10% 11- 6
+10
+ 21m 4T 21 nmm
1 (1 (1-3)m
m
1-3 3
By taking 17-
110/(1-814200 2m 44m6 2mm we get
48
4 AA/RTL10/111199 9) 2m 4md + 2mm2) + (1 3)m
10g
+ 10T /1 - (1 - 9-mm 4/mo"
0
Now have two types of parameter choice. If we take B 1 (batch size free), we need to take B that have
min(1, o/T
VNL
25

---

]L((W))||
2/ RT + 4md + 2md²) + 22110 a (RL)
+ 100
T
+ (OL(RR))
RL
4 mmm""
T E [VV[(W))]]]
T Σ
so (1
/ + 48 28a)
(RL) 1/4
(o(RL)
VT
71/4
VRLT
T1/4
1
If we take 8 as an arbitrary constant in (0, 1), then we will need to take B = T, so that
[L(W()) (
V + 48 + 282)
RL
VT
VT
T4/2
VT
1 1/m8
Proof of Lemma 3.2 We denote B = AAT B
T X2 the result after Newton-Schultz iterations, Xo
and O UV
where B UEV is the SVD decomposition with 01 2 02 2 0m the
[B||
singular values.
known that Newton-Schultz converges quadratically, so we have
|& ||2 = X - O112 S
X
We now bound XX, O 2/- We know that Xo
UΣV BB
UEV
a
X, O|12 = UEVT UVT
Σ -
Σ
= U 1v
=
2
Where last equality due to the fact that || |l2 is unitary invariant. The matrix Σ with values Øe
is diagonal
1 on the diagonal. From that observation we get that
Σ
σi = max (1 (1- Ti
mm
1
max
01
01
k
For the Frobenius norm, get similar analysis. X,
O F
E
I
since -| is unitary
19
invariant, so we just need calculate all in all we have
I F
is known that
61
Σ
S /r
Σ
so
F
01
2
Σ 1
Σ
11-1
<
= vF
Proof of Theorem 3.8 for any layer E [Z] in the following, for simplicity of notation, ignore the index j and use (t) instead. By Lemma A.3 the low-rank optimization block 1 in Algorithm ] guaranteed to converge; we denote by T, E N the time index / at which we enter block 1 for the 4th
time (i.e.. |G, (T) < <2), for E N. Furthermore, we recall that Gf & Ww,
(o()):: when clear from the context, we omit j from W,, and use instead Vw,f (0 vf (W(t). Consider the
SVD decomposition the gradient W (0T,) = UT TD) T For t [T,, T - 1]. we define the projected gradient as (1) (r())) where p(T)(r) = U(T.) |.: r] ussin the exact truncated-SVi calculation (in Block 1). For simplicity, refer to (T) as (T,), and
26

---

denote P(T.) QT.)) Q(T.), Next, let he - f(W(o) f(W(T:-1)). and T7. denote the learning rate. Then,
h+1 f(w(ctl) s(W((()))
= / (w() m (wat)) -w())) - ( (W(T...)
$ f(Wlo) - (Wi)H)))
7/ ,0(0)
Ank (0()) 113
f(w()) f (WiT.
Tt |
3 n
+ 2
mm
h4 " ( ||3
n mm
Brp
where (1) follows Eq] 7 and (2) follows Lemma
A.3
Thus, summing over t [Ta, Tı41), we get
T-
Σ
(h141 - h4) s
T-n-1
Σ
1t
(1)
5
T
+
Bnfn
2
v mm
2
T=
Tin--
T-n--
5
1
3nfn
hT S hT,
+
(3)
4
2
v m
Assume a constant learning rate = 17. and define T: := Tit1 have,
T, For each interval [Tg, Ta+1 -
1).
Tinl-
55 110
B1221
G < kT, NT. + Ti
2
v m
Summing over all I,...,N, N,
Σ
n
T1411
(4)
1/ S Σ
N
(hr. - hT,.)
+
Σ
N
T
5
n1
+
Bifn
2
m
2
E
i=1
1=1
5
110
3122
hT.
ITNIN +T
Vm
2
where T = Em Te is the total number of iterations. Using hT, hT I M, we get,
Tan-1
4M <
50
B3
+
2T
+ 3mm
Vm
Let
2M V 3n
Then, the average squared norm of the projected gradients satisfies,
1
4M
50
G
+
3mm
T
T
m
4M
50
2M
+ 2
+ Bn
(V
2M
T
\m
Bn
VT
Bm
VT)
/M
50
2M
+ 2.
+ 2,3
V2M
73/2
Vm
Bn
VT
4/2MBn
10g
2 2M Bn
T3/2
m
VT
27

---

Thus, the bound becomes,
T Σ0 G
10
4v2M3n
10o vm
+
2/2M
2M Bn
T3/2
VT
Recall that for P(T.), it holds that for some E (0, 1]. Then,
-QTT)
S <a]G(T))>
v a T
Q 1T
From Lemma B.3 in II, under 1j S Amax"
we get,
E
WE [T. T..)) (1) || < E
Hence,
T Σ
1 N-1
(T411- T, G
TN
N-1
$ (1 TN
Σ
-
(TH11
1
i=0
1
422MMB
100
2 Bn
(1 TNTN
TN2
Vm
VTN
4/2MBn
10o
2/2M
(4)
(1
TN
(1 [2
Accordingly, for any > 0, if Tn satisfies
4 Bn
10c
2/2MBB
S E, (1-a)T
(1 -1-a0 TT
(1
a sufficient condition for this hold is,
{ { (12 2/5
/6/2MBm 2/3
30g
TN 2 max
a)-a)m/mi
(1 -1-
follows that
N-1 T
min 0StSTN
|
| B Se.
Tn
i=0
Thus, there exists an iteration index E [0, Tw] such that
G
| =,
which, by definition, implies that the algorithm reaches e-critical point.
This concludes that Algorithm achievess an e-critical point.
In the following, provide the auxiliary lemma that is used in the proof of Theorem 3.8
Lemma A.3 (Convergence of the Inner Fixed Low-Rank Optimization). Consider the same setting and assumptions as Theorem 3.8 Then, the second time = T E in which Algorithm 7 enters Block 1 (where it updates the projection matrix) happens for finite E e N.
28

---

Proof. By the B-smoothness of f, we have
f(W(+))) - f(W())+++))) + (G w(( w + 8
-w
2
Substituting the update rule W(++)) = W 00, we get
f(w(+1)) s s(((() n ( (0) + Buh
110 (2) 113.
(5)
2
Since (( = (p(()(r) G and O E (Pnge(P(())) (r) ) it holds that (G (0) = (G )
By equation (2.8) in 41 we have
1 1 + 5 | - 1
- .00) -
4
2
Now bound || |(
|| M |F ||( V -M
< || V]((")))l + - -M
(by assumptions 1)-(2)).
Substituting this into the previous inequality gives
5
(1)
G
12
(6)
2
mm
Substituting into equation
we obtain
f(((+1) S f(w())) - T/t 1 |/
5.
n10
BTg 2
|000
m
2
Since by definition |000 < get
f(w(L-)) - f(w())-
Th (
5
nho
Bn,n
+
+
4
Vm
2
For constant step size 7/2 = 1/. summing over 1 to T. get
((W < s(W()))
5
noT
BTFnT
+
2
m
2
Rearranging and using f(w())- ff" S M, conclude
_
1
Σ
( 13+
4M
10
+
2m
T
nT
B Additional Information
Definition B.1. (Reversibility [9 A neural network ø that maps the input æ to output y o(x;; 8) reversible, there exists L(x; 8) so that y = L(x; f)x, and the backpropagated gradient 9n. satisfies 9x I,T (ax; 0)gy, where 9y is the backpropagated gradient at the output y. T.(x;;))
depends on the input æ and weight 0 in the network o.
29

---

Several critical observations regarding Algorithm warrant attention. Initially, in order to minimize memory consumption, Algorithm| imppleeents a per-layer weight update during the process of backpropagation, advocated by contemporary studies, see, e.g., 45 This approach contrasts with conventional optimizers, which typically update all weights after backpropagation by retaining the complete gradients in memory. method potentially marked by significant inefficiency. Should there be a desire to generate an adapter (i.e., a parallel low-dimensional LoRA-type model) subsequent to fine-tuning, this can be achieved with efficiency through the following steps. Firstly, the training
weights gap A W Tuned W Pretmined is computed, where Fune-Tuned denotes the model weight upon process completion, and Pretrained refers to the original model weight. Subsequently, ''Adaptor _ rank(A) is determined utilizing a matrix ranking algorithm, followed by the resolution of min AER"**AMp BBRrk || ABI through any optimization algorithm (e.g., gradient descent). It is noteworthy that any solution to this matrix factorization optimization problem is well-known as global optimum [54].
C Update Step Rule Formulation
Definition C.I. [Subspace-Awart Momenn-Orthogonalizaali (SUMO)) SUMO formulates the subsequent gradient update rules. Refer
1 = Q o00 (0) Vw (W..))
M AM + - 3)
o (++1) Orthogonalliztionn (MM((M())) (M
SUMO
W w
with Q E Ruxr and R, E R2 denoting projection matrices, T E N representing the subspace update interval, 17 indicating the learning rate, & constituting stochastic batch, and Orthogonalization SVD(A) as the operator that resolves the following through Singular Value De- composition (SVD), described in
arg min O
A FF either O lor 00" =1}.
D Additional Experiments
In Table 7 we evaluated SUMO and state-of-the-ar memory-efficient fine-tuning methods on the MAWPS[55) dataset using the LLaMA2-7B model. We report results across two rank settings (32 and 128), comparing training time, memory usage, and task accuracy. SUMO consistently achieves superior accuracy while maintaining competitive efficiency in both memory and time (comparing to Galore).
Table Fine-tuning LLaMA2-7B on MAWPS
Methods
Rank
Time(h) J
Memory (GB) 1
Accuracy (%)
LoRA
32
0.40
14.36
45,80
DoRA
32
0.69
15.01
44.96
GaLore
32
2.59
15.15
58.40
SUMO
32
1.83
13.86
58.47
SUMO (SVD)
32
1.56
13.86
61.23
LoRA
128
0.45
15.64
65.97
DoRA
128
0.72
16.17
66.81
GaLore
128
2.61
15.79
64.29
SUMO (Newton-Shultz5)
128
1.78
14.12
64.41
SUMO (SVD)
128
1.62
14.12
68.03
30

---

Comparison with Muon. Table 8 below complements Table 2 by comparing full fine-ture tuning.
vanilla Muon, and our SUMO; results are reported mean standard deviation when available.
Table 8: Additional comparison to Table 2: full fine-tuning vs. vanilla Muon and our SUMO (mean std when reported).
Model
Memory
CoL.A
STS-B
MRPC
RTF
SST-2
Full Fine-Tuning
747M
62.24
90.92
91.30
79.42
94.57
Muon Full bine Tuning
458M
61.19
90,98
92.14
80.83
94.71
SUMO (Newton Schulz5, rank=4)
197M
61.81
0.02
90.81 0.013
92.43 1. 0.034
79.33 1. 0.031
94.14 0.028
SUMO (SV,, rankw4)
197M
62.32 0.015
91.05
0.007
93.48
0.022
81.08 +: 0.019
94.33 0.01
SUMO (Newton-Schulz5, rank=8)
198M
61.73 + 0.021
90.77 = 0.032
91.93 0.04
79.66 0.03
94.13 0.025
SUMO (SVD, man=8)
198M
61.69 0.014
91.11 + 0.02
93.72 0.018
81.38 + 0.011
94.83 0.01
These results show that our SUMO, achieves better performance with a significantly smaller memory footprint compared to Muon full fine-tuning approach.
Hyperparameters Grid search. To evaluate the impact of Subspace Update Frequency (K) and Ranks (r), we performed grid search during the pretraining of the LLLMA 130M model on the C4 dataset. This specifie setup allows for direct comparison with the Galore method. Table: Perplexity results from grid-search of Subspace Update Frequency (K) and Ranks (r) for the LLaMA 130M model pretrained on the C4 dataset. Values presented Galore/SUMO.
Table 9: Perplexity from grid-search over Subspace Update Frequency (K) and Rank (r) for LLaMA-130M on C4. Values are Galore/SUMO
Update Frequency
Rank 128
Rank 256
Rank 512
100
29.7/28.27
27.9/26.74
27.4/26.73
250
18.1/27.86
26.5/24.87
26.2/24.82
500
27.2/25.91
25.6/24.98
25.3/24.31
lk
26.8/25.83
25.1/25.42
24.8/24.93
D.1 Details of benchmarks in Table
Specifically, the pretrained models is evaluated on the following tasks:
Perplexity: Measured on the C4 dataset 46)
Commonsense Reasoning: BoolQ [6
RTE (5]. HellaSwag (HS) (58] Winogrande (WG)
[59], OpenBookQA (OBQA) 66 ARC- Easy (ARC-E), and ARC-Challenge (ARC-C)
611
Physical and Scientific Reasoning: PIQA
62 SciQ 63 and MathQA 64
D.2 Details of Fine-Tuning on GLUE
We fine-tune the pre-trained RoBERTa Base model the GLUE benchmark using the model provided by the Hugging Face. Table I0 and, we detail the hyper parameters used in fine-tuning.
MNLI
SST-2
MRPC
CoLA
QNLI
QQP
RTE
STS-B B
Batch Size
16
16
16
32
16
16
16
16
# Epochs
30
30
30
30
30
30
30
30
Learning Rate
IE-05
1E-05
3E-05
3E-05
IE-05
IE-05
1E-05
IE-05
Rank Contig.
r
Projection back scale
4
Max Seq. Len.
512
Table 10: Hyperparameters of fine-tuning RoBERTa base for the comparison in Table 2 with respect
only to rank=4.
31

---

MNLI
SST-2
MRPC
CoLA
QNLI
QQP
RTE
STS-B
Batch Size
16
16
16
32
16
16
16
16
# Epochs
30
30
30
30
30
30
30
Learning Rate
1E-05
2h-05
2E-05
1E-05
1E-05
2H-05
2F-05
3E-05
Rank Conlig.
r == 8
Projection back scale
Max Seq. Len.
5T2
Table 11: Hyperparameters of fine-tuning RoBERTa base for the comparison in Table 2 with respect
only to rank=8.
D.3 Comparison with vanilla Muon
Following we conducted a full fine-tumi linettuning experiment with vanilla Muon complete the comparison for Table in the paper, as presented below.
Model
Memory
CoLA
STS-B
MRPC
RTE
SST-2
Full aue-Tunnng
747M
62.24
90,92
91.30
79.42
94.57
Muon Full Fine Tuning
458M
61.19
90.98
92.14
80.83
94.71
SUMO (Newton-Setee Schuuii,, rank=4)
197M
61.81 0.00
90.81 +: 0.013
92.43
0.034
79.33 +: 0.031
94.11 = 0.028
SUMO (SVD, rank=4)
197M
62.32 1: 0.015
91.05 1: 0.007
93.48
0.022
81.08
0.019
94.93 1: 0.01
SUMO Sewttnn-chull5,, rankw8)
198M
61.73 + 0.021
90.77 + 0.032
91.93 0.04
79.66
0.03
94.13 0.025
SUMO (SVD, rank-8
198M
61.69 0.011
91.11 + 0.02
93.72 0.018
81.38 + 0.011
91.83 = 0.01
Table 12: Comparison with vanilla Muon fine-tuning results.
32


Reviews:
[{"review_id": "I16nY6xBTP", "reviewer": "Reviewer_FdsP", "summary": "Based on the observation that the gradient rank decreases during the training, the paper proposes a novel algorithm that can achieve both memory efficiency and fast convergence rate. The mathematics is solid, but the results are somehow unclear and weak.", "questions": "1. in table 3, why the memory efficiency doesn't increase from 350M model to 1B model?\n2. in the introduction and in the section 3.1, the authors claims both reference [2] and themselves found that the gradient rank decreases through the training, is [2] from the authors?\n3. SVD of a large matrix will take a huge amount of computation, when the model size increases. Can you actually apply this method to the frontier models, like GPT 4o?\n4. There are many assumptions in the theoretical analysis, how reliable are these assumptions in a real application.", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths: \n1. the insight and motivation of this paper are very insightful and interesting, the algorithms are simple but clear and correct. \n2. The mathematics of the analysis and derivations is good, easy to follow and understand. The remarks are particularly interesting and insightful. \n\nWeaknesses: \n1. The paper claims \"substantially improves convergence rates\", but there is no direct evidence in the results section demonstrates the convergence rate, although various final results are better than the baselines. \n2. It also claims \"SUMO effectively mitigates approximation errors\", but only verifies it in the mathematics analysis, but not in the results or ablation study. It would be wonderful to demonstrate this point in an ablation study. I believe this is actually a missed opportunity. \n3. The Experiments and results section is weak and unclear, particularly tables, probably the authors have no space to explain them. It makes the paper looks like not ready. \n4. Why there are 4 different models are used in experiments? Besides the size of the model, are there any other reasons, such as why zero-shot using Phi-2 and 8-shot using LLaMA? why not use both in both tests?", "was_revised": true, "final_justification": "I appreciate the authors addressed most of my questions and argued about my concerns. They accepted that there is an important typo in their original manuscript, and their new results may lead to more concerns, all of these makes me worry the result section of their paper is not mature and ready. Although I have this concern but I like the insight and motivation in this paper, so  I keep my score.", "rebuttal": "We sincerely thank the reviewer for their positive and encouraging feedback on our work. We are glad they found the motivation insightful, the algorithms clear and correct, and the mathematical analysis easy to follow. We appreciate the suggestions regarding the empirical validation of our claims and agree that clearer experiments and ablations would strengthen the work. We address these points in detail below.\n\n**W1 convergence rate**: We want to clarify that Figure 2 in Section 4 addresses this point. The figure presents a convergence plot on the QNLI task, where we directly compare the optimization steps required for SUMO (SVD), SUMO (Newton-Schulz5), and GaLore. The results show that SUMO with SVD converges approximately 1.6× faster than both baselines, reaching similar or higher final accuracy in significantly fewer steps. This provides direct visual and quantitative evidence supporting our claim of substantially improved convergence rates. We will add similar learning rate figures for other Glue tasks to the appendix. \n\n**W2 mathematical and empirical evidence for mitigation of the approximation error**: We thank the reviewer for the helpful comment regarding empirical support for SUMO’s effectiveness in mitigating approximation errors. Below, we highlight and clarify points in the paper that demonstrate this point: \n\na. Mathematical Foundation (Equation 2):\n The approximation error associated with the Newton-Schulz orthogonalization, as derived in Formula (2), illustrates that this error is inherently dependent upon the condition number of the moment matrix. Given a fixed number of iterations (denoted by i), the approximation error bound grows exponentially with respect to the condition number: ‖E_i‖_F ≤ √r ( 1 - 1/κ )^(2i). \nThis establishes a direct and quantifiable relationship between the numerical accuracy of orthogonalization and the spectral condition of the matrix, making the mitigation of approximation errors a critical aspect of the optimization process.\n\nb. Visual Evidence (Figure 1): Empirically, the impact of this mathematical relationship is visually substantiated in Figure 1. The condition number of the first-order moment matrices during training is shown to increase significantly, especially in early optimization steps, underscoring scenarios wherein approximation errors would notably affect convergence performance when using Newton-Schulz orthogonalization. These conditions confirm the relevance of addressing approximation errors for practical training stability and efficiency.\n\nc. Ablation Study (Table 2): To concretely demonstrate this theoretical insight empirically, we conducted an ablation study explicitly comparing SUMO utilizing exact SVD orthogonalization against its counterpart employing Newton-Schulz5 orthogonalization across the GLUE benchmark tasks. As summarized in Table 2, the exact SVD variant consistently outperforms the Newton-Schulz5 variant across all GLUE tasks, despite identical optimization settings and model configurations.\n\nThis performance gap directly reflects the practical advantage gained by reducing approximation errors through exact SVD-based orthogonalization. For instance, the QNLI task clearly demonstrates accelerated convergence and higher accuracy for SUMO (SVD) over SUMO (Newton-Schulz5), which can be attributed explicitly to the superior numerical accuracy in optimization steps, mitigating errors arising from high condition numbers during moment orthogonalization.\n\nThrough this expanded explanation, we reinforce both mathematically and empirically the critical importance of using exact SVD orthogonalization within SUMO to effectively reduce approximation errors, thus achieving enhanced optimization performance.\n\n**W3 Elaboration on experiments section**: We thank the reviewer for this feedback regarding the clarity of the Experiments and Results section. We therefore kindly elaborate in the following. \n\nTable 2 (GLUE): SUMO consistently outperforms other memory-efficient methods while using less memory. It exceeds SOTA accuracy at ranks 4 and 8 on multiple tasks.\n\nTable 3 (Pretraining): On C4, SUMO improves perplexity by >0.5% and reduces memory to ~1.25 at the lowest.\n\nTables 4–5 (Reasoning): On GSM8K, SUMO achieves 2%+ higher accuracy in both 0- and 8-shot settings.\n\nThese results collectively demonstrate that SUMO improves memory efficiency, convergence, and performance across all settings. We will incorporate these explanations and additional clarifications in the final version.\n\nNew results: We extend our evaluation to diverse commonsense and reasoning tasks using zero-shot tasks, followed by the exact evaluation protocol described in Table 4 of Apollo (Zhu et al. SGD-like Memory, AdamW-level Performance). SUMO-pretrained models achieve lower perplexity and outperform AdamW on downstream benchmarks.\n\n**Table 7**: Zero-shot evaluation of LLaMA-350M models pretrained with sequence length 1024 across reasoning tasks\n|Method|Memory|Perplexity|BoolQ|RTE|HS|WG|OBQA|ARC-E|ARC-C|PIQA|SciQ|MathQA|Avg.|\n|-----------|------|---------|----------|----------|----------|----------|---------|----------|----------|----------|---------|----------|----------|\n|AdamW|1.37G|16.30|0.4917|0.4693|0.3688|0.5233|0.332|0.3729|0.2449|0.6534|0.609|0.2064|0.4272|\n|APOLLO|0.34G|15.64|0.5373|0.4698|0.3850|0.4925|0.322|0.3788|0.2483|0.6681|*0.624*|0.2127|0.4406|\n|APOLLO-Mini|0.15G|16.12|0.5376|0.4562|0.3707|0.5217|*0.324*|0.3758|0.2312|0.6638|0.619|0.2224|0.4374|\n|SUMO|0.18G|*15.49*|*0.5479*|*0.4709*|*0.3937*|*0.5313*|0.321|*0.3832*|*0.2496*|*0.6709*|0.623|*0.2246*|*0.4416*|\n\nIn addition to Table 3, we include a comparison of our results with state-of-the-art methods, AdamSN[1] and LDAdam. The comparison shows that while other methods have a slight advantage on smaller models, SUMO is highly effective at scale, achieving a perplexity of 14.68 on the 1B model (vs. AdamSN's 14.96) with superior memory efficiency (2.87G vs. 6.17G).\n\n **W4 experiment configuration**\n We thank the reviewer for this question. We use four models to evaluate our method under diverse and complementary scenarios (classification, regression, zero-shot, 8-shot) as follows:\n\nRoBERTa-base: for GLUE, to benchmark against well-established fine-tuning baselines.\nLLaMA (60M–1B): for pre-training experiments, to study memory-efficient scaling trends.\nPhi-2 (2.7B): for zero-shot GSM8K, because Phi-2 is optimized explicitly for reasoning-heavy tasks and provides a strong baseline for evaluating generalization without in-context examples.\nLLaMA (3B): for 8-shot GSM8K, since LLaMA is widely adopted for few-shot evaluation, where its strong in-context learning ability is well documented.\n\nImportantly, our method is model-agnostic and would behave identically when applied to either model in both regimes. The chosen setup follows standard experimental protocols commonly used for each task in prior and related work.\n\nIn addition, following we have expanded tables 4 and 5 from the paper to complete an 0/8-shot evaluation for both models.\nTable 4/5: Zero-shot and 8-shot evaluations on GSM8K dataset for Phi-2, 2.7B and LLaMA, 3B\n \n|Model|Accuracy(0-shot)|Accuracy(8-shot)|\n|-------------------|------------------|------------------|\n|Phi-2(2.7B)Base|15.16%|19.37%|\n|Phi-2(2.7B)GaLore|52.24%|71.48%|\n|Phi-2(2.7B)LoRA|42.8%|67.24%|\n|Phi-2(2.7B)SUMO|*54.13%*|72.35%|\n|LLaMA(3B)Base|11.52%|17.93%|\n|LLaMA(3B)GaLore|33.66%|74.9%|\n|LLaMA(3B)LoRA|30.78%|68.3%|\n|LLaMA(3B)SUMO|35.21%|76.7%|\n\nSUMO demonstrates the strongest generalization, outperforming other methods across both models in zero-shot and few-shot scenarios.\n\n**Q1 Typo Correction**. Thanks! The memory size is 2.87 GB.\n\n**Q2** Please note that the low-rank analysis here differs from the one in [2]. In [2], the authors considered the gradient, while here we consider the momentum.\n\n**Q3 Methodology Clarification**. A key contribution of SUMO is applying exact SVD within a low-dimensional subspace onto which gradients are projected, rather than on full-sized matrices. This is both practical and beneficial: exact SVD in this compact space captures the true dominant directions, improving gradient conditioning and accelerating convergence. As shown in Block 2 of Algorithm 1, this is feasible because (1) gradients in LLMs often exhibit low-rank structure, and (2) SUMO confines SVD to small matrices (e.g., rank 8 or 16), resulting in minimal overhead.\nAdditionally, as detailed in Section 3.2 and Algorithm 1 (Block 1), SUMO employs Randomized Truncated SVD, which differs from GaLore, to efficiently update the low-rank projection every few hundred steps. This optimization reduces the computational cost from O(mn²) to O(mnr + mr²), making the method more scalable. As shown in Table 3, SUMO effectively pretrains LLaMA models with up to 1 billion parameters, achieving lower perplexity and memory usage compared to competing methods. Although scaling to advanced models like GPT-4o is a future goal, SUMO's adaptive low-rank design positions it favorably for this type of goal.\n\n**Q4 Reliability of theoretical assumptions** The assumptions presented are valid in practical applications. For instance, in Lemma 3.3, we assume an L-Lipschitz gradient and a bounded Newton–Schulz approximation error, both of which are commonly made in related work. We also provide a bound on this error in Lemma 3.2 and show it can be controlled through sufficient iterations. Figure 1 illustrates that large condition numbers often arise in practice, supporting the practical relevance of this bound.\nThe smoothness and boundedness assumptions in Theorem 3.8 are similarly standard, appearing in works such as GaLore, AdaRankGrad, and Muon. While we acknowledge that these assumptions may not hold perfectly in all settings, even then we still observe strong empirical performance across multiple networks and benchmarks.\n\nWe hope this addresses all your concerns, and we would be happy to further clarify open issues"}, {"review_id": "iZyxZix7P7", "reviewer": "Reviewer_1DmJ", "summary": "The paper introduces SUMO (Subspace-Aware Moment-Orthogonalization), an optimizer that exploits the empirical observation that first-order momentum tensors in large-scale neural-network training quickly collapse to a very low-rank subspace. SUMO periodically (every K steps) identifies this subspace via randomized SVD, performs an exact orthogonalization of the momentum inside the reduced space, limits norm growth with a lightweight scaling gate, and then projects the update back to the full parameter space. The authors (i) prove that Newton–Schulz-style approximate orthogonalization leaves a non-negligible residual when the momentum matrix is ill-conditioned, (ii) provide convergence bounds that show strictly tighter guarantees when the SVD error δ → 0, and (iii) implement a practical algorithm that adds negligible wall-time overhead. Experiments on GLUE fine-tuning (RoBERTa-Base), LLaMA-style pre-training (60 M–1 B parameters), and GSM8k reasoning show consistent memory reductions (20–85 %), 1.2–1.6 × step-speedups, and modest accuracy/perplexity gains over AdamW, LoRA, GaLore, and Muon baselines.", "questions": "1. How sensitive is SUMO to the SVD period K and rank r? Could you provide a small grid-search on GLUE?\n2. What is the per-step wall-clock overhead (ms) of the randomized SVD and subspace rotation on A100/H100?\n3. Does the “low-rank momentum” observation hold for attention and MLP weight matrices individually, or only when concatenated?\n4. Can you share preliminary results on >7 B-parameter models or multi-node training?", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. It identifies and formalizes the low-rank momentum phenomenon, then tailors an orthogonalization scheme to it.\n2. It supplies error bounds linking SVD precision, condition number, and convergence.\n3. It uses randomized SVD and subspace recycling to keep runtime overhead low.\n4. It covers three model scales, three task families, and compares to four strong baselines.\n5. It delivers both memory savings and quality improvements without extra hyper-parameters.\n\nWeaknesses:\n1. Experiments has only single-seed runs; training token budgets differ across methods; no wall-clock timing.\n2. “invertible layer” and rank-1 convergence are not justified for attention/FFN layers.\n3. No study of subspace update period K, rank r, or γ (norm limiter).\n4. All experiments are single-GPU; no distributed data-parallel tests.", "was_revised": true, "final_justification": "Thanks for the author's efforts in addressing my concern. Since most of my concerns were well clarified, I would like to raise my score accordingly.", "rebuttal": "We thank the reviewer for the detailed and thoughtful review. We are glad that the core ideas, theoretical contributions, and practical benefits of SUMO were appreciated. We also acknowledge the valuable concerns raised about hyperparameter sensitivity, runtime metrics, and scalability, which we address in our response.\n\n**W1 adding std for different seeds run**. We thank the reviewer for the suggestion. While our initial experiments followed community-standard single-seed evaluation (e.g., on GLUE), we observed consistent performance trends across tasks and datasets. To further demonstrate stability, we conducted 10 runs on CoLA, STS-B, MRPC, RTE, SST-2, and MNLI, and 5 runs on QNLI and QQP, using seeds drawn from a Gaussian distribution. We now report standard deviations to support the robustness of our results.\n\n**Table 2** Extended: Evaluation comparison of SUMO against state-of-the-art memory-efficient fine-tuning methods on the GLUE benchmark using the pre-trained RoBERTa-Base model. For comparison,we provide detailed results for SUMO using both SVD and Newton-Schulz5 orthogonalizations (ablation study).The table shows that SUMO consistently outperforms the low-rank fine-tuning methods and is relatively stable across runs, as indicated by the small standard deviations.\n|Model|Memory|CoLA|STS-B|MRPC|RTE|SST-2|MNLI|QNLI|QQP|\n|---|---|---|---|---|---|---|---|---|---|\n|Full Fine-Tuning|747M|62.24|90.92|91.30|79.42|94.57|87.18|92.33|92.28|\n|LoRA (rank=4)|257M|61.38|90.57|91.07|78.70|92.89|86.82|92.18|91.29|\n|GaLore (rank=4)|253M|60.35|90.73|92.25|79.42|94.00|87.00|92.24|91.06|\n|SUMO (Newton-Schulz5, rank=4)|197M|61.81 ± 0.02|90.81 ± 0.013|92.43 ± 0.034|79.33 ± 0.031|94.14 ± 0.028|86.90 ± 0.021|92.22 ± 0.04|91.26 ± 0.018|\n|SUMO (SVD, rank=4)|**197M**|**62.32 ± 0.015**|**91.05 ± 0.007**|**93.48 ± 0.022**|**81.08 ± 0.019**|**94.93 ± 0.01**|**87.36 ± 0.023**|**93.25 ± 0.036**|**91.67 ± 0.015**|\n|LoRA (rank=8)|264M|**61.83**|90.80|91.90|79.06|93.46|86.94|92.25|91.22|\n|GaLore (rank=8)|257M|60.06|90.82|92.00|79.78|94.38|87.17|92.20|91.11|\n|SUMO (Newton-Schulz5, rank=8)|198M|61.73 ± 0.021|90.77 ± 0.032|91.93 ± 0.04|79.66 ± 0.03|94.13 ± 0.025|87.20 ± 0.01|92.22 ± 0.031|91.36 ± 0.02|\n|SUMO (SVD, rank=8)|**198M**|61.69 ± 0.014|**91.11 ± 0.02**|**93.72 ± 0.018**|**81.38 ± 0.011**|**94.83 ± 0.01**|**87.59 ± 0.012**|**93.65 ± 0.037**|**91.73 ± 0.018**|\n\nToken Budget: We ensured a fair comparison of token budgets by matching the total number of optimization steps across all methods. Since SUMO converges faster, as shown in Figure 2, this means that fewer training tokens are generally required. \n\nAs for the wall-clock timing, please refer to the following Table 6, presented in the (supplementary material) appendix D, \"Additional Experiments\" where we reported training times. \n\n|Methods|Rank|Time(h)|Memory(GB)|Accuracy(%)|\n|---|---|---|---|---|\n|LoRA|32|0.40|14.36|45.80|\n|DoRA|32|0.69|15.01|44.96|\n|GaLore|32|2.59|15.15|58.40|\n|SUMO(Newton-Schulz5)|32|1.83|13.86|58.47|\n|SUMO(SVD)|32|1.56|13.86|61.23|\n|LoRA|128|0.45|15.64|65.97|\n|DoRA|128|0.72|16.17|66.81|\n|GaLore|128|2.61|15.79|64.29|\n|SUMO(Newton-Schulz5)|128|1.78|14.12|64.41|\n|SUMO(SVD)|128|1.62|14.12|**68.03**|\n\n**W2 Low‑Rank Gradient Behavior in FFN Layers.**\n Weights in feed‑forward (FFN) blocks, especially the “project‑up” layer in Transformers, naturally exhibit a low-rank gradient structure during training. This phenomenon is theoretically substantiated in GI-rank collapse theory [1], which shows gradients in fully-connected architectures collapse to low rank, even under nonlinearity and standard activations. Separately, in the JoMA framework [2] and the GaLore study, it is shown that Transformer FFN gradients converge to extremely low rank (sometimes rank‑1), in practice during training, despite the layers not being strictly invertible [3] (in Appendix Sec. B). These findings support our empirical observations and validate why our SUMO method, leveraging low-rank orthogonalization, remains effective in FFN and attention layers. This will be clarified in the main text. \n\nIllustration for the moment singular values decay is presented in Figure 1 (b), in the manuscript.\n\n**W3 Study of subpsace update period K, rank r**\nPlease refer to our reply to your Q1 (below), where we provided a grid-search. \n\n**W4 Single GPU**. Our focus in this paper was to evaluate SUMO under constrained compute settings (single GPU) to emphasize its practical benefits in memory-limited environments. Nonetheless, the core SUMO operations, low-rank projections, and subspace-aware updates are orthogonal to the data-parallel training paradigm and can be naturally integrated into distributed frameworks.\nTo clarify, SUMO is implemented in PyTorch and compatible with torch.distributed. We are currently conducting multi-GPU experiments, which we plan to release with our open-source code.\n\n**Q1 Hyperparameters Grid search.**: To evaluate the impact of Subspace Update Frequency (K) and Ranks (r), we performed a grid search during the pretraining of the LLaMA 130M model on the C4 dataset. This specific setup allows for a direct comparison with the Galore method.\nTable: Perplexity results from a grid-search of Subspace Update Frequency (K) and Ranks (r) for the LLaMA 130M model pretrained on the C4 dataset. Values are presented as Galore/SUMO.\n|Update Frequency|Rank = 128|Rank = 256|Rank = 512|\n|---|---|---|---|\n|100|29.7/**28.27**|27.9/**26.74**|27.4/**26.73**|\n|250|28.1/**27.86**|26.5/**24.87**|26.2/**24.82**|\n|500|27.2/**25.91**|25.6/**24.98**|25.3/**24.31**|\n|1k|26.8/**25.83**|**25.1**/25.42|**24.8**/24.93|\n\nCompared to Galore, SUMO demonstrates notably less sensitivity to variations in rank and update frequency. This suggests greater hyperparameter robustness and potentially easier tuning, achieving competitive or superior perplexity with more stable performance across the tested range.\n\n**Q2 per-step wall-clock overhead (ms) of the randomized SVD**. In the following Table, we show the time measurements for subspace computations on the GULE MRPC dataset, comparing sumo, and Galore on A100 GPU: \n| Subspace update role                        | Average time (ms) of subspace update - Initial Rank 8 |\n|---------------------------------------------|-------------------------------------------------------|\n| GaLore SVD Block                            | 0.6610947                                           |\n| SUMO Randomized SVD    | 0.1232528                                            |\n\nIn addition, the average time of the moment's subspace transformation/rotation is measured at 0.023574 ms.\n\n**Q3 low-rank momentum observation**. This low-rank momentum phenomenon occurs independently within MLP layers, as supported by our proof Theorem 3.1. The concatenation of layers is not a prerequisite for the observation. Further empirical evidence of (first order) moment singular value decay can be found in Figure 1 (b) of the manuscript. In this Figure, singular value decay of the momentum at step 100 shows that the top 4 singular values capture the vast majority of the information.\nPlease refer to our response to W2 for a further detailed theoretical explanation. \n\n**Q4 high scale results**. We are working on a larger scale evaluation, and hope to report the results as soon as we have them. Please note though that our paper focuses on developing a method to perform well on hardware with limited resources.\n\nWe hope this addresses all your concerns, and we would be happy to further clarify open issues.\n\n[1] Bradley T. Baker et al. Low-Rank Learning by Design: the Role of Network Architecture and Activation Linearity in Gradient Rank Collapse.\n\n[2] Tian, Y. et al. Demystifying multilayer transformers via joint dynamics of MLP and attention. In The Twelfth International Conference on Learning Representations.\n\n[3] Jiawei Zhao et al. GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection."}, {"review_id": "ZoDwadEv2k", "reviewer": "Reviewer_AznZ", "summary": "This paper introduces SUMO (Subspace-Aware Moment-Orthogonalization), a novel optimization algorithm designed to accelerate the training of LLMs while simultaneously enhancing memory efficiency. SUMO's core innovation is the use of an exact SVD to perform moment orthogonalization within a dynamically adapting low-dimensional subspace. This method ensures that optimization steps are better aligned with the spectral characteristics of the loss function.", "questions": "1. Your complexity analysis notes that SVD can require approximately twice the FLOPs of Newton-Schulz5 in certain low-rank settings. However, your results show a significant speedup in terms of optimization steps (Figure 2) and total training time (Table 6). Could you provide more insight into the wall-clock time trade-off and confirm if the per-step overhead of SVD is consistently overcome by the faster convergence across all tested scenarios? And I find the results in Table 6 somewhat puzzling: the reported training times (in hours) are GaLore (2.59) > SUMO (Newton-Schulz5) (1.83) > SUMO (SVD) (1.56), which appears to contradict the relative algorithmic complexity of the methods.\n2. Since SUMO is built upon GaLore and Muon, a natural baseline for comparison would be these two methods. However, I did not observe such comparisons in the main experimental results presented in Tables 1, 2, and 3. I am particularly curious about how SUMO performs relative to Muon.", "limitations": "Types:\n\n1. Are the two square brackets in line 140 missing corresponding citations?\n\n2. In **Algorithm 1，“Q_t** *←* Truncated_Randomized_SVD(**G_t**)” should be “**Q^{(t)}** *←* Truncated_Randomized_SVD(**G^{(t)}**)”. And other G_t should be G^{(t)} in Algorithm 1.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "1. **Strengths**\n    1. The paper provides substantial theoretical backing for its approach. These theoretical claims are validated through empirical evaluations.\n    2. The paper successfully argues that the computational cost of SVD, typically a concern, is manageable and even preferable within the proposed low-rank framework. \n    3. The authors have conducted a thorough empirical study across a diverse set of tasks, including pre-training, fine-tuning, and long-context reasoning.\n2.  **Weaknesses**\n    1. My primary concern is that it seems SUMO is essentially Muon with GaLore. In fact, GaLore stands for Gradient Low-Rank Projection and can be combined with various optimization methods. The GaLore mentioned in the paper specifically refers to Algorithm 2: Adam with GaLore from [1]. Therefore, SUMO differs from Muon with GaLore only in some minor modifications. The authors should provide a more detailed discussion of these differences to clearly articulate the novelty and distinct contributions of SUMO compared to Muon with GaLore.\n    2. The pre-training experiments are conducted on up to 1.1B and 13.1B tokens. While the results are promising, validating the method’s effectiveness at a larger training scale would strengthen the claims regarding scalability. In practice, I have observed that many improvements appear effective during the early stages of training, but their advantages often diminish—or even reverse—after exposure to a substantially larger number of training tokens.\n\n[1] Jiawei Zhao, et. al. Galore: Memory-efficient llm training by gradient low-rank projection. ICML 2024.", "was_revised": true, "final_justification": "I appreciate the authors’ detailed response and the supplementary experiments provided, which help clarify several points. While some concerns remain, such as the lack of an ablation study isolating the specific impact of Block 1.1 and questions about scaling behavior with increasing training tokens. The paper presents a meaningful advancement, and the core ideas are well-supported by theory and experiments. With further refinements and more thorough documentation in the final version, this work will make a valuable contribution to the field. Therefore, I raise my rating.", "rebuttal": "We thank the reviewer for the helpful feedback and recognition of our theoretical and empirical contributions. We appreciate the concerns about the distinction from prior work and large-scale validation. These are important points, and we provide clarifications and new comparisons below.\n\n**W1 distinction from Muon.**\n We thank the reviewer for raising this critical point, as it provides an excellent opportunity to clarify the significant distinctions and novel contributions of SUMO compared to Muon combined with GaLore. While both SUMO and Muon incorporate gradient orthogonalization, and GaLore provides gradient low-rank projection, their combination does not simply equate to SUMO.\nSpecifically, SUMO differs notably from Muon combined with GaLore in the following key aspects:\n\nExact Orthogonalization via SVD: Muon uses Newton-Schulz iterations to approximate the spectral orthogonalization of gradient moments. However, as we demonstrate analytically in Lemma 3.2, this approximation can introduce significant errors, particularly in ill-conditioned optimization landscapes often encountered during the training of large language models (LLMs). In contrast, SUMO utilizes the exact Singular Value Decomposition (SVD) for orthogonalization, which is feasible because we are operating in a low-dimensional space. This approach completely eliminates the approximation error associated with Newton-Schulz methods.\n\nSubspace-Aware Dynamic Adaptation: GaLore, by definition, uses a fixed-rank gradient projection but does not dynamically adapt the optimization subspace or adjust its orthogonalization based on spectral information. SUMO explicitly performs SVD-based orthogonalization within a dynamically adapted, low-dimensional subspace. This subspace awareness ensures optimal alignment with the dominant spectral characteristics of the loss landscape, significantly improving convergence stability and speed, as validated empirically.\n\nMoment Transformation Across Adaptive Subspaces: SUMO uniquely maintains consistency in optimization by transforming the first-order moments between dynamically changing subspaces (Algorithm 1 Block 1.1). This subspace transformation is critical to retaining accurate optimization steps and is a novel component that does not exist in either Muon or GaLore independently.\n\nTheoretical and Empirical Justification: Our theoretical analysis clearly distinguishes SUMO from Muon by quantifying the exact impact of Newton-Schulz approximation errors on convergence (Lemma 3.3, Remark 3.7). Empirically, we provide direct ablation studies comparing SUMO (exact SVD) against SUMO (Newton-Schulz) and GaLore alone. Results consistently show SUMO’s superior performance, attributed explicitly to eliminating orthogonalization errors via exact SVD and the adaptive subspace optimization strategy.\n\nThese distinct innovations highlight that SUMO is not merely a combination of existing methods but introduces significant theoretical and practical advances that directly address the shortcomings of Muon and GaLore when used individually or in combination.\n\n**W2 Scalability of SUMO.** \nWe appreciate the reviewer's feedback regarding the need for large-scale validation. A central motivation for our research on efficient fine-tuning and optimization is to democratize pretraining, making it accessible to smaller laboratories beyond the major industry players like OpenAI and Google. In this context, demonstrating consistent performance improvements across up to 13 billion tokens is quite significant. This scale exceeds the capacities of many academic settings and provides strong evidence of the method's robustness and scalability within realistic computational budgets. Additionally, we believe that the faster convergence of SUMO is valuable even for a larger number of tokens.\n\n\n**Q1 explaining convergence time result.**\nWhile SVD may have higher FLOPs per step in some low-rank settings compared to Newton-Schulz5, our empirical results demonstrate that SUMO (SVD) consistently achieves faster convergence in terms of optimization steps (Figure 2) and lowers total training time (Table 6). This is primarily due to the improved stability and accuracy of exact SVD-based orthogonalization, which leads to a more direct and efficient path to convergence in the highly anisotropic loss landscapes of LLMs.\n\nThe seemingly contradictory results in Table 6, where SUMO (SVD) has a lower training time (1.56 hours) than GaLore (2.59 hours) and SUMO (Newton-Schulz5) (1.83 hours), despite the FLOPs discussion, are due to early stopping based on validation performance. Faster convergence means that SUMO (SVD) reaches stopping cretearia on validation performance (e.g., perplexity or accuracy on a validation set) in significantly fewer optimization steps. Therefore, even with a potentially higher per-step computational cost, the reduced number of total steps leads to a substantial decrease in overall wall-clock training time. This faster convergence rate consistently overcomes the per-step overhead of SVD across all tested scenarios, making SUMO (SVD) more efficient in practice.\n\n**Q2 Comparison with GaLore:**\n We would like to clarify that GaLore is a primary baseline in all our key experiments. We compare against it directly in Table 2 (GLUE), Figure 2 (QNLI convergence), Table 3 (LLaMA pre-training), and Tables 4, and 5, where SUMO consistently shows superior performance and memory efficiency.\n\n**Comparison with Muon.** Thank you for your suggestion. We opted not to include the original Muon because it operates under a different paradigm than SUMO. While SUMO is designed for low-rank, memory-efficient optimization, Muon updates the entire parameter space, leading to memory usage comparable to full fine-tuning, which is already addressed in our baselines. \n\nWe do provide a fair comparison: SUMO (Newton-Schulz5) incorporates Muon’s core approximation within our low-rank framework. The weaker performance of SUMO (Newton-Schulz5) compared to SUMO (SVD) underscores the advantages of exact orthogonalization, especially in memory-constrained environments. Additionally, applying Muon to all parameters introduces significant computational overhead, making it less suitable for the efficiency-focused scope of our work. We hope this clarifies our experimental design.\n\nWe would kindly like to share the following new experimental results are presented with the two other reviewers.\n\n**Hyperparameters Grid search**: To evaluate the impact of Subspace Update Frequency (K) and Ranks (r), we performed a grid search during the pretraining of the LLaMA 130M model on the C4 dataset. This specific setup allows for a direct comparison with the Galore method.\nTable: Perplexity results from a grid-search of Subspace Update Frequency (K) and Ranks (r) for the LLaMA 130M model pretrained on the C4 dataset. Values are presented as Galore/SUMO.\n|Update Frequency|Rank = 128|Rank = 256|Rank = 512|\n|---|---|---|---|\n|100|29.7/**28.27**|27.9/**26.74**|27.4/**26.73**|\n|250|28.1/**27.86**|26.5/**24.87**|26.2/**24.82**|\n|500|27.2/**25.91**|25.6/**24.98**|25.3/**24.31**|\n|1k|26.8/**25.83**|**25.1**/25.42|**24.8**/24.93|\n\nCompared to Galore, SUMO demonstrates notably less sensitivity to variations in rank and update frequency. This suggests greater hyperparameter robustness and potentially easier tuning, achieving competitive or superior perplexity with more stable performance across the tested range.\n\n**Per-step wall-clock overhead (ms) of the randomized SVD** In the following Table, we show the time measurements for subspace computations on the GULE MRPC dataset, comparing sumo, and Galore on A100 GPU: \n| Subspace update role                        | Average time (ms) of subspace update - Initial Rank 8 |\n|---------------------------------------------|-------------------------------------------------------|\n| GaLore SVD Block                            | 0.6610947                                           |\n| SUMO Randomized SVD    | 0.1232528                                            |\n\nIn addition, the average time of the moment's subspace transformation/rotation is measured at 0.023574 ms.\n\nMoreover, we extend our evaluation to diverse commonsense and reasoning tasks using zero-shot tasks, followed by the exact evaluation protocol described in Table 4 of Apollo (Zhu et al. SGD-like Memory, AdamW-level Performance). SUMO-pretrained models achieve lower perplexity and outperform AdamW on downstream benchmarks.\n\n**Table 7**: Zero-shot evaluation of LLaMA-350M models pretrained with sequence length 1024 across reasoning tasks\n|Method|Memory|Perplexity|BoolQ|RTE|HS|WG|OBQA|ARC-E|ARC-C|PIQA|SciQ|MathQA|Avg.|\n|-----------|------|---------|----------|----------|----------|----------|---------|----------|----------|----------|---------|----------|----------|\n|AdamW|1.37G|16.30|0.4917|0.4693|0.3688|0.5233|0.332|0.3729|0.2449|0.6534|0.609|0.2064|0.4272|\n|APOLLO|0.34G|15.64|0.5373|0.4698|0.3850|0.4925|0.322|0.3788|0.2483|0.6681|**0.624**|0.2127|0.4406|\n|APOLLO-Mini|0.15G|16.12|0.5376|0.4562|0.3707|0.5217|**0.324**|0.3758|0.2312|0.6638|0.619|0.2224|0.4374|\n|SUMO|0.18G|**15.49**|**0.5479**|**0.4709**|**0.3937**|**0.5313**|0.321|**0.3832**|**0.2496**|**0.6709**|0.623|**0.2246**|**0.4416**|\n\nWe hope this addresses all your concerns, and we would be happy to further clarify open issues."}]

Author rebuttal:

Meta-review: The paper proposes a new optimizer based on moment orthogonalization and demonstrates practical benefits when applying it to LLM training.  After a rebuttal and discussion process in which the authors provided additional experimental results, all three reviewers lean toward accept and indicate that most, but perhaps not all, of their concerns were addressed.  There is a consensus that the ideas and their technical development are a strength, though there may be room for further experimental validation.  The Area Chair agrees with the unanimous reviewer rating.
Decision: Accept (poster)