Paper: ReLoRA: High-Rank Training Through Low-Rank Updates (ICLR-2024)

Reviewer comment (Experiments):
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.

Author response:
### 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  |

Reviewer comment (Experiments):
In Table 3, why isn't full fine-tuning included?

Author response:
### 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 |

Reviewer comment (Reproducibility):
For the Control Baseline, how are the trainable parameters chosen?

Author response:
### 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                |

Reviewer comment (Experiments):
ReLoRA outperforms LoRA for fine-tuning can be due to the Warm start?

Author response:
### 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.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Theory):
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).

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Theory):
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?

Author response:
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.

Reviewer comment (Evaluation):
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?

Author response:
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.

Reviewer comment (Evaluation):
How can ReLoRA speed up the training, as the forward and backward computation is not sparsified (or even more computation)?

Author response:
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.

Reviewer comment (Experiments):
How is the performance and further speedup when you combine the ReLoRA and low-precision quantization, as mentioned in the paper?

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Writing):
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.

Author response:
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.

Reviewer comment (Novelty):
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.

[1] Exploring Low Rank Training of Deep Neural Networks. Siddhartha Rao Kamalakara, Acyr Locatelli, Bharat Venkitesh, Jimmy Ba, Yarin Gal, Aidan N. Gomez. 2022

Author response:
We actually looked into this paper after the submission and we think it’s a great example of how low-rank training could be beneficial for improving computational efficiency of training large models. We would like to highlight some similarities and differences between our methods.

Before applying factorization, Kamalakara et al. first train the network in a regular high-rank way (Section 3.4. Pre-training). Specifically, they train full-rank for 40K steps and then train low-rank for 200K more steps, which is roughly on-par with our 5K/20K ratio. This goes inline with our results and makes their final method close to our Warm Start + LoRA baseline.

Our work differs notably in these key areas: 1) Kamalakara et al. perform an SVD decomposition of the weights after the initial full-rank training and only keep low-rank versions of the weights 2) They use a relatively high rank which comprises 62.5% of the hidden state size, while we use the rank of 128, which is about 6% of the 1.3B model hidden size. In our studies, we found little to no difference between ranks 128 and 512 for the 1.3B model. We attribute it to the weight merging procedure which can aggregate to a higher rank than just a single update. One more difference between the approaches is that Kamalakara et al. reduce the learning rate for their method, in contrast to ReLoRA, where we increase it. This decision seems to be counter-intuitive for us, as generally models with more trainable parameters use smaller learning rates.

Differences in the experiment scale:

Additional differences between our experimental setups include the fact that they use very small sequence length (256 instead of 512) and a very small pre-training dataset (One Billion Words), which is on par with our 70M-scale experiments, but 20x smaller than our largest 1.3B experiments. Finally, and touching upon your question of scale, the largest experiment in the study seems to be less than 300 TPUv2 hours, compared to our largest experiment being 640 A100 hours. In terms of theoretical flops (16-bit), TPUv2 is 1.7 times slower than A100: 180 vs 312 Tflops. In total, the largest experiment in our paper used 3.6 more compute than in Kamalakara et al.