Table of Contents
Fetching ...

Head Pursuit: Probing Attention Specialization in Multimodal Transformers

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

TL;DR

Head Pursuit investigates how individual attention heads in generative transformers encode specific semantic attributes, proposing a sparse-signal framework that uses SOMP with the unembedding dictionary to identify concept-aligned heads. The approach enables dataset-level analysis and targeted interventions, demonstrating that inverting or amplifying a small fraction of specialized heads can controllably suppress or enhance targeted content across language and vision-language tasks. The main contributions include a principled sparse-decomposition method, empirical evidence of cross-modal head specialization, and practical demonstrations of targeted control in QA, toxicity mitigation, image classification, and captioning. This work advances interpretability and controllability of large-scale generative models by revealing an interpretable, manipulable linear structure within attention layers that requires minimal intervention to effect meaningful changes.

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.

Head Pursuit: Probing Attention Specialization in Multimodal Transformers

TL;DR

Head Pursuit investigates how individual attention heads in generative transformers encode specific semantic attributes, proposing a sparse-signal framework that uses SOMP with the unembedding dictionary to identify concept-aligned heads. The approach enables dataset-level analysis and targeted interventions, demonstrating that inverting or amplifying a small fraction of specialized heads can controllably suppress or enhance targeted content across language and vision-language tasks. The main contributions include a principled sparse-decomposition method, empirical evidence of cross-modal head specialization, and practical demonstrations of targeted control in QA, toxicity mitigation, image classification, and captioning. This work advances interpretability and controllability of large-scale generative models by revealing an interpretable, manipulable linear structure within attention layers that requires minimal intervention to effect meaningful changes.

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.
Paper Structure (36 sections, 3 equations, 10 figures, 17 tables, 1 algorithm)

This paper contains 36 sections, 3 equations, 10 figures, 17 tables, 1 algorithm.

Figures (10)

  • Figure 1: Overview of our method. Given a language or vision-language model and a target property defined by text (e.g., colors, sentiments, numbers), we score all attention heads according to how well they align with interpretable directions from a fixed dictionary, using a method based on Matching Pursuit mallat1993matching (1). We then select the top-$k$ heads (2) and intervene by rescaling their contribution to the residual stream (3), either enhancing or suppressing the attribute in the model’s output.
  • Figure 2: Question answering performance of Mistral-7B on TriviaQA. F1 scores are shown for samples associated (country) or not (other) with the target attribute, normalized to the base model accuracy without intervention. Random baselines are displayed as medians and interquartile ranges.
  • Figure 3: Classification results under different head selection strategies: (light blue) 16 heads with highest variance ratio explained by SOMP; (dark blue) 32 heads with highest explained variance ratio; (yellow) 16 random heads, with the same layer-wise counts of top 16; (orange) 32 random heads, with the same layer-wise counts of top 32.
  • Figure 4: (left) Jaccard similarity between sets of top-$16$ LLaVA heads selected with SOMP over different datasets; (right) Classification accuracy on a target dataset, denoted by row, when the top-$16$ heads are selected with SOMP on a source dataset, denoted by column. Accuracy is normalized w.r.t. the base accuracy on target dataset.
  • Figure 5: Effect of inhibitory interventions on LLaVA attention heads for Flickr30k captioning. Bars show the normalized presence of target attributes (colors, sentiments, quantity) under different head-selection strategies (SOMP and random). Random baselines are reported in terms of medians and interquartile ranges.
  • ...and 5 more figures