PaperlayerView exact revision on arXiv

Source-linked AI summary

Ranking and Partitioning of Attention Heads for Sparse and Continual Reinforcement Learning Fine-tuning

Ruizhe Shi, Zaiwen Yang, Qiubo Wang, Banghua Zhu, Maryam Fazel

arXiv:2610.21336v1cs.LGcs.CL

TL;DR

RL fine-tuning can be costly when many language-model parameters are updated, and continual updates can interfere with previously learned capabilities. The paper ranks attention heads by gradient norm for sparse tuning and partitions heads across tasks. It reports that higher-ranked heads outperform lower-ranked ones and that partitioning improves retention while preserving adaptation.

  • Problem

    RL fine-tuning can be expensive when many parameters are updated, while sequential task updates can interfere with previously learned capabilities.

    Source: §1 · p. 1

  • Method

    The paper ranks attention heads by gradient norm for sparse RL tuning and partitions heads across sequential tasks.

    Source: §1 · p. 1

  • Results

    Across several tasks, tuning higher-ranked heads consistently outperforms lower-ranked heads while often matching or exceeding denser baselines, and head partitioning reduces forgetting while preserving adaptation to new tasks.

    Source: §1 · p. 1, §7 · p. 9

  • Takeaways & Limitations

    Gradient norm characterizes the training benefit of a head subset, while head partitioning provides a way to reduce forgetting in continual RL fine-tuning.

    Source: §7 · p. 9

    Experiments are restricted to relatively small language models, limited tasks, and settings without long reasoning traces.

    Source: §7 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

RL fine-tuning can be costly when many parameters are updated, motivating sparse updates to attention heads. The paper ranks heads by gradient norm and partitions them across tasks, with theoretical analysis guiding both designs.

1 INTRODUCTION

  • Updating a large fraction of language-model parameters increases gradient-memory and distributed-synchronization costs, motivating sparse RL fine-tuning.Prior work suggests adaptation can often be achieved by updating restricted model components.Source: §1 · p. 1
  • Attention heads are motivated as units for selection because some are more important or specialized, while sparse task-specific subsets can support high-level capabilities.Prior pruning studies find many heads removable with little degradation and a smaller subset playing more important roles.Source: §1 · p. 1
  • Sequential task updates can interfere with previously learned capabilities, and parameter isolation may still leave indirect functional interference.This motivates careful selection of trainable parameters for continual learning.Source: §1 · p. 1
  • The work studies attention-head ranking and partitioning for sparse and continual RL fine-tuning.It uses gradient norm to select heads for sparse training and allocates heads across sequential tasks.Source: §1 · p. 1
  • The paper analyzes gradient norm's role in optimization and the second-order structure of multi-head attention to characterize parameter and functional isolation.The analysis is presented as theoretical guidance for both ranking and partitioning.Source: §1 · p. 2

2 RELATED WORKS

Related work motivates sparse updates, continual-learning subspaces, and attention heads as structured units. This paper distinguishes its approach by using static head-level gradient norms and connecting them to stability–plasticity.

Sparse optimization in deep learning and LMs

  • Sparse updates to attention modules or a single layer have obtained competitive performance, and can improve distributed-RL communication efficiency.The paper situates head selection within broader sparse optimization work.Source: §2 · p. 2

Subspace continual learning

  • Continual-learning methods constrain updates using subspaces, while related parameter-selection methods use sensitivity, activation differences, or dynamic masks.Unlike Lin et al. (2026), this work uses static head-level gradient norms.Source: §2 · p. 2

Theoretical understanding of continual learning

  • Prior analyses connect forgetting to overlap between task-induced tangent subspaces and study parameter isolation through Hessian-based arguments; this work connects stability–plasticity directly to gradient norms.Source: §2 · p. 2

Mechanism of attention heads

  • Attention heads are heterogeneous: many are redundant, while smaller subsets can specialize or be critical to particular capabilities.These findings motivate treating heads as structured units.Source: §2 · p. 2

3 PRELIMINARIES

The preliminaries define head-wise parameter blocks and gradient-based quantities used to study sparse updates and task interactions. They also specify the multi-head self-attention setting.

Notations

  • The trainable attention parameters are partitioned into disjoint head blocks, and subsets of heads index corresponding parameter and vector blocks.The notation also defines head support as the heads with nonzero update blocks.Source: §3 · p. 2
  • In RL fine-tuning, each training example is a prompt, and its empirical loss is induced by rollouts from that prompt.The text also specifies Euclidean and matrix operator norms.Source: §3 · p. 2

Gradient norm

  • For two tasks, the head-set gradient norms are defined as the norms of their respective loss gradients restricted to that set.The losses are assumed twice continuously differentiable on the relevant training regions.Source: §3 · p. 3

Second-order terms of loss Hessian

  • The directional second-order derivative measures gradient-norm change along its own normalized gradient direction; the mixed derivative measures task-A gradient-norm change along task B's direction.Source: §3 · p. 3

Multi-head self-attention

  • The setup considers a multi-head self-attention layer with an input matrix and H heads, with per-head query and value projections.Source: §3 · p. 3

4 RANKING THE ATTENTION HEADS

The section motivates gradient-norm ranking as a way to select useful attention heads for sparse RL fine-tuning, then evaluates its theoretical basis and empirical performance. The analysis links gradient norm to training progress and reports that top-ranked heads often match or exceed denser tuning while retaining advantages over lower-ranked heads.

4.1 THEORETICAL UNDERSTANDINGS

  • Gradient norm characterizes training efficiency under a fixed head budget through both initial loss-decrease rate and finite-time progress.The top-k heads maximize the initial loss-decrease rate; the finite-time result is expressed through time-averaged gradient norm.Source: §4.1 · p. 3
  • Under the stated second-order condition, time-averaged gradient norm remains controlled by its initial value, so ranking need not be recomputed during training.In the experiments, the ranking remains instructive for hundreds of training steps.Source: §4.1 · p. 4
  • For group-sparse Adam, heads are selected by first-moment magnitudes, and top-k heads approximately optimize Adam’s local quadratic objective.The initial gradient norm serves as a natural surrogate for the first-moment magnitude.Source: §4.1 · p. 4

4.2 HEAD-SPARSE LEARNING EXPERIMENTS

  • Larger learning rates help bottom-50% tuning in some cases, but its gap to top-50% tuning generally remains.The comparison uses a fixed training configuration and the same trainable-parameter budget.Source: §4.2 · p. 5
  • Within-task head rankings are consistently stable, while cross-task rankings differ more but remain small overall, suggesting a shared intrinsic hierarchy.With large training sets, the authors note rankings could be estimated from a small data subset to reduce computation cost.Source: §4.2 · p. 5
  • Gradient-norm and direct-performance layer rankings are highly consistent across the QWEN3-BASE series and exactly match for QWEN3-4B-BASE.The rankings use gradient norms computed on a 3k-subset of DAPO-Math-17k.Source: §4.2 · p. 5, Table 3 · p. 5
  • Performance generally improves as more top-ranked heads are included, while a relatively small head budget can already achieve the major performance.Figure 1 compares top 10%, 30%, 50%, 70%, and 90% head tuning with full-attention tuning across training, in-distribution test, and out-of-distribution test sets.Source: §4.2 · p. 5, Figure 1 · p. 6
  • Top-50% tuning outperforms bottom-50% tuning across all model–task pairs and evaluation splits, and often matches or exceeds full-attention and full tuning, despite updating only half the attention heads.Source: §4.2 · p. 4

5 PARTITIONING THE ATTENTION HEADS

The section analyzes how to allocate attention heads across sequential tasks to balance learning a new task with retaining an earlier one. Experiments show that partitioning generally reduces forgetting while preserving adaptation, with allocation strategies trading off retention and new-task learning.

5.1 THEORETICAL UNDERSTANDINGS

  • Task-B learning improves with its restricted gradient norm, while the change in task-A loss is bounded by the interaction between task-A and task-B gradients on the updated heads.The analysis frames this as a stability–plasticity trade-off.Source: §5.1 · p. 6
  • For task B, the allocation should favor heads with small task-A but large task-B gradient norms; this motivates balancing high-gradient heads across tasks.The discussion notes that unupdated heads' task-A gradient norms may remain controlled by their initial values when gradients change smoothly.Source: §5.1 · p. 6
  • Initial-checkpoint rankings can guide continual training over hundreds of steps, and resetting stage-I heads largely degrades stage-I performance while maintaining stage-II performance.The reset result suggests approximate functional isolation across allocated head subsets.Source: §5.1 · p. 7, §5.2 · p. 8

5.2 SEQUENTIAL CONTINUAL LEARNING EXPERIMENTS

  • The experiments compare Top–Bottom and Round-Robin for two tasks, and use Shared allocation for three tasks, with 25% of heads per task and 25% shared.The authors present these strategies as demonstrations, not an exhaustive or optimal allocation scheme.Source: §5.2 · p. 7
  • Head partitioning generally reduces forgetting while preserving adaptation; Top–Bottom typically retains the first task better, whereas Round-Robin adapts slightly better to the second.Across three-task sequences, the approach also shows minimal degradation on previous tasks.Source: §5.2 · p. 7, Figure 2 · p. 7, Figure 3 · p. 8

6 SECOND-ORDER STRUCTURE OF SPARSE ATTENTION-HEAD UPDATES

This section examines second-order effects in attention-head updates. It shows that distinct heads have separate functional-Hessian blocks and describes a sharper bound for a concentrated-attention regime.

6 SECOND-ORDER STRUCTURE OF SPARSE ATTENTION-HEAD UPDATES

  • The analysis considers one multi-head self-attention layer on a fixed rollout batch, treating subsequent computation as part of the loss and using a Gauss–Newton Hessian decomposition.The result is stated for the attention parameters of that layer.Source: §6 · p. 9
  • For multi-head self-attention, the functional Hessian has zero cross-head blocks, so a selected head set's operator norm is the largest norm among its individual head blocks.The paper contrasts this isolation with shared parameterizations such as standard LoRA and unstructured sparse subsets.Source: §6 · p. 9
  • When attention concentrates on similar token representations, the paper states that a sharper single-head bound is possible under bounded input and weight norms and a bounded loss gradient.The stated regime requires most attention mass to fall on tokens whose representations are close.Source: §6 · p. 9

7 CONCLUSION

The conclusion summarizes gradient-based head ranking and partitioning as approaches to sparse and continual RL fine-tuning, while identifying limits in model, task, and theoretical scope. The appendices include proofs supporting the ranking and continual-learning analyses.

7 CONCLUSION

  • The paper studies gradient-based attention-head ranking and partitioning, reporting efficient sparse tuning with high-ranked heads and reduced forgetting under continual head partitioning.Experiments are limited to relatively small language models, a limited task set, and settings without long reasoning traces; the analysis may not extend directly to MLP layers, and the second-order theory covers one attention layer.Source: §7 · p. 9

A.1 OMITTED PROOFS IN SECTION 4.1

  • Selecting the k heads with the largest gradient norms maximizes the initial loss decrease.This is the conclusion stated in the omitted proof of Theorem 6.Source: §A.1 · p. 14

A.2 OMITTED PROOFS IN SECTION 5.1

  • The omitted proof of Theorem 9 relates accumulated task-B gradient norms to the decrease in task-B loss over the training interval.The passage states that this relation gives the desired bound.Source: §A.2 · p. 15

A.3 OMITTED PROOFS IN SECTION 6

The omitted proofs characterize attention-head second-order structure using centered attention moments and head-specific parameter blocks. They also derive bounds under bounded-input and concentrated-attention conditions.

A.3 OMITTED PROOFS IN SECTION 6

  • The proof uses attention-weighted means and centered moments; its stated lemma assumes bounded token representations and concentrated attention on a set of similar representations.The concentration condition places at least 1 − ε of the attention mass on the set, whose representations are within ρ of one another.Source: §A.3 · p. 15
  • Under bounded input and weight norms, the cited second-derivative formulas express nonzero functional-Hessian blocks through second and third centered attention moments.The proof notes that these blocks are linear in the second moment, third moment, or both.Source: §A.3 · p. 16
  • Each head output depends only on its own parameter block, making the functional Hessian block diagonal across heads and its operator norm the maximum norm of a diagonal block.This is the proof's basis for the head-wise Hessian structure.Source: §A.3 · p. 16

B.1 SUPPLEMENTAL TRAINING CURVES

The appendix supplements the paper with sparse-training curves, continual-learning accuracy and reset results, parameter and ranking-distance tables, and layer-wise distributions of selected heads.

B.1 SUPPLEMENTAL TRAINING CURVES

  • Figures 7–9 plot sparse RL training accuracy for QWEN3-1.7B, QWEN3-4B, and MISTRAL-7B across the listed head-tuning and denser-tuning methods.The QWEN3 figures include top-50%, bottom-50%, and random-50% head tuning, full-attention tuning, and full tuning; the MISTRAL figure includes top-50%, bottom-50%, and full-attention tuning.Source: Figure 7 · p. 18, Figure 8 · p. 18, Figure 9 · p. 18

B.2 SUPPLEMENTAL NUMERICAL RESULTS

  • Tables 4–9 report accuracy across two- and three-stage continual RL task sequences, with parenthesized changes from the preceding stage.The sequences cover Math, Coding, and IF on QWEN3-1.7B and QWEN3-4B.Source: Table 4 · p. 19, Table 5 · p. 19, Table 6 · p. 20
  • Table 10 presents trainable-parameter counts and FP32-equivalent gradient memory, using 4 bytes per parameter.Source: Table 10 · p. 21
  • Table 11 reports Spearman distances between head rankings, defining lower distances as greater similarity and an expected distance of 1 for independent uniform random rankings.Within-task distances use disjoint subsets; cross-task distances use training sets.Source: Table 11 · p. 21
  • Table 12 reports test-accuracy changes after resetting stage-I-updated heads following stage II; its caption says the results suggest adaptations on disjoint head sets are functionally localized.The caption identifies stage-II performance that does not decrease after reset as highlighted.Source: Table 12 · p. 22

B.3 LAYER DISTRIBUTIONS OF TOP-RANKED HEADS

B.4 HYPERPARAMETER TUNING RESULTS

The section presents QWEN3 sparse RL training curves for comparing top- and bottom-ranked head tuning across learning rates on Math and IF tasks.

B.4 HYPERPARAMETER TUNING RESULTS

  • Figure 17 plots QWEN3 sparse RL accuracy curves for Math and IF, fixing top-50% head tuning at 3ˆ10´6 and testing bottom-50% tuning at 3ˆ10´6, 5ˆ10´6, and 7ˆ10´6.Top-50% tuning is shown as a horizontal line with an error bar for mean±std over three random rollout seeds.Source: Figure 17 · p. 25

C IMPLEMENTATION DETAILS

The implementation section describes the benchmark datasets, base models, training and evaluation settings, and computational setup for the experiments.

Datasets and benchmarks

  • Math and instruction-following use 1k training subsets and disjoint 1k in-distribution test subsets, with MATH-500 and IFBench as out-of-distribution tests.Coding uses MBPP for training, MBPP+ for in-distribution testing, and APPS-intro-test for out-of-distribution testing; MBPP+ uses the same problems as MBPP with strengthened tests.Source: C IMPLEMENTATION DETAILS · p. 26

Models

  • The experiments use QWEN3-1.7B, QWEN3-4B, and MISTRAL-7B-INSTRUCT-V0.3, with model selection limited by computation resources.Source: C IMPLEMENTATION DETAILS · p. 26

Training and evaluation hyper-parameters

  • Training uses 16 prompts with four responses each per rollout, task-specific generation limits, and 400–600 rollout steps; evaluation reports pass@1 with greedy decoding and task-specific verifiers.The training passage also specifies optimizer, learning rates, clipping range, and other hyperparameters; evaluation generates one response per example with a maximum length of 4,096 tokens.Source: C IMPLEMENTATION DETAILS · p. 26

Computation resources and codebase

  • Experiments run on NVIDIA A100 and RTX A6000 GPUs, typically taking about one day; the codebase is based mainly on miles and the sparse-head implementation is not yet fully optimized.The authors state that realizing its full efficiency benefits should be straightforward with modern training infrastructure.Source: C IMPLEMENTATION DETAILS · p. 26

Related papers

Loading 2610.21336v1…