Head Pursuit: Probing Attention Specialization in Multimodal Transformers

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

This paper introduces Head Pursuit, a method that applies Simultaneous Orthogonal Matching Pursuit (SOMP) over a transformer's unembedding matrix to identify attention heads specialized in narrow semantic domains. By editing as few as 1% of heads, the method can reliably suppress or enhance targeted concepts in both language and vision-language model outputs without any fine-tuning.

Problem

Prior interpretability tools (e.g., Logit Lens, Attention Lens) are heuristic, applied to individual examples, and cannot quantitatively rank heads by concept relevance across a dataset. The Attention Lens requires training separate linear probes per head, making it unscalable to billion-parameter models. There is no principled, training-free method to identify and intervene on semantically specialized attention heads across both unimodal LLMs and multimodal VLMs.

Method

The method decomposes each attention head's activation matrix H ∈ ℝ^{n×d} as a sparse linear combination of dictionary atoms drawn from the model's unembedding matrix D ∈ ℝ^{v×d}. Specifically, Simultaneous Orthogonal Matching Pursuit (SOMP) is applied jointly over n samples, iteratively selecting the dictionary atom (token direction) maximally correlated with residuals across all samples, then refitting via restricted least squares. The fraction of head variance explained by the concept-restricted SOMP reconstruction ranks heads by specialization. Intervention is performed by rescaling selected head contributions to the residual stream: α = −1 to suppress, α > 1 (e.g., 5) to enhance a target attribute. No additional training is required; head activations are computed on held-out training data, evaluated on a disjoint test set.

Key Contributions

• Reframes SOMP as a multi-sample generalization of the Logit Lens, providing a principled sparse decomposition of attention head outputs over the unembedding dictionary.
• Demonstrates consistent semantic specialization of attention heads in both unimodal LLMs (Mistral-7B) and generative VLMs (LLaVA-NeXT-7B, LLaVA-NeXT-13B, Gemma3-12B, Qwen2.5-VL-7B).
• Shows that inverting ~1% of heads identified by SOMP suppresses targeted concepts in QA, toxicity mitigation, image classification, and image captioning tasks, with significantly stronger specificity than Logit Lens baselines.
• Shows that amplifying specialized heads (α > 1) enhances target concept generation by >60% while preserving caption quality (CIDEr >80% of baseline).
• Reveals that head overlap between datasets follows semantic similarity (e.g., MNIST↔SVHN share digit heads; EuroSAT↔RESISC45 share remote-sensing heads).

Results

• TriviaQA (Mistral-7B, country names): Inverting 8 heads (0.8% of total) degrades country-answer F1 substantially more than non-country F1; random control shows no significant effect; Logit Lens degrades both domains equally, confirming SOMP's specificity.
• Toxicity mitigation (Mistral-7B): With 32 SOMP heads, normalized toxic generation count drops to 0.66 (RTP) and 0.49 (TET); Logit Lens achieves 0.71/0.68; random heads maintain or increase toxicity (~1.0–1.13).
• Image classification (LLaVA-NeXT-7B, 6 datasets): Inverting top-32 SOMP heads significantly disrupts accuracy across all datasets; 32 random heads at equivalent layers have near-zero impact. Cross-dataset transfer degrades performance in proportion to semantic similarity.
• Image captioning (Flickr30k, LLaVA): Inhibiting 16 SOMP heads reduces color/sentiment/quantity keyword count to near zero while CIDEr remains >80% of original. Enhancing 32 heads increases target keyword presence by >60%.

Limitations

• SOMP imposes a linearity assumption that may not capture nonlinear structure in head representations.
• Method quality depends on the completeness and accuracy of the semantic keyword dictionary; noisy or incomplete lists can bias head selection.
• Intervention mechanism is coarse — global rescaling/inversion of entire head outputs rather than position- or modality-selective steering.
• Evaluated primarily on suppression of concept generation; more complex tasks (e.g., multi-hop reasoning) are not addressed.

Relevance to Vision-Language Models

This paper provides a mechanistic interpretability tool directly applicable to generative VLMs, showing that the LLM backbone within VLMs exhibits the same concept-specialized head structure as unimodal language models. The method requires no VLM-specific modifications — it applies the unembedding matrix of the language backbone to decompose multimodal attention heads — making it immediately portable to LLaVA-style architectures and validated on LLaVA-NeXT-7B/13B, Gemma3-12B, and Qwen2.5-VL. For researchers tracking VLMs, this work advances the line connecting linear representation hypotheses in LLMs to multimodal settings and offers a lightweight, training-free alternative to prompt engineering or LoRA-based steering for concept control in image captioning and classification.