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

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

cs.LGOct 15, 2026

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

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

Organizations: University of Washington · Tsinghua University · University of Hong Kong · Amazon, Inc.

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.

Figures & tables

Model	Task	Eval split	Base	Top 50%	Bottom 50%	Rand 50%	Full-attn	Full
Qwen3-1.7B	Math	In-dist test	26.6	36.7 $\pm 1.0$	30.9 $\pm 1.0$ ( $\downarrow 5.8$ )	38.6	33.1	35.3
		Out-dist test	70.2	76.7 $\pm 0.9$	73.6 $\pm 1.4$ ( $\downarrow 3.1$ )	79.2	74.2	73.6
Qwen3-1.7B	IF	In-dist test	41.7	57.3 $\pm 0.8$	47.0 $\pm 0.2$ ( $\downarrow 10.3$ )	51.6	59.4	54.6
		Out-dist test	17.7	22.2 $\pm 2.7$	15.9 $\pm 0.8$ ( $\downarrow 6.3$ )	19.7	20.0	19.7
Qwen3-1.7B	Coding	In-dist test	53.7	59.5 $\pm 1.2$	57.8 $\pm 0.8$ ( $\downarrow 1.8$ )	58.5	58.2	61.6
		Out-dist test	33.6	35.2 $\pm 1.4$	33.9 $\pm 0.4$ ( $\downarrow 1.3$ )	35.6	34.8	34.4

Table 1: Sparse RL accuracy (%) of Qwen3 series. Parenthesized arrows are the performance drops compared with top 50%. The top-2 results are highlighted in pink. Top/bottom 50% performance is reported as mean

\pm

std over three random rollout seeds

\{32,42,52\}

Model	Task	Eval split	Base	Top 50%	Bottom 50%	Full-attn
Mistral-7B-Instruct-v0.3	IF	In-dist test	27.4	69.9	53.5 ( $\downarrow 16.4$ )	56.9
		Out-dist test	18.3	21.0	20.3 ( $\downarrow 0.7$ )	25.3

Table 2: Sparse RL accuracy (%) of Mistral-7B-Instruct-v0.3 . Parenthesized arrows are the performance drops compared with top 50%. The highest results are highlighted in pink.

Model	Performance ranking ( Zhang et al., 2026 )	Gradient norm ranking
Qwen3-1.7B-Base	$10>12>1>7>24$	$12>10>7>24>1$
Qwen3-4B-Base	$16>14>11>24>2$	$16>14>11>24>2$
Qwen3-8B-Base	$16>15>8>2>0$	$15\approx 16>0>8>2$

Table 3: Math candidate-layer rankings for Qwen3-Base series. Rankings are obtained using the gradient norms of attention heads in each layer, computed on a

3

k-subset of DAPO-Math-17k .

Panel of Figure 1 — Figure 1: Sparsity–performance trade-off on training, in-distribution test, and out-of-distribution test sets. Each panel includes top 10, 30, 50, 70, and 90% attention-head tuning plus full-attention tuning. All jobs share exactly the same training configuration.

Figure 2: Continual RL accuracy of Qwen3-1.7B (Math

\rightarrow

Coding). The model is first trained on Math task and then on Coding task. Shaded regions mark stages.

Figure 3: Continual RL accuracy of Qwen3-4B (Coding

\rightarrow

IF). The model is first trained on Coding task and then on IF task. Shaded regions mark stages.

Figure 4: Continual RL accuracy of Qwen3-4B (Math

\rightarrow

IF). The model is first trained on Math task and then on IF task. Shaded regions mark stages.

Figure 5: Continual RL accuracy of Qwen3-1.7B (Math

\rightarrow

Coding

\rightarrow

IF). The model is first trained on Math task, then on Coding task, and finally on IF task. Shaded regions mark stages.

Figure 6: Continual RL accuracy of Qwen3-4B (Math

\rightarrow

\rightarrow

Coding). The model is first trained on Math task, then on IF task, and finally on Coding task. Shaded regions mark stages.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Sparse RL training accuracy curves of Qwen3-1.7B . Methods include top 50%, bottom 50%, and random 50% attention-head tuning, full-attention tuning, and full tuning. Full tuning is shown as a final-score horizontal line.

Figure 8: Sparse RL training accuracy curves of Qwen3-4B . Methods include top 50%, bottom 50%, and random 50% attention-head tuning, full-attention tuning, and full tuning. Full tuning is shown as a final-score horizontal line.

Figure 9: Sparse RL training accuracy curves of Mistral-7B . Methods include top 50% and bottom 50% attention-head tuning, and full-attention tuning.

Method	Train stage	Eval split	Math metric	Coding metric
Top 50% $\rightarrow$ bottom 50%	$\mathrm{I}$	In-dist test	37.5	56.6
		Out-dist test	75.6	37.0
	$\mathrm{II}$	In-dist test	35.8 ( $\downarrow 1.7$ )	56.1
		Out-dist test	75.2 ( $\downarrow 0.4$ )	37.7
Full-attn continual	$\mathrm{I}$	In-dist test	33.1	54.8
		Out-dist test	74.2	35.2

Appendix

Table 4: Continual RL accuracy (%) of Qwen3-1.7B (Math

\rightarrow

Coding). The model is first trained on Math task and then on Coding task. Parenthesized arrows are performance changes from the previous stage. The highest performance across all methods is highlighted in pink. The performance before training on the corresponding task is highlighted in gray.

Method	Train stage	Eval split	IF metric	Math metric
Top 50% $\rightarrow$ bottom 50%	$\mathrm{I}$	In-dist test	58.0	24.6
		Out-dist test	20.3	70.6
	$\mathrm{II}$	In-dist test	57.7 ( $\downarrow 0.3$ )	30.5
		Out-dist test	20.7 ( $\uparrow 0.3$ )	73.8
Round robin 50%+50%	$\mathrm{I}$	In-dist test	52.5	27.2
		Out-dist test	19.0	70.0

Appendix

Table 5: Continual RL accuracy (%) of Qwen3-1.7B (IF

\rightarrow

Math). The model is first trained on IF task and then on Math task. Parenthesized arrows are performance changes from the previous stage. The highest performance across all methods is highlighted in pink. The performance before training on the corresponding task is highlighted in gray.

Method	Train stage	Eval split	Coding metric	IF metric
Top 50% $\rightarrow$ bottom 50%	$\mathrm{I}$	In-dist test	69.0	49.9
		Out-dist test	51.4	24.3
	$\mathrm{II}$	In-dist test	68.3 ( $\downarrow 0.7$ )	53.1
		Out-dist test	52.6 ( $\uparrow 1.2$ )	23.0
Round robin 50%+50%	$\mathrm{I}$	In-dist test	68.0	49.8
		Out-dist test	42.9	25.7

Appendix

Table 6: Continual RL accuracy (%) of Qwen3-4B (Coding

\rightarrow

IF). The model is first trained on Coding task and then on IF task. Parenthesized arrows are performance changes from the previous stage. The highest performance across all methods is highlighted in pink. The performance before training on the corresponding task is highlighted in gray.

Method	Train stage	Eval split	Math metric	IF metric
Top 50% $\rightarrow$ bottom 50%	$\mathrm{I}$	In-dist test	59.3	49.6
		Out-dist test	88.4	22.7
	$\mathrm{II}$	In-dist test	60.0 ( $\uparrow 0.7$ )	53.6
		Out-dist test	87.2 ( $\downarrow 1.2$ )	26.0
Round robin 50%+50%	$\mathrm{I}$	In-dist test	53.2	50.3
		Out-dist test	86.0	23.3

Appendix

Table 7: Continual RL accuracy (%) of Qwen3-4B (Math

\rightarrow

IF). The model is first trained on Math task and then on IF task. Parenthesized arrows are performance changes from the previous stage. The highest performance across all methods is highlighted in pink. The performance before training on the corresponding task is highlighted in gray.

Train stage	Eval split	Math metric	Coding metric	IF metric
$\mathrm{I}$	In-dist test	34.5	55.6
	Out-dist test	76.2	34.9
$\mathrm{II}$	In-dist test	32.4 ( $\downarrow 2.1$ )	59.3	41.7
	Out-dist test	75.0 ( $\downarrow 1.2$ )	35.2	14.0
$\mathrm{III}$	In-dist test	33.6 ( $\uparrow 1.2$ )	58.7 ( $\downarrow 0.6$ )	47.3
	Out-dist test	74.8 ( $\downarrow 0.2$ )	34.0 ( $\downarrow 1.2$ )	17.7

Appendix

Table 8: Continual RL accuracy (%) of Qwen3-1.7B (Math

\rightarrow

Coding

\rightarrow

IF). The model is first trained on Math task, then on Coding task, and finally on IF task. Parenthesized arrows are performance changes from the previous stage. The performance before training on the corresponding task is highlighted in gray.

Train stage	Eval split	Math metric	IF metric	Coding metric
$\mathrm{I}$	In-dist test	55.7	50.2
	Out-dist test	86.8	21.3
$\mathrm{II}$	In-dist test	53.6 ( $\downarrow 2.1$ )	57.7	66.4
	Out-dist test	83.0 ( $\downarrow 3.8$ )	24.0	43.5
$\mathrm{III}$	In-dist test	53.9 ( $\uparrow 0.3$ )	56.7 ( $\downarrow 1.0$ )	67.2
	Out-dist test	85.0 ( $\uparrow 2.0$ )	23.3 ( $\downarrow 0.7$ )	49.2

Appendix

Table 9: Continual RL accuracy (%) of Qwen3-4B (Math

\rightarrow

\rightarrow

Coding). The model is first trained on Math task, then on IF task, and finally on Coding task. Parenthesized arrows are performance changes from the previous stage. The performance before training on the corresponding task is highlighted in gray.

Model	50% Attn	Full Attn	Full
Qwen3-1.7B	0.176B / 0.656 GiB	0.352B / 1.312 GiB	1.721B / 6.410 GiB
Qwen3-4B	0.472B / 1.758 GiB	0.944B / 3.516 GiB	4.022B / 14.985 GiB
Mistral-7B-Instruct-v0.3	0.671B / 2.500 GiB	1.342B / 5.000 GiB	7.248B / 27.001 GiB

Appendix

Table 10: Trainable parameters and FP32-equivalent gradient memory (4 bytes per parameter).

	Within-task			Cross-task
Model	Math	IF	Coding	Math & IF	Math & Coding	IF & Coding
Qwen3-1.7B	0.025	0.026	0.051	0.044	0.084	0.081
Qwen3-4B	0.067	0.040	0.047	0.076	0.121	0.104

Appendix

Table 11: Spearman distances (

d:=1-\text{Spearman correlation}

) between head rankings, where lower values indicate more similar rankings, identical rankings have distance

0

, and two independent uniformly random rankings have expected distance

1

. For Math and IF, within-task distances are averaged over the three pairwise comparisons among disjoint 1k-subsets of DAPO-Math-17k and RLVR-IFEval . For Coding, the within-task distance is computed between two disjoint 487-subsets of MBPP . Cross-task distances are computed using the training sets.

			Stage $\mathrm{I}$ task		Stage $\mathrm{II}$ task
Model	Task sequence	Head allocation	In-dist	Out-dist	In-dist	Out-dist
Qwen3-1.7B	Math $\rightarrow$ Coding	Partition	$-10.7$	$-6.0$	$+1.3$	$-4.0$
Qwen3-1.7B	IF $\rightarrow$ Math	Partition	$-16.4$	$-4.3$	$+1.1$	$+1.8$
		Round robin	$-11.6$	$-1.67$	$-3.3$	$-0.4$
Qwen3-4B	Coding $\rightarrow$ IF	Partition	$-4.0$	$-6.4$	$+0.4$	$+0.0$
		Round robin	$-0.8$	$+6.0$	$+0.6$	$+2.3$

Appendix

Table 12: Head-reset experiments after continual RL. After completing stage

\mathrm{II}

, we reset the heads updated during stage

\mathrm{I}

to their base-model values. We report the resulting test accuracy (%) changes. Stage-

\mathrm{II}

performance that does not decrease after the reset is highlighted in orange. The results suggest that the adaptations learned on the disjoint head sets are functionally localized.

Figure 10: Layer-wise distribution of selected attention-head groups for Qwen3-1.7B , computed on the training set of Math.

Figure 11: Layer-wise distribution of selected attention-head groups for Qwen3-1.7B , computed on the training set of Basic Coding.

Figure 12: Layer-wise distribution of selected attention-head groups for Qwen3-1.7B , computed on the training set of Instruction-following.

Figure 13: Layer-wise distribution of selected attention-head groups for Qwen3-4B , computed on the training set of Math.

Figure 14: Layer-wise distribution of selected attention-head groups for Qwen3-4B , computed on the training set of Basic Coding.

Figure 15: Layer-wise distribution of selected attention-head groups for Qwen3-4B , computed on the training set of Instruction-following.

Figure 16: Layer-wise distribution of selected attention-head groups for Mistral-7B-Instruction-v0.3 , computed on the training set of Instruction-following.

Panel of Figure 17 — Figure 17: Sparse RL training accuracy curves of Qwen3 series for Math and IF tasks. The learning rate for top-50% attention-head tuning is fixed at $3\times 10^{-6}$ , while bottom-50% attention-head tuning is evaluated with learning rates of $3\times 10^{-6}$ , $5\times 10^{-6}$ , and $7\times 10^{-6}$ . Top-50% attention-head tuning is shown as a horizontal line with an error bar indicating mean $\pm$ std over three random rollout-seeds.

Continual Learning for LLMs

Continual learning for large language models studies how models absorb new knowledge and skills while retaining previously learned capabilities.

See all 231 papers

Can Capability Enhancement be Localized?
Oct 15, 2026Anh Tuan Nguyen, Anna Bair, Zico KolterAttention Head Analysis Multi-Head Attention
Carnegie Mellon University
ReCoLoRA: Spectrum-Aware Recursive Consolidation for Continual LLM Fine-Tuning
Jul 4, 2026Wentao LuContinual Learning for LLMs Low-Rank Adaptation
When Can Attention Heads Be Statically Defined?
Sep 28, 2026Weixian Waylon Li, Yintao Tai, Marcio Fonseca +1Attention Head Analysis Efficient Attention
University of Edinburgh, United Kingdom · Chamber of Deputies, Brazil

Oct 15, 2026J/K move · Enter open · S save

Can Capability Enhancement be Localized?

Anh Tuan Nguyen, Anna Bair, Zico Kolter

Carnegie Mellon University

Attention Head Analysis Multi-Head Attention

Open paper AskPDF arXiv

Transformer-based language models have shown strong capabilities across many tasks and domains. However, prior work shows that, for a given task, not all attention heads matter equally. Specifically, they suggest that many redundant heads can be safely removed with little-to-no effect, while a small set of important heads, when ablated, severely damages the model's performance on the given task. In this work, we investigate a complementary aspect, called Capability Enhancement Localization (CEL): given a model and a task, does there exist a small set of heads whose removal improves the task performance? We first investigate this question empirically across models from different families. Using a simple two-stage routine that screens candidate heads and evaluates their joint ablations, we identify sparse head combinations whose removal improves task-specific performance without further fine-tuning. Motivated by these observations, we ask theoretically whether such a phenomenon can naturally arise from training dynamics, and under what conditions. To this end, we construct a simplified multi-head softmax attention model and two tasks with conflicting features, and show that gradient descent from small random initialization provably concentrates the conflicting features into particular heads. We further show that ablating such heads improves performance on one task while harming the other, even though training converges to a global optimum of the joint objective. This establishes a possible training-induced theoretical explanation for CEL.

Abstract

Continual Learning for LLMs

Explore similar work

Can Capability Enhancement be Localized?

ReCoLoRA: Spectrum-Aware Recursive Consolidation for Continual LLM Fine-Tuning

When Can Attention Heads Be Statically Defined?