Capability Localization in Transformers

• Capability localization is the process of identifying sparse internal components, such as attention heads or neurons, that selectively modulate specific high-level capabilities in Transformer models.
• Compressed sensing and neuron scoring methodologies efficiently isolate task-specific subcomponents, leading to significant performance drops on targeted benchmarks upon their ablation.
• Empirical results demonstrate that localized modules can robustly control skills like math reasoning, code generation, and linguistic behaviors while minimally affecting unrelated tasks.

Capability localization is the identification of sparse internal model components whose targeted intervention selectively modulates a high-level capability. In current Transformer research, the localized components are typically attention heads or MLP neurons, and the capabilities include mathematical reasoning, code generation, and linguistic behaviors. Recent work frames the topic as a shift away from localizing individual facts toward localizing recurring computational structure across datasets: one line of work argues that individual knowledge cannot be localized with high fidelity and reliability, while another shows that many capabilities are highly localized to small subsets of attention heads and can be found efficiently through compressed sensing (Huang et al., 28 Feb 2025, Bair et al., 11 Feb 2026).

1. Conceptual scope and relation to knowledge localization

A central distinction in the literature is between individual knowledge and capability. Individual knowledge refers to the encoding of a single factual triplet $(s,r,o)$, whereas commonality refers to the shared inductive structure across many samples, such as the chain of reasoning required to solve grade-school math problems or detect positivity in movie reviews. Huang et al. report that prior attempts to localize single facts through dispersed parameters, parameter layers, or parameter chains lack both fidelity and reliability when tested under prompt rewrites and causal intervention. Their reported fidelity criterion asks whether identified parameters remain the same when the same fact is paraphrased several times; existing fact-localization methods achieve only $7$–$37\%$ overlap. Their reliability criterion asks whether zeroing, doubling activations, or editing layers consistently strengthens or weakens the fact in the output; these interventions often behave no better than random baselines (Huang et al., 28 Feb 2025).

By contrast, the head-level study of Transformer capabilities treats a task as a measurable behavioral function of the model and asks whether a small set of attention heads disproportionately controls that function. In this setting, a capability is localized if ablating a small number of heads causes a large drop on the target benchmark while largely preserving unrelated performance. The reported results show that zeroing out as few as five task-specific heads can degrade the capability of interest by up to $65\%$ while largely preserving performance on unrelated tasks, suggesting a modular organization in which specialized capabilities are implemented by sparse, functionally distinct components (Bair et al., 11 Feb 2026).

This division of labor motivates the modern use of the term. Capability localization does not primarily ask where a single proposition is stored; it asks which internal subcomponents repeatedly instantiate a computational skill across many examples.

2. Formal problem statements and units of localization

At the attention-head level, the formalization begins with a pre-trained Transformer LLM $M$ with $L$ layers and $H$ heads per layer, for a total of $N=L\cdot H$ attention heads $h_{l,i}$. A capability or task $T$ is associated with an evaluation set $7$0 and a baseline accuracy $7$1. For each head $7$2, the marginal performance drop under knockout is defined as

$7$3

The goal is to identify the $7$4 heads whose individual $7$5 values are largest. The compressed-sensing formulation introduces an unknown contribution vector $7$6 with $7$7 and assumes that $7$8 is $7$9-sparse because only a small number of heads matter for the capability. Random multi-head ablations are encoded by a binary matrix $37\%$0, producing measurements

$37\%$1

where $37\%$2 and $37\%$3 captures non-linear residuals. Recovery is then performed with Lasso:

$37\%$4

Large negative entries in $37\%$5 indicate heads whose removal lowers accuracy the most (Bair et al., 11 Feb 2026).

At the neuron level, Huang et al. define a capability-specific dataset $37\%$6 and score each MLP neuron by its aggregate gradient contribution to the correct answer across the dataset. A capability neuron is a hidden unit whose aggregate contribution is consistently large after discounting global mean and variance. The method, called Commonality Neuron Localization (CNL), thresholds neuron scores according to

$37\%$7

where $37\%$8 and $37\%$9 are the empirical mean and variance of scores over all MLP neurons, and $65\%$0 is the default threshold, with ablations also reported at $65\%$1 and $65\%$2. Consistency across random dataset splits is measured by overlap and IoU between the selected neuron sets (Huang et al., 28 Feb 2025).

Taken together, these formulations suggest that capability localization is not confined to a single representational granularity. It appears at least in attention heads and in MLP neurons, although the two papers operationalize the target in different ways.

3. Dataset-level localization of capability neurons

CNL localizes neurons by exploiting repetition across a dataset rather than attribution for a single instance. The procedure chooses a capability-specific dataset, splits it into two equal subsets to gauge consistency, computes a score for each hidden neuron via integrated gradients over the probability of the correct next token, applies the thresholding rule above, and then measures overlap and IoU between the two selected sets. If consistency is high, the union of selected neurons is treated as the capability-neuron set, and causal tests are performed by enhancement or erasure (Huang et al., 28 Feb 2025).

On GSM8K with Llama 2-7B, the reported localization fidelity is overlap $65\%$3, IoU $65\%$4, with the selected neurons comprising $65\%$5 of all MLP parameters. The same report states that prior fact-localization methods KN, ROME, and KC achieve only $65\%$6–$65\%$7 overlap and identify much larger fractions of the network, namely $65\%$8–$65\%$9. These results are used to support the claim that individual knowledge cannot be localized faithfully, whereas dataset-level commonality can be localized as capability (Huang et al., 28 Feb 2025).

The causal tests are equally central. In the enhancement setting on GSM8K with Llama 2-7B and $M$0, fine-tuning only the located $M$1 neurons for one epoch yields an average $M$2 pp gain on GSM8K, compared with random $M$3 pp$M$4 and “w/o located” $M$5 pp$M$6. At epoch $M$7, the gain remains $M$8 pp relative to random; GPTJ-6B shows similar patterns. In the erase setting, zeroing the located neurons causes an average $M$9 pp drop on GSM8K, whereas random erasure produces $L$0 pp (Huang et al., 28 Feb 2025).

The cross-data experiments further constrain the interpretation. Using GSM8K-identified neurons, enhancement improves Meta_Math by $L$1 pp and erasure reduces it by $L$2 pp, while effects on IMDB and Emotion are negligible. The reciprocal experiments with Emotion or Code25K neuron sets show that each localized set primarily affects its own capability domain. Pairwise overlaps indicate that tasks sharing sub-skills, such as multiple-choice comprehension for GSM8K and Emotion, share neuron subsets with overlaps around $L$3–$L$4, while dissimilar tasks are nearly orthogonal (Huang et al., 28 Feb 2025).

These results situate CNL as a dataset-conditioned method for identifying sparse MLP subspaces associated with transferable skills rather than isolated facts.

4. Compressed-sensing localization of attention heads

The compressed-sensing approach identifies task-specific attention heads through strategic knockouts and a small number of model evaluations. Each measurement corresponds to ablating a subset of heads $L$5, running the model on the evaluation set, and recording the resulting accuracy. Two constructions of $L$6 are used: Bernoulli sampling, where $L$7 i.i.d., and stratified sampling, which enforces approximately equal ablation frequency for every head. Under standard RIP-type assumptions, the paper states that exact recovery of a $L$8-sparse $L$9 is possible with $H$0 measurements; empirically, $H$1–$H$2 suffices even for $H$3 (Bair et al., 11 Feb 2026).

The experiments cover Llama-3.1-8B $H$4 heads$H$5, Llama-3.2-3B $H$6, Llama-3.2-1B $H$7, Qwen-2.5-7B $H$8, and Qwen-2.5-3B $H$9. The target capabilities are mathematical reasoning on GSM8K and Arithmetic, code generation on MBPP and HumanEval, and linguistic behaviors on Swearing and Rhyming. Specificity is assessed with HellaSwag, BoolQ, ARC-Challenge, and MMLU. Identification uses 100-sample subsets for efficiency, with final evaluation on full test sets (Bair et al., 11 Feb 2026).

Ablating only five stratified-CS heads yields the following maximum target-task degradations while largely preserving unrelated performance:

Capability	Maximum target-task drop	General-task change
Math (GSM8K)	$N=L\cdot H$0	$N=L\cdot H$1–$N=L\cdot H$2
Code (MBPP)	$N=L\cdot H$3	$N=L\cdot H$4
Swearing	$N=L\cdot H$5	$N=L\cdot H$6
Rhyming	$N=L\cdot H$7	$N=L\cdot H$8

The same ablations that cripple the target tasks cause near-zero degradation, around $N=L\cdot H$9 to $h_{l,i}$0, on HellaSwag, BoolQ, ARC, and MMLU. Heads found on GSM8K also degrade Arithmetic, and MBPP heads also degrade HumanEval, indicating dataset-level generalization within a capability family (Bair et al., 11 Feb 2026).

The method is also presented as efficient. A greedy one-shot procedure requires $h_{l,i}$1 evaluations to obtain about a $h_{l,i}$2 drop on GSM8K, whereas stratified compressed sensing with $h_{l,i}$3 evaluations achieves $h_{l,i}$4. Random-Bernoulli compressed sensing with $h_{l,i}$5 achieves $h_{l,i}$6, and the stratified variant with $h_{l,i}$7 matches the greedy result. The ablation curves show diminishing returns after roughly five heads: for one math case, ablating the first head causes about a $h_{l,i}$8 drop, the second adds about $h_{l,i}$9, and later heads contribute progressively less (Bair et al., 11 Feb 2026).

5. Empirical regularities and modular organization

Across five models and four distinct capabilities, the head-level study concludes that only a handful of attention heads are responsible for high-level skills and that Transformer architectures appear to implement specialized computations in sparse, functionally distinct subcomponents. The neuron-level study arrives at a related claim from a different angle: recurring capability structure produces stable, compact neuron masks, whereas individual facts do not. Read jointly, these results suggest a modular, sparse organization in which capabilities can be localized at multiple internal levels, even though the papers analyze different components and intervention types (Bair et al., 11 Feb 2026, Huang et al., 28 Feb 2025).

The specificity results are important for interpretation. Large target-task drops accompanied by small changes on unrelated benchmarks imply that the localized subcomponents are not merely high-activation or high-salience units in a generic sense. Likewise, the cross-dataset transfers in both papers indicate that localization is not restricted to a single benchmark: GSM8K heads generalize to Arithmetic, MBPP heads generalize to HumanEval, and GSM8K neurons transfer strongly to Meta_Math while leaving IMDB and Emotion largely unchanged (Bair et al., 11 Feb 2026, Huang et al., 28 Feb 2025).

The literature also records exceptions to a naive modularity view. The head-level study identifies a small set of “universal” heads, such as L1H29 in Llama-8B, whose ablation degrades all tasks broadly and produces degenerate output modes. It also reports that smaller models sometimes rely on shared format-level heads, including knowledge-MC heads in WMDP and MMLU at 3B and 1B scale, rather than fully task-specific heads. This indicates that localization need not imply complete isolation, uniqueness, or the absence of shared infrastructure (Bair et al., 11 Feb 2026).

A common misconception is therefore that capability localization means a capability is stored in exactly one place. The reported findings are more qualified: the dominant contributors are sparse, but there can be shared heads, universal heads, and distributed residual interactions.

6. Applications, limitations, and open directions

Both research lines place capability localization within interpretability, intervention, and safety. The head-localization paper identifies three application classes: interpretability, by pinpointing where a capability “lives”; targeted model editing or unlearning, by removing or modifying specific skills without wholesale finetuning; and AI safety, by disabling hazardous capabilities such as malicious code generation or toxic language through head ablation. The neuron-localization paper presents CNL as relevant to parameter-efficient fine-tuning, modular capability editing, and more interpretable deployment of LLMs (Bair et al., 11 Feb 2026, Huang et al., 28 Feb 2025).

The current methods also have explicit caveats. The compressed-sensing formulation relies on a first-order additive approximation, with higher-order interactions treated as noise. It identifies only head-level contributions, while lower-level circuits, neurons, and MLP layers may also matter. Its hyperparameters require tuning, with reported settings of $T$0–$T$1, $T$2–$T$3, and a regularization parameter $T$4, and the empirical validation is limited to instruction-tuned Llama and Qwen models with GQA. The neuron-localization framework is dataset-conditioned by construction, uses integrated gradients with $T$5 integration steps, and thresholds scores with a chosen $T$6 value; this means that the discovered neuron set is tied to a particular capability dataset and scoring protocol (Bair et al., 11 Feb 2026, Huang et al., 28 Feb 2025).

The stated future directions are correspondingly technical. For attention heads, the proposed agenda includes deeper mechanistic analysis of the identified heads, including what queries, keys, and values they attend to; extension to other components such as FFN layers and neurons or to other modalities such as vision and audio; and provable analysis of head interactions beyond first order. For dataset-level neuron localization, the results motivate further work on modular capability editing and parameter-efficient fine-tuning grounded in stable capability-specific subspaces (Bair et al., 11 Feb 2026, Huang et al., 28 Feb 2025).

As presently defined in the LLM interpretability literature, capability localization is therefore a causal and operational notion: a capability is localized to the extent that a small, identifiable subset of internal components can be found whose manipulation predictably changes that capability while leaving unrelated behavior comparatively intact.