ReLoRA: High-Rank Training Through Low-Rank Updates
Authors: Vladislav Lialin
Sherin Muckatira
Namrata Shivagunde
Anna Rumshisky
Venue: ICLR 2024 poster (2024)
Keywords: language models
pre-training
training efficiency
parameter-efficient fine-tuning
lora
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: ReLoRA is a parameter-efficient method that can be used during model pre-training stage. We demonstrate it's efficacy by training LMs with up to 1B parameters and show significant speed and memory improvements

Abstract: Despite the dominance and effectiveness of scaling, resulting in large networks with hundreds of billions of parameters, the necessity to train overparameterized models remains poorly understood, while training costs grow exponentially. In this paper, we explore parameter-efficient training techniques as an approach to training large neural networks. We introduce a novel method called ReLoRA, which utilizes low-rank updates to train high-rank networks. We apply ReLoRA to training transformer language models with up to 1.3B parameters and demonstrate comparable performance to regular neural network training. ReLoRA saves up to 5.5Gb of RAM per GPU and improves training speed by 9-40% depending on the model size and hardware setup. Our findings show the potential of parameter- efficient techniques for large-scale pre-training. Our code is available on GitHub.

Published as a conference paper at ICLR 2024
RELORA: HIGH-RANK TRAINING THROUGH LOW-RANK UPDATES
Vladislav Lialin .. Sherin Muckatira', Namrata Shivagunde', and Anna Rumshisky
'University of Massachusetts Lowell
Eleuther AI Ammaoon
ABSTRACT
Despite the dominance and effectiveness of scaling, resulting in large networks with hundreds of billions of parameters, the necessity to train overparameterized models remains poorly understood, while training costs grow exponentially. this paper, explore parameter-efficier training techniques as approach training large neural networks. We introduce novel method called ReLoRA, which utilizes low-rank updates train high-rank networks. We apply ReLoRA training transformer language models with up to 1.3B parameters and demonstrate comparable performance to regular neural network training. ReLoRA saves up to 5.5GB of RAM per GPU and improves training speed by 9-40% depending on the model size and hardware setup. Our findings show the potential of parameter- efficient techniques for large-scale pre-training. Our code is available on GitHut
INTRODUCTION
Over the past decade, the machine learning field has been dominated by the trend of training increasingly overparameterized networks adopting the "stack more layers" approach (Krizhevsky et al., 2012 He et al. [016)] Kaplan et al. 2020). The definition of large network has evolved from models with 100 million (Simonyan and Zisserman 2015) et al. 2018) hundreds of billions (Brown et al. 2020: Chowdhery et al., 2022) of parameters, which has made computational costs associated with training of such networks prohibitive to most of the research groups. Despite this, the necessity to train models which can have orders of magnitude more parameters than the training examples (Brown et al. 2020
Chowdhery et al. 2022; Fedus et al. (022). is poorly understood Allen-Zhu et al. 2019: Zhang et al. 2021).
theoretically (Jacot et al.) 2018:
Alternative approaches to scaling, such as more compute-efficient scaling optima (Hoffmann et al. 2022), retrieval-augmenneed models (Khandelwal et al. [2000] Borgeaud et al. 2022), and the simple approach of training smaller models for longer (Tovronn et al., (023). have offered new trade-offs. However, they do not bring closer to understanding why need overparameterized models and rarely democratize the training of these models. For example, training RETRO (Borgeaud et al. 2022) requires complex training setup and infrastructure capable of quickly searching over trillions of tokens, while training LLaMA-7B (Touvron et al. 2023) still requires hundreds of GPUs.
In contrast, approaches like zero-redundancy optimizers (Rajbhandari al. 2020), 16-bit training (Micikevicius et al.) 2018), 8-bit inference (Dettmers et al.) 2022), and parameter-efficient fine-tuning (PEFT) (Lialin al., (023) have played crucial role making large models more accessible. Specifically, PEFT methods have enabled fine-tuning of billion-scale language or diffusion models on consumer hardware. This raises the question: Can these approaches also benefit pre-training?
Our Contribution In this study, we introduce ReLoRA which uses individually low-rank updates that aggregate during the training process to train a high-rank network. We empirically demonstrate that ReLoRA performs high-rank update and achieves performance similar to regular neural network training. The components of ReLoRA include initial full-rank training of the neural network (similar
"oorespondance to vlad. 1 ia 1 @n@gmai com
githuh.co//guiaaice///llnnn

---

Published as a conference paper at ICLR 2024
5.0
250M
250M ReLoRA
(99M trainable)
Loss
Trainable Params
2500
5000
7500
10000
12500
15000
17500
20000
250
100
0
2500
5000
7500
10000
12500
15000
17500
20000
Step
Figure 1: Training loss for 250M models. ReLoRA learns a high-rank network through a sequence of low-rank updates. It outperforms networks with the same trainable parameter count and achieves similar performance to training a full network at 100M+ scale. The efficiency of ReLoRA increases with the model size, making a viable candidate for multi-billion-panee piitt-aillion--raaeetee training.
to Frankle et al. (2019)) LoRA training, restarts, a jagged learning rate schedule, and partial optimizer resets. We evaluate ReLoRA on transformer language models up to 1.3B parameters. Finally, observe that the efficiency of ReLoRA increases with model size, making a viable option for efficient training of cetiiiiillion-parmmette networks.
METHOD
We are interested in the rank of the sum of two matrices: rank(A + B) < rank(A) + rank(B). We know that for matrix A, rank(A) dim(A). there exists a B, rank(B) < dim(BB such that sum of them has a higher rank than either A or B.
We want to exploit this property to make flexible parameter-efficitee training method. We start with LoRA (Hu et al. (022) which is parameter-efficien fine-tuning method based on the idea of low-rank updates. LoRA can be applied to any linear operation parameterized through W e mxxx Specifically, LoRA decomposes the weight update SW into a rank-r product W Wp as shown in Equation| where
a fixed scaling factor usually equal to
oW WB WA Rinxx WB E Rxxuut
In practice, LoRA is usually implemented by adding new trainable parameters WA and WB, which could be merged back into the original parameters after training. Thus, these implementations are restricted by the rank max Wn Fank(WAW:
could restart LoRA, meaning we merge WA and WB during training and reset the values of these matrices, we could increase the total rank of the update. Doing this multiple times brings the total neural network update to:
T
Tx Σ L=
AW
SW. +
SW.
+
oW, SWAWt + ... ... #WW.. (2)
However, implementing restarts not trivial practice and requires several modifications to the optimization procedure. Unlike plain stochastic gradient descent, Adam (Kingma and Ba 2015) update guided mainly by the first and second moments of the gradient accumulated over the previous steps. In practice, Adam's B and B2 usually very high 0.9 0.999. This means that
2

---

Published as a conference paper at ICLR 2024
Learning Rate Multiplier
0.6
0.4
0.2
0.0
2000
4000
6000
8000
10000
Step
Figure 2: Jagged cosine scheduler used in ReLoRA. As base for our scheduler, we follow a standard cosine decay schedule as in Touvron al. (2023). On every optimizer reset, we set the learning rate to zero and perform quick (50-100 steps) learning rate warm-up back to the cosine schedule.
after the merge-ande and-reinn-,, continuing to use old gradient moments for W will guide it in the same direction as W]] and optimize the same subspace.
To resolve this issue, ReLoRA performs a partial reset of the optimizer state during merge-and-reine via magnitude pruning. To avoid loss diverging after an optimizer reset it also sets the learning rate with subsequent warm-up (Figure 2p. Our ablation studies (Table show that both of these modifications are required to improve the performance over LoRA. Finally, in our experiments, we found that in the case of training from scratch (random initialization) short full-rank training needed to "warm start" ReLoRA. All of this allows ReLoRA to achieve performance comparable to full-rank training, especially large transformer networks, by only training a small set of parameters at time. ReLoRA is described in Algorithm
Enhancing computational efficiency Unlike other low-rank training techniques (Schotthöfer et al. 2022, Sui al. 2023) Kamalakara et al. 2022), ReLoRA follows the LoRA approach by maintaining the frozen weights of the original network and adding new trainable parameters. At first glance, this may appear computationally inefficient; however, the differentiation between frozen and trainable parameters plays crucial role in parameter-efffiient fine-tuning (Lialin et al. 2023).
By reducing the number of trainable parameters, ReLoRA significantly reduces the memory spent on the optimizer states and enables the utilization of larger batch sizes, maximizing hardware efficiency. Additionally, reduces the bandwidth requirements in distributed setups, which are often the limiting factor in large-scale training. Furthermore, since the frozen parameters not being updated between restarts, they can be kept in a low- cow-prccision quantized format (Dettmers et al. 2023), further reducing their memory and computational impact.
Locally Low-Rank Training: Intuition Multiple studies suggest that neural network training is either completely low-rank or has multiple phases with initially high-rank and subsequent low-rank al.) (202T) shows that as the model becomes larger when
training. For example, AAghjaanynn cxample,/Ahhaaany is pre-trained for longer, the rank of the update needed to learn downstream task reduces. Arora al.) (2019) finds that SGD is biased towards low-rank solutions. The existence of Lottery Tickets early in training (Frankle et al. (2119) also partially supports this hypothesis, since training lottery ticket network could effectively be seen as low-rank approximation to the regular training process. Our empirical analysis (Section 4 shows that pre-trained neural networks exhibit high-rank updates over long trajectories (Figure However, for a sufficiently small trajectory, the training can be effectively approximated by low-rank update, which goes in-line with the findings of Boix-Adserd jet al. (2023). Given the above results, we speculate that neural network training is locally low-rank, which directly motivates ReLoRA.
3

---

Published as a conference paper at ICLR 2024
Algorithm 1 ReLoRA. 0 is model parameters, 0 is model parameters with linear layers replaced with ReLoRA, M and V are Adam optimizer states, 77 learning rate, and g is the reinit frequency.
Require: 6, M, V. 9. 1)
1: for t in warm start steps do
2: Update M, V, (Regular training for warm start)
3: end for
4: for layer in model layers do
5: if layer is linear then
6: layer +- ReLoRA(W" W W3))
7: Freeze W
8: end if
9: end for
10: for in training steps do
11: Update 0, M, V (Training step with ReLoRA)
12: if (OD q) = 0 then
for in model layers do
14:
ifl is linear then
W + (W" + sW W
W kaiming_inie (_t__innitWta A) W 0
17:
Mw: <- prune( Mw: ): Wx < (V
18:
end if
end for
20: Start 77 warmup
21: end if
22:
end for
23:
return 0
3 EXPERIMENTS
To evaluate the effectiveness of ReLoRA, apply it train transformer language model on the C4 dataset (Raffel et al. 2020) using various model sizes: 60M, 130M, 250M, 350M, and 1.3B.
In all experiments, train without data repetition (single epoch) least compute-optimal amount of data, estimated using Chinchilla Scaling Laws (Hoffmaan et al. 2022).
Architecture and training hy perparameteess Our architecture is based on transformer Vaswani cct al. 2017) and closely resembles LLaMA (Touvron et al. (2023). Namely, we use pre-normalization. RMSNorm Zhang and Sennrich 2019), SwiGLU activations (Shazeer 2020), Eh fully-connected
hidden state size Louvron et al. 2023). and rotary embeddings (Su et al. (021)) We select the number of pre- mee-raining tokens based the Chinchilla scaling laws (Hotfmann al. 2022 Architecture and training hyperparameters are presented Table
For all LoRA and ReLoRA experiments,
rank 128 initial experiments showed it to have the best perplexity/memory trade-off. You can find additional recommendations on ReLoRA hyperparameter selection Appendix A
We perform additional experiments comparing different
Params
Hidden
Heads
Layers
Learning rate
Batch size
Seq. len.
Data amount
60M
512
le-3
122K
256
1.2B
130M
768
12
le-3
154K
256
2.6B
250M
768
16
24
5e-4
590K
512
6.8B
350M
1024
24
5e-4
590K
512
6.8B
1.3B
2048
24
32
4e-4
786K
512
23.1B
Table Hyperparameters of the language models trained in this study. Batch size and data amount specified in tokens.
4

---

Published as a conference paper at ICLR 2024
60M
130M
250M
350M
1.3B
Full training
33.81 (60M)
23.65 (130M)
22.39 (250M)
18.66 (350M)
16.83 (250M)
Control
36,52 (43M)
27.30 (72M)
25.43 (99M)
23.65 (130M)
21.73 (250M)
LoRA
47.44 (43M)
34.17 (72M)
36.60 (99M)
57.11 (125M)
LoRA + Warm Start
34.73 (43M)
25,46 (72M)
22.86 (99M)
19.73 (125M)
18.23 (250M)
ReLoRA
34.46 (43M)
25.04 (72M)
22.48 (99M)
19.32 (125M)
17.27 (250M)
Training tokens
1.2B
2.6B
6.8B
6.8B
23.1B
Table 2: Language model perplexity when trained using each of the above methods. Number of trainable parameters for each model in (brackets). Control baseline is full-rank training model with the same total number of parameters as the number trainable parameters in low-rank training.
CoLA
STS-B B
MRPC
RTE
SST2
MNLI
QNLI
QQP
Avg
Full-rank pretrained
35.43
83.85
76,96
64.26
88.99
70.98
83.38
84.49
73.54
Not pretrained
7.59
22.73
67.00
51.15
82.61
60.04
67.92
78.40
54.68
ReLoRA
31.07
83.33
78.43
60.65
89.45
72.27
83.93
86.01
73.14
Table 3: Applying ReLoRA to fine-tune 350M models pre-trained full-rank and using ReLoRA. We observe minimal differences between the models.
rank choices for the 1.3B model in Section 4.1 We use bfloat 16 for all floating point operations and FlashAttention (Dao et al. 2022) for effective attention computation.
ReLoRA and baselines setup our experiments. ReLoRA replaces all attention and fully- connected network parameters, while updating the embeddings and normalization layers full-rank. Since ReL.RR- Awrapped models have fewer trainable parameters than full-rank training, we include Control baseline, which is full-rank transformer with the same number of trainable parameters as ReLoRA.
We initialize ReLoRA from a checkpoint of full-rank training 5,000 update steps and reset it every 5,000 steps thereafter, 3 times total we reach 20K steps. After each reset, 99% the optimizer state is pruned based magnitude, and the loss is warmed up for the next 100 iterations. ReLoRA parameters are reinitialized following LoRA best practices, Kaiming initialization (He et al. 2015) for A-matrix, and zeros for B-matrix.
Scaling up to 1.3B After initial results at 130M and 350M model sizes, we applied ReLoRA to train a 1.3B parameter language model. As a baseline, we pre-trained a 1.3B model from scratch on 23B tokens. We performed multiple ReLoRA runs starting from 2K, 5K, and 10K checkpoints. In most of the experiments, we continued using r = 128 and our additional experiments show minimal difference between rank 128 and 512 (hidden size is 2048). Section 4. describes these experiments in detail.
4 RESULTS
Parameter-efficient pre-training Our results are presented in Table and nificantly outperforms LoRA training demonstrating the effectiveness our proposed modifications (ablated in Section (). Additional pre-training loss figures available in Appendix
and Figure [ ReLoRA sig-
Furthermore, ReLoRA achieves similar performance to full-rank training in both upstream and downstream tasks (Table 3
High-rank training through low-rank updates To determine whether ReLoRA performs higher rank update than LoRA, we plot the singular value spectrum of the learned update to the warm-start
Note that the absolute values of GLUE results expected to be quite far from state-of-the-art, because models pre-trained roughly 20 times less data than T5 BERT.
5

---

Published as a conference paper at ICLR 2024
Projections
Projections
V Projections
Down Projections
3.0
ReLoRA
ReLoRA
ReLoRA
ReLoRA
LoRA
2.5
LoRA
Full-rank
2.0
Full-rank
2.0
Full-rank
Full-rank
Frequency
training
training
training
training
1.5
1.5
1.5
1.5
1.0
1.0
1.0
1.0
0.5
0.5
0.5
0.5
0.0
0.0
0.0
0.0
1.0
1.5
0.0
1.0
Singular Value
Singular Value
Singular Value
Singular Value
Figure 3: Singular values spectra of the weight difference between ReLoRA and LoRA at 5,000 iterations (warm start) and 20,000 iterations. ReLoRA exhibits closer resemblance full-rank training than to IooLoR LoRA, indicating its effectiveness in approximating full-rank behavior. 350M models.
1000
750
Full-rank Training
500
ReLoRA
250
0
Wa
Wk
Wy
Wop
Wdown
Figure 4: The number of singular values >0.1 in weight matrices of the learned update. 350M models.
weights. Specifically, the difference between warm-start weights and the final weights for ReLoRA. LoRA, and full-rank trained models. Figure illustrates significant qualitative differences between LoRA and ReLoRA for the singular values of AWq. AWK K. AWv, and AWdown- While most of the singular values for LoRA are zero (Figure 4 with a noticeable number of exceptionally high values above 1.5, ReLoRA exhibits higher distribution mass between 0.1 and 1.0, reminiscent of full-rank training.
Additionally, computed the number of singular values less than 0.1 for LoRA, ReLoRA, and full-rank training. Our results (Figure 4 show that ReLoRA has a much smaller number of near- zero singular values than LoRA, closer to full-rank training. This observation emphasizes the significance of high-rank updates and demonstrates that ReLoRA does accomplish high-rank update by performing multiple low-rank updates. We also perform ReLoRA component ablation (Table and discuss it in S
4.1 SCALING TO 1.3B
Our best this model size starts after 10K step start (33% of the total update steps). We train ReL.RA with rank r = 128, learning rate 5e-4, 100 steps Ir warmup, and 50 steps restarts warmup. The results are presented in the Figure Ig Table 4 ReLoRA clearly outperforms LoRA throughout the training with the gap between the methods increasing from 0,56 at 15K steps to 0.96 30K steps. At the end of the training, ReLoRA is able reach perplexity of 17.24, only 0.44 higher than full-rank training. You can find additional recommendations on ReLoRA hyperparameter selection in Appendix
Varying ReLoRA rank this experiment we wanted to evaluate if 128 is still applicable to the model of this size (hidden size 2048) if it needs to be increased. To do that, used an early checkpoint for the warm start (5K out of 30K steps). This was beneficial for the comparison, as at this point loss changes quickly which makes any differences in training dynamics more evident. We train these models for an additional 10K iterations. Unexpectedly, we found very little difference between ranks 128 (ppl. 19.16) and 512 (ppl. 19.00).

---

Published as a conference paper at ICLR 2024
3.6
3.4
1B ReLoRA (250M trainable)
Loss
18 LoRA (250M trainable)
250M
3.0
2.8
Trainable Params
0
5000
10000
15000
20000
25000
1300
268
5000
10000
15000
20000
25000
Step
Figure 5: Training loss at 1.3B scale and the associated baselines. ReLoRA outperforms LoRA throughout training and the gap between the methods increases over time.
1.3B @ 15K steps
1.3B @20K steps
1.3B @ 30K steps
Full training
17.67 (250M)
17. 17.00 (250M)
16.83 (250M)
Control
22.67 (250M)
22.00 (250M)
21.73 (250M)
LoRA + Warm Start
18.50 (250M)
18.38 (250M)
18.23 (250M)
ReLoRA
17.94 (250M)
17.64 (250M)
17.27 (250M)
Training tokens (billions)
11.8
15.7
23.1
Table 4: Results at 1.3B scale. Number of trainable parameters for each model in (brackets).
Negative results: Online ReLoRA Intu- itively. more frequent ReLoRA resets can lead better performance, as they, principle, can learn a higher rank update. Usually, for ev- ery ReLoRA reset, we would also perform optimizer reset and learning rate scheduler re- warmup (Section T.. However, our experi- ments observed that very high ReLoRA reset rates lead to worse performance.
250M @ 15k steps)
1.3B (@ steps)
ReLoRA Online ReLoRA
27. 76.6 29.31
17.36
17.80
Table 5: Online ReLoRA
Online ReLoRA resolves this issue quite elegantly - it merges LoRA parameters very frequently (e.g., every 100 iterations) while keeping the optimizer reset rate at 2-5K iterations. Unexpectedly, we found that it performs worse than regular ReLoRA at both 250M and 1.3B scales (Table 55).
ReLoRA Training Speedup Training ReLoRA took 440 A100-hours, saving 56 A100-hours compared to full-rank training. part of the speedup was due the ability to use times larger microbatch size. When training with the same microbatch size, ReLoRA improved RAM consumption from 27.8Gb 22.3Gb saving 5,5Gb of GPU RAM. Overall, in the 8xA100 setup. combining the warm start and o training time, 1.3B-ReL.ORA took 86 hours (wall clock) to train compared to 93.5 hours to train 1.3 model full-rank on the same amount of data. This yields relative speed improvement of 9%.
We additionally observed that ReLoRA speedup is significantly hardware-dependdnt (Table In our early experiments on 2xRTX3090, relatively budget setup of 6xA6000 Ada, we estimated 152 hours of wall-clock training time for the IB full-rank model and 119 hours for the ReLoRA model with 33% warm start. This saves 33 hours yielding a relative speedup of 21%. We attribute the difference to the GPU memory speed. ReLoRA more effectively utilize oww-handwidth memory has less trainable parameters.
estimated the speedup of 42%. In more practical, but still
7

---

Published as a conference paper at ICLR 2024
Restarts
Optimizer Reset
Jagged Schedule
Warm Start
Perplexity (L)
X
x
X
34.17
x
X
34.25
X
X
(diverged)
X
34.29
X
29.77
X
X
V
25,46
V
25.04
Regular training
23.65
Table 6: Ablation studies of ReLoRA (130M models). Restarts and warm starts are essential for good performance. Restarts and optimizer resets without a jagged schedule causes the model to diverge.
8xA100
6xA6000 (Ada)
2x3090
Full-rank throughput
137 ex/sec
84 ex/sec
8.8 ex/sec
ReLoRA throughput
157 ex/sec
124 ex/sec
17.8 ex/sec
Immediate speedup
15%
48%
102%
Warm-start adjusted ReLoRA throughput
149 ex/sec
111 ex/sec
14.8 ex/sec
Total speedup
9%
32%
51%
Table Performance metrics in different hardware configurations. Warm start adjustment assumes 33% of full-rank training before switching to ReLoRA.
4.2 ABLATION STUDIES
We conduct ablation studies on all four crucial components of ReLoRA: restarts, jagged schedule, optimizer resets, and warm starts, utilizing the 130M-sized model. The results are presented in T In this section, we will focus on and analyze certain combinations of these components.
LoRA ReLoRA, without the aforementioned components, is essentially equivalent to training a low-rank network parameterized by LoRA. This approach yields remarkably high perplexity, indicating that a simple matrix decomposition has significantly different training dynamics from full-rank training.
Adding restarts and optimizer resets ReLoRA, without jagged schedule and optimizer reset, performs similarly LoRA because old optimizer states force the newly initialized parameters into the same subspace the prior weights, limiting the model's 's capacity. However, doing a naive optimizer reset with ReLoRA causes the model to diverge. training and has positive impact on the mixture. In our initial experiments, we also observed that combination of partial optimizer reset and jagged scheduler allows for quicker warmup, as low as 50 steps, instead of hundreds of steps required when the optimizer is initialized from scratch.
jagged schedule helps to stabilize
Warm start The warm start shows the most significant improvement, dropping perplexity by almost 10 points. To investigate whether post-warmup training contributes to the loss, we measured the perplexity of the warmed-up network, which equals 27.03. It outperforms all low-rank methods except for our final ReLoRA recipe but still demonstrates significant difference from the final network. This demonstrates the importance of early training, similar to concept of the lottery ticket hypothesis with rewinding (Frankle al. 2019). In experiments, unless specified otherwise. we performed warm start for about 1/4 of the total training updates.
RELATED WORK
Scaling versus Efficiency The relationship between overparametrization and neural network trainability and generalization has been extensively studied (Zhang et al. 2017 Belkin al. 2018
8

---

Published as a conference paper at ICLR 2024
Frankle and Carbin 2019 Nakkiran et al.
2019 Singh al. 2021), yet it remains mystery Zhang
et al. ).
Moreover, scaling laws (Kaplan et al. 2020 Ghorbani et al. 202
Hoffmann et al. 2022 demonstrate
a simple and strong power-law dependence between network size and its performance across a variety of modalities. This finding not only supports overparametrization but also encourages the training of extraordinarily resource-intensive neural networks (Brown et al. 2020; Chowwdeer et al.) 2022
Fedus et al., 2022). Nonetheless, the Lottery Ticket Hypothesis (Frankle al., 2019) suggests that overparanetrrizatoon could, in principle, be minimized.
Parameter-efficient fine-tuning Aghajanyan et al. (202I) found that pre-training reduces the amount of change to the network required to learn a new task through fine-tuning. Le.. larger networks or networks pre-trained on more data require smaller modifications in terms of the rank of the range to learn new task. This explains the success of parameter-efficient fine-tuning methods (Lialin et al. 2023) and has also motivated the development of low-rank fine-tuning methods such as LoRA (Hu al. 2022) and Compacter (mahabadi et al. 202Tb.
Low-rank neural network training Training low-rank representations have been explored in the context of CNN compression, regularization, and efficient training (Idelbayev and Curreira-Perpinatt
Jaderberg et al., 2014 Sui et al., (2033] Schotthöfer et al. [022][Lin Lin et al. [2020] Yuan et al. 2021 Zhao al. (023). However, most of these methods are either specitic to CNNs, not scale well, or have not been evaluated on large transformers (Vaswani et al. 2017) with hundreds of millions of parameters, which can benefit greatly from efficient training. While transformers have been shown to have a low-rank internal dimensionality and representations (Aghajanyan et al. [021]] Wang et al. 2020), the study by [hooanapalli et al. (2020) demonstrated that the low rank key and query projections in multi-head attention bottlenecks the performance of transformers. Our experiments (Section (6) also demonstrate that low-rank transformers perform significantly worse compared to the full-rank baseline and ReLoRA
2020:
CONCLUSION
In this paper, demonstrate that parameter-efficient fine-tuning methods can be adapted for pre- training large language models. We first examined the limitations of a low-rank matrix factorization (LoRA) approach and observed that it struggles to effectively train high-performing models. To address this issue, we proposed ReLoRA, which leverages the rank of the sum property to train a high-rank network through multiple low-rank updates. Similar to the lottery ticket hypothesis with rewinding, ReLoRA employs a full-rank training warm start before transitioning to ReLoRA. During training, ReLoRA periodically merges its parameters into the main parameters of the network, performs optimizer reset and learning rate re-warmup.
transformer
We demonstrated that ReL.RR consistently outperforms LoRA for training large transformer models. Our largest experiment demonstrated 9% wall-clock time reduction in 8xA setup and much larger (20 40%) speed improvements on cheaper hardware. Further, our results show similar performance to regular training making ReLoRA promising candidate for improving the efficiency of large model training. Our further studies will focus on improving ReLoRA performance, efficiency, applying it to larger models, and applying it to continued pre-training of existing large language models.
ACKNOWLEDGMENTS
This paper has been a journey and we are sincerely grateful to everyone who supported us. We would like to express our gratitude to Stability,ai, Eleuther.ai, and the Google Cloud for Research Program for providing computational resources essential for this research; Eric Lehman and Artem Krivosheey, for supporting this project from the very beginning.
Special thanks Jason Phang, Hailey Schoelkopf, Enrico Shippole, and Stella Biderman for their technical advice and assistance with computational resources. Our experiments at billion parammeer scale wouldn't be possible without your support.
This work was funded in part by an NSF CAREER award (1S-115))72) and an Amazon Alexa AI research award to Anna Rumshisky.

---

Published as a conference paper at ICLR 2024
REFERENCES
A. Aghajanyan, S. Gupta, and Zeulemoyer. Intrinsic dimensionality explains the effectiveness of language model fine-tuning, In Proceedings of the 59th Anmual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference Naturul Lunguage Processing (Volume 1: Long Pupers), pages 7319-7328 Online, Aug. 2021. Association for Computational Linguistics. doi: 10.18653/v1/2022 Lacl- 568. URLhttps: acl httpa////aaaaathhooogoooo wttp:..gg////0 202g acl-2la.cllooo 568
Z. Allen-Zhie Zhu, Y. Li. and Z. Song. A convergence theory for deep learning via over-voramann parameterization. In K. Chaudhuri and R. Salakhutdinov, editors. Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 242-252. PMLR, 09-15 Jun 2019. URL htt.:httt:///cccc. .tt.eed...........//////77.
zhu 9a h
Arora, N. Cohen, W. Hu, and Y. Luo. Implicit regularization deep matrix factorization, 2019.
M. Belkin, D. J. Hsu, S. Ma, and S. Mandal. Reconciling modern machine-lcamiing practice and the classical dias-variance trade-off. Proceedings of the National Academy of Sciences, 116: 15849 15854, 2018.
Bhojanapalli, Yun, S. Rawat, S. Reddi, and S. Kumar. Low-rank bottleneck multi- attention models. In International Conference on Machine Learning, pages 864 873. PMLR, 2020.
E. Addera, E. Littwin, E. Abbe, S. Bengio, and J. Susskind, Transformers learn through gradual rank increase, 2023.
Borgeaud, A. Mensch, J. Hoffmann, T. Cai, E. Rutherford, K. Millican, G. Van Den Driessche, J.-B. Lespiau, B. Damoc, A. Clark, D. De Las Casas, A. Guy, J. Menick, R. Ring. T. Hennigan, S. Huang. L. Maggiore, C. Jones, A. Cassirer, A. Brock, M. Paganini, G. Irving, O. Vinyals, S. Osindero, K. Simonyan, J. Rac, E. Elsen, and L. Sifre. Improving language models by retrieving from trillions of tokens. In K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Researck, pages 2206-2240, PMLR, 13-23 Jul 2022. URL https::
mlr. ptw.s/.ess.....................e v1 62/ borgeaud22a.. htm.
T. Brown, B. Mann, N. Ryder. M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan. P. Shyam, G. Sastry. A. Askell, S. Agarwal. A. Herbert- Voss. G. Krueger. T. Henighan. R. Child, A. Ramesh, D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray. B. Chess, J. Clark, C. Berner, McCandlish, A. Radford, I. Sutskever, and D. Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing
Systems, volume 33, pages 1877-1901. Curran Associates, Inc., 2020. URL httpss:/prroceddii 8bfbaac Pataaaaceee2224aepe.......
heurips.cc/papen / 2020 /ttp://0000//470006ffcco
A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung, C. Sutton, S. Gehrmann, P. Schuh, K. Shi, S. Tsvyashchenko, J. Maynez, A. Rao, P. Barnes, Y. Tay, N. M. Shazeer, V. Prabhakaran, F. Reif, N. Du, B. C. Hutchinson, R. Pope, J. Bradbury, J. Austin, M. Isard, G. Gur-Ari, P. Yin, T. Duke, A. Levskaya, Ghemawat, S. Dev, H. Michalewski, X. Gareía, V. Misra, K. Robinson, Luan, H. Lim, B. Zoph, A. Spiridonov. R. Sepassi, D. Dohan, S. Agrawal, M. Omernick, A. M. Dai, T. S. Pillai, M. Pellat, A. Lewkowycz, E. Moreira, R. Child, O. Polozov, K. Lee, Z. Zhou, X. Wang. B. Saeta, M. Díaz, O. Firat, M. Catasta, J. Wei, K. S. Heier-Hellssernn D. Eck, J. Dean, S. Petrov, and N. Fiedel. Palm: Scaling language modeling with pathways. ArXiv, abs/2204,022 12, 2022.
L. Fedus, Zhou, Ippolito.
Dao, D. Y. Fu, Ermon, A. Rudra, and C. Re. Flashattention: Fast and memory-efficient exact attention with 10 1O-awareness. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Newral Informaiin Processing Systems, 2022.
T. Dettmers, M. Lewis, Y. Belkada, and L. Zettlemoyer. (PPT3.in8(( 8-bit matrix multiplication for transformers scale. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems, 2022. URL. http:://ppp...aaaaaccccceoeeeorp
wtt.:/ewwweee..tteeeeeeee owt.oormmmmmmidiiiiiiiaaaaaaaaa CWqHoxaD
T. Dettmers, A. Pagnoni, A. Holtzman, and L. Zettlemyyrr Qlora: Efficient finctuning of quantized llms. ArXiv, abs/2305.1433 14314, 2023. httphhttt//aaa www.::essaaaammmimaaeeaae http:echh.........orrp..s....eee
A. Edalati, S. Tahaci, Kobyzev, V. Nia. J. Clark, and M. Rozagholizadeh. Krona: Parameter efficient tuning with kronecker adapter. ArXiv, abs/2212.10650, 2022. http:ttppe//ppp..eeattttiiissees
o org/ ID: 254926823
W. Fedus, B. Zoph, and N. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. J. Mach. Learn. Res., 23(1), jan 2022. ISSN 1532-4435.
10

---

Published as a conference paper at ICLR 2024
J. Frankle and M. Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In Interna- tional Conference on Learning Representations, 2019, URLihtpps//oooeneeeeev....s wtwp://rrmwieo.notooooomummmmm
[J1 -33333CF7
J. Frankle, G. Karolina Dziugaite, D. M. Roy, and M. Carbin. Stabilizing the lottery ticket hypothesis. arXiv e-prints, pages arXiv- 1903, 2019.
B. Ghorbani, O. Firat, M. Freitag, A. Bapna, M. Krikun, X. García, C. Chelba, and C. Cherry. Scaling laws for neural machine translation. ArXiv, abs/2109.07740, 2021.
performm K. He, X. Zhang. S. Ren, and Sun. Delving deep into rectifiers: Surpassing human-level performance imagenet classification. CoRR, abs/1502.01852, 2015. URU/pttp://arxL..or abs/1502...... 01852
K. Hc, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. 2016 JEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 770-778, 2016.
J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. de las Casas, L. A. Hendricks, J. Welbl, A. Clark, T. Hennigan, E. Noland, K. Millican, G. van den Driessche, B. Damoc, A. Guy, S. Osindero, K. Simonyan, E. Elsen, O. Vinyals, W. Rae, and L. Sifre. An empirical analysis of compute-optimal large language model training. A. H. Oh, Agarwal, Belgrave, K. Cho, editors, Advances in Neural Information Processing Systems, 2022. URL
LBBeRULOAPR
E. J. Hu, yelong shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen. LoRA: Low-rank adaptation of large language models. In International Conference Learning Representations, 2022. URL https://cppeneevie..oo hotp:///ooooommoodeeeeeooooo
forum?
Idelbayev and A. Carreira-Perpinar Low Lank compression neural nets: Learning the rank of each layer. In 2020 JEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). pages 9046-8056, 2020. doi: 10. ALL11//CCPR2220022000000000
A. Jacot, F. Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Proceedings of the 32nd International Conference Neural Information Processing Systems, NIPS' 18, page 8580 8589, Red Hook, NY, USA, 2018. Curran Associates Inc.
M. Jaderberg. Vedaldi, and A. Zisserman. Speeding up convolutional neural networks with low rank expansions. In Proceedings of the British Machine Vision Conference. BMVA Press, 2014. doi: http:
S. R. Kamalakara, A. F. Locatelli, B. Venkitesh, J. Ba, Y. Gal, and A. N. Gomez. Exploring low rank training of deep neural networks. ArXiv, abs/2209. 13569, 2022. URL http://aap..smmannccccccc.....ee org/ DOrpusID:2525455:0
Kaplan, S. McCandlish, Henighan, T. Brown, Chess, Child, S. Gray, Radford, Wu, and D. Amodei. Scaling laws for neural language models. ArXiv, abs/2001 108661, 2020.
U. Khandelwal, O. Levy, D. Jurafsky, L. Zettlemoyer, and M. Lewis. Generalization through memorization: Nearest neighbor language models. In International Conference on Learning Representations, 2020. URL httts:://cpenrewe
net forum? wtw.:mmmddddiki-eEKe---
D. P. Kingma and Ba. Adam: A method for stochastic optimization. CoRR, abs/1412.6990) 2015.
A. Krizhevsky, I. Sutskever, and G. E. Hinton. Imagenet classification with deep convolutional neural networks. In F. Percira, C. Burges, L. Bottou, and K. Weinberger, editors, Advances Neural Information Processing Systems, volume 25. Curran Associates, Inc., 2012. URL ht htttttpp://proccedinge neurips. cc/ h
Paper pddf
V. Lialin, V. Deshpande, and A. Rumshisky. Scaling down to scale up: A guide to parameter-efficien fine-tuning, 2023.
R. Lin, C.-Y. Ko, Z. He, Chen, Y. Cheng, H. Yu, Chesi, and N. Wong. Hotcake: Higher order tucker articulated kernels for deeper enn compression. In 2020 IEEE 15th International Conference on Solid-Star & Integrated Circuit Technology (ICSICT), pages 1-4, 2020. doi: 10. 10.1 00.100/CCSICC14987220022222
R. K. mahabadi, J. Henderson, and S. Ruder. Compacter: Efficient low-rank hypercomplex adapter layers. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advonces in Neural Information Processing Systems, 2021. IRR./hhtss:///peeeet
net. wttpre/wwwwweeeeeeeeeeeeee-eeeee http://pww...eeee.ee//omooooooo/

---

Published as a conference paper at ICLR 2024
P. Micikevicius, S. Narang. J. Alben, G. Diamos, Elsen, D. Garcia, B. Ginsburg. M. Houston, O. Kuchaiev, G. Venkatesh, and H. Wu. Mixed precision training. International Conference on Learning Representations, 2018. WRR.bttps://ppezeee
net http://wwwwoooooeoooooooooaaaaa/
Nakkiran, G. Kaplun, Y. Bansal, Yang. Barak, Sutskever. Deep double descent: where bigger models and more data hurt. Journal of Statistical Mechanicx: Theory and Experiment. 2021, 2019.
A. Radford, K. Narasimhan, T. Salimans, I. Sutskever, et al. Improving language understanding by generative pre-training. 2018.
C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Revearch, 21 (140):1-67, 2020. http:t:: http://www..oggppppp////////////
org/ httpp/ppppsss//pp//ppppp
html
S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-16, 2020. doi: 10.1109/SC41400
405,20SCC440..22200000222 202.02002..002....
Schotthöfer, E. Zangrando, Kusch, G. Coruti, and Tudisco. Low-rank lottery tickets: finding efficient low- rank neural networks via matrix differential equations. Koyejo, Mohamed, A. Agarwal, Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 20051- 20063. Curran Associates, Inc., 2022. URL
wtwpppprrppppppppppppppp..ep Files/
N. Shazeer. Glu variants improve transformer, 2020.
K. Simonyan and A. Zisserman. Very deep convolutional networks for large-scale image recognition. In Y. Bengio and Y. LeCun, editors, 3rd International Conference Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL/btt::/aazx.v..e. abs/
1409. 1556
Singh, Bachmann, and Hofmann. Analytic insights into structure and rank of neural network hessian maps. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL
net wwt/foooumee http//oooommmmooogorrrro LM77N
J. Su, Y. Lu, S. Pan, B. Wen, and Y. Liu. Roformer: Enhanced transformer with rotary position embedding. ArXiv, abs/2104.09864,, 2021.
Y. Sui, M. Yin, W. Yang. Y. Gong, J. Xiao, H. Phan, D. Ding, X. Xu, S. Liu, Z. Chen, and B. Yuan. ELRT: Towards efficient low-rank training for compact neural networks, 2023. /RLL/ttpp:///openeewics C3926 9m8b88)
net fwrumorrmmiii
Y.-.. Sung. J. Cho, and M. Bansal. Lst: Ladder side-tuning for parameter and memory efficient transfer learning. ArXiv, abs/2206.06522, 2022. / 249642544]
H. Touvron, Lavril, G. Izacard, X. Martinet, M.-., Lachaux, Lacroix, Rozière, Goyal, Hambro, F. Azhar, A. Rodriguez, A. Joulin, E. Grave, and G. Lample. Llama: Open and efficient foundation language models. arXiv preprint 233.2:2300..1977,.22
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, t. Kaiser, and I, Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
S. Wang, B. Z. Li, M. Khabsa, H. Fang, and H. Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.0 04768, 2020.
X. Yuan, P. H. P. Savarese, and M. Maire. Growing efficient deep networks by structured sparsification. In International Conference Learning Representations, 2021. URL httppss///ppeeeeeee.../ net/ Forum? wwtp:dddwwwwwwwweee
B. Zhang and Sennrich. Root mean square layer normalization. In H. Wallach, H. Larochelle, Beygelzimer, F. d'Alché-Bin Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL
neur... cc/ paper
Ei /ttpa//aaae//1//////////22222222
C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals. Understanding deep learning requires rethink- ing generalization. In International Conference on Learning Representations, 2017. URL. httpa 7/apen net httpeeeeee..ooomm//m///aaooooaoa
12

---

Published as a conference paper at ICLR 2024
C. Zhang, Bengio. M. Hardt, B. Recht, and O. Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64:107 115, 2021.
Zhao, Y. Zhang, Chen, Schäfer, and Anandkumar. Inrank: Incremental low-rank learning. arXiv preprint arx:v22366. /250, 2023.
A A PRACTICAL GUIDE TO RELORA
In this section, we wanted to answer most common questions on hyperparameter selection. Especially how select ReLoRA-speciff hyperparameters to reliably get better performance than LoRA. In all of our experiments, we applied LoRA/ReLoRA to all of the linear layers in the model: kqy-projection layers, FFN layers and other projections, except for logits and embeddings.
We observed that E {64, }28} works well for all of the networks, up to IB. One small, but important hyperparameter change from full-rank training ReLoRA-training that was crucial for the performance was increased learning rate. ReLoRA (and LoRA) requires 1.5 2 times larger learning rate than regular training/fine-tuning to achieve similar performance.
When taking about ReLoRA-see A-seciific hyperparameters, we did not observe significant dependence on optimizer pruning percentage as long as it's larger than 90%. Larger pruning rates can lead to slightly better performance at the cost of possible loss instabilities during the reset. We tested several ReLoRA reset rates with 350M and 1.3B models and found that 2K iterations reset rate performed consistently well in both pre-training and fine-tuning experiments and always led to better performance than no resets. In general, we observed good results with reset rates 2K-5K.
B RELORA FOR FINE-TUNING
We apply ReLoRA to fine-tune T5-base (220M parameters) and TS-large (770M parameters) on the GLUE benchmark. We use the same type of learning rate scheduler as in ReL.oRA pre-training and prune 90% of the low magnitude optimizer states during each LoRA merge-and-reimt (restart). The batch size is equal to 128 examples and the learning rate tuned (from lc-4 to Se-4) on each model and dataset combination. We perform additional ReLoRA ablation studies using the T5-Large model and QNLI dataset. Specifically, we explore different ReLoRA ranks, optimizer state pruning rates, and the total number of ReLoRA resets.
Method
SST-2
MILI
QNLI
QQP
RTE
STS-B
MRPC
CoLA
Avg
Adapters"
94.2
86.4
93.1
88.9
75.1
91.1
88.9
64.4
85.3
Prompt Tuning"
90.3
82.5
92.5
88.5
59.5
74.6
0.0
72.2
Ladder Side Tuning"
94.1
85.6
93.3
88.8
71.9
90.7
90.4
58.1
84.1
Compacter*
93.9
86.1
92.9
90.4
76.3
91.0
91.5
64.4
85.8
KronA'
86.3
93.2
90.6
63.3
86.1
Full fine-tuning
93.6
86.2
92.8
91.7
74.8
90.1
92.7
63.4
85.7
LoRA
93.92
86.12
91.95
90.62
78.34
89.96
90.52
60.04
85.18
ReLoRA
94.15
85,96
91.68
87.2
77.74
89.88
90.03
59.92
84.57
Full fine-tuning (TS-L)
94.7
91.6
89.9
78.9
90.6
57.0
85.0
LoRA (T5-L)
95.59
89,
93.98
91.44
85,92
90.89
92.90
63.77
87.99
ReLoRA (T5-L)
95.7
89.06
93.68
91.04
84.72
90.53
90.57
61.72
87.47
Table ReLoRA fine-tuning does not outperform LoRA. Results with and are T5-base results from Sung et al. 20022 and Edalati et al. (2022) respectively.
ReLoRA fine-tuning ablations Table shows the results of varying ReLoRA hyperparameters. A rank of 64 seems to provide the best performance. The results indicate that the model's performance remains largely unaffected
when 99% of the optimizer states reset. Our analysis of the jagged cosine learning rate scheduler's 's impact on classification accuracy in the QNLI dataset suggests that resets adequate (reset rate 4000).
13

---

Published as a conference paper at ICLR 2024
Rank
Acc.
Pruning
Acc.
Reset rate
Fresets
Acc.
16
94.05
85%
94.51
6000
94.38
94.16
92%
94.33
4000
94.73
64
94.55
95%
94.31
2000
5
94.34
128
94.44
99%
94.56
1000
11
94.33
Table 9: ReLoRA fine-tuning ablations. We apply ReLoRA fine-tune T5-large on the QNLI dataset and vary LoRA rank (r), optimizer state pruning percentage, and reset frequency of ReLoRA. Reset rate means the number of iterations between ReLoRA resets.
C LEARNING CURVES OF MODELS PRE-TRAINED IN THE STUDY
In this section we present additional training loss plots for all of the models from Table 2 60M: Figure 130M: Figure 250M: Figure S 350M: Figure 1.3B: Figure I0
D RANKS OF 130M MODELS
Figures
and 12 show spectral properties for 130M model.
SMALLER WARM START PERIOD
demonstrates that ReLoRA consistently outperforms the warmed-started LoRA baseline. To provide more contrasting example, we performed additional pre-training experiments starting from just 2K warm-started network. Figure 13 shows a significant performance gain with ReLoRA over LoRA by 1.4 ppl points (ppl 23.64 vs 25.08). While the absolute performance of ReLoRA is lower compared to full-rank training this context, these experiments validate our initial hypothesis that LoRA restarts positively impact performance.
Table 2
4.4
60M
60M ReLoRA
4.2
(40M trainable)
Loss
4.0
60M Warm+ttrtttart
(40M trainable)
3.8
40M
3.6
Trainable Params
2000
4000
6000
8000
10000
40
2000
4000
6000
8000
10000
Step
Figure 6: 60M experiments training loss
14

---

Published as a conference paper at ICLR 2024
4.50
130M
4.25
130M ReLoRA
4.00
(72M trainable)
130M LoRA+WarmStar
3.75
trainable)
72M
3.25
Trainable Params
5000
20000
60
2500
5000
7500
10000
17500
20000
Step
Figure 7: 130M experiments training loss
250M
4.5
250M ReLoRA
(99M trainable)
Loss
99M
3.5
Trainable Params
0
2500
5000
7500
10000
12500
15000
17500
20000
250
100
2500
5000
7500
10000
12500
15000
17500
20000
Step
Figure 8: 250M experiments training loss
4.0
350M
3.8
350M ReLoRA
(130M trainable)
Loss
350M LarA+Wartttrrt
3.4
(130M trainable)
130M
3.2
3.0
Trainable Params
7500
10000
12500
20000
350
2500
5000
7500
10000
12500
15000
17500
20000
Step
Figure 9: 350M experiments training loss
15

---

Published as a conference paper at ICLR 2024
3.4
ReLoRA
(250M trainable)
Loss
3.2
1B LoRA (250M trainable)
250M
Trainable Params
15000
1300
258
15000
20000
Step
Figure 10: 1.3B experiments training loss
Projections
K Projections
V Projections
Down Projections
3.0
3.0
3.0
3.0
ReLoRA
ReLoRA
ReLoRA
ReLoRA
2.5
LoRA
2.5
LoRA
2.5
LoRA
2.5
LoRA
Frequency
2.0
training
Full-rank
2.0
training
Full-rank
2.0
training
Full-rank
2.0
training
Full-rank
1.5
1.5
1.5
0.0
0.0
0.5
1.5 Singular Value
1.0
2.0
0.0
0.5
1.5 Singular Value
1.0
2.0
0.0
0.5
1.5 Singular Value
1.0
2.0
0.0
0.5
1.0
1.5
2.0
Singular Value
Figure 11: Singular values spectra of the weight difference between ReLoRA and LoRA at 5,000 iterations (warm start) and 20,000 iterations. ReLoRA exhibits a closer resemblance to full-rank training than LoRA, indicating its effectiveness in approximating full-rank behavior. 130M models.
600
Full-rank Training
400
ReLoRA
200
LoRA
Wo
Wx
Wv
Wap
Waten
Figure 12: The number of singular values < 0.1 in attention and FCN matrices of the learned update. 130M models.
16

---

Published as a conference paper at ICLR 2024
Eval loss
ReLoRA 350M warm start
LoRA 350M 2K warm start
3.4
3,2
updato_stee sstpp
5k
10k
15k
20k
Figure 13: ReLoRA significantly outperforms LoRA when started from an early (2K steps) check- point.
17


Reviews:
[{"review_id": "NMnj89VXHe", "reviewer": "Reviewer_vHWf", "summary": "This work proposes ReLoRA, a parameter-efficient method for pretraining neural networks that iteratively trains low-rank adapters, merges them, and trains new sets of adapters for the same parameters. The proposed method is validated on language model pretraining on C4 for sizes ranging from 60M to 1.3B parameters and finetuning T5-base and T5-large; using ReLoRA results in significant memory savings with similar quality of the final model compared to full-rank training.\n\n--- \nPost-rebuttal update: dear authors, thank you for the response! I am happy to see a more detailed study of the speedup, and I encourage you to include it into the revised version. However, since no revision of the paper was published during the discussion period, I am inclined to keep my score of a weak accept: I am in favor of accepting the paper as is, even though I believe it could benefit from additional clarifications and improvements in presentation (as noted by me and other reviewers).", "questions": "* While [1] is not strictly relevant to the submission's topic, I think it studies a similar set of questions, and I would be curious to hear authors' opinion of that paper.\n\n[1] Exploring Low Rank Training of Deep Neural Networks. Siddhartha Rao Kamalakara, Acyr Locatelli, Bharat Venkitesh, Jimmy Ba, Yarin Gal, Aidan N. Gomez. 2022", "limitations": "", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "* The proposed approach is simple to implement and use, yet quite efficient in practice\n* Overall, the paper is clearly written and easy to understand\n* I particularly appreciated the authors including negative results in their submission: the community could significantly benefit from a more widespread use of that practice, yet the majority of papers does not report failed experiments.", "weaknesses": "* It is not fully clear to me if the proposed technique will indeed scale to model sizes which are currently considered to be most capable (starting from 7B parameters and reaching up to 70B in most cases). While I understand that even current results are quite promising, it might be the case that at larger scales, low-rank training would more severely affect the capabilities of the model, which would make the work less impactful for the community.\n* Similarly to the above question, I think that a more principled way to conduct experiments would be to train ReLoRA until it achieves the same perplexity as the full-rank baseline. If we aim to achieve the same quality as the standard approach, it is good to know how many additional iterations with ReLoRA would be necessary for that. I would also expect that the wall-clock time gains become less pronounced after such a comparison.\n* The work could benefit from a bit more polish and proofreading: there are multiple typos and incomplete sentences throughout the paper. For example, see \"hparam\" -> \"hyperparameter\" and \"Flash attention\" -> \"FlashAttention\" in page 4, \"resented\" -> \"presented\" in page 5 and \"ReLoRA clearly outperforms LoRA at At\" in page 8.", "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": "JCfmJqz5Te", "reviewer": "Reviewer_x3dX", "summary": "This paper proposes a parameter-efficient training method called ReLoRA. Based on the basic solution of LoRA, ReLoRA cyclically merges the trainable parameters W_A and W_B into the original parameters and reinitializes the parameters, thus training high-rank networks. The authors show the method's potential in parameter-efficient pre-training and the immediate replacement of LoRA. Experimental results show memory and computation reduction on various model scale setups.", "questions": "- What is the relationship of parameters’ rank and task performance? Is it possible to quantify the respective RANK needed for pre-training and fine-tuning?\n- Can equation 3 prove that \\delta W achieves a higher rank compared to LoRA? Is it possible to add the relationship between higher rank and random initialization?\n- From Figure 1 and Table 2, is it possible to conclude that the pattern of low-rank updates is mainly in the middle and late stages of model training?\n- How can ReLoRA speed up the training, as the forward and backward computation is not sparsified (or even more computation)?\n- How is the performance and further speedup when you combine the ReLoRA and low-precision quantization, as mentioned in the paper?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper points out the low-rank limitations of LoRA on parameter updating and proposes a merge-initialization approach, called ReLoRA, to make some improvements to LoRA. The experimental results show that this improvement can be applied in the model pre-training process to reduce memory and computational resources overhead. In the fine-tuning stage, it can also be used as an alternative to LoRA to achieve better performance on some downstream tasks.", "weaknesses": "The main weakness is whether to keep high-rank training or use LoRA. As far as I understand, methods such as LoRA for parameter-efficient fine-tuning are proposed based on the low-rank nature of the fine-tuning process, so it seems that there is no need for high-rank training during the fine-tuning phase (on the other hand, merging parameters may contaminate the original pretrained checkpoints); While the authors claim the potential of ReLoRA in the training process, it also emphasizes that the warm-start of the full-parameter training has a significant impact on the training effect, so is it possible that such a technique with a low-rank adaptor is unsuitable or unnecessary for pre-training? Clarification of the following questions may help.", "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": "Fet3kp7EFl", "reviewer": "Reviewer_DYvD", "summary": "The paper introduces ReLoRA (Recurrent Low Rank Adaptation) for training large language models that optimizes the training process to be more computationally efficient while maintaining performance. The main idea is utilizing multiple low-rank updates to effectively train high-rank networks, leveraging the 'rank of sum' property. The design incorporates multiple components like full-rank training warm start, periodic parameter merging, jagged learning rates, partial optimizer restarts. The authors show that ReLoRA can replace LoRA for fine tuning performance.", "questions": "Discussion of Limitations:\n- The paper could benefit from a more thorough discussion of the limitations of ReLoRA. Understanding the scenarios in which ReLoRA may not perform as expected or could be improved is crucial for future research and practical applications.\n\nExplanation of Design Choices:\n- Some design choices in ReLoRA, such as the specific strategy for periodic parameter merging and the choice of low-rank updates, could be explained in more detail. Providing the rationale behind these choices and discussing potential alternatives would strengthen the paper.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Originality\n- The paper introduces a unique combination of low-rank updates, full-rank training warm start, and periodic parameter merging. This approach not only reduces the computational resources required for training but also ensures the model’s effectiveness across various tasks. \n\nQuality\n- The authors perform decent amount of experimentation and thorough analysis. The authors have provided a comprehensive evaluation of the method on large transformer models (1.3B params), ensuring the validation of their approach. \n\nClarity\n- The paper is well-structured and clearly written. I did not have a problem following the proposal.", "weaknesses": "Performance Similarity to LoRA:\n- Based on the empirical results, ReLoRA, exhibits performance (perplexities) that is very similar to that of LoRA. This raises questions about the practical necessity of the more complex ReLoRA framework (multiple new hyperparameters and intricate design choices) when simpler alternatives provide comparable results. Especially because the training speedups are marginal ~9%. \n\n- The authors could strengthen their contribution by providing clearer and more substantial evidence of scenarios where ReLoRA significantly outperforms LoRA, helping to justify the additional complexity.", "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": "YB7MDWQco6", "reviewer": "Reviewer_xXqo", "summary": "The authors propose ReLoRA, a parameter-efficient method that can be applied to both fine-tuning and pre-training. ReLoRA utilizes low-rank updates that aggregate to train high-rank networks. Experiments show that ReLoRA can speed up the training process and kind of reduce the memory consumption.", "questions": "1. In Figure 1, what's Loss for, training of validation? Why can Loss demonstrate \"similar performance\"?\n2. In Table 3, why isn't full fine-tuning included?\n3. For the Control Baseline, how are the trainable parameters chosen?\n4. ReLoRA outperforms LoRA for fine-tuning can be due to the Warm start?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "!. The orientation is good. With several existing PEFT methods, it's important to implement parameter-efficient methods to improve the pre-training process. \n2. Extensive experiments are done. Several models are used and various ablation studies are included.\n3. The method is relatively simple and effective, with some ingeniously designed tricks.", "weaknesses": "1. The writing is not good. Figure 1 lacks notes on Loss. The method part is quite confusing. Warm start isn't even mentioned except in Algorithm 1. In 3.1 \"Architecture and training hyperparameters\", it seems that a hand-made architecture is designed and used, while in \"Scaling up to 1.3B\", BERT is mentioned, so what exactly is the architecture? In \"trained on 8×A100GPUs (or more)\", this \"or more\" seems the experiment details are not so clear.\n2. The method includes warm start at first, which means the memory consumption of ReLoRA is equal or close to full training at first. One has to afford the huge memory consumption at first, so the memory problem has not been solved. As shown in the results, full training outperforms ReLoRA on almost all tasks (or average), so as far as I'm concerned, ReLoRA is just sacrificing the overall performance for speeding up. Also, it seems unfair if comparing with LoRA for pre-training, because LoRA isn't designed for pre-training. However, in Table 2, results of LoRA and ReLoRA are quite close, so I doubt the efficiency of ReLoRA. Directly using LoRA+warmstart can be a simpler choice.\n3. It seems ReLoRA introduces a lot of new hyperparameters, making the param tuning more complex than regular training or LoRA.", "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)