LoRA Done RITE: Robust Invariant Transformation Equilibration for LoRA Optimization

Thank you for the positive feedback!

it would be interesting to present the best learning rate for different strategies in this work after the search.

Thank you for the suggestions. In the revised version, we have included the best learning rates for different optimizers on different datasets in Appendix A.9 (Table 9).

Reviewer vefd,

Thanks again for the positive feedback! We have included the best learning rates for different optimizers on different datasets in the appendix as you suggested.

Additionally, we follow the suggestions of Reviewer zd3u and vefd to make the following modifications.

We include the visualization of the loss curve and the weight updates. This shows that the training loss of LoRA-RITE decreases faster and both the factor A and B can be learned effectively.

We also added experiments to show that that LoRA-RITE performs better than the baselines across different ranks and is competitive with full-finetuning.

Please feel free to let us know if you have any additional questions!

Thanks, All authors of Submission 8978

Thank you for the review comments!

Convergence analysis should be conducted to better justify the proposed method, e.g. the norm of A and B, like in figure 1.

Thank you for the suggestion. We have added visualization of update magnitudes of the proposed method in Appendix A.7 (Figure 3 and 4). Further, to make it more clear to visualize the difference between different methods, we plot the relative update magnitude (update norm divided by parameter norm) which is a better indicator of how much the parameters are updated. The figures show that for conventional optimizers, factor A barely changes, while LoRA-RITE is able to learn factor A effectively.

Additionally, we have added the training loss curves in Appendix A.6. The figures show that the training loss of LoRA-RITE decreases faster than the baseline methods.

The analysis over experimental results is limited, e.g., for some datasets, the proposed method demonstrates significant performance gain compared to LoRA (Adam), what is the property of the dataset such that the proposed optimization can result in such improvement?

We are still continually investigating the property of the dataset where LoRA-RITE can have significant performance gain. One observation we have made so far is that LoRA-RITE usually performs better when there is a gap between the baseline methods and full finetuning, while the improvement is usually milder if their performance is already close. This can be seen from the tables below.

In the table above, we observe a larger gap between full finetuning and lora in CausalEffectClassification, CorefereferenceResolution and Global, and the gap between LoRA-RITE and baselines is also larger on those tasks. Similar observations can be made in the following table, except the second and fourth column where LoRA-RITE outperforms Full Finetuning, which is probably due to some regularization effect.

As analyzed in at line 322, from time complexity perspective, the proposed method is r times slower than Adam, but in Table 4, it is only 8% slower, it will be interesting to elaborate more on the inconsistency.

The training speed measured in Table 4 includes the time for the model’s forward and backward passes. Since the complexity of the forward and backward passes is at least $\Omega(nm)$, the overhead of LoRA-RITE is relatively small. We have added these explanations to the new version of the paper.

Dear Reviewer wRHA,

Thanks again for your constructive suggestions!

We follow your suggestion to include the visualization of the loss curve and the weight updates. We also added experiments to show that that LoRA-RITE performs better than the baselines across different ranks and is competitive with full-finetuning.

We believe the addition of these results further verifies the effectiveness of LoRA-RITE and will be very helpful to the readers.

Please feel free to let us know if you have any additional questions!

Thanks, All authors of Submission 8978

Thank you for the detailed review!

The paper lacks visual illustrations of loss curves to cross-validate the effectiveness of their method in accelerating convergence.

Thank you very much for the suggestion. We have added the training loss curves in Appendix A.6 (Figure 2). The figures show that the training loss of LoRA-RITE decreases faster than other methods.

The authors do not provide visual evidence of how the magnitude of $A$ updates after applying their method

Thanks for the suggestion! We have added visualizations of update magnitudes of Lora-RITE to Appendix A.7 (Figure 3 and 4). Further, to make it more clear to visualize the difference between different methods, we plot the relative update magnitude (update norm divided by parameter norm) which is a better indicator of how much the parameters are updated. The figures show that for conventional optimizers, factor A barely changes, while LoRA-RITE is able to learn factor A effectively.

The authors do not address potential numerical instability issues. Specifically, their algorithm involves inverting the matrix $R_A$, which could be zero at the initial training step.

We thank the reviewer for the question. We would like to clarify that at the initial training step it is $R_B$ that is 0 instead of $R_A$.

There are two places in the algorithm that could have potential numerical instability and we did carefully address the numerical issues in our implementation as discussed below:

1. For inverting the matrix $R_A$ and $R_B$, instead of directly inverting $R_A$, we apply the pseudoinverse (we have added this note to the revised paper). At the initial training step, $R_B$ is 0, and so is the gradient of $A$. After multiplying with the pseudoinverse of $R_B$, we get that the unmagnified gradient of $A$ is also 0 to avoid NaNs.

2. For computing the inverse square root of $\bar{V_A}$ and $\bar{V_B}$, we add a small $\epsilon I$ to $\bar{V_A}$ and $\bar{V_B}$ respectively, before taking the inverse square root. Since $\bar{V_A}$ and $\bar{V_B}$ are already unmagnified, this does not impact transformation invariance.

We have included these details in the revision.

The conclusion and limitations section is somewhat lacking, as it only discusses the limitations of the base LoRA method rather than potential limitations of their proposed method.

Thanks for the suggestion. We have added the potential limitations of LoRA-RITE to the revised version: The work focuses on addressing transformation invariance when the optimization problem can be written in the form of f(AB), and this assumption may not hold for parameter-efficient structures beyond LoRA.

In terms of contributing to the LLM community, open-sourcing the related code would be highly appreciated.

Thank you for pointing out this. We plan to release the related code after it passes internal review.

Dear Reviewer zd3u,

Thanks again for your insightful suggestions.

We follow your suggestion to include the visualization of the weight updates. Additionally, the extra experiments asked by Reviewer wRHA shows that LoRA-RITE performs better than the baselines across different ranks and is competitive with full-finetuning. We believe these results further verify the effectiveness of LoRA-RITE.

We also improve the writing to include details on addressing potential numerical instability issues and to discuss the limitation of our method. We believe these modifications will greatly help the readers.

Please feel free to let us know if you have any additional questions!

Thanks, All authors of Submission 8978