Can Capability Enhancement Be Localized to Specific Attention Heads?

Given a model and a task, is there a small set of attention heads whose deletion, without fine-tuning, improves performance on that task? The authors call this Capability Enhancement Localization (CEL).

• Task scenario: The authors state that the same parameter set can often support multiple tasks via natural language instructions or a few examples, so it is necessary to understand the contribution of components such as attention heads to individual tasks.
• Existing knowledge: The authors state that many heads can be pruned after training with little change in overall performance; there are also a few task-critical heads whose ablation harms the corresponding task. CEL asks the opposite direction.
• Existing special cases: Held and Yang (2023), Jin et al. (2024), and Nguyen et al. (2025) observed improvements from head pruning in target language, knowledge-conflict QA, or mathematical reasoning. The authors want not to presuppose a source of interference and to search for such sparse ablations across broader tasks and models.
• Two things this paper does: Use a disjoint discovery set for a two-stage search, with a head budget typically 3, then measure gains on a held-out set; and, in a simplified attention model with feature conflict, provide a mechanism induced by gradient descent.

The introduction and Related Work divide the discussion into three categories. The authors only write the differences from this paper when explicitly comparing.

Redundancy, specialization, and capability localization

• Removable heads: Michel et al. (2019) and Voita et al. (2019) show that many attention heads can be removed after training with little loss in task performance. The introduction also writes that Clark et al. (2019b) observed similar attention patterns and overlapping behavior within the same BERT layer.
• Specialized computation: Olsson et al. (2022) discuss induction mechanisms corresponding to in-context learning; Wu et al. (2025) discuss long-context retrieval; Zhang et al. (2024) discuss the handling of operands and operators in arithmetic.
• Localization studies: Meng et al. (2022), Panigrahi et al. (2023), and Huang et al. (2025) localize factual associations, post-fine-tuning task skills, and a small group of neurons related to shared capabilities, respectively. Bair et al. (2026) use LASSO to find task-critical attention heads; the authors state that jointly ablating a few of these heads reduces target task performance while roughly preserving unrelated task performance.

Performance improvements after head pruning

• Cross-lingual: Held and Yang (2023) identify heads related to cross-lingual interference, and pruning them can improve target language performance.
• Knowledge conflicts: Jin et al. (2024) study conflicts between internal knowledge and external context, and write that pruning selected heads can improve question answering. The introduction also writes that this work distinguishes heads biased toward internal weights or external context.
• Mathematical reasoning: Nguyen et al. (2025) observe that pruning certain heads can improve mathematical reasoning, and use different pruning configurations to generate diverse answers for best-of-N selection.

Circuits and implications of localization for intervention

• Circuits: Wang et al. (2022) identify an attention head circuit for indirect object identification; Conmy et al. (2023) use activation patching to automatically discover computational subgraphs.
• Localization does not equal editing location: Hase et al. (2023) find that causal tracing results do not indicate which MLP layers are best for factual editing.
• Authors' qualification of their own experiments: The authors state that head ablation measures the effect of a specified intervention and does not reveal the computation performed by the removed heads; the theory, in a simplified model, explains how training produces heads whose deletion benefits a task.

The experimental baseline is the full model under the same protocol, not another pruning algorithm. Task data are Arithmetic (Brown et al., 2020), MBPP (Austin et al., 2021), HellaSwag (Zellers et al., 2019), BoolQ (Clark et al., 2019a). The authors state that this paper turns to the enhancement regime, i.e., finding small head sets whose joint removal, without fine-tuning, improves task performance, and adds a training-dynamics explanation that prior work ignored.

The empirical part uses two-stage intervention selection to fix a small head set on the discovery set, then measures the ablation gain on the holdout set. The theory part analyzes gradient descent on three-head diagonal softmax attention with two feature-conflicting tasks.

Ablation definition and selection objective

• Head set: The model has L layers, H heads per layer, for N=LH total. An intervention S satisfies that the number of heads is at most k, and k is much smaller than N; this paper typically uses k=3.
• Ablation method: M\S zeroes the outputs of heads in S at every token, before the attention output projection. All other parameters are unchanged, with no fine-tuning.
• Gain: Accuracy is the proportion of examples on the dataset judged correct by the task scorer. The test gain is Δ_T,Dtest(S)=Acc_T,Dtest(M∖ S)−Acc_T,Dtest(M), i.e., the difference in accuracy between the ablated model and the full model. If the gain is greater than 0, S is said to be beneficial on that test set.
• Data isolation: Selection uses only a discovery set disjoint from the test set, yielding subsets with head counts between 1 and k, fixed before testing. The authors want the gain to come from the task, not just from fitting the discovery set.

Two-stage routine

• Ranking: Algorithm 1 takes candidate count p=10 and budget k=3. On the discovery set, it compares higher gain, then lower harm, then higher accuracy, then fewer heads, then lexicographic order by head identity. Harm is the number of examples the full model answers correctly and the ablated model answers incorrectly.
• First stage: Each head is evaluated individually on the discovery set, and the top p single-head interventions form the candidate pool.
• Second stage: Enumerate subsets of size 1 to k in that pool; those of size at least 2 are evaluated on the discovery set, reusing single-head results. The highest-scoring one is frozen.
• Holdout comparison: On the test set, compare the ablated model for that subset with one matched full-model run. Test examples do not participate in scoring or selection. The authors state that this search is limited to the pool selected by single-head effects and does not necessarily find the subset with the highest gain among all heads.

Simplified model and joint objective

• Two tasks: U, V, W are independent Rademacher variables. Fix a, b, c greater than 0, and set t=a/2. Task A's input is (U, V, 0) and target is bV; task B's input is (U, 0, W) and target is aU+cW. On coordinate 0, A's ideal prediction should ignore U, while B should give coefficient a.
• Architecture: One layer with three heads. The query for coordinate j is along a fixed direction, the memory pairs are the positive and negative vectors of that direction, and the three heads share a fixed mask. The authors state that this gives each head the same access to each feature, with allocation not predetermined.
• Diagonal coefficients: Query, key, value, and output projections are all diagonal. The effective coefficient is c_{hj}=o_{hj}v_{hj}\tanh(q_{hj}k_{hj}), i.e., the product of output, value, and the tanh of the query–key product. The three head coefficients sum to w, and the scalar prediction is the inner product of w with the input.
• Risk: Observations are noiseless, and the analysis uses population squared risk. Proposition 4.1 writes the two-task risk as the sum of squared coordinate errors, and the joint loss is one quarter of the sum of the two. The optimal aggregate coefficient is (a/2, b, c). The authors state that the objective constrains only the per-feature sums, not how the heads split them; deleting head h subtracts that head's coefficient vector from w.

Training and the two theorems

• Training: Four scalars per head–feature, initialized as small-scale α times an independent standard normal vector. Fixed step size, full-batch gradient descent. No sign conditions, parameter tying, or rebalancing during training.
• Localization: Theorem 4.2 states that for almost all initializations, sufficiently small α, and sufficiently small fixed step size, the trajectory is bounded and converges, and the aggregate coefficient tends to (a/2, b, c). Each feature has one winning head, with labels independent and uniform over the three heads; the other heads' contributions remain at order α to the fourth. Different features can share the same winning head, so localization alone does not guarantee that deleting a head is beneficial.
• Beneficial deletion: Theorem 4.3 defines event E as the winning head of the conflicting feature not simultaneously carrying the two task-specific features, and gives Pr(E)=4/9; the appendix writes this probability as (2/3) squared. On E, deleting that head without changing other parameters makes task A's risk decrease by about t squared and task B's risk increase by about 3t squared after the checkpoint.
• Authors' claim about the limit: The authors state that in the small-initialization limit, both full task risks tend to t squared; after deleting the selected head, task A's risk is 0 and task B's is 4t squared, and the joint objective increases by half of t squared. The authors state that this intervention improves one task by leaving the joint-optimal tradeoff.

Experimental setup

• Models: Llama-3.2-3B-Instruct, Qwen2.5-3B-Instruct, SmolLM3-3B-Instruct. Table 2 writes LLaMA-3.2-3B as 28 layers, 24 heads per layer, 672 total ablatable heads; Qwen2.5-3B and SmolLM3-3B are both 36 layers, 16 heads per layer, 576 total. SmolLM3 is evaluated in non-thinking mode.
• Data scale: In Table 3, the discovery/holdout sets are Arithmetic 2,000/1,000, HellaSwag 2,000/10,042, BoolQ 2,000/3,270, and MBPP 374/500. The holdout sets for HellaSwag and BoolQ are the official validation sets; the MBPP holdout set is the official test set; both Arithmetic parts are drawn from the official validation set of 20,000 items and are disjoint.
• Selection: For each model–task pair, Algorithm 1 is run on the discovery set with p=10 and k=3. Subsets are fixed before testing, and test examples do not participate in scoring or selection.
• Baselines and metrics: The baseline is the full model under the same task protocol. Table 1 reports holdout accuracy, Gain (percentage points), and Rel. error reduction. Correctness is determined by the task scorer.
• Runs: Each Slurm job uses one NVIDIA L40S, two CPUs, 24 GB system memory, and batch size 8.
• Not stated in the paper: number of repetitions, random seeds, layer and head numbers of selected heads, and the formula for Rel. error reduction. The paper does not state that an LLM judge is used.

Main results

Per Table 1. The left of the arrow is the full-model holdout accuracy, the right is the post-ablation accuracy, and the parentheses give the Gain (percentage points). SmolLM3 is in non-thinking mode.

| Task | SmolLM3 | Qwen2.5-3B-Instruct | Llama-3.2-3B-Instruct | | Arithmetic | 70.40%→83.50% (+13.10) | 89.50%→90.90% (+1.40) | 68.00%→74.30% (+6.30) | | BoolQ | 82.57%→84.92% (+2.35) | 80.21%→82.35% (+2.14) | 78.72%→81.25% (+2.54) | | HellaSwag | 75.75%→77.05% (+1.29) | 75.02%→76.66% (+1.63) | 71.64%→73.78% (+2.14) | | MBPP | 50.40%→52.00% (+1.60) | 60.40%→62.20% (+1.80) | 47.80%→53.00% (+5.20) |

• Authors' scope: The authors state that these interventions improve holdout accuracy on Arithmetic, BoolQ, HellaSwag, and MBPP, but the gains vary by model–task pair, so they do not claim that ablation helps every model on every task.
• Relative error reduction: Table 1 also reports Rel. error reduction. In that column, the minimum is 3.23% for SmolLM3 on MBPP, and the maximum is 44.26% for SmolLM3 on Arithmetic. The main text does not separately explain this column.
• Implication of the search: The authors state that the two-stage process does not necessarily find the subset with the highest gain among all heads; the results in the table establish an empirical phenomenon in the tested settings and do not explain why training produces these removable components.

Quantitative statements in the theory section

• Not another set of language-model experiments: Theorem 4.3 gives the probability of event E as 4/9. E is the event that the winning head carrying the conflicting feature is not among the winning heads of the two task-specific features.
• After the checkpoint: For almost all initializations falling in E, sufficiently small α, and iterations after the checkpoint, the risk of task A decreases as t squared plus quartic small terms, and the risk of task B increases as 3t squared plus quartic small terms.
• Authors' interpretation of the limit: The authors state that both full risks tend to t squared; after deleting the selected head, task A's risk is 0 and task B's is 4t squared, and the joint objective increases by half of t squared. The authors state that this leaves the joint-optimal tradeoff and do not assert that this mechanism is responsible for the gains in Table 1.

Other ablations and analyses

• The paper does not report numbers after varying p, k, discovery-set sampling, or learning rate.
• All 12 gains in Table 1 are positive; the table has no model–task pair with gain not greater than 0.
• The paper does not report the identities of the selected heads, so the specific layers and heads deleted cannot be verified from the main text.

Limitations and future directions noted by the authors

• Extrapolation: Results only concern the evaluated models, datasets, and inference protocols, and do not establish that a beneficial ablation exists for every model–task pair (Section 5).
• Discovery set stability: Selection and evaluation use disjoint examples, but the stability of the selected subset and its gains across different discovery samples remains under-characterized (Section 5).
• Search: The two-stage process does not guarantee an optimal subset; single-head screening may exclude combinations that only yield benefits in joint ablation; gains depend on the discovery data and candidate pool; search failure cannot establish that no beneficial sparse intervention exists (Section 5).
• Theoretical scope: The analysis is limited to fixed feature encodings, diagonal projections, small Gaussian initialization, and full-batch gradient descent on a noiseless population objective. The authors state it does not establish that this mechanism explains the gains in the evaluated language models, nor that it covers all CEL (Section 5).
• Search follow-up: The authors include as future work searches that account for head interactions, to address beneficial combinations that single-head screening might miss (Section 6).
• Theoretical follow-up: The authors include as future work characterizing conditions beyond fixed feature encodings and diagonal parameterization, and allowing interactions between learned representations and features (Section 6).

Experimental coverage

• The tested models are Llama-3.2-3B-Instruct, Qwen2.5-3B-Instruct, SmolLM3-3B-Instruct, and the tasks are Arithmetic, BoolQ, HellaSwag, MBPP, totaling 12 model–task pairs (Section 3.2, Table 1).
• Discovery and held-out set sizes are in Table 3; head budget k=3, single-head candidate count p=10, and the two data parts are disjoint (Algorithm 1, Section 3.1).
• SmolLM3 is evaluated in non-thinking mode; each job uses one NVIDIA L40S, batch size 8 (Section 3.2, Table 1, Table 2).
• The theoretical object is three heads, two equally weighted tasks, diagonal projections, population squared risk, and full-batch gradient descent (Section 4).
• The paper does not report the number of repetitions, the layers and indices of the selected heads, or results varying p or k. The experimental intervention is zeroing head outputs; the authors state in Section 2 that this does not reveal the computation of the deleted heads.

CEL asks: can deleting a small number of attention heads, without further fine-tuning, improve held-out performance on a given task.

• Question: Starting from redundant heads and task-critical heads, the authors instead look for sparse head sets whose deletion improves a specified task, without presupposing that interference comes from language, knowledge conflicts, or mathematical reasoning.
• Method: The discovery set and test set are disjoint. They first evaluate single heads and keep the top 10, then select among subsets of size 1 to 3, freezing before testing. The ablation zeroes these heads' outputs before the attention output projection, at every token.
• Held-out numbers: For three 3B instruction models on Arithmetic, BoolQ, HellaSwag, MBPP, all 12 gains in Table 1 are positive. SmolLM3's Arithmetic goes from 70.40% to 83.50% (+13.10 percentage points, Rel. error reduction of 44.26%). According to Table 1, the smallest gain cell is that model's HellaSwag at +1.29 percentage points. Llama-3.2-3B-Instruct's MBPP goes from 47.80% to 53.00% (+5.20 percentage points).
• Theory: Two equally weighted tasks have inconsistent ideal coefficients for the same input coordinate. Three-head diagonal softmax attention is trained with full-batch gradient descent from small Gaussian initialization. The authors state in Theorem 4.2 that each feature concentrates into one winning head; Theorem 4.3 states that with probability 4/9, deleting the head carrying only the conflicting feature reduces task A's risk and increases task B's risk, even though the aggregated coefficients are already optimal for the joint objective.
• Authors' conclusion: The authors state that joint optimization reaching optimality can still coexist with head deletion improving a single task. They also state that this is only one possible mechanism and do not claim it caused the gains in the aforementioned language models.