Ranking and splitting attention heads improve language model fine-tuning
Ranking and Partitioning of Attention Heads for Sparse and Continual Reinforcement Learning Fine-tuning
Machine LearningComputation and Language
Summary
Fine-tuning large language models can be tricky and costly because many parts can change unpredictably. The authors found a way to rank different parts of the model’s attention system by how much they influence learning, allowing them to focus on just a few key pieces. They also divide these important parts when teaching the model multiple tasks, helping it remember old tasks better while learning new ones. Their tests on math, coding, and instruction tasks showed this approach works well compared to tuning many parts at once.
What this means in practice
- •For machine learning engineers: Optimize fine-tuning of language models by training selected attention heads for efficient learning and adaptation.
- •For ai product teams: Enhance multi-task model updates by separating attention heads per task to maintain old skills while adding new ones.
Authors
Ruizhe Shi, Zaiwen Yang, Qiubo Wang, Banghua Zhu, Maryam Fazel
Abstract
We study sparse and continual reinforcement learning fine-tuning of language models through attention-head ranking and partitioning. For sparse learning, we rank attention heads by gradient norm to identify a small subset of heads that can be trained effectively. For continual learning, we partition heads across tasks to reduce cross-task interference and preserve previously learned capabilities. Theoretically, we relate gradient norm to training effectiveness, and further characterize why sparse, separated head updates can remain stable. Empirically, across math, instruction-following, and coding tasks, tuning higher-ranked heads consistently outperforms tuning lower-ranked ones while often matching denser baselines, and head partitioning improves retention while preserving adaptation to new tasks.