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

For straightforward and easily understandable experiments, we utilize the combination of recently proposed CP-CHILD-A and CP-CHILD-B datasets, which involves the detection of colonic polyps [28]. Both datasets combined comprises of 9,500 colonoscopy RGB images obtained from 1,600 patients using the Olympus PCF-H290DI and FUJIFILM EC-530wm. The task involves identifying colonic polyps, which is challenging since it is not trivial to distinguish polyps from other problematic colonic tissues such as lesions, ulcerative colitis, and inflammatory bowel disease. Furthermore, many of the polyps presented in the dataset do not appear fully in the picture and are only visible in the corners, thus increasing the challenge. Further details regarding data splits of these datasets are provided in Table 1.

To showcase the capabilities of the proposed method, we employ the most commonly used transformer-based computer vision architecture, Vision Transformer-Base/16 (ViT-B/16) [8]. As suggested by [8], we use this model for images of size $224\times 224$ with image tokens of size $16\times 16$, resulting in a total of $196$ tokens. We modify the final linear layer of the model to accommodate the two-class classification problem tackled in this work.

We use two ViT-B/16 models where the first one is randomly initialized and trained from scratch while the second one is pretrained in a supervised fashion. Apart from those two models, in order to capture various properties of self-supervised learning methods that have seen increased use in medical imaging problems in recent years, we employ two additional models pretrained using (1) DINO [4], a discriminative SSL framework, and (2) MAE [11], which relies on a generative approach. The pretraining for the aforementioned models is performed on the ImageNet dataset [21], a large-scale dataset containing natural images.

Given an image $\texttt{X}\in\mathbb{R}^{D\times W\times H}$ and its categorical association $\bm{y}\in\mathbb{R}^{M}$, sampled from a dataset $(\texttt{X},\bm{y})\sim\mathcal{D}$, where $y_{c}=1$ and $y_{m}=0$ for all $m\in\{0,\ldots,M\}\backslash{c}$, let $g_{\theta}(\cdot)$ represent a vision transformer with parameters $\theta$ that maps an image to a set of prediction likelihoods, denoted as $g(\theta,\mathtt{X})=\bm{y}^{\prime}$, where $\sum_{i=1}^{M}y_{i}^{\prime}=1$. If $\arg\max(\bm{y}^{\prime})=c$, then the classification is considered correct.

When utilized as an input for a vision transformer, the image X undergoes a process of patchification (also called tokenization), transforming it into a set of tokens as shown in Figure 1. In this study, we adopt ViT-B/16 which employs $16\times 16$ patches. Given that the input images have a resolution of $224\times 224$, the input X is patchified into $196$ tokens denoted as $\mathtt{X}=[\bm{x}_{1},\bm{x}_{2},\ldots,\bm{x}_{196}]$. As shown in Figure 1, in the setup of ViT, a special token called [cls] token is prepended on $\mathtt{X}$, thus increasing the number of tokens by one. This token is then used for the purpose of making classification using a linear layer after the final transformer encoder layer.

Token Insight is designed to progressively identify the most critical token contributing to the model’s prediction. To achieve this, at each step, it removes one token at a time, measures the change in prediction, and identifies the token causing the largest drop in prediction confidence for the correct class. Subsequently, after examining each token, the one resulting in the highest confidence change when removed is permanently discarded. This process is then repeated to discover remaining critical tokens until the prediction ultimately changes. In the case of polyp detection problem, Token Insight finds the tokens that lead to the prediction of the \sayPolyp class and the token discarding operation is carried out repeatedly until the prediction changes from \sayPolyp-positive to \sayPolyp-negative for the images considered. Token Insight finds its origins in the works of [15, 18]. However, unlike those approaches which stop after a single iteration, we continue to search for critical tokens greedily and discard them until the prediction is eventually changed, thus identifying all of the critical tokens that contribute to the prediction.

Formally, we can represent Token Insight as follows: given an input image represented as $\mathtt{X}=[\bm{x}_{k}]_{k\in\{1,\ldots,196\}}$, denote by $\mathtt{X}^{(t_{1},\ldots,t_{i})}$ the image where the tokens $\{t_{1},\ldots,t_{i}\}$ have been removed. At the next iteration, $i+1$, the most critical token is the token $t_{i+1}$ that gives rise to the largest drop in prediction confidence:

The iterative search described above stops when $\arg\max(g(\theta,\mathtt{X}^{(t_{1},\ldots,t_{i})})_{c})\neq c$. Note that Token Insight does not alter or remove the [cls] token during the process described above and, in order to generate the prediction confidences [cls] token is used according to the description of [8] throughout the process.

An illustration of the token discarding operation of Token Insight is provided in Figure 2 where the example image is initially predicted with $0.99$ confidence as having a polyp. Using Token Insight, critical tokens are identified and discarded iteratively until the prediction eventually drops down to $0.44$ for the polyp-positive class, thus identifying tokens that contribute to the prediction made by the model.

Relation to occlusion-based methods –  Note that, on the surface, Token Insight appears to be similar to the occlusion-based methods which mask certain regions of the input to measure the change in the prediction [29]. The primary difference between Token Insight and these methods is that during the token discarding process the model is not affected by the potential spurious signals such as missingness bias introduced by the masking operation (i.e., the bias that is introduced by the mask) [13]. Since the tokens are completely removed, the model prediction is not influenced by the removed tokens. Unfortunately, since the token removal operation is only available for transformer-based models, Token Insight is not usable for CNNs and other non-transformer DNN architectures.

As described in Section 2.2 and Section 2.1, we employ four ViT-B/16 models, three of which are pretrained on the ImageNet dataset, and train them on the combined CP-Child dataset. To discover the most appropriate model with the highest performance, we perform extensive training efforts and in what follows, we detail the training routine that leads to the identification of the best model.

For pretrained models (supervised, DINO, MAE), training is performed for $25$ epochs using Stochastic Gradient Descent (SGD) with an initial learning rate of $10^{-3}$ and momentum of $0.05$. The model trained from scratch undergoes a longer training period ($50$ epochs) to compensate for the lack of pretraining and uses an increased learning rate of $0.035$, weight decay of $10^{-5}$, and momentum of $0.075$. All models employ a batch size of $32$ and use cosine annealing, which reduces the learning rate after each epoch using the method described in [6].

In the work of [28], pretrained ResNets are reported to attain an accuracy of approximately $99\%$ on the CP-Child dataset. With pretrained ViTs, we replicate these results, achieving comparable performance. The model trained from scratch exhibits slightly reduced performance at $95.3\%$ due to the absence of pretraining. Nonetheless, this model still delivers commendable results on this dataset, making these models suitable for a study on predictive reasoning.

In order to investigate various properties of the proposed approach as well as models pretrained in different ways, we apply Token Insight to the images containing polyps in the CP-Child-B dataset.

Number of tokens discarded –  For all four models, we investigate the number of tokens discarded via Token Insight to examine the model’s reliance on tokens for making a polyp-positive prediction. Doing so, in Figure 3(b), we present the number of tokens discarded to change a polyp-positive prediction into a polyp-negative one. As it can be seen, both MAE and the model trained from scratch discard more tokens before the prediction changes, while models pretrained with supervised learning and DINO discard fewer, and hence rely on more tokens to make a prediction. This implies that treating all transformer-based medical AI models the same would be a mistake since the number of tokens used for confirming a prediction may vary significantly depending on how the model is trained.

The most impactful token – Expanding our investigation, in Figure 3(b), we illustrate the influence of the most impactful tokens across the dataset, as measured by their impact on the prediction confidence when removed. Complementing our previous insights, we observe that individual tokens in the models pretrained with supervised learning, as well as DINO, have more influence on the prediction compared to MAE and the model trained from scratch. Once again, this observation reveals that it is possible to have two models sharing the same architecture, but where one heavily relies on fewer tokens, whereas the other focuses on a broader area. Consequently, depending on the medical imaging task at hand, practitioners may prefer one type of model over another.

Combining the two experiments discussed above, we present Figure 3(c) which shows confidence drops per removed token over the number of tokens discarded for each model. Graphs presented in Figure 3(c) can be seen as a summarizing views of Figure 3(b) and Figure 3(b), where the goal is to reveal the different behaviors of various models. As can be seen, all models exhibit an initial stage in which confidence drops only moderately as the initial few tokens are removed, followed by a steep drop in confidence. MAE and scratch see a relative robustness in the decrease of prediction confidence as a function of the number of tokens removed, compared to DINO and supervised pretrained models.

Identifying shortcut learning – In Figure 4, we present several examples showcasing the application of Token Insight on images containing polyps. Note that, while Token Insight accurately identifies regions containing polyps, thereby indicating that the model learns the correct signals, there are instances (particularly for MAE and Scratch) where certain tokens are identified as important but are not medically relevant to the prediction. This suggests that the model under consideration is indeed making predictions based on signals that are not intended in the dataset, also identified by [26]. As demonstrated, the primary use-case of the proposed approach is to uncover such cases and to discover predictions based on erroneous signals, thereby enhancing the trustworthiness of AI-based medical imaging.