Head Attribution in Transformers

• Head attribution is a method to assign causal or explanatory responsibility to specific model heads, such as transformer attention or classifier heads.
• It employs techniques like head masking, integrated gradients, contrastive relevance, and spectral analysis to measure performance impact and guide model pruning.
• Recent advances highlight its applications in reasoning preservation, safety control through backdoor detection, and steering model behavior, with ongoing refinements addressing limitations.

Searching arXiv for the cited papers to ground the article in current sources. arXiv search query: (Sheth et al., 27 Apr 2026) OR (Tan et al., 20 Apr 2026) OR (Chen et al., 29 Jan 2026) OR (Xu, 22 May 2026) OR (Zhan et al., 10 Jun 2025) OR (Hao et al., 2020) OR (Yu et al., 26 Sep 2025) OR (Farvardin et al., 26 Mar 2026) Head attribution, in current machine-learning literature, denotes the assignment of explanatory or causal responsibility to a “head.” In transformer research this usually means an attention head, but related work also uses the term for classifier-head design and for dense attribution on human head images. Across these usages, head attribution has been instantiated as expected performance degradation under head masking, integrated-gradients attribution on self-attention matrices, contrastive relevance on wrong-versus-correct token preferences, influence-function scores restricted to head parameter subspaces, spectral signatures of specialized head circuits, behavior-specific steering scores, intrinsic class attribution in classifier heads, and canonical semantic coordinates for human head pixels (Sheth et al., 27 Apr 2026, Hao et al., 2020, Tan et al., 20 Apr 2026, Chen et al., 29 Jan 2026, Xu, 22 May 2026, Zhan et al., 10 Jun 2025, Farvardin et al., 26 Mar 2026, Pozdeev et al., 4 Nov 2025).

1. Definitions and problem formulations

A common formulation treats an attention head $h$ in layer $\ell$ as a functional unit whose importance is the loss increase caused by intervening on that head. In "Pruning via Causal Attribution Preserves Reasoning Performance in LLMs" (Sheth et al., 27 Apr 2026), the causal attribution score is

$$\Delta(h_\ell)=\mathbb{E}_{(x,y)\sim \mathcal{D}_{\text{cal}}}\!\left[\mathcal{L}\!\left(f^{\text{mask}(h_\ell)}(x),y\right)-\mathcal{L}\!\left(f(x),y\right)\right],$$

with a robust estimate

$$\widehat{\Delta}(h_\ell)=\mathrm{median}\{\Delta_k(h_\ell)\}_{k=1}^{K}-\widehat{\Delta}_{\text{base}}.$$

Here head attribution is explicitly interventional: a head is important because masking it degrades reasoning performance on a calibration set.

A second formulation attributes predictions to token–token interactions inside self-attention. In "Self-Attention Attribution: Interpreting Information Interactions Inside Transformer" (Hao et al., 2020), a head-level attribution matrix is defined by integrated gradients over the attention matrix,

$$\mathrm{Attr}_h(A)=A_h\odot \int_{\alpha=0}^{1}\frac{\partial \mathbf F(\alpha A)}{\partial A_h}\,d\alpha,$$

and the scalar head score is

$I_h=\mathbb{E}_x[\max(\mathrm{Attr}_h(A))].$

Under this view, head attribution does not ask only whether a head matters, but which interactions mediated by that head matter.

A third formulation shifts the target of attribution from inference behavior to training origin. In "Mechanistic Data Attribution: Tracing the Training Origins of Interpretable LLM Units" (Chen et al., 29 Jan 2026), Mechanistic Data Attribution restricts influence functions to a head-specific subspace $\theta_{\text{sub}}$ and a head-specific probe $f_{\text{probe}}$,

$$\mathcal{I}_{\text{MDA}}(z)\approx - \nabla_{\theta_{\text{sub}}} f_{\text{probe}}(\theta)^\top \hat{H}_{\theta_{\text{sub}}}^{-1}\nabla_{\theta_{\text{sub}}}\mathcal{L}(z,\theta).$$

This makes head attribution a question about which training samples caused a head to emerge and behave as it does.

Other recent formulations are explicitly behavior-centric. "Backdoor Attribution: Elucidating and Controlling Backdoor in LLMs" (Yu et al., 26 Sep 2025) defines per-head causal effects on backdoor outputs through counterfactual activation substitution, while "DEAL: Disentangling Transformer Head Activations for LLM Steering" (Zhan et al., 10 Jun 2025) defines a behavioral relevance score for each head via quantized latent representations and AUC-ROC discrimination between behavior-aligned and behavior-violating responses. These formulations make head attribution operational for safety analysis and inference-time steering.

2. Causal attribution for pruning and reasoning preservation

The most direct interventional treatment of head attribution appears in Causal Attribution Pruning (CAP) (Sheth et al., 27 Apr 2026). CAP is a three-stage, training-free procedure: it computes causal head attribution scores by masking heads on a reasoning calibration set, converts head scores into normalized importance factors $I(h_\ell)\in[1.0,10.0]$, and uses these factors to define the weight-level pruning metric

$\ell$0

Attention projections $\ell$1 are partitioned by head region, while MLP weights receive a layer-average importance $\ell$2. A global threshold $\ell$3 is then chosen so that the fraction of weights with $\ell$4 equals the target sparsity within $\ell$5, producing unstructured weight-level sparsity (Sheth et al., 27 Apr 2026).

The calibration set is reasoning-centric: GSM8K, StrategyQA, and ARC-Challenge, with $\ell$6 examples per task, total calibration size $\ell$7–$\ell$8, stratified by length quartiles and fixed $\ell$9. CAP is forward-only and training-free; head scores are cached per model plus calibration configuration, which allows multiple sparsity levels to be explored without recomputing attribution and reduces pruning time by $$\Delta(h_\ell)=\mathbb{E}_{(x,y)\sim \mathcal{D}_{\text{cal}}}\!\left[\mathcal{L}\!\left(f^{\text{mask}(h_\ell)}(x),y\right)-\mathcal{L}\!\left(f(x),y\right)\right],$$0. The paper states that median aggregation over $$\Delta(h_\ell)=\mathbb{E}_{(x,y)\sim \mathcal{D}_{\text{cal}}}\!\left[\mathcal{L}\!\left(f^{\text{mask}(h_\ell)}(x),y\right)-\mathcal{L}\!\left(f(x),y\right)\right],$$1 disjoint subsets is crucial for stable head rankings (Sheth et al., 27 Apr 2026).

Empirically, CAP is evaluated on Llama-3-8B-Instruct and Mistral-7B-Instruct at 10%, 20%, and 50% sparsity on GSM8K, StrategyQA, and ARC-Challenge. On Llama-3-8B-Instruct at 20% sparsity, ARC-Challenge accuracy rises from 43.9% with Wanda to 70.8% with CAP, a relative accuracy gain of $$\Delta(h_\ell)=\mathbb{E}_{(x,y)\sim \mathcal{D}_{\text{cal}}}\!\left[\mathcal{L}\!\left(f^{\text{mask}(h_\ell)}(x),y\right)-\mathcal{L}\!\left(f(x),y\right)\right],$$2; at the same sparsity, CAP is also better on StrategyQA but worse on GSM8K, and its perplexity is slightly worse than Wanda, 9.57 versus 8.70 (Sheth et al., 27 Apr 2026). At 50% sparsity, CAP collapses on Llama-3, with GSM8K 0.5%, StrategyQA 6.9%, ARC-Challenge 30.1%, and perplexity 428.2, which the paper attributes to coarse MLP attribution. By contrast, on Mistral-7B-Instruct at 50% sparsity, CAP outperforms Wanda on GSM8K, StrategyQA, and ARC-Challenge, with perplexity on a similar scale, 10.5 versus 9.82 (Sheth et al., 27 Apr 2026).

Heatmaps of $$\Delta(h_\ell)=\mathbb{E}_{(x,y)\sim \mathcal{D}_{\text{cal}}}\!\left[\mathcal{L}\!\left(f^{\text{mask}(h_\ell)}(x),y\right)-\mathcal{L}\!\left(f(x),y\right)\right],$$3 further show that fewer than 20% of heads have high causal scores. In Llama-3-8B, critical heads concentrate in layers 11–14 and near boundaries, while in Mistral-7B critical heads concentrate in early layers 0–5 and are more diffuse across deeper layers. This suggests that moderate sparsity can be achieved by protecting a small set of causally important heads while pruning aggressively in low-impact regions (Sheth et al., 27 Apr 2026).

3. Attribution of interactions and failures inside transformers

Integrated-gradients attribution on self-attention matrices provides a head-level view of how information interactions support a prediction. In AttAttr (Hao et al., 2020), the attribution matrix $$\Delta(h_\ell)=\mathbb{E}_{(x,y)\sim \mathcal{D}_{\text{cal}}}\!\left[\mathcal{L}\!\left(f^{\text{mask}(h_\ell)}(x),y\right)-\mathcal{L}\!\left(f(x),y\right)\right],$$4 is used to identify important heads, construct attribution trees, and generate adversarial patterns. Head pruning guided by attribution score tracks functional importance substantially better than pruning by average attention score. On MNLI, when only 10% of heads are kept, AttAttr achieves about 60% accuracy, whereas the Taylor expansion baseline from Michel et al. is about 40%. The same work reports that keeping only the top-2 heads per layer retains about 97% of original accuracy, while removing the top-2 heads per layer causes extreme accuracy drops (Hao et al., 2020).

The same paper uses aggregated per-layer attribution matrices to construct attribution trees that expose hierarchical interactions across layers. In MNLI, lower layers emphasize local interactions, mid layers aggregate to tokens such as “supplement” and “extra,” and top layers emphasize cross-sentence semantic matching. In SST-2, the tree traces sentiment through interactions that aggregate toward tokens such as “matched” and “spirit.” The paper also shows that high-attribution patterns can be used as adversarial triggers: on MNLI, inserting extracted patterns can reduce class accuracy from 82.87% to 0.8% in one reported setting, and on SST-2, one trigger lowers positive-class accuracy from 92.79% to 54.95% (Hao et al., 2020).

A contrastive formulation appears in "Contrastive Attribution in the Wild: An Interpretability Analysis of LLM Failures on Realistic Benchmarks" (Tan et al., 20 Apr 2026). There the attribution target is the logit difference

$$\Delta(h_\ell)=\mathbb{E}_{(x,y)\sim \mathcal{D}_{\text{cal}}}\!\left[\mathcal{L}\!\left(f^{\text{mask}(h_\ell)}(x),y\right)-\mathcal{L}\!\left(f(x),y\right)\right],$$5

with relevance propagated through layers and internal states by AttnLRP. The implemented analysis operates at hidden-state granularity rather than directly at head granularity, but the paper explicitly states that the same machinery can be refined to attention heads, MLP blocks, or individual neurons. It defines per-state relevance approximately as

$$\Delta(h_\ell)=\mathbb{E}_{(x,y)\sim \mathcal{D}_{\text{cal}}}\!\left[\mathcal{L}\!\left(f^{\text{mask}(h_\ell)}(x),y\right)-\mathcal{L}\!\left(f(x),y\right)\right],$$6

and builds cross-layer attribution graphs for long contexts (Tan et al., 20 Apr 2026).

The empirical study spans Qwen3 models of 0.6B, 1.7B, 4B, and 8B, plus Olmo-3-7B-Think checkpoints, on IFEval, GAIA2 / ARE, MATH, and EvalPlus. The paper classifies attribution outcomes into M-IA, NC-IA + M-AG, and NC-IA + AG, and identifies recurring patterns such as Underweight Relevant Tokens, Overweight Irrelevant Tokens, and URT + OIT. In a representative IFEval failure, “<<” briefly surpasses “First” in relevance at a mid layer, but late-layer residual additions favor “First” at the final layer, illustrating late-layer bias amplification overriding mid-layer context signals. The paper also states that many MATH failures likely require neuron-level or feature-level analysis, because coarse hidden-state attribution is often not enough to explain numerical reasoning errors (Tan et al., 20 Apr 2026).

4. Spectral and circuit-level discovery of attributed heads

A distinct line of work avoids labels and attribution gradients by identifying heads through activation geometry. "Spectral Probe-Circuits: A Three-Step Recipe for Identifying Attention-Head Circuits in Pretrained Transformers" (Xu, 22 May 2026) defines, for each head, a participation ratio

$$\Delta(h_\ell)=\mathbb{E}_{(x,y)\sim \mathcal{D}_{\text{cal}}}\!\left[\mathcal{L}\!\left(f^{\text{mask}(h_\ell)}(x),y\right)-\mathcal{L}\!\left(f(x),y\right)\right],$$7

from the singular-value spectrum of the head-output activation matrix, and a time-integrated spectral signal

$$\Delta(h_\ell)=\mathbb{E}_{(x,y)\sim \mathcal{D}_{\text{cal}}}\!\left[\mathcal{L}\!\left(f^{\text{mask}(h_\ell)}(x),y\right)-\mathcal{L}\!\left(f(x),y\right)\right],$$8

This signal ranks heads doing sustained content-dependent computation without labels or attribution gradients (Xu, 22 May 2026).

The method is explicitly three-step: per-head spectral ranking, task-pattern screening, and group ablation against a matched-random control. It is validated across TS-51M, Karpathy’s 124M model, the Pythia family, OLMo 1B, and OLMoE 1B-7B. For induction, the recipe identifies circuits of 2–6 heads in some models and 3–11 heads across the broader set, with a 94–100% drop in synthetic-induction top-1 after ablation. The paper further reports that the same computation identifies the seed-specific circuit on each of six independent seeds of a 51M-parameter probe model (Xu, 22 May 2026).

An important scaling result is that the fraction of heads doing identifiable specialized computation is conserved at 17–19% across the Pythia family from 124M to 410M, while specific induction circuits stay 3–11 heads and are therefore sublinear in total head count. The paper also emphasizes that best-class ranking can be confounded by first-token / BOS dominance, so for sink-heavy regimes it screens all heads for capability-specific selectivity, such as induction-selectivity $$\Delta(h_\ell)=\mathbb{E}_{(x,y)\sim \mathcal{D}_{\text{cal}}}\!\left[\mathcal{L}\!\left(f^{\text{mask}(h_\ell)}(x),y\right)-\mathcal{L}\!\left(f(x),y\right)\right],$$9, before causal ablation (Xu, 22 May 2026). This distinguishes general specialization from task-causal structure.

5. Data-, behavior-, and safety-level attribution

Mechanistic Data Attribution makes head attribution explicitly historical: it asks which training samples caused a head to form. MDA represents each mechanism by the three-tuple $$\widehat{\Delta}(h_\ell)=\mathrm{median}\{\Delta_k(h_\ell)\}_{k=1}^{K}-\widehat{\Delta}_{\text{base}}.$$0, where $$\widehat{\Delta}(h_\ell)=\mathrm{median}\{\Delta_k(h_\ell)\}_{k=1}^{K}-\widehat{\Delta}_{\text{base}}.$$1 is a monitoring metric, $$\widehat{\Delta}(h_\ell)=\mathrm{median}\{\Delta_k(h_\ell)\}_{k=1}^{K}-\widehat{\Delta}_{\text{base}}.$$2 selects the head subspace, and $$\widehat{\Delta}(h_\ell)=\mathrm{median}\{\Delta_k(h_\ell)\}_{k=1}^{K}-\widehat{\Delta}_{\text{base}}.$$3 measures functional behavior. For induction heads, $$\widehat{\Delta}(h_\ell)=\mathrm{median}\{\Delta_k(h_\ell)\}_{k=1}^{K}-\widehat{\Delta}_{\text{base}}.$$4; for previous-token heads, $$\widehat{\Delta}(h_\ell)=\mathrm{median}\{\Delta_k(h_\ell)\}_{k=1}^{K}-\widehat{\Delta}_{\text{base}}.$$5. Using EK-FAC to approximate subspace curvature, the paper ranks training samples by their influence on head-specific probe objectives and shows, through deletion and augmentation in Pythia 14M, 31M, 70M, and 160M, that high-influence samples causally modulate the emergence of induction and previous-token heads (Chen et al., 29 Jan 2026). It reports that the top 10% of samples contribute about 50% of total influence mass, that high-influence samples are dominated by repetitive structural data such as XML, LaTeX, repeated identifier patterns, and long repeated strings, and that interventions targeting induction-head formation induce a concurrent change in in-context learning capability. The paper states that this is direct causal evidence for the long-standing hypothesis regarding the functional link between induction heads and ICL (Chen et al., 29 Jan 2026).

Backdoor Attribution turns head attribution into a safety-control mechanism. BAHA decomposes attention outputs into per-head contributions $$\widehat{\Delta}(h_\ell)=\mathrm{median}\{\Delta_k(h_\ell)\}_{k=1}^{K}-\widehat{\Delta}_{\text{base}}.$$6, averages poisoned activations to obtain $$\widehat{\Delta}(h_\ell)=\mathrm{median}\{\Delta_k(h_\ell)\}_{k=1}^{K}-\widehat{\Delta}_{\text{base}}.$$7, and defines the Average Causal Indirect Effect through counterfactual substitution. Across Llama-2-7B-chat and Qwen-2.5-7B-Instruct with label-modification, fixed-output, and jailbreak backdoors, the paper finds that backdoor-positive heads are sparse: ablating a minimal $$\widehat{\Delta}(h_\ell)=\mathrm{median}\{\Delta_k(h_\ell)\}_{k=1}^{K}-\widehat{\Delta}_{\text{base}}.$$8 of total heads is sufficient to reduce the Attack Success Rate by over 90%. It then constructs a Backdoor Vector

$$\widehat{\Delta}(h_\ell)=\mathrm{median}\{\Delta_k(h_\ell)\}_{k=1}^{K}-\widehat{\Delta}_{\text{base}}.$$9

from the top-attributed heads and shows that a 1-point intervention on single representation can boost ASR on clean inputs up to $$\mathrm{Attr}_h(A)=A_h\odot \int_{\alpha=0}^{1}\frac{\partial \mathbf F(\alpha A)}{\partial A_h}\,d\alpha,$$0 or suppress ASR on triggered inputs down to $$\mathrm{Attr}_h(A)=A_h\odot \int_{\alpha=0}^{1}\frac{\partial \mathbf F(\alpha A)}{\partial A_h}\,d\alpha,$$1 (Yu et al., 26 Sep 2025). This establishes head attribution not only as diagnosis but as control.

DEAL addresses behavior-relevant heads for inference-time steering by training a per-head vector-quantized autoencoder on last-token head activations, partitioning the latent space into $$\mathrm{Attr}_h(A)=A_h\odot \int_{\alpha=0}^{1}\frac{\partial \mathbf F(\alpha A)}{\partial A_h}\,d\alpha,$$2 semantic units quantized with a shared codebook of size $$\mathrm{Attr}_h(A)=A_h\odot \int_{\alpha=0}^{1}\frac{\partial \mathbf F(\alpha A)}{\partial A_h}\,d\alpha,$$3, adding a supervised contrastive loss, and then using an autoregressive prior over the discrete codes of positive examples to score each head by AUC-ROC discrimination between behavior-aligned and behavior-violating responses (Zhan et al., 10 Jun 2025). These scores are used for both selection and importance weighting in steering interventions. On TruthfulQA, DEAL$$\mathrm{Attr}_h(A)=A_h\odot \int_{\alpha=0}^{1}\frac{\partial \mathbf F(\alpha A)}{\partial A_h}\,d\alpha,$$4 improves over ITI by an average relative 20% across seven models and reaches an 81.5% improvement on Qwen2.5-7B-Instruct; the same head selections show strong zero-shot generalization on MQuAKE and CLUTRR (Zhan et al., 10 Jun 2025). The paper interprets this as evidence that the selected heads encode behavior-relevant information more robustly than heads chosen by linear probe accuracy or other heuristics.

6. Architectural and geometric extensions beyond transformer attention heads

Head attribution has also been built directly into network architecture. In "End-to-end Feature Alignment: A Simple CNN with Intrinsic Class Attribution" (Farvardin et al., 26 Mar 2026), FA-CNN replaces the usual linear classifier head with a 3D global average pooling head and constrains the backbone with dampened skip connections,

$$\mathrm{Attr}_h(A)=A_h\odot \int_{\alpha=0}^{1}\frac{\partial \mathbf F(\alpha A)}{\partial A_h}\,d\alpha,$$5

so that penultimate channels are arranged per class. The classifier head computes

$$\mathrm{Attr}_h(A)=A_h\odot \int_{\alpha=0}^{1}\frac{\partial \mathbf F(\alpha A)}{\partial A_h}\,d\alpha,$$6

The paper proves that FA-CNN penultimate feature maps are identical to Grad-CAM saliency maps and that these maps morph layer-by-layer over depth toward penultimate class activations (Farvardin et al., 26 Mar 2026). In this setting, attribution is intrinsic to the classifier head rather than added post hoc.

A different use of the phrase appears in "Densemarks: Learning Canonical Embeddings for Human Heads Images via Point Tracks" (Pozdeev et al., 4 Nov 2025). Here head attribution refers to assigning every pixel in a human head image a canonical 3D coordinate in a shared unit cube,

$$\mathrm{Attr}_h(A)=A_h\odot \int_{\alpha=0}^{1}\frac{\partial \mathbf F(\alpha A)}{\partial A_h}\,d\alpha,$$7

augmented by latent cube features, landmark anchoring, and segmentation constraints. The model is trained on point tracks from CelebV-HQ, starting from about 35K videos and filtering to about 32K videos, with up to 400 tracks per video. The representation supports common semantic parts, face/head tracking, and stereo reconstruction, and is reported to be robust to pose variations while covering the entire head, including hair. On dense warping quality for same-person pairs on Nersemble, DenseMarks achieves MAE 3.68 and RMSE 5.9, compared with DINOv3 MAE 7.6 and RMSE 12.69, Hyperfeatures MAE 8.26 and RMSE 13.29, Fit3D MAE 12.75 and RMSE 21.83, and Sapiens MAE 14.88 and RMSE 24.12 (Pozdeev et al., 4 Nov 2025). In this usage, attribution is geometric and semantic rather than causal in the transformer sense.

7. Limitations, controversies, and research directions

Current head-attribution methods differ sharply in what they regard as the attributed object: a causal contribution to loss, a relevance score for a contrastive logit, a training-data origin, a latent behavioral code, or a spatially aligned class map. This plurality is productive, but it also means that scores are not directly interchangeable. CAP shows strong reasoning preservation at 10–20% sparsity, yet it is limited by coarse MLP attribution at 50% sparsity and performs poorly in preliminary tests on Mixture-of-Experts architectures; the paper recommends extending causal attribution to neurons in MLPs and exploring task-specific or multi-task calibration (Sheth et al., 27 Apr 2026). Contrastive LRP highlights realistic failure dynamics, but it relaxes strict conservation and often becomes diffuse on MATH, where the authors state that neuron-level or feature-level analysis is likely required (Tan et al., 20 Apr 2026).

Methodological costs are also significant. MDA requires EK-FAC-based influence estimation over mechanistically selected subspaces and is therefore expensive; its validity depends on Hessian approximation and local linearization, and the paper notes that high-influence samples modulate the rate of circuit emergence more than the absolute possibility of emergence (Chen et al., 29 Jan 2026). Spectral Probe-Circuits avoids labels and gradients, but its second stage still depends on pre-defined task patterns, and companion work is motivated partly by the observation that, for composed tasks, pattern selectivity can decouple from task-causal structure (Xu, 22 May 2026). DEAL requires behavior-labeled contrastive data and per-head VQ-AE training; its causal interpretation is stronger than linear probing but remains tied to intervention outcomes rather than a full structural causal model (Zhan et al., 10 Jun 2025).

Architectural and geometric variants have their own constraints. FA-CNN provides positive attribution only, because post-ReLU features and Grad-CAM equivalence do not encode negative evidence; it also ties channel allocation to the number of classes and is slightly below the strongest ResNet baselines under full training recipes (Farvardin et al., 26 Mar 2026). DenseMarks depends on point-tracker quality, manual choices of landmark anchors in the cube, and training on human talking heads; performance on non-human heads or heavy occlusion is not guaranteed (Pozdeev et al., 4 Nov 2025).

Taken together, these works suggest that head attribution is moving from descriptive saliency toward mechanistic intervention. The strongest recent results are those that combine a precise attribution target with a direct manipulation: head masking for reasoning preservation, group ablation for induction circuits, data deletion and augmentation for circuit emergence, counterfactual substitution for backdoor heads, and inference-time steering on behavior-relevant heads (Sheth et al., 27 Apr 2026, Xu, 22 May 2026, Chen et al., 29 Jan 2026, Yu et al., 26 Sep 2025, Zhan et al., 10 Jun 2025). A plausible implication is that future progress will come less from a single universal score than from task-aligned attribution schemes that can be validated by causal control.