Source-linked AI summary
Is it Time to Replace CNNs with Transformers for Medical Images?
Christos Matsoukas, Johan Fredin Haslum, Magnus Söderberg, Kevin Smith
TL;DR
The paper asks whether vanilla ViTs can replace CNNs for medical diagnosis despite limited data. It compares the models across three medical datasets and initialization strategies, finding that ImageNet-pretrained ViTs match CNNs and self-supervised ViTs slightly outperform them.
Problem
The paper investigates whether ViTs can replace CNNs for medical imaging when datasets are small and labels may be less reliable.
Method
The study compares vanilla ViTs and CNNs under similar conditions across three medical datasets, using random, supervised ImageNet, and self-supervised initialization.
Results
ImageNet-pretrained ViTs perform comparably to CNNs, while self-supervised pretraining followed by fine-tuning lets ViTs outperform CNNs by a small margin.
Takeaways & Limitations
ViTs can reliably replace CNNs for small medical datasets when transfer learning is used, with self-supervision providing a marginal additional advantage.
Abstract
from arXiv · showhide
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.
1. Introduction
The paper examines whether vanilla vision transformers can replace CNNs for medical imaging despite limited data, and compares their performance under common initialization strategies. It finds that ImageNet-pretrained ViTs match CNNs, while self-supervised pretraining gives ViTs an advantage.
- Motivation: Medical imaging datasets are smaller and often have less reliable labels, making the data requirements of ViTs a central concern.Evidence from natural images suggested ViTs needed very large-scale pretraining to outperform CNNs.
- Motivation: Transfer learning is commonly used because ImageNet-pretrained CNNs generally outperform models trained from scratch on medical images.It can improve final performance and reduce training time.
- Study aim: The paper tests whether ViTs can replace CNNs as plug-and-play models for medical imaging using three publicly available datasets and limited hyperparameter tuning.The experiments target practitioners with conventional datasets and limited computational budgets.
- Findings: ViTs pretrained on ImageNet perform comparably to CNNs when data is limited.This evaluates supervised ImageNet pretraining as a standard transfer-learning strategy.
- Findings: Transfer learning favours ViTs under standard training protocols, and self-supervised pretraining followed by fine-tuning lets ViTs outperform CNN counterparts.The paper also highlights improved explainability through transformer attention maps.
2. Related Work
Prior medical-imaging work used transformers mainly in custom architectures and specialized tasks, while the benefits of transfer learning and self-supervision for vanilla ViTs remained uncertain. This paper addresses those gaps by comparing pure ViTs with CNNs in medical diagnosis.
- Medical vision transformers: Medical applications of ViTs have focused mainly on segmentation, with fewer studies addressing registration or detection.Existing approaches generally combine transformer or attention modules with convolutional feature extractors.
- Medical vision transformers: Prior medical-imaging studies did not evaluate pure, off-the-shelf vision transformers.This leaves the plug-and-play use of vanilla ViTs insufficiently studied.
- Transfer learning: Whether ViTs benefit from transfer learning in medical domains remained unexplored, despite established transfer-learning benefits for CNNs.The source of CNN gains—feature reuse versus improved initialization and weight statistics—was disputed.
- Self-supervised learning: Self-supervised pretraining achieved strong results for CNNs on chest X-rays and skin lesions, but its combination with ViTs had not been evaluated in medical imaging.The unresolved comparison concerned whether self-supervised ViTs could outperform CNN counterparts.
3. Methods
The study compares parameter- and compute-matched vanilla CNN and ViT models across three medical datasets and three initialization strategies. Evaluation uses repeated experiments with standard training procedures and task-specific metrics.
- Models: RESNET50 and DEIT-S with 16^16 tokens were selected as comparable CNN and ViT models in parameters, memory requirements, and compute.The comparison minimizes architectural and resource differences.
- Initialization strategies: The experiments use random initialization, supervised ImageNet transfer learning, and target-dataset self-supervised pretraining followed by fine-tuning.Self-supervised pretraining is initialized from supervised ImageNet weights.
- Datasets: APTOS 2019 classifies 3,662 retinal images into five diabetic-retinopathy severity categories.Its split is adjusted because of the dataset’s small size.
- Evaluation: Results are reported as medians with standard deviations over five repetitions using quadratic Cohen Kappa, recall, and ROC-AUC for the three tasks.The metrics correspond respectively to APTOS2019, ISIC2019, and CBIS-DDSM.
4. Experiments
The experiments show that randomly initialized CNNs outperform ViTs in the low-data setting, whereas supervised ImageNet pretraining closes the gap. Self-supervised pretraining improves both models and gives ViTs a small advantage.
- Random initialization: CNNs outperform randomly initialized ViTs by a large margin across all evaluated tasks.The result is attributed in the discussion to ViTs’ lack of inductive bias in limited-data settings.
- ImageNet pretraining: Supervised ImageNet pretraining was tested by initializing both models with ImageNet weights before fine-tuning on medical tasks.This directly evaluates whether the standard CNN transfer-learning procedure also works for vanilla ViTs.
- Self-supervised pretraining: Both ViTs and CNNs perform better with self-supervised pretraining, with ViTs outperforming CNNs by a small margin.The reported advantage is observed after self-supervised pretraining and supervised fine-tuning.
- Overall comparison: Self-supervised pretraining followed by fine-tuning produces the best overall performance, with ViTs holding a small advantage over comparable CNNs.The advantage is consistent with the paper’s main experimental conclusion.
5. Discussion
Across medical-image datasets, CNNs are stronger from scratch, while transfer learning narrows the gap and self-supervised pretraining gives ViTs a small advantage. ViTs also provide finer, built-in attention visualizations, although larger unlabeled datasets may offer further opportunity.
- CNNs outperform ViTs when trained from scratch in the low-data regime across medical-image datasets.This pattern is consistent with the view that transformers lack CNN-like inductive bias.
- Self-supervised pretraining yields small but consistent improvements for both architectures, with the best overall performance from self-supervised ViTs.The strong self-supervision advantage previously reported for ViTs in larger natural-image datasets does not appear in this low-data regime.
- Large unlabeled medical-image collections could enable self-supervised learning when only a small fraction of images are labeled.
- Transfer learning bridges the performance gap between CNNs and ViTs, producing similar performance.
- Interpretability: ViT attention maps provide finer, built-in visualizations of model focus than CNN saliency methods.Examples localize attention to lesion boundaries, hemorrhages and exudates, and dense breast regions.
6. Conclusion
The paper concludes that vanilla transformers can reliably replace CNNs for medical imaging with little effort. They match CNN performance on small datasets when transfer learning is used, and can marginally outperform CNNs with self-supervised pretraining while offering high-resolution saliency maps.
- Vanilla transformers can reliably replace CNNs on medical image tasks with little effort.
- ViTs require transfer learning to reach CNN-level performance on small medical datasets.
- Self-supervised pretraining lets ViTs marginally outperform CNNs with limited samples, with the margin expected to grow as sample size increases.
- ViTs add built-in high-resolution saliency maps for understanding model decisions.