Class-Discriminative Attention Maps for Vision Transformers

TL;DR

The results suggest that existing importance estimators may not provide sufficient class-sensitivity, and class-discriminative attention maps (CDAM) are shown to be highly class-discriminative and semantically relevant, while providing compact explanations.

Abstract

Importance estimators are explainability methods that quantify feature importance for deep neural networks (DNN). In vision transformers (ViT), the self-attention mechanism naturally leads to attention maps, which are sometimes interpreted as importance scores that indicate which input features ViT models are focusing on. However, attention maps do not account for signals from downstream tasks. To generate explanations that are sensitive to downstream tasks, we have developed class-discriminative attention maps (CDAM), a gradient-based extension that estimates feature importance with respect to a known class or a latent concept. CDAM scales attention scores by how relevant the corresponding tokens are for the predictions of a classifier head. In addition to targeting the supervised classifier, CDAM can explain an arbitrary concept shared by selected samples by measuring similarity in the latent space of ViT. Additionally, we introduce Smooth CDAM and Integrated CDAM, which average a series of CDAMs with slightly altered tokens. Our quantitative benchmarks include correctness, compactness, and class sensitivity, in comparison to 7 other importance estimators. Vanilla, Smooth, and Integrated CDAM excel across all three benchmarks. In particular, our results suggest that existing importance estimators may not provide sufficient class-sensitivity. We demonstrate the utility of CDAM in medical images by training and explaining malignancy and biomarker prediction models based on lung Computed Tomography (CT) scans. Overall, CDAM is shown to be highly class-discriminative and semantically relevant, while providing compact explanations.

Figures (24)

• Figure 1: Extending attention maps (AM), the proposed class-discriminative attention maps (CDAM) quantify and visualize input features that are relevant for the target class in Vision Transformers (ViT). We visualize importance scores obtained with CDAM for a linear classifier with ViT-S/8, trained with DINO, as a backbone caron2021emerging. For details, refer to \ref{['sec:AS-class']}. Orange and blue colors correspond to positive and negative values, respectively.
• Figure 2: Graphical scheme of class-discriminative attention maps (CDAM). (a) CDAMs are obtained by first propagating an image through the transformer blocks and using its latent representation $l=\texttt{[CLS]}'$ to infer a class score $f_c$ or similarity score $g(l,l_c)$ (see \ref{['definition-class']} or \ref{['definition-concept']}). (b) Detailed view of the final transformer block, adapted from dosovitskiy2020image.
• Figure 3: Visual comparison of different importance estimators which obtain explanations w.r.t. different output classes. Similar heat maps regardless of target classes imply that importance estimators are not class discriminative. Attention maps (AM) are, by design, not class discriminative. Orange and blue colors correspond to positive and negative values, respectively.
• Figure 4: Class-Discriminative Attention Maps (CDAM) for user-defined concepts (\ref{['sec:AS-similarity']}). The concept embedding $l_c$ has been obtained by averaging the latent representations of 30 images that include the common concept. In each instance, from left to right: Sample image, AM, and CDAM. Orange and blue correspond to positive and negative values, respectively.
• Figure 5: When targeting different classes (left) or concepts (right), the resulting CDAMs are distinct and align with the corresponding objects. Note that green parts of vegetables are shown to activate the zuccini class (left). The wrinkles in the elephant's trunk appear to be related to the zebra concept, semantically mistaken by the model as zebra stripes.