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

## Abstract

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 observed during the fine-tuning of LLMs, particularly focusing on the attention mechanism: (1) Different Impact, optimizing the 𝐖 v subscript 𝐖 𝑣 \mathbf{W}_{v} bold_W start_POSTSUBSCRIPT italic_v end_POSTSUBSCRIPT matrix significantly improves performance over optimizing the 𝐖 k subscript 𝐖 𝑘 \mathbf{W}_{k} bold_W start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT matrix. Fine-tuning only the 𝐖 q subscript 𝐖 𝑞 \mathbf{W}_{q} bold_W start_POSTSUBSCRIPT italic_q end_POSTSUBSCRIPT and 𝐖 v subscript 𝐖 𝑣 \mathbf{W}_{v} bold_W start_POSTSUBSCRIPT italic_v end_POSTSUBSCRIPT matrices is computationally efficient, delivering results that are comparable to, or even better than, fine-tuning all three matrices 𝐖 q subscript 𝐖 𝑞 \mathbf{W}_{q} bold_W start_POSTSUBSCRIPT italic_q end_POSTSUBSCRIPT , 𝐖 k subscript 𝐖 𝑘 \mathbf{W}_{k} bold_W start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT , and 𝐖 v subscript 𝐖 𝑣 \mathbf{W}_{v} bold_W start_POSTSUBSCRIPT italic_v end_POSTSUBSCRIPT . (2) Efficient Convergence, employing distinct learning rates for these matrices is crucial for optimal performance, with a higher learning rate for the 𝐖 v subscript 𝐖 𝑣 \mathbf{W}_{v} bold_W start_POSTSUBSCRIPT italic_v end_POSTSUBSCRIPT matrix expediting convergence. However, theoretical analyses of these phenomena are still relatively limited. We present a theoretical analysis of these phenomena from two perspectives: (i) Generalization, where we demonstrate that fine-tuning only 𝐖 q subscript 𝐖 𝑞 \mathbf{W}_{q} bold_W start_POSTSUBSCRIPT italic_q end_POSTSUBSCRIPT and 𝐖 v subscript 𝐖 𝑣 \mathbf{W}_{v} bold_W start_POSTSUBSCRIPT italic_v end_POSTSUBSCRIPT improves generalization bounds, enhances memory efficiency, and (ii) Optimization, where we emphasize that the feature learning of the attention mechanism is efficient, particularly when using distinct learning rates for the matrices, which leads to more effective fine-tuning. 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.

## 1 Introduction

Large Language Models (LLMs) are often built on Transformer architectures [] and possess a large number of parameters, enabling them to generalize across a broad range of general tasks [, , , , ]. However, achieving optimal performance on specific tasks typically necessitates fine-tuning these pre-trained models. Despite the formidable capabilities of LLMs, the fine-tuning process is resource-intensive, requiring significant computational power, storage, and time due to the large scale of model parameters involved. Fine-tuning all the parameters of a large language model, known as full fine-tuning, is highly computationally expensive. To reduce
the computational cost, various parameter-efficient fine-tuning (PEFT) methods have been proposed [, , , , ], which only fine-tune a small number of (extra) model parameters.

A fundamental component of transformers is the attention mechanism, particularly the interactions among the query matrix $\mathbf{W}_{q}$, the key matrix $\mathbf{W}_{k}$, and the value matrix $\mathbf{W}_{v}$. During the fine-tuning of LLMs involving the attention mechanism, two interesting phenomena have been observed: (1) Different Impact—optimizing the $\mathbf{W}_{v}$ matrix is pivotal for enhancing performance, significantly more so than adjustments to the $\mathbf{W}_{k}$ matrix, which exhibit limited impact on the outcomes. Additionally, fine-tuning only the $\mathbf{W}_{q}$ and $\mathbf{W}_{v}$ matrices often yields results that are comparable to or surpass those achieved by fine-tuning all three matrices $\mathbf{W}_{q}$, $\mathbf{W}_{k}$, and $\mathbf{W}_{v}$, which also reduces the number of tunable parameters by approximately 1/3, offering computational benefits (Section 3). (2) Efficient Convergence—using the same
learning rate for $\mathbf{W}_{q}\&\mathbf{W}_{k}$ and $\mathbf{W}_{v}$ is not optimal for efficient convergence. In fact, it is essential to apply distinct learning rates for the $\mathbf{W}_{q}$, $\mathbf{W}_{k}$, and $\mathbf{W}_{v}$ components to ensure optimal fine-tuning performance. Specifically, the learning rate for $\mathbf{W}_{v}$ should generally be higher than that for $\mathbf{W}_{q}$ and $\mathbf{W}_{k}$ to facilitate efficient convergence (Section 4).

While certain empirical guidelines, such as the original Low-Rank Adaptation (LoRA) [], explore which weight matrices in transformers are suitable for the application of LoRA, comprehensive theoretical analyses of these phenomena are still limited. This includes aspects such as selecting appropriate weight types for fine-tuning and optimizing learning rate settings. In this paper, we seek to explore the issue from the following two aspects.

- •
Generalization: advantages of fine-tuning $\mathbf{W}_{q}\&\mathbf{W}_{v}$ over $\mathbf{W}_{q},\mathbf{W}_{k},\mathbf{W}_{v}$ together. We perform a thorough theoretical analysis to demonstrate the advantages. To be more specific, we employ information-theoretic approaches [, , , ] to establish the generalization bounds of fine-tuning pre-trained models with attention mechanism (See Theorem 1 for details). This indicates that fine-tuning $\mathbf{W}_{q}\&\mathbf{W}_{v}$ instead of $\mathbf{W}_{q},\mathbf{W}_{k},\mathbf{W}_{v}$ reduces the number of parameters, while improving generalization bounds and potentially providing memory benefits.
- •
Optimization: convergence analysis of attention mechanism with varying learning rate settings. To further investigate the aforementioned phenomena, we examine the optimization process of the attention mechanism. First, we discuss the learning dynamics in transformers in Case 1, suggest that $\mathbf{W}_{v}$ may experience instances of inefficient learning during fine-tuning process for downstream tasks. This naturally leads to the hypothesis that accelerating the learning of $\mathbf{W}_{v}$ in the early stages could potentially induce $\mathbf{W}_{k}$ and $\mathbf{W}_{q}$ to begin learning earlier. Additionally, by using scaling arguments for large width-$n$ networks [, ], we illustrate (Theorem 2) that the feature learning of attention mechanism is efficient when the learning rate for $\mathbf{W}_{v}$ should be generally much larger than that of $\mathbf{W}_{q}\&\mathbf{W}_{k}$ in fine-tuning.

Building on our experimental and theoretical insights, one can develop new algorithms to improve the effectiveness (e.g., storage, and time) of fine-tuning. Experimental results for our strategy (in Section 5) on benchmark datasets [] and open source pre-trained models [, ]
verify that the method can visibly influence fine-tuning efficiency(^1^11The code is available at [https://github.com/chen123CtrlS/LightweightAtt](https://github.com/chen123CtrlS/LightweightAtt)). We do not make direct comparisons with various parameter-efficient fine-tuning methods, as our strategy is primarily intended to demonstrate how theoretical analysis can effectively guide experimental procedures.

## 2 Preliminaries and Background

In this section, we first describe the core components of our study by reviewing some basic notations. The transformer model [] serves as the backbone of most state-of-the-art pre-trained models. For clarity, we briefly outline its key equations, focusing on the self-attention function, as follows.

Self-attention. Given a sequence of $m$ vectors $\mathbf{C}\in\mathbb{R}^{m\times d_{in}}$ over which we would like to perform attention and a query vector $\mathbf{x}\in\mathbb{R}^{d_{in}}$, that is, the input is [$\mathbf{C},\mathbf{x}$] $\in\mathbb{R}^{(m+1)\times d_{in}}$. The conventional attention function can be expressed as(^2^22For simplicity, we focus on the last vector of input in a single-head self-attention. Our analysis is readily generalizable to multi-head self-attention.):

$$ $\text{Attn}(\mathbf{x}\mathbf{W}_{q},\mathbf{C}\mathbf{W}_{k},\mathbf{C} \mathbf{W}_{v})\\ =\text{softmax}\left(\frac{\mathbf{x}\mathbf{W}_{q}\mathbf{W}_{k}^{T}\mathbf{C }^{T}}{\sqrt{d_{out}}}\right)\mathbf{CW}_{v},$ (1) $$

where $\mathbf{W}_{q}$,$\mathbf{W}_{k}$,$\mathbf{W}_{v}\in\mathbb{R}^{d_{in}\times d_{out}}$ are query, key and value (projection) matrices.

A unified framework for parameter-efficient fine-tuning.
Building on the work of [], we consider a unified framework that establishes connections among various parameter-efficient fine-tuning methods. Specifically, we reinterpret these methods as modifications applied to specific hidden states within pre-trained models, the composition function:

$$ $\mathbf{h}\leftarrow l_{1}\mathbf{h}+l_{2}\Delta\mathbf{h},$ (2) $$

where $l_{1},l_{2}$ are coefficients, $\mathbf{h}$ is denoted as the hidden representation to be directly modified and $\Delta\mathbf{h}$ is a modification vector. Additionally, $\mathbf{h}$ and $\mathbf{x}$ can represent the attention output and input respectively. Here, we will present two special cases:

- •
LoRA. LoRA [] injects trainable low-rank matrices into transformer layers to approximate the weight updates. Instead of directly adjusting the full weight matrix $\mathbf{W}\in\mathbb{R}^{d_{in}\times d_{out}}$, LoRA represents its update with a low-rank decomposition $\mathbf{W}+\Delta\mathbf{W}=\mathbf{W}+\mathbf{AB}$, where $\mathbf{A}\in\mathbb{R}^{d_{in}\times r},\mathbf{B}\in\mathbb{R}^{r\times d_{
out}}$ are tunable parameters. For a specific input $\mathbf{x}$, LoRA modifies the projection output $\mathbf{h}$ as (where $s\geq 1$ is a tunable scalar hyperparameter):
$\displaystyle\mathbf{h}\leftarrow\mathbf{h}+s\Delta\mathbf{h},\ \ \ \Delta
\mathbf{h}:=\mathbf{xAB}.$
(3)
- •
Prefix tuning. Prefix tuning [] prepends $r$ tunable prefix vectors to the keys and values of the attention mechanism at every layer. Specifically, two sets of prefix vectors $\mathbf{P}_{k},\mathbf{P}_{v}\in\mathbb{R}^{r\times d_{out}}$ are concatenated with the original key $\mathbf{CW}_{k}$ and value $\mathbf{CW}_{v}$, attention is then applied to the prefixed keys and values as(^3^33Without loss of generalization, we ignore the softmax scaling factor for ease of notation.):
$\displaystyle\mathbf{h}\leftarrow(1-\lambda(\mathbf{x})\mathbf{h}+\lambda(
\mathbf{x})\Delta\mathbf{h},\ \ \ \Delta\mathbf{h}:=\text{softmax}(\mathbf{xW}
_{q}\mathbf{P}_{k}^{T})\mathbf{P}_{v}\triangleq\text{softmax}(\mathbf{xA})
\mathbf{B},$
(4)
where $\lambda(\mathbf{x})=\frac{\sum_{i}\text{exp}(\mathbf{xW}_{q}\mathbf{P}_{k}^{T}
)_{i}}{\sum_{i}\text{exp}(\mathbf{xW}_{q}\mathbf{P}_{k}^{T})_{i}+\sum_{j}\text
{exp}(\mathbf{xW}_{q}\mathbf{W}_{k}^{T}\mathbf{C}^{T})_{j}}$ is a scalar that represents the sum of normalized attention weights on the prefixes. We derive a detailed equivalent form of Prefix tuning to establish its connection with LoRA in Appendix B.1.

###### Remark 1 .

By defining $\mathbf{A}=\mathbf{W}_{q}\mathbf{P}_{k}^{T},\mathbf{B}=\mathbf{P}_{v}$ in Eq.(4), we can establish a connection with LoRA in Eq.(3). Notably, if we replace the softmax attention with linear attention here, the two are equivalent to some extent. Intuitively, in the attention mechanism, $\mathbf{A}\ (\mathbf{W}_{q}\mathbf{P}_{k}^{T})$ is responsible for generating attention scores, while $\mathbf{B}\ (\mathbf{P}_{v})$ utilizes these attention scores to produce the target content. Therefore, during fine-tuning, query, key, and value are likely to exhibit varying degrees of importance. This may also provide theoretical insights for recent works [, ], which empirically observed an asymmetry where the project-down matrix $\mathbf{A}$ is responsible for extracting features from the input, while the project-up matrix $\mathbf{B}$ utilizes these features to generate the desired output in LoRA fine-tuning.

$\Theta$ Notation. For the convergence analysis, we adopt the following notation to describe the asymptotic behavior as the width $n$ increases, similar to those in [, ]. Given sequences $c_{n}\in\mathbb{R}$ and $d_{n}\in\mathbb{R}^{+}$, we write $c_{n}=O(d_{n})$ and $c_{n}=\Omega(d_{n})$ to mean $c_{n}<\kappa d_{n}$ or $c_{n}>\kappa d_{n}$, respectively, for some constant $\kappa>0$. We denote $c_{n}=\Theta(d_{n})$ when both $c_{n}=O(d_{n})$ and $c_{n}=\Omega(d_{n})$ hold, implying that $c_{n}$ and $d_{n}$ grow at comparable rates. For vector sequences $c_{n}=(c_{n}^{i})_{1\leq i\leq k}\in\mathbb{R}^{k}$ (for some $k>0$), we write $c_{n}=O(d_{n})$ when $c_{n}^{i}=O(d_{n}^{i})$ for all $i\in[k]$, and analogous notation applies for other asymptotic bounds. Finally, when the sequence $c_{n}$ is a vector of random variables, convergence is understood to refer to convergence in the second moment (i.e., $L_{2}$ norm).

###### Remark 1 .

## 3 Advantages and Generalization Analysis

This section, we show our first interesting observation (Different Impact) in fine-tuning the attention mechanism and the storage benefit of fine-tuning only $\mathbf{W}_{q}$ and $\mathbf{W}_{v}$ in Section 3.1. Afterwards,
we give a mutual information based generalization bounds of fine-tuning only $\mathbf{W}_{q}$ and $\mathbf{W}_{v}$ in Section 3.2, which provide a better generalization error.

### 3.1 Empirical Advantages of Fine-Tuning only Query,Value Matrices

To explore the Different Impact, we focus our study on adapting only the attention weights for downstream tasks, while freezing the other modules to ensure simplicity and parameter efficiency. Furthermore, we investigate the impact of adapting different types of attention weight matrices in a Transformer, as outlined below. We present our empirical results using LoRA to fine-tune a set of language models (Roberta-base [] and Llama3.1-8b []) across various benchmarks []. Further details on the experimental setup and additional empirical results can be found in Appendix C.1.

**Table 1: Performance comparison across different $r$ values and weight types. To enable a fair comparison, we initialize the weights for all tasks with the original pretrained weights. Test accuracy of Roberta-base (R) and Llama3.1-8b (L) fine-tuning on SST2, QNLI, QQP, MNLI, with sequence length T = 128 and half precision (FP16). All values are averaged over 3 random seeds. The best result is shown in bold, the second best result is shown in $\underline{\text{underline}}$, and the third best result is shown with $\text{double}\ \underline{\underline{\text{underlines}}}$.**
|  | Weight Type | $\mathbf{W}_{q}$ | $\mathbf{W}_{k}$ | $\mathbf{W}_{v}$ | $\mathbf{W}_{q},\mathbf{W}_{k}$ | $\mathbf{W}_{q},\mathbf{W}_{v}$ | $\mathbf{W}_{q},\mathbf{W}_{k},\mathbf{W}_{v}$ |
| --- | --- | --- | --- | --- | --- | --- | --- |
| SST2(R) | $r=4$ | $0.904$ | $0.902$ | $0.913$ | $\underline{\underline{0.919}}$ | $\bm{\underline{0.920}}$ | $\underline{\bm{0.920}}$ |
| $r=8$ | $0.914$ | $0.906$ | $\underline{\underline{0.918}}$ | $0.915$ | $\underline{0.919}$ | $\bm{0.922}$ |  |
| $r=16$ | $0.907$ | $0.905$ | $0.916$ | $\underline{\underline{0.917}}$ | $\underline{0.921}$ | $\bm{0.923}$ |  |
| QNLI(R) | $r=4$ | $0.854$ | $0.835$ | $\underline{\underline{0.878}}$ | $0.866$ | $\bm{0.888}$ | $\underline{0.887}$ |
| $r=8$ | $0.857$ | $0.841$ | $\underline{\underline{0.875}}$ | $0.866$ | $\underline{0.889}$ | $\bm{0.895}$ |  |
| $r=16$ | $0.854$ | $0.840$ | $\underline{\underline{0.875}}$ | $0.867$ | $\underline{\bm{0.890}}$ | $\underline{\bm{0.890}}$ |  |
| QQP(R) | $r=4$ | $0.812$ | $0.804$ | $\underline{\underline{0.828}}$ | $0.823$ | $\underline{0.838}$ | $\bm{0.843}$ |
| $r=8$ | $0.812$ | $0.806$ | $\underline{\underline{0.828}}$ | $0.823$ | $\underline{0.840}$ | $\bm{0.844}$ |  |
| $r=16$ | $0.812$ | $0.804$ | $\underline{\underline{0.831}}$ | $0.823$ | $\underline{0.839}$ | $\bm{0.844}$ |  |
| QQP(L) | $r=8$ | $0.864$ | $0.845$ | $0.865$ | $\underline{\underline{0.866}}$ | $\underline{\bm{0.874}}$ | $\underline{\bm{0.874}}$ |
| $r=16$ | $0.864$ | $0.845$ | $\underline{\underline{0.869}}$ | $0.867$ | $\underline{\bm{0.874}}$ | $\underline{\bm{0.874}}$ |  |
| MNLI(R) | $r=4$ | $0.748$ | $0.733$ | $\underline{\underline{0.807}}$ | $0.772$ | $\underline{0.820}$ | $\bm{0.828}$ |
| $r=8$ | $0.749$ | $0.733$ | $\underline{\underline{0.809}}$ | $0.778$ | $\underline{0.820}$ | $\bm{0.827}$ |  |
| $r=16$ | $0.750$ | $0.734$ | $\underline{\underline{0.810}}$ | $0.780$ | $\underline{0.824}$ | $\bm{0.828}$ |  |
| MNLI(L) | $r=8$ | $0.802$ | $0.660$ | $\underline{\underline{0.862}}$ | $0.814$ | $\underline{\bm{0.871}}$ | $\underline{\bm{0.871}}$ |
| $r=16$ | $0.803$ | $0.663$ | $\underline{\underline{0.863}}$ | $0.815$ | $\underline{\bm{0.871}}$ | $\underline{\bm{0.871}}$ |  |

Table 1 provides a detailed comparison of the impact of fine-tuning different weight matrices ($\mathbf{W}_{q},\mathbf{W}_{k},\mathbf{W}_{v}$) across various rank values $r$ and weight update strategies in LoRA fine-tuning on tasks like SST2, QNLI, QQP, and MNLI. As seen in the table, we can see a clear trend where solely updating the $\mathbf{W}_{v}$ matrix outperforms just learning the $\mathbf{W}_{q},\mathbf{W}_{k}$ matrix. Interestingly, the combination of fine-tuning both $\mathbf{W}_{q}$ and $\mathbf{W}_{v}$ often leads to performance that matches or even exceeds that achieved by fine-tuning all three matrices $\mathbf{W}_{q},\mathbf{W}_{k},$ and $\mathbf{W}_{v}$. This pattern is consistently observed across various tasks and rank values, further emphasizing the importance of these two matrices over $\mathbf{W}_{k}$ during fine-tuning.

Computational benefits. Here, we show that the reduced amount of adapted parameters by (roughly) $1/3$ provides computational gains. The key benefit of parameter-efficient method is to save memory during training, storage and communication []. Fine-tuning $\mathbf{W}_{q}\&\mathbf{W}_{v}$ alone as opposed to both $\mathbf{W}_{q}\&\mathbf{W}_{v}$ and $\mathbf{W}_{k}$ reduces the number of parameters by $1/3$, when the dimensions of $\mathbf{W}_{q}$, $\mathbf{W}_{k}$, and $\mathbf{W}_{v}$ are the same.

### 3.2 Information-Theoretic Generalization Bounds

In the previous part, we establish that the Different Impact among $\mathbf{W}_{q}$, $\mathbf{W}_{k}$, and $\mathbf{W}_{v}$ during fine-tuning. Some studies [, , ] often treat $\mathbf{W}_{q}$ and $\mathbf{W}_{k}$ as a single unit ($\mathbf{W}_{qk}=\mathbf{W}_{q}\mathbf{W}_{k}^{T}$), however, the benefits of fine-tuning $\mathbf{W}_{q}\&\mathbf{W}_{v}$ alone, rather than fine-tuning $\mathbf{W}_{q}\&\mathbf{W}_{v}$, and $\mathbf{W}_{k}$ together, have yet to be further clarified. Therefore, we will further analyze this issue from an information-theoretic generalization perspective.

Recently, information-theoretic generalization bounds [, , , ] have been introduced to analyze the expected generalization error of learning algorithms. A key benefit of these bounds is that they depend not only on the data distribution but also on the specific algorithm, making them an ideal tool for studying the generalization behavior of models trained using particular algorithms.

Generalization error.
We let $\mathcal{Z}=\mathcal{X}\times\mathcal{Y}$ be the instance space and $\mu$ be an unknown distribution on $\mathcal{Z}$, specifying random variable $Z$. Here, $\mathcal{X}$ denotes the feature space and $\mathcal{Y}$ is the label space. Suppose one observes a training set $S_{N}\triangleq(Z_{1},...,Z_{N})\in\mathcal{Z}^{N}$, with $N$ i.i.d. training examples drawn from $\mu$. In the information-theoretic analysis framework, we let $\mathcal{W}$ be the space of hypotheses related to the model, and a stochastic learning algorithm $\mathcal{A}$ which takes the training examples $S_{N}$ as its input and outputs a hypothesis $W\in\mathcal{W}$ according to some conditional distribution $Q_{W|S_{N}}$. Given a loss function $\ell:\mathcal{W}\times\mathcal{Z}\rightarrow\mathbb{R}^{+}$, where $\ell(w,Z)$ measures the “unfitness” or “error” of any $Z\in\mathcal{Z}$ with respect to a hypothesis $w\in\mathcal{W}$. We take $\ell$ as a continuous function and assume that $\ell$ is differentiable almost everywhere with respect to $w$.
The goal of learning is to find a hypothesis $w$ that minimizes the population risk, and for any $w\in\mathcal{W}$,
the population risk is defined as $L_{\mu}(w)\triangleq\mathbb{E}_{Z\sim\mu}[\ell(w,Z)].$ However, since only can partially observe $\mu$ via the sample $S_{N}$, we instead turn to use the empirical risk, defined as $L_{S_{N}}(w)\triangleq\frac{1}{N}\sum_{i=1}^{N}\ell(w,Z_{i})$. Then the expected generalization error of $\mathcal{A}$ is defined as

$$ $\widetilde{\text{error}}(\mathcal{A})\triangleq\mathbb{E}_{W,S_{N}}[L_{\mu}(W) -L_{S_{N}}(W)],$ $$

where the expectation is taken over $(S_{N},W)\sim\mu^{N}\otimes Q_{W|S_{N}}$.

Consider the following variations of fine-tuning algorithms: tuning both $\mathbf{W}_{k}$ and $\mathbf{W}_{q}\&\mathbf{W}_{v}$ matrices (as in classic attention mechanism in fine-tuning), tuning only $\mathbf{W}_{q}\&\mathbf{W}_{v}$:

###### Definition 1 (Fine-tuning algorithms) .

Recalling A unified framework for parameter-efficient fine-tuning, we can model the fine-tuning process of the attention mechanism as $\mathbf{h}+\Delta\mathbf{h}=\mathbf{x}\mathbf{W}+\mathbf{x}\Delta\mathbf{W}$. Let $\mathbf{W}=\{\mathbf{W}_{i}\}_{i=1}^{L}$ be a set of abstract parameter matrices related to a pretrained model, where each $\mathbf{W}_{i}$ is associated with the parameters $\mathbf{W}_{q}^{i},\mathbf{W}_{k}^{i},\mathbf{W}_{v}^{i}$. The indices $1,...,L$ represent the layers of the model where these parameters are to be fine-tuned. Let $\mathcal{I}\subseteq\{1,...,L\}$ denote the subset of layers selected for fine-tuning. Given a fine-tuning training set $S_{N}$, let $r$ denote the chosen lora-rank, and assume each tuned parameter is quantized to $q$ bits. Define the following algorithmic frameworks for selecting an adaptation $\Delta\mathbf{W}=\{\Delta\mathbf{W}_{i}\}_{i=1}^{L}$ (with other details left open to choice).
(1) $\mathcal{A}_{QKV}$: For each $i\in\mathcal{I}$, optimize $\{\mathbf{W}_{q}^{i},\mathbf{W}_{k}^{i},\mathbf{W}_{v}^{i}\}_{i\in\mathcal{I}}$ to fit the data $S_{N}$.
(2) $\mathcal{A}_{QV}$: For each $i\in\mathcal{I}$, optimize $\{\mathbf{W}_{q}^{i},\mathbf{W}_{v}^{i}\}_{i\in\mathcal{I}}$ to fit the data $S_{N}$.

Then we have the following theorem to bound the generalization error using the information-theoretic generalization framework.

###### Theorem 1 (Generalization bounds on adapting $\mathbf{W}_{q}\&\mathbf{W}_{v}$ and/or $\mathbf{W}_{k}$ ) .

Consider the algorithms of Definition 1. Assume the loss $\ell(\mathbf{W},Z)$ is R-subGaussian under $(\Delta\mathbf{W},{{Z}})\sim P_{\Delta\mathbf{W}\mid\mathbf{W}}\times\mu$. Then,

$$ $\displaystyle\widetilde{\text{error}}(\mathcal{A}_{QV})\leq\sqrt{\frac{4R^{2}} {N}qr\sum_{i\in\mathcal{I}}(d_{in}+d_{out})},$ $\displaystyle\widetilde{\text{error}}(\mathcal{A}_{QKV})\leq\sqrt{\frac{6R^{2} }{N}qr\sum_{i\in\mathcal{I}}(d_{in}+d_{out})},$ $$

where $\mathbf{W}_{q}^{i},\mathbf{W}_{k}^{i},\mathbf{W}_{v}^{i}\in\mathbb{R}^{d_{in}
\times d_{out}}$. See Appendix B.2 for a proof.

###### Remark 2 (Discussion of the advantages) .

We can evaluate the empirical risk ($L_{S_{N}}$) by observing the model’s performance on the dataset we have. If the generalization error (Theorem 1) is determined, it is at least possible to estimate the population risk ($L_{\mu}$). This generalization bound increases with the number of parameters being tuned, which grows as a function of $r$ and the dimensions of the parameter matrices. In Table 1, we know that with the same $r$ value, fine-tuning $\mathbf{W}_{q}\&\mathbf{W}_{v}$ consistently achieves results comparable to or even surpassing those of fine-tuning $\mathbf{W}_{q},\mathbf{W}_{k},\mathbf{W}_{v}$. This reduces the number of parameters for the same $r$,
while improving generalization bounds and potentially providing memory benefits.

###### Definition 1 (Fine-tuning algorithms) .

###### Theorem 1 (Generalization bounds on adapting 𝐖 q & 𝐖 v subscript 𝐖 𝑞 subscript 𝐖 𝑣 \mathbf{W}_{q}\&\mathbf{W}_{v} bold_W start_POSTSUBSCRIPT italic_q end_POSTSUBSCRIPT & bold_W start_POSTSUBSCRIPT italic_v end_POSTSUBSCRIPT and/or 𝐖 k subscript 𝐖 𝑘 \mathbf{W}_{k} bold_W start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT ) .

###### Remark 2 (Discussion of the advantages) .

## 4 Convergence Analysis in Optimization

In Section 3, we have already demonstrated the generalization performance of the attention mechanism during fine-tuning. Our focus will now shift toward optimizing convergence efficiency. Some optimization observations have also been reported in previous works [, , ], such as: Li et al. [] provide theoretical analyses of learning dynamics in transformers and observes a roughly two-stage process of self-attention. Meanwhile, He et al. [] empirically show that the attention mechanism, particularly the value vector, stores the largest amount of memories and has the greatest influence during fine-tuning. However, there is not yet a satisfactory explanation for why this phenomenon occurs or how it can be effectively leveraged. In this section, we will explore these questions in more depth.

### 4.1 An Insight into Inefficient Learning for Value Matrix

We first discuss the optimization process of attention mechanism in the following simple case.

###### Case 1 .

Omitting the scale factor for qualitative analysis in Eq.(1), we obtain:

$$ $\text{Attn}(\mathbf{x}\mathbf{W}_{q},\mathbf{C}\mathbf{W}_{k},\mathbf{C} \mathbf{W}_{v})=\text{softmax}\left(\mathbf{x}\mathbf{W}_{q}\mathbf{W}_{k}^{T} \mathbf{C}^{T}\right)\mathbf{C}\mathbf{W}_{v}.$ $$

Intuitively, if $\mathbf{W}_{q},\mathbf{W}_{k},\mathbf{W}_{v}$ are initialized as random matrices close to zero and trained simultaneously, then in the initial step, $\nabla_{\mathbf{W}_{k}}L(\nabla_{\mathbf{W}_{q}}L)$ contains the term $\mathbf{W}_{q}(\mathbf{W}_{k})$, which is close to 0. By contrast, $\nabla_{\mathbf{W}_{v}}L$ contains the softmax-normalized attention weights. Therefore, during the initial steps (in training), $\mathbf{W}_{v}$ intuitively grows at a much faster rate than $\mathbf{W}_{k}(\mathbf{W}_{q})$.

The work of [] empirically exhibits Case 1 with an approximately two-stage phenomenon: (1) In stage 1 (initial steps), the norms of $\mathbf{W}_{k}$ and $\mathbf{W}_{q}$ remain close to zero across all layers, while the norm of
$\mathbf{W}_{v}$ increases significantly, accompanied by rapid changes in its orientation. (2) In stage 2, the norms of $\mathbf{W}_{k}$ and $\mathbf{W}_{q}$ begin to grow significantly, though much later than the $\mathbf{W}_{v}$ matrices. Briefly, in this case, $\mathbf{W}_{v}$ reaches a certain level of learning during training before $\mathbf{W}_{k}$ and $\mathbf{W}_{q}$ begin to learn. This suggests that when fine-tuning the model for downstream tasks, there may also be instances of inefficient learning in $\mathbf{W}_{v}$. Additionally, is there a fine-tuning strategy that could facilitate more effective learning for downstream tasks? For instance, accelerating the learning of $\mathbf{W}_{v}$
in the early stages could potentially induce earlier learning in $\mathbf{W}_{k}$
and $\mathbf{W}_{q}$.

Next, we present the second interesting phenomenon Efficient Convergence. We use the General Language Understanding Evaluation (GLUE, []) to evaluate the fine-tuning performance of different fine-tuning strategies, which consists of several language tasks that evaluate the understanding capabilities of language models. Using LoRA, we fine-tune Roberta-base from the RoBERTa family [] and Llama3.1-8b [] on MNLI, QQP, QNLI, and SST2 tasks with varying
learning rates $(\eta_{QK},\eta_{V})$ to identify the optimal combination. Other empirical details are provided in Appendix C.1. We present our empirical results using LoRA to fine-tune language models, as visualized in the heatmaps (Figure 1 and Figure 2).

Figure: Figure 1: The test accuracy of RoBERTa-base fine-tuning was evaluated over 3 epochs for MNLI, QQP, and QNLI, and 6 epochs for SST-2, with a sequence length $T=128$ and using half-precision (FP16). The LoRA hyperparameters were set to $\alpha=r=8$. All reported values represent the average results across 3 random seeds. We highlight (1) the best overall accuracy and (2) the values where $\eta_{V}/\eta_{QK}=1$. These values are shown in red. For better visualization, when accuracy is lower than a fixed threshold, we set it to threshold.
Refer to caption: x1.png

Figure: Figure 2: Left: The test accuracy of Llama3.1-8b fine-tuning was evaluated over 800 steps for MNLI. Key values like Figure 1 are also shown in red. Right: The training loss over 800 steps for MNLI fine-tuning on Llama3.1-8b, showing comparison between two optimal learning rate $(\eta_{QK},\eta_{V})$ settings in Left: (1) with $\eta_{V}=\eta_{QK}$ (2) with $\eta_{V}>>\eta_{QK}$.
Refer to caption: x5.png

In Figure 1 and Figure 2, we observe that (1) test accuracy consistently reaches its maximum for certain sets of learning rates where $\eta_{QK}<\eta_{V}$, outperforming the standard practice of setting $\eta_{QK}$
and $\eta_{V}$ equal. (2) More interestingly, the gap between the optimal choice of learning rates overall and the optimal choice when $\eta_{QK}=\eta_{V}$ varies across different tasks. This is probably due to the fact that harder task (like MNLI) requires more efficient feature learning. Additionally, we compare two optimal learning rate $(\eta_{QK},\eta_{V})$ settings in Figure 2 (Left), the $\eta_{V}>>\eta_{QK}$ setting has a better convergence than $\eta_{V}=\eta_{QK}$ setting in Figure 2 (Right).

It is also important to note that due to limited computational resources in our experiments, we use a sequence length of $T=128$ and fine-tune for only 3 epochs on MNLI and QQP. Therefore, it is expected that our test accuracies may be lower than those reported by Hu et al. [], where the authors fine-tune RoBERTa-base with a sequence length of $T=512$ (for MNLI) and for more epochs (30 for MNLI). We do not include confidence intervals for clearer visualization, however, the fluctuations remain within acceptable limits. See Figure 2 (Right) for instance. In Appendix C.2, we provide additional results including the training loss.

###### Case 1 .

### 4.2 Convergence Analysis for Learning Rate

It naturally raises the question of why $\eta_{QK}$ and $\eta_{V}$ should be set differently. In practice, most state-of-the-art models have a large width (embedding dimension), making it worthwhile to examine the training dynamics as the width approaches infinity.

Starting with a Toy setting. Revisiting Definition 1, we have $\Delta\mathbf{h}=\text{softmax}(\mathbf{xA})\mathbf{B}$. In the case of a linear attention mechanism, we instead have $\Delta\mathbf{h}=\mathbf{xA}\mathbf{B}$. Then consider the following toy setting

$$ $\displaystyle f(x)=x(W^{*}+a^{T}b),$ $$

where $W^{*}\in\mathbb{R}^{n\times 1}$ are the fixed(^4^44Here, we primarily focus on the case of $\Delta\mathbf{W}$ to provide insightful theoretical results.) pre-trained weights, $b\in\mathbb{R},a\in\mathbb{R}^{1\times n}$ are adaptation weights, $x\in\mathbb{R}^{n}$ is the model input (This corresponds to $r=1$ in Definition 1). The training goal is to minimize the loss $\mathcal{L}(\theta)=\frac{1}{2}(f(x)-y)^{2}$ where $\theta=(a,b)$ and $(x,y)$ is an input-output datapoint(^5^55To simplify the analysis, we assume that the fine-tuning dataset consists of a single sample, though our analysis can be easily generalized to multiple samples. All conclusions remain essentially valid when $(a,b)$ are matrices.). Similar to LoRA, we generally aim to initialize the product $a^{T}b$ to zero, ensuring that fine-tuning starts from the pre-trained model. This requires at least one of the weights, $a\ (\text{related to}\mathbf{W}_{q}\&\mathbf{W}_{k})$ or $b\ (\text{related to}\mathbf{W}_{v})$, to be initialized to zero. If both are initialized to zero, $\mathbf{W}_{q}\&\mathbf{W}_{k}$ learning cannot occur efficiently in init steps, as discussed in Section 4.1 (More detailed initialization settings are shown in Appendix B.3).

And we assume that $x=\Theta(1)$, meaning that the input coordinates remain of the same order as the width increases. In the subsequent analysis, we examine how the fine-tuning dynamics evolve as the model width
$n$ increases.

To streamline the analysis, we assume $W^{*}=0$, a common simplification that can be applied without loss of generality. This assumption is implemented by setting $\hat{y}=y-xW^{*}$. We denote the fine-tuning step by using subscript $t$. Let $U_{t}=f_{t}(x)-y$, the gradients are then computed as:

$$ $\displaystyle\frac{\partial\mathcal{L}}{\partial a_{t}}=xU_{t}b_{t},\ \ \ \frac{\partial\mathcal{L}}{\partial b_{t}}=xa_{t}^{T}U_{t}.$ $$

And at step $t$ with learning rate $\eta_{a},\eta_{b}>0$, we have

$$ $\displaystyle\Delta f_{t}$ $\displaystyle\triangleq f_{t}(x)-f_{t-1}(x)$ $\displaystyle=-\underbrace{\eta_{a}||x||^{2}U_{t-1}b_{t-1}^{2}}_{\delta_{t}^{1 }}-\underbrace{\eta_{b}(xa_{t-1}^{T})^{2}U_{t-1}}_{\delta_{t}^{2}}+\underbrace {\eta_{a}\eta_{b}||x||^{2}(xa_{t-1}^{T})U_{t-1}^{2}b_{t-1}}_{\delta_{t}^{3}}.$ $$

###### Remark 3 .

The output update is influenced by three key terms. The first two items $\delta_{t}^{1},\delta_{t}^{2}$ (order one in $\eta_{a}/\eta_{b}$) represent linear contributions to the update, meaning they result from changes in the model output when either
$a$ is updated with $b$ held constant, or vice versa. The last item $\delta_{t}^{3}$ (order two in $\eta_{a}\eta_{b}$) corresponds to a multiplicative update that captures the combined effects of changes in both $a$ and $b$. As we scale the width(^6^66This property is generally satisfied in practice when the model width is large (e.g., $n\approx 800$ for Roberta-base and $n\approx 4000$ for Llama3.1-8b).), the desirable feature updates are such that $\Delta f_{t}=\Theta(1)$, ensuring they remain unaffected by this scaling (the updates do not explode with width, see x for more details). Ideally, we aim for both $\delta_{t}^{1}$ and $\delta_{t}^{2}$ to be $\Theta(1)$. If this condition isn’t met, it indicates that either $a$ or $b$ is not being updated efficiently. For example, if $\delta_{t}^{1}=o(1)$, it suggests that as $n\to\infty$, the model behaves as if $a$ is essentially fixed, with only $b$ being trained. We say that the feature learning in the attention mechanism is efficient when $\delta_{t}^{i}=\Theta(1)$ for $i\in\{1,2\}$ and all $t>1$, it means
that both $a$ and $b$ parameter updates significantly contribute to the change in $f_{t}(x)$. We will see that when both $\delta_{t}^{1}$ and $\delta_{t}^{2}$ are $\Theta(1)$, the term $\delta_{t}^{3}$ is also $\Theta(1)$.

Let us assume that we train the model with gradient descent with learning rate $\eta_{a}=\Theta(n^{c_{a}}),\eta_{b}=\Theta(n^{c_{b}})$ for some $c_{a},c_{b}\in\mathbb{R}$. In the study by Yang et al. [], it is noted that the training dynamics primarily involve operations such as matrix-vector products and the summation of vectors or scalars. Given the nature of these operations, it is easy to see that any quantity
in the training dynamics should be of order $n^{\gamma}$ for some $\gamma\in\mathbb{R}$. We write $v=\Theta(n^{\gamma[v]})$, for any quantity $v$ in the training dynamics. When $v$ is a vector, we use the same notation when all entries of $v$ are $\Theta(n^{\gamma[v]})$ (See Appendix B.4 for the formal definition of $\gamma$).

With reference to the method of Hayou et al. [], we start from the initialization in Starting with a Toy setting, we have $f_{0}(x)=0$. Feature learning of attention mechanism is efficient when $\delta_{t}^{i}=\Theta(1)$ for $i\in\{1,2\}$ and all $t>1$, and $f_{t}(x)=\Theta(1)$ for $t>1$. This can be interpreted as:

$$ $\displaystyle\left\{\begin{array}[]{l}c_{a}+1+2\gamma[b_{t-1}]=0\quad\left( \delta_{t}^{1}=\Theta(1)\right)\\ c_{b}+2\gamma[xa_{t-1}^{\top}]=0\quad\left(\delta_{t}^{2}=\Theta(1)\right)\\ \gamma[xa_{t-1}^{\top}]+\gamma[b_{t-1}]=0\quad\left(f_{t-1}(x)=\Theta(1)\right ),\end{array}\right.$ $$

which, after simple calculations, implies that $c_{a}+c_{b}=-1$. Notice that the above also leads to the $c_{a}+c_{b}+1+\gamma[xa_{t-1}^{\top}]+\gamma[b_{t-1}]=0\ (\delta_{t}^{3}=
\Theta(1))$. This is only a necessary condition. In the following section, we will provide theoretical conclusions in the toy model setting that offer guidance for real-world experiments.

###### Theorem 2 (Efficient fine-tuning in attention mechanism (Informal)) .

In the case of Starting with a Toy setting, with $\eta_{a}=\Theta(n^{-1})$ and $\eta_{b}=\Theta(1)$, we have for all $t>1$, $i\in\{1,2,3\}$,$\delta_{t}^{i}=\Theta(1)$. In other words, the feature learning of attention mechanism is efficient when $\eta_{QK}(\eta_{a})=\Theta(n^{-1}),\eta_{V}(\eta_{b})=\Theta(1)$. We denote $\eta_{V}/\eta_{QK}$ as $\lambda$.
We refer the reader to Appendix B.5 for more details on the proof.

###### Remark 4 .

In practice, Theorem 2 implies that the learning rate for $\mathbf{W}_{v}$ should be generally much larger than that of $\mathbf{W}_{q}\&\mathbf{W}_{k}$ in fine-tuning. We verify that this scaling is valid for general neural network models in Section 4.1. Naturally, the optimal ratio $\lambda$ depends on the architecture and the fine-tuning task through the constants in ‘$\Theta$’. This represents a limitation of the asymptotic results, as they do not provide insights into how the task and neural architecture influence these constants. We will further address this issue in our future work.

###### Remark 3 .

###### Theorem 2 (Efficient fine-tuning in attention mechanism (Informal)) .

###### Remark 4 .

## 5 An Example of Improving Fine-tuning Methods

Based on all our exciting insights, it becomes intuitive to design lightweight attention-based fine-tuning improvements, particularly for downstream tasks. To illustrate how theoretical analysis effectively guides experimental procedures, we propose an example method where we freeze the $\mathbf{W}_{k}$ and fine-tuning the $\mathbf{W}_{q}\&\mathbf{W}_{v}$ using different learning rates. This procedure is reported in Figure 5.

Experimental setup. We conduct experiments on widely adopted benchmark datasets [] and Roberta-base model []. We selected two mainstream baselines: Full Fine-tuning and LoRA. Additionally, we adapt only the attention weights for downstream tasks, keeping the other modules frozen to maintain simplicity and validate the theoretical guidance through experiments. Moreover, as discussed in Remark 4, the optimal ratio $\lambda$ depends on the architecture and the fine-tuning task. In our experiments, we evaluated the performance for $\lambda$ values of 2, 4, and 8 (one can also determine a general optimal ratio through experiments, and even apply different settings across different layers of the model). We report the average results based on 3 random
seeds, as shown in Table 2. The hyperparameter settings for the experiments
can be found in Appendix C.1.2.

**Table 2: Comparison of fine-tuning methods across GLUE benchmark. We report results on development set, Pearson correlation for STS-B, Matthew’s correlation for CoLA, average accuracy for MNLI (matched and mismatched), and accuracy for other tasks. The best results on each dataset are shown in bold and the second best results are shown in underline. The QKV(QV) setting refers to fine-tuning $\mathbf{W}_{q},\mathbf{W}_{k},\mathbf{W}_{v}(\mathbf{W}_{q},\mathbf{W}_{v})$. It is noted that the total number of parameters in the Roberta-base model is 124.65M. $\lambda$ means $\eta_{V}=\lambda\eta_{Q}$ and $r$ is the LoRA rank, and a larger $\lambda$ does not necessarily lead to better performance.**
| Method | #Param (M) | RTE | STS-B | MRPC | CoLA | MNLI | SST-2 | QQP | QNLI |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Full Fine-tune (QKV) | 21.85 | 73.64 | 90.49 | 84.55 | 60.34 | 86.68 | 93.23 | 90.48 | 92.37 |
| LoRA (QKV) $r=8$ | 1.62 | 70.76 | 90.25 | 85.04 | 58.03 | 86.70 | 93.92 | 89.15 | 92.17 |
| LoRA (QKV) $r=16$ | 2.07 | 70.39 | 90.25 | 86.03 | 58.04 | 86.78 | 93.92 | 89.26 | 92.18 |
| Full Fine-tune (QV) $\lambda=2$ | 14.76 | 73.53 | 91.01 | 86.02 | 60.57 | 62.03 | 93.11 | 90.56 | 91.96 |
| Full Fine-tune (QV) $\lambda=4$ | 14.76 | 72.29 | 90.56 | 87.01 | 61.88 | 35.44 | 91.05 | 89.81 | 88.85 |
| Full Fine-tune (QV) $\lambda=8$ | 14.76 | 72.29 | 90.02 | 88.97 | 61.86 | 35.44 | 84.75 | 85.93 | 50.54 |
| LoRA (QV) $r=8,\lambda=2$ | 1.48 | 71.84 | 90.37 | 86.02 | 58.54 | 86.85 | 94.03 | 89.47 | 92.33 |
| LoRA (QV) $r=8,\lambda=4$ | 1.48 | 75.09 | 90.83 | 87.01 | 59.56 | 86.95 | 94.04 | 90.09 | 92.86 |
| LoRA (QV) $r=8,\lambda=8$ | 1.48 | 76.13 | 90.75 | 88.97 | 61.88 | 86.93 | 93.46 | 90.01 | 92.34 |
| LoRA (QV) $r=16,\lambda=2$ | 1.77 | 70.39 | 90.46 | 86.03 | 58.55 | 86.83 | 94.38 | 89.77 | 92.33 |
| LoRA (QV) $r=16,\lambda=4$ | 1.77 | 76.17 | 91.05 | 87.99 | 60.06 | 87.19 | 94.03 | 90.30 | 92.73 |
| LoRA (QV) $r=16,\lambda=8$ | 1.77 | 72.92 | 90.96 | 89.95 | 59.31 | 87.31 | 93.92 | 90.43 | 92.95 |

Results. We leverage our theoretical results (Theorem 1 and Theorem 2) to enhance the efficiency of existing fine-tuning methods, such as Full Fine-tune and LoRA, on downstream tasks. As shown in Table 2, the improved fine-tuning approach not only outperforms the original version but also significantly reduces the number of parameters. For instance, on the MRPC task, LoRA (QV) $r=16,\lambda=8$ (1.77M) achieves better performance compared to Full Fine-tune (QKV) (21.85M) and LoRA (QKV) $r=16$ (2.07M). This series of experiments clearly demonstrates that our theoretical insights effectively enhance fine-tuning algorithms, particularly in terms of memory usage and optimization efficiency. Moreover, these theoretical results can guide the improvement of other fine-tuning algorithms and even aid in the design of more efficient ones.

## 6 Conclusion and Limitation

In this paper, we present our key findings in fine-tuning attention mechanism: Different Impact—optimizing the $\mathbf{W}_{v}$ matrix significantly improves performance compared to the $\mathbf{W}_{k}$ matrix. Fine-tuning only the $\mathbf{W}_{q}$ and $\mathbf{W}_{v}$ matrices is computationally efficient and can yield results that match or surpass fine-tuning all three matrices $\mathbf{W}_{q}$, $\mathbf{W}_{k}$, and $\mathbf{W}_{v}$. Efficient Convergence—using distinct learning rates for these matrices is essential for optimal performance, with a higher learning rate for $\mathbf{W}_{v}$ speeding up convergence. While theoretical analysis of these phenomena is limited, this paper provides insights from two angles: Generalization—fine-tuning only $\mathbf{W}_{q}$ and $\mathbf{W}_{v}$ improves generalization and memory efficiency, and Optimization—using different learning rates enhances the efficiency of feature learning in the attention mechanism, leading to more effective fine-tuning. Our analysis provides a theoretical foundation for the configuration and improvement of lightweight algorithms in LLMs fine-tuning. However, further studies are required on (i) how the task and neural architecture influence the optimal ratio $\lambda$ and (ii) whether the results about attention hold true for tasks beyond natural language processing. These studies will further deepen our understanding of attention-based fine-tuning in LLMs.

## Appendix A More Related Works

Attention mechanism analysis.
A key component of transformers is the attention mechanism, which dates back to []. Initially designed to capture long-range signals in sequential inputs by mixing individual tokens, it has also been utilized to capture general structures in input data. After the fully-attention-based language
model has appeared [, ], the research community gets interested in the functionality and benefits of the attention. For instance, transformers implicitly favor hierarchical interpretations of input sequences [], the computational graphs tend to be tree-structured [, ]. Theoretical analysis of training dynamics sheds light on how to identify key tokens [], select a few relevant tokens only (which is called localized attention) or select many tokens uniformly [], and learn topic structure []. Besides, considerable works [, , ] try to understand in-context learning capabilities from the perspective of gradient descent with attention.

Scaling for neural networks. Scaling refers to the process of enlarging a specific ingredient of a model to enhance its overall performance []. The method is straightforward: extend the width or depth of a neural network towards infinity, analyze how this limit is influenced by hyperparameters like the learning rate and initialization variance during training, and then establish well-founded choices for these hyperparameters to achieve a specific objective [, , , , , , , , ]. In the theory of scaling of neural networks, one usually tracks the asymptotic behaviour of key quantities as we scale some model ingredient, it is a standard approach used to derive scaling rules for initialization [], activation function [], network parametrization []. In this paper, we are interested in scaling model capacity via the width n for the fact that most state-of-the-art pre-trained models have large width. Examples of the infinite-width limit can be found in studies focused on initialization methods [, ], or more comprehensive approaches to network parameterization. For instance, Yang et al. [] introduced µP, a parameterization technique for neural networks that guarantees feature learning in the infinite-width limit, providing specific scaling rules for both architecture and learning rates to optimize feature learning [, ].

Parameter-efficient fine-tuning. Fine-tuning all the parameters of a large language models, known as full fine-tuning, is highly computationally expensive. To reduce
the computational cost, various parameter-efficient fine-tuning (PEFT) methods have been proposed [], which only fine-tune a small number of (extra) model
parameters. PEFT methods can be divided into two categories from the perspective of whether extra parameters are involved: (1) extra-parameter methods, freeze all of the original parameters of an LLM and insert a set of learnable parameters to optimize the model input or model layers suach as adapter tuning [] , prompt tuning
[] and prefix tuning []; (2) intra-parameter methods freeze most of the original parameters of an LLM and only tune a small number of parameters of the LLM such as LoRA []. Furthermore, He et al. [] present a unified framework that
establishes connections between PEFT methods and Zhu et al. [] formally identify and investigate asymmetry in the roles of low-rank adapter matrices in
LoRA fine-tuning.

## Appendix B Omitted Proofs and Additional Results

### B.1 The connection between Prefix tuning and LoRA.

Here, we provide an alternative view of Prefix tuning (without loss of generalization, we ignore the softmax scaling factor for ease of notation):

$$ $\displaystyle\text{Attn}(\mathbf{x}\mathbf{W}_{q},\text{concat}(\mathbf{P}_{k} ,\mathbf{C}\mathbf{W}_{k}),\text{concat}(\mathbf{P}_{v},\mathbf{CW}_{v})$ $\displaystyle=\text{softmax}(\mathbf{x}\mathbf{W}_{q}\text{concat}(\mathbf{P}_ {k},\mathbf{C}\mathbf{W}_{k})^{T})\begin{pmatrix}\mathbf{P}_{v}\\ \mathbf{CW}_{v}\\ \end{pmatrix}$ $\displaystyle=(1-\lambda(\mathbf{x}))\text{softmax}(\mathbf{x}\mathbf{W}_{q} \mathbf{W}_{k}^{T}\mathbf{C}^{T})\mathbf{C}\mathbf{W}_{v}+\lambda(\mathbf{x}) \text{softmax}(\mathbf{x}\mathbf{W}_{q}\mathbf{P}_{k}^{T})\mathbf{P}_{v}$ $\displaystyle=(1-\lambda(\mathbf{x}))\overbrace{\text{Attn}(\mathbf{x}\mathbf{ W}_{q},\mathbf{CW}_{k},\mathbf{CW}_{v})}^{{\color[rgb]{0,0,1}\text{standard attention}}}+\lambda(\mathbf{x})\overbrace{\text{Attn}(\mathbf{x}\mathbf{W}_{q },\mathbf{P}_{k},\mathbf{P}_{v})}^{{\color[rgb]{0.0,0.5,0.0}\text{independent of C}}},$ $$

where $\lambda(\mathbf{x})=\frac{\sum_{i}\text{exp}(\mathbf{xW}_{q}\mathbf{P}_{k}^{T}
)_{i}}{\sum_{i}\text{exp}(\mathbf{xW}_{q}\mathbf{P}_{k}^{T})_{i}+\sum_{j}\text
{exp}(\mathbf{xW}_{q}\mathbf{W}_{k}^{T}\mathbf{C}^{T})_{j}}$ is a scalar that represents the sum of normalized attention weights on the prefixes. Notice that the first term in blue represents the original attention mechanism without prefixes, while the second term in green introduces a position-wise adjustment independent of $\mathbf{C}$. It provides an alternative perspective on Prefix tuning, where a position-wise modification is applied to the original attention output $\mathbf{h}$ via linear interpolation:

$$ $\displaystyle\mathbf{h}\leftarrow(1-\lambda(\mathbf{x})\mathbf{h}+\lambda( \mathbf{x})\Delta\mathbf{h},\ \ \ \Delta\mathbf{h}:=\text{softmax}(\mathbf{xW} _{q}\mathbf{P}_{k}^{T})\mathbf{P}_{v}\triangleq\text{softmax}(\mathbf{xA}) \mathbf{B}.$ $$

### B.2 Proof of Theorem 1

The origin form of the mutual information based bound is predicated on a sample-specific MI, which quantifies the shared information between the output variable $\mathbf{W}$ and the input sample set $S_{N}$. The following lemma shows the result:

###### Lemma 1 .

(Xu and Raginsky [, Theorem 1.]). Assume the loss $\ell(\mathbf{W},Z)$ is R-subGaussian for any $\mathbf{W}\in\mathcal{W}$, then

$$ $\widetilde{\text{error}}(\mathcal{A})\leq\sqrt{\frac{2R^{2}}{N}I(\mathbf{W};S_ {N})},$ $$

where $I(\mathbf{W};S_{N})=D_{KL}(Q_{\mathbf{W},S_{N}}\|Q_{\mathbf{W}}\otimes Q_{S_{N
}})$ is the mutual information and $D_{KL}$ denotes the $KL$ divergence.

Unroll the terminal parameters’ mutual information $I(\mathbf{W};S_{N})$ to the full trajectories’ mutual information will get:

###### Lemma 2 .

Let Definition 1 hold, then $I(\mathbf{W}+\Delta\mathbf{W};S_{N}|\mathcal{A})\leq I(\Delta\mathbf{W};S_{N}|
\mathcal{A},\mathbf{W})$.

###### Proof.

$$ $\displaystyle I(\mathbf{W}+\Delta\mathbf{W};S_{N}|\mathcal{A})$ $\displaystyle\leq I(\mathbf{W},\Delta\mathbf{W};S_{N}|\mathcal{A})$ (*) $\displaystyle=I(\mathbf{W};S_{N}|\mathcal{A})+I(\Delta\mathbf{W};S_{N}| \mathcal{A},\mathbf{W})$ (**) $\displaystyle=I(\Delta\mathbf{W};S_{N}|\mathcal{A},\mathbf{W}).$ $$

where Eq. (*) is by the data processing inequality (e.g., $Z-(X,Y)-(X+Y)$ form a Markov chain then $I(X+Y,Z)\leq I(X,Y;Z)$), Eq. (**) is by the chain rule of the mutual information, and $I(\mathbf{W};S_{N})=0$ for $\mathbf{W}$ is independent of $S_{N}$.
∎

Then combine Lemma 1 and Lemma 2 , we can get:
$\widetilde{\text{error}}(\mathcal{A})\leq\sqrt{\frac{2R^{2}}{N}I(\Delta\mathbf
{W};S_{N}|\mathcal{A},\mathbf{W})}.$
We consider the case of tuning $\mathbf{W}_{q}\&\mathbf{W}_{v}$ only first. Applying the above results, note that here

$$ $\displaystyle I(\Delta\mathbf{W};S_{N}|\mathcal{A}_{QV},\mathbf{W})$ $\displaystyle=I(\{\mathbf{W}_{q}^{i},\mathbf{W}_{v}^{i}\}_{i\in\mathcal{I}};S_ {N}|\mathcal{A}_{QV},\mathbf{W}),$ $$

where we have used the data processing inequality (DPI), noting that the $\mathbf{W}_{k}^{i}$ are here considered fixed constant matrices as they are not trained.

We can now bound this expression as

$$ $\displaystyle I(\{\mathbf{W}_{q}^{i},\mathbf{W}_{v}^{i}\}_{i\in\mathcal{I}};S_ {N}|\mathcal{A}_{QV},\mathbf{W})\leq H(\{\mathbf{W}_{q}^{i},\mathbf{W}_{v}^{i} \}_{i\in\mathcal{I}})\leq 2qr\sum_{i\in\mathcal{I}}(d^{i}+k^{i}),$ $$

where $\mathbf{W}_{q}^{i},\mathbf{W}_{k}^{i},\mathbf{W}_{v}^{i}\in\mathbb{R}^{d_{in}
\times d_{out}}$, and we have noted that mutual information is upper bounded by discrete entropy, and entropy in turn is upper
bounded by the uniform distribution over its possible support set (q bits in each of $r\sum_{i\in\mathcal{I}}(d_{in}+d_{out})$ dimensions). The bounds for the other algorithms are similar.

###### Lemma 1 .

###### Lemma 2 .

###### Proof.

### B.3 Initialization Discussion

Following standard initialization schemes (e.g., LeCun Init and He Init [, ]), one generally consider a Gaussian initialization of the weights as follows: $a_{i}\sim\mathcal{N}(0,\sigma_{a}^{2}),\quad b\sim\mathcal{N}(0,\sigma_{b}^{2})$ (The Gaussian distribution can be substituted with any other distribution that has finite variance). Revisiting Starting with a Toy setting, $a\in\mathbb{R}^{1\times n},b\in\mathbb{R}$. Thus, one should set $\sigma_{a}^{2}=\Theta(n^{-1}),\sigma_{b}^{2}=0$ to ensure $xa^{T}$ does not explode with width ($xa^{T}=\Theta(1)$), for a non-zero
initialization for $a$. This is justified by
the Central Limit Theorem (See [] for more technical details). And if we choose
a non-zero initialization for $b$, one should make sure that
$\sigma_{b}^{2}=\Theta(1),\sigma_{a}^{2}=0$. And we will consider these two initialization schemes to show our theoretical understanding.

### B.4 Gamma Function

Why introduce the Gamma function?
In Section 4.2, the learning rate $\eta_{a}=\Theta(n^{c_{a}}),\eta_{b}=\Theta(n^{c_{b}})$ for some $c_{a},c_{b}\in\mathbb{R}$. And in Appendix B.3 we assume that the init weights are also scale polynomially with n, it is evident that preactivations, gradients, and weight updates all exhibit asymptotic polynomial growth in n.
Operations.
We write $v=\Theta(\gamma[v])$ to capture it, and some elementary operations (Given two real-valued variables $v_{1},v_{2}$):

- •
Multiplication. $\gamma[v_{1}\times v_{2}]=\gamma[v_{1}]+\gamma[v_{2}].$
- •
Addition. Generally, we have $\gamma[v_{1}+v_{2}]=\max(\gamma[v_{1}],\gamma[v_{2}])$. The only instance where this does not hold is when $v_{1}=-v_{2}$. This is typically a zero-probability event if the two variables are random variables that are not perfectly correlated, which is the case in most scenarios where we apply this formula (Appendix B.5).

### B.5 Proof of Theorem 2

Theorem 2.[Efficient fine-tuning in attention mechanism (Informal)]
In the case of Starting with a Toy setting, with $\eta_{a}=\Theta(n^{-1})$ and $\eta_{b}=\Theta(1)$, we have for all $t>1$, $i\in\{1,2,3\}$,$\delta_{t}^{i}=\Theta(1)$. In other words, the feature learning of attention mechanism is efficient when $\eta_{QK}(\eta_{a})=\Theta(n^{-1}),\eta_{V}(\eta_{b})=\Theta(1)$.

###### Proof.

In Section 3, we say that the
feature learning of attention mechanism is efficient when $\delta_{t}^{i}=\Theta(1)$ for all $t,i\in\{1,2,3\}$. Using the elementary formulas from Appendix B.4, we can get (for all $t$):

$$ $\displaystyle\left\{\begin{array}[]{l}\gamma[\eta_{a}]+1+2\gamma[b_{t-1}]=0 \quad\left(\delta_{t}^{1}=\Theta(1)\right)\\ \gamma[\eta_{b}]+2\gamma[xa_{t-1}^{\top}]=0\quad\left(\delta_{t}^{2}=\Theta(1) \right)\\ \gamma[\eta_{a}]+\gamma[\eta_{b}]+1+\gamma[xa_{t-1}^{\top}]+\gamma[b_{t-1}]=0 \quad\left(\delta_{t}^{3}=\Theta(1)\right).\end{array}\right.$ $$

Simple calculations yield $\gamma[\eta_{a}]+\gamma[\eta_{b}]=-1$. Further consider the gradient update from $t-1$ to $t$, the recursive formulas are given by:

$$ $\displaystyle\left\{\begin{array}[]{l}\gamma[xa_{t}^{\top}]=\max\left(\gamma[ xa_{t-1}^{\top}],\gamma[\eta_{a}]+1+\gamma[b_{t-1}]\right)\\ \gamma[b_{t}]=\max\left(\gamma[b_{t-1}],\gamma[\eta_{b}]+\gamma[xa_{t-1}^{\top }]\right)\end{array}\right.$ $$

Starting from $t=1$. In both initialization schemes discussed in Appendix B.3, we have to set $\gamma[\eta_{b}]=0$ and $\gamma[\eta_{a}]=-1$ to ensure that $\gamma[f_{t}]=\gamma[xa_{t}^{T}]+\gamma[b_{t}]=0$:
(1) $\sigma_{a}^{2}=\Theta(n^{-1}),\sigma_{b}^{2}=0$. We have $\gamma[xa_{1}^{T}]=\gamma[xa_{0}^{T}]=0,\ \gamma[b_{1}]=\gamma[\eta_{b}(xa_{0}
^{T})y]=\gamma[\eta_{b}]$. Therefore, for $t=2$, $\gamma[xa_{2}^{T}]=\max(0,\gamma[\eta_{a}]+1+\gamma[\eta_{b}])=\max(0,0)=0,\
\gamma[b_{2}]=\max(\gamma[\eta_{b}],\gamma[\eta_{b}]+0)=\gamma[\eta_{b}]$, this holds for $t\geq 1$ by induction.
(2) $\sigma_{a}^{2}=0,\sigma_{b}^{2}=\Theta(1)$. We have $\gamma[b_{1}]=\gamma[b_{0}]=0,\ \gamma[xa_{1}^{T}]=\gamma[\eta_{a}||x||^{2}U_{
0}b_{0}^{2}]=\gamma[\eta_{a}]+1$. Therefore, for $t=2$, $\gamma[b_{2}]=\max(0,\gamma[\eta_{b}]+\gamma[\eta_{a}]+1)=\max(0,0)$,$\gamma[xa_{2}^{T}]=\max(\gamma[\eta_{a}]+1,\gamma[\eta_{a}]+1+0)=\gamma[\eta_{
a}]+1$,this holds for $t\geq 1$ by induction.
To sum up, setting $\eta_{QK}(\eta_{a})=\Theta(n^{-1}),\eta_{V}(\eta_{b})=\Theta(1)$ ensures efficient fine-tuning in attention mechanism.
∎

###### Proof.

## Appendix C Extension to Experiments

### C.1 Empirical Details

#### C.1.1 GLUE Tasks with RoBERTa

For our experiments with Roberta-base models, finetuned on GLUE tasks, we use the following setup:

Tasks. MNLI, QQP, SST2, QNLI

Training Algorithm. AdamW with $\beta_{1}=0.9,\beta_{2}=0.99,\epsilon=1e-8$, linear schedule, no warmup.

Targert Modules for Fine-tuning. ‘query’, ‘key’ and ‘value’.

Learning rate.
(1) For Table 1, $\eta_{QK}=\eta_{V}=5e^{-5}$.
(2) For Figure 1,
$\eta_{QK}=\{2e^{-5},5e^{-5},1e^{-4},2e^{-4},4e^{-4},8e^{-4}\}$,
$\eta_{V}=\{1e^{-4},2e^{-4},4e^{-4},8e^{-4},1e^{-3},2e^{-3}\}$

GPUs. Nvidia A800.

Other Hyperparameters. Sequence length $T=128$, train batch size $batchsize=32$, number of train , number of random seeds $s=3$.
(1) For Table 1, epochs E = 6 (E = 10
for SST2).
(2) For Figure 1, epochs E = 3 (E = 6
for SST2).

#### C.1.2 Training Hyperparameters

Training hyperparameters.

**Table 3: Training hyperparameters for different datasets. More details can be seen in our code.**
| Corpus | length | learning rate | batch size | epochs |
| --- | --- | --- | --- | --- |
| RTE | 128 | 1e-04 | 32 | 20 |
| MRPC | 128 | 1e-04 | 32 | 20 |
| STS-B | 128 | 1e-04 | 32 | 20 |
| CoLA | 128 | 1e-04 | 32 | 20 |
| SST-2 | 128 | 1e-04 | 32 | 10 |
| QNLI | 128 | 1e-04 | 32 | 10 |
| QQP | 128 | 1e-04 | 32 | 10 |
| MNLI | 128 | 1e-04 | 32 | 10 |

- •
For Full Fine-tune (QKV) and LoRA (QKV), we use $\eta_{Q}=\eta_{K}=\eta_{V}=\text{1e-04}$.
- •
For the improved methods, we use $\eta_{Q}=\text{1e-04}$, $\eta_{V}=\lambda\times\text{1e-04}$.

The hyperparameter settings here differ from those in Table 1 and Figure 1, so the results may show slight variations.

#### C.1.3 MNLI Task with Llama3.1-8b

For our experiments with Llama3.1-8b models, finetuned on MNLI, we use the following setup:

Training Algorithm. AdamW with $\beta_{1}=0.9,\beta_{2}=0.999,\epsilon=1e-6$, constant schedule.

Targert Modules for Fine-tuning. ‘q_proj, k_proj, v_proj’.

Learning rate grid.
For Table 1, $\eta_{QK}=\eta_{V}=1e^{-5}$.
Else:
$\eta_{QK}=\{1e^{-6},5e^{-6},1e^{-5},5e^{-5},1e^{-4}\}$,
$\eta_{V}=\{1e^{-6},5e^{-6},1e^{-5},5e^{-5},1e^{-4},2e^{-4},4e^{-4}\}$

Hyperparameters. LoRA rank $r=16,\alpha=16$, and dropout 0.1. Precision FP16. Sequence length $T=128$, train batch size $batchsize=128$.

GPUs. Nvidia A800.

### C.2 Empirical Results

#### C.2.1 GLUE Tasks Train Loss

Figure: Figure 3: The train loss of RoBERTa-base fine-tuning. Other settings are same to Figure 1.
Refer to caption: x7.png

#### C.2.2 MNLI Llama3.1-8b

Figure: Figure 4: The train loss of Llama3.1-8b fine-tuning. Other settings are same to Figure 2.
Refer to caption: x11.png

#### C.2.3 Algorithm Framework

Figure: Figure 5: A brief diagram outlining how our theoretical insights guide the experiments.
Refer to caption: x12.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
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).
soundness: 3
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.
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.
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.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
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 `λ`?
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
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.
soundness: 1
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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
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?
soundness: 1
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.
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.
weaknesses: See my questions.

Meta-review: 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.

Decision: Reject