Flat-LoRA: Low-Rank Adaption over a Flat Loss Landscape

Thank you for your valuable and detailed comments! We are glad that you find our paper good written, with innovative ideas and credible experimental performance. In the following, we provide our point-by-point response and hope our response helps address your concerns. We also look forward to the subsequent discussion which may further help solve the current issues.

W1: LoRA only fine-tunes a very small number of parameters, which is unlikely to overfit on a specific task and can inherently maintain good generalizability.

A1: Thanks for your insightful question. We agree with you that LoRA is less likely to overfit on a specific task due to its fewer trainable parameters. However, pursuit flatness can still help generation performance even when not in an overfitting regime. To validate this, we perform experiments with LoRA rank=1 and find that Flat-LoRA can still consistently outperform LoRA, although the improvements may not be as large as with a higher rank.

We hypothesis that when the lora rank is small, e.g. rank=1, LoRA may attempt to find the lowest loss in the corresponding rank-1 space, which may appear sharp in the full parameter space. This phenomenon is illustrated in the first column of Figure 5 in the paper, where we observe that with a small LoRA rank, the loss landscape in the merged weight space tends to be significantly sharper.

Moreover, the noise injection approach in Flat-LoRA can also be readily incorporated into full fine-tuning to provide significant generalization improvement with minimal additonal memory and computation cost:

W2: The limited enhancement on the GLUE benchmark does not prove the lack of generalizability of methods like LoRA; More datasets and models should be compared. For example: the SuperGLUE benchmark, SQuAD, XSum, CNN/Dailymail, and some LoRA training on Stable Diffusion.

A2: One reason for our limited improvements on GLUE is that we utilize a much higher baseline to evaluate the real effectiveness of our Flat-LoRA, leaving less room for improvement. In fact, our LoRA baseline is even much better than the optimized LoRA-GA (NeurIPS'24) performance [1], with 88.97 (our LoRA) on MRPC datasets compared to 68.38 (LoRA) and 85.29 (LoRA-GA) reported in [1]. Our improvements on GSM8K and Human-Eval are much significant (i.e. +3.18% and +1.37%) as the tasks are generally harder.

SuperGLUE We conduct experiments on more challenging SuperGlUE datasets with T5-base. The training settings remain the same as described in Section 4.1 of the paper. The results in the following table shows that Flat-LoRA outperforms LoRA with more significant improvements compared to the results on the GLUE datasets.

SQuAD, XSum, CNN/Dailymail We then conduct experiments on SQuAD, XSum and CNN/Dailymail datasets using T5-base model. The training settings remain consistent with those described in Section 4.1 of the paper, except that we adopt 3 training epochs. From the results below, we observe that Flat-LoRA consisitently outperforms LoRA, but the improvement is relatively small. We hypothesize that flatness may not be particularly beneficial for tasks like summarization and question answering but may confer greater advantages for tasks such as mathematics and code-related problems (see Q4 below and Section 4.3 in the revision). Moreover, the improvements brought by Flat-LoRA require minimal additional memory and computation cost.

Stable Diffusion We finetune SDXL with the pipeline of Dreambooth with the popular scripts implemented by HuggingFace. As shown in Figure 3 in the revision, Flat-LoRA exhibits better personalization than LoRA while maintaining better generation ability. For instance, in the second column, the image generated by Flat-LoRA includes a distinctive square behind the bird, aligning more closely with the "icon" feature present in the training images (top row). Furthermore, Flat-LoRA more effectively preserves the concept of eyes, whereas in columns 1, 3, and 5, the birds generated by LoRA are missing eyes.

We have included the above experiments into Section 4.4 and Appendix C, G in the revision.

[1] Wang, Shaowen, Linxi Yu, and Jian Li. "LoRA-GA: Low-Rank Adaptation with Gradient Approximation." NeurIPS'24.

W3: There is a lack of extensive comparisons, such as with methods like DyLoRA[1] and AdaLoRA[2].

A3: Thanks for your suggestion. We have added the DyLoRA and AdaLoRA baselines in the following table and Table 5 in the revision accordingly.

W4: More relevant articles, such as AWP[3], should be cited, which is an effective method to improve generalization.

A4: Thanks for providing the relevant works. We have added into citation accordingly.

Q4: Performance on the alpaca dataset

A5: Due to the limited rebuttal time and computation resource, we fine-tune Alpaca dataset with Llama2-13B model to also meet larger LLM experiments requested by Reviewer pg6m. From the results in the following table, we observe that Flat-LoRA consistently outperforms LoRA. We find that the improvements on DROP and HumanEval are more significant (+0.71% and +1.83%), perhaps indicating that flatter minima better support math-related and coding-related tasks. We have added these results into Section 4.3 in the revision.

Accodringly, we have evaluated our method using the suggested Matthews Correlation Coefficient metric on CoLA dataset. From the results in the following table, we observe, again, Flat-LoRA consistently outperforms other methods under the the matthews correlation metric. We have included this results in the updated version.

Thanks again for your time and effort.

Thank you for your valuable and detailed comments! We are glad that you find our paper well written with reasonable core idea. In the following, we provide our point-by-point response and hope our response helps address your concerns. We also look forward to the subsequent discussion which may further help solve the current issues.

W1: The paper lacks mathematical rigor in several places, particularly in equations 8 and 9: (1) var(X) should be a covariance matrix; (2) var(W'_{i,:}X) should be a scalar; (3) These dimensions are inconsistent and cannot be equated

A1: Thank you for pointing out these issues. In equations 8 and 9, we aim to characterize the effects of the varianace of the activation after adding random weight perturbations. We assumed activations have a similar magnitude of variance and used var(X) to represent it. In the revision, we use $\text{var}(\mathbf{x}\_i)$ for clearer notation as a scalar describing the variance of each element. We also rewrite var(W'_{i,:}X) to $\text{var}(\mathbf{W}'\_{j,:}\mathbf{x})=(1+\sigma^2) \|\mathbf{W}'\_{j, :}\|_2^2\cdot \mathrm{var}(\mathbf{x}\_i)$ to ensure that the output is a scalar. We hope our revision (equations 8-12) could make the statement more rigor.

W2: Insufficient experimental validation: (1) Should test on larger models (e.g., LLAMA 13B or 70B) (2) Should evaluate on SuperGLUE

A2: Following your suggestions, we fine-tune Llama-2-13B model on Alpaca and evaluate InstructEval metrics. From the results in the following table, we observe that Flat-LoRA consistently outperforms LoRA. We find that the improvements on DROP and HumanEval are more significant (+0.71% and +1.83%), perhaps indicating that flatter minima better support math-related and coding-related tasks.

We then evaluate Flat-LoRA on more challenging SuperGlUE datasets with T5-base. The training settings remain the same as described in Section 4.1 of the paper. The results in the following table shows that Flat-LoRA outperforms LoRA with more significant improvements compared to the results on the GLUE datasets.

We have added these results into Section 4.3 and Appendix C in the revision.

W3: Lacks necessary ablation studies, particularly regarding σ^2

A3: Thank you for your suggestion. We vary $\sigma^2$ on fine-tuning CIFAR-10/100 with CLIP ViT-B/32. We find that the optimal results are achieved when $\sigma^2$ is 0.10 or 0.15. We have included it into Appendix E.

W4: In Equation 7, n should be sqrt(n), as large n values would result in negligibly small epsilon values added to the weight matrix

A4: We apologize for the error that the factor in Line 208 should be $1/\sqrt{n}$ rather than $n$, and we have corrected it in the revision. We agree that the magnitude should be scaled by a factor of $1/\sqrt{n}$ given the input dimension $n$ to avoid negligibly small $\boldsymbol{\epsilon}$ with large $n$. In Equation 7, $n$ represents the order of the variance rather than the standard deviation, and the factor should be squred as $n$. In Equation 8-12 of the revision, we also show that such scaling would increase the variance of the activation by a factor of $\sigma^2$, which is not related to the input dimension $n$ as we expected.

W5: No memory usage results are reported; Table 3 lacks full-tuning baseline results

A5: We report the memory usage for fine-tuning GSM8K datasets using Llama 2-7B model. The experiments are conducted with BF16 mixed-precision training and a micro-batch size of 2, running on an NVIDIA GeForce RTX 4090 GPU. We observe that Flat-LoRA brings very little additonal memory cost compared to LoRA, comfirming its effectiveness on maintaining memory efficiency. We have included it into Appendix D.

We complement the full fine-tuning baselines for Table 3 as follows: when the gap between LoRA and full FT is large, Flat-LoRA can bring more significant improvement over LoRA.

Thank you very much for your thorough review and prompt feedback. We address your further concerns below:

A missing square in Eq. 11&12:

Indeed, they should be:

Eq. 11:

$=||\mathbf{W}'\_{j, :}||\_2^2\cdot \mathrm{var}[\mathbf{x}\_i]+n\cdot \frac{\sigma^2}{n} ||\mathbf{W}'\_{j, :}||^2\_2\cdot \left(\mathrm{var}[\mathbf{x}\_i]+\mathbb{E}^2[\mathbf{x}\_i]\right)$

Eq. 12:

$$=(1+\sigma^2) ||\mathbf{W}'\_{j, :}||\_2^2\cdot \mathrm{var}[\mathbf{x}\_i]+\sigma^2||\mathbf{W}'\_{j, :}||^2\_2\cdot\mathbb{E}^2[\mathbf{x}\_i].$$

$\boldsymbol{\epsilon}$ in Eq. 7:

$\boldsymbol{\epsilon}$ is defined as a matrix with the same dimensions as $\mathbf{W}$, representing the noise added to $\mathbf{W}$. As you correctly point out, the scalar $0$ in Eq. 7 is should be $\mathbf{0}$. To clarify this further, we rewrite the definition of the distribution for each element of $\boldsymbol{\epsilon}$ as follows:

$\boldsymbol{\epsilon}\_{i,j} \sim \mathcal{N}\left (0, \frac{\sigma^2}{ {n} } ||\mathbf{W}'\_{i, :}||_2^2 \right).$

To avoid any misunderstanding that $\boldsymbol{\epsilon}$ represents a vector, we would replace it with notation $\bf N$ as a random matrix in the revised manuscript.

Thank you again for your valuable feedback, which have greatly helped improve our manuscript. We will incorporate these corrections into the revised version.

Thank you for your valuable and detailed comments! We are glad that you find our idea of optimizing to reach a flat landscape in the full parameter space to be innovative and well-justfied with practical implications. In the following, we provide our point-by-point response and hope our response could address your concerns. We also look forward to the subsequent discussion which may further help solve the current issues.

W1: Althougth the paper emphasizes the computational overhead and the minimal memory, the perturbation generation and its integration into the mixed-precision training could be simplified or clarified furthermore.

A1: Thank you for your suggestion. As you mentioned, the main contribution of this paper is the efficient and effective generation of perturbations to optimize the flatness of loss landscape in the full parameter space. For mixed precision training, which is widely adopted during LLM training, we offer an easy approach to seamlessly integrate the perturbation injection process into the precision casting, introducing no addtional memory cost. However, our main approach is to efficiently store the perturbation based on filter norms and random seed, which is more general and does not require mixed-precision training.

We have made the statement more clear in the revision (Page 5).

W2: The method concentrates mostly on linear layers in the transformer-based models. Despite the fact that the authors acknowledge this as a limitation, extending such approach to other parameter categories would make the method more versatile.

A2: Following your suggestion, we extend the injection of random weight perturbation to all layers, referred to as "Flat-LoRA (all)". Specifically, we additionally add perturbations to layernorm layers, biases, and class embeddings, etc. We generate noise based on the absolute weight $|\mathbf{W}|$. From the results in the following table, we observe that Flat-LoRA (all) indeed improves performance, though the improvement is not as large as Flat-LoRA (Linear) over LoRA. We include this extension in Appendix A.

W3.: Although the comparisons with approaches such as SAM are very insightful, deeper analysis with more recent variations of sharpness-aware algorithms could strengthen the study and contribution.

A3: Following your suggestion, we also compare Flat-LoRA and LoRA with two SAM's variants, ASAM and GSAM. From the results in the following table, we observe that optimizing the sharpness over the full parameter space $\mathbf{W}$ generally provides better performance than on the LoRA space $\mathbf{A}$ and $\mathbf{B}$. We have added these experiments into Appendix B.

Q1: Could the perturbation generation strategy be optimized or adapted to incorporate other noise categories (e.g., adversarial perturbations)?

A4: Yes. It is possible to incorporate adversarial perturbations with random weight perturbations, i.e., perturbing the weights with random perturbations before applying adversairal ones, to enhance generalization performance, as shown in R-SAM [1]. It's also possible to alternatively use random and adversarial perturbations to improve the efficiency of sharpness-aware training within a certain computation budget. However, we note that incoporating adversarial weight perturbation doubles the time training cost, requiring two gradient step per iteration, and introduces a hard copy of the perturbation weights in memory, which may be undesirable for PEFT training setting.

[1] Liu, Yong, et al. "Random sharpness-aware minimization." Advances in Neural Information Processing Systems (2022).

Thank you for your valuable and detailed comments! We are glad that you find our paper well written with clearly explained insights, and that you consider the idea of storing random seeds for memory efficiency and integrating into mixed-precision training to be very clever. In the following, we provide our point-by-point response and hope our response helps address your concerns. We also look forward to the subsequent discussion which may further help solve the current issues.

W1: On line 125, the sentence "the LoRA matrices. better initialization strategy for LoRA parameters." should probably use a comma instead of a period.

A1: Thanks for your careful reading. We have fixed the sentence to "initialize the LoRA matrices, which prvoides a better initialization for LoRA parameters".

W2: The work on adding full-rank noise has been done in Noisetune (https://arxiv.org/abs/2202.12024) and LORASC (https://arxiv.org/abs/2407.01491), especially in LoRASC, where the exact same approach is used, adding a full rank noise to each LoRA optimization process.

A2: Thank you for providing the relevant works; we have included them in our citations. Here, we clarify the differences between our Flat-LoRA and Noisetune and LoRASC:

1. Output: Flat-LoRA produces a task-specific LoRA for each task, enabling efficient deployment. In contrast, both Noisetune and LoRASC generate full-rank weights for each task. In LoRASC, LoRA parameters are merged into the weights each epoch, resulting in high-rank updates.
2. Objective: We consider the flatness of weights after merging the LoRA parameters, while LoRASC generates noise based solely on LoRA parameters. Flat-LoRA is more aware of the loss landscape for the final output weights.
3. Method: We remove the weight noise after calculating the perturbed gradient for each step. However, Noisetune does not require recovering unperturbed weights and adds weight noise only once before training, while LORASC injects once before each training epoch. Moreover, we carefully consider the width factor $n$ into noise generation, and analyze the effects to activations.
4. Memory: We address memory issues when introducing weight noise and propose strategies to manage it carefully. In comparison, LoRASC needs to store a copy of perturbed weights $\widetilde{W}$, which is not memory efficient for PEFT fine-tuning.

W3: (1) The proposed Flat-LoRA series offers limited improvement across various LoRA variants and tasks, with most gains being around a few tenths of a percentage in accuracy (except for significant improvements in gsm8k and Human-Eval). It would be helpful to include tasks like apaca and other SFT tasks. (2) I’m particularly interested in the practical significance of this work—flat local optima should theoretically bring stronger generalization, such as supporting OOD (e.g., Alpaca, instruct-eval, etc.) and robustness (e.g., image-R, image-C) evaluations, proving that a flatter local optimum could enhance out-of-domain generalization.

A3: Thank you for your detailed suggestions. (1) We appreciate your recognition of our significant improvements on GSM8K and Human-Eval. One reason for our limited improvements on GLUE is that we utilize a much higher baseline to evaluate the real effectiveness of our Flat-LoRA, leaving less room for improvement. In fact, our LoRA baseline is even much better than the optimized LoRA-GA (NeurIPS'24) performance [1], with 88.97 (our LoRA) on MRPC datasets compared to 68.38 (LoRA) and 85.29 (LoRA-GA) reported in [1].

To further evaluate the effectiveness of our approach, we conduct experiments on more challenging SuperGlUE datasets with T5-base. The training settings remain the same as described in Section 4.1 of the paper. The results in the following table shows that Flat-LoRA outperforms LoRA with more significant improvements compared to the results on the GLUE datasets.

(2) Following your suggestion, we evaluate OOD and robustness performance on Alpaca and corruption datasets, respectively.

For Alpaca, we fine-tune with Llama-2-13B model and evaluate InstructEval metrics. From the results in the following table, we observe that Flat-LoRA consistently outperforms LoRA. We find that the improvements on DROP and HumanEval are more significant (+0.71% and +1.83%), perhaps indicating that flatter minima better support reasoning and coding tasks.

For the corrpution datasets, we fine-tune CLIP ViT-B/32 on CIFAR-100 and test the model on OOD CIFAR-100-C datasets. We report the results across different corruption levels (from 1 to 5). As shown in the following table, we observe that Flat-LoRA outperforms LoRA more as the corruption level increases. This shows that a flatter local optimum could enhance out-of-domain generalization, as you suggested.

We have added these results into Section 4.3 and Appendix C and F in the revision.

[1] Wang, Shaowen, et al. "LoRA-GA: Low-Rank Adaptation with Gradient Approximation." NeurIPS'24.

Q1: What do you think is the essential difference or advantage between adding perturbations to weights and adding perturbations to data samples? Is adding perturbations to data samples simpler?

A4: Adding perturbations to weights encourages the network to locate at a flatter minima, leading to better generalization performance. In contrast, adding perturbations to data samples improves the network's robustness to input data perturbation, such as adversarial perturbation in adversarial training. However, adding perturbations to data can lead to the distribution shift of data and ofen results in worse generalization performance on clean data, which is discussed in many adversairal training works, e.g. [1-3].

[1] Raghunathan, Aditi, et al. "Understanding and Mitigating the Tradeoff between Robustness and Accuracy." ICML'20.

[2] Xing, Yue, et al. "On the generalization properties of adversarial training." AISTATIS'21.

[3] Moayeri, Mazda, et al. "Explicit tradeoffs between adversarial and natural distributional robustness." NeurIPS'22.

Dear Reviewer,

Could you kindly respond and indicate whether authors have addressed your concerns?

Thanks, AC

Thank you for the author's response. I understand the differences between the proposed work and Noisetune as well as LORASC mentioned by the author. However, I am not quite clear about the exciting aspects brought by this work. Being different doesn't necessarily constitute a contribution, though I acknowledge that there are some clever ideas in the author's design. Still, I don’t think it represents a significant advancement in understanding for the community.

That said, I find the OOD and robustness experiments to be very positive. At the very least, they demonstrate a certain level of practical significance for this work.

In summary, I will stick to my current rating. I believe this paper is roughly at the level of a weak acceptance.