PaperlayerView exact revision on arXiv

Source-linked AI summary

Can Capability Enhancement be Localized?

Anh Tuan Nguyen, Anna Bair, Zico Kolter

arXiv:2610.21294v1cs.LG

TL;DR

The paper asks whether a sparse set of attention heads can be removed to improve a model’s performance on a task, complementing work on redundant and task-critical heads. It searches candidate ablations empirically and analyzes a simplified attention model with conflicting tasks. Sparse ablations improve held-out accuracy in the evaluated settings, while the theory gives one possible training-induced mechanism under specific assumptions.

  • Problem

    Prior work identifies redundant and task-critical heads, while the paper asks whether a sparse head ablation can instead improve task performance without further fine-tuning.

    Source: §1 · p. 2

  • Method

    The paper uses a two-stage empirical search and studies training dynamics in a simplified softmax attention model jointly trained on two tasks with conflicting features.

    Source: §1 · p. 2, §4.1 · p. 6

  • Results

    Sparse ablations improve held-out accuracy across the evaluated settings, including a 13.10 percentage-point gain for SmolLM3 Arithmetic, from 70.40% to 83.50%.

    Source: §3.3 · p. 4, §6 · p. 9

  • Takeaways & Limitations

    The theory establishes one possible mechanism in which gradient descent concentrates features in heads, so removing a conflicting-feature head can help one task while harming another at a joint-objective optimum.

    Source: §6 · p. 9

    The theoretical mechanism is established only under fixed feature encoding, diagonal projections, small Gaussian initialization, and full-batch gradient descent on a noiseless population objective, and is not shown to explain the empirical gains.

    Source: §5 · p. 9

Abstract

from arXiv · show

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 \textit{redundant} heads can be safely removed with little-to-no effect, while a small set of \textit{important} heads, when ablated, severely damages the model's performance on the given task. In this work, we investigate a complementary aspect, called \textit{Capability Enhancement Localization} (CEL): \textit{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 \textit{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 \textit{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.

1 INTRODUCTION

Prior work shows that attention heads can be redundant or task-critical, motivating the question of whether removing a sparse set can instead improve task performance. The paper investigates this Capability Enhancement Localization question empirically across models and theoretically through a simplified attention model.

1 INTRODUCTION

  • The paper frames CEL as a complement to findings on redundant and task-critical heads, which can respectively be removed with little effect or harm task performance when ablated.Prior work also identifies specialized head roles, including induction, long-context retrieval, arithmetic processing, and preferences between internal knowledge and external context.Source: §1 · pp. 1–2
  • Capability Enhancement Localization (CEL) asks whether ablating a sparse set of attention heads can improve task performance without further fine-tuning.The question complements findings that some heads are redundant while others are important for particular tasks.Source: §1 · pp. 1–2
  • The empirical study finds evidence that sparse head ablations can significantly improve task-specific performance across multiple model families and task configurations.The work also seeks a theoretical explanation for how such beneficial ablations can arise during training.Source: §1 · p. 2
  • In a simplified softmax attention model trained on two tasks with conflicting features, gradient descent can concentrate those features in a particular head, whose ablation helps one task while harming the other, even at a joint-objective optimum.The authors present this as one possible mechanism for CEL.Source: §1 · p. 2

2 RELATED WORK

Related work examines attention-head redundancy and specialized functions, task-critical components, and beneficial pruning in selected settings. This paper focuses on whether sparse head removal can enhance task performance more broadly and on a possible training-based explanation.

Attention-head functionality

  • Prior studies report that sparse task-critical head ablations can substantially degrade target-task performance while largely preserving unrelated-task performance.Other localization work studies components associated with factual associations, task-specific skills, and shared capabilities.Source: §2 · p. 2

Performance enhancement through head pruning

Mechanistic interpretability and circuit discovery

  • Mechanistic interpretability can localize computations, but such localization need not predict the effect of changing a model; the experiments measure ablation effects without identifying removed heads’ computations.The paper’s simplified theory instead shows how training can produce a head whose removal benefits one task.Source: §2 · p. 3

3 EMPIRICAL MOTIVATION: SPARSE BENEFICIAL HEAD ABLATION

The section defines beneficial head ablation as a positive task-specific performance gain from removing a sparse set of heads without fine-tuning, and describes a two-stage search with held-out evaluation. Across the tested model–task pairs, the selected ablations improve held-out accuracy; a simplified theory offers a possible mechanism but does not explain the empirical language-model gains.

3.1 HEAD ABLATION AND EVALUATION

  • The intervention targets at most k heads from the model’s total set, with k typically set to 3.The total number of heads is the number of layers times the number of heads per layer.Source: §3.1 · p. 3
  • A head subset is beneficial when its ablation yields a positive empirical gain, measured as ablated-model accuracy minus intact-model accuracy under the same task protocol.The ablation zeros selected heads’ outputs before the attention output projection; remaining parameters are unchanged and no fine-tuning is performed.Source: §3.1 · p. 3

3.2 EXPERIMENTAL SETUP

  • The two-stage search screens individual heads, evaluates joint subsets from a small candidate pool, selects by discovery-set ranking, then freezes the subset for held-out evaluation.Selection uses discovery examples disjoint from test examples; the evaluated models are Llama-3.2-3B-Instruct, Qwen2.5-3B-Instruct, and SmolLM3-3B-Instruct, across math, code generation, and language-skill tasks.Source: §3.2 · pp. 3–4

3.3 RESULTS

  • The simple search may miss the best subset among all heads, and the simplified theory supplies a possible mechanism without identifying what caused the language-model gains.The search considers subsets drawn from heads selected by individual-head effects.Source: §3.3 · pp. 4–5
  • Held-out accuracy improves across Arithmetic, BoolQ, HellaSwag, and MBPP; SmolLM3 Arithmetic rises from 70.40% to 83.50%, a gain of 13.10 percentage points.The gains vary across model–task pairs, so the results support beneficial ablations in the tested settings, not every model–task pair.Source: §3.3 · p. 4, Table 1 · p. 5

4 A THEORETICAL EXPLANATION OF THE CAPABILITY ENHANCEMENT LOCALIZATION PHENOMENON

The paper studies whether conflicting tasks can induce feature localization across attention heads, making head ablation beneficial for one task despite joint optimality. A simplified softmax model shows how small-random-initialization gradient descent can produce this outcome under specific conditions.

4.1 TASKS, ATTENTION MODEL, AND TRAINING

  • Task A ignores coordinate 0 while task B needs it, creating conflict; coordinates 1 and 2 carry task-specific signals under noiseless population risks.The analysis therefore isolates conflict and training dynamics rather than finite-sample overfitting.Source: §4.1 · p. 5
  • The analysis uses a three-head attention model with fixed feature directions, paired memory tokens, diagonal projections, and joint population-risk training.The fixed encoding supplies the features; training determines their allocation across heads.Source: §4.1 · pp. 5–6

4.2 JOINT OPTIMALITY DOES NOT DETERMINE HEAD ALLOCATION

  • The predictor is linear in the input, but its effective coefficients are nonlinear functions of trainable weights, and their optimal aggregate does not fix their allocation across heads.The shared coordinate’s joint optimum is a/2, while task-specific coefficients b and c are fit exactly.Source: §4.2 · p. 7

4.3 GRADIENT DESCENT LOCALIZES FEATURES

  • Under small Gaussian initialization and sufficiently small fixed-step gradient descent, each feature becomes concentrated in one head while contributions from other heads remain small.The winning-head assignments depend on initialization; localization alone does not ensure beneficial ablation.Source: §4.3 · p. 8

4.4 LOCALIZATION ENABLES BENEFICIAL HEAD ABLATION

  • With probability 4/9, deleting the head carrying the conflicting feature reduces task A’s risk to zero but raises task B’s risk to 4t^2 and the joint objective by t^2/2.These conclusions hold in the small-initialization limit and also at the limiting trained model.Source: §4.4 · pp. 8–9

5 DISCUSSION AND LIMITATIONS

The empirical findings are limited to evaluated models, datasets, and inference protocols, and the stability of selected subsets and gains remains insufficiently characterized. The theoretical account is a possible mechanism under restrictive modeling and training assumptions, not an established explanation of the empirical gains.

5 DISCUSSION AND LIMITATIONS

  • The empirical evidence covers three model families and four benchmarks, so it does not establish beneficial ablations for every model–task pair.Stability across discovery samples remains insufficiently characterized despite disjoint selection and evaluation examples.Source: §5 · p. 9
  • The two-stage search does not guarantee an optimal subset and may miss combinations whose benefits appear only through joint ablation.Gains depend on discovery data and the candidate pool; an unsuccessful search does not establish that no beneficial sparse intervention exists.Source: §5 · p. 9
  • The theoretical mechanism assumes fixed feature encoding, diagonal projections, small Gaussian initialization, and full-batch gradient descent on a noiseless population objective.It does not establish that this mechanism explains evaluated language-model gains or all instances of CEL.Source: §5 · p. 9

6 CONCLUSION AND FUTURE WORK

The paper reports empirical CEL across three model families and four benchmarks, and presents a training-dynamics mechanism in a simplified attention model. It identifies empirical and theoretical directions that remain open, and describes the authors’ use of AI tools.

6 CONCLUSION AND FUTURE WORK

  • Gradient descent from small Gaussian initialization can concentrate features in individual heads, allowing removal of a head carrying a conflicting feature to help one task while harming another, even at a joint-objective global minimum.The empirical search also finds sparse head subsets whose ablation improves held-out accuracy across the evaluated models and benchmarks.Source: §6 · p. 9
  • Empirical follow-up could improve search for interacting heads; theoretical follow-up could extend beyond fixed feature encoding and diagonal parameterization to learned representations and feature interactions.The paper describes its empirical and theoretical investigation as initial.Source: §6 · p. 9

AI USE STATEMENT

  • The authors used AI tools for experiment implementation, theoretical brainstorming and proof assistance, and wording and grammar edits, while reporting that they verified the theoretical results.They say they drafted paragraphs without LLMs, revised the AI-assisted wording, and supplied citations and references by hand.Source: AI USE STATEMENT · p. 10

A ADDITIONAL INFORMATION FOR EXPERIMENT SETTING

The additional experiment-setting information identifies model details and dataset sizes. It defines the head-count notation and notes SmolLM3’s evaluation mode.

A ADDITIONAL INFORMATION FOR EXPERIMENT SETTING

  • Table 2 gives model details and defines H as query heads per layer and N = LH as the total number of ablatable heads.The table caption identifies the table’s scope and notation.Source: Table 2 · p. 13
  • SmolLM3 is evaluated in non-thinking mode for convenience.Source: Table 2 · p. 13
  • Table 3 reports dataset sizes for discovery Ddist and held-out evaluation Dtest.Source: Table 3 · p. 13

B ADDITIONAL RESULTS AND OMITTED PROOFS FOR SECTION 4

This section provides detailed proofs for the results presented in Section 4.

B ADDITIONAL RESULTS AND OMITTED PROOFS FOR SECTION 4

B.0.1 OMITTED PROOF FOR SECTION 4.2

The omitted proof derives head-feature responses, the aggregate predictor, task-specific population risks, and the joint objective for the three-head model. It expresses the predictor as linear in input features while its effective coefficients remain nonlinear functions of trainable weights.

B.0.1 OMITTED PROOF FOR SECTION 4.2

  • For each head h and feature j, the head’s contribution is x_j c_hj e_j, with c_hj = o_hj v_hj tanh(q_hj k_hj).The proof derives the response from the attention mechanism and output projection.Source: §B.0.1 · pp. 13–14
  • Summing head-feature coefficients gives w_j, so the model prediction is fθ(x) = w⊤x, linear in inputs but with effective coefficients that are nonlinear functions of trainable weights.The proof identifies this distinction as important for the later localization result.Source: §B.0.1 · pp. 13–14
  • The joint population objective is formed by averaging the two task risks, and the proposition organizes the proof around head responses, predictor form, population risks, and the joint optimum.The supplied passage introduces the objective but does not state its expanded form or optimum values.Source: §B.0.1 · p. 13
  • The task-specific population risks are R_A(w) = w_0^2 + (w_1 − b)^2 and R_B(w) = (w_0 − a)^2 + (w_2 − c)^2.Source: §B.0.1 · pp. 13–14

B.1 ADDITIONAL RESULTS AND OMITTED PROOFS FOR SECTION 4.3

With small initialization and gradient-flow dynamics, each feature develops a unique early leader among its heads; the appendix characterizes this competition through finite escape times.

B.1 ADDITIONAL RESULTS AND OMITTED PROOFS FOR SECTION 4.3

  • For each feature, Lemma B.1 studies three heads’ four-parameter quartets initialized at small scale α, with nonzero coordinates and one uniquely smallest magnitude.Source: §B.1 · pp. 14–15
  • Each auxiliary ODE has a finite escape time, and the parameter product eventually becomes positive as the weights balance near escape.Source: §B.1 · p. 15
  • For each feature, the three heads’ escape times are almost surely distinct, giving a uniquely defined earliest leader head.Source: §B.1 · pp. 14–15

3. The variable ˆh0, ˆh1, ˆh2 are independent across features

This section analyzes how discrete gradient descent under small initialization turns early head competition into feature localization, using escape-time, approximation, and stability analyses. It establishes independent, uniform winning-head assignments across the three heads, convergence to exact fitting, and a finite-checkpoint residual bound.

Lemma B.1

  • For each feature, three heads almost surely have a unique early leader, but this gradient-flow result alone does not show that the leader permanently wins.The early leader is identified by the heads’ different escape times; the remark explicitly limits the result to continuous-time analysis.Source: §3 · p. 17

Lemma B.2 (Discrete GD follows the auxiliary dynamics before escape)

  • Before the selected heads’ escape times, normalized discrete gradient descent tracks the auxiliary ODE with error O(α^4 + h_α), which is O(α^2) for fixed η, p, and r.The result requires a positive residual bounded above and a time horizon strictly before the selected heads’ escape times; step sizes need only be positive and bounded as specified.Source: §3 · p. 17

Lemma B.3 (No residual overshoot)

  • Under the stated weight bounds and sufficiently small step size, each GD update preserves a positive residual and does not increase it.The lemma’s conclusion is that the next residual remains between half the current residual and the current residual.Source: §3 · p. 19

Lemma B.4 (Bounded invariant region for the winning head)

  • A candidate winning head remains in a bounded invariant region if it starts with positive, approximately balanced parameters and the stated conditions on the next update hold.The result keeps all four updated parameters below R and within [−B, B].Source: §3 · p. 19

Lemma B.5 (Remaining residual clock)

  • Once a winning head is established, the remaining residual-weighted training time is finite, leaving insufficient time for losing heads to catch up and localizing each feature to a particular head.This conclusion combines the remaining-clock bound with the preceding leading-head, discrete-GD tracking, and stability results.Source: §3 · p. 21

3. The variable ˆh0, ˆh1, ˆh2 are independent across features

  • The parameter trajectory is bounded and converges to a finite limit that exactly fits each feature, while each residual stays positive and non-increasing.At convergence, w∞ = β for each feature.Source: §3 · pp. 21, 23
  • Each feature selects one of three heads, with independent uniform winner labels across features; the losing heads remain bounded during training.The localization result holds for sufficiently small initialization and a fixed step size below a threshold depending only on the model parameters.Source: §3 · pp. 21, 23, 24
  • For any fixed M > 0, the residual at the stated finite checkpoint can be bounded by α^M for sufficiently small α.The checkpoint is defined as Nα = ⌈log(1/α) ηα²⌉.Source: §3 · pp. 21, 24

B.2 ADDITIONAL RESULTS AND OMITTED PROOFS FOR SECTION 4.4

Under a specified winner-allocation event, the proof’s intervention deletes h0 without changing other parameters, and the reported bounds are positive for sufficiently small initialization.

B.2 ADDITIONAL RESULTS AND OMITTED PROOFS FOR SECTION 4.4

  • Pr(E)=4/9; under Theorem 4.2’s assumptions with t=a/2, the prescribed intervention deletes h0 without changing other parameters.Source: §B.2 · p. 25
  • On E, every head other than h⋆ loses feature 0, while h⋆ loses features 1 and 2.Source: §B.2 · p. 25
  • G_A=t^2+O_G(α^4) and D_B=3t^2+O_G(α^4); both are positive for sufficiently small α, and D_B measures task-B ablation damage.Source: §B.2 · p. 25

Related papers

Loading 2610.21294v1…