From Attention to Prediction Maps: Per-Class Gradient-Free Transformer Explanations

Ronen Schaffer, Tomer Michaeli, Meir Yossef Levi

Keywords: explainable ai, vision transformer

Abstract: The Vision Transformer (ViT) has become a standard model architecture in computer vision, especially for classification tasks. As such, explaining ViT predictions has attracted significant research efforts in recent years. Many methods rely on attention maps, which highlight \emph{where} in the image the network directs its attention. In this paper, we introduce Prediction~Maps -- a novel explanation method that complements attention maps by revealing \emph{what} the network sees. Prediction maps visualize how each patch token within a given layer is associated with each possible class. This is done by utilizing the classification head at the output of the network, originally trained to be fed with the class token at the last layer. Specifically, to obtain the prediction map of a particular layer, we apply the classification head to every patch token within that layer. We show that prediction maps provide complementary information to attention maps and illustrate that combining them leads to state-of-the-art explainability performance. Furthermore, since our proposed method is neither gradient- nor perturbation-based, it offers superior computational and memory efficiency compared to competing methods.

Primary Area: interpretability and explainable AI

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 5091