Head Pursuit: Probing Attention Specialization in Multimodal Transformers

NeurIPS 2025 (Spotlight) | Paper |BibTeX

Lorenzo Basile, Valentino Maiorca, Diego Doimo, Francesco Locatello, Alberto Cazzaniga

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.

Installation

uv sync
source .venv/bin/activate

Project Structure

src/headpursuit/
├── apply_somp.py              # Table 1 - Find specialized heads via SOMP (Matching Pursuit)
├── intervention/
│   ├── textgeneration.py       # §4.1-4.2 - Language experiments (QA, toxicity)
│   └── imageunderstanding.py   # §5.1-5.2 - Vision experiments (classification, captioning)
├── evaluation/                 # Compute metrics (F1, CIDEr, toxicity classifier)
├── plot/                       # Generate paper figures
├── encode/                     # Extract attention head activations
└── constants/                  # Keyword dictionaries (WORDS['colors'], WORDS['sentiments'], etc.)

src/residual/                   # Residual stream tracing utilities (from ResiDual library)
├── tracing/                    # Model-specific tracers (Mistral, LLaVA, Gemma3, Qwen)
└── data/dataset.py             # Load HF datasets

Quick Start

Full pipeline for Figure 2 (TriviaQA country suppression):

# Step 1: Encode activations (extract head representations)
python -m headpursuit.encode.text --model mistral --dataset triviaqa

# Step 2: Score heads with SOMP
python -m headpursuit.apply_somp --model mistral --dataset triviaqa --property countries

# Step 3: Run intervention (rescale top-8 heads)
python -m headpursuit.intervention.textgeneration \
    --model mistral --dataset triviaqa --property countries --k 8 --alpha -1

# Step 4: Evaluate and plot
python -m headpursuit.evaluation.triviaqa --model mistral
python -m headpursuit.plot.triviaqa --model mistral  # → Generates Figure 2

Adapt for other experiments:

• Table 2 (toxicity): --dataset rtp/tet --property val --k 32
• Figure 5 (captioning): Use intervention/imageunderstanding.py with --dataset flickr30k --property colors --k 16
• Figure 3 (classification): --dataset mnist --property val (and similar for other datasets)
• Appendix models: Add --model gemma3, --model qwen25vl, or --model llava_next_13

Supported Models

Understanding SOMP

Simultaneous Orthogonal Matching Pursuit (SOMP) decomposes attention head outputs using the model's unembedding matrix as a dictionary. Unlike Logit Lens (which analyzes single tokens), SOMP processes entire datasets to find heads that consistently activate on semantic concepts (e.g., color words). The Explained Variance Ratio (EVR) scores heads—high EVR indicates specialization. See Algorithm 1 in the paper for details.

Custom Properties

Define your own keyword dictionaries:

# In src/headpursuit/constants/__init__.py:
WORDS['animals'] = ['dog', 'cat', 'bird', 'fish', 'elephant']

# Then use: --property animals

Citation

If you use this work in your research, please cite:

BibTeX:

@inproceedings{basile2025headpursuit,
  title={Head Pursuit: Probing Attention Specialization in Multimodal Transformers},
  author={Lorenzo Basile and Valentino Maiorca and Diego Doimo and Francesco Locatello and Alberto Cazzaniga},
  booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
  year={2025}
}

Please also cite ResiDual, from which large parts of this codebase are adapted:

BibTeX:

@article{basile2025residual,
  title     = {ResiDual Transformer Alignment with Spectral Decomposition},
  author    = {Lorenzo Basile and Valentino Maiorca and Luca Bortolussi and Emanuele Rodolà and Francesco Locatello},
  journal   = {Transactions on Machine Learning Research},
  year      = {2025},
  url       = {https://openreview.net/forum?id=z37LCgSIzI},
  note      = {}
}

License

See LICENSE file.

Acknowledgments

The authors acknowledge the Area Science Park supercomputing platform ORFEO made available for conducting the research reported in this paper, and the technical support of the Laboratory of Data Engineering staff.