LoRA Done RITE: Robust Invariant Transformation Equilibration for LoRA Optimization
Authors: Jui-Nan Yen
Si Si
Zhao Meng
Felix Yu
Sai Surya Duvvuri
Inderjit S Dhillon
Cho-Jui Hsieh
Sanjiv Kumar
Venue: ICLR 2025 Oral (2025)
Keywords: optimization
LoRA
Primary area: optimization
TL;DR: Improve the optimization of LoRA using adaptive matrix preconditioning method with transformation invariance

Abstract: Low-rank adaption (LoRA) is a widely used parameter-efficient finetuning method for LLM that reduces memory requirements. However, current LoRA optimizers lack transformation invariance, meaning the updates depending on how the two LoRA factors are scaled or rotated. This deficiency leads to inefficient learning and sub-optimal solutions in practice. This paper introduces LoRA-RITE, a novel adaptive matrix preconditioning method for LoRA optimization, which can achieve transformation invariance and remain computationally efficient. We provide theoretical analysis to demonstrate the benefit of our method and conduct experiments on various LLM tasks with different models including Gemma 2B, 7B, and mT5-XXL. The results demonstrate consistent improvements against existing optimizers. For example, replacing Adam with LoRA-RITE during LoRA fine-tuning of Gemma-2B yielded 4.6% accuracy gain on Super-Natural Instructions and 3.5% accuracy gain across other four LLM benchmarks (HellaSwag, ArcChallenge, GSM8K, OpenBookQA).

Published as a conference paper at ICLR 2025
LORA DONE RITE: ROBUST INVARIANT TRANSFOR- MATION EQUILIBRATION FOR LORA OPTIMIZATION
Jui-Nan Yen +1
Si Si
Zhao Meng
Felix Yu
Sai Surya Duvvuri 3
Inderjit S. Dhillon
Cho-Jui Hsich 12
Sanjiv Kumar 2
ABSTRACT
Low-rank adaption (LoRA) is widely used parameter-efficient finetuning method for LLMs that reduces memory requirements. However, current LoRA optimiz- ers lack transformation invariance, which leads to weight updates that depend on how the two LoRA factors are scaled or rotated. This deficiency leads to ineffi cient learning and sub-optimal solutions in practice. This paper introduces LoRA- RITE, a novel adaptive matrix preconditioning method for LoRA optimization, which achieves transformation invariance while being computationally efficient. We provide theoretical analysis to demonstrate the benefit of our method and con- duct experiments on various LLM tasks with different models including Gemma 2B, 7B, and mT5-XXL. The results demonstrate consistent improvements over ex- isting optimizers. For example, replacing Adam with LoRA-RITE during LoRA fine-tuning of Gemma-2B yields 4.6% accuracy gain on Super-Natural Instruc- tions and 3.5% accuracy gain across four other LLM benchmarks (HellaSwag, ArcChallenge, GSM8K, OpenBookQA).
INTRODUCTION
Low-Rank Adaptation (LoRA) (Hu et al. 2022) is a popular parameter-efficien method for fine- tuning Large Language Models (LLMs). By freezing the pretrained weights and injecting trainable low-rank matrices into each layer, LoRA significantly reduces memory requirements and mitigates overfitting in limited data settings. More formally, letting W E R be weight matrix in an LLM, LoRA freezes W and introduces a low-rank matrix Z added to W, where represented by the multiplication of two rank-r matrices A and B, i.e.
Z = AB' E R A Rm B RaXT, r & min(m,n).
(1) LoRA factors in this paper. Recent research has
The matrices A and B will be referred explored numerous variations and improvements over the classic LoRA algorithm (Valipour et al.,
2023
Zhang et al. 2023b Liu al. 2024 Yaras al. 2024).
Despite being widely used in practice, we find that applying standard optimizers to LoRA leads to updates that are not "transformation invariant". By definition of LoRA in (I). the same update Z can be decomposed in multiple ways, i.e., Z = A, = B, Ideally, an optimizer should yield the same update to Z regardless of the specific factorization. However, commonly used optimizers with diagonal preconditioners like Adam (Kingma & Bal (014)). Adagrad (Duchi et al. (011). RMSProp (Tieleman & Hinton] 2012), and even second-order methods like Shampoo (Gupta et al., 2018) and CASPR (Duvvuri et al. 2024), violate this principle when applied to LoRA. This violation not only presents a mathematical inconsistency but also leads to significant inefficiencies during training. In practice, receiving substantial updates while the other remains nearly fixed. Although this can be partially mitigated by some recently proposed approaches such employing different learning rates for the two factors (Hayou et al. 2024), we ask the question: is there a more principled way to design an optimizer that inherently enforces transformation invariance for LoRA?
observe that one LoRA factor often dominates the optimization process,
To address this challenge, we first prove that any form of diagonal preconditioner cannot achieve transformation invariance, which motivates the use of matrix preconditioners. However, existing
Work done while Google. 'UCLA "Google 3UT Austin

---

Published as a conference paper at ICLR 2025
matrix preconditioners like Shampoo and CASPR lack transformation invariance and introduce sig- nificant computational and memory overhead. To overcome these limitations, we propose LoRA- RITE (Robust Invariant Transformation Equilibration), novel optimizer designed specifically for LoRA optimization. LoRA-RITE employs a transformation-invaraan preconditioner on the low- rank side, achieving transformation invariance without incurring substantial overhead. Furthermore, we demonstrate how to maintain this property when incorporating first and second moments, cru- cial for the practical effectiveness of adaptive optimization methods. Empirical evaluations across various datasets and models confirm the effectiveness of the proposed algorithm.
The contribution of this paper can be summarized below:
We propose LORA-RITE, the first adaptive matrix preconditioning optimizer for LoRA that is caansformationninvaaiant the property that is lacking in most existing optimizers when applied to LoRA. Theoretically, we provide a convergence analysis for our method. Despite utilizing matrix preconditioners, LORA-RITE achieves little overhead in both memory and time compared to first-order optimizers, especially when the LoRA rank (r) is significantly smaller than the original matrix dimensions m and n.
The proposed optimizer leads to significantly improved performance across multiple datasets and architectures. For instance, when applied to the GSM8K (Cobbe et al. 2021) dataset with Gemma 7B IT model (Gemma Team al. 2024), LORA-RITE achieves
55.50% accuracy rate. This surpasses the widely-used Adam optimizer (Kingma & Ba 2014) by substantial margin (48.37%) and even outperforms the second-best optimizer on this dataset, Lamb (You et al. 2020) (50.64%), by approximately 5%.
TRANSFORMATION INVARIANCE FOR LORA OPTIMIZATION
We now introduce the concept of transformation invariance in LoRA training and demonstrate that most existing optimizers, when applied to LoRA, do not satisfy this property. This deficiency leads inefficient learning in practice.
2.1 DEFINITION OF TRANSFORMATION INVARIANCE
As introduced in 1. LoRA adds low-rank matrix Z AB to the original weight matrix W and learns the LoRA factors A E TR B E R to minimize the fine-tuning loss. Observe that many different LoRA factors (A1, B B),, (Ag, B2) can represent the same finetuned weight,
= A A,
(2)
When an optimizer is applied to train LoRA, will produce different updates, 8A,,8B, or A B2. based on the specific parameterization used. Even though (A1, B1) and (A2, B2) rep- resent the same finetuned weight Z, the updates using different parameterizations can produce dif- ferent updates to Z. This suggests serious inconsistency and implies that the update could be suboptimal under some parameterizations. Based on this observation, we propose that LoRA opti- mization should ensure transformation invariance, defined as follows:
Definition 1 (Transformation Invariance). Let (A1, B,) be a pair of LoRA factors and let A A, R, B2 B, RR for some invertüble matrix R. An optimizer exhibits transformation invariance if its updates, (oA,, BB) and (oA2, 82), satisfy
(Ai + ( + &B,) = (A2 + 5A2) +8BB) : + 8Z.
This means the optimizer should produce the same update, 8Z, to the fine-tuned weights for any equivalent LoRA factorizations. As special case, introduce scalar scale invariance below in Definition
2 as a weaker version of transformation invariance, which only requires that updates remain equivalent when the LoRA factors scaled up or down by scalar factor. Formally, define it as:
Definition 2 (Scalar Scale Invariance). Let (A,, Bi) be a pair of LoRA factors and let A sA1, Bg = (1/8) B, for some nonzero scalar constant 8. An optimizer exhibits scalar scale invari- ance its updates, (5A1, (BB, and (, ò 8B2), satisfy
(Ai + (B,) + 8B,) = (A, + 8A, 1.0BB++0000 +
Surprisingly, we will show that most commonly used LoRA optimizers do not even satisfy this weaker form of transformation invariance.
2

---

Published as a conference paper at ICLR 2025
2.2 EXISTING OPTIMIZERS ARE NOT SCALAR SCALE INVARIANT
We now show that neither gradient descent nor Adam are scalar scale invariant, and in fact, almost all the existing optimizers
scalar scale invariant when applied to LoRA optimization.
Assume loss is f (2), ) and the gradient of the loss function f with respect to Z, A1, B1, A2. B2 to be V == ff VA,
if/ V == 0f/OBi. V == / and VB2 == Of/0B,
respectively. Recall Z = A, = A, and by chain rule, we have
VA == 0 0f * 0Z / Of f B1 VZBi,
(4)
Similarly, we have VB, VZ A, = VZB2, and VB, VZ A2. For gradient descent,
8A1 = A1, 8B, == -7/BB1. 8A, = - 5B2 == B7,
where y is the learning rate. To test scalar scale invariant, let A₂ sA1 B3 = (1/s)Br, we have
AA B2 - (1//s) = (1/8)6A1 (since B₂ (1 (1/))B1)
(5)
iB V A s7) VZ A soBI (since A₃ sA1).
(6)
Consequently, from (5) and (6:
A²B2 + 0A A28B2 + B
A A (1/3))A A + O
# B, + BJ A + +
and so 3 does not hold for arbitrary Therefore, gradient descent is not scalar scale invariant, and the gradient can be arbitrary large for the LoRA factors when s goes or infinity.
Can this issue be mitigated by adaptive updates such as Adam? The answer is no. To see this, let
8A1 := V (V O VA1)
be the Adam update for A1, where @ denotes elementwise multiplication. Note that we omit the momentum term for brevity. Defining updates of A2, B₁ and BN similarly, we have
7 A2
(I/s)nVAi
A = (V 2 O V A22
(V/8)(VA, @ 112 = 6A1
and similarly 8B, = 8B As result,
A.B. B + A + 8A
B. T + (1/s) + S + SA,OB
* A. + B A.OB + 8A,
thus failing satisfy scalar scale invariance. Actually, can that most of the existing opti- mizers, such as Adagrad (Duchi et al., 2011]. RMSProp (Tieleman & Hinton 2012), and Shampoo (Gupta al. 2018) are not scalar scale or transformation invariant.
2.3 BENEFITS OF TRANSFORMATION INVARIANCE
Why is transformation invariance important? Beyond the mathematical argument that different pa- rameterizations of the same weight update should be equivalent, we demonstrate that transformation invariance leads to more efficient feature learning. The concept of efficient feature learning, intro- duced in (Hayou al. 20244, describes the asymptotic training behavior of LoRA the network width grows. As discussed earlier, for LoRA, the update to the matrix = AB can be decom- posed into three parts
zz= = (A (B)(B + +BB - AB = SAB T + ASB + T
where the 5A5B is typically negligible as it depends on the square of the learning rate. Efficient feature learning requires that both é dAB x and A x are of magnitude O(r²) = 0(1) with respect to the network width where is the input embedding. On the contrary, if the magnitude of SAB x is dependent on n: Onn"), then it can explode if a > and diminish if a < 0, as the network width grows.
Hayou et al. (2024) show that conventional optimizers do not satisfy efficient feature learning. This can be seen from Figure T where the weight norm for factor B changes significantly while the weight norm for factor A barely changes. More discussions about the this unbalanced training dynamic are in Appendix (A.8..
Under mild assumptions, we can show that a transformation-iivvaiant optimizer guarantees efficient feature learning. The proof is given in Appendix A.3
3

---

Published as a conference paper at ICLR 2025
Theorem 1. Any optimizer that is transformation ianammaaaatio-iivvannnnn and uses the same update rule for both A and B will achieve efficient feature learning.
Beyond the efficient learning guarantee, in practice when training LoRA with existing optimizers, is often the case that only one of the LoRA factors is updated properly, while the other remains almost unchanged,
shown in Figure T This is also consequence of lacking scalar scale invari- ance, as the initial scales for the two LoRA factors can be very different (one typically initialized from O while other from Gaussian random).
10
10
Adam
ScaledAdam Shampoo
100
10-4-
Adam ScaledAdam
Shampoo
2000
4000
6000
8000
10-5
2000
4000
6000
8000
Train Step
Train Step
(a) Frobenius norm of LoRA factor
(b) Frobenius norm of LoRA factor
Figure 1: The norms of the LoRA factors A and B change differently across different training steps.
OUR PROPOSED OPTIMIZER
We now present our proposed algorithm, which satisfies transformation invariance, and achieves significant improvements over previous LoRA optimizers many empirical tasks.
3.1 DIAGONAL PRECONDITIONING IS NOT ENOUGII FOR TRANSFORMATION INVARIANCE
In this section show that diagonal preconditioning is not enough to achieve transformation in- variance. Recall that the LoRA factors are A E Rmxr and B E Rexr, When updating A, most existing optimizers utilize the following update:
vec(6A) X VV),
where vec(-) lists the elements of matrix as vector in column- majoor order and X R m is symmetric positive definite preconditioning matrix. Diagonal preconditioning methods like Adam and Adagrad assume X a diagonal matrix, while matrix preconditioning methods such as Shampoo and CASPR use non-diagonal X.
Consider the LoRA pairs (A1, Bi) and (A2, B2) such that: A = A, BN = B.R where R is invertible matrix. By (4. the gradient with respect to the LoRA factors are:
A V VZB. VARR
VB, VZ A. V.
Let us consider the when vec( A) = V in update 7 and apply transpose on both sides to get:
1. For A R
substitute vec(ōA) 8AT and
8A = - -BX, òB VBX - TX
(9)
In this case, have two equivalent LoRA pairs (A1, B1), (Ag, B2) with their corresponding preconditioners X, X,, then from and we have
A A VZ , X A B, + A SA V B and
(10)
SA,B. + A +8A28B] = B nA VZ - X.AIV2
= (R -T,T R - A VZ / /V2B, X. A VZ.
(11)

---

Published as a conference paper at ICLR 2025
For 100 and 111 to be equal for any arbitrary 17.
that it is necessary to have
VZB,X, VZ- R X A VZ,
(12)
as it is the only term quadratic to However, there might not exist a pair of diagonal preconditioners X1; X2 o 0 such that the above condition is satisfied for arbitrary V Z A, and B, because R can be full matrix chosen in an udversarial fashion, such that R X X2 is non non-diagonal and thus different from the left hand side. Consequently, we conclude necessary adopt matrix preconditioning to achieve transformation invariance.
3.2 ACHIEVING TRANSFORMATION INVARIANCE
To achieve transformation invariance, we begin by recognizing that the LoRA factors, A and B, can be decomposed into their respective orthogonal bases and magnitudes:
A = UARA, B = UBRB.
where UA and UB can be obtained through polar decomposition.
Note that the gradients of A and B. VA VZB and VB = VZ A, depend on both the basis and the magnitude. To achieve transformation invariance, we introduce the concept of "unmagnified gradients" distinguished from standard gradients by the symbol V:
VA VZUB VAR V VZ VBR
(13)
where R! and R are the pseudo-invers, of RA and Rg respectively. These unmagnified gradi- ents, relying solely on the column spaces of A and B, remain invariant transformations of the LoRA factors. This invariance forms the cornerstone of our algorithm's ability to achieve transfor- mation invariance.
A
Adaptive preconditioning methods like Adam have demonstrated superiority over non-adaptive methods like SGD. Furthermore, established earlier, matrix preconditioning is crucial for achiev- ing transformation invariance. Therefore, we propose utilizing these unmagnified gradients for adap- tive matrix preconditioning. Additionally, we only do one-sided preconditioning, on the shorter side that is of size to ensure low time and memory complexity of the proposed method.
Since our update rule is symmetric for A and B, for brevity, from now on we only describe the update rule for A. For simplicity, let's first discuss the case without momentum. We propose the update rule for
LORA-RITE: òA = VVAAAVVVAV)AAAAAAAAAAAA
(14)
where 77 is the learning rate. This update can be broken down into two parts. The first part
AVA 112
resembles the adaptive preconditioning mechanism in Adagrad, but employs matrix operations in- stead of element- wise operations. Crucially, the use of unmagnified gradients ensures this term remains consistent across all equivalent LoRA pairs, up to the choice the basis.
The second part (R adjusts the magnitude of the update for different LoRA pairs. Since (R ) B = UT this effectively takes out the magnitude of B in TAB We thus have
A = - A, V A2) =
A 1 (BJ)² -1//ZUU, A VAa) 1/2U - VA, -/1 )t
= -U,VZ(UB, A,, VAAU -((" BA (B
(15)
Similarly, 6B, 8B, A (A,) Consequently, we have
A,B,T + B A &A = A, + +AABBT A28 + 8A28B
This demonstrates that proposed method satisfies transformation invariance. Note that this sim- plified update rule does not yet incorporate accumulated first and second moments, which will be addressed in the following paragraphs.
5

---

Published as a conference paper at ICLR 2025
Incorporating second moment. Adaptive optimizers typically employ accumulated second ments for preconditioning. A naive approach might involve replacing the VA T VA term in (14)
with its accumulated sum over training iterations E (1.2...... T::
I Σ
VAA
However, since each V., is computed with respect to the basis at specific step, directly summing them is mathematically incorrect. Instead, we must account for the varying basis each step. To achieve this, we accumulate the second moment as follows:
VA. PA. VA. PA. VA, VA.
(16)
where VA.-, is the accumulated second moment based on the previous basis at step t 1, and PA. (UB.) UB. transforms it the basis step f. During the adjustment,
Tr(PA, VA.., PA)) s Tr(VA..)).
indicating a potential loss information from the accumulated second moment. To quantify this loss, for symmetric positive definite matrices E1, E2 IR"""" we define
dx 1 E2) = max (A, A, E2) I min E UE,U II,
U
where À; EE) is the i-th eigenvalue of E, and U E Bexr is an orthogonal matrix that reflects our freedom choose the basis. We then define the "escaped mass" as
PA: PA,-1 d, (VA,,,,,, PA, PA,VA.,,PAl
(17)
To compensate for this, we add PA, I preconditioner, ensuring that VA, tPA, monotonically increases under suitable choice of basis, even though the choice basis does not influence the actual update.
Finally, our unmagnified preconditioned step, when incorporating second moment, can be written as SA. V A. VA. + PP.
(18) Note that similar to Adam, we can turn (16) into the Exponential Moving Average (EMA) form, where we multiple the first term by 1 32 and the second term by 32, with the hyper-parameter 82 E (0, 1) controls the decay rate.
Incorporating first moment. Similar the second moment, the first moment must also be ad- justed for changes in the basis using projection matrix. The update rule for maintaining the first moment can then be written
MA. MAM.. PI. (1 - S.,
Our final proposed update rule, incorporating both first and second moment, is
8A, M (RB)*.
(19)
Algorithm LORA-RITE
1: Initialize: unmagnified first and second moment MA. 0. VA. =0
2: for 1.... do
Compute the gradient A.:;
4: Compute polar decomposition of the LoRA factor Bi: B, UB, RB.
Compute the unmagnified gradient VA, VA, and PA. (UB.) T
Update unmagnified second moment VA. PA. VA. PA, + (VA.) TA.
7: Update the escaped mass PA. - d, (VA, A. PA. VA,, PA.);
Compute the unmagnified precondition step SA. V. + +PA,)) 1/2,
9: Update the unmagnified first moment MA. 3, MA
Pl, + 1-1))
10: Update model parameters Ac+1 = A, MA. (R
11:
end for
Our proposed algorithm, LoRA-RITE (Robust Invariant Transformation Equilibration for LoRA training), is summarized as Algorithm I where we show the updates for A, and update for B can be derived in the same way. Note that have shown that the main update rule of (18) satisfies transformation invariance, and this property can be extended even after adding the first and second moment into the algorithm, as shown in the following theorem (proof in Appendix).
Theorem 2. In Algorithm | every unmagnified term is consistent across all equivalent LoRA pairs. Consequently, Algorithm transformation invariant.

---

Published as a conference paper at ICLR 2025
Time and Space Complexity. The time and space complexity of our algorithm is similar to first order methods like Adam when r & m, n. In each iteration of Algorithm |] the dominant com- putational costs arise from (1) coooppsccomooitiinnn for m-by-r and n-by-r matrices which takes O(nr² + m)r) time, (2) matrix inverses and roots for r-by-r matrices which takes O(r)) time, and (3) matmuls with time complexity O(nr² mr?). Thus, the overall complexity per step is O(mr² + 22² + ))) It is only times slower than Adam, and since r is very small, this overhead negligible when comparing with the back-propagating time, The memory cost of our method is O(mr + nr) which is the same as Adam. We summarize the time and space complexity of our method versus some commonly used optimizers in Table the Appendix.
3.3 THEORETICAL ANALYSIS
Following previous work (uupta et al. 2018: Feinberg et al. 2023). we provide a convergence analysis of the proposed algorithm within the online optimization framework (Hazan et al.) 2016
Shalev-Shwartz et al. 2012). In online convex optimization setting, a parameter O, E K. is chosen iteratively, where K. is a convex decision set. After each decision- Or, a convex loss function /a is revealed, potentially chosen adversarially. The regret accumulated by the algorithm up to step T is defined as
Regrety- = =E
(=1 f((O,) min DEK -== f((0).
Σ
In the online convex optimization analysis, we bound the first-order term VO, (0; 0r) where 0* represents an arbitrary minimizer, and then use convexity of ft to connect it to the loss function. However, due the inherent structure of LoRA, loss functions ft, t {1...., T} are not convex with respect to Therefore, we directly bound the first-order term instead.
We assume for the fine-tuned weight Z of each layer, the convex decision set imposes the following constrains: || < DA, || p < DB, where || denotes the Frobenius norm. Additionally, we G. Following previous work, analyze convergence in the simplified scenario where the first moment is omitted and the second moment is summation, similar Adagrad. For LORA-RITE, our theoretical analysis yields the following result:
assume the gradient satisfies | Zl
Theorem 3. LORA-RITE satisfies:
1 T V (0, T 1
VO,
0.t1) = 1C
where ni is a fixed constant learning rate.
This theorem shows that the method either converges to particular stable solution or just move around directions that does not change the function value, suggesting a form of convergence. To further strengthen the guarantee, we introduce an additional assumption:
Assumption 1. Let A.
=(VA, + PA, 1/2 be the unmagnified preconditioner PA. =
(UB.) UB.., and QA (R 1 B._,' then have
-1 X
QA X-1 A 0% < X -PA, X A.-
AA,, PT.1
A
and
Tr(PA,XA) PX. T + c -A, 1/2)
This assumption essentially constrains the change in UB. and RB, to be relatively smooth. Under this assumption, we can establish the following stronger convergence result:
Theorem 4. Under Assumption D our proposed method satisfies:
1
v0, - -
T Our analysis closely resembles that one-sided matrix Adagrad. The key idea is to have a change of variable for both A and B such that all the quantities get replace by its unmagnified counterparts.
Compared to one-sided matrix Adagrad, which has a regret bound of (GG(DA + D3g I-1/2) which is higher than DBGDADDT-11/ the regret bound of LORA-RITE, The above inequality is using D3 + D3 2 2DADB. where the difference between both sides of inequality large when the two LoRA factors exhibit imbalance in magnitudes- DA and DR- This advantage is particularly relevant because previous work has shown that LoRA factors often exhibit such imbalances (Hayou et al. 2024). which can also be seen in Figure] I providing an explanation for the strong empirical performance of our method.
7

---

Published as a conference paper at ICLR 2025
4 RELATED WORK
Related Optimizers. Adaptive first-order optimizers like Adagrad (Duchi et al. 2011) utilize ac- cumulated second moments, essentially diagonal preconditioners, scale updates for each coordi- nate. This approach, adopted by optimizers like Adam (Kingma & Bal 2014) and RMSProp (Tiele) man & Hinton 2012), has become standard for training deep neural networks, including LoRA, and many other similar first-order methods have also been developed in the literature (Loshchilov & Hutter 2017 Chen et al. 2024). However, discussed in Section 3.1 these methods lack transfor- mation invariance when applied to LoRA.
Several higher-order preconditioners have shown promise in various training scenarios (Shi al., 2023). For example, Shampoo (Gupta et al. 2018) approximates the full second moment matrix using a Kronecker product, leading to the following preconditioned gradient:
L LGG 1/4 L- GG R G
(20)
where L E Rmxxm, R E Rnxx are the left and right preconditioner matrices, and G E IR x
is
:
the gradient. Many other higher-order methods follow this framework (Martens & Grosse 2015
Morwani et al. 2024, Duvvuri et al. 2024). These methods incur O(m² + n²²) additional memory overhead and require periodic computation of roots of L and R with O + n²*) computational cost. This complexity significantly exceeds that of our proposed method, as demonstrated in 7 Comparing (20) and (18) reveals that our method applies preconditioning only to the low-rank side of LoRA, resulting in negligible overhead. Furthermore, unlike our provably monsffommaiooninvvaiiin approach, Shampoo-based methods lack this property.
LARS You et al. 2017) and Lamb (You et al. 2020) are layer-wise adaptive optimization methods originally designed for large batch training. They dynamically adjust the update norm for each weight matrix based on its current norm, which ensures scalar scale invariance. Nonetheless, they still lack transformation invariance.
large language models (LLMs) grow in size, full fine-tuning on down- stream tasks becomes increasingly intturce-ittensive. Parameter-efficieet ods such as (Houlsby et al. 2019: He et al. 2022b a Lester et al., 2021 Li & Liang 2021) have emerged to address this issue by reducing the number of trainable paramters. As popular PEFT algorithm, LoRA (Hu et al. 2022) has been the subject of extensive research, with numerous varia- tions and improvements proposed. One line of research focuses on dynamically adjusting the LoRA rank during training. This includes DyLoRA (Valipour et al.] 2023). IncreLoRA (Zhang et al.) 2023a), and AdaL.oRA (Zhang al. (0033)) Another approach involves enhancing LoRA perfor- mance through the addition of extra scaling matrices, which includes DoRA (Liu et al. 2024) and DeepLoRA (Yaras al., (2024). These directions are orthogonal to our work.
Variants of LoRA.
fine-tuning (PEFT) meth-
Regarding LoRA optimization, Hayou et al. (2024) highlight the limitations of traditional optimiz- ers they fail achieve efficient feature learning. To address this issue, they propose LoRA+ which uses two different learning rates 7A and 1) for LoRA weights. However, this leads to an extra hyperparameter to be tuned in practice. In contrast, Zhang Pilanci (2024) propose the use of matrix preconditioning methods to achieve efficient feature learning. They propose the use of Riemannian gradient descent for LoRA optimization. As far as we know, Riemannian gradient de- scent is the only method in the literature that satisfies transformation invariance. However, similar to gradient descent, Riemannian gradient descent does not incorporate momentum and adaptivity, so it performs worse than Adam in their experiments. To improve the performance, they propose to combine Riemannian gradient descent with element-wise Adam, which becomes ScaledAdam. However, this combination makes ScaledAdam no longer transformation invariant.
EXPERIMENTAL RESULTS
We evaluate the proposed LoRA optimizer against other optimizers across range of datasets. This includes the Super-Natural Instructions dataset, a comprehensive collection of diverse NLP tasks, as well as four standard LLM benchmarking datasets.
We compare the following optimizers:
Adam Kingma & Ba 2014): The most widely used default optimizer for LoRA finetuning.
8

---

Published as conference paper at ICLR 2025
Table 1: Experimental results on the Super-Natural instruction dataset.
Cause Effect Classification
Coreference Resolution
Title Generation
Data to Text
Model
Optimizer
Global
Adam
58.93
77.06
51.30
55.52
50.51/74.54
LoRA+
58.84
76.08
51.32
55.68
49.76/74.20
ScaledAdam
58.71
77.55
51.16
55.69
49.40/74.01
Gemma-2B
Shampoo
58.11
77.17
51.30
55,48
50.79/74.74
Lamb
60.97
80.69
52.26
55.85
53.53/76.43
LoRA-RITE
61.26
82.02
52.26
55.98
55.11/77.12
Adam
67.17
86.05
51.58
55.38
58.46/78.17
LoRA+
65.50
86.67
51.51
55.34
ScaledAdam
65.79
85.05
51.61
55,40
57.32/77.92
Gemma-7B
Shampoo
66.29
85.62
51.86
55.43
57.99/78.21
Lamb
69.62
86.57
51.87
55.5
57.79/78. 18
LORA-RITE
71.26
88.14
52.17
55.62
59.71/79.05
Table 2: Experimental results on LLM benchmarking datasets.
Model
Optimizer
HellaSwag
ArcChallenge
GSM8K
OpenBookQA
Avg.
Adam
83.76
45.31
24.26
64.0
54.33
LoRA+
83.75
45.31
23.65
64.4
54.28
Gemma-2B
ScaledAdam
83.52
45.22
23.96
64.8
54.38
Shampoo
83.26
44.88
23.35
63.6
53.77
Lamb
86.60
47.35
26.76
68.0
57.18
LORA-RITE
87.28
49.06
30.10
68.8
58.81
Adam
94.07
54.78
48.37
77.60
68.71
LoRA+
93.99
54.01
48.75
77.60
68.59
Gemma-7B
ScaledAdam
93.31
52.90
48.07
75.80
67.52
Shampoo
94.15
52.47
49.05
76.80
68.12
Lamb 95.11
69.80
50.64
83.20
74.69
LORA-RITE
95.59
71.76
55.50
84.80
76.91
LoRA+ (Hayou et al. 2024): Adam with different learning rates for and B. We set learning rate of B be 4 times large than A, which is the value they used for decoder models.
ScaledAdam (Zhang & Pilanci 2024): A variant of Adam designed for LoRA optimization. Shampoo (Gupta et al.] 2018): A well-known adaptive matrix preconditioning method. To obtain similar training time the other methods, the block size is set 512 and the preconditioners are updated every 100 steps.
Lamb (You al. 2020): A variant Adam that normalizes the updates for each layer based on the norm of the parameters.
LORA-RITE: Our proposed optimizer that is transformation invariant.
For each optimizer applied on each data set, we search for the best learning rate from 2 10 to 2 10-2 The other hyperparameters are listed the Appendix. For most of the experiments we chose rank r 16 for LoRA, based on the ablation study over the rank. We conduct experiments on Gemma (Gemma Team et al. 2024) 2B, 7B, and mT5-XXL (Xue et al. 2021) using TPUs.
Results on Super-Natural Instruction Dataset. The Super-Natural instruction dataset (Wang) al., 2022) contains collection of 1600+ NI.P tasks, including both classification and genera- the
tion tasks. We use 10% split of the data for validation. Following (Wang et al. 2022), exact match accuracy to evaluate classification and ROUGE-L score to evaluate generation tasks.
Table presents the performance of individual fine-tuning two classification and two genera- tion tasks for 2,000 steps. It also includes the performance of fine-tuning on the global training set of over 1,600 tasks for 10,000 steps, reporting both exact match accuracy and ROUGE-L score evaluated on the global validation set. As shown in Table [ our proposed method demonstrates su- perior performance across both classification and generation tasks. Compared Adam, method

---

Published as a conference paper at ICLR 2025
Table Ablation study on different ranks and different model architectures.
Gemma-2B (rank=4)
Gemma-2B (rank=16)
mTS-XXL (rank=4)
mTS-XXL (rank=16)
Adam
63.00
64.0
72.00
72.20
ScaledAdam 63.00
64,8
70.80
74.60
Lamb
67.80
68.0
70.40
73.40
LORA-RITE 70.40
68.8
74.80
75.00
Table Number of training steps per second for different optimizers. Shampoo preconditioner updated once every 100 steps. LoRA-RITE has small overhead compared with first-order methods.
Adam
LoRA+
ScaledAdam
Shampoo
Lamb
LoRA-RITE
Gemma-2B
0.948
0.930
0.917
0.837
0.929
0.878
Gemma-7B
0.120
0.120
0.114
0.112
0.116
0.114
achieves 2.3% to 4.9% accuracy improvements on the classification tasks and also shows significant improvements in the global training setting. Furthermore, we found that Lamb performs well on some of the data sets but there's 's still significant gap between Lamb and LoRA-RITE. Since Lamb enforces scalar scale invariance but not transformation invariance, this result implicitly suggests that transformation invariance crucial for achieving optimal performance.
Results on other LLM Benchmarking Datasets. We also evaluate the performance common LLM benchmarking datasets, including HellaSwag (Zellers et al. 2019), AreChallenge (Clark al.) 2018), GSM8K (Cobbe et al. (02T), and OpenBookQA (Mihaylov et al. (018). The summary information of these datasets is the Appendix. The results are presented in Table 2 observe that the trend is similar to the SuperNatural instruction results, where LORA-RITE achieves the best performance on all the data sets, and Lamb usually the second best optimizer.
We can
Ablation Study. We conduct an ablation study on the choice of different LoRA ranks and model architectures. Specifically. we considered rank 4 and 16 on both Gemma 2B (decoder only) and mT5-XXL cencoder-decoder
the OpenBookQA dataset. As can from Table 3 our pro- posed method performs consistently well across different LoRA ranks. Furthermore, our method be successfully applied to mT5-XXL which has encoder-decoder architecture, showing the generalizability of the proposed optimizer.
Training Speed Comparison. We compare the training speed of different optimizers. Table 4 shows the number of training steps per second for different optimizers with LoRA rank 16 on the OpenBookQA dataset using TPUv5e. As we can see, LoRA-RITE is only 8% slower than Adam on Gemma 2B, while the difference decreases to 5% when model size increases to 7B. Also, Shampoo is slower than LORA-RITE in this case despite the fact that it recomputes the preconditioner with much lower frequency (once every 100 steps). This is due to our approach of preconditioning only the low-rank side of the LoRA factors.
6 CONCLUSIONS
Current LoRA optimization techniques lack transformation invariance, which implies that equiva- lent LoRA parameterizations can yield significantly different updates. This hinders efficient feature learning and often leads suboptimal solutions in practice. We introduce a novel, transformation- invariant optimization algorithm with comparable time and memory overhead to Adam. Our algo- rithm consistently achieves higher accuracy than existing LoRA optimizers across diverse datasets and models.
Limitations. Although this work introduces better optimizer for LoRA, it is important to ac- knowledge that LoRA itself has limitations. For instance, LoRA has smaller representational power and may result minor performance decrease compared to full fine-tuning. Also, how to se- lect rank to strike a good trade-off between efficiency and accuracy may be non-trivial in practice. The work focuses on addressing transformation-invan -nvvariance when the optimization problem can be written in the form of fABB and this assumption may not hold for other parameter-efffiient structures beyond LoRA. Applying LORA-RITE to ensure transformation invariance for the other more complicated LoRA variants will be interesting future direction.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank Sashank Reddi and Corinna Cortes for useful feedback. The code for our project is available at
com http.///gggee..ynns11//Looooo. RITE
REFERENCES
Guanzheng Chen, Fangyu Liu, Zaiqiao Meng, and Shangsong Liang. Revisiting crrameteer-fficien tuning: Are we really there yet? In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 2612-2626, 2022.
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang. Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, et al. Symbolic discovery of optimization algorithms. Advances in neural information processing systems, 36, 2024.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv:
0545011, 100.00553.55
2018.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168 2021.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Sai Surya Duvvuri, Fnu Devvrit, Rohan Anil, Cho-Jui Hsich, and Inderjit Dhillon. Caspr: Com- bining axes preconditioners through kronecker approximation for deep learning. In Forty-first International Conference on Machine Learning. 2024.
Vladimir Feinberg, Xinyi Chen, Y Jennifer Sun, Rohan Anil, and Elad Hazan. Sketchy: Memory- efficient adaptive regularization with frequent directions. arXiv preprint arXlv:2302.03764, 2023.
Thomas Mesnard Gemma Team, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju. Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, Pouya Tafti, Léonard Hussenot, and et al. Gemma. 2024.
doi: 10347400KAGLLE/M/300L URL https: www wwggle.......................... com/m.
3301
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor opti- mization. In International Conference Machine Learning, pp. 1842-1850. PMLR, 2018.
Soufiane Hayou, Nikhil Ghosh, and Bin Yu. Lorat: Efficient low rank adaptation of large models.
In Forty-first International Conference on Machine Learning, 2024.
Elad Hazan et al. Introduction to online convex optimization. Foundations and Trends in Opti- mization, 2(3-4):157-325, 2016
Junxian He, Chunting Zhou, Xuezhe Ma, Taylor Berg-Kirkpatricke and Graham Neubig. Towards unified view of parfieter-efffcient transfer learning. In International Conference on Learning Representations, 2022a.
Shwai He, Liang Ding, Daize Dong, Jeremy Zhang, and Dacheng Tao. Sparseadapter: An easy approach for improving the parameter-efficiency of adapters. In Findings of the Association for Computational Linguistics: EMNLP 2022, pp. 2184-2190, 2022b.
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, An- drea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficient transfer learning for nlp. In International conference on machine learning, pp. 2790-2799, PMLR, 2019.
Edward J Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation large language models. International Conference on Learning Representations, 2022.
11

---

Published as a conference paper at ICLR 2025
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Brian Lester, Rami AI-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Pro- cessing, pp. 3045-3059, 2021.
Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. In Proceedings of the 59th Anmual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference Natural Language Processing (Volume 1: Long Papers),
pp. 4582 0092,499,22 2021.
Shih- Yang Liu, Chien-Yi Yi Wang, Hongxu Yin, Pavlo Molchanov, Yu-Chiang Frank Wang, Kwang- Ting Cheng, and Min-Hung Chen. Dora: Weight-decomposed low-rank adaptation. arXiv preprint vrXiv:2402.09353 2024.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2017. URL https::/ www.pp..ssssanncccolsoaaaaaa o.t.oa/aa.e/////o..caaaa.c...oo. orq/
CorpusID:53392277
James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factored approxi- mate curvature. In International conference on machine learning. pp. 2408-2417. PMLR, 2015.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? new dataset for open book question answering. In EMNLP, 2018
Depen Morwani, Itai Shapira, Nikhil Vyas, Eran Malach, Sham Kakade, and Lucas Janson. A new perspective on shampoo's preconditioner. arXiv preprint arXiv:2466. 2024.
Shai Shalev-Shwarzz et al. Online learning and online convex optimization. Foundations and Trends in Machine Learning. 4(2):107-194, 2012.
Hao-Jun Michael Shi, Ysung-Hsien Lec, Shintaro Iwasaki, Jose Gallego-Posada. Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, and Michael Rabbat. A distributed data- data-parallel py- torch implementation the distributed shampoo optimizer for training neural networks at-scale. arXiv preprint arXiv:2309.0649) 2023.
T. Tieleman and G. Hinton. Lecture 6.5 RMSProp: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural networks for machine learning. 2012.
Mojtaba Valipour, Mehdi Rezagholizadeh, Ivan Kobyzev, and Ali Ghodsi. Dylora: Parameter- efficient tuning of pre-trained models using dynamic search-free free low-rank adaptation. In Pro- ceedings of the 17th Conference of the European Chapter of the Association for Computational Linguistics, pp. 3274-3287, 2023.
Yizhong Wang. Swaroop Mishra, Pegah Alipoormolabashi, Yeganeh Kordi, Amirreza Mirzaei, Atharva Naik, Arjun Ashok, Arut Selvan Dhanasekaran, Anjana Arunkumar, David Stap, et al. nupper-naturalinstuution Generalization via declarative instructions on 1600+ nlp tasks. Pro- ceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp- 5085-5109, 2022.
Linting Xue, Noah Constant, Adam Roberts, Mihir Kale, Rami AI-Rfou, Aditya Siddhant, Aditya Barua, and Colin Raffel. mt5: A massively multilingual pre-trained text-to-text transformer, 2021. URL
Can Yaras, Peng Wang, Laura Balzano, and Qing Qu. Compressible dynamics in deep overpa- rameterized low-rank learning & adaptation. In Forty-first International Conference on Machine Learning, 2024.
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arXiv: 1708 1708.03888. 2017.
12

---

Published as a conference paper at ICLR 2025
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference Learning Representattonss 2020.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a ma- chine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
Fangzhao Zhang and Mert Pilanci. Riemannian preconditioned lora for fine-tuning foundation mod- els. Forty-first International Conference on Machine Learning, 2024.
Feiyu Zhang, Liangzii Li, Junhao Chen, Zhouqiang Jiang, Bowen Wang, and Yiming Qian. In- crelora: Incremental parameter allocation method for parameter-efficient fine-tuning. arXiv preprint arXiv:2308. S.Ki::2308./0433 12043, 2023a.
Qingru Zhang, Minshuo Chen, Alexander Bukharin, Pengcheng He, Yu Cheng, Weizhu Chen, and Tuo Zhao. Adaptive budget allocation for parameter-effie efficceet fine-tuning. In The Eleventh Inter- national Conference on Learning Representations, 2023b.
13

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 HYPERPARAMETERS
Table 5 shows our hyperparameters. We set weight decay and dropout probability to as our early experiments suggest that setting non-zero value does not improve the performance of the baselines.
Table 5: The setting for hyperparameters.
Hyperparameter
Value
Learning rate
2 10
6 to * 10
Weight decay
0
Dropout prob
LoRA target
Oproj- Kproj Uprojs Oproj
LoRA rank
16
LoRA
16
Batch size
16
Train step
2000
LR schedule
Linear decay
Warmup step
100
Evaluation period
100
Momentum 31
0.9
Second moment /32
0.999
A.2 DATASETS
Table shows the summary information of the LLM benchmarking datasets. We use the test evaluate ArcChallenge, as it is much larger than the development set.
Table 6: Summary information of the LLM benchmarking datasets.
Dataset
#Train
#Dev
#Test
Split for Eval
HellaSwag
39905
10042
10003
Dev
AreChallenge
1119
299
1172
Test
GSMSK
7473
NA
1319
Test
OpenBookQA
4957
500
500
Dev
A.3 PROOF OF THEOREM I
Let ||A = (nn", |B||| = (n²)), VZ =(nn), = =((n")) where is the learning rate
and n the network width. Since
AT from chain rule we know VA vZB and VB V A. Since the update rule is symmetric, we can express the updates as
| 0(n aa tyb
[|8811| 0(n ((n=htuvetzcddd
the update rule is scalar scale invariant, then for any A n² A1, B2 = n B, we have
| Bill = A B2ll.
which means
za + (y + 1)b + zc + d = xaa + 8) + (v 1)(b 6) zc d,
thus x5 - (y 1)8 = for all 5, which means y = x Consequently, we have
| B 0 ma [ 01200zbb+cct
14

---

Published as a conference paper at ICLR 2025
Table Time and space complexity comparison for LoRA optimization.
Algorithm
Time Complexity
Space Complexity
Forward/Backward
((mm)
((nm) Om + n n2y²)
Full Matrix Adagrad (Duchi et al.
2011
(m n²rra)
Adam (Kingma & Ba
2014
O + nr)
O + nr)
Lamb You et 2020
O nr)
O + nr)
Shampoo Guppta et |
O + n² +
(( + n² + +
KFAC (Martens & Grosse
2015
(m² + n1² r8)
O + n2 + 12)
ScaledAdam (Zhang & Pilanci
2024
O nr2)
O(mr + nr)
LORA-RITE: (our proposed)
O(mr² + nr?)
Omrr nr + +r2)
Similarly, we have
A | = (n (y+++b++++++++e (v+1))++o+)) = =nz (mm
Since these two are equal, we can achieve efficient feature learning
| | |||| = ||AA|||| ||||=||| = 0(1),
where æ: is the input vector, by selecting proper learning rate 17.
A.4 PROOF OF THEOREM
In Algorithm I for two equivalent LoRA pairs (A,, Bi), (A2, B2), their obtained basis UB, and UB could be different.
Consequently, introduce the concept of consistency, which means matrix output of f(A, B) is the same across all equivalent LoRA pairs, up to the choice of the basis.
Definition (Consistency). For matrix XA R we call it consistent if
XA, UB.
UA, E mR
for all equivalent LoRA pairs (A1, B1), ) (A2, B2). Similarly, for HA Rrrx, we call it consistent
if
UB, AA, UB, = UB, HA, uxn for all equivalent LoRA pairs (A1, B1), (A2, B2).
Additionally, to simplify the notation, from now on we use A, to stand for A,, B₁ and use Å, B to stand for A2. B2-
Then, we proceed to prove Theorem First, one should note the fact that
UBU U U
for all equivalent LoRA pairs. Consequently,
UB(VA)]VALL UBU - VUUBUR - VZ V UJ U (VA) VAUT
(21)
B
which shows that (VA) V is consistent.
Second, we can prove VA, is consistent by mathematical induction. The base case VA. 0 is consistent. Assuming VAL. is consistent, then
UB... PA. (U.)
-UB... (UB,) UB... VAL (UB, (UB.))
(22)
-UR... R U A._. VAL, U UR.., (Un,)) (U,
-UB.,, PA. V V P (U (
15

---

Published as a conference paper at ICLR 2025
which means PA. VA. PP, and thus
VA, = PA, VA,1 PA, + AA, VVA
consistent.
Lastly, since
UB (VA 1) = UUBVA +PA -1/
we have
SAUUB VA.(V, + 1)
-VZBB (VA, +PA, -
VZ +PA.B BB, UB)
-VZ. B. VA. UB +PA. B. UB) )
U U
UUU
Thus, can similarly proof by induction that both SA, and M A, are consistent, which completes our proof.
A.5 PROOF OF THEOREM 3
For convenience, for matrix g2xrr H Rr we define
X H TxXXXX 11/2.
We also utilize the following lemma for online optimization.
Lemma (Lemma 5.13 Hazan et al. (0116).. For online optimization, if 0, is updated as 0.+1 0, Xg,, then we have
V0. (0: 0²)
1
00
0 x-
Σ (gt) X.g:
2m
0, -
-.).
2m
1=2
Lemma 2 (Lemma 5.13, 5.14
Hazan et al. (0166). For arbitrary matrix G, E Rmxr H.
E G.T G, we have
G w/ 2 1/3
Proof of Theorem 3
Since we are preconditioning each layer independently, all three terms in Lemma can written as summation over the L layers. For simplicity, from now omit the summation and the subscript for layers.
For our method, the preconditioner X A, is as follows,
X A, R (VA. + IPA.I)
We define the unmagnified preconditioner
X (VA. + IA.1)
Then for the factor, we have
vec A.)
vec(8A.)) Σ Tr(VA, 6A.) = Σ A
(23)
T VAL 2n
AA V -
16

---

Published as a conference paper at ICLR 2025
where the inequality comes from Lemma and the fact that
XA. = (VA. + +AA.1) 1/2 3 E VA) 1/2.
This completes our proof.
A.6 PROOF OF THEOREM 4
To prove Theorem 4 we start by bounding the escape mass.
Lemma 3. For the escaped mass PA in Algorithm 7 we have
PAT Σ
-
Proof of Lemma]
Since PA.PA, I I, by Ostrowski's inequality have
(VA.,..) 2 (PAA, VA.,, PX,)
for all i. Consequently, we have
d (VA, V.L PA,) Tr(VA Tr(PA, VAL_, PX.)
Thus,
T PAT
Tr(VA, .) Σ Tr(PA, PA.)
T
Tr((VA,) VAJ -[VAA))=
Proof of Theorem
We prove Theorem by by bounding the right-hand side of Lemma 1 Since the first term is a constant and we already bound the second term in Theorem here we only need bound the third term.
For the third term, we have
A.
A.
(14, A.) 2 X QA, Sh QA, 0% S PD x PA. PX
A
<D
where the last inequality comes from our assumption. Consequently, since
Tr A. ) Tr A. PX.) 2 Tr A. + +c(p/2 - -PA. 1))
we have
T
T
1 - PA. A. P < Σ
Tr A. PA. A. PA. TT Ar +cp
<TH( "
Tc s (V2
V A (A)
1=1
where the last inequality comes from Lemma
Summing up the bound for the second and the third term, we get
T V2 12+ c DD TT Σ VA,
(2n) + n
17

---

Published as a conference paper at ICLR 2025
Choosing = = we have
T (2 v2+ c) Tr(
-ODADGTTT
i=1
which completes the proof.
A.7 TRAINING LOSS CURVE VISUALIZATION
To cross-validate the effectiveness of - we plot the training loss curve of each method the Super-Natural instruction dataset and the OpenBookQA dataset. Figure 2 shows that LoRA- RITE has the lowest training loss, which demonstrates the effectiveness of our method.
Adam
Adem
14
LORA+
ScaledAdam
Scaledhdam
Shampoo
Shampoo
Lamb
0.200
Lamb
LoRA.
0.15
1.0
0.010
0.005
0.000
200
400
600
600
Steps
Train Steps
Super-Natural instruction
(b) OpenBookQA
Figure 2: The training loss curve for the Super-Natural instruction dataset and the OpenBookQA dataset.
A.8 UPDATE MAGNITUDE VISUALIZATION
To visualize the update magnitude the two LoRA factors, we plot the update norm divided by the weight norm, || || and || ||
Figure and Figure 4 show that for conventional optimizers, factor barely changes, while LoRA- RITE able to learn the factor A effectively. This demonstrates the importance of transformation invariance.
Adam
LoRA+
LoRA
ScalesAdam
ScaledAdam
Shampoo
10F
Shampoo
[[|]]]
LORA-RITE
[[88]] [[8]
LoRA RITE
10-
10-2
[[4|]]
10
10-1
10
2000
4000
6000
3000
2000
4000
6000
Train
Steps
Update magnitude of the A factor
Update magnitude of the B factor
Figure 3: The update magnitude of A and B for the Super-Natural instruction dataset.
A.9 ABLATION STUDY ON DIFFERENT RANKS
To study the effect of different LoRA ranks, we conduct additional ablation study on different datasets.
18

---

Published as a conference paper at ICLR 2025
10
Adam
LORA+
Scaled/dam
Shamooo
109
Shampoo
[00|]] HAII
[181)
LaRU-RIT
10
10
10
10
200
400
600
800
Train Steps
Steps
(a) Update magnitude of the factor
(b) Update magnitude of the B factor
Figure 4: The update magnitude of A and B for the OpenBookQA dataset.
Table 8: Ablation study on different ranks on Gemma-2B on the LLM benchmarking datasets.
Optimizer
Rank
HellaSwag
ArcChallenge
GSM8K
OpenBookQA
Avg.
== 4
81.83
42.32
20.92
63.0
52.02
Adam
16
83.76
45.31
24.26
64.0
54.33
r = 64
84.56
46.67
26.08
67.0
56.08
r
81.95
44.80
21.15
63.0
52.73
ScaledAdam
r = 16
83.52
45.22
23.96
64.8
54.38
64
84.42
48.21
26.61
67.0
56.56
4
86.01
46.67
25.25
67.8
56.43
Lamb
16
86.60
47.35
26.76
68.0
57.18
r 87.83
47.53
29.04
62.8
56.80
87.08
49.57
29.49
70.4
59.14
LORA-RITE
r = 16
87.28
49.06
30.10
68.8
58.81
64
87.89
49.91
31.46
68.8
59.52
As we can see from Table 8, higher rank generally improves LoRA performance, approaching full fine-tuning. This explains why the performance gap between LoRA-RITE and other methods rows at higher ranks, they all converge towards the results of full fine-tuning.
Additionally, one can observe that LoRA has inherent regularization properties. As noted in pre- vious research (Chen et al. 2022), this means that sometimes a lower rank can actually lead to better performance. This effect depends on factors like model generalization and training data size. This explains why LoRA-RITE achieves better performance at rank 4 instead of 16 and why Lamb achieves better performance at rank 16 than rank 64.
A. BEST LEARNING RATE FOR DIFFERENT OPTIMIZERS
we list the best learning rate for each optimizer the LLM benchmarking datasets. We observe that PORA-RITE and Lamb usually prefer a larger learning rate than the other baselines.
In Table 9
19

---

Published as a conference paper at ICLR 2025
Table 9: Best Learning Rate for Different Optimizers on LLM benchmarking datasets.
Model
Optimizer
HellaSwag
ArcChallenge
GSM8K
OpenBookQA
Adam
le-5
5e-5
le-5
5e-5
LoRA+
le-5
5e-5
le-5
5e-5
Gemma-2B
ScaledAdam
5e-5
5e-5
le-5
2e-4
Shampoo
le-5
5e-5
5e-5
5e-5
Lamb
5e-3
5e-3
5e-3
5e-3
LoRA-RITE
2e-4
1e-3
2e-4
2e-4
20


Reviews:
[{"review_id": "xm5Fb4OIMv", "reviewer": "Reviewer_zd3u", "summary": "The paper introduces an enhancement to the existing LoRA method. LoRA, which stands for Low-Rank Adaptation, is a parameter-efficient fine-tuning technique widely employed in training Large Language Models. However, a significant limitation of LoRA is that matrix $A$ is typically much larger than matrix $B$, leading to reduced training efficiency of LoRA. The authors leverage the concept of transformation invariance to formally capture this phenomenon, demonstrating that LoRA weight updates remain consistent across different decomposition results. They further define a weaker version, termed \"scalar scale invariance,\" which only examines whether decompositions with varying magnitude distributions in matrices $A$ and $B$ yield different weight updates. The authors demonstrate that transformation invariance is crucial for efficient feature learning, propose a method to preserve this invariance while considering optimizer implementation, and conduct experiments across various models to illustrate the advantages of their approach.", "questions": "Will your method be susceptible to numerical instability? If not, why? If so, how will you address this issue?\n\nIn terms of contributing to the LLM community, open-sourcing the related code would be highly appreciated.", "limitations": "", "rating": 10, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The authors take a significant step beyond simple intuitions by delving into the core of training efficiency, proposing a novel approach to fully achieve transformation invariance in LoRA models. For soundness, they provide rigorous proofs for all mathematical statements in the paper and conduct extensive experiments across various LoRA advancement methods and benchmarks to demonstrate the superiority of their method. In terms of contribution, given that LoRA is a widely-used training method and their work is computationally inexpensive, individuals training Large Language Models and potentially other AI models can greatly benefit from their advancements.", "weaknesses": "The paper lacks visual illustrations of loss curves to cross-validate the effectiveness of their method in accelerating convergence. Although they state that matrix $A$ remains nearly identical during training, they do not provide visual evidence of how the magnitude of $A$ updates after applying their method, which could further validate its effectiveness. Additionally, the authors do not address potential numerical instability issues. Specifically, their algorithm involves inverting the matrix $R_A$, which could be zero at the initial training step. The conclusion and limitations section is somewhat lacking, as it only discusses the limitations of the base LoRA method rather than potential limitations of their proposed method.", "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": "a7yaoCsopJ", "reviewer": "Reviewer_wRHA", "summary": "This paper points out the issue of standard LoRA training on the decomposed matrices A and B and then enforces the transformation invariant property into optimization, which further improve the representation capability for better performance", "questions": "See weaknesses\n\nQuestions:\n\n1. As analyzed in at line 322, from time complexity perspective, the proposed method is r times slower than Adam, but in Table 4, it is only 8% slower, it will be interesting to elaborate more the inconsistency.\n\n2. At table 3, when LoRA rank is higher, the improvement over other methods become smaller, I wonder the authors have any insight of that?\n\n3. Why the LAMB paper is cited twich in different way? (line 642, 646)\n\n4. At line 345, LoRA-Rite should be LoRA-RITE for consistency.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The paper propose a new optimization approach that retains the transformation invariant for the LoRA-type fine-tuning, which is widely used in the large model fine-tuning. Moreover, the paper presented the convergence analysis theoretically. \n\n2. The experimental results demonstrate significant improvement with marginal computation increase, results in a better trade-off compared to all SOTA methods.", "weaknesses": "1. As the paper focuses on the optimization, a convergence analysis should be conducted to better justify the proposed method, e.g. the norm of A and B, like in figure 1.\n\n2. The analysis over experimental results are limited, e.g., for some datasets, the proposed method demonstrates significant performance gain compared to LoRA (Adam), what is property of dataset such that the proposed optimization can results such 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": {}}, {"review_id": "3H9p6cksoZ", "reviewer": "Reviewer_vefd", "summary": "This paper focuses on improving LoRA, a novel adaptive matrix preconditioning method, and names the proposed method as LoRA-RITE for LoRA optimization. Authors argue that LoRA’s ability to decompose the fine-tuned weight Z in multiple ways is not the best approach for optimization, and they propose a new way for invariant decomposition. Authors have applied the new optimization function, named LoRA-RITE, on Gemma for several datasets and show strong performance on multiple LLM benchmarks.", "questions": "Please check weaknesses section. Authors mentioned that they are searching for the learning rate between 2e-6 and 2e-2; it would be interesting to present the best learning rate for different strategies in this work after the search.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "* The overall mathematical proofs are reasonable and appear to be correct.\n* With the implementation of LoRA-RITE, authors have shown significant improvement over other optimizers on different language benchmarks.", "weaknesses": "I do not have major technical concerns for this paper. It is relatively solid in both performance and implementation analysis. Authors mentioned that they are searching for the learning rate between 2e-6 and 2e-2; it would be interesting to present the best learning rate for different strategies in this work after the search.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)