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

TL;DR

This paper provides theoretical insights into fine-tuning attention mechanisms in large language models, highlighting the importance of selectively tuning certain matrices and assigning different learning rates to improve efficiency and performance.

Contribution

It introduces a novel understanding of the importance of tuning specific attention matrices and learning rate strategies, leading to more efficient fine-tuning methods for LLMs.

Findings

Optimizing the value matrix W_v yields better performance than optimizing the key matrix W_k.

Fine-tuning only W_q and W_v can match or surpass full fine-tuning results.

Higher learning rates for W_v accelerate convergence and enhance performance.

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 explore two remarkable phenomena related to the attention mechanism during the fine-tuning of LLMs (where $\mathbf{W}_q$, $\mathbf{W}_k$, and $\mathbf{W}_v$ denote the weights of the query, key, and value layers, respectively). The first phenomenon, termed "Unequal Importance of Attention Matrices", highlights the impact of fine-tuning different weight matrices. It shows that optimizing the $\mathbf{W}_v$ matrix yields significantly better performance than optimizing the $\mathbf{W}_k$ matrix. Fine-tuning only the $\mathbf{W}_q$ and $\mathbf{W}_v$ matrices is computationally efficient while delivering results comparable to, or even better than fine-tuning all three matrices ($\mathbf{W}_q$, $\mathbf{W}_k$, and $\mathbf{W}_v$). The second phenomenon,"Attention Matrices with Customized Learning Rate Lead to Better Convergence", emphasizes the importance of assigning distinct learning rates to these matrices. Specifically, a higher learning rate for the $\mathbf{W}_v$ matrix compared to $\mathbf{W}_q$ and $\mathbf{W}_k$ accelerates convergence and improves performance. Building on these insights, we propose a new strategy that improves fine-tuning efficiency in terms of both storage and time. Experimental results on benchmark datasets validate the effectiveness of this approach, supporting our theoretical findings. Our analysis lays the theoretical groundwork for configuring and improving algorithms in LLMs fine-tuning.