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.
Figures & tables
Model
Task
Intact acc.
Ablated acc.
Gain (pp)
Rel. error reduction
SmolLM3
Arithmetic
70.40%
83.50%
+13.10
44.26%
BoolQ
82.57%
84.92%
+2.35
13.51%
HellaSwag
75.75%
77.05%
+1.29
5.34%
MBPP
50.40%
52.00%
+1.60
3.23%
Qwen2.5-3B-Instruct
Arithmetic
89.50%
90.90%
+1.40
13.33%
BoolQ
80.21
82.35
+2.14
10.82%
Table 1: Held-out performance of head ablations. SmolLM3 is evaluated in non-thinking mode.
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Model
Layers L
Heads H
Total Heads N
LLaMA-3.2-3B
28
24
672
Qwen2.5-3B
36
16
576
SmolLM3-3B
36
16
576
Appendix
Table 2: Model details. Here, H denotes query heads per layer, and N=LH is the total number of ablatable heads. In our experiments, for convenience, SmolLM3 is evaluated in non-thinking mode.
Dataset
Discovery
Held-out
Arithmetic
2,000
1,000
HellaSwag
2,000
10,042
BoolQ
2,000
3,270
MBPP
374
500
Appendix
Table 3: Dataset sizes for discovery Ddist and held-out evaluation Dtest .
Large language models (LLMs) exhibit a wide range of capabilities, including mathematical reasoning, code generation, and linguistic behaviors. We show that Transformer architectures contain small subsets of attention heads that are necessary for certain capabilities. Zeroing out as few as five task-specific heads can degrade performance by up to 60% on standard benchmarks measuring the capability of interest, while largely preserving performance on unrelated tasks. We introduce a compressed sensing-based method that exploits the sparsity of these heads to identify them via strategic knockouts and a small number of model evaluations. We validate these findings across Llama and Qwen models ranging from 1B to 14B parameters and a diverse set of capabilities including mathematical abilities and code generation, revealing a modular organization in which specialized capabilities are dependent on sparse, functionally distinct components. Overall, our results suggest that capability localization is a general organizational principle of Transformer language models, with implications for interpretability, model editing, and AI safety. Code is released at https://github.com/locuslab/llm-components.
Multi-head attention enables transformer models to represent multiple attention patterns simultaneously. Empirically, head specialization emerges in distinct stages during training, while many heads remain redundant and learn similar representations. We propose a theoretical model capturing this phenomenon, based on the multi-index and single-location regression frameworks. In the first part, we analyze the training dynamics of multi-head softmax attention under SGD, revealing an initial unspecialized phase followed by a multi-stage specialization phase in which different heads sequentially align with latent signal directions. In the second part, we study the impact of attention activation functions on performance. We introduce the Bayes-softmax attention, which achieves optimal prediction performance in this setting.
M. Sagitova, O. Duranthon, L. Zdeborová
Statistical physics of computation laboratory, École Polytechnique Fédérale de Lausanne, Switzerland
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.
Ruizhe Shi, Zaiwen Yang, Qiubo Wang +2
University of Washington · Tsinghua University · University of Hong Kong +1