GMAR: Gradient-Driven Multi-Head Attention Rollout for Vision Transformer Interpretability

TL;DR

This work tackles the interpretability gap in Vision Transformers by introducing GMAR, a gradient-driven framework that assigns per-head importance scores and integrates them into a weighted attention rollout. By computing class-specific gradients to derive head weights $w_h$ (normalized via $L1$ or $L2$) and applying these weights during rollout, GMAR yields head-aware, interpretable maps that better reflect regions driving predictions. Empirical results on a ViT-Large-Patch16-224 model fine-tuned on Tiny-ImageNet show GMAR outperforming traditional Attention Rollout and providing more faithful explanations, with qualitative maps that focus on task-relevant objects and regions. The approach offers a practical, scalable enhancement to ViT interpretability, with potential applicability to larger datasets and more advanced architectures.

Abstract

The Vision Transformer (ViT) has made significant advancements in computer vision, utilizing self-attention mechanisms to achieve state-of-the-art performance across various tasks, including image classification, object detection, and segmentation. Its architectural flexibility and capabilities have made it a preferred choice among researchers and practitioners. However, the intricate multi-head attention mechanism of ViT presents significant challenges to interpretability, as the underlying prediction process remains opaque. A critical limitation arises from an observation commonly noted in transformer architectures: "Not all attention heads are equally meaningful." Overlooking the relative importance of specific heads highlights the limitations of existing interpretability methods. To address these challenges, we introduce Gradient-Driven Multi-Head Attention Rollout (GMAR), a novel method that quantifies the importance of each attention head using gradient-based scores. These scores are normalized to derive a weighted aggregate attention score, effectively capturing the relative contributions of individual heads. GMAR clarifies the role of each head in the prediction process, enabling more precise interpretability at the head level. Experimental results demonstrate that GMAR consistently outperforms traditional attention rollout techniques. This work provides a practical contribution to transformer-based architectures, establishing a robust framework for enhancing the interpretability of Vision Transformer models.

Figures (4)

• Figure 1: Overview of the Gradient-Driven Multi-Head Attention Rollout (GMAR) algorithm structure as applied to Vision Transformers, detailing the computation of head weights in multi-head gradients and illustrating the attention mechanism across layers and heads to enhance interpretability.
• Figure 2: Comparing attention head weights, with the highest and lowest values highlighted in red and yellow
• Figure 3: Overlay images generated for 10 randomly selected images from the Tiny ImageNet 200 dataset. The methods used include Grad-CAM, Attention Rollout, and GMAR with L1 and L2 norms, highlighting the regions of interest identified by each approach.
• Figure 4: Visualization of the attention map, weighted attention map, and difference map.