Identifying Critical Tokens for Accurate Predictions in Transformer-based Medical Imaging Models

TL;DR

The paper addresses interpretability in transformer-based medical imaging by introducing Token Insight, a token-discarding approach that identifies the tokens driving predictions without altering model architecture. It greedily removes tokens to maximize drops in the correct-class confidence, preserving the $[cls]$ token and relying solely on prediction changes to quantify token importance. Experiments on the CP-CHILD-A/B polyp dataset with ViT-B/16 models pretrained in supervised and self-supervised fashions (DINO, MAE) or trained from scratch reveal varying token reliance and instances of shortcut learning, demonstrating the method's ability to uncover both medically relevant cues and spurious signals. This approach enhances transparency and clinical trust by localizing influential input components without model modification, while also highlighting computational costs scaling as $O(N^2)$ for $N$ tokens (with ViT-B/16 having $N=196$), and suggesting directions for theoretical bounds and efficiency improvements.

Abstract

With the advancements in self-supervised learning (SSL), transformer-based computer vision models have recently demonstrated superior results compared to convolutional neural networks (CNNs) and are poised to dominate the field of artificial intelligence (AI)-based medical imaging in the upcoming years. Nevertheless, similar to CNNs, unveiling the decision-making process of transformer-based models remains a challenge. In this work, we take a step towards demystifying the decision-making process of transformer-based medical imaging models and propose Token Insight, a novel method that identifies the critical tokens that contribute to the prediction made by the model. Our method relies on the principled approach of token discarding native to transformer-based models, requires no additional module, and can be applied to any transformer model. Using the proposed approach, we quantify the importance of each token based on its contribution to the prediction and enable a more nuanced understanding of the model's decisions. Our experimental results which are showcased on the problem of colonic polyp identification using both supervised and self-supervised pretrained vision transformers indicate that Token Insight contributes to a more transparent and interpretable transformer-based medical imaging model, fostering trust and facilitating broader adoption in clinical settings.

Figures (4)

• Figure 1: An overview of the ViT architecture and the tokenization of image patches are illustrated.
• Figure 2: An example of procedure of Token Insight where critical tokens ($t$) that contribute to the prediction made by the model are iteratively discovered based on their impact on prediction. Identified tokens are discarded after each iteration ($i$) until the prediction changes from positive to negative. For clarity, we provide the prediction confidence ($y'_c$) below each each image where this result is obtained using the image with discarded tokens (black patches). The token that is removed at each iteration is highlighted in green.
• Figure 3: (a) Number of tokens discarded to change predictions from polyp-positive to polyp-negative for CP-Child dataset. (b) Maximum confidence change measured with a single token during the process of (a). Outliers are omitted to enhance visual clarity. (c) Change in polyp-positive confidence based on token removal for all images in the CP-Child test dataset with black lines indicating average number of tokens discarded measured across all images.
• Figure 4: Illustrations of Token Insight maps highlighting the most impactful tokens for a polyp-positive prediction based on models pretrained with various methods.