Interpretability-Aware Vision Transformer

TL;DR

The paper addresses the interpretability gap in Vision Transformers by introducing IA-ViT, a three-component model (feature extractor, predictor, interpreter) trained with an interpretability-aware objective. The interpreter uses a single-head self-attention to simulate the predictor and provide faithful explanations, guided by a joint loss that combines cross-entropy, knowledge-distillation-based simulation, and attention distribution alignment via MMD. Empirical results show IA-ViT delivers competitive accuracy while producing high-quality explanations and improved fairness on CelebA, with ablation analyses confirming the necessity of the simulation and regularization terms. This approach demonstrates that integrating interpretability into the training objective yields reliable, task-consistent explanations without compromising performance, facilitating trustworthy deployment of ViTs in high-stakes domains.

Abstract

Vision Transformers (ViTs) have become prominent models for solving various vision tasks. However, the interpretability of ViTs has not kept pace with their promising performance. While there has been a surge of interest in developing {\it post hoc} solutions to explain ViTs' outputs, these methods do not generalize to different downstream tasks and various transformer architectures. Furthermore, if ViTs are not properly trained with the given data and do not prioritize the region of interest, the {\it post hoc} methods would be less effective. Instead of developing another {\it post hoc} approach, we introduce a novel training procedure that inherently enhances model interpretability. Our interpretability-aware ViT (IA-ViT) draws inspiration from a fresh insight: both the class patch and image patches consistently generate predicted distributions and attention maps. IA-ViT is composed of a feature extractor, a predictor, and an interpreter, which are trained jointly with an interpretability-aware training objective. Consequently, the interpreter simulates the behavior of the predictor and provides a faithful explanation through its single-head self-attention mechanism. Our comprehensive experimental results demonstrate the effectiveness of IA-ViT in several image classification tasks, with both qualitative and quantitative evaluations of model performance and interpretability. Source code is available from: https://github.com/qiangyao1988/IA-ViT.

Figures (4)

• Figure 1: Illustration of attention maps and predictive distributions from both the CLS patch and other image patches.
• Figure 2: IA-ViT consists of three major components: feature extractor, predictor, and interpreter. Both the predictor and the interpreter generate the class prediction for this cat image. KD is applied on the two logits in the simulation objective. The attention weights in SA and MSA are aligned via MMD during the training process for better explanations.
• Figure 3: Quantitative explanation performance comparison in terms of differences between insertion and deletion scores.
• Figure 4: Examples of attribution maps generated by the baseline explanation methods for ViT and the attention weights from the interpreter of IA-ViT.