Vision transformer token routing improves dermoscopic lesion detection
MedTokenBudget: Lesion-Preserving Token Routing for Dermoscopic Image Classification
Computer Vision and Pattern RecognitionArtificial Intelligence
Summary
Images used for diagnosing skin problems often focus on a small lesion area, but current AI systems process every part of the image equally. The authors propose MedTokenBudget, which learns to keep only the important patches containing the lesion while discarding the rest. This method uses extra information about lesion locations to guide the AI in focusing on meaningful image parts, improving efficiency without losing key diagnostic details. They tested this approach on a skin lesion dataset and found it kept more lesion information than other token-reduction methods.
What this means in practice
- •For medical imaging software developers: Build more efficient skin lesion classifiers that focus on lesion areas to reduce computation without losing diagnostic information.
- •For healthcare data engineers: Implement token selection methods to optimize storage and processing of dermoscopic images by retaining lesion-relevant patches selectively.
Tested on one dataset.
Authors
Zhexiang Li
Abstract
Dermoscopy classifiers built on Vision Transformers process all image patches uniformly, although diagnostic evidence is concentrated in the lesion region. Existing token pruning methods reduce tokens using generic saliency or similarity signals, but rarely ask whether the retained subset still contains the lesion. This paper introduces MedTokenBudget, a supervised post-backbone token routing framework that learns to construct compact lesion-enriched representations when auxiliary lesion masks are available. Its Lesion-Aware Token Scoring (LATS) module fuses attention entropy, feature norm, and local feature contrast through a learned scorer, then routes the top-$K$ patches under a target budget. LATS is trained with budget curriculum learning, diversity regularization, attention distillation, and lesion-mask supervision. The trained router is evaluated with a lesion retention rate that directly measures how much ground-truth lesion evidence survives the token budget. On ISIC 2019, mask-supervised LATS consistently outperforms Random and ToMe at headline budgets while retaining substantially more lesion patches. Code is provided for reproducibility, and complete tabulated results are included in the supplementary material.