Parameter and Memory Efficient Pretraining via Low-rank Riemannian Optimization
Authors: Zhanfeng Mo
Long-Kai Huang
Sinno Jialin Pan
Venue: ICLR 2025 Poster (2025)
Keywords: Neural Network Optimization
Parameter Efficient Pretraining
Low-rank Optimization
Primary area: optimization
TL;DR: A Riemannian optimizer for pretraining fully low-rank parameterized language models.

Abstract: Pretraining large language models often requires significant computational resources and memory due to their vast parameter amount. An effective approach to enhance parameter efficiency in both training and inference is to parameterize each full-size weight as the product of two trainable low-rank factors. While low-rank fine-tuning has achieved great success, low-rank pretraining remains challenging as it requires learning extensive knowledge from scratch under the restrictive low-rank parameterization. During standard low-rank pretraining, separately optimizing the low-rank factors introduces redundant information from the full gradient, which hinders the learning process. To achieve efficient yet effective low-rank pretraining, we propose a **Lo**w-rank **R**iemannian **O**ptimizer (**LORO**). At each LORO update step, the low-rank factor pairs are jointly updated to ensure their full-size product moves along the steepest descent direction on the low-rank manifold, without the need to compute any memory-intensive full-size matrices or gradients. Hence, our LORO finds low-rank models that achieve high performance comparable to full-size pretrained models, while significantly reducing memory usage and accelerating both training and inference. A LLaMA 1B model pretrained with LORO achieves a perplexity score of 2\% better than the full-size baseline, with a 54\% reduction in model memory, a $\times1.8$ speedup in training, and a $\times2.2$ speedup in inference. The code is available on https://github.com/mzf666/LORO-main.

Published as a conference paper at ICLR 2025
PARAMETER AND MEMORY EFFICIENT PRETRAINING VIA LOW-RANK RIEMANNIAN OPTIMIZATION
Zhanfeng Mo
Long-Kai Huang
Sinno Jialin Pan
Nanyang Technological University, Singapore Tencent Al Lab & The Chinese University of Hong Kong
zhanfengo01@ntu.eem
hlongkai @gmai 1 com;
sinnopanecutk.edu
ABSTRACT
Pretraining large language models often requires significant computational re- sources and memory due to their vast parameter amount. An effective approach to enhance parameter efficiency in both training and inference is to parameterize each full-size weight as the product of two trainable low-rank factors. While low- rank fine-tuning has achieved great success, low-rank pretraining remains chal- lenging as requires learning extensive knowledge from scratch under the restric- tive low-rank parameterization. During standard low-rank pretraining, separately optimizing the low-rank factors introduces redundant information from the full gradient, which hinders the learning process. To achieve efficient yet effective low-rank pretraining, we propose Low-rank Riemannian Optimizer (LORO). At each LORO update step, the low-rank factor pairs jointly updated ensure their full-size product moves along the steepest descent direction on the low-rank manifold, without the need to compute any memory inttensive full-size matrices or gradients. Hence, our LORO finds low-rank models that achieve high per- formance comparable full-size pretrained models, while significantly reducing memory usage and accelerating both training and inference. A LLaMA IB model pretrained with LORO achieves perplexity score of 2% better than the full-size baseline, with 54% reduction in model memory, a x1.8 speedup in training, and a x2.2 speedup in inference. The code is available on GitHut")
INTRODUCTION
In recent years, large language models (LLMs) (Brown et al. 2020 Touvron et al. 2023a b Team et al. [2024a]] b). whose parameter sizes often reach hundreds of billions, have obtained remarkable performance across a wide range of applications. Typically, a pretraining and fine-tuning paradigm is adopted for building powerful LLMs from scratch: model is first pretrained on large-scale unsupervised corpora (Devlin et al. 2019 Brown et al. 2020) to acquire general semantics and smaller downstream dataset to enhance its domain- specific capabilities (Hu et al. 2022 Team et al. 2024b Ramesh et al. 2024).
extensive knowledge. and then fine- tunet
An extensive amount of trainable parameters results in substantial memory footprint and high computational costs during pretraining and fine-tuning. To alleviate these issues, various memory- efficient training techniques have been proposed to reduce memory usage and accelerate training. These approaches stem from either an engineering perspective, such as gradient checkpointing (Chen al.) 2016), memory offloading (Rajbhandari al. 2020),
optimization perspective,
including low-rank optimizers (Hu et al.) 2022| et al.
2023: Hayou et al. 2024) Lialin et al.)
2024) and low-precision optimizers (Dettmers al.) 2021 Li al. [202]] Zhang et al. (2024b). Since low-rank adaptation (LoRA) (Hu et al. (222)) has gained increasing attention, becoming one of the most popular parameter-efficien size pretrained weight W E Rdxd with W + BA, where B Rddxr and A E Brxd are two trainable low-rank factors, while W is frozen during fine-tuning to reduce training memory foot- print. However, these methods depend on full-size pretrained weights, making them inapplicable to
tuning methods. At each layer, LoRA replaces the full-

---

Published as a conference paper at ICLR 2025
pretraining scenarios. Furthermore, they ultimately produce models with full-size weights, which do not result in any memory reduction or speedup in the inference phase.
In this work, we aim to enhance efficiency during both the training and inference by exploring the pretraining of low-rank parameterized language models, where each weight matrix is parameterized as W - BA, the product of two rank-r matrices. In practice, such low-rank parameterization enjoys great potential in improving both training and inference efficiency, as it reduces the parameter size at each layer from de to 2rd and decreases the FLOPS from (dd) to O(rd)). Theoretically, a low-rank model can achieve a x x(2r/d) memory reduction and x (d/2rr inference speedup.
However, pre-training low-rank language models from scratch using standard stochastic gradient descent meth- ods is much more challenging. As shown in (Kamalakara al. 2022 Zhao al.) 2024), separately optimizing each low-rank factor of the model may result in a significant performance drop compared to full-size models. There- fore, as compromise, recent works consider pretrain- ing language models with full-size weights while utiliz- ing low-rank or sparse updates to enhance training effi- ciency: ReLoRA (Lialin et al. 2024) periodically merges low-rank update into the full-size weight after a full-size warm-up training phase; GaLore (Zhao et al.] 2024) up- dates the full-size model weights using low-rank yet full- size gradients; SLTrain (Han et al. 2024) proposes sparse plus low-rank pretraining, where the full-size weights are decomposed into sum of sparse and low-rank matrices. However, these methods can merely produce full-size
Purpkesity va Memory Efficiancy on LLaMA-119
Lowerank
1300
-00
SUTrain
- cen
Aull 402
201
LORD
-200
ICO)
Figure 1: Visualization of the perplex- ity (log scaled) and memory- -fficinccf trade-off of pretraining LLaMA-1B on C4 (Table D Methods closer to the bottom-left corner and associated with smaller bubble sizes exhibit superior performance and memory efficiency. models, leading to limited or no memory reduction or acceleration during inference.
While recent studies Zhao et al., 2024, Han et al. 2024) observe that the rank of the resulting weight matrix W obtained from full-size pretraining methods is typically high (close to d). this does not necessarily mean that low-rank pretraining approaches (learning BA instead of W) are incapable of achieving satisfactory performance in learning foundation models. We argue that with careful design of the learning process for the low-rank factors A and B, it is possible to train a model with performance comparable to full-size models. Therefore, this paper aims to address the research question: Can we train a low-rank parameterized language model from scratch to achieve performance comparable the full-size baseline? In essence, low-rank pretraining is more chal- lenging than fine-tuning. Given the pretrained full-size weights, fine-tuning only requires capturing the minor dopain-specific knowledge through low-rank adaptation. In contrast, learning exten- sive pretraining knowledge from scratch under the restrictive low-rank parameterization requires a more careful optimization process. As noted in the discussions of (3) in Section 3 optimizing the low-rank factors separately introduces redundant information from the full gradient, which hinders the feature learning process. In this scenario, the separate optimization of B and A fails to direct their product BA along the steepest descent direction of the loss on the rank-r matrix manifold. This necessitates considering the intricate geometry of the manifold, prompting us to optimize the low-rank factors jointly through Riemannian optimization. While recent works (Savostianova al.) 2023 Schotthöfer et al.) 2022) have attempted to train low-rank vision classifiers using Riemannian gradient methods, they either require evaluating full-size gradients or rely on SVD-like parameter- ization, which introduces Orr²) additional parameters, making them less preferable for pretraining language models. To the best of our knowledge, this paper is the first to explore the pretraining of low-rank language models through the lens of Riemannian optimization.
Contributions. We propose the Low-rank Riemannian Optimizer (LORO) for effective low-rank pretraining. At each LORO update step, each pair of low-rank factors is jointly updated to ensure their product is transported along the Riemannian gradient of the rank-r matrix manifold to pursue effective descent in the model loss. All LORO'S updates are performed directly within the low- rank parameterization, without the need to evaluate the costly full-size weights or gradients. To our knowledge, we are the first to show that pretraining low-rank language models with Riemannian gra- dients achieve perplexity scores comparable to full-size pretraining, based on extensive experiments conducted on LLLMA models ranging from 60M to 1B in size. On LLLMA- IB, LORO achieves a
2

---

Published as a conference paper at ICLR 2025
perplexity score of 2% better than the full-size baseline, with 54% less model memory cost and of- fers a 1.8 speedup in training and a x2.2 speedup in inference. Within the same parameter budget, we further show that LORO outperforms its full-size counterparts, highlighting LORO's capacity to unblock the potential of low-rank parameterization for enhanced representation learning.
2 RELATED WORKS
Low-rank fine-tuning. Pioneered by LoRA (Hu et al., 2022). which achieves parameter-efficient
fine-tuning by restricting optimization exclusively to the low-rank factors BA, various low-rank fine-tuning methods have been proposed to enhance downstream adaptation (Zhang al., 2023
Ding et al. 2023: Hayou et al.
2024; Zhang & Pilanci 2024; Xia et al.) 2024). improve cross-
task generalization (Agiza al.
2024
Huang et 2003 Buehler & Buehler [2024). and increase
fine-tuning efficiency (Chen et al. 2023] Dettmers et al. 2023 Gamal & Rabusseau 2023 Kooh-
payegani et al. 2024 Valipour et al. 2003) However, all these low-rank fine-tuning methods require access to the full-size pretrained weights so the model can be adapted within low-rank parame- ter space. When directly applied pretraining, the absence of pretrained full-size weights-where most knowledge resides-ooe often results in significant performance drops.
Low-rank pretraining. While a few works have attempted to train low-rank vision models (Sui et al. [223 Khodak et al., 2021] Savostianova et al., (2023), training low-rank parameterized lan-
guage models from scratch has been observed to be challenging and remains an open problem Ka-
malakara et al.) 2022. Zhao et al. 2024). Existing methods circumvent these challenges by employ-
ing low-rank or sparse updates while still using full-size weights, such as updating full-size weights with low-rank gradients (Zhao et al., 2024; Zhang et al. 2024b), periodically merging low-rank factors into full-size weights (Lialin et al. 2024), or using a low-rank plus sparse weight parameter- ization (Han et al. 2024). As result, these methods provide little to improvement in inference efficiency. In contrast, LORO is the first to directly tackle language model pretraining using low- rank parameterized weights.
Optimization on low-rank matrix manifolds. Abundant theories and algorithms (Absil & Os- eledets 2015 Absil et al.. 2009. Mishra al.) (014) have been proposed to perform Riemannian optimization on low-rank matrix manifolds for traditional applications, such as matrix completion Vandereyckenn 2013 Bian al. 2023) and regressions (Bleakley & Yamanishi 2009] Yuan al.,
2007] Amit et al. 2007). Recent works (Savostianova et al.) [023]] Schotthöfer et al. 2022) aim to train low-rank vision classifiers using Riemannian gradient flow (Koch & Lubich 2007). How-
ever, these methods either require evaluating full-size gradient and weight matrices or rely on SVD- like parameterization, which introduces parameter overhead, making them incompatible with low- rank parameterized language models. In contrast, LORO operates under low-rank parameterization, avoiding full-size matrix evaluations and enabling efficient pretraining on language models.
3 PRELIMINARIES
Low-rank pretraining. The rank-r parameterization of a d-by-d matrix is defined as W = BA, where E Rdxr and A Rixd are the associated rank-r factors. When r < d, BA is called the low-rank parameterization of W. In this work, we focus on training the decoder-only transformer- based language models (Vaswani et al. 2017 Touvron et al., 2023a) from scratch using low-rank parameterization. Recall that the two fundamental building blocks of the transformer model are the linear projection f-ro)()) and the self-attention function facta() (.), which are defined as
Sproj(X) (X) A XW p. faten (X) softmax X (XW, W x /Va/ XW..
(1)
where X E Ra is the hidden feature, n is the token number, d is the hidden dimension, and Wp W,, Wx. W. E dxd denote the trainable projection, query, key, and value matrices, respectively. For notation simplicity, we assume that d remains constant across all layers, though all derivations in this work also apply to scenarios with varying dimensions. As the composition of multiple layers of linear projections, self-attention functions, and various operations (e.g., activation, normalization, and positional embedding), the target language model can be parameterized by the stack of trainable matrices across all N layers, denoted as (W,,.., Wx). Therefore, the rank-r parameterized
3

---

Published as a conference paper at ICLR 2025
language model is implemented by parameterizing each trainable matrix as a rank-r factorization W; = B,A; with B; E pakr and A, E Rrxd and is represented by (BiAi..... By) In practice, we set < d for parameter efficiency. During pretraining. the rank-r model is trained from scratch on unsupervised corpus data for the next token prediction, defined as:
min L(WI,'.. WN), s.t. W; B,A, B, Rddr A E ERnx I,,,,,,
where T(-): Rdxd x... x Raxd - R- denotes the smooth next token prediction loss (i.e., the logarithm of perplexity score (Radford & Narasimhan 2018)). At the 1-th training step with learning rate 7j > O, the full-size weight updated by the full gradient W < V/V:: L. In contrast, during the standard low-rank pretraining, the low-rank factors updated separately, which leads to
W (B. VB; () D = W 7 (B.B.T Vw: Vw; ATTA
(3)
where the last line is derived from the chain-rule B B W VA! and - The update direction in (3) is the sum of two sketched versions of the full gradient, each projected onto the low-rank subspaces spanned by the columns of B and the rows of A.. When these two subspaces overlap. the sketched gradient introduces redundant information from the full gradient, hindering the learning of effective representations. From a geometric perspective, separately opti- mizing the low-rank factors overlooks the intricate structure of the low-rank parameterization and fails to guide their product toward the steepest loss descent on the low-rank matrix manifold To address this issue, we recast (2) into Riemannian optimization problem:
min BRERaXT, A
L(B BN)
min L(B ByAx), (B:A))EM,
(4)
where M, denotes the manifold of all d-by-d rank-r matrices. In essence, this reformulation
counts for the geometric structure of M, by treating the product (B.A;) as a whole: it emphasizee that the principle for updating B; and A, is to guide the product A.) towards the steepest descent direction within the manifold, rather than optimizing them separately in the ambient space.
Optimization low-rank matrix manifold. We briefly review the retraction-based Riemannian optimization on manifolds and refer interested readers to (Lee) 2019, Absil & Oseledets 2015) for rigorous definitions and a more comprehensive review. Imagine an ant on the Earth's ' surface seek- ing the highest point. Although it wants to climb straight up along the steepest path an infeasible option that would send it into space- must navigate the curved surface, choosing the best direction from the feasible tangents at its current location to remain on the Earth. This thought experiment helps illustrate the difference between optimization the flat full-size matrix space (e.g., the ambi- ent space) and on the curved rank-r matrix manifold M, (e.g., the Earth).
As visualized in Figure 2 when optimizing in the ambient space, we evaluate the negative Euclidean gradient (blue vec- tor) at each step by backpropagation, then update the current weight by moving it small step towards the blue vector. How- ever, this update is infeasible on Mr. as it moves the current weight out of the manifold. Instead, aim update along the black dashed curve on the manifold, which guides us toward smaller loss values. To this end, we need to identify the Rie- mannian gradient (red solid vector) within the tangent space (gray plane) that aligns most closely with the Euclidean gra- dient. This can be done by projecting the blue vector onto the tangent space. We then update the weight for a small step along the red solid vector and retract it back to the manifold (via the red dashed vector) to avoid violating the rank con- straint. This process is known as the retraction-based Rieman- nian optimization step. involving: 1. Updating with the pro- jected gradient and 2. Retracting back to the manifold.
I TwtM,
FL(W()
w
grid Z(W5)
Mr
SVD,()
W
Figure 2:
Retraction- Rie- mannian optimization, where the grey plane is the tangent space, VL(W') and grad L(W*) denote the Fuclidean and Riemannian gra- dient at W respectively.
Thanks to the embedded manifold theories, Proposition T (see (Lee 2019 Example 8. 8.14) and (Absil al. 2009 Equation 3.37)) shows how compute the Riemannian gradient via orthogonal projec- tion onto the tangent space, while Proposition 2 (see (Lewis & Malick 2008 Lemma 2.1)) shows that retraction M, can be done by rank-r singular value decomposition (SVD).
See Appendix C for more detailed discussions the issues in standard low-rank pretraining.
4

---

Published as a conference paper at ICLR 2025
Proposition 1 (Tangent Space and Riemannian Gradient on M,). M, is smooth submanifold em- bedded in the ambient space Rdxd For any base point W E M, with singular value decomposition
W -UVV with U, V E Raxr being orthogonal, the tangent space at W is given by
Tw & {UZVT + +UV +UV ZER U V Raxr -U - = (5)
and the orthogonal projection to Tw.,, is given by P(G) UUTG +
-UU
When equipped with the matrix inner-produs prreuuprouuctt metric, M, becomes Riemannian manifold. Then, the Riemannian gradient of any smooth function L(-) : Rdxdd + E evaluated at W, denoted as grad L(W), can be calculated by projecting the Euclidean gradient ((W) to Mp;
grad I(W)- = P(VL(W)) = -U VL(W) + VL(WVVV' - UU VL(W)VV".
(6)
Proposition 2 (Retraction to M, via SVD). Informally, for any base point W E My and tangent vector G E Mr, a retruction is function R(-) that maps the shifted base point (W + G) back to MM. A feasible retraction on M, is R(W + G) SVD,(W + GG, where SVD,(-) (-) is the rank-r singular value decomposition, projecting the updated base point from Raxd back to M,-
4 METHODOLOGGY
4.1 RIEMANNIAN UPDATE UNDER LOW-RANK PARAMETERIZATION
Proposition 1 and Proposition 2 enable us to perform Riemannian low-rank pretraining on M, via
w + H(W; ') grad L(W!)) SVD, (w: - P. (Vw:L))
(7)
where P:(.) denotes the orthogonal projection to the tangent space M,,, as defined in Proposition For each trainable weight, we: 1. evaluate the Euclidean gradient Vw; L. and derive the tangent space projection P:(-) from the SVD of base point W; = U, V 2. find the Riemannian gradient Pv, L) by projecting Vw, L to the tangent space; 3. transport the base point along the negative Riemannian gradient; 4. retract the updated base point back to M, via the SVD,(+) projection.
However, these update steps encounter several practical issues that render them incompatible with low-rank parameterization: First, the full-size gradient Vw, I. is inaccessible during backpropaga- tion (Rumelhart et al.. 1988); we can only obtain VB, and V VA,L, the gradients of the low-rank factors. Second, in steps I, 2, and 4, we need to manipulate the metory-intensive full-size matrix W, either by performing SVD or matrix multiplication, which compromises memory efficiency.
This motivates us to revise the update steps of 77 into a "low-rank parameterization friendly" form by avoiding any computations with full-size d-by-d matrices (c.g., the weight W, and gradient Vw, L) and ensuring all operations can be performed under the low-rank parameterization. To achieve this objective, our focus on the following two goals: Eliminating the need for the full-size gradient when calculating the Riemannian gradient through tangent space projection; 2. Keeping the Riemannian update in its factorized form to avoid the necessity of performing SVD on full-size matrices. To achieve the 1st goal, we propose Proposition 3 which shows that the Riemannian gradient on M, can be computed from VB, L and L.
Proposition 3 (Riemannian Gradient under Low-rank Parameterizzaton), Following the previous notation, for any low-rank parameterized weight W = BA E MM. where B E Rdxr and A E Rr the Riemannian grudient of any smooth function I(.)) : Rdxd + R. evaluated at W is
grad L(W)- = B(B VAL+ VBL (AA
B(B B) VAL AAL (8)
The proof of Proposition can be obtained from Proposition T based on two key observations: 1. Let UEVT be the SVD of BA, the column space of B coincides with that of U, and the row space of coincides with the column space of V. 2. Backpropagation gives V A = VBL and B VwL = L. The proof details are postponed to the Appendix B In practice, for each weight matrix W. (8) involves only multiplications of r-by-d and r-by-- matrices, as well as the inversion of r-by-r matrices, eliminating the need to compute any expensive d' b-by-d matrices or the d-by-d gradient V Ww, L. Additionally, all terms in (8) can be efficiently computed from the readily available variables B,, A;, VB, T., and V A, obbaiieed during backpropagation.
5

---

Published as a conference paper at ICLR 2025
Although Proposition] 3 provides memory- approach for caleulating the Riemannian gradi- ent, the resulting grad L(W) remains a full-size d-by-d matrix, which compromises the feasibility of the subsequent Riemannian update and retraction. To address this issue and achieve the 2nd goal, we propose Proposition 4 which shows that the successive update steps can be achieved within a factorized form, avoiding the need to compute any full-size weight gradient.
Proposition (Riemannian Update under Low-rank Parameterization). Following the previous notations. for any low-rank parameterized weight W = BA E M,, let B = Qg and A = QARA a be the QR deccooopossiio of B and A Then the Riemannian updated weight W A (W - 1i grad L(W)) admits a low-rank factorization W' B'A', with B' = (QB. QB) U. 1/2 and A' VI (QA, QA) Specifically, U.E.V V. determined by the rank-r SVD of the following 2r-by-2r matrix, i.e.,
U.E.V. = SVD,
7)
RA 0
E ER2222
, R n R VAL.QA ER
(9)
R
and Q RB is the QR decomposition of the d-by-r matrix (VBL R. RA is the QR decomposition of the d-by-r matrix AV R Qa ALL A
QB Q BBL R while
The proof of Proposition can be derived from Proposition using the tangent space decomposition /Vanderyycken] 2013 Bian al.) 2023) and the proof details are postponed to Appendix [B] The key to avoiding SVD on full-size d-by-d matrices lies in factorizing the d-by-d tangent vector into product of d-by-2r orthogonal matrix, 2r-by-2r full-rank matrix, and 2r-by-d orthogonal matrix, enabling SVD to be applied to the smaller 2r-by-2r matrix instead. Notably, the revised Riemannian update Proposition 4 only involves QR decomposition of d-by-r matrix, inversion of an r-by-r matrix, SVD on a 2r-by-2r matrix, and multiplication of d-by-2r matrices.
4.2 LOW-RANK RIEMANNIAN OPTIMIZER
In Proposition 4 we have derived the "low-rank parameterization friendly" Riemannian update:
B+I1 (QB, U.) A,+1 E!'vV (QA,, Q
(10)
where QB. QB. Q4. Q U., E., and V. can be directly computed from Bi, A. VB and VAyL, without evaluating any full-size weight or gradient. Although memory-efficient, the QR decomposition, SVD, and matrix inversion involved in (T0) impede the computational efficiency of low-rank pretraining. This motivates us to develop rough but efficient approximation of the exact Riemannian update, which enables us to amortize the computational cost across multiple steps by periodically applying the exact update while primarily using the approximated update.
To this end, simple approach is to skip the costly Riemannian SVD) retraction and directly update the weight along the Riemannian gradient. Unfortunately, the Riemannian gradient in (8) has maximum rank of 2r, making incompatible with the rank-r parameterization. To establish an approximate update that can be applied to B, and A, respectively, we highlight the differences between the standard low-rank update (3) and Riemannian update (8) in blue:
BB VA LDLL -A) + +(n"), Riemannian update- = (B T B VALL (I (AAA" +V (A(A!)) "AA))
Standard low-rank update 77
Ontoggnnal projection
In the Riemannian update, the gradients are additionally normalized by multiplying them with ei- ther (B: -1 or (A;A;")- Additionally, V A is projected onto the orthogonal space of the row space of A& (the middle blue term), thereby removing the redundant gradient components. This motivates us to approximate this normalization by scaling the gradient terms in proportion to the spectral norm of either (B: BB) or (A The orthogonal projection can be omit- ted, as it does not affect the gradient scale. Notice that [ (
((B B
or (B. (B)
or (B where || .l2 is the spectral norm and (x(-) denotes the k-th largest singular value of matrix. We estimate o, (B.)² by the average of the squared singular values, i.e., (1/r) Ei= o (B which equals (1/r) I|@ where || |r denotes the Frobenius norm. As- suming each row of B contains extreme values and its norm is upper-bounded by a constant,

---

Published as a conference paper at ICLR 2025
obtain B = (dd). Thus, approximate the scaling factor as r/d avoid explicitly computing
||
We set the approximated Riemannian update step as
B B + B 17 (r/d) L. A ++1 A - 7- (r/d) VA.L.
(1I)
We now propose the Low-rank Riemannian Optimizer): for every K training steps, LORO periodically employs an exact low-rank Riemannian update step (equation I0), while in between, it uses the approximated low-rank Rie- mannian update (equation [T]. In practice, LORO accu- mulates the gradients of B, and A using the momentum strategy as the Adam optimizer (Kingma& & Ba 2015). A pseudo-code of LORO is outlined in Algorithm| | and a PyTorch code is attached in Algorithm [
Algorithm Low-rank Ricmannian Optimizer
Input: current updato step trainable low-rank fac- toes B, A E IRehr, backward gradients Viy Varh, learning
v, update frequency K.
Oulput: updated rank factors B AA if mod K & U then
B
A. +1 - A -/t/////
else
(Qa. Rg) +- QIQ)) (Q,, R,,) QR(AST)
We refer the
7g Vn; T....R,,, R2, R,,,%,,7 7,, V. R3
readers to Appendix C for more detailed discussions on how LORO benefits low-rank pretraining.
(Q Rn) +- QRíZZ Q Q.
Q R.A) -QRIZ Q,Q, Za)
S RgR] .Rt V Qa
Complexity Analysis. As shown in Algorithm I the algorithmic complexity of the approximate LORO up- date step is ((rd) as it only involves element-wise op- erations. On the other hand, the exact update step in- cludes QR decomposition of d-by-r matrices with a com- plexity of Orrdd), inversion of r-b-- triangular matrices with complexity of O(r²), SVD of 2r-by-2r matrices at O(r²), and matrix multiplications with a complexity of d))d). Suppose the frequency of the exact update is K, then the average complexity of LORO is O(r d)/K rd). In practice, frequency K in the same order as r, which simplifies the LORO complexity to O(r² + rd). There- fore, when comparing with the standard Euclidean low-rank pretraining with ((rd) complexity, LORO does not introduce evident computational overhead.
U.E.V! +- SVD, Rp
R
B - (Qg. Qn) U.
A V. (Q.:. Q2))
end if return B A/+
set the
5 EXPERIMENT RESULTS
5.1 EFFICIENT LOW-RANK PRETRAINING
Experiment setup. In this section, we evaluate the effectiveness of our proposed LORO in training low-rank language models from scratch. All the experiments are implemented in PyTorch (Paszke) ct al., 2019) and conducted on NVIDIA 40G A100 GPUs For fair comparison with existing
baselines, we follow a standard experiment setup used in (Zhao et al.] 2024 Lialin et al. 2024). Specifically, adopt the ILaMA--ased language model (Touvron et al. 2023b) equipped with
RMSNorm (Zhang & Sennrich 2019) and SwiGLU activations (Shazeer, 2020). We follow the same model configurations as (Zhao et al., 2024] Table 5) for the LLaMA-60MP
7330M-350M and
-1B model and we conduct all the pretraining experiments using BF16 (bfloat 16) format. We train all the models on the C4 (Colossal Clean Crawled Corpus) dataset (Raffel et al. 2019). a large-scale cleaned dataset designed for language models pretraining. The models are trained on sufficiently large volume of data without repetition to simulate the pretraining setting in practice.
LORO configurations. We consider the low-rank parameterization that represents each target weight matrix as the product of a rank-r tall matrix and a rank-r wide matrix as defined in equa- tion 2 In this paper, we apply low-rank parameterization to all weights in attention functions (e.g. query-, key-, value-, and out-projection) and all linear projection layers (e.g. down-, up-, and gate- projection) with the same rank Following the setup in (Zhao et al. 2024: Han et al., 2024), we keep the embedding layer and the output language model head to be full-size. We fix the LORO up- date frequency K 500 step for all pretraining experiments, and we only tune the learning rate. We refer the readers to Appendix for implementation details of the experiments in this subsection.
Baselines. We compare our LORO against multiple baselines that involve low-rank structures. 'Full- size' denotes pretraining the full-size model with the Adam optimizer (Kinmma & Bal (015). 'Low- rank' is traditional low-rank training approach (Kamalakara et al. 2022). 'LoRA' (Hu et al.)
Since this study primarily focuses on comparing performance using BF16 across various pertaining meth- ods, we did conduct experiments on larger models due to limited computational resources.
7

---

Published as a conference paper at ICLR 2025
Table 1: Results of low-rank pretraining on C4 dataset, including validation perplexity (PPL), the
number of parameters in millions (Param), the estimated total memory cost in GB (Mem), where 1GB contains 10 bytes. r, d1, and d2 denote the rank, the hidden- and intermediate-dimensii the LLaMA model. Results of the baselines are reported from (Zhao al. 2024 Han et al. 2024).
of
Model (# Token)
LLaMA-G0M (1. (B)
LLaMA- 300 (2.28)
LLaMA-S50M (6.48)
LLaMA-IB (13.18)
dd
/28/512/188
256
2018
/2736
512/ 2018/
Metbod
PPL (1)
Param
Mcm (1)
PPL (.)
Param
Mem (1)
PPL (L)
Param
Mem (4)
PPL
Param
Mcm (J)
Full-sine
34.06
58
0.35
24.36
134
0.81
18.80
368
2.21
15.56
1339
8.04
Low-rank
78.18
43
0.24
45.51
94
0.57
37.41
185
1.11
142.53
609
3.66
LoRA
33.92
134
25.58
19.21
1339
ReLoRA
37.04
58
0.36
29.37
134
0.84
29.08
368
1.85
18.33
1339
6.34
GaLore
34.88
134
368
15.64
1339
4.76
SLTrain
34.15
44
0.26
26.04
97
0.60
19.12
194
1.24
16.14
646
4.16
LORO
33.96
43
0.24
24.59
94
0.57
18.84
185
1.11
15.19
609
3.66
Table 2: Training efficiency metrics evaluated on 1x40G A100 GPU, including token batch size (TBS), ratio of memory cost reduction (MemRd), and throughput (Tokens / sec).
Model
Method
PPL (1)
TBS
Mem (J)
MemRd (f)
Tokens / sec (t)
Speedup (t)
Full-size
18.80
16K
2.21
0%
35264
x1
LLaMA 350M
GaLore
18.95
16K
1.59
28%
35427
x1
LORO
18.84
16K
1.11
50%
40640
x1.15
Full-size
15.56
8K
8.04
0%
7890
x1
LLaMA IB
GaLore
15.64
8K
4.76
41%
7886
x0.98
LORO
15.19
3.66
54%
14351
x.882
2022 is low-rank fine-tuning method that reparameterizes each weight as W BA, where W a full-size matrix that remains fixed during training, and B and A are trainable low-rank factors. In pretraining scenarios, the full-size matrix W randomly initialized. 'ReLoRA' (Lialin et al. 2024) is a variant of LoRA that are adapted to pretraining by periodically absorbing the updated BA into W. 'GaLore' (Zhao et al. 2024) leverages a low-rank gradient to update the full-size weights 'SLTrain' (Han et al. (2244) parameterizes each weight as the sum of low-rank factors and a sparse matrix, with both components being jointly optimized. However, all baselines, except for 'Low-rank', work with full-size models and do not engage in the challenging low-rank pretraining.
LORO finds high- perfreerforing low-rank Models. Table shows that, our LORO consistently out- performs all other memory-efficient baselines in terms of both perplexity and total memory cost (including the memory usage of model parameters and optimizer states, as estimated in (Han et al.) 2024 Appendix C)) across all settings. For LLaMA-1B model, our LORO achieves perplexity score that is 2% lower than the full-size baseline. These results validate the existence of competitive low-rank models within the low-rank matrix manifold, which exhibit comparable superior per- formance compared to the full-size baselines. Moreover, our LORO offers a practical approach for training low-rank models from scratch without sacrificing performance.
LORO enhances training efficiency. Section T shows that a low-rank model with d1 = d2 is expected to achieve (2r/d) -timmestmes reduction in memory and (d/2r)-times speedup inference. In practice, we set r = /4 across all runs to pursue an approximate x0.5 FLOPS reduction and about x0.5 model size reduction. As shown in Table 2 for a LLaMA-IB model, LORO achieves x x1.82 training speedup and 54% reduction in memory compared to the full-size baseline, while GaLore does not exhibit practical speedup as it still works on the full-size model. We observe that the training speedup of LORO is less evident on LLaMA-350M model. This is because the data loading process introduces additional overhead and execution time, which overshadow the efficiency gains provided by LORO smaller models.
LORO enhances inference efficiency. In Table 3 we empirically evaluate the inference effi- ciency of the low-rank LLaMA-1B model obtained by LORO with varying batch sizes and se- quence lengths. Compared to the full-size baseline, our LORO achieves a x2.08 ~ x4.04 infer- ence speedup and a 5% ~ 17% reduction in maximum memory allocation, across all settings. The speedup of LORO is evident for small batch sizes, and it converges to x2 as the batch size increases. This occurs full-size baselines benefit more from GPU parallelization at larger batch sizes.
8

---

Published as a conference paper at ICLR 2025
Table Inference efficiency metrics on 1x40G A0O GPU, including maximum memory allocation in GB (MaxMem), memory reduction (MemRd), and total time of processing 500 batches.
LLaMA
Batch size = 512
Butch size 256
Sog
MaxMem (L)
MemRd (*) Time ())
Speedup (?)
Seq len
MaxMem (J)
MemRd (f)
(1)
Speedup (
Full-size
256
6.67
LORO
128
25.07
5%
1,57
x2.14
256
25.10
5%
3.07
x2.17
Full-size
14.49
1.79
128
14.52
4.69
LORO
13.08
10%
x2.27
128
13.11
10%
1.54
x3.04
Full-sine
32
8.49
1.41
64
8.52
3.28
LORO
7.09
17%
0.43
x3.32
64
7.12
16%
0.81
x4.04
LLaMA
Batch size = 128
Batch size = 64
IB
Sog len
MaxMem
MemRd (*)
Time (.)
Speedup (1)
Seq lon
MaxMem (4)
MemRd
Time
Speedup
Full-size
512
26.52
13.21
1024
26.53
63
LORO
25.12
5%
x2.09
1024
12.79
Full-size
256
14.53
6.67
512
14.54
14.68
LORO
256
13:13
10%
3.1
x2.15
512
13,14
10%
6.31
x2.33
8.54
256
7.59
LORO
128
7.14
10%
1.61
x2.32
256
7.15
16%
3,23
x2.35
Figure 3: Linear interpolation between LORO
Table 4: Comparison between LORO and other pre- trained full-size architectures within the same pa- rameter budget.
pretrained models and full-size baselines.
Interpolated LLaMA: (1 a/Woo aWauuiio
Model size
-e- 60M
Model
LLaMA 130M
LLaMA 350M
-a-- 130M
Token
2.2
6.4
r/d,
256.
708/ 2048
256/ 1024 /2736
Metiod
(P).
Pirau
Mai (L)
Mem
Pull-size
24.36
134
0.81
18,80
568
2.21
Low-can
45.51
94
0.57
37.41
185
1.11
Full- size Shallow
28.49
94
0.57
21.87
185
1.11
Pull size Slim
25.34
94
0.57
21.40
185
1.11
0.0
0.2
Interpoiation Rate
0.4
0.6
0.5
1.0
I DRD Slim MEP
1. LORO
25.16
24.59
94
94
0.57
0.57
18,84
18.95
185
185
1.11
LORO finds flatter minima. As Table 1 shows that LORO is able to find minima within the low-rank matrix manifold with high performance comparable to the full-size models, we aim to compare between the low-rank minima found by LORO, denoted by (B; and the full-size minima, denoted by (W:). To this end, we probe the loss landscape of these minima by via linear model interpolation (Lucas et al.) 2021 Gueta et al. 2023). Specifically, for a 0....1...., 1), we visualize the perplexity score the model with weights (1 - a B A + aW; As shown in Figure 3 steep loss barrier is observed between the low-rank
models and the full-size pretrained models, indicating that the LORO minima do not reside in the same basin as the full-size minima. Furthermore, the loss barrier exhibits asymmetry, with the side closer to the full-size models being significantly steeper than the side closer to the low-rank models. This implies that LORO can find flatter low-rank minima, which may explain why sometimes outperforms the full-size baselines as shown in Table 2
LORO explores effective low-rank structures. We compare the performance of LORO low-rank pretrained models against various full-size counterparts within the same parameter budget. Specif- ically, we evaluate three full-size variants of the original model: 1) 'Full-size Shallow,' with fewer layers than the full model, 2) "Full-size Slim,' with a reduced hidden dimension but the same num- ber of layers, and 3) 'LORO + Slim MLP, where LORO is applied only to the attention weights, while the MLP layers are slimmer but retain full-size matrices. Their configurations are detailed in Appendix A As reported in Table the LORO low-rank pretrained models outperform full-size variants, demonstrating that LORO unlocks the potential of low-rank structures as more effec- tive parameterization scheme. We attribute this to the fact that, within the same parameter budget, low-rank structures benefit from larger hidden dimension, which enhances their ability learn high-dimensional features.
5.2 ABLATION STUDIES
Ablation on the LORO rank r. In Table we evaluate the perplexity scores of LORO pretrained LLaMA-60M models with varying ranks. As anticipated, the performance decays r decreases. In practice, setting = d /4 is preferable, as achieves a 50% reduction in both memory and FLOPS while maintaining high performance.

---

Published as a conference paper at ICLR 2025
Table 5: Ablation on LORO rank
Figure 4: Ablation on LORO exact update frequency K.
Training curves Valdation Curves LORO frankc 256) on LLaMA-130M
Model
LLaMA 60M
# Token
1.1 B
F-ll size
di do Method
512/ 1376
LORD. 250
121
Rank
PPL (1)
Param
Mcm (L)
19T
LORO K- 200
1060.
75a
LOGO,
Full-size
512
34.06
58
0.35
Lowerant
Low-rank
128
78.18
LORD. 500
LORO
256
31.54
0.31
LORD, K-
LOND, K- 2002
LORO
128
33.96
43
0.24
Lowrurk
LORO
64
39.40
0.22
LORO
32
49.88
35
0.21
4003
15000
Table 6: Extending LORO to fine-tuning RoBERTa-base models on GLUE benchmark. Results of
the baselines reported from (Zhao et al.
2024
Han et al. 2024).
RoBERTa-buse
Memory
CoLA
STS-B
MRPC
RTE
SST2
MNLI
QNLI
QQP
Avg
Full-size
747M
62.24
90.92
91.3
79.42
94.57
87.18
92.33
92.28
86.28
LoRA, r 4
257M
61.38
90.57
91.07
78.70
92.89
86.82
92.18
91.29
85.61
Galore,
253M
60.35
90.73
92.25
79.42
94.04
87.00
92.24
91.06
85,89
LORO,
257M
61.58
90.75
92.60
78.69
94.14
86.15
92.58
90.64
85.89
LoRA,
264M
61.83
90.80
91.90
79.06
93.46
86.94
92.25
91.22
85.93
Galore, =
257M
60.06
90.82
92.01
79.78
94.38
87.17
92.2
91.11
85.94
SLTrain, r =8
60.35
90.74
92.38
79.42
94.15
86.53
92.40
91.27
85.93
LORO, r 8
257M
63.10
90.97
93.12
78.92
94.61
86.65
92.75
91.00
86.39
Ablation on the LORO exact update frequency K. In Figure 4 we compare the convergence of different runs of LORO with varying exact update frequencies.
observed that for K > 750 when the LORO exact update is overly lazy, the training curve tends to explode and fail to converge to performing minima. In practice, we recommend setting K = 500, as significantly reduces training overhead and results in competitive low-rank models with performance comparable to full- size baselines across LLaMA models of varying sizes, ranging from 60M to 1B parameters.
5.3 EFFICIENT LOW- RANK FINE-TUNING
Following the experiment setup in ZZhao et al. 2024 Section 5.4), we extend our LORO to fine- tune the pretrained RoBERTa-ber base model (Liu et al. 2019) GLUE datasets (Wang et al. 2019). Specifically, we reparameterized each weight in the query- and key-projections as W + BA, where W is the full-size pretrained weight and B and A are low-rank factors. Then, we apply LORO to B and A exclusively and freeze the remaining parameters. The hyperparameters are detailed in Appendix A In Table 6) we compare the fine-tuning performance of LORO with three base- lines, including 'LRAA' (Hu et al. 2022), 'Gal (Zhao et al. 2024), and 'SLTrain" (Han et al.) 2024). Although LORO achieves decent performance, its outperformance over other baselines in fine-tuning is less significant than in pretraining. We argue that, with full-size pretrained weights provided, fine-tuning requires learning only a small amount of domain-specifie knowledge, which can be easily captured through low-rank adjustments in each weight matrix. As result, the potential for improvement with different low-rank techniques is limited, and the performance gains become less noticeable.
CONCLUSION
In this paper, we proposed LORO for efficient yet effective pretraining of low-rank parameterized language models. Specifically, LORO optimizes the full-size product of the low-rank factors toward the steepest gradient descent direction on the low-rank matrix manifold, without computing any ex- pensive full-size weight or gradient. Extensive experiments demonstrate that LORO can discover competitive low-rank models with performance comparable to full-size baselines, while providing significant memory reduction and acceleration in both training and inference. This work demon- strates the potential low-rank structures in language model pretraining, inspiring us to explore efficient low-rank pretraining methods for vision and multi-modal generative models.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
The research work described in this paper was conducted in the JC STEM Lab of Machine Learning and Symbolic Reasoning funded by The Hong Kong Jockey Club Charities Trust. Sinno Pan also thanks the support of the Microsoft Research Asia collaborative research grant.
REFERENCES
A. Absil and Oseledets. Low-rank retractions: survey and new results. Comput. Optim. Appl., 62(1):5-29, September 2015. ISSN 0926-6003, doi: 10. 10.1007/110889-01499977 URL
https: /top:::/doo... doi -t-..///......00/11008------1010 / -14-9714-1
P.A. Absil, R. Mahony, and R. Sepulchre. Optimization Algorithms Matrix Manifolds. Princeton
University Press, 2009. ISBN 9781400830244. URL
googl e com. com.
books?ld=NSQGRn
Ahmed Agiza, Marina Neseem, and Sherief Reda. Mtlora: Low-rank adaptation approach for effi- cient multi-task learning. In Proceedings of the IEEE/CVE Conference on Computer Vision and Paitern Recognition, pp. 16196-16200 2024.
Yonatan Amit, Michael Fink, Nathan Srebro, and Shimon Ullman. Uncovering shared struc- tures in multiclass classification. In Proceedings of the 24th International Conference Ma- chine Learning, ICML '07 pp. 17-24, New York, NY, USA, 2007. Association for Com- puting Machinery. ISBN 9781595937933.. /tdp..oog/1.1115//2222222 127344991.2
doi: 10.1145/127346.1227447
URL https:
Fengmiao Bian, Jian-Feng Cai, and Rui Zhang. A sinannddttioeeddemmnnmmmnn gradint descent algo- rithm for low-rank matrix recovery, 2023. URL sttps://axxiv..or//abaaaa 02543
Kevin Bleakley and Yoshihiro Yamanishi. Supervised prediction of drug target interactions using bipartite local models. Bioinformatics, (3118)333972403, September 2009. ISSN 1367-4803. doi: 10.1099/hoinfoomatie////..p. Sboinformatics/bteee
URL. https: ///e///////////// org 1t.g/11.109.3
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sand- hini Agarwal, and et al. Language models are few-shot learners.
H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural In- formation Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc.,
2020. URL alee/ file/1457cc 24570066000000971111188888888880
neurips ctt.cp/.ppppp//eeeeeeeeeeeeeeeee files//aper/202c/
Eric L.. Buehler and Markus J. Buehler. X-lora: Mixture of low-rank adapter experts, a flexible framework for large language models with applications in protein mechanics and molecular de-
sign, 2024. UtLptt http://aaxvvv... hbtpg//b//242...22 07148
Patrick Chen, Hsiang-Fu Yu, Inderjit Dhillon, and Cho-Jui Hsich. Drone: Data- aware low-rank compression for large nlp models. In M. Ranzato, A. Beygelz- imer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural In- formation Processing Systems, volume 34, pp. 29321-29334. Curran Associates, Inc.,
2021. URL httppss//prooeeddng.s. neurippp.....u.pp. cc wwtppper.ppppppeeee ceneerffileeeee files//appr/20211 File/ Fble://6eeett11C000eccCC11111112
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. ArXiv, abs/1604.00117 2016. URL
brg/CorpusID: 15865278
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-cond long-context large language models. ArXiv, abt/2339..12300 2023. URL
https://api remanticon sescnticscho ar..org/ooppus

---

Published as a conference paper at ICLR 2025
Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block- wise quantization. ArXiv, abs/2110.02861, 2021. htt.rsss/app..ssmmnniiccchhllll
#ttp/oorpusD::2840552
Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized llms. arXiv preprint arXiv:2305. 14314, 2023.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pp. 4171-4186 Minneapolis, Minnesota, June 2019. Association for Com- putational Linguistics. doi: 10.18653/11/11--1 1423. URL https: http:://aclaathologgy... org/ N19- 1423
Ning Ding, Xingtai Lv, Qiaosen Wang. Yulin Chen, Bowen Zhou, Zhiyuan Liu, and Maosong Sun. Sparse low-rank adaptation of pre-trained language models. In The 2023 Conference on Empir- ical Methods in Natural Language Processing, 2023. URL ptctps://cpenrevieet
Eorum?id=jxqa 7F /FE
Marawan Gamal and Guillaume Rabusseau. ROSA: Random orthogonal subspace adaptation. In
Workshop on Efficient Systems for Foundation Models ICML2203, 2023. URL https://
openreview nwt Seee:/oorumiddd44P91 pb63.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neu- ral networks. In International Conference on Artificial Intelligence and Statistics, 2010. URL https://api www.steissccootoo.......o
waag//orrusssD... : 5575601
Almog Gueta, Elad Venezian, Colin Raffel, Noam Slonim, Yoav Katz, and Leshem Choshen. Knowledge is region in weight space for fine-tuned language models. In The 2023 Conference on Empirical Methods in Natural Language Processing, 2023. htttpss///ppereereet net/ net/ forum?ormmmmmm vi BnrryPbb
Andi Han, Jiaxiang Li, Wei Huang, Mingyi Hong, Akiko Takeda, Pratik Jawanpuria, and Bamdev Mishra. SItrain: sparse plus low-rank approach for parameter and memory efficient pretraining, 2024. URL/attps://axxx org abs /20..24006..0 02214
Soufiane Hayou, Nikhil Ghosh, and Bin Yu. LoRA+: Efficient low rank adaptation of large
models. In Forty-first International Conference on Machine Learning, 2024. URL https:
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. In International Con- ference on Learning Representations, 2022. URL httt:s://ppereeiieeet /ownneviie.e...ccoooo.o net. forum?
id=nZee/KeeFYf
Chengsong Huang, Qian Liu, Bill Yuchen Lin, Chao Du, Tianyu Pang, and Min Lin. Lorahub: Efficient cross-tass generalization via dynamic loRA composition. In R0-FoMo:Robustner of Few-shol and Zerv-shol Learning in Large Foundation Models, 2023. URL https:// opeerrvieww net. ftntt/formmid=ll/// I
Ajay Jaiswal, Lu Yin, Zhenyu Zhang, Shiwei Liu, Jiawei Zhao, Yuandong Tian, and Zhangyang Wang. From galore to welore: How low-rank weights non-uniformly emerge from low-rank gradients. 2024. URL //atps:///axii. org/ abs//bb//2// 2407 11239
Siddhartha Rao Kamalakara, Acyr Locatelli, Bharat Venkitesh, Jimmy Ba, Yarin Gal, and Aidan N. Gomez. Exploring low rank training of deep neural networks, 2022. URL https:///axrtt
org/abs/220S abaaabs20000 13569
Mikhail Khodak, Neil A. Tenenholtz, Lester Mackey, and Nicolo Fusi. Initialization and regulariza- tion of factorized neural layers. In International Conference on Learning Representations, 2021. URL
JT
12

---

Published as a conference paper at ICLR 2025
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA. May 7-9, 2015, Conference Track Proceedings. 2015. U //arxiv.org/abait 612.69800666
Othmar Koch and Christian Lubich. Dynamical low-rank approximation. SIAM Journal on Matrix Analysis and Applications, 29(2)434-454, 2007. doi: 10.1137/050039703. URL https:// doi. /tt.:o///1..113705069999999
Soroush Abbasi Koohpayegani, Navaneet K L, Parsa Nooralinejad, Soheil Kolouri, and Hamed Pirsiavash. NOLA: Compressing loRA using linear combination of random basis. In The Twelfth International Conference on Learning Representations, 2024. URL lttpps:///ppeneeeeees Peet/Eorunn.dd=T
J.M. Lee. Introduction to Riemannian Manifolds. Graduate Texts in Mathematics. Springer Interna- tional Publishing, 2019, ISBN 9783319917542. URL httpt:tttpoooottoo http://ooooooo..oll.o... com. sg/
Adrian S. Lewis and Jérôme Malick. Alternating projections on manifolds. Math. Oper. Res., 33 (1):216-234, February 2008, ISSN 0364-765X.
doi: 10. 1287/morr (t077//mor..00..2911 URL https:
1287/moor 1070. 0291
Bingrui Li, Jianfei Chen, and Jun Zhu. Memory efficient optimizers with 4-bit states. In
Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https:
// //oeereeveew net/forum? d=nN8
Vladislav Lialin, Sherin Muckatira, Namrata Shivagunde, and Amna Rumshisky. ReloRA: High- rank training through low-rank updates. The Twelfth International Conference on Learning Representations, 2024. URL htRpsttt::
net/ htt.//orrumm---...nnppppp.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zeutlemoee, and Veselin Stoyanov. Roberta: A robustly optimized BERT pretraining approach. CoRR, abs/1907.11692, 2019. htL httptttpppaaa......tt///o abs 1907. 2bbbb/1999.11222
James Lucas, Juhan Bae, Michael R. Zhang, Stanislav Fort, Richard S. Zemel, and Roger Baker Grosse. Analyzing monotonic linear interpolation in neural network loss landscapes. ArXiv, 00s/2104.11044, 2021.
org/CorppsIID: URL https::/ www.seicscnnnoooooeo...e. http:://aai.eeenniiiiccoooooooo httpaa/iicccco...oooo////oooooop
233346934
Bamdev Mishra, Gilles Meyer, Silvère Bonnabel, and Rodolphe Sepulchre. Fixed-rank matrix fac- torizations and riemannian low-rank optimization. Comput. Stat., 29(3-4):511-211, June 2014. ISSN 0943-4062.
doi: 10. 10007//001880013-46666 URL https: nttps:///ci.org//11 s0C 000180-013-0
Guillaume Olikier, André Uschmajew, and Bart Vandereycken. Gauss-southwell type descent methods for low-rank matrix optimization. ArXiv, abs/2306.00897, 2023. URL https:
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Ed- ward Yang, Zach De Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. PyTorch: an imperative style, high--performnnce deep learning library. Curran Associates Inc., Red Hook, NY, USA, 2019.
Alec Radford and Karthik Narasimhan. In Improving Language Understanding by Generative Pre- Training, 2018. URL BtLL:tts::/// http://aapp.s.paaiieeeaeee semanticscho ar org http://orppiDDDD 49313245
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou. Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. CoRR, 1tt//1900.108.... 2019. http:httttt/tttt wrw. abs 910. 10683.
13

---

Published as a conference paper at ICLR 2025
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: memory optimizations toward training trillion parameter models. In Proceedings of the International Conference for High Performance Computing. Networking, Storage and Analysis, SC '20. IEEE Press, 2020. ISBN 0117221199866
Amrutha Varshini Ramesh, Vignesh Ganapathiraman, Issam H. Laradji, and Mark Schmidt. Block- LLM: Memory-efficient adaptation of LLMs by selecting and optimizing the right coordinate blocks, 2024. URL https: opw.review..eeee net http://ww...t//rommmmomaoooo-o/o
David E. Rumelhart, Geoffrey E. Hinton, and Ronald J. Williams. Learning representations by back- propagating errors, pp.
MIT Press, Cambridge, MA, USA, 1988. ISBN 0262010976.
Dayana Savostianova, Emanuele Zangrando, Gianluca Ceruti, and Francesco Tudisco. Robust low- rank training via approximate orthonormal constraints. In Thirty-seventh Conference Neu- ral Information Processing Systems, 2023, URL
NJPSVVOU3R
Steffen Schotthöfer, Emanuele Zangrando, Jonas Kusch, Gianluca Ceruti, and Francesco Tudisco. Low-rank lottery tickets: finding efficient low-rank neural networks via matrix differential equa- tions. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL httpthtss///ppeeetitet www.:oruttfoomoooooo
Id=IILJOKWZMy
Noam Shazeer. GLU variants improve transformer. CoRR, abs/2002.05202, 2020. URL https:
//arxxvvorq/////aa.... wtt.:/rg/aabsa/oo..ooo 05202
Yang Sui, Miao Yin, Wanzhao Yang, Yu Gong, Jinqi Xiao, Huy Phan, Ding Ding, Xiaozhong Xu,
Shan Liu, Zhenzhong Chen, and Bo Yuan. ELRT: Towards efficient low-rank training for compact neural networks, 2023. URL /RLL/tttps///pre.reeeee....etttt www.nneee/oooummmmm
9m8bB m
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, and et al. Gemma: Open models based gemini research and technology, 2024a. URL htt.t:ttt////iiiiii
abs/ 082955
Gemma Team, Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhu- patiraju. Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, and al. Gemma 2: Improving open language models at a practical size, 2024b. btttts://aaaattt
org/abs/2405 00118
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, and et al. Llama: Open and efficient foundation language models, 2023a. URL
13971
Hugo Touvron, Louis Martin, Kevin Stone. Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, and et al. Llama Open foun- dation and fine-tuned chat models, 20236. htt.:tss::///tt //axxi org
Mojtaba Valipour, Mehdi Rezagholizadeh, Ivan Kobyzev, and Ali Ghodsi. Dylora: Parameter ef- ficient tuning of pre-trained models using dynamic search-free low-rank adaptation, 2023, URL
https::/aaxxis org/ abs 2210.022220 07558
Bart Vandereycken. Low-rank matrix completion by riemannian optimization. SIAM Journal on Optimization, 23(2):1114-123, 2013. doi: 10. 11371 10845768. URL https:///o. http::://dooooogg/... org/
10.1137/110845763
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gameet (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL 1tt.://aae233S/eeeeeeooooaaaaaaa
fi 11le/33fee24354ddeee)
14

---

Published as a conference paper at ICLR 2025
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: multi-task benchmark and analysis platform for natural language understanding. In International Conference Learning Representations, 2019. URLpsttss///ppeeeee///
net:
Xin Wang, Yu Zheng, Zhongwei Wan, and Mi Zhang. Svd-Ilm: Truncation-aware singular value decomposition for large language model compression, 2024. URL https://axxiv.oq//
Noo/2403.07377
Wenhan Xia, Chengwei Qin, and Elad Hazan. Chain of lora: Efficient fine-tuning of lan-
guage models via residual learning. ArXiv, abs/2401.04155 2024. URL https://api.
semanticschole org/ Corpus ID: 266899736
Ming Yuan, Ali Ekici, Zhaosong Lu, and Renato Monteiro. Dimension Reduction and Coefficient Estimation in Multivariate Linear Regression. Journal of the Royal Statistical Society Series B;
Statistical Methodology, 69(3):32--346, 05 2007. ISSN 1369-7412. doi: 10.11 1467-9868. 14,11111,.466-9888 9868 2007 00591.x 200.08088.200000.11....
/tt:: 2007.00591.x URL wttp::://oo..og//....o.. wttp:/i..r./.11111111.1111111.11
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Curran Associates Inc., Red Hook, NY, USA, 2019.
Fangzhao Zhang and Mert Pilanci. Riemannian preconditioned loRA for fine-tuning foundation models. In Forty-first International Conference Machine Learning, 2024. URL https: //openreview.
Qingru Zhang, Minshuo Chen, Alexander Bukharin, Pengcheng He, Yu Cheng, Weizhu Chen, and Tuo Zhao. Adaptive budget allocation for parameter-efficient fine-tuning. In The Eleventh In- ternational Conference on Learning Representations, 2023. URL
Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates gain more. In Workshop Efficient Systems for Foundation Models II @ JCML2024, 2024a.
net / forum?
id=TonkMNfx9d
Zhenyu Zhang, Ajay Jaiswal, Lu Yin, Shiwei Liu, Jiawei Zhao, Yuandong Tian, and Zhangyang Wang. Q-galore: Quantized galore with int4 projection and layer-adaptive low-rank gradients, 2024b. URL httttt://ttxii..ttts
abs/2407 wttps///.rrr/b//2....ooo 08296
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient Ilm training by gradient low-rank projection, 2024.
15

---

Published as a conference paper at ICLR 2025
A
IMPLEMENTATION DETAILS
A.1 PRETRAINING LLAMA MODELS C4 DATASET
The implementation details introduced in this subsection apply the experiments in Table 1 Table 2 Table 3 Table 4 Table 5 Figure y and Figure 4 To achieve a fair comparison with other follow the same model architecture and pretraining hyperparameters as in (Zhao et al.] 2024| Appendix C.I). For all runs, we set the max data sequence length as 256 with a batch size of 512 (i.e., a token batch size of 131K). For fair comparison in terms of the pretraining efficiency, we run all the experiments on 1x NVIDIA 40G A100 GPU, without using model parallelization or data parallelization. To achieve the 512 training batch size, we adjust the per-device batch size and vary the number of gradient accumulation steps according to the size of the LLaMA models. Based on (Zhao et al., 2024 Appendix C.I), we employ a learning rate warmup starting from 0 during the first 10% of the pretraining steps, followed by cosine annealing scheduler that decays to 10% of the maximum learning rate. We initialize the low-rank factors with Xavier initialization (Glorot& & Bengio 2010) across all experiments.
pretraining baselines,
We implement LORO following the pseudo-code attached in 2 Following (Zhao et al. 2024 Tou- vron et al. 2023b), impose low-rank parameterization to all the weight matrices in the self- attention modules (including q-proj, k.proj, v-proj, and o-proj) and linear projection mod- ules (including down_proj, up up-proj, and gate.pro and we update them with LORO, while we apply the standard Adam optimizer to update the token_emb, 1m_head, and other layer nor- malization modules. All modules share the same learning rate. Specifically, for all LORO pretrain- ing results reported in Section [.1] we set the LORO exact update frequency to K = 500 and the learning rate to 0.01 for the LLLMA-60M, -130M, and -350M models, while for ILLMA- BB we set K 200 and the learning rate to 0.005. To stabilize the training process, at each LORO exact update step, we apply a 5-step linear learning rate warmup and we refresh the Adam statistics.
A.2 FINE-TUNING ROBERTA MODELS ON C4 DATASET
For fair comparison on fine-tuning with other baselines, we follow the same model architecture, datasets and evaluation processes as in (Zhao et al. 2024, Appendix D.1). We report Matthew's correlation score for CoLA, the Pearson correlation score for STS-B, matched and mismatched accuracy for MNLI, Fl score for MRPC, and accuracy for other tasks. Specifically, we fine-tune the pretrained RoBERTa-base checkpoint from Huggingface on the GLUE benchmark. For each weight matrices in q-proj and k_proj, adopt the standard LoRA parameterization W ABA, where W denotes the full-size pretrained weight, the LoRA scale hyperparameter, and and A are the trainable low-rank factors. During fine-tuning, we apply LORO to update the low-rank factors in the same manner as in pretraining, while keeping all other parameters unchanged. The training hyperparameters are attached in Table
Table 7: Model architecture and pretraining settings. dN and d2 denotes the hidden- and ciermeeiite-dimensio the LLaMA model. Tokens' denotes the amount of training tokens, PPD-BZZ' denotes the per-device batch size, and '# Grad Accum denotes the number of gradient accumulation steps.
Model
Params
d1
Heads
Layers
Steps
Tokens
PD-BSZ
Grad Accum
60M
58M
512
1376
10K
1.1B
256
2
130M
134M
768
2048
12
20K
2.2B
128
Shallow
94M
820
1720
12
20K
2.2B
128
4
Slim
94M
640
1250
12
20K
2.2B
128
4
Slim MLP
94M
768
928
12
20K
2.2B
128
4
368M
368M
1024
2736
16
60K
6.4B
64
8
185M Shallow
185M
1024
1880
16
60K
6.4B
64
8
185M Slim
185M
736
1640
16
60K
6.4B
64
8
185M Slim MI.P
185M
1024
940
16
60K
6.4B
64
8
1339M
2048
5461
24
32
100K
13.1B
32
16
16

---

Published as a conference paper at ICLR 2025
Table 8: Hyperparameters of LORO in fine-tuning RoBERTa experiments.
CoLA
STS-B B
MRPC
RTE
SST-2
MNLI
QNLI
QQP
Batch Size
64
16
16
16
16
16
16
16
# Epochs
30
40
20
30
40
30
30
30
Learning Rate
4e-4
4c-4
2c-4
5e-4
5e-5
5e-5
4e-4
5e-5
LoRA
32
8
8
8
LORO Freq K
100
200
100
100
100
100
100
100
Batch Size
64
16
16
16
16
16
16
16
# Epochs
30
40
20
30
40
30
30
30
Learning Rate
4e-4
4e-4
2e-4
5e.4
5e-5
5e-5
4e-4
5e-5
LoRA
16
16
8
8
16
16
LORO Freq
20
100
100
100
500
200
100
200
A.3 PYTORCH IMPLEMENTATION OF LORO
Algorithm 2 Low-rank Riemannian Optimizer (L.RO)) in PyTorch (Paszke et al. 2019)
def LORO_step (B, k3)
[d, r1, [d, r] if st_step d, 3.shape
!-
Approximate Riemannian update
B ir
d) B.grad
à ir (r d) g.grad
else:
dB, di B.grad, g.grad 3 dB: [d, x), dı: [r, d]
Qb, Rb torch qr (5) & Ob: (d, r], Rb: [r, r]
Qa, Ra torch.qr (A) Qa: (d, r], Ra: (r, E]
dB Pa_ini [d, Fl, dA_ Rb_inv: [d, ri
Ra
torch. selve (Ra. dB.T).I
Rb
torch. solve (Rb. d
0b_: [d, rl, Rb_: [r, r), Q=_: [d, ri, Ra_; [r,
Qb_, Rb_ torch.qr qr (dB
(dB__dBB______ Ra_ ob @b. dB Ra__a_iov)
Qa Ra_ torch. Rb__ib - Qa Qa.T Rb__
= Rb 8 Ra.T 1r di RR_Rt_inv.T
Qa [r, r]
$ sVD on f22, 2r1, U: (2r, zl, sig: fr,el, V: [2r,2]
U, Sig, V torch. (mat ([IS, ir Ra_,T), [- ir * Rb, zeros]l))
mat [[bb, @b 2 siq.sqrt. () sig. sqrt ()
(10a, Qa_1) Exact Riemannian update
return A
B PROOF DETAILS
Proposition 3 (Riemannian Gradient under Low-rank Parameterization), Following the previous
notation, for any low-rank parameterized weight W = BA E Mr. where B E gdxr and e grxd the Riemannian gradient of any smooth function L(-) Rdxd + evaluated W is
grad L(W) B VAL. VBL. (AAT) B
B(B B VAL. (AAT)- (12)
Proof of Proposition We first need to show that, for any ambient point G E Rdxdd the orthogonal projection to the tangent space M, is given by
P(G) B(B -B G A B BB G GAATAATT--A
Then, the proof can be completed by taking G = V L,, and showing that Vw AT = VBL. and B V L.
According to Proposition I LLee 2019
Example 8.14) and Absil et al. 2009 Equation 3.37)), the UU G + G- VVT UU G VV where W = UEV is
othorgonal operator P(G) the singular value decomposition (SVD) of W. Suppose the QR decompositions of B and A
17

---

Published as a conference paper at ICLR 2025
QgRg- B and QARA A and the SVD of the full-rankr-by- matrix RgR is U ER V Then, it holds that U ( ( column-wise orthogonal and ΣR _x xr diagonal. By the uniqueness SVD, holds that QBUR U and VAVR V. which implies
UEVT, where (QBUR). (V) E Ra are
UU RUUURURQU R Qs (QB
QB QRQR
Q
(13)
=(QBRB (R Q QBR (R B
(14)
Similarly, we have tion above proves that the row-space of B coincides with that of U, and the column-space of A coincides with that of V. Therefore, the orthogonal projection the tangent space can be reparam- eterized under the basis of and A.
which completes the proof. In intuition, the deriva-
Proposition (Riemannian Update under Low-rank Parameterization). Following the previous notations, for any low-rank parameterized weight W = BA E Mp. let B QnRB and A aa be the QR-decomposition of and A such that QB. QA E R are column- wise orthogonal, and R.B RA E RR are upper-triangllar Then the Riemannian updated weight, defined by W SVD,(W n grad L(W)), admits a low-rank factorization W = B'A', with B' (QB, QB) U.E! and A' = Σ I (QA, Qa) -peciiicalll, U.E.V. is determined by the rank-r SVD of the following 2r- matrix, i.e.
U.E.V. = SVD, (_
S
-1
E
R 2rx2
RERR R VAL QAERTM (15)
RB
and Qg Rg is the QR decomposition of the d-by-r matrix (V - Q A is the QR decomposition of the d-by-r matrix A R QQQ V L
QB Q
VBL R.)) while
Proof of Proposition 4 This proof follows a similar spirit of orthogonal tangent vector decompo-
sition as in (Vandereycken, 2013 Section 3) and (Bian et al., [023] Section 2.2.3). To complete the proof, we need to perform such decomposition under the unnormalized basis spanned by B and A, such that the full-size weight and gradient terms be canceled out using the relation V L. A = VLL and B VwL VAL.
As the low-rank representation of the Riemannian gradient is provided in Proposition 3 only need to manipulate the terms in the updated weight W grad L(W) such that can be de- composed as UEVT, where U, V E Rdd22r are column-wiss othorgonal matrices. Thereby, the Riemannian SVD retraction can be performed on a smaller 2r-by-2r matrix instead, since
SVDV( - VT.
(16)
To complete the proof, we only need to identify the representation of Ü, Σ, V by orthogonal tangent vector decomposition, under low-rank parameterization. Based on Proposition 3 it holds that
grad L(W)
=B(B B) TT TW + V A A AAAA(AATTAA - B BB VwL. ATAAAA
=Q=Q] VwL V QT QBQ TWL-QqQ.
-Q V + (1- Q=Q VwL QxQ. + QaQ VwL (I Q4))).
where the latter two terms are orthogonal the column space of and A T respectively. To cancel the full-size gradient terms, we leverage the relations Q VwL= = R T and VwL-QA= VBL R 1 which shows that
grad L(W)
-Q VAL QAQT + +(I -Q2QT) VBL R QT Q VAL -I VVL-(I-QAQ)) (I-QAQ.).
18

---

Published as a conference paper at ICLR 2025
Therefore, the updated weight can be represented by
W 1) grad L(W)
Qn (R (QnR V QQQQ4
+1-Qu VBL RA" Q7 QgR VAL (I-QAQI)
Q (RRRA- R VAL Q4) Q Rg Q
QB R
Q
=QB QB) R n R VAL QA
1 R
-7 R
Notice that (QB, Qg) and (QA, Q,) are column- orthogonal matrices, making them desir- able choice of the aforementioned U and V. and hence completes the proof.
C FURTHER DISCUSSIONS ON THE MECHANISM OF LORO
In this section, provide supplementary analysis and discussion of the LORO mechanism. Specifically, Section C.T analyzes the limitations of standard low-rank pretraining, while Section C.2 discusses how LORO effectively overcomes these limitations, leading to significant improve- ments.
C.I LIMITATIONS IN STANDARD LOW-RANK OPTIMIZER
In this subsection, present additional theoretical and empirical evidence demonstrating that the standard low-rank optimizer (i.c. independently optimizing low-rank factors) is inefficient for achieving satisfactory training performance.
In Section C.I.I our empirical evidence demonstrates that: 1. the standard low-rank optimizer produces spiky training curves and high perplexity across various initialization schemes and learning rates, 2. it learns low-rank weights with high condition numbers (i.e., the ratio of the largest singular value to the smallest nonzero singular value), and 3. it continues to yield poor performance even when combined with optimizer state refreshing.
In Section C.T.2 our theoretical analysis indicates that: the standard low-rank gradient in (3) fails to preserve adequate information from the full-size gradient within the tangent space, and 2. nearly-stationary point obtained through standard low-rank gradient descent with a near-zero gradi- ent norm may be significantly far away from true stationary point on the manifold.
Based on these observations and analysis, we conjecture that the failure of the standard low-rank optimizer stems from the following factors: 1. the standard low-rank gradient contains redundant full-size gradient information, preventing it from accurately approximating the steepest descent di- rection the manifold, 2. the lack of proper scaling in the standard low-rank gradient causes de- viations from the steepest descent direction on the manifold, and 3. the standard low-rank gradient biases towards ill-conditioned weights, leading to instability during training.
C.1.1 EMPIRICAL. EVIDENCE
Standard low-rank optimizer generally leads to poor performance. To validate that pretrain- ing low-rank language model is generally hard. we train a rank-128 LLaMA-60M model with standard low-rank optimizer under different initialization schemes and learning rates. Specifically, we initialize the low-rank factors with Xavier initialization (Glorot & Bengio 2010), or standard Gaussian initialization with standard deviation varying in {0.1, 0.01, 0.001]. and we tried different learning rates in {0.01, 0.005, 0.001}. We report the training curve and the final perplexity of the models.
As shown in 5 in all settings, the training curve of the vanilla low-rank optimizer is spiky and unstable, leading unsatisfactory evaluation perplexity. This suggests that independent updates of the low-rank factors are insufficient for good performance across different configurations. Training low-rank language model match the full-size baseline is generally challenging.
19

---

Published as a conference paper at ICLR 2025
1w.. 2.80
3.85 1108..1
10
331
Figure 5: Ablation study the impact of learning rate and initialization schemes standard low- rank training for a rank-128 LLaMA-600 -0M model. The y-axis (log-scale) represents perplexity, with the gray horizontal dashed line indicating the performance of full-size training.
raaiing ourves and VaNo303n Curves LUMO VOrK
125 LLOING 0UM
10
Taining Cave
Val Less CLaVE
(--QUX for
DCS Pig
+333
GE2 Eig scc
000% mny
1.0
714
1.80
1100
241
Pnatraning StADE
Peoraaiing Ptipe
Figure 6: Ablation study on the impact of optimizer state refreshing on standard low-rank training for a rank-128 LLaMA-60M model. The y-axis (log-scale represents perplexity, with the gray horizontal dashed line indicating the performance of full-size training.
Standard low-rank optimizer leads to poor performance even when combined with optimizer state refreshing. To assess whether optimizer state refreshing can enhance the standard low-rank optimizer, we perform an ablation study by varying the refreshing frequency in the standard low- rank optimizer. Following the setup in Section 5.I trained rank-128 LLaMA-60M models using vanilla low-rank gradient descent with different learning rates (chosen from {0.01, 0.005, 0.001)) and optimizer state refreshing frequencies (chosen from {0, 100, 500, 1000)).
As shown in Figure 6. all trials exhibit spiky training curves and unsatisfactory performance. This validates that applying optimizer state refreshing to standard low-rank gradient descent still leads to bad training performance. We conjecture that this is because the standard low-rank gradient does not actively explore new optimization subspace, and the training process biases towards low-rank weights with increasing condition numbers. In this case, refreshing the past momentum statistics resided does not help in exploring beneficial optimization subspace.
Standard low-rank optimizer learns ill-conditioned weights. To understand the cause of the instability and spiky training curves in the standard low-rank optimizer, we inspect the condition number and norm of the low-rank factors throughout the training process. Specifically, we train rank- 128 LLaMA--0M models with Xavier initialization and learning rate of 0.01, using different optimizers: the standard low-rank optimizer, the full-size Adam optimizer (Kingma & Bal 2015), and LORO. We visualize the evolution of the singular values, Frobenius norm, and condition num- bers of the learned weights at various layers throughout the training process.
As shown in Figure 13, Figure 14 and Figure [3] the standard low-rank gradient results in model weights with significantly higher condition numbers compared to LORO or full-size Adam. We conjecture that the standard low-rank gradient update introduces imbalances among the singular values of the low-rank factors, thereby impairing the stability of the feature learning process.
C.1.2 THEORETICAL EVIDENCE
Standard low-rank gradient fails to approximate the full-size gradient well. To understand the difference between the training dynamic of the standard low-rank optimizer and its full-size
20

---

Published as a conference paper at ICLR 2025
counterpart, we examine the gap between the standard low-rank gradient and the full-size gradient. Following (3), we define the standard low-rank gradient as Z - BB" V + V A, where V is the full-size gradient and W = BA is the low-rank weight. We show that Z fails to adequately preserve the full-size gradient information. In contrast, LORO adopts the Riemannian gradient, which provides the best approximation of the full-size gradient within the tangent space Mw,,, thereby preserving the full-size gradient information more effectively than the standard low-rank gradient.
According to Proposition 1 and Proposition S the orthogonal projection onto the tangent space is
P(G) = UU
UGG)=U'' G GVV - -UU GVVT. = PBG + GPA - PPPPAP
(17)
where PB A BB and PA A AT(AA))) are the orthogonal projection matrices onto the column space of B and the row space of A. We first study how well the full-size gradient can be approximated within the tangent space w.r.t Frobenius norm -l... For any tangent vector G resides in the tangent space, the full-size gradient approximation error equals to
|G Lw || P)V + [P-(Vw
(18)
where (.) is the orthogonal projection to the complemented subspace of the tangent space. This implies that the full-size gradient approximation error is minimized at G L), which is exactly the Riemannian gradient as shown in Proposition 3. Therefore, the Riemannian gradient represents the steepest gradient descent direction on the manifold, as provides the best approxi- mation of the full-size gradient within the tangent space.
Now, we can study how well the standard low-rank gradient Z approximates the full-size gradient. Since Z is an element in the tangent space (can be proven by showing P(Z) = Z), the full-size gradient approximation error of Z is determined by IZ L)W, which equals
[[[((B T B) I)B V + VW (AAA(((AAT-- I)A + PB WP All (19)
This error arises from two main factors: the scale of gradient components, BB V. and Vw T.A A, is not properly normalized to match that of PBG and GP and 2. the redundant gradient information, PB V --whch represents the gradient component in the intersection of the column space of B and the row space of A-is double-counted across the two sketched gradient terms. Our conclusions align with the findings in (Zhang et al. 2024a), which suggests that down scaling the learning rates of some specific modules (e.g., the low-rank weights) is beneficial.
Therefore, instead of updating B and A independently, we should employ LORO to update the low- rank weight W (BA) in the direction of the negative Riemannian gradient, which represents the steepest descent direction on the manifold. This approach ensures that the full-size gradient information is maximally exploited.
Nearly-zero standard low-rank gradient norm does not indicate a true stationary point. To understand the performance gap between the standard low-rank optimizer and its full-size counter- part, we study and compare their stationary point conditions. As mentioned in (Absil et al.) 2009) Vandereycken, 2013; Olikier et al., 2023), in Riemannian optimization, nearly-zero Riemannian gradient norm indicates that the parameter is close to a stationary point on the manifold, showing that the training process nearing convergence.
When using the standard low-rank optimizer, however, while the norm of the standard low-rank gradient is nearly zero, the learned parameter can significantly deviate from true stationary point on the manifold. This is because
I[ I) = LA +
I BV > (a,(A)² + o,) L
(20)
where (-) denotes the r-th largest singular value (i.e. the smallest nonzero singular value). This further implies
IP(Vw (a, (A) + ta, (B)""
(21)
Consequently, when B and A are ill-conditioned (i.e., when the condition numbers a (A)oo,(A) and 01 (B)/o, (B) are large) or nearly singular (i.e., (A) and a, (B) are small), the Riemannian gradient norm can remain large even if the norm of the standard low-rank gradient is nearly-zero, This indicates the inferior convergence of the standard low-rank optimizer. To mitigate this issue, adopt the Riemannian gradient to optimize the product of the low-rank factors.
21

---

10000 12500 15000 17500 20000 Published as a conference paper at ICLR 2025
C.2 EFFECTIVENESS LORO
In this section, we present additional theoretical and empirical evidence showing that LORO can overcome the limitations of standard low-rank training. Furthermore, we show that both the exact and approximate LORO updates are necessary and effective.
In Section C.2.1 our empirical evidence shows that: 1. the exact LORO update is necessary and effective, the approximate LORO update necessary and effective, LORO benefits from optimizer state refreshing.
In Section C.2.2. our theoretical evidence shows that: the exact LORO update balances the singular values of low-rank factors, and 2. the exact LORO update potentially reduces the full-size gradient approximation error the successive approximate LORO updates.
Based on our observations and analysis, argue that the performance improvement stems from the LORO update itself, because: 1. LORO approximates the Riemannian gradient better than the standard low-rank gradient, and it benefits from utilizing more information from the full-size gradient, 2. LORO effectively explores new optimization subspace which reduces the condition number of the low-rank weights, leading the accumulative error during successive approximate LORO steps.
stable training process, and 3. LOROO minimizes
C.2.1 EMPIRICAL EVIDENCE
The exact LORO update is necessary and effective. To validate the effectiveness of the exact LORO update, we conduct a finer ablation study on the exact LORO update frequency. Following Section 5.1] train rank-256 ALaMA-130m models with K E {10, 50, 100, 250, 500, 650, 750). We visualize the training curve and report the final perplexity.
As shown in Figure spiky and the evaluation perplexity exhibits a slight improvement. However, when the LORO exact update is overly lazy, iee. K > 650, the learning curve explodes. Our results show that while the exact LORO update is necessary, it does not need to be applied at every update step. In practice performing the LORO update every 500 steps achieves satisfactory performance.
the LORO exact update is getting more frequent, the loss curve is less
Training curves and Validation Curves of LORO (rank 256) LLaMA-130M
Full-iize
LORD, K 650
Trairing Perpicatt
104
Full-sizo
Validation Perpleni pelllittyy
45
LORO. K 250
LORO. 500
LORD. K- 500
40
LORO, K 100
10x
LORD. K=
LORO. K
LORD,
LORD. 50
LORO.X=10
10p
30
101
25
500
1000
1500
2000
2500
3000
2500
5000
7500
Pretraining Stops (0, 3000]
Pretraining
Figure 7: Ablation study on the impact of exact LORO update frequency K on LORO for training rank-256 LLaMA-130M models. The y-axis (log-scale) represents perplexity, with the gray hori- zontal dashed line indicating the performance of full-size training.
The approximate LORO update necessary and effective. To validate the necessity of using the approximate LORO steps, train rank-256 LLaMA-130M models with LORO, both with and without learning rate down-scaling between successive exact LORO updates.
As shown in Figure the training curve either fails to converge yields poor results when the approximate LORO steps are omitted. These results imply that the exact LORO update is necessary and effective. The performance improvement stems from the exact LORO update. In practice, LORO with K = 500 can achieve satisfactory training performance without the need to apply the exact LORO update at every step.
22

---

Published as a conference paper at ICLR 2025
Training curves and Validation Curves of LORO (rank = 256, K 500) on LLaMA-130M
108
Low-rank
10a
LORO we approx
training Perploxty
10a
..... LORO approx
Low-rank
10a
LORD apprrx
500
1000
1500
2000
2500
3000
2500
5000
7500
10000
12500
15000
17500 20000
Pretraining ning Steps in [0, 3000]
Pretraining Steps
Figure 8: Ablation study on the impact of learning rate down-scaling in approximate LORO updates for training rank-256 LLaMA-I 30M models. The y-axis (log-scale) represents perplexity, with the gray horizontal dashed line indicating the performance of full-size training.
LORO benefits from optimizer state refreshing. To study how optimizer state refreshing affects LORO, conduct ablation study the optimizer state refreshing. Following the setup in Sec- we compare the training curve and validation perplexity of rank-256 LLaMA- 1LLMA-330M models trained with LORO, both with and without refreshing the optimizer states periodically. As shown in Figure 9 both configurations lead to satisfactory performance. While refreshing the optimizer states introduces spikes in the early stages of training, eventually improves the validation loss of the LORO results. This suggests that LORO benefits from optimizer state refreshing.
tion 5.T
We conjecture that the stability of LORO against optimizer state refreshing arises because the exact LORO step shifts the low-rank factors (B, A) to a new optimization subspace with a reduced condi- tion number. This conjecture is partially supported by the analysis in Section C.2.2) In this scenario, the optimizer state retains momentum statistics accumulated many steps ago, originating from an obsolete subspace that significantly deviates from the current one. This indicates that dropping the old momentum statistics can help reduce accumulated error during the approximate LORO steps and encourage exploration in the new subspace.
Training curves and Validation Curves of LORO (rank = 256, K 500) on LLaMA-130M
10a
Full-sizo
Full-size
LORO win refeshing
60
LORD refeshing
Tranning Perplexty
104
LORO rafreshing
lolidation Perplexiey
LORD w/ refrashing
10a
45
10r
30
10a
500
1000
1500
2000
2500
3000
2500
5000
7500
10000
12500
15000
17500
20000
Pretralinin ning Stops [0, 3000]
Pretraining Steps
Figure 9: Ablation study on the impact of optimizer state refreshing on LORO updates for training rank-256 256 LLAMA-130M models. The y-axis (log scale) represents perplexity.
C.2.2 THEORETICAL EVIDENCE
Exact LORO update balances the singular values of low-rank factors. discussed Section C.I.2 in contrast to the standard low-rank gradient, LORO adopts Riemannian gradients, which pro- vides the best approximation of the full-size gradient within the tangent space. To better understand the mechanism of the exact LORO update, need to inspect how the exact LORO update affects the singular values of low-rank weights.
We show that the exact LORO update implicitly balances the singular values of the low-rank factors. According to (10) and Algorithm the exact LORO update ensures (B( az (A) Vo (W),
23

---

Published as a conference paper at ICLR 2025
for i = 1,..., 7. Therefore, the exact LORO update minimize each (B) (,)2 among all the low-rank factors (B. that satisfy 0, (B)) = o,()). This implies LORO prevents (A) + o, (B) to be overly small, and it can potentially reduce the condition number of the weights and stabilize the training process. As discussed Section smaller condition numbers of B and A also indicate better convergence property.
Exact LORO update potentially reduces error in approximate LORO updates. As discussed in Section C.I.I the exact LORO plays a subtle yet crucial role in successful low-rank training. To understand how few exact LORO steps can stabilize the whole training process, we study how an exact LORO step affects the successive approximate LORO steps.
Resuming from the derivations above, we show that the singular value balancing effect in the exact LORO step minimizes an upper bound of the full-size gradient approximation error. Specifically, the full-size gradient approximation eтror can be bounded by
[ B) - I)B' VwL + LW ((AAT) I)A + PRVWLPAlle ||P (VwL))l <[|| -bble + +II A ( V + LIIF;
(22)
where As denote the singular value matrices of B and A. Suppose Aw is the top-r singu- lar values matrix of the low-rank weight W = (BA), among all the choices of (B, A) satisfying
АgdA- Aw. the bound (ijI le + I|l l is minimized when AB AA 1/2 This
condition satisfied by the exact LORO steps, implying that the exact LORO update potentially re- duces the accumulative full-size gradient approximation error in the successive approximate LORO
updates.
In summary, conjecture that LORO benefits from two key factors: 1. The exact LORO update explores the new subspace by balancing the singular values of the low-rank factors. Thereby, reduces the condition numbers, leading smaller accumulative gradient approximation errors in the successive approximate LORO updates. Section C.2. provides supportive empirical evidence for this conjecture. The approximated LORO update downscales the standard low-rank gradients, preventing the rapid accumulation of redundant gradient information and avoiding undesired loss spikes before the next exact LORO update.
D FURTHER ABLATION STUDIES AND DISCUSSIONS
Ablation study on the LORO scaling factor. To study how the learning rate scaler affects the approximate LORO steps, we conducted an additional ablation study on the learning rate scaler scheme of low-rank factors in the approximated LORO steps. Following the same configuration as Section 5. we train the rank-122 LLaMA 60M model with LORO. using learning rate scaler choosing from {1, 0.5, 0.25, 0.1, 0.01), where 0.25 = r/n = 128/512 is exactly the Reimannian down-scaling rate we used (1T). We report the final perplexity of the resulting models. As shown in Figure 10 our LORO exhibits robustness across range of learning rate scalers (i.e. from 0.1 to 0.5).
Training curves and Validation Curves of LORO (rank 128, K 500) on LLaMA-60M
10a
scaler 0.01
10'
0.1
Training Perplexity
104
scaler 0.25 (LGRO)
scaller 0.5
Validation Perplexity
scaler
scaler 0.1
10a
Ir scaler 0.25 (LORO)
10a
scalle =
10a
10a
1500
2000
2500
3000
7000
10000
Pretraining Steps in [0, (..0000
Pretraining Steps
Figure 10: Ablation study on the impact of learning rate scaler on approximate LORO updates for training rank-128 LLaMA-60M models. The y-axis (log scale) represents perplexity.
24

---

Published as a conference paper at ICLR 2025
Comparing LORO against preconditioned LoRA. We conduct further experiments to compare our approximate LORO update against the preconditioned LoRA (denoted by "PrecondLoRAA) proposed in Zhang & Pilanci 2024). We design two experiments to study the following research questions: 1. Is Precond originally proposed for low-rank fine-tuning, sufficient to achieve strong performance low-rank pretraining? 2. Is the Precondl.ok update step more ideal than the in low-rank pretraining?
approximate LORO update (I1
Specifically, we train rank-128 LLaMA-60M models under the following configurations:
"PrecondLoRA" Following (hhang & Pilanci dated by: B < B LL.B. (AA ) and A
2024). both the low-rank factors are up- 7/(B B)-1. VA T.
"PrecondLoRA+ exact-LORO", We replace the approximate LORO step in Equation (12) with the PrecondLoRA update above, while we keep the exact LORO step in every K 500 steps unchanged.
To ensure a fair comparison with LORO, we use the same learning rate = 0.01, the same cosine learning decay schedule, and fix the frequency of optimizer state refreshing 500.
In Figure I we visualize the training curve between PrecondLoRA, "PrecondLoRA+ exact- LORO", and our original LORO. Our observations summarized as follows:
LORO > PrecondLoRA: Interestingly, while PrecondLoRA achieves a faster descent in both training and validation loss during the early stages of model training, it eventually exhibits a surge in validation loss. We conjecture that PrecondLoRA bias towards ill- conditioned low-rank factors, making the matrix inversion in the PrecondLoRA update numerically unstable. In low-rank fine-tuning, where the rank r is as small as r = 8, 16, 32, the numerical instability in the X matrix inversion is severe. However, in low-rank pretraining scenarios, where the rank r is as large as r = 128 or 256, the instability of matrix inversion becomes evident and hinders the final performance.
LORO >> "PrecondLoRA + exact-LORO": It is observed that the validation loss of "Pre- condL.RRA exact-L.ORO" exceeds that of LORO during the early training stages but eventually falls behind. We conjecture that performing matrix inversion at every step is overly frequent, which compromises the performance of PrecondLoRA. approximate LORO step demonstrates better performance, without the need to compute the expensive matrix inversion at every update step.
contrast, our
"ott exact-LORO" > PrecondLoRA: It observed that "PrecondLoRA+ exact-LORO" does not exhibit a surge in validation loss at the end of the training. This implies that our exact LORO update is effective in alleviating the instability Precond- LoRA training.
In conclusion, for low-rank pretraining, our observation indicates that our approximate/exact LORO update is preferable to PrecondLoRA update which initially designed for low-rank fine-tuning.
Training curves and Validation Curves of LORO (rank 128, K 500) on LLaMA-60M
10a
Full-size
PrecondiLoRA
Training Perpiextty
10a
PrecondLoRRA + exact-LORO
LORO
Velication Perplcxit
PracondLoRA LORO (ours)
10a
50
10a
80
10a
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
Pretraining Steps
Pretraining Steps
Figure 1I: Comparing LORO against preconditioned LoRA Zhang & Pilanci 2024 on training rank- 128 LLaMA-60M models. The y-axis (log scale) represents perplexity.
25

---

Published as a conference paper at ICLR 2025
Comparing exact LORO update against standard SVD projection. To study whether replac- ing the exact LORO update with a simple low-rank SVD projection is sufficient for good training performance, conduct further experiments to evaluate the effectiivneess of the simple low-rank projection. We train rank- 128 I. LLaMA-60M models, where we replace the exact IORO update with simple SVD projection every K steps (chosen from (500, 100, 10)). Specifically, the projection update comprises of (U, S, V) < SVD, and f UvS, < VSVT. To ensure a fair comparison with LORO, we use the same learning rate = 0.01, the same cosine learning decay schedule, and the frequency of optimizer state refreshing as 500.
As shown in Figure 12 for all K E {10, 100, 500), it is observed that the exact SVD step exhibits less satisfactory performance. As the exact update gets more frequent (K decreases from 500 10), the training curve and validation loss get worse. On the contrary, Figure 7 shows that our LORO enjoys improved training stability and performance under increasing K. This suggests that: 1. using simple SVD projection update (retraction without Riemannian update) is insufficient for good training, and 2. using Riemannian update before SVD projection leads to improved performance.
We conjecture that this is because the low-rank SVD projection plays crucial role in exploring an effective new optimization subspace for the low-rank factors B and A. Therefore, before performing the low-rank SVD projection balance the singular values,
need carefully determine the successive updated version of the low-rank weight (BA). From this perspective, LORO benefits from using the Riemannian gradient, which updates the product (BA) in the direction of the steepest descent on the low-rank manifold. In future work, we will focus on developing more advanced techniques for exploring better optimization subspaces for low-rank pretraining.
Training curves and Validation Curves of LORD (rank 128, K 500) LLaMA-60M
10a
Full-size
44
SVD, K- 10 SVD, SVD, Km 500
Training Perplexity
10r
Jalidation Hesplettt
LORO (ours).
40
10a
38
Ful-ssie
SVD, K
10'
J4
SVD, K= 100
32
LORD (ours), K 500
10a
4000
B000
10000
2000
B000
10000
Protraining Steps
Protraining Stops
Figure 12: Comparing exact LORO update against simple SVD projection on training rank-128 LLaMA-60M models. The y-axis (log scale) represents perplexity.
LORO finds competitive low-rank language models with comparable performance to the full- size baselines. Based on our experimental results and observations, we empirically validate the existence of strong low-rank language models capable of achieving performance comparable their full-size counterparts. This area remains under-explored in the existing literature. In this paper, we have demonstrated that LORO consistently achieves competitive performance that is comparable to the full-size basclines across various scenarios, indicating that its effectiveness is not merely a coincidence or the result of specific experimental setups or hyperparameter tuning.
Generally, full-size training provides satisfactory stability and strong performance, even with stan- dard optimizers. However, these advantages come at the cost of high memory and computation de- mand. From an environmentally friendly and economic perspective, our L.ORO approach achieves a better balance between efficiency and performance.
In pretraining, the fact that "low-rank language model can be pretrained to achieve high perfor- mance comparable to the full-size baseline" neither entirely unexpected nor surprising. As shown in Figures 2 and Figure 3 of (Jaiswal et al. 2024), full-size pretrained LLaMA models demonstrate low-rank structures in both their weight matrices and gradients. As shown in (Chen et al.) (02T)) and (Wang et al. 2024). low-rank compressed language models also exhibit competitive perfor- mance against the full-size baseline. These observations suggest the feasibility of training low-rank parameterized language models to achieve performance comparable to their full-size counterparts.
26

---

Published as a conference paper at ICLR 2025
In fine-tuning, the fact that "low-rank fine-tuning achieves comparable or better performance than fall-parameter fine-tuning" has also been reported in (Hu et al., 2022). (hang et al. 2023). and (Zhao et al. 2024). In fine-tuning scenarios where the downstream dataset (e.g., GLUE) relatively small, full full-parameter fine-tuning can lead to overfitting, resulting in suboptimal performance. More- over, in such cases, full-size pretrained checkpoint is typically available, allowing new knowledge to be incorporated through a low-rank adaptation step. This makes it feasible to adopt a full-size pretrained model downstream tasks within a low-rank subspace.
In summary, "low-rank pretraining" and "low-rank fine-tuning" are two distinct settings, with our work focusing on the former. Our empirical results and analysis suggest that the potential of low- rank language models was under-explored in existing literature. In future work, we aim to dedicate more effort to developing advanced low-rank pretraining methods and theoretical tools for analyzing the dynamics of low-rank pretraining.
27

---

Published as a conference paper at ICLR 2025
Condition Frobenius LLaMA--00.layers...o wwtp:/aayer...mpupppppp... mmmppupppproj
1912
THE MARORR MAR
low LORDLears)
LORD]aun)
20F
8300
6002
Condition Number and Frobenlus Norm of rank-128 LLaMA-60M. www.-600..ayyys..ssssos www.ei..ttet.a.....ss wwt...p_ttooo
7x.12
20F
moberta Rem
Method LORO
Nuthod LORDSeees
Herted
200
conDtawrs)
2003
4000
6300
R302
10020
2030
6200
6006
2030
10932
2000
4200
6002
6000
10003
Taaking Stops
Traiking Strpe
ving Stope
Condition Frobenlus rank-128 LLaMA-60M.layers 7.mpp.uppro up_pro
167
Podbutiaa Norm
Tul <it-
Fullebe Lowviink LOND(ours)
Lokk))our)
Herwo
Fillioo
Tranning
Candition Number and Frobenius Norm rank-1 128 LLaMA 60M. www..aayes....see. wwtt:.77.ffffff...... attn..q_roojj
74.12
Prdestto Morm
Mpphod LORD
Mathed Falalis
Merlod
LOROtsurd
cOROcaund
2000
4000
$300
0303
13020
2000
4000
5026
6020
10002
2000
4300
6203
6000
10003
Takning Staps
Thaking Stope
Thaiiin
Figure 13: Visualization of the condition number and Frobenius norm for linear and self-attention layers in rank- 128 LLaMA-60M model.
28

---

Published as a conference paper at ICLR 2025
Visualizing singular value distribution rank-128 LLaMA-60M. www.:laaers..m.......
Siggia vokees wirgeaswokes spM for 1ark
Singga vikues ejoo Lowanoa
Robss KpMes20 for Use
-3
15
376
359
100
122
2206
7300
800
228
3030
528
10
24
31
16
310
111
120
330
3100
1206
1290
$000
12k
X0V
1269
5289
Visualizing singular value distribution of rank-128 LLaMA--6M.layers.7.4.
720
35.
+15
10
200
216
300
130
24
320
230
12
3000
5297
7060
X090
Stngstor AUAKS OERO Tor LONDLovrsl
501g2/0 Kdues 04 THE LORDIOUS)
Nuoss SURUAN for LDDDtemm)
30
33
51
300
15
115
Teaf 9rp
Teihro Megs
Figure 14: Visualization of the singular value distribution for the linear layers in rank-128 LLaMA- 60M model. The y-axis represents the order singular values, the x-axis represents the training steps. Deeper color indicates a larger singular value is observed.
29

---

Published as a conference paper at ICLR 2025
Visualizing singular value distribution rank-128 LLaMA-60M.ayerrs.... attaatttttpr.e
Seogarrwkes spM for Loe
Singda vikues epo For LaLaaaa
Robis 42M0420 for Use mook
11
100
122
2206
7200
8002
228
3030
5288
5,09
XX
24
-1
16
0000
-1
310
114
330
3100
1206
1290
$000
12k
X0V
5263
5299
Visualizing singular value distribution of rank-128 M0AMA--0M.layyrs.7..oo attp___tttt__ooo
-20
25
22
10
19
200
216
130
200
300
12
3000
5297
7000
Stngstor PJAKS OERO Tor LORDLovrs!
50122/0 Kdues 04 TH LORD
Tuons SURUAN for LDDDtemm)
30
2.15
300
-25
125
Teaf 9r
Thihro Megs
Figure 15: Visualization of the singular value distribution for the self-attention layers in rank-128 LLaMA-60M model. The y-axis represents the order of singular values, the x-axis represents the training steps. Deeper color indicates a larger singular value observed.
30


Reviews:
[{"review_id": "a5mSFbkXnm", "reviewer": "Reviewer_fYLm", "summary": "The paper focuses on memory-efficient pretraining for decoder-only transformer models, particularly in the pretraining setting. It introduces the Low-rank Riemannian Optimizer (LORO), which maintains a low-rank parameterization throughout training and inference. A central challenge in standard low-rank pretraining is that separately optimizing low-rank factors introduces redundant gradient information, which hinders learning; LORO addresses this by using Riemannian optimization to jointly update factors along the low-rank manifold’s steepest descent. Through periodic exact updates and low-cost approximations, LORO reduces memory usage, increases training speed, and achieves lower perplexity than baseline methods for low-rank transformer training.", "questions": "Refer to the weaknesses section and please share your views on this.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "I think the main formulation of the paper, which is to use Riemannian optimization to jointly update factors along the low-rank manifold’s steepest descent, is a novel and interesting approach to the problem of low-rank pretraining. The paper is generally well-written and the math is clear and easy to follow. I particularly liked the method section, where the authors devlop the solution by systematically addressing the memory and computational overheads. \n\nThe experimental results in the pre-training setting are strong and show that LORO can achieve competitive performance with reduced memory and computation.", "weaknesses": "The main problem I see with the paper is that the authors do not provide a rigorous analysis supporting their central claim that independent updates of the low-rank factors are inefficient. They state that separately optimizing the factors B and A introduces redundant information from the full gradient, which hinders learning, but this claim is not systematically studied or empirically validated. I think this is a major weakness of the paper and should be addressed in a revision. Specifically, it would be good to see some evidence that this phenomenon is common and that it is a significant problem in practice. Additionally, it would be good to also see how LORO is able to overcome this issue and why it is more efficient than other methods in this regard.\n\nMinor: Some other related works that can be cited:\n\n1. Adam-mini: Use Fewer Learning Rates To Gain More, Zhang et al\n2. BlockLLM: Memory-Efficient Adaptation of LLMs by Selecting and Optimizing the Right Coordinate Blocks: Ramesh et al", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SrJSfSx4Xs", "reviewer": "Reviewer_4NrY", "summary": "This paper proposes a novel pretraining method that combines Riemannian optimization with Low-Rank Adaptation (LoRA). The authors aim to enhance the update process of the low-rank factors by integrating Riemannian optimization, which purportedly allows for more effective use of information during parameter updates and reduces redundancy caused by subspace overlap. The paper also introduces a well-designed mechanism to avoid full-size weight and gradient computation, reducing the computational cost during pretraining and resulting in inference-time acceleration for the proposed LORO method.", "questions": "See above.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Overall, the paper presents a clear theoretical framework and demonstrates impressive experimental results. However, I remain skeptical about the results.", "weaknesses": "1. As part of the motivation, the authors assert that \"separately optimizing the low-rank factors\noverlooks the intricate structure of the low-rank parameterization and fails to guide their product\ntoward the steepest loss descent on the low-rank matrix manifold.\" This statement appears to be\nlargely intuitive. I believe that this claim requires a clearer analytical breakdown and more\nsubstantial theoretical justifications.\n\n2. The paper mentions that to stabilize training, they apply a 5-step linear learning rate warmup and\nrefresh the Adam statistics at each exact LORO update step. In prior work, changing the\nparameters of the Adam optimizer required careful handling, as restarting the optimizer state\noften led to loss spikes. However, the proposed method seems to avoid this issue. It would be\nbeneficial if the authors could provide a more detailed explanation or proof regarding why LORO\nexhibits robustness to this situation, particularly in comparison to previous methods where loss\nspikes were problematic.\n\n3. The authors indicate that not every step in training involves a Riemannian update. From the\nablation study provided, it appears that using K= 500 and K= 250 and yields similar performance.\nThis raises questions about the necessity of the Riemannian update step. Specifically, I am\nconcerned that the performance improvement may not stem from Riemannian optimization itself,\nbut rather from the fact that the Riemannian update step refreshes the optimization subspace. It\nis possible that standard low-rank steps are already sufficient for good training performance. A\nmore detailed ablation study isolating the effects of the Riemannian updates would help clarify\ntheir contribution to the overall performance of the method.\n\n4. In the non-Riemannian update steps, the authors adopt the update rule in (10). The authors claim that this is a reasonable approximation of the Riemannian update step. However, intuitively, this appears to simply scale down the learning rate for low-rank pretraining,\nwhich would be equivalent to standard low-rank pretraining. Previous work has shown that vanilla\nlow-rank pretraining often underperforms, so it is unclear why LORO’s approximate update step\nachieves strong empirical results. I would appreciate further clarification or experiments that\nexplain why this approximation works well in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HaBMQdbOBd", "reviewer": "Reviewer_37oB", "summary": "This paper propose a Riemannian update scheme for low-rank pretraining of large laguage models(LLM). The proposed scheme utilize a Riemannian gradient step (for every $K$ iterations) on the low-rank manifold (represented by SVD) for efficient Riemannian optimization. Numerical experiments are provided to show the efficiency of the proposed approach comparing to the existing low-rank pretraining approaches.", "questions": "I have following comments and concerns regarding the paper:\n\n1. The effecitveness of the proposed approach is still under-explored: The authors claim that the goal of the proposed approach is to (Line 82) \"train a low-rank parameterized language model from scratch to achieve performance comparable to the full-size baseline\". According to the scaling law (which might not be 100\\% applicable here), the primary indicator of the final model performance should be the model size, and this approach is using **exactly** same number of parameters as the low-rank approach. This makes me really wonder if the low-rank approach is highly underestimated in Figure 1 and Table 1. Maybe it's due to initialization/random seed in the Galore paper. The authors might want to comment on this, especially consider that their update is identical to low-rank except for the Riemannian step in every $K$ steps\n2. Continuing above point, since the algorithm only update Riemannian step for every $K$ steps, I would expect to see a clear improvement at the Riemannian step (since other steps are same as low-rank, and low-rank is performing badly), say in Figure 4 (and the curve of low-rank should also be included in Figure 4). However seems that the lose is just decreasing steadily for $K=250$ or $500$. Again, the authors might want to comment on this.\n3. One additional experiment I would recommend doing is that, updating the Riemannian step for every iteration or very frequently, and see if it clearly outperforms low-rank/sltrain. And then say that due to the commputational cost, oen could jsut update the Riemannian step for every $K$ iterations.\n4. The inference time of the proposed method should also be the same as the low-rank approach?\n\nI think overall this is an interesting work. But it's still hard for me to believe that the every-$K$-steps Riemannian update could be such a miracle to turn the badly-performing low-rank approach to SOTA, considering no additional model parameters introduced. I'm open to discussion and would like to hear the authors' comment on this.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The paper is well-organized. The proposed algorithm is theoretically motivatived by Riemannian optimization and does not require much more computational overhead by updating the Riemannian step every $K$ steps. The proposed approach yields good performance in the experiments in the paper.", "weaknesses": "(Please respond to the \"Questions\" section directly)\n\nThe effecitveness of the proposed approach is still under-explored both theoretically and numerically; Comparing the low-rank approach, I still have some doubts on the reliability of this work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pvzG4n1NzH", "reviewer": "Reviewer_9v9P", "summary": "The paper proposes to parameterize the weights as a low-rank matrix for parameter and memory efficient pretraining of LLMs. Instead of separately optimizing two factors using (stochastic) gradient descent, the paper leverages Riemannian geometry of low-rank matrices, and propose efficient gradient computation and update strategies. The pretraining performance is evaluated on LlaMA 60M to 1B models.", "questions": "1. How did the experiment initialize the B and A?\n\n2. Have you tried different scaling factors for B and A, and whether using different scaling factors can result in different performance?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper proposes a novel Riemannian optimization strategy for optimizing low-rank matrices.\n\n2. The proposed gradient and update strategies only through the factorized form are interesting.\n\n3. The performance is strong compared to the baselines and is comparable to the full-rank pretraining.", "weaknesses": "My biggest concern on this work is regarding the empirical simplification using approximated Riemannian update. The paper proposes to use a scaling factor for both the update of B and A. This results in an update which exactly matches the standard update (the same scaling factor can be merged with learning rate). In summary, it appears that for most iterations, Algorithm 1 simply perform standard gradient descent update for the low rank factors and periodically update with Riemannian gradient. This raises the question of whether the periodic Riemannian gradient update results in such a performance improvement.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)