This paper introduces a fine-tuning strategy for LLMs that leverages the unequal importance of attention matrices and customized learning rates to enhance efficiency, demonstrating through theoretical analysis and experiments on GLUE benchmarks that fine-tuning only Wq and Wv with higher learning rates for Wv can match or exceed full fine-tuning performance with fewer parameters.
Large Language Model, Parameter-Efficient Fine-Tuning, Transformer, Fine-tuning, Efficiency, Optimization
Xinhao Yao, Hongjin Qian, Xiaolin Hu, Gengze Xu, Wei Liu, Jian Luan, Bin Wang, Yong Liu
Gaoling School of Artificial Intelligence, Renmin University of China, Beijing Key Laboratory of Research on Large Models and Intelligent Governance, BAAI, XiaoMi, Engineering Research Center of Next-Generation Intelligent Search and Recommendation, MOE
Generated by grok-3
Background Problem
Large Language Models (LLMs) based on Transformer architectures excel in generalization across tasks but require resource-intensive fine-tuning for specific applications due to their vast parameter count. This paper addresses the computational burden of fine-tuning by focusing on the attention mechanism, specifically the query (Wq), key (Wk), and value (Wv) matrices. It investigates two key issues: the unequal importance of these matrices during fine-tuning and the impact of customized learning rates on convergence efficiency, aiming to reduce computational costs while maintaining or improving performance.
Method
The paper proposes a fine-tuning strategy based on two phenomena in the attention mechanism of Transformers:
- Unequal Importance of Attention Matrices: It suggests that fine-tuning only Wq and Wv matrices often matches or exceeds the performance of fine-tuning all three matrices (Wq, Wk, Wv), reducing parameter count by approximately 1/3. This is supported by theoretical generalization bounds using information-theoretic approaches (Theorem 1).
- Customized Learning Rates for Convergence: It advocates for distinct learning rates, with a higher rate for Wv compared to Wq and Wk, to accelerate convergence. This is backed by convergence analysis in a toy setting and asymptotic scaling arguments for large-width networks (Theorem 2). The strategy involves freezing Wk and fine-tuning Wq and Wv with tailored learning rates (λ = ηV/ηQK), integrated into methods like LoRA and DoRA for parameter efficiency.
Experiment
The experiments were conducted on benchmark datasets from the GLUE suite (e.g., SST2, QNLI, QQP, MNLI) using pre-trained models like Roberta-base and Llama3.1-8b. The setup compared fine-tuning strategies (full fine-tuning, LoRA, DoRA) under different configurations of attention matrices (QKV vs. QV) and learning rate ratios (λ=2,4,8). Due to resource constraints, experiments used a sequence length of T=128 and limited epochs (3-6), which may underrepresent performance compared to standard setups. Results showed that fine-tuning only Wq and Wv with customized learning rates often outperformed or matched full QKV fine-tuning, with significant parameter reduction (e.g., LoRA QV r=16, λ=8 achieved better MRPC results with 1.77M parameters vs. 21.85M for full QKV). However, the setup lacks comprehensive comparison with other PEFT methods and does not explore diverse tasks or architectures, limiting the robustness of conclusions. The results align with theoretical expectations but may be overly optimistic due to cherry-picked configurations.
Further Thoughts
While the paper offers valuable insights into attention mechanism fine-tuning, I am concerned about the scalability of the proposed strategy across diverse tasks and architectures beyond NLP and the GLUE benchmark. The reliance on toy settings for convergence analysis (Theorem 2) and the lack of clarity on how λ (learning rate ratio) should be adapted for different contexts suggest a gap in practical applicability. This connects to broader research in PEFT, such as adapter layers or prompt tuning, where task-specific tuning strategies often outperform generic approaches—could the unequal importance of matrices vary similarly? Additionally, exploring the interplay between attention matrices and emergent abilities in LLMs (e.g., in-context learning) could reveal whether fine-tuning Wv more aggressively impacts capabilities beyond task performance. Future work should also consider robustness to adversarial inputs or safety alignment, as uneven updates to attention components might introduce biases or vulnerabilities, an area underexplored in current PEFT literature.