Pith. sign in

REVIEW 4 major objections 6 minor 30 references

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

T0 review · 4 major / 6 minor · reviewed 2026-08-10 · deepseek-v4-flash

Pith's one-line read Token Insight identifies the image tokens a ViT truly depends on by removing the token that most lowers prediction confidence, repeating until the prediction flips.

desk verdict A clean, well-specified token-removal method for ViT interpretability, but the paper's central causal claim rests on an untested premise and the experiments need baseline controls. read the letter →

arxiv 2501.15452 v1 pith:EHHQEKTQ submitted 2025-01-26 cs.CV cs.AI

classification cs.CVcs.AI
keywords tokendiscardingvisiontransformerinterpretabilitymedicalimagingpolypdetectionshortcutlearningself-supervisedpredictionconfidence
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

The paper introduces Token Insight, a procedure that finds which image patches (tokens) a vision transformer actually depends on for its prediction. It works by repeatedly deleting the token whose removal causes the largest drop in confidence in the predicted class, continuing until the model changes its answer; the deleted set then represents the evidence the model used. The authors apply this to colonic polyp detection with four ViT-B/16 models (scratch, supervised, DINO, MAE pretraining) and show that models differ in how many tokens they depend on and which regions they treat as decisive. The central promise is that this gives a more transparent view of transformer-based medical AI and can flag shortcut learning, where predictions rely on signals unrelated to the disease. A key stated advantage over occlusion-based methods is that token removal avoids the missingness bias that masked inputs introduce.

What carries the argument

The central mechanism is iterative token discarding: the ViT input is patchified into 196 tokens plus the [cls] token, and at each iteration every remaining non-[cls] token is separately deleted, the model's confidence in the original predicted class is recomputed, and the token producing the largest confidence drop is permanently discarded; the procedure halts when the predicted class changes. Token Insight inherits the idea from earlier works that use token removal or masking for interpretability, but extends them by continuing greedily rather than stopping after one step. The method's claimed advantage over occlusion-style masking is that deleting a token entirely removes its contribution rather than replacing it with a mask value that can introduce missingness bias.

What would settle it

Run the same greedy removal procedure on a model but compare the confidence drop caused by deleting each 'critical' token against the drop caused by deleting the same number of randomly chosen tokens across many images; if random deletions produce drops as large as the flagged tokens, the method is measuring the generic effect of missing input rather than token-specific evidence.

Watch

Extended reading notes

Core claim

Token Insight is a greedy, model-agnostic procedure for a vision transformer: given an image tokenized into patches, at each step it evaluates every remaining token, removes the one whose deletion produces the largest drop in prediction confidence for the correct class, and repeats until the model's argmax flips away from that class. The set of tokens discarded in this way is presented as the critical evidence for the prediction. Applying this to polyp detection, the paper reports that supervised and DINO-pretrained ViTs need fewer tokens to abandon a polyp-positive call, and their individual tokens have larger confidence impact, whereas MAE-pretrained and scratch-trained models rely on a broader set of tokens. The paper further shows example cases where the critical tokens lie outside the medically relevant polyp region, which it interprets as evidence of shortcut learning. The authors claim that relying on prediction-confidence changes alone, using only native token-discarding, makes the method applicable to any transformer model without additional modules.

Load-bearing premise

The whole ranking rests on the premise that measuring how much confidence drops when a token is deleted reflects that token's true causal contribution, with no distorting effect from deleting a token that the model never saw deleted during training.

Editorial extensions

If this is right

  • For a given transformer, Token Insight produces a ranked set of critical tokens; the size of that set is a per-model measure of how concentrated or diffuse the evidence for a prediction is.
  • Because the procedure requires only forward passes of the unmodified model, it can be applied to any ViT-style architecture regardless of pretraining scheme.
  • The reported differences between pretraining methods suggest that model selection for medical imaging should consider not just accuracy but how many and which tokens a model leans on.
  • If a model's critical tokens fall outside the disease region, Token Insight can flag that the model may be using spurious correlations or shortcut learning.
  • Treating all transformer-based medical models alike is a mistake: token reliance varies substantially with training procedure.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • This is an editorial inference: the greedy, confidence-drop-based ranking is a heuristic, and the paper itself notes that theoretical guarantees for finding the smallest influential token set are left for future work, so the identified tokens may be one of many possible sufficient sets rather than the unique causal evidence.
  • The same procedure could be used as a data-cleaning diagnostic: by running Token Insight over a training set and inspecting tokens that repeatedly appear critical in polyp images but are not annotated as polyp, one could generate candidate shortcut regions for manual review.
  • Because removing tokens changes the effective input distribution, the claimed absence of missingness bias is not fully settled by the fact that tokens are deleted; comparing Token Insight with zero-padding or inpainting the same tokens on the same model would give a direct test.
  • The O(N^2) forward-pass cost limits application to high-resolution images; reusing cached attention or gradient information could make the method scale.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

4 major / 6 minor

Summary. The paper proposes Token Insight, a method for identifying critical image tokens in ViT-based medical image classifiers. Given an input for which the model predicts the correct class, Token Insight greedily removes the token whose deletion causes the largest drop in the predicted confidence for that class, repeats the procedure until the argmax changes, and reports the removed tokens as critical for the prediction. The method is applied to four ViT-B/16 models (trained from scratch, supervised pretrained, DINO pretrained, MAE pretrained) on the CP-CHILD colonic polyp dataset. The reported experiments compare the number of tokens discarded and per-token confidence changes across the four models, and qualitative examples are used to claim that Token Insight can reveal shortcut learning.

Significance. If its causal interpretation were supported, Token Insight would be a simple and generally applicable tool for probing ViT predictions, requiring no extra modules and using only the model's own confidence. The method is precisely defined and reproducible in its operational form, and the comparison of models with different pretraining strategies is a useful empirical question. The paper also deserves credit for formulating the missingness-bias concern explicitly, even though its treatment of that concern is incomplete. However, the central interpretive claims about causal token importance and shortcut learning are not yet supported by the evidence presented, so the contribution is currently more of a well-specified method proposal than a validated interpretability result.

major comments (4)
  1. [§2.3 (Relation to occlusion-based methods)] The claim that token removal is free of missingness bias is not established. Deleting a token changes the sequence length and removes information, producing inputs that differ from the training distribution; the resulting confidence drop may reflect generic sensitivity to missing context rather than the causal importance of the removed token. Because Figures 3a, 3b, and the shortcut-learning discussion in Figure 4 all depend on this premise, the paper should include a random-removal baseline (confidence drops for randomly selected token subsets at matched cardinalities) and, ideally, repeated-seed estimates. Without such controls, the ranking of 'critical tokens' is not distinguishable from a generic effect of information loss.
  2. [§3.2 (Identifying shortcut learning)] The claim that certain tokens are 'not medically relevant' and therefore indicate shortcut learning is qualitative and not validated. There is no ground-truth labeling of intended versus spurious signal, no quantitative overlap of critical tokens with expert or pathology annotations, and no head-to-head comparison with a model known to rely on a controlled spurious artifact. To support this claim, the authors should provide either quantitative evaluation against polyp/pathology masks or a synthetic shortcut experiment (e.g., adding a known artifact pattern and testing whether Token Insight highlights it).
  3. [§3.1–§3.2] Each of the four models appears to be trained only once, and all cross-model comparisons rest on that single run. The reported differences in tokens discarded (Figure 3a) and maximum confidence change (Figure 3b) between MAE/scratch and supervised/DINO models could be due to training stochasticity or hyperparameter choices. The paper should report multiple seeds with error bars and, if appropriate, a statistical test for the model comparisons.
  4. [Conclusion, first paragraph] The statement that 'the set thus obtained forms a good approximation to the smallest set of highly influential tokens' is not supported by any experiment or bound in the paper. The greedy algorithm is described, but no comparison to exhaustive search, no lower or upper bounds, and no ablation against random subsets are presented. Either remove this claim or supply the promised computational evidence.
minor comments (6)
  1. [Figure 2 caption] The caption contains a duplicated word: 'below each each image' should be 'below each image'.
  2. [§2.3] The formal greedy selection rule does not specify how ties are broken when multiple removed tokens yield the same confidence, and the stopping criterion would be clearer if the exact condition (e.g., arg max of the full softmax vector) were restated in a displayed equation.
  3. [§2.2–§2.3] The manuscript does not explain how the ViT handles variable-length token sequences after token removal, particularly with respect to learned positional embeddings; this detail is important for reproducibility of the implementation.
  4. [§3.1] Only the from-scratch model's test accuracy is reported (95.3%), with other models described only as 'comparable'; please report the actual test accuracies of all four models.
  5. [§4] The statement that computational cost scales as O(N^2) is ambiguous: the number of forward passes is O(N^2) in the greedy procedure, but each forward pass itself is O(N^2) in attention cost, so the total computational cost should be stated more precisely.
  6. [Throughout] The dataset name is written inconsistently as 'CP-Child' and 'CP-CHILD'; please standardize the spelling.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity: Token Insight's critical-token notion is an explicit operational definition, and the paper's reported results are computed directly from that definition rather than from a fitted parameter or a self-citation chain.

full rationale

Token Insight defines criticality transparently in Section 2.3: "the most critical token is the token t_{i+1} that gives rise to the largest drop in prediction confidence," and every reported outcome (number of tokens discarded, confidence drops, and token maps) is calculated directly from that definition. This is an operational definition rather than a hidden equivalence between separate constructs: the paper does not fit a parameter to a subset of data and then predict a closely related quantity, nor does it invoke an author-specific uniqueness theorem. The only in-house citation, [17], is a survey used for background on SSL frameworks and is not load-bearing. The contested claim that token removal avoids missingness bias is an assumption about the causal validity of the ranking, not a circularity; the absence of a random-removal control is an empirical limitation of the broader shortcut-learning interpretation, not a case where the output equals the input by construction. The derivation chain is self-contained because the method's output is, by the paper's own equations, the greedy confidence-drop ranking, and this is stated explicitly rather than concealed.

Assumptions & free parameters 0 free parameters · 4 assumptions · 0 invented entities

The central claim rests on the equivalence between confidence drop and causal importance, and on the validity of token-depleted inputs. These are domain assumptions, not derived facts. The greedy approximation is explicitly admitted by the authors to lack a proof. There are no fitted parameters in the method itself.

assumptions (4)
  • domain assumption Prediction confidence is a faithful proxy for the model's decision-making process.
    Token Insight defines importance as the drop in g(theta, X)_c, so all criticality claims inherit this assumption; invoked throughout Section 2.3.
  • domain assumption Removing tokens from a ViT input produces a valid input whose behavior is causally comparable to the original.
    The method's validity depends on this, asserted in Section 2.3 without empirical support.
  • ad hoc to paper Greedy sequential removal recovers the smallest or at least a representative set of critical tokens.
    The final paragraph of Section 4 says this is supported only by unshown computational evidence; no theoretical bound is given.
  • domain assumption CP-CHILD dataset labels and train/test splits are reliable for polyp classification.
    Standard empirical assumption; the paper relies on the public dataset as ground truth.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Identifying Critical Tokens for Accurate Predictions in Transformer-based Medical Imaging Models." pith.science (2026). https://pith.science/paper/EHHQEKTQ

@misc{pith2026250115452,
  author       = {Pith},
  title        = {Pith review of: Identifying Critical Tokens for Accurate Predictions in Transformer-based Medical Imaging Models},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/EHHQEKTQ}},
  note         = {Machine review of arXiv:2501.15452}
}
read the original 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

Figures reproduced from arXiv: 2501.15452 by the authors.

Figure 1
Figure 1. An overview of the ViT architecture and the tokenization of image patches [PITH_FULL_IMAGE:figures/full_fig_p002_1.png] view at source ↗
Figure 2
Figure 2. An example of procedure of Token Insight where critical tokens ( [PITH_FULL_IMAGE:figures/full_fig_p004_2.png] view at source ↗
Figure 3
Figure 3. (a) Number of tokens discarded to change predictions from polyp-positive [PITH_FULL_IMAGE:figures/full_fig_p007_3.png] view at source ↗
Figures from the paper (1 more)
Figure 4
Figure 4. Figure 4: Illustrations of Token Insight maps highlighting the most impactful tokens [PITH_FULL_IMAGE:figures/full_fig_p008_4.png]

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

30 extracted references · 11 canonical work pages

  1. [1]

    Bai, B., Liang, J., Zhang, G., Li, H., Bai, K., Wang, F.: Why attentions may not be interpretable? In: Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining. pp. 25–34 (2021)

  2. [2]

    Bastings, J., Filippova, K.: The elephant in the interpretability room: Why use attention as explanation when we have saliency methods? arXiv preprint arXiv:2010.05607 (2020)

  3. [3]

    arXiv preprint arXiv:2210.09461 (2022)

    Bolya, D., Fu, C.Y., Dai, X., Zhang, P., Feichtenhofer, C., Hoffman, J.: Token merging: Your vit but faster. arXiv preprint arXiv:2210.09461 (2022)

  4. [4]

    In: Proceedings of the IEEE/CVF International Conference on Computer Vision

    Caron, M., Touvron, H., Misra, I., J´ egou, H., Mairal, J., Bojanowski, P., Joulin, A.: Emerging properties in self-supervised vision transformers. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 9650–9660 (2021)

  5. [5]

    In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition

    Chefer, H., Gur, S., Wolf, L.: Transformer interpretability beyond attention visu- alization. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. pp. 782–791 (2021)

  6. [6]

    In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition

    Chen, X., He, K.: Exploring simple siamese representation learning. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 15750–15758 (2021)

  7. [7]

    In: Proceedings of the IEEE/CVF International Conference on Com- puter Vision

    Chen, X., Xie, S., He, K.: An empirical study of training self-supervised vision transformers. In: Proceedings of the IEEE/CVF International Conference on Com- puter Vision. pp. 9640–9649 (2021)

  8. [8]

    arXiv preprint arXiv:2010.11929 (2020)

    Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al.: An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 (2020)

Show all 30 references
  1. [9]

    Nature Machine In- telligence 2(11), 665–673 (2020)

    Geirhos, R., Jacobsen, J.H., Michaelis, C., Zemel, R., Brendel, W., Bethge, M., Wichmann, F.A.: Shortcut learning in deep neural networks. Nature Machine In- telligence 2(11), 665–673 (2020)

  2. [10]

    In: Proceedings of the IEEE/CVF International Conference on Computer Vision

    Haurum, J.B., Escalera, S., Taylor, G.W., Moeslund, T.B.: Which tokens to use? investigating token reduction in vision transformers. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 773–783 (2023) 10 Kang et al

  3. [11]

    In: Proceedings of the IEEE/CVF Conference on Com- puter Vision and Pattern Recognition

    He, K., Chen, X., Xie, S., Li, Y., Doll´ ar, P., Girshick, R.: Masked autoencoders are scalable vision learners. In: Proceedings of the IEEE/CVF Conference on Com- puter Vision and Pattern Recognition. pp. 16000–16009 (2022)

  4. [12]

    In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (2016)

    He, K., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image recognition. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (2016)

  5. [13]

    arXiv preprint arXiv:2204.08945 (2022)

    Jain, S., Salman, H., Wong, E., Zhang, P., Vineet, V., Vemprala, S., Madry, A.: Missingness bias in model debugging. arXiv preprint arXiv:2204.08945 (2022)

  6. [14]

    In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition

    Long, S., Zhao, Z., Pi, J., Wang, S., Wang, J.: Beyond attentive tokens: Incorporat- ing token importance and diversity for efficient vision transformers. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 10334–10343 (2023)

  7. [15]

    arXiv preprint arXiv:2310.07819 (2023)

    Madsen, A., Reddy, S., Chandar, S.: Faithfulness measurable masked language models. arXiv preprint arXiv:2310.07819 (2023)

  8. [16]

    Matsoukas, C., Haslum, J.F., S¨ oderberg, M., Smith, K.: Is it time to replace cnns with transformers for medical images? arXiv preprint arXiv:2108.09038 (2021)

  9. [17]

    arXiv preprint arXiv:2305.13689 (2023)

    Ozbulak, U., Lee, H.J., Boga, B., Anzaku, E.T., Park, H., Van Messem, A., De Neve, W., Vankerschaver, J.: Know your self-supervised learning: A survey on image-based generative and discriminative training. arXiv preprint arXiv:2305.13689 (2023)

  10. [18]

    Advances in Neural Information Processing Systems 34, 24898–24911 (2021)

    Pan, B., Panda, R., Jiang, Y., Wang, Z., Feris, R., Oliva, A.: ia − red2: Interpretability-aware redundancy reduction for vision transformers. Advances in Neural Information Processing Systems 34, 24898–24911 (2021)

  11. [19]

    arXiv preprint arXiv:2202.12015 (2022)

    Renggli, C., Pinto, A.S., Houlsby, N., Mustafa, B., Puigcerver, J., Riquelme, C.: Learning to merge tokens in vision transformers. arXiv preprint arXiv:2202.12015 (2022)

  12. [20]

    In: International Conference on Learn- ing Representations (2021)

    Rigotti, M., Miksovic, C., Giurgiu, I., Gschwind, T., Scotton, P.: Attention-based interpretability with concept transformers. In: International Conference on Learn- ing Representations (2021)

  13. [21]

    International Journal of Computer Vision 115(3), 211–252 (2015)

    Russakovsky, O., Deng, J., Su, H., Krause, J., Satheesh, S., Ma, S., Huang, Z., Karpathy, A., Khosla, A., Bernstein, M., Berg, A.C., Fei-Fei, L.: ImageNet large scale visual recognition challenge. International Journal of Computer Vision 115(3), 211–252 (2015)

  14. [22]

    Serrano, S., Smith, N.A.: Is attention interpretable? arXiv preprint arXiv:1906.03731 (2019)

  15. [23]

    Medical Image Analysis p

    Shamshad, F., Khan, S., Zamir, S.W., Khan, M.H., Hayat, M., Khan, F.S., Fu, H.: Transformers in medical imaging: A survey. Medical Image Analysis p. 102802 (2023)

  16. [24]

    International Conference on Learning Representations (2015)

    Simonyan, K., Zisserman, A.: Very deep convolutional networks for large-scale image recognition. International Conference on Learning Representations (2015)

  17. [25]

    arXiv preprint arXiv:2305.09617 (2023)

    Singhal, K., Tu, T., Gottweis, J., Sayres, R., Wulczyn, E., Hou, L., Clark, K., Pfohl, S., Cole-Lewis, H., Neal, D., et al.: Towards expert-level medical question answering with large language models. arXiv preprint arXiv:2305.09617 (2023)

  18. [26]

    Sun, S., Koch, L.M., Baumgartner, C.F.: Right for the wrong reason: Can inter- pretable ml techniques detect spurious correlations? In: International Conference on Medical Image Computing and Computer-Assisted Intervention. pp. 425–434. Springer (2023)

  19. [27]

    Advances in neural information pro- cessing systems 30 (2017) Identifying Critical Tokens for Accurate Predictions 11

    Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, L., Polosukhin, I.: Attention is all you need. Advances in neural information pro- cessing systems 30 (2017) Identifying Critical Tokens for Accurate Predictions 11

  20. [28]

    BMC Medical Imaging 20, 1–14 (2020)

    Wang, W., Tian, J., Zhang, C., Luo, Y., Wang, X., Li, J.: An improved deep learning approach and its applications on colonic polyp images detection. BMC Medical Imaging 20, 1–14 (2020)

  21. [29]

    In: Computer Vision–ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part I 13

    Zeiler, M.D., Fergus, R.: Visualizing and understanding convolutional networks. In: Computer Vision–ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part I 13. pp. 818–833. Springer (2014)

  22. [30]

    In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition

    Zeng, W., Jin, S., Liu, W., Qian, C., Luo, P., Ouyang, W., Wang, X.: Not all tokens are equal: Human-centric visual analysis via token clustering transformer. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 11101–11111 (2022)

Pith tools

Reviewed August 10, 2026 · model on record in the stance chip above.