ReLoRA: High-Rank Training Through Low-Rank Updates

Thank you for acknowledging the contributions of our study, and for recognizing that we propose a simple and effective technique for a much-needed application of parameter-efficient methods to pre-training. We also appreciate you highlighting our extensive experiments and ablation studies performed to verify and evaluate our method. Overall, we believe the positive aspects highlighted in your review outweigh the negatives, and we respectfully ask that you reconsider your assessment. We hope that the clarity of writing, which can be easily improved in the final camera-ready version, can be seen as a minor issue. We address your remaining concerns and provide some clarifications in the comments below. ### Architecture, and experimental setup As we mention in Section 3.1, the model architecture we use for all experiments is the LLaMA architecture. Thank you for pointing out the confusing part in the section “Scaling up to 1.3B”. We will make sure to clarify that we are only comparing our models to BERT-base and BERT-Large in size, not in the architecture. “8xA100GPUs (or more)” – in some of the experiments with the 1.3B model, we had to use more than one A100 node to speed up the experimentation. However, all efficiency comparisons (throughput and memory consumption estimates) were made in the 8xA100 setup. We will make sure to clarify this. ### Warm start While we’re not gaining any improvements for the initial model warm start, it is short (~25%), so we are saving significantly on compute in total. Please refer to this table for the concrete speedups for the 1.3B model: | | 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 (33% warm-start) | 149 ex/sec | 111 ex/sec | 14.8 ex/sec | | Total speedup | 9% | 32% | 51% | Also, warmup may not be necessary in some applications. For example, in continued pre-training, where an existing model (e.g., LLaMA-7B) is trained on additional large-scale data or tuned for a specific narrow domain (e.g. medical), a warmup is not required, since the model is already pre-trained. With ReLoRA, this leads to immediate speedups of 15% for A100 and 48% for A6000. ### Is ReLoRA sacrificing the overall performance for speed up? We believe that achieving efficiency gains often requires small compromises in performance. This is evident in techniques such as distillation [6], quantization [7], approximate attention [8,9], and PEFT [10,11]. Notably, ReLoRA consistently surpasses LoRA in all tested scenarios. The performance gap between ReLoRA and LoRA widens with increasing model size, ranging from just 0.27 perplexity points in a 60M model to a significant 0.96 in a 1.3B model (more examples in the table below). We think this result by itself is impactful for the ML community where LoRA is used for compute- and data-heavy tasks like RHLF [12]. | | 60M | 130M | 250M | 350M | 1.3B | |--------------------|--------|--------|--------|--------|--------| | LoRA + Warm Start | 34.73 | 25.46 | 22.86 | 19.73 | 18.23 | | ReLoRA | 34.46 | 25.04 | 22.48 | 19.32 | 17.27 | | Δ | 0.27 | 0.42 | 0.38 | 0.41 | 0.96 | ### Additional hyperparameters We address the issue of additional hyperparameters in Appendix A. Specifically, we find that ReLoRA is not very sensitive to rank r and the optimizer pruning percentage. Default values (128 and 0.9, respectively) work great in all scenarios. 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. While additional hyperparameters can be annoying, it’s not uncommon for deep learning methods to introduce new ones. E.g, in case of Adam, beta1 and beta2 do not exist in regular SGD, yet Adam consistently outperforms SGD with default beta1 and beta2.

### Question 1: In Figure 1, what's Loss for, training of validation? Why can Loss demonstrate "similar performance"? We used training loss in this figure for clearer visualization. Loss is a standard measure for language model performance, since perplexity is just exp(loss). It's commonly used to evaluate language models (see references [2,3,4,5], for example). We would like to emphasize that our pre-training experiments were all single-epoch, so the model never sees the same example twice. Therefore, showing training loss is effectively the same as validation loss. All language modeling result tables show the validation set perplexity. We also provide downstream evaluation in Table 8. ### Question 2: In Table 3, why isn't full fine-tuning included? We will include full fine-tuning results in Table 3 in the camera-ready. Here’s an updated table. | | | | | | | | | | |---------------|------|------|------|------|------|------|------|------| | Full fine-tuning* | 93.6 | 86.2 | 92.8 | 91.7 | 74.8 | 90.1 | 92.7 | 63.4 | 85.7 | | LoRA | 94.0 | 86.3 | 92.6 | 90.6 | 72.2 | 89.8 | 90.2 | 60.3 | 84.5 | | ReLoRA | 94.8 | 86.3 | 93.0 | 90.8 | 80.1 | 89.9 | 91.4 | 61.1 | 85.9 | | | | | | | | | | | | | Full fine-tuning (T5-large) | 94.7 | 89.1 | 91.6 | 89.9 | 78.9 | 90.6 | 88.9 | 57.0 | 85.0 | | LoRA (T5-large) | 95.9 | 89.6 | 94.4 | 91.4 | 85.6 | 90.5 | 90.7 | 60.3 | 87.3 | | ReLoRA (T5-large) | 96.1 | 89.5 | 94.6 | 91.6 | 83.4 | 90.8 | 91.4 | 62.1 | 87.4 | ### Question 3: For the Control Baseline, how are the trainable parameters chosen? The control baseline refers to training a full-rank model with a total parameter count equal to the trainable parameter count used in low-rank training. We determine the number of layers and hidden size of the control model based on this total number of parameters. Here are our configuration choices: | Model Size | Hidden Size | Intermediate Size | Num Attention Heads | Num Hidden Layers | |----------------|---------------|-------------------|---------------------|-------------------| | 1.3B | 2048 | 5461 | 32 | 24 | | 350M | 1024 | 2736 | 16 | 24 | | 250M | 768 | 2560 | 16 | 24 | | 99M | 640 | 1708 | 10 | 12 | | 130M | 768 | 2048 | 12 | 12 | | 70M | 512 | 1368 | 8 | 12 | ### Question 4: ReLoRA outperforms LoRA for fine-tuning can be due to the Warm start? There is no warm start in our experiments that use ReLoRA for fine-tuning. Thank you for pointing out this detail, we will clarify this in the camera-ready. Thank you again for your thorough review. We hope that our comments address your concerns and we are confident that the minor issues can be easily addressed in the final version. Given the above, we sincerely hope you are able to raise your score. ### References [1] LLaMA: Open and Efficient Foundation Language Models, Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, Guillaume Lample. 2023 [2] A Neural Probabilistic Language Model. Yoshua Bengio, Réjean Ducharme, Pascal Vincent, Christian Jauvin. 2003 [3] Recurrent neural network based language model.Tomas Mikolov, Martin Karafia, Lukas Burget, Jan “Honza” Cernocky, Sanjeev Khudanpur. 2010 [4] Exploring the Limits of Language Modeling, Rafal Jozefowicz, Oriol Vinyals, Mike Schuster, Noam Shazeer, Yonghui Wu. 2016 [5] Language Models are Unsupervised Multitask Learners, Alec Radford, Jeff Wu, Rewon Child, D. Luan, Dario Amodei, Ilya Sutskever. 2019 [6] DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. Victor Sanh, Lysandre Debut, Julien Chaumond, Thomas Wolf. 2019 [7] BitNet: Scaling 1-bit Transformers for Large Language Models. Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang, Huaijie Wang, Lingxiao Ma, Fan Yang, Ruiping Wang, Yi Wu, Furu Wei. 2023 [8] Reformer: The Efficient Transformer. Nikita Kitaev, Łukasz Kaiser, Anselm Levskaya. 2020 [9] Efficient Transformers: A Survey. Yi Tay, Mostafa Dehghani, Dara Bahri, Donald Metzler. 2020 [10] BitFit: Simple Parameter-efficient Fine-tuning for Transformer-based Masked Language-models. Elad Ben Zaken, Shauli Ravfogel, Yoav Goldberg. 2021 [11] Scaling Down to Scale Up: A Guide to Parameter-Efficient Fine-Tuning. Vladislav Lialin, Vijeta Deshpande, Anna Rumshisky. 2023 [12] QLoRA: Efficient Finetuning of Quantized LLMs. Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, Luke Zettlemoyer. 2023

Thank you for your review, for acknowledging the novelty and importance of our work. We would like to address your concerns. ### ReLoRA vs LoRA performance ReLoRA consistently outperformed LoRA + Warm Start in all scenarios, with the difference between ReLoRA and LoRA increasing with model size from just 0.27 perplexity points (60M) to 0.96 (1.3B), which is a very large perplexity difference. | | 60M | 130M | 250M | 350M | 1.3B | |--------------------|--------|--------|--------|--------|--------| | LoRA + Warm Start | 34.73 | 25.46 | 22.86 | 19.73 | 18.23 | | ReLoRA | 34.46 | 25.04 | 22.48 | 19.32 | 17.27 | | $Δ$ | 0.27 | 0.42 | 0.38 | 0.41 | 0.96 | ### Speedups The ~9% training speedup, while seemingly marginal, is significant in the context of large-scale language model training, where even small percentages can translate to substantial resource savings. Additionally, the speedup depends a lot on the size of the network and the hardware setup. The table below describes speedup for training a 1.3B model in different hardware setups. As you can see, in 6xA6000 Ada, we see a 48% immediate speedup (not accounting for warm start) and 32% speedup when the warm start is accounted for. | | 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 (33% warm-start) | 149 ex/sec | 111 ex/sec | 14.8 ex/sec | | Total speedup | 9% | 32% | 51% | We would like to thank you for your review and ask to raise your score given the clarifications above.

We appreciate your feedback and recommendations! Here we would like to respond to your points. > “Is it possible that such a technique with a low-rank adaptor is unsuitable or unnecessary for pre-training” We believe there might have been some confusion regarding the primary focus of our paper. ReLoRA is conceived primarily as a method for efficient **pre-training**, rather than fine-tuning. We show that a sequence of low-rank updates can be used to pre-train models from an early checkpoint, saving substantial compute. This is the first study to apply parameter-efficient methods to pre-training, which offers considerable savings in compute without significantly compromising the quality of the model. Fine-tuning results are complementary to the main point of the paper. We apologize if this was not sufficiently clear in our original submission, and we hope that this clarification would allow you to reassess the paper. Below, we address specific questions raised in your review ### Question 1: Is it possible to quantify the respective RANK needed for pre-training and fine-tuning? Thank you for this question, it’s very relevant. We can measure the number of singular values above a certain threshold (0.1) in the learned update (same way as we did for Figure 4). Here’s what we find: * Pre-training: 47% of singular values > 0.1 * Fine-tuning: 4% of singular values > 0.1 Experimental setup: - We used ReLoRA of rank=128 for both pre-training and fine-tuning - Pre-training resets LoRA parameters every 5K iterations (3 times total) - Fine-tuning resets LoRA every 1K iterations (2 times total) - For fine-tuning, we used the SST2 dataset and achieved 90% accuracy (full fine-tuning achieves 90% as well). Number of singular values > 0.1 (Average across layers) * Pre-Training (Q Projections): 31.2 * Fine-Tuning (Q Projections): 467.4 * Pre-Training (V Projections): 19.2 * Fine-Tuning (V Projections): 365.2 * Pre-Training (Up Projections): 59.5 * Fine-Tuning (Up Projections): 568.0 * Pre-Training (Down Projections): 57.0 * Fine-Tuning (Down Projections): 539.5 As you can see, we see a drastic difference between pre-training and fine-tuning, which is probably caused by much higher complexity and the amount of data in the pre-training task, as compared to the downstream classification task (SST2).

### Question 2: Can equation 3 prove that $Δ W$ achieves a higher rank compared to LoRA? Is it possible to add the relationship between higher rank and random initialization? Thanks for raising this question. We would like to explain our motivation behind ReLoRA and why we think proving this could be difficult. **Intuition for why we need optimizer resets** Let’s assume that training a model using unrestricted gradient descent (regular training) performs a high-rank update (higher than a single LoRA update). This assumption is verified in Figure 4. Let’s now learn two LoRA updates sequentially using ADAM optimizer. After learning the first update, we perform a merge-and-reinit operation on the model weights where merge is $$ \begin{align*} W^{t} &= W^{t-1} + s W^{t-1}_A W^{t-1}_B \\\\ W^{t}_A &= 0 \\\\ $W^{t}_B &= N\left(0, \frac{2}{h}\right) \end{align*} $$ At the next iteration, model update using Adam will be $$ \begin{align*} m^{t+1} &= \beta_1 m^{t} + (1 - \beta_1) g^{t} \\\\ v^{t+1} &= \beta_2 v^{t} + (1 - \beta_2) (g^{t})^2 \\\\ \hat{m}^{t+1} &= \frac{m^{t+1}}{1 - (\beta_1)^{t+1}} \\\\ \hat{v}^{t+1} &= \frac{v^{t+1}}{1 - (\beta_2)^{t+1}} \\\\ \\\\ W^{t+1}_A &= W^{t}_A - \eta \frac{\hat{m}^{t+1}}{\sqrt{\hat{v}^{t+1}} + \epsilon} \end{align*} $$ $W^{t+1}_B$ update is similar. Now, accounting for the high values of $\beta_1$ and $\beta_2$ we can say that the update at step t+1 will be highly similar to the previous update steps, because $m^{t}$ and $v^{t}$ are carried on from the previous update. Because of the high similarity between them, which will carry on over many update iterations due to high value of the betas, intuitively, we expect the learned update to either not change the rank at all or only minimally increase it, which motivates our optimizer reset strategy presented in the paper. And is verified by our ablation study (Table 4). **Intuition for why proving rank increase is hard** Let’s assume that instead of ADAM we use SGD. In this case, there is no dependence between the updates. $$ \begin{align*} W^{t+1}_A &= 0 + \eta \frac{\partial L}{\partial W_A}(W^{t+1}_A) \\\\ W^{t+1}_B &= N\left(0, \frac{2}{h}\right) + \eta \frac{\partial L}{\partial W_B}(W^{t+1}_B) \end{align*} $$ At this point, $W^{t+1} = W^{t} + W^{t+1}_A W^{t+1}_B$ it likely to be dominated by the initialization of $W^{t}_B = N(0, \frac{2}{h})$. However, the longer we train the network (let’s assume we trained it for 1000 steps), the less the update relies on the initialization of $W_B$ and more on the dynamics of neural network training. $$ \begin{align*} W^{t+1000}_A &= 0 + 1000 \eta \sum_i \frac{\partial L}{\partial W_A}(W^{t+i}_A) \\\\ W^{t+1000}_B &= N\left(0, \frac{2}{h}\right) + 1000 \eta \sum_i \frac{\partial L}{\partial W_B}(W^{t+i}_B) \end{align*} $$ (assuming constant learning rate for simplicity) As far as we know, as of now, no reliable theory has been proposed to explain NN training dynamics here, and developing it lies outside of the scope of our study. We would be grateful for any references that you could point to here. Finally, as our experiments show, a combination of LoRA resets with optimizer resets and learning rate re-warmup is needed for ReLoRA to work. This means that any proof that does not use these conditions will be insufficient. We believe that ReLoRA is self-sufficient as a purely empirical method and that the experiments we report are valuable due to consistent results across model scales from very small (70M) to relatively large (1.3B) scale. While having some theoretical guarantees is desirable, deep learning's history is peppered with contradictions to established theories. For instance, in 2012-2017, many in the ML community doubted the usefulness and reliability of deep learning, arguing that non-convex optimization with SGD wasn't feasible and the large number of parameters would lead to overfitting and poor generalization. Yet, these concerns were proven unfounded with the advancements post-2018. It's also worth noting that many ICML, ICRL, and NeurIPS papers include mathematical proofs with unrealistic simplifications, making them irrelevant to actual use cases. We aim to avoid this pitfall and ensure our paper remains relevant to real-world applications.

### Question 3: 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? This is consistent with our current understanding. A similar pattern was also observed for the lottery ticket hypothesis, where for larger models an initial warmup training of the full network was required [5], but the rest of the training could be approximated through a low-rank lottery ticket training. ### Question 4: How can ReLoRA speed up the training, as the forward and backward computation is not sparsified (or even more computation)? Similarly to PEFT methods (e.g., LoRA), ReLoRA significantly reduces the memory spent on optimizer states and gradients, allowing the use of larger microbatch sizes, which increases the efficiency. Further, distributed training is frequently bottlenecked by the GPU communication speed, even for ~1B-sized models. ReLoRA drastically reduces the communication volume to synchronize the gradients as the number of trainable parameters is much smaller. Finally, the update itself is much faster, as it requires less computation within ADAM: computing new optimizer states, computing the update, and applying the update to ReLoRA parameters. This question is discussed in-depth in this PEFT survey [4]. In total, ReLoRA speeds up training in the same way PEFT does, and its speedup is dependent on the model size, with larger models getting bigger speedups. ### Question 5: How is the performance and further speedup when you combine the ReLoRA and low-precision quantization, as mentioned in the paper? We did perform some experiments with 4-bit QLoRA instead of LoRA with the 1.3B model. In our observations, we saw virtually no difference between quantized and non-quantized models in terms of the loss, but unfortunately we observed a slowdown. After talking to QLoRA authors, we concluded that this behavior is expected for models of this size. The reasons are as follows: the normal float 4 (nf4) format stores floats in only 4 bits, but does computation in bfloat16, so for every multiplication, there is an added overhead of converting the numbers to bf16-floats and performing a 16-bit operation. For the 1.3B model in our case, using the quantized weights only wins about 2Gb GPU memory, not allowing to increase batch size significantly enough to improve the throughput. Again, we would like to thank you for your review. We hope that our responses address your concerns and we would like to ask you to increase your score. [1] Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning, Armen Aghajanyan, Luke Zettlemoyer, Sonal Gupta, 2020 [2] Few-Shot Parameter-Efficient Fine-Tuning is Better and Cheaper than In-Context Learning, Haokun Liu, Derek Tam, Mohammed Muqeeth, Jay Mohta, Tenghao Huang, Mohit Bansal, and Colin Raffel, 2022 [3] Implicit Self-Regularization in Deep Neural Networks: Evidence from Random Matrix Theory and Implications for Learning, Charles Martin, Michael Mahoney [4] Scaling Down to Scale Up: A Guide to Parameter-Efficient Fine-Tuning. Vladislav Lialin, Vijeta Deshpande, Anna Rumshisky. 2023. [5] Stabilizing the Lottery Ticket Hypothesis, Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M. Roy, Michael Carbin, 2020

Thank you for your thorough review! Also, big thanks for highlighting some issues with the text. These sections were probably added last and we hope that the rest of the text and the general organization is clear. We will address them and do another round of proofreading. We will now respond to your concerns, and then discuss the paper [1] you mention. ## Will the method scale to larger model sizes? Regarding whether ReLoRA will be able to achieve the same loss as full-rank training at 7B or even 70B scale, we strongly believe so. Multiple papers suggest that full-rank NN training is locally low-rank. For example, [2] show that as the model becomes larger or when pre-trained for longer, the rank of the change needed to learn a downstream task reduces. [3] finds that SGD is biased towards low-rank solutions. The existence of Lottery Tickets early in training [4] also partially supports this hypothesis, since training a lottery ticket network could effectively be seen as a low-rank approximation to the regular training process. While there is no formal proof yet that NN training is locally low-rank, our results both 1) provide a practical contribution to more efficient training and 2) support the growing evidence that despite their high dimensionality, neural network optimizations may intrinsically lie in far lower dimensional spaces. We had a similar question about an early version of this paper. At the time, the biggest model we had trained was 350M and the question was if this would scale beyond 1B parameters. In this iteration of the paper, we were able to successfully scale ReLoRA to 1.3B. ## Comparing ReLoRA training efficiency to full-rank training Thanks for raising this question! We would like to highlight several important details. First, for ReLoRA, as for all PEFT methods the speedup increases with model size. For example, at 130M-scale using ReLoRA can actually be significantly slower than regular training. 350M is roughly when ReLoRA throughput matches regular training and at 1.3B, we already see some significant improvement: 15% immediate speedup (not accounting for warm start) and 9% total speedup as reported in the paper. Second, the speedup is significantly hardware-dependent with ReLoRA being more effective at cheaper hardware. For details please look at this table: | | 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 (33% warm-start) | 149 ex/sec | 111 ex/sec | 14.8 ex/sec | | Total speedup | 9% | 32% | 51% | Throughput is measured in examples per second. This table will be added into the paper to clarify the speedups. [1] Exploring Low Rank Training of Deep Neural Networks. Siddhartha Rao Kamalakara, Acyr Locatelli, Bharat Venkitesh, Jimmy Ba, Yarin Gal, Aidan N. Gomez. 2022 [2] Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning, Armen Aghajanyan, Luke Zettlemoyer, Sonal Gupta, 2020 [3] Implicit Regularization in Deep Matrix Factorization, Sanjeev Arora, Nadav Cohen, Wei Hu, Yuping Luo, 2019 [4] Linear Mode Connectivity and the Lottery Ticket Hypothesis, Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M. Roy, Michael Carbin, 2020 [5] Implicit Self-Regularization in Deep Neural Networks: Evidence from Random Matrix Theory and Implications for Learning, Charles Martin, Michael Mahoney. 2018 We are again very thankful for your review and would like to ask you to increase your score if your concerns were sufficiently addressed.