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

1 Introduction

Large Language Models (LLMs) are often built on Transformer architectures [Vaswani et al., 2017] and possess a large number of parameters, enabling them to generalize across a broad range of general tasks [Likhomanenko et al., 2021; Touvron et al., 2021; Dosovitskiy et al., 2021; Min et al., 2022]. However, achieving optimal performance on specific tasks typically necessitates fine-tuning these pretrained 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 finetuning, is highly computationally expensive. To reduce the computational cost, various parameter-efficient fine-tuning (PEFT) methods have been proposed [Ding et al., 2023; Houlsby et al., 2019; Lester et al., 2021; Li and Liang, 2021; Hu et al., 2022], 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 matrixWq (Wq), the key matrixWk (Wk), and the value matrixWv (Wv).

∗Work done during the collaborative project with XiaoMi.

†Corresponding Author.

During the fine-tuning of LLMs involving the attention mechanism, two interesting phenomena have been observed 1 : (1) Unequal Importance of Attention Matrices-optimizing the Wv is pivotal for enhancing performance, significantly more so than adjustments to the Wk, which exhibit limited impact on the outcomes. Additionally, fine-tuning only the Wq and Wv often yields results that are comparable to or surpass those achieved by fine-tuning all three matrices Wq, Wk, andWv, which also reduces the number of tunable attention parameters by approximately 1/3, offering computational benefits (Section 3). (2) Attention Matrices with Customized Learning Rate Lead to Better Convergence-using the same learning rate for Wq&Wk and Wv is not optimal for efficient convergence. In fact, it is essential to apply distinct learning rates for the Wq, Wk, and Wv components to ensure optimal fine-tuning performance. Specifically, the learning rate forWv should generally be higher than that for Wq andWk to facilitate efficient convergence (Section 4).

While certain empirical guidelines, such as the original Low-Rank Adaptation (LoRA) [Hu et al., 2022], 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. Reflecting on the attention equation itself (Section 2): (1) In linear algebra, two matrices multiplied without an intermediate activation can be equivalent to a single matrix. Some studies [Noci et al., 2022; Bao et al., 2024] often treat Wq and Wk as a single unit (Wqk = WqWT k ), however, the benefits of fine-tuning Wq&Wv alone have yet to be further clarified. (2) Considering the scenario where the values of Wq, Wk, and Wv approach zero, the gradients ofWq&Wk tend to diminish towards zero. In contrast, the gradient ofWv remains non-zero due to the influence of softmax normalization. Driven by the above motivations, this paper delves into the issue from the following two perspectives.

1 Extended version and code, are available at https://github.com/ XiaoMi/EfficientFT.

• Generalization: advantages of fine-tuning Wq&Wv over Wq,Wk,Wv together. We perform a thorough theoretical analysis to demonstrate the advantages. To be more specific, we employ information-theoretic approaches [Xu and Raginsky, 2017; Polyanskiy and Wu, 2019; Wang and Mao, 2022; Zhu et al., 2024] to establish the generalization bounds of fine-tuning pre-trained models with attention mechanism (See Theorem 1 for details). This indicates that fine-tuning Wq&Wv instead of Wq,Wk,Wv 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, suggesting that Wv may experience instances of inefficient learning during downstream task fine-tuning. This naturally leads to the hypothesis that accelerating the learning of Wv in the early stages could potentially induce Wk and Wq to begin learning earlier. Additionally, by using scaling arguments for large width-n networks [Yang et al., 2022; Hayou et al., 2024b], we illustrate (Theorem 2) that the feature learning of attention mechanism is efficient when the learning rate for Wv should be generally much larger than that ofWq&Wk 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 [Wang et al., 2018] and open source pre-trained models [Liu et al., 2019; AI@Meta, 2024] verify that the method can visibly influence fine-tuning efficiency. 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.

A summary of the main theoretical analyses. 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, storage-friendly), we give Theorem 1 (information-theoretic generalization bounds), showing that with the same r value, fine-tuning Wq&Wv consistently achieves results comparable to or even surpassing those of fine-tuning Wq,Wk,Wv. This reduces the number of parameters for the same r, while improving generalization bounds and potentially providing memory benefits. In Section 4 (optimization, time-friendly), 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 Wv should be generally much larger than that of Wq&Wk 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.

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 pretrained 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 C ∈ Rm×din over which we would like to perform attention and a query vector x ∈ Rdin , that is, the input is [C, x] ∈ R(m+1)×din . Conventional attention can be expressed as 2 :

where Wq,Wk,Wv ∈ Rdin×dout are query, key and value (projection) matrices.

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

where l1, l2 are coefficients, h is denoted as the hidden representation to be directly modified and ∆h is a modification vector. Moreover, h and x can represent the attention output and input respectively. Here, we present two special cases:

• LoRA. LoRA [Hu et al., 2022] injects trainable low-rank matrices into transformer layers to approximate the weight updates. Instead of directly adjusting the full weight matrix W ∈ Rdin×dout , LoRA represents its update with a low-rank decomposition W + ∆W = W + AB, where A ∈ Rdin×r ,B ∈ Rr×dout are tunable parameters. For a specific input x, LoRA modifies the projection output h as (where s ≥ 1 is a tunable scalar hyperparameter):

2 For simplicity, we focus on the last vector of input in a singlehead self-attention. Our analysis is readily generalizable to multihead self-attention.

Prefix tuning. Prefix tuning [Li and Liang, 2021] prepends r tunable prefix vectors to the keys and values of the attention mechanism at every layer. Specifically, two sets of prefix vectors Pk, Pv ∈ Rr×dout are concatenated with the original key CWk and value CWv, attention is then applied to the prefixed keys and values as 3 :

where α(x) = ∑i exp(xWqPTk )i ∑i exp(xWqPTk )i+∑j exp(xWqWT k CT )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 connect it with LoRA in Appendix A.1.

Remark 1. By defining A = WqPTk ,B = Pv 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, A (WqPTk ) is responsible for generating attention scores, while B (Pv) 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 [Zhu et al., 2024; Hayou et al., 2024a], which empirically observed an asymmetry where the project-down matrix A is responsible for extracting features from the input, while the project-up matrix B utilizes these features to generate the desired output in LoRA.

Θ Notation. We use standard asymptotic notation to describe behavior as the width n grows, following conventions in [Yang et al., 2022; Hayou et al., 2024b]. Given sequences cn ∈ R and dn ∈ R+, we write cn = O(dn) and cn = Ω(dn) to mean cn < κdn or cn > κdn, respectively, for some constant κ > 0. We denote cn = Θ(dn) when both cn = O(dn) and cn = Ω(dn) hold, implying that cn and dn grow at comparable rates. For vector sequences cn = (cin )1≤i≤k ∈ Rk (for some k > 0), we write cn = O(dn) when cin = O(din ) for all i ∈ [k], and analogous notation applies for other asymptotic bounds. Finally, when the sequence cn is a vector of random variables, convergence is understood to refer to convergence in the second moment (i.e., L2 norm).

3 Advantages and Generalization Analysis

In this section, we show our first interesting observation (Unequal Importance of Attention Matrices) in fine-tuning the attention mechanism and the storage benefits of fine-tuning onlyWq andWv (Section 3.1). Next, we give mutual information based generalization bounds of fine-tuning only Wq andWv (Section 3.2), providing a better generalization error.

3.1 Empirical Advantages

To explore the Unequal Importance of Attention Matrices, 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 [Liu et al., 2019] and Llama3.18b [AI@Meta, 2024]) across various benchmarks [Wang et al., 2018]. Further details on the experimental setup and additional empirical results are in Appendix B.1.

3 Without loss of generalization, we ignore the softmax scaling factor for ease of notation.

Table 1 provides a detailed comparison of the impact of fine-tuning different weight matrices (Wq,Wk,Wv) across various rank values r and weight update strategies in LoRA fine-tuning on tasks like SST2, QNLI, QQP, and MNLI. We can see a clear trend where solely updating the Wv matrix outperforms just learning theWq,Wk matrix. Interestingly, the combination of fine-tuning bothWq andWv often leads to performance that matches or even exceeds that achieved by fine-tuning all three matrices Wq,Wk, and Wv. This pattern is consistently observed across various tasks and rank values, further emphasizing the importance of these two matrices over Wk 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 [Lialin et al., 2023]. Fine-tuningWq&Wv alone as opposed to both Wq&Wv and Wk reduces the number of parameters by 1/3, when the dimensions of Wq, Wk, and Wv are the same. Moreover, we discuss in Appendix B.2 about why fine-tune Wq&Wv instead of Wk&Wv.

3.2 Information-Theoretic Generalization Bounds

In the previous part, we establish that the Unequal Importance of Attention Matrices among Wq , Wk, and Wv during fine-tuning. Some studies [Noci et al., 2022; Bao et al., 2024] often treat Wq andWk as a single unit (Wqk = WqWT k ), however, the benefits of fine-tuning Wq&Wv alone, rather than fine-tuningWq&Wv, andWk 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 [Xu and Raginsky, 2017; Russo and Zou, 2019; Steinke and Zakynthinou, 2020; Wang and Mao, 2022] 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 Z = X × Y be the instance space and µ be an unknown distribution on Z, specifying random variable Z. Here, X denotes the feature space and Y is the label space. Suppose one observes a training set SN ≜ (Z1, ...,ZN) ∈ Z N , with N i.i.d. training examples drawn from µ. In the information-theoretic analysis framework, we let W be the space of hypotheses related to the model, and a stochastic learning algorithm A which takes the training examples SN as its input and outputs a hypothesis W ∈ W according to some conditional distribution QW|SN . Given a loss function ℓ : W ×Z → R+, where ℓ(w, Z) measures the "unfitness" or "error" of any Z ∈ Z with respect to a hypothesis w ∈ W. We take ℓ as a continuous function and assume that ℓ 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 ∈ W, the population risk is defined as Lµ(w) ≜ EZ∼µ[ℓ(w,Z)]. However, since only can partially observe µ via the sample SN, we instead turn to use the empirical risk, defined as LSN (w) ≜ 1N ∑Ni=1 ℓ(w, Zi). Then the expected generalization error of A is defined as

where the expectation is taken over (SN,W) ∼ µN ⊗ QW |SN .

Consider the following variations of fine-tuning algorithms: tuning bothWk andWq&Wv matrices (as in classic attention mechanism in fine-tuning), tuning only Wq&Wv:

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 h + ∆h = xW + x∆W. Let W = {Wi}Li=1 be a set of abstract parameter matrices related to a pretrained model, where each Wi is associated with the parameters Wiq ,Wik ,Wiv . The indices 1, ...,L represent the layers of the model where these parameters are to be fine-tuned. Let I ⊆ {1, ...,L} denote the subset of layers selected for finetuning. Given a fine-tuning training set SN, 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 ∆W = {∆Wi}Li=1 (with other details left open to choice). (1) AQKV : For each i ∈ I, optimize {Wiq ,Wik ,Wiv }i∈I to fit the data SN. (2) AQV : For each i ∈ I, optimize {Wiq ,Wiv }i∈I to fit the data SN.

Then we use the information-theoretic generalization framework to bound the generalization error:

Theorem 1 (Generalization bounds on adapting Wq&Wv and/or Wk). Consider the algorithms of Definition 1. Assume the loss ℓ(W, Z) is R-subGaussian under (∆W, Z) ∼ P∆W|W × µ. Then (See Appendix A.2 for a proof),

where Wiq ,Wik ,Wiv ∈ Rdin×dout .

Remark 2 (Discussion of the advantages). We can evaluate the empirical risk (LSN ) 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µ). 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-tuningWq&Wv consistently achieves results comparable to or even surpassing those of fine-tuningWq,Wk,Wv. This reduces the number of parameters for the same r, while improving generalization bounds and potentially providing memory benefits.

4 Convergence Analysis in Optimization

In Section 3, we have already demonstrated the generalization performance of the attention mechanism during finetuning. Our focus now shift toward optimizing convergence efficiency. Some optimization observations have also been reported in previous works [Hu et al., 2022; Li et al., 2023; He et al., 2024], such as: [Li et al., 2023] provide theoretical analyses of learning dynamics in transformers and observes a roughly two-stage process of self-attention. [He et al., 2024] 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

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:

Intuitively, if Wq,Wk,Wv are initialized as random matrices close to zero and trained simultaneously, then in the initial step, ∇Wk L(∇Wq L) contains the termWq(Wk), which is close to 0. By contrast, ∇Wv L contains the softmaxnormalized attention weights. Therefore, during the initial steps (in training), Wv intuitively grows at a much faster rate than Wk(Wq ).

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

Next, we present the second interesting phenomenon Attention Matrices with Customized Learning Rate Lead to Better Convergence. We use the General Language Understanding Evaluation (GLUE, [Wang et al., 2018]) 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 [Liu et al., 2019] and Llama3.1-8b [AI@Meta, 2024] on MNLI, QQP, QNLI, and SST2 tasks with varying learning rates (ηQK, ηV ) to identify the optimal combination. Other empirical details are provided in Appendix B.1 and we evaluate the LLaMA3.1-8B model on more complex benchmarks in Appendix B.3. We present our empirical results using LoRA to fine-tune language models, as visualized in the heatmaps (Figure 1 and Figure 2).

In Figure 1 and Figure 2 (Appendix B.3), we observe that (1) test accuracy consistently reaches its maximum for certain sets of learning rates where ηQK < ηV , outperforming the standard practice of setting ηQK and ηV equal. (2) More interestingly, the gap between the optimal choice of learning rates overall and the optimal choice when ηQK = ηV varies across different tasks. This is probably due to the fact that harder task (like MNLI) requires more efficient feature learning. We compare two optimal learning rate (ηQK, ηV ) settings in Figure 2 (Left), the ηV >> ηQK setting has a better convergence than ηV = η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., 2022], 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 B.3, we provide additional results including the training loss.

4.2 Convergence Analysis for Learning Rate

It naturally raises the question of why ηQK and ηV should be set differently. In practice, the large width (embedding dimension) of state-of-the-art models makes it valuable to examine training dynamics as width approaches infinity.

Starting with a Toy setting. Revisiting Definition 1, we have ∆h = softmax(xA)B. In the case of a linear attention mechanism, we instead have ∆h = xAB. Then consider the following toy setting:

where W ∗ ∈ Rn×1 are the fixed 4 pre-trained weights, b ∈ R, a ∈ R1×n are adaptation weights, x ∈ Rn is the model input (This corresponds to r = 1 in Definition 1). The training goal is to minimize the loss L(θ) = 12 (f(x) - y)2 where θ = (a, b) and (x, y) is an input-output datapoint 5 . Similar to LoRA, we generally aim to initialize the product aT b to zero, ensuring that fine-tuning starts from the pretrained model. This requires at least one of the weights, a (related toWq&Wk) or b (related toWv), to be initialized to zero. If both are initialized to zero, Wq&Wk learning cannot occur efficiently in init steps, as discussed in Section 4.1 (More detailed initialization settings are in Appendix A.3).

And we assume that x = Θ(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 ˆ y = y - xW ∗ . We denote the fine-tuning step by using subscript t. Let Ut = ft(x) - y, the gradients are then computed as:

And at step t with learning rate ηa, ηb > 0, we have

Remark 3. The output update is influenced by three key terms. The first two items δ1 t , δ2 t (order one in ηa/η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 δ3 t (order two in ηaηb) corresponds to a multiplicative update that captures the combined effects of changes in both a and b. As we scale the width 6 , the desirable feature updates are such that ∆ft = Θ(1), ensuring they remain unaffected by this scaling

4 Here, we primarily focus on the case of ∆W to provide insightful theoretical results.

5 To 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.

6 This property is generally satisfied in practice when the model width is large (e.g., n ≈ 800 for Roberta-base and n ≈ 4000 for Llama3.1-8b).

(the updates do not explode with width, see x for more details). Ideally, we aim for both δ1 t and δ2 t to be Θ(1). If this condition isn't met, it indicates that either a or b is not being updated efficiently. For example, if δ1 t = o(1), it suggests that as n →∞, 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 δi t = Θ(1) for i ∈ {1, 2} and all t > 1, it means that both a and b parameter updates significantly contribute to the change in ft(x). We will see that when δ1 t , δ2 t are Θ(1), the term δ3 t is also Θ(1).

Let us assume that we train the model with gradient descent with learning rate ηa = Θ(nca ), ηb = Θ(ncb ) for some ca, cb ∈ R. In the study by [Yang et al., 2022], 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γ for some γ ∈ R. We write v = Θ(nγ[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 Θ(nγ[v] ) (See Appendix A.4 for the formal definition of γ).

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

which, after simple calculations, implies that ca + cb = -1. Notice that the above also leads to the ca+cb+1+γ[xa⊤t-1 ]+ γ[bt-1] = 0 (δ3 t = Θ(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 ηa = Θ(n-1 ) and ηb = Θ(1), we have for all t > 1, i ∈ {1, 2, 3},δi t = Θ(1). In other words, the feature learning of attention mechanism is efficient when ηQK(ηa) = Θ(n-1 ), ηV (ηb) = Θ(1). We denote ηV /ηQK as λ. We refer the reader to Appendix A.5 for more details on the proof.

Remark 4. In practice, Theorem 2 implies that the learning rate for Wv should be generally much larger than that of Wq&Wk in fine-tuning. We verify that this scaling is valid for general neural network models in Section 4.1. Naturally, the optimal ratio λ depends on the architecture and the finetuning task through the constants in 'Θ'. 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 future.

5 An Example of Improving Fine-tuning

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 Wk and fine-tuning theWq&Wv using different learning rates. This procedure is reported in Figure 5. In Appendix B.2, we discuss how to set the ratio λ?

Experimental setup. We conduct experiments on widely adopted benchmark datasets [Wang et al., 2018] and Robertabase model [Liu et al., 2019]. We selected mainstream baselines: Full Fine-tuning, LoRA [Hu et al., 2022] and DoRA [Liu et al., 2024]. 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. In our experiments, we evaluated the performance for λ 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 B.1 and the base model performance for each task can be seen in Table 2 and Appendix B.3. We also extend ablation experiments on Mistral-7B [AI@Mistral, 2023] in Appendix B.3.

Results. We leverage our theoretical results (Theorem 1 and Theorem 2) to enhance the efficiency of existing finetuning methods, such as Full Fine-tune, LoRA [Hu et al., 2022] and DoRA [Liu et al., 2024], 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, λ = 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: Unequal Importance of Attention Matrices and Attention Matrices with Customized Learning Rate Lead to Better Convergence. While theoretical analysis of these phenomena is limited, this paper provides insights from two angles: Generalization-fine-tuning only Wq and Wv 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 how task type and architecture affect the optimal learning rate ratio λ. These studies will further deepen our understanding of attention-based fine-tuning.

Ethical Statement

There are no ethical issues.

Acknowledgments

We sincerely appreciate the anonymous reviewers for their valuable suggestions. This research was supported by National Natural Science Foundation of China (No.62476277), National Key Research and Development Program of China (NO. 2024YFE020320), CCF-ALIMAMA TECH Kangaroo Fund(No.CCF-ALIMAMA OF 2024008), and HuaweiRenmin University joint program on Information Retrieval. We also acknowledge the support provided by the fund for building worldclass universities (disciplines) of Renmin University of China and by the funds from Beijing Key Laboratory of Big Data Management and Analysis Methods, Gaoling School of Artificial Intelligence, Renmin University of China, from Engineering Research Center of NextGeneration Intelligent Search and Recommendation, Ministry of Education, from Intelligent Social Governance Interdisciplinary Platform, Major Innovation & Planning Interdisciplinary Platform for the "DoubleFirst Class" Initiative, Renmin University of China, from Public Policy and Decisionmaking Research Lab of Renmin University of China, and from Public Computing Cloud, Renmin University of China.