Theoretical Insights into Fine-Tuning Attention Mechanism: Generalization and Optimization

Authors: Xinhao Yao, Hongjin Qian, Xiaolin Hu, Gengze Xu, Yong Liu
Primary area: foundation or frontier models, including LLMs
Keywords: Large Language Models, Attention mechanism, Fine-tuning, Generalization and Optimization

Abstract:
Large Language Models (LLMs), built on Transformer architectures, exhibit remarkable generalization across a wide range of tasks. However, fine-tuning these models for specific tasks remains resource-intensive due to their extensive parameterization.  
In this paper, we investigate two remarkable phenomena related to the attention mechanism during the fine-tuning of LLMs. The first phenomenon, termed “Unequal Importance of Attention Matrices,” highlights the impact of fine-tuning different weight matrices. It shows that optimizing the $\mathbf{W}_v$ matrix yields significantly better performance than optimizing the $\mathbf{W}_k$ matrix. Fine-tuning only the $\mathbf{W}_q$ and $\mathbf{W}_v$ matrices is computationally efficient while delivering results comparable to, or even better than fine-tuning all three matrices ($\mathbf{W}_q$, $\mathbf{W}_k$, and $\mathbf{W}_v$). The second phenomenon, “Attention Matrices with Customized Learning Rate Leads to Better Convergence,” emphasizes the importance of assigning distinct learning rates to these matrices. Specifically, a higher learning rate for the $\mathbf{W}_v$ matrix compared to $\mathbf{W}_q$ and $\mathbf{W}_k$ accelerates convergence and improves performance. Building on these insights, we propose a new strategy that improves fine-tuning efficiency in terms of both storage and time. Experimental results on benchmark datasets validate the effectiveness of this approach, supporting our theoretical findings. Our analysis lays the theoretical groundwork for configuring and improving lightweight algorithms in LLMs fine-tuning.

Decision: reject
Reviewer scores: 5, 6, 5, 3

Reviews and discussion:
<thread>
creation date:
1730096695714

thread_type:
official_review

<thread_object>
creator: reviewer_txen

<content>
summary:
The paper investigates two phenomena observed during the fine-tuning of Transformer LLMs, focusing on the attention mechanism. The first phenomenon is the "Different Impact," where optimizing the Wv (value) matrix significantly improves performance more than optimizing the Wk (key) matrix. The second is "Efficient Convergence," where using distinct learning rates for different matrices is crucial for optimal performance, with a higher learning rate for Wv being beneficial for faster convergence. The paper provides theoretical analysis on these phenomena from the perspectives of generalization and optimization. Based on this, this paper proposes only finetuning Wq and Wv, and with a larger learning rate for Wv.

strengths:
The observation that fine-tuning only Wq and Wv​ matrices yields performance gains comparable to finetuning of Wq, Wk and Wv is interesting. Trying to analyze and uncover the reason can benefit the research commnity.

weaknesses:
See my questions.

questions:
1. My biggest concern with this paper is that the key observation seems to originate from the LoRA paper: *"Adapting both \( W_q \) and \( W_v \) gives the best performance overall"* (this sentence is quoted from the LoRA paper), and even the experiment setup is similar to Table 5 in the LoRA paper. Deriving key observations from other works isn’t an issue (or saying we share the same observation/insights), but I believe it should be directly and explicitly pointed out, rather than presented as your own observation, concerning the lora paper is very famous for several years and you do cite it. Or could you clarify what differences there are between your work and the LoRA paper? Please provide a detailed comparison between their findings and those in the LoRA paper, highlighting any novel contributions or deeper insights their work provides. This would help clarify the paper's originality and contribution to the field.

2. As you pointed out on line 205, \( W_q \) and \( W_k \) can be treated as a single unit. In linear algebra, two matrices multiplied without an intermediate activation can be equivalent to a single matrix. This could explain why fine-tuning only \( W_q \) and \( W_v \) achieves comparable accuracy to tuning all matrices. If this is the case, have you tried fine-tuning only \( W_k \) and \( W_v \)? It seems that neither the LoRA paper nor your work has conducted this experiment. Could you please include this experiment in your study or explain why not to perform it. This could provide valuable insights and strengthen the paper's analysis.

3. What are the main conclusions of the theoretical analyses in Sections 3 and 4? I found it difficult to follow your proofs. Please include a clear summary of the main conclusions from their theoretical analyses at the end of Sections 3 and 4. I would suggest to provide more intuitive explanations to help readers better understand the theoretical analyses and proofs.

4. Section 5 appears to be simply fine-tuning \( W_q \) and \( W_v \) with a search for \lambda. Do you have any insights on how to determine \lambda? could you provide guidelines or heuristics for determining an appropriate \lambda value based on your theoretical and empirical results?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730293586932

thread_type:
official_review

<thread_object>
creator: reviewer_f8uv

<content>
summary:
This paper explores fine-tuning strategies for the attention mechanism from a theoretical perspective, focusing on generalization and optimization issues, and claims to provide theoretical insights to guide algorithm design. It has two major propositions, one is fine-tuning WqWv matrix is more efficient and can achieve comparable results as fine-tuning WqWkWv matrix. Another is fine-tuning Wq and Wv should use different learning rates for more efficient convergence.

strengths:
The idea of using different local learning rates for Wq and Wv is interesting. As V is timed with attention logits, q and v should have more different gradient distributions than q and k. Using different local learning rates for these two types of matrices is reasonable and worth exploring.

weaknesses:
1. The writing and presentation are rather poor. Both sentence level and logical level polishment are suggested for this work. 
    E.g. In line 16~22, "In this paper, we investigate two remarkable phenomena ... with a higher learning rate for the Wv matrix expediting convergence." The ordering of the sentences and the way of structuring the arguments add unnecessary difficulty to reading. This kind of problem happens across the whole manuscript.
    A suggestion is, use multiple statement sentences instead of clauses. A rewritten version of these sentences is " In this paper, we investigate two remarkable phenomena associated with attention mechanism, during the fine-tuning of LLMs. One is 'Different Impact' and another is 'Efficient Convergence'. ......" Another point is that the naming of "Different Impact" and "Efficient Convergence" doesn't help with understanding. Consider other ways to name them.

2. For the "Different Impact", the authors claimed that fine-tuning Wq&Wv is favored over Wq,Wk,Wv together. However, this argument is not supported by their experimental results in Table 1. It can be seen from the table that fine-tuning Wq,Wk,Wv together still gets the best performance for most of the cases. I wonder if the authors are writing the arguments and conducting experiments separately without any discussion. The authors should summarize their findings according to the experiment results.

questions:
1.  In line 160~193, the author writes "As seen in the table, we can see a clear trend where solely updating the Wv matrix outperforms just learning the Wq,Wk matrix. Interestingly, the combination of fine-tuning both Wq and Wv often leads to performance that matches or even exceeds that achieved by fine-tuning all three matrices Wq,Wk, and Wv". This statement is not supported by the results in Table 1, as solely updating the Wv is inferior to updating Wq,Wk together, and fine-tuning all three matrices Wq,Wk, and Wv achieves the best result in most cases. The authors should first rewrite the whole section of "Different Impact" to have a consistent argument with the experiment results, and investigate more tasks to see whether there are certain types of tasks in which "different impact" is true.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730703531023

thread_type:
official_review

<thread_object>
creator: reviewer_ftn2

<content>
summary:
This paper presents theoretical and empirical insights into improving the efficiency of fine-tuning large language models. The key contributions are as follows:

1. Different Impact in Fine-Tuning: The paper demonstrates that fine-tuning only the value (`Wv`) and query (`Wq`) matrices of the attention mechanism is computationally efficient and can yield results comparable to or better than fine-tuning all three matrices (`Wq`, `Wk`, `Wv`). It suggests that `Wv` has a significantly greater impact on performance compared to `Wk`.

2. Efficient Convergence: The authors show that using distinct learning rates for different matrices, particularly a higher rate for `Wv`, expedites convergence and optimizes the fine-tuning process. This insight provides a practical approach to achieving better fine-tuning results with fewer computational resources.

3. Theoretical Analysis: The paper provides a theoretical analysis from the perspectives of generalization and optimization. It establishes that fine-tuning `Wq` and `Wv` enhances memory efficiency and generalization bounds while optimizing convergence dynamics by accelerating learning for specific matrices.

4. Proposed Strategy: Building on these insights, the authors propose a strategy to improve fine-tuning efficiency, both in terms of storage and time. Experimental results on benchmark datasets validate the effectiveness of this approach.

strengths:
- Originality: The originality of this paper lies in its focused approach to fine-tuning the attention mechanism of large language models. By exploring the selective fine-tuning of `Wv` and `Wq` matrices, the authors introduce a novel method that challenges the conventional approach of fine-tuning all attention matrices (`Wq`, `Wk`, `Wv`). The theoretical insights into the distinct roles of these matrices, combined with empirical validation, provide a fresh perspective on optimizing the attention mechanism. Additionally, the proposed use of differentiated learning rates for convergence further demonstrates creativity in enhancing the fine-tuning process.

- Quality: The quality of the paper is reflected in its theoretical analysis and empirical validation. The authors employ both information-theoretic methods and optimization theory to support their claims, which significantly strengthens the credibility of their contributions.

- Clarity: The paper is generally well-written, making the complex ideas accessible to a broad audience. 

- Significance: The significance of the paper lies in its potential to influence future research and practical implementations of fine-tuning large language models. This contribution is highly relevant for both academia and industry, as it provides a more resource-efficient approach to adapting pre-trained models for downstream tasks.

weaknesses:
1. **Lack of Base Model Performance for Each Task**: The paper does not provide the performance of the base model before fine-tuning for each task, making it challenging to evaluate the true effectiveness of the fine-tuning methods. Including these baseline results would help contextualize the improvements made through fine-tuning.

2. **GLUE Evaluation Is Too Simple for LLaMA3.1-8B**: The use of the GLUE benchmark to evaluate the LLaMA3.1-8B model is insufficient, as GLUE tasks are relatively simple compared to the capabilities of such a large model. Evaluating on more challenging tasks, such as coding or mathematical problem, could make the results more convincing and demonstrate the robustness of the proposed approach.

3. **Lack of Ablation Studies**: The paper could benefit from more extensive ablation studies to isolate the effects of different components of the proposed method, such as the impact of learning rate scaling and the specific contribution of each matrix (`Wq`, `Wv`). This would provide a clearer understanding of the factors contributing to the observed performance gains.

4. **Lack of Guidance on Hyperparameter `λ`**: The paper does not provide sufficient guidance on how to choose the hyperparameter `λ`, which controls the relative learning rates of different matrices. Without explicit guidelines or a heuristic for selecting `λ`, practitioners may find it difficult to replicate the reported results or apply the method effectively in different contexts.

questions:
1. **Baseline Performance**: Could you provide the performance metrics of the base model before fine-tuning for each of the tasks? This would help in better understanding the relative improvements brought by your fine-tuning strategy.

2. **More Challenging Evaluations**: Have you considered evaluating the LLaMA3.1-8B model on more complex benchmarks, such as tasks involving coding or mathematical problem? Including such challenging tasks would make your results more comprehensive and convincing.

3. **Ablation Studies**: It would be helpful if you could add more ablation studies to isolate the effects of different components of the proposed method, such as the specific impact of `Wq` vs. `Wv` fine-tuning or the effect of learning rate scaling. This would provide a clearer picture of what drives the observed improvements.

4. **Hyperparamete `λ`**: Could you provide more guidance on how to choose the hyperparameter `λ`?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730707674519

thread_type:
official_review

<thread_object>
creator: reviewer_swzh

<content>
summary:
This paper investigates and addresses the resource-intensive nature of fine-tuning Large Language Models (LLMs) on transformer architectures, focusing specifically on the attention mechanism. The study identifies two key phenomena:

1. **Different Impact**: The research demonstrates that optimizing the value matrix (Wv) in the attention mechanism leads to better performance improvements than optimizing the key matrix (Wk). Moreover, fine-tuning just the query (Wq) and value (Wv) matrices—not only reduces computational demands but also yields results that are on par with or surpass full matrix optimization (including Wq, Wk, Wv).

2. **Efficient Convergence**: The paper highlights the importance of employing distinct learning rates for different matrices, particularly using a higher learning rate for Wv to speed up the convergence process.

The contributions of the paper are both theoretical and practical. Theoretically, it offers a new analysis on how selective fine-tuning of the attention matrices can enhance generalization bounds and improve memory efficiency. Practically, it proposes a fine-tuning strategy that optimizes model training time and resource use. The effectiveness of this approach is supported by experimental results on benchmark datasets, which validate the proposed methods and lay the groundwork for developing more efficient algorithms for fine-tuning LLMs.

strengths:
- **Importance of Understanding Attention Mechanisms During Fine-Tuning**: The challenge of gaining a deeper understanding of the attention mechanism during fine-tuning is a critical one. The approach developed in this paper has the potential to serve as a plug-and-play solution for achieving improved accuracy-efficiency trade-offs in LLM fine-tuning.

- **Empirical and Theoretical Contributions**: This paper offers both empirical and theoretical analyses to elucidate the behavior of the attention mechanism during fine-tuning. The insights provided here could make a valuable contribution to the field, supporting the development of enhanced fine-tuning techniques that optimize the accuracy-efficiency trade-off.

weaknesses:
- **Generalizability of the Proposed Approach**: The primary concern is the generalizability of the proposed approach. Specifically, the authors could enhance the analysis by demonstrating that the proposed method consistently improves LLM performance across diverse scenarios. To this end, it would be beneficial to include performance results under more complex, open-ended generation tasks, such as MT-Bench or comparable challenging benchmarks. Additionally, considering variations in model behavior, evaluating the approach on a broader range of LLMs, such as Mistral, would further strengthen the claim of general applicability.

- **Visualization for Better Insight**: To deepen the analysis and understanding of the findings, visualizing the learned attention distributions across different settings could be valuable. By examining how attention distributions vary under different configurations, the authors could offer a more nuanced understanding of the observed effects, shedding light on the underlying phenomenon.

questions:
- How will the proposed method perform on other PEFT techniques, such as DoRA [1]?

[1] Liu, Shih-Yang, et al. "Dora: Weight-decomposed low-rank adaptation." arXiv preprint arXiv:2402.09353 (2024).
</content>
</thread_object>

</thread>
<thread>
creation date:
1732336298502

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response (part1)

comment:
Dear everyone,

We sincerely appreciate the detailed and constructive feedback provided by all reviewers. We are encouraged to note that the reviewers have recognized:

>- Reviewer SWzh: This paper offers both **empirical and theoretical analyses** to elucidate the behavior of the attention mechanism during fine-tuning. **The insights provided here could make a valuable contribution to the field, supporting the development of enhanced fine-tuning techniques that optimize the accuracy-efficiency trade-off**.
>- Reviewer FTN2: The authors employ **both information-theoretic methods and optimization theory** to support their claims, which significantly strengthens the credibility of their contributions..... This contribution is highly relevant for **both academia and industry**, as it provides a more resource-efficient approach to adapting pre-trained models for downstream tasks.
>- Reviewer f8uV: The idea of using different local learning rates for Wq and Wv is interesting. As V is timed with attention logits, q and v should have more different gradient distributions than q and k. **Using different local learning rates for these two types of matrices is reasonable and worth exploring.**
>- Reviewer TXeN: The observation that fine-tuning only Wq and Wv matrices yields performance gains comparable to finetuning of Wq, Wk and Wv is interesting. **Trying to analyze and uncover the reason can benefit the research commnity**.
>
We have addressed all the questions raised by the reviewers through detailed clarifications, providing separate responses to each reviewer (Following **Reviewer f8uV's suggestion**, we use colored text to show the modifications and improvements in the rebuttal revision). Additionally, we would like to address some common concerns in a consolidated global response.

>**(1) The main contribution of this work.**  
>
>Our primary objective is to provide a general **theoretical framework** that reveals the underlying mechanism behind the  phenomena observed during the fine-tuning of LLMs involving the attention mechanism. Building on our experimental and theoretical insights, **one can develop new algorithms** to improve the effectiveness (e.g., storage, and time) of fine-tuning. 
>
> According to the traditional statistical learning viewpoint, performance can be defined by the sum of **optimization error and generalization error**. Our theoretical analyses in Sections 3 and 4 correspond to generalization and optimization, respectively.
>
>- In Section 3 (generalization), we give **Thereom 1** (Information-theoretic genralization bounds), showing that with the same r value, fine-tuning $W_q,W_v$ consistently achieves results comparable to or even surpassing those of fine-tuning $W_q,W_k,W_v$. This reduces the number of parameters for the same r, while improving generalization bounds and potentially providing memory benefits.  (**storage-friendly**)
>
>- In Section 4 (optimization), we discuss the learning dynamics in fine-tuning attention mechanism, and we illustrate (**Theorem 2**) that the feature learning of attention mechanism is efficient when the learning rate for $W_v$ should be generally much larger than that of $W_q,W_k$ in fine-tuning. (**time-friendly**)
>
>Building on our experimental and theoretical insights, example in Section 5 is presented solely to illustrate how theoretical analysis can guide experimental procedures effectively.

>**(2) How to choose good $\lambda$ values from both theoretical and empirical aspects?** 
>
>Theoretically:
>
>As discussed in **Remark 4**, the optimal ratio $\lambda$ depends on the architecture and the fine-tuning task via the constants in $\Theta$ in **Theorem 2**.  This is a limitation of these asymptotic results since they do not offer any insights on how the constants are affected by the task and the neural architecture.
>
>Empirically:
>
>However, we can still employ some heuristic methods, such as: we can select an appropriate range by conducting a certain amount of experiments, as shown in **Figure 1**, it seems that a ratio of order $2^1-2^4$ is optimal. （That is why we do a simple search for $\lambda$ (2,4,8) in Section 5.）
>
>Moreover, $\lambda$ should not be too large; otherwise, as shown in the MNLI subplot in Figure 1, the model's performance will collapse ($\eta_{QK}=4e^{-4},\eta_V=2e^{-3}$).
</content>
</thread_object>

</thread>


Metareview:
metareview:
While the paper addresses an important topic—fine-tuning strategies for attention mechanisms in large language models—and presents some interesting observations, the contributions appear incremental and insufficiently distinct from prior work. The theoretical insights provided lack intuitive explanations and fail to clearly articulate novel contributions beyond known observations in the literature. Furthermore, the paper's writing is frequently unclear. Additionally, the lack of adequate baselines, ablation studies, and broader evaluations on diverse and challenging benchmarks limits the generalizability of the claims.