Head Pursuit: Probing Attention Specialization in Multimodal Transformers

AI summary

Overview

Research area: Mechanistic interpretability of large generative transformers, spanning unimodal language models and vision-language models (VLMs).

Technical level: Intermediate. The paper assumes familiarity with transformer attention heads, the residual stream, and the Logit Lens, but the core idea is explained clearly through an analogy to sparse signal recovery.

Scope: The paper proposes a sparse-coding based method (Head Pursuit, built on Simultaneous Orthogonal Matching Pursuit) for identifying which individual attention heads in text-generative transformers specialize in specific semantic or visual attributes, and shows that editing a small fraction of those heads can suppress or enhance the targeted concepts.

What This Paper Is About

Large language and vision-language models work well, but it remains unclear how their internal components, particularly individual attention heads, organize and represent knowledge. Existing interpretability tools such as the Logit Lens are applied heuristically and usually to one example at a time, which makes it hard to generalize findings across a dataset or to quantify how much a specific component shapes the model's output. This paper reframes the practice of probing intermediate activations with the final decoding layer as sparse signal recovery, so that attention heads can be ranked by their relevance to a target concept across many samples at once, and then edited to change model behavior.

Key Contributions

1. A principled reframing of interpretability as sparse recovery. The authors connect an established sparse recovery algorithm, Matching Pursuit (and its multi-sample variant SOMP), to standard interpretability tools such as the Logit Lens, showing that the Logit Lens is equivalent to a single step of Matching Pursuit on an individual example.

2. A method for identifying concept-specialized heads. By restricting the model's unembedding matrix to the rows of tokens associated with a target semantic area and running SOMP, the fraction of a head's variance explained becomes a measure of specialization, allowing heads to be ranked and selected.

3. Demonstration that head specialization is actionable. Negating (sign-inverting) relevant heads causes targeted degradation in task performance, while amplifying them promotes the generation of specific attributes, across both language and vision-language tasks.

4. Validation across modalities and tasks. The approach is tested on question answering and toxicity mitigation in language, and on image classification and image captioning in vision-language settings.

Main Findings

• Semantically coherent specialized heads exist. Applying SOMP to Mistral-7B attention heads on TriviaQA prompts yields interpretable groupings: L18.H27 surfaces tokens like COVID, Soviet, Obama, Biden, Clinton; L24.H20 surfaces British, American, European, German, English; L25.H14 surfaces December, July, April, October, February; L30.H28 surfaces 9, 1, 3, 7, five. Aggregated Logit Lens on the same heads returns noisier and more redundant tokens, for example Sunday, breakfast, Oct, October for the "Months" head and 8, u, u, n for the "Numbers" head.

• Very few heads can carry a concept. Editing as few as 1% of the heads, selected by the method, can reliably suppress or enhance targeted concepts in the model output.

• Question answering shows targeted, not generic, damage. On Mistral-7B evaluated on TriviaQA with the F1 score, performance on country-related examples noticeably degrades when the signs of 8 heads (0.8% of the total) or more are inverted. Performance on the remaining examples declines more gradually. Country names account for over 6% of the answers in the test split. Random control heads (10 independently sampled sets) show no significant impact, and Logit Lens-selected heads degrade performance roughly equally inside and outside the targeted domain, i.e. they are relevant but not concept-specific.

• Toxicity mitigation works with an incomplete keyword list. Normalized counts of toxic generations after intervention (lower is better): on RealToxicityPrompts, SOMP gives 0.83 at 8 heads, 0.67 at 16 heads, and 0.66 at 32 heads, versus 0.91, 0.79, and 0.71 for Logit Lens and 1.02, 1.00, and 1.13 for random heads. On Thoroughly Engineered Toxicity, SOMP gives 0.83, 0.68, and 0.49 at 8, 16, and 32 heads, versus 0.81, 0.73, and 0.68 for Logit Lens and 0.97, 0.95, and 0.95 for random heads. A held-out lexical metric on strictly disjoint toxic keywords also improves: on RTP, SOMP reaches 1.00, 0.78, and 0.72 at 8, 16, and 32 heads; on TET, 0.80, 0.65, and 0.56.

• Image classification confirms head-level specialization in a VLM. For LLaVA-NeXT-7B on MNIST, SVHN, GTSRB, EuroSAT, RESISC45, and DTD, inverting the top 32 heads identified by the method significantly disrupts classification accuracy relative to the no-intervention baseline, while inverting 32 random heads at equivalent layers has substantially lower or no impact. At k = 16 the pattern is similar except for DTD, whose performance is only weakly affected, suggesting higher head redundancy on that task.

• Head selections reflect semantic similarity between datasets. Jaccard similarity between top-16 head sets shows substantial overlap for MNIST and SVHN (both digit recognition) and for EuroSAT and RESISC45 (both remote sensing). Cross-dataset interventions degrade performance more when source and target datasets are semantically related; GTSRB performance drops notably when intervening with heads selected from MNIST or SVHN, which the authors attribute to the presence of numerical symbols in all three datasets.

• Captioning can be steered in both directions. On Flickr30k with LLaVA, inhibiting as few as 16 heads almost completely removes attribute-related keywords for colors, sentiments, and quantity, while CIDEr always exceeds 80% of the original. Enhancing the same heads with α = 5 increases the presence of target concepts by more than 60% in all three cases with 32 heads, again with only marginal effect on caption quality. Example outputs include "A small dachshund wearing a pink sweater" becoming "A small dachshund wearing a sweater" under color inhibition, and becoming "A black and brown dog wearing a pink sweater" under color enhancement.

• The structure is consistent and controllable. The authors conclude that attention layers contain a highly interpretable and manipulable linear structure, consistent with evidence that large-scale models represent high-level concepts in approximately linear residual subspaces.

Methodology in Plain English

The researchers treat each attention head's output into the residual stream as a matrix of activations across a dataset, and then ask: can this be rebuilt as a sparse combination of known, meaningful directions? They use the model's own unembedding matrix as that dictionary, because each row of it corresponds to a token that can be read back out as natural language.

They then apply Simultaneous Orthogonal Matching Pursuit (SOMP), a multi-sample extension of Orthogonal Matching Pursuit, which greedily picks the dictionary direction most correlated with the residual across all samples, adds it to a support set, refits the coefficients with least squares, and repeats until a set sparsity level is reached (typically 50 iterations in the experiments). Working across all samples at once, rather than one example at a time, is what makes the analysis stable at dataset level. This is a direct generalization of the Logit Lens: the Logit Lens is one Matching Pursuit step on one example.

To find heads for a target concept, they restrict the dictionary to the unembedding rows for words in that concept (country names, toxic words, class names such as colors, sentiments, or quantity) and rank heads by the fraction of variance explained by the SOMP reconstruction. Intervention is deliberately simple: selected heads are rescaled, with α = −1 for inhibition (sign inversion) and α > 1 for enhancement (α = 5 in the captioning experiments), affecting both the head's direct contribution to the residual stream and the indirect flow of its information to later layers. Every experiment includes a random control made of disjoint heads matched in size and layer distribution, reported over 10 independently sampled sets. Head representations used for selection come from training data strictly disjoint from the evaluation data. Experiments ran on a single NVIDIA H100 GPU with 80GB VRAM using HuggingFace transformers checkpoints.

Why This Matters

Impact on research. The paper gives interpretability a mathematically grounded alternative to heuristic probing: instead of inspecting individual examples and hoping the pattern generalizes, it decomposes head behavior into sparse, dataset-level, human-readable directions. It also connects two previously separate literatures, sparse coding signal recovery and mechanistic interpretability, and shows the Logit Lens is a special case of a more general family of methods.

Real-world applications.

• Safer text generation: suppressing toxic heads to reduce harmful content in deployed chatbots without retraining.
• Content moderation and auditing: identifying which internal components are responsible for a given attribute, which could support model inspection and compliance reviews.
• Controlled image captioning: making captions include or exclude colors, sentiments, or quantity terms for assistive technology, accessibility, or dataset balancing.
• Domain-specific degradation testing: deliberately disrupting classification of a specific visual domain to probe where a model is fragile.

Industry relevance. The interventions require no additional training, no separate probes per head, and no gradient computation, which makes the method cheap enough to scale to billion-parameter models, a limitation the paper explicitly attributes to the Attention Lens approach. For teams deploying or fine-tuning large multimodal models, head-level rescaling offers a lightweight alternative to prompt engineering or fine-tuning for steering model behavior.

Future Directions

• More selective and fine-grained interventions. The authors suggest rescaling heads only at specific input positions or modalities, such as disabling heads only over image patch tokens while preserving text understanding in a VLM.
• Multimodal-output settings. Adapting the technique to image generation, where selectively enhancing or suppressing heads during decoding could steer generated images toward or away from particular semantic attributes, as a controllable alternative to prompt engineering or fine-tuning.
• Addressing the linearity assumption. SOMP imposes a linear structure that may not fully capture the nonlinear structure of head representations, an open limitation.
• Improving the semantic dictionary. The reliability of the recovered heads depends on the quality and coverage of the keyword lists; incomplete or noisy lists can bias selection toward spurious or underrepresented concepts.

Target Audience

Interpretability and mechanistic interpretability researchers; engineers working on safety, alignment, or content control for deployed language and multimodal models; and practitioners interested in training-free methods for steering generative model behavior. Readers should be comfortable with transformer architecture and the residual stream, though the sparse recovery formulation is presented in a self-contained way with pseudocode in the appendix for SOMP. The paper is not written for a general audience.

Authors’ abstract

Language and vision-language models have shown impressive performance across a wide range of tasks, but their internal mechanisms remain only partly understood. In this work, we study how individual attention heads in text-generative models specialize in specific semantic or visual attributes. Building on an established interpretability method, we reinterpret the practice of probing intermediate activations with the final decoding layer through the lens of signal processing. This lets us analyze multiple samples in a principled way and rank attention heads based on their relevance to target concepts. Our results show consistent patterns of specialization at the head level across both unimodal and multimodal transformers. Remarkably, we find that editing as few as 1% of the heads, selected using our method, can reliably suppress or enhance targeted concepts in the model output. We validate our approach on language tasks such as question answering and toxicity mitigation, as well as vision-language tasks including image classification and captioning. Our findings highlight an interpretable and controllable structure within attention layers, offering simple tools for understanding and editing large-scale generative models.