[Paper Review] Is it Time to Replace CNNs with Transformers for Medical Images?

Christos Matsoukas, Johan Fredin Haslum+2 more|arXiv (Cornell University)|Aug 20, 2021

AI in cancer detectionComputer Science23 references105 citations

TL;DR

The paper compares vanilla CNNs and vision transformers (ViTs) for medical imaging, showing ViTs can match or exceed CNN performance when pretrained (ImageNet or self-supervised) and do so with added explainability from attention maps.

ABSTRACT

Convolutional Neural Networks (CNNs) have reigned for a decade as the de facto approach to automated medical image diagnosis. Recently, vision transformers (ViTs) have appeared as a competitive alternative to CNNs, yielding similar levels of performance while possessing several interesting properties that could prove beneficial for medical imaging tasks. In this work, we explore whether it is time to move to transformer-based models or if we should keep working with CNNs - can we trivially switch to transformers? If so, what are the advantages and drawbacks of switching to ViTs for medical image diagnosis? We consider these questions in a series of experiments on three mainstream medical image datasets. Our findings show that, while CNNs perform better when trained from scratch, off-the-shelf vision transformers using default hyperparameters are on par with CNNs when pretrained on ImageNet, and outperform their CNN counterparts when pretrained using self-supervision.

Motivation & Objective

• Assess whether off-the-shelf ViTs can replace CNNs for medical image diagnosis without extensive redesign.
• Evaluate performance of CNNs vs ViTs under different initialization strategies (random, ImageNet pretraining, self-supervised pretraining).
• Determine if transfer learning and self-supervision benefits seen in natural images extend to medical imaging.
• Analyze interpretability advantages of ViTs via attention maps.

Proposed method

• Standardized comparison between ResNet50 (CNN) and DeiT-S (ViT) with 16x16 tokens under three initialization strategies.
• Three medical datasets: APTOS2019 (diabetic retinopathy), ISIC2019 (skin lesions), CBIS-DDSM (mammography).
• Training: Adam optimizer, base lr 1e-4, warm-up, lr decay, 256x256 inputs, common augmentations.
• Five repetitions per setting with selection of best validation checkpoint.
• Self-supervised pretraining with DINO on target data after ImageNet initialization, followed by supervised fine-tuning.

Experimental results

Research questions

• RQ1 Can vanilla ViTs serve as plug-and-play replacements for CNNs in medical image analysis under modest data regimes?
• RQ2 Do ImageNet pretraining and self-supervised pretraining help ViTs reach CNN-level or better performance on medical tasks?
• RQ3 Does self-supervised pretraining provide a larger advantage to ViTs compared to CNNs in medical imaging?
• RQ4 What are the interpretability advantages of ViTs compared to CNNs in medical imaging?

Key findings

• ViTs pretrained on ImageNet perform comparably to CNNs when data is limited.
• Transfer learning with ImageNet initialization helps ViTs match CNN performance on medical tasks.
• Self-supervised pretraining plus supervised fine-tuning yields the best performance, with ViTs slightly outperforming CNNs in this setting.
• Random initialization yields CNNs outperforming ViTs in all datasets.
• ViTs offer improved interpretability through attention maps, providing finer-grained saliency than Grad-CAM alone.