C AN T RANSFORMERS R EPLACE CNN S FOR M EDICAL I MAGE C LASSIFICATION ?
Anonymous authors
Paper under double-blind review
A BSTRACT
Convolutional Neural Networks (CNNs) have reigned for a decade as the de facto approach to automated medical image diagnosis, pushing the state-of-the-art in classiﬁcation, detection and segmentation tasks.
Recently, vision transformers (ViTs) have appeared as a competitive alternative to CNNs, yielding impressive levels of performance in the natural image domain, while possessing several interesting properties that could prove beneﬁcial for medical imaging tasks.
In this work, we explore whether it is feasible to switch to transformer-based models for medical image classiﬁcation as well, or if we should keep working with CNNs – can we trivially replace CNNs with transformers?
We consider this question in a series of experiments on several standard medical image benchmark datasets and tasks.
Our ﬁndings show that, while CNNs perform better if trained from scratch, off-the-shelf vision transformers can perform on par with CNNs when pretrained on ImageNet, both in a supervised and self-supervised setting.
I NTRODUCTION
Breakthroughs in medical image analysis over the past decade have been largely fuelled by convolutional neural networks (CNNs).
CNN architectures have served as the workhorse for numerous medical image analysis tasks including breast cancer detection (Wu et al., 2019), ultrasound diagnosis (Christiansen et al., 2021), and diabetic retinopathy (De Fauw et al., 2018).
Whether applied directly as a plug-and-play solution or used as the backbone for a bespoke model, CNNs such as
ResNet (He et al., 2016) are the dominant model for medical image analysis.
Recently, however, vision transformers have gained increased popularity for natural image recognition tasks, possibly signalling a transition from convolution-based feature extractors to attention-based models (ViTs).
This raises the question: can ViTs also replace CNNs in the medical imaging domain?
In the natural image domain, transformers have been shown to outperform CNNs on standard vision tasks such as I MAGE N ET classiﬁcation (Dosovitskiy et al., 2020), as well as in object detection (Carion et al., 2020) and semantic segmentation (Ranftl et al., 2021).
Interestingly, transformers succeed despite their lack of inductive biases tied to the convolution operation.
This may be explained, in part, by the attention mechanism central to transformers which offers several key advantages over convolutions: it explicitly models and more efﬁciently captures long-range relationships (Raghu et al., 2021), and it has the capacity for adaptive modeling via dynamically computed self-attention weights that capture relationships between tokens.
In addition, it provides a type of built-in saliency, giving insight as to what the model focused on (Caron et al., 2021).
Over the years, we have gained a solid understanding of how CNNs perform on medical images, and which techniques beneﬁt them.
For example, it is well-known that performance drops dramatically when training data is scarce (Cho et al., 2015).
This problem is particularly acute in the medical imaging domain, where datasets are smaller and often accompanied by less reliable labels due to the inherent ambiguity of medical diagnosis.
For CNNs, the standard solution is to employ transfer learning (Morid et al., 2020): typically, a model is pretrained on a larger dataset such as I MAGE N ET (Deng et al., 2009) and then ﬁne-tuned for speciﬁc tasks using smaller, specialized datasets.
For medical imaging, CNNs pre-trained on I MAGE N ET typically outperform those trained from scratch, both in terms of ﬁnal performance and reduced training time despite the differences between the two domains (Morid et al., 2020; Raghu et al., 2019; Tajbakhsh et al., 2016). However, recent studies show that, for CNNs, transfer from I MAGE N ET to medical tasks is not as useful as previously thought (Raghu et al., 2019; Neyshabur et al., 2020).
Supervised transfer learning requires many annotations.
This can be problematic for medical tasks where annotations are costly and require specialized experts.
Self-supervised approaches, on the other hand, can learn powerful representations by leveraging the intrinsic structure present in the image, rather than explicit labels (He et al., 2020; Chen et al., 2020).
Although self-supervision performs best when large amounts of unlabelled data are available, it has been shown to be effective with CNNs for certain medical image analysis tasks (Azizi et al., 2021; Sowrirajan et al., 2021).
In contrast to CNNs, we know relatively little about how vision transformers perform at medical image classiﬁcation.
Are vanilla ViTs competitive with CNNs?
Despite their success in the natural image domain, there are questions that cast doubt as to whether that success will translate to medical tasks.
Evidence suggests that ViTs require very large datasets to outperform CNNs – in Dosovitskiy et al.
the beneﬁts of ViTs only became evident when Google’s private 300 million image dataset, JFT-300M, was used for pretraining.
Reliance on data of this scale may be a barrier to the application of transformers for medical tasks.
As outlined above, CNNs rely on techniques such as transfer learning and self-supervision to perform well on medical datasets.
Are these techniques as effective for transformers as they are for CNNs?
Finally, the variety and type of patterns and textures in medical images differ signiﬁcantly from the natural domain.
Studies on CNNs indicate that the beneﬁts from transfer diminish with the distance from the source domain (Azizpour et al., 2016; Raghu et al., 2019; Neyshabur et al., 2020)
Can ViTs cope with this difference?
In this work we explore whether it is feasible for ViTs to replace CNNs for medical image classiﬁcation, and if there is an advantage of doing so.
Given the enormous computational cost of considering the countless variations in CNN and ViT architectures, we limit our study to prototypical CNN and ViT models over a representative set of well-known publicly available medical image benchmark datasets.
Through these experiments we show that:
• ViTs pretrained on I MAGE N ET perform comparably to CNNs for classiﬁcation tasks on medical images.
While CNNs outperform ViTs when trained from scratch, ViTs receive a larger boost in performance from pretraining on I MAGE N ET .
• Despite their reliance on large datasets, self-supervised ViTs perform on par or better than CNNs, if I MAGE N ET pretraining is utilized, albeit by a small margin.
• Features of off-the-shelf I MAGE N ET pre-trained ViTs are versatile. k -NN tests show selfsupervised ViT representations outperform CNNs, indicating their potential for other tasks.
We also show they can serve as drop-in replacements for segmentation.
These ﬁndings, along with additional ablation studies, suggest that ViTs can be used in medical image analysis, while at the same time potentially gaining from other properties of ViTs such as built-in explainability.
The source code to reproduce our work is included in the supplementary material, and we will release it at the time of publication.
R ELATED W ORK
Transformers in vision problems.
Following the success of transformers (Vaswani et al., 2017a) in natural language processing (NLP), attention-based models captured the interest of the vision community and inspired numerous improvements to CNNs (Hu et al., 2018; Zhang et al., 2020).
The ﬁrst work to show that vanilla transformers from NLP can be applied with minimal changes to largescale computer vision tasks and compete with standard CNNs was from Dosovitskiy et al.
(2020), which introduced the term vision transformer , or ViT 1 .
Like their NLP counterparts, the original vision transformers required an enormous corpus of training data to perform well.
To overcome this problem, Touvron et al.
introduced D EI Ts, which proposed a distillation token in addition to the CLS token that allows transformers to learn more efﬁciently on smaller datasets.
As the number of applications of vision transformers and architectural variations have exploded (Khan et al.
(2021) provides a good review), DeiT and the original ViT (Dosovitskiy et al., 2020) have become the standard benchmark architectures for vision transformers.
Vision transformers in medical imaging.
Given the very recent appearance of vision transformers, their application in medical image analysis has been limited.
Segmentation has seen the most applications of ViTs.
Chen et al.
(2021a),
Chang et al.
(2021) and Hatamizadeh et al.
(2021) independently suggested different methods for replacing CNN encoders with transformers in U-Nets (Ronneberger et al., 2015), resulting in improved performance for several medical segmentation tasks.
Vanilla transformers were not used in any of these works, however.
In each, modiﬁcations to the standard vision transformer from Dosovitskiy et al.
(2020) are proposed, either to adapt to the U-Net framework (Chen et al., 2021a) or to add components from CNNs (Chang et al., 2021).
Other transformer-based adaptations of CNN based-architectures suggested for medical segmentation tasks include Zhang et al.
(2021), who combined pyramid networks with transformers, and
Lopez Pinaya et al.
(2021) who used transformers in variational auto encoders (VAEs) for segmentation and anomaly detection in brain MR imaging.
Only a handful of studies have tackled tasks other than segmentation, such as 3-D image registration (Chen et al., 2021b) and detection (Duong et al., 2021).
The only work so far, to our knowledge, applying transformers to medical image classiﬁcation is a CNN/transformer hybrid model proposed by Dai et al.
(2021) applied to a small (344 images) private MRI dataset.
Notably, none of these works consider pure, off-the-shelf vision transformers – all propose custom architectures combining transformer/attention modules with components from convolutional feature extractors.
Improved initialization with transfer- and self-supervised learning.
Medical imaging datasets are typically orders of magnitude smaller than natural image datasets due to cost, privacy concerns, and the rarity of certain diseases.
A common strategy to learn good representations for smaller datasets is transfer learning.
For medical imaging, it is well-known that CNNs usually beneﬁt from transfer learning with I MAGE N ET (Morid et al., 2020) despite the distinct differences between the domains.
However, the question of whether ViTs beneﬁt similarly from transfer learning to medical domains has yet to be explored.
While transfer learning is one option to initialize models with good features, another more recent approach is self-supervised pre-training.
Recent advances in self-supervised learning have dramatically improved performance of label-free learning.
State-of-the-art methods such as DINO (Caron et al., 2021) and BYOL (Grill et al., 2020) have reached performance on par with supervised learning on I MAGE N ET and other standard benchmarks.
While these top-performing methods have not yet been proven for medical imaging, there has been some work using earlier self-supervision schemes on medical data.
Azizi et al.
(2021) adopted SimCLR (Chen et al., 2020), a self-supervised contrastive learning method, to pretrain CNNs.
This yielded state-of-the-art results for predictions on chest X-rays and skin lesions.
Similarly, Sowrirajan et al.
(2021) employed MoCo (He et al., 2020) pre-training on a target chest X-ray dataset, demonstrating again the power of this approach.
While these works exhibit promising improvements thanks to self-supervision, they have all employed
CNN-based encoders.
It has yet to be shown how self-supervised learning combined with ViTs performs in medical imaging, and how this combination compares to its CNN counterparts.
M ETHODS
The main question we investigate is whether prototypical vision transformers can be used as a dropin alternative to CNNs for medical diagnostic tasks.
More concretely, we consider how each model type performs on various tasks in different domains, with different types of initialization, and across a range of model capacities.
To that end, we conducted a series of experiments to compare representative ViTs and CNNs for various medical image analysis tasks under the same conditions.
Hyperparameters such as weight decay and augmentations were optimized for CNNs and used for both CNNs and ViTs – with the exception of the initial learning rate, which was determined individually for each model type using a grid search.
To keep our study tractable, we selected representative CNN and ViT model types as there are too many architecture variations to consider each and every one.
For CNNs, an obvious choice is the
R ES N ET family (He et al., 2016), as it is the most common and highly cited CNN backbone, and recent works have shown that R ES N ET s are competitive with more recent CNNs when modern training methods are applied (Bello et al., 2021).
The choice for a representative vision transformer is less clear because the ﬁeld is still developing.
We considered several options including the original
ViT (Dosovitskiy et al., 2020), SWIN transformers (Liu et al., 2021), CoaT (Xu et al., 2021), and focal transformers (Yang et al., 2021), but we selected the D EI T family (Touvron et al., 2021) for the following reasons: (1) it is one of the earliest and most established vision transformers, with the most citations aside from the original ViT, (2) it is similar in spirit to a pure transformer, (3) it was the ﬁrst to show that vision transformers can compete on mid-sized datasets with short training times, (4) it retains the interpretability properties of the original transformer.
As mentioned above, CNNs rely on initialization strategies to improve performance for smaller datasets.
Accordingly, we consider three common initialization strategies: randomly initialized weights, transfer learning from I MAGE N ET , and self-supervised pretraining on the target dataset.
Using these models and initialization strategies, we consider the following datasets and tasks:
Medical image classiﬁcation.
Five standard medical image classiﬁcation datasets were chosen to be representative of a diverse set of target domains.
These cover different imaging modalities, color distributions, dataset sizes, and tasks, with expert labels.
• APTOS 2019 –
In this dataset, the task is classiﬁcation of diabetic retinopathy images intocategories of disease severity (Kaggle, 2019).
APTOS 2019 contains 3,662 high-resolution retinal images.
• CBIS-DDSM – A mammography dataset containing 10,239 images.
The task is to detect the presence of masses in the mammograms (Sawyer-Lee et al., 2016; Lee et al., 2017; Clark et al., 2013).
• ISIC 2019 –
Here, the task is to classify 25,333 dermoscopic images among nine different diagnostic categories of skin lesions (Tschandl et al., 2018; Codella et al., 2018; Combalia et al., 2019).
CheXpert – This dataset contains 224,316 chest X-rays with labels over 14 categories of diagnostic observations (Irvin et al., 2019).
PatchCamelyon – Sourced from the Camelyon16 segmentation challenge (Bejnordi et al., 2017), this dataset contains 327,680 patches of H&E stained WSIs of sentinel lymph node sections.
The task is to classify each patch as cancerous or normal (Veeling et al., 2018).
Medical image segmentation.
Although there has been some recent work applying vision transformers to medical image segmentation, all works thus far have proposed hybrid architectures.
Here, we evaluate the ability of vanilla ViTs to directly replace the encoder of a standard segmentation framework designed for CNNs.
In particular, we use D EEP L AB 3 (Chen et al., 2017) featuring a
R ES N ET 50 encoder.
We merely replace the encoder of the segmentation model with D EI T-S and feed the outputs of the CLS token to the decoder of D EEP L AB 3.
As D EEP L AB 3 was designed and optimized for R ES N ET s, CNNs have a clear advantage in this setup.
Nevertheless, these experiments can shed some light into the quality and versatility of ViT representations for segmentation.
Note that, unlike the classiﬁcation experiments, all models were initialized with I MAGE N ET pretrained weights.
We consider the following medical image segmentation datasets:
• ISIC 2018 – This dataset consists of 2,694 dermoscopic images with their corresponding pixel-wise binary annotations.
The task is to segment skin lesions (Codella et al., 2019; Tschandl et al., 2018).
• CSAW-S – This dataset contains 338 images of mammographic screenings.
The task is pixel-wise segmentation of tumor masses (Matsoukas et al., 2020).
Experimental setup.
We selected the R ES N ET family (He et al., 2016) as representative CNNs, and D EI T 2 (Touvron et al., 2021) as representative ViTs.
Speciﬁcally, we test R ES N ET -18,
R ES N ET -50, and R ES N ET -152 against D EI T-T, D EI T-S, and
D EI T-B.
These models were chosen because they are easily comparable in the number of parameters, memory requirements, and compute.
They are also some of simplest, most commonly used, and versatile architectures of their respective type.
Note that ViTs have an additional parameter, patch size, that directly inﬂuences the memory and computational requirement.
We use the default 16 × 16 patches to keep the number of parameters comparable to the CNN models.
[Table]
Using these models [Table] we compare three commonly used initialization strategies:
Randomly initialized weights (Kaiming initialization from He et al.
Transfer learning using supervised I MAGE N ET pretrained weights, 3.
Self-supervised pretraining on the target dataset , following I MAGE N ET initialization.
After initialization using the strategies above, the models are ﬁne-tuned on the target data using the procedure described below.
Early experiments indicated that initialization using self-supervision on the target dataset consistently outperforms self-supervision on I MAGE N ET .
We also determined that
DINO performs better than other self-supervision strategies such as BYOL (Grill et al., 2020).
Training procedure.
For supervised training, we use the A DAM optimizer (Kingma & Ba, 2014) 3 .
We performed independent grid searches to ﬁnd suitable learning rates, and found that 10 − 4 works best for both pretrained CNNs and ViTs, while 10 − 3 is best for random initialization.
We used these as base learning rates for the optimizer along with default 1,000 warm-up iterations.
When the validation metrics saturated, the learning rate was dropped by a factor of 10 until it reached its ﬁnal value of 10 − 6 .
For classiﬁcation tasks, images were resized to 256 × 256 with the following augmentations applied: normalization; color jitter which consists of brightness, contrast, saturation, hue; horizontal ﬂip; vertical ﬂip; and random resized crops.
For segmentation, all images were resized to 512 ×
and we use the same augmentation methods as for classiﬁcation except for
C SAW -S, where elastic transforms are also employed.
For each run, we select the checkpoint with highest validation performance.
For self-supervision, pretraining starts with I MAGE N ET initialization, then applies DINO (Caron et al., 2021) on the target data following the default settings – except for three small changes determined to work better on medical data: (1) the base learning rate was set to 10 − 4 , (2) the initial weight decay is set at 10 − 5 and increased to 10 − 4 using a cosine schedule, and (3) we used an EMA of 0.99.
The same settings were used for both CNNs and ViTs; both were pre-trained for 300 epochs using a batch size of 256, followed by ﬁne-tuning as described above.
The classiﬁcation datasets were divided into train/test/validation splits (80/10/10), with the exception of APTOS2019, which was divided 70/15/15 due to its small size.
For the segmentation datasets, we split the training set into train/validation with a 90/10 ratio.
For ISIC2018 and CSAW-S, we used the provided validation and test sets for model evaluation.
Evaluation.
Unless otherwise speciﬁed, each experimental condition was repeated ﬁve times.
We report the median and standard deviation of the appropriate metric for each dataset: Cohen Kappa,
Recall, ROC-AUC, and IoU. For classiﬁcation, we measure performance of the:
• initialized model representations, through a k -NN evaluation protocol ( Caron
et al., 2021) • ﬁnal model after ﬁne-tuning on the target dataset.
The k -NN evaluation protocol is a technique to measure the usefulness of learned representations for some target task; it is typically applied before ﬁne-tuning.
For CNNs, it works by freezing the pretrained encoder, passing test examples through the network, and performing a k -NN classiﬁcation using the cosine similarity between embeddings of the test images and the k nearest training images ( k = 200 in our case). For ViTs, the principle is the same except the output of the CLS token is used instead of the CNN’s penultimate layer.
[Figure]
MethodDINO 0.
APTOS2019,
[Figure]
ImageNet Initialization MethodDINO 0.
1.0 DDSM, ROC-AUC
Random ImageNet Initialization MethodDINO 0.
1.0 ISIC2019, Recall
Random ImageNet Initialization MethodDINO 0.
1.0 CheXpert, ROC-AUC
[Figure]
Initialization MethodDINO 0.
[Figure]
1.0 Camelyon, ROC-AUC
[Figure]
E XPERIMENTAL R ESULTS
Are randomly initialized vision transformers useful?
Our ﬁrst experiment compares the performance of ViTs against CNNs when initialized with random weights (He et al., 2015).
The results in
Table 1 and Figure 1 indicate that in this setting, CNNs outperform ViTs across the board.
This is in line with previous observations in the natural image domain, where ViTs trained on smaller datasets are outperformed by similarly-sized CNNs, a trend that was attributed to the vision transformer’s lack of inductive bias (Dosovitskiy et al., 2020).
Note that the performance gap appears to shrink as the number of training examples n increases.
However, since most medical imaging datasets are of modest size, the usefulness of randomly initialized ViTs appears to be limited.
Does pretraining transformers on I MAGE N ET work in the medical image domain?
To deal with the smaller size of medical imaging datasets, CNNs are often initialized with I MAGE N ET pretrained weights.
This typically results in better performance than random initialization (Morid et al., 2020), while recent works have questioned the usefulness of this approach for CNNs (Raghu et al., 2019; Neyshabur et al., 2020).
We investigate if ViTs beneﬁt from I MAGE N ET pre-training in the medical domain and to which extent.
To test this, we initialize all models with supervised I MAGE N ET -pretrained weights and then ﬁne-tune on the target data using the procedure described in Section 3.
The results in Table 1 and Figure 1 show that both CNNs and ViTs beneﬁt from I MAGE N ET initialization.
However, ViTs appear to beneﬁt more from transfer learning, as they make up for the gap observed using random initialization, performing on par with their CNN counterparts.
Do transformers beneﬁt from self-supervision in the medical image domain?
Recent selfsupervised learning schemes such as DINO and BYOL achieve performance near that of supervised learning.
When used for pretraining in combination with supervised ﬁne-tuning, they can achieve a new state-of-the-art (Caron et al., 2021;
Grill et al., 2020).
While this phenomenon has been demonstrated for CNNs and ViTs on big data in the natural image domain, it is not clear whether self-supervised pretraining of ViTs helps for medical imaging tasks.
To test if ViTs beneﬁt from selfsupervision, we adopt the learning scheme of DINO, which can be readily applied to both CNNs and ViTs (Caron et al., 2021).
The results reported in Table 1 and Figure 1 show that both ViTs and CNNs perform better with self-supervised pretraining.
ViTs appear to perform on par, or better than CNNs in this setting, albeit by a small margin.
Are pretrained ViT features useful for medical imaging, even without ﬁne-tuning?
When data is particularly scarce, features from pre-trained networks can sometimes prove useful, even without ﬁne-tuning on the target domain.
The low dimensional embeddings from a good feature extractor can be used for tasks such as clustering or few-shot classiﬁcation.
We investigate whether different types of pretraining in ViTs yields useful representations.
We measure this by applying the k -NN evaluation protocol described in Section 3.
The embeddings from the penultimate layer of the CNN and the CLS token of the ViT are used to perform k -NN classiﬁcation of test images, assigning labels based on the most similar examples from the training set.
We test whether outof-domain I MAGE N ET pretraining or in-domain self-supervised DINO pretraining (Caron et al., 2021) yields useful representations when no supervised ﬁne-tuning has been done.
The results [Table] [Table] appear in Figure 1 and Table 3 of the Appendix A, where we compare embeddings for D EI T-S and
R ES N ET 50.
Pretraining with I MAGE N ET yields surprisingly useful features for both CNNs and
ViTs, considering that the embeddings were learned on out-of-domain data.
We observe even more impressive results with in-domain self-supervised DINO pretraining.
ViTs appear to beneﬁt more from self-supervision than CNNs – in fact, ViTs trained without any labels using DINO perform similar or better than the same model trained from scratch with full supervision.
The results indicate that representations from pretrained vision transformers, both in-domain and out-of domain, can serve as a solid starting point for transfer learning or k -shot methods.
Do ViTs learn meaningful representations for other tasks like segmentation?
Segmentation is an important medical image analysis task.
While previous works have shown that custom attentionbased segmentation models can outperform CNNs (see Section 2), here we ask: can we achieve comparable performance to CNNs by merely replacing the encoder of a CNN-based segmentation model with a vision transformer? This would demonstrate versatility of ViT representations for other tasks.
We consider D EEP L AB 3 (Chen et al., 2017), a mainstream segmentation model with a R ES N ET 50 encoder.
We simply replace the R ES N ET 50 with D EI T-S, and train the model as prescribed in Section 3.
Despite the model being designed for R ES N ET 50 and D EI T-S having fewer parameters connected to the D EEP L AB 3 head, we observe that ViTs perform on par or better than CNNs in our segmentation tasks.
The results in Table 2 and Figure 2 clearly indicate that vision transformers are able to produce high quality embeddings for segmentation, even in this disadvantaged setting.
How does capacity of ViT models affect medical image classiﬁcation?
To understand the practical implications of ViT model size for medical image classiﬁcation, we conducted an ablation study comparing different model capacities of CNNs and ViTs.
All models were initialized with I MAGE N ET pretrained weights, and we followed the training settings outlined in Section 3.
The results are presented in Figure 3 and Table 4 in Appendix B. Both CNNs and ViTs seem to beneﬁt similarly from increased model capacity, within the margin of error.
For most of the datasets, switching from a small ViT to a bigger model results in minor improvements, with the exception of ISICwhere model size appears to be an important factor for optimal performance and APTOS where DeiT-T performs better than its larger variants. Perhaps the tiny size of APTOS2019 results in diminishing returns for large models.
Other ablation studies.
We also considered the effect of token size on ViT performance.
Tablein Appendix C shows that reducing patch size results in a slight boost in performance, but at the cost of increasing the memory footprint.
Finally, we investigate how ViTs attend to different regions of medical images.
In Figure 5 of Appendix D we compute the attention distance , the average distance in the image across which information is integrated in ViTs.
We ﬁnd that, in early layers, some heads attend to local regions and some to the entire image.
Deeper into the network, attention becomes more global – though more quickly for some datasets than for others. [Figure]
D ISCUSSION
Our investigation of prototypical CNNs and ViTs indicate that CNNs can readily be replaced with vision transformers in medical imaging tasks without sacriﬁcing performance.
The caveat being that employing some form of transfer learning is necessary.
Our experiments corroborate previous ﬁndings in the natural image domain and provide new insights, which we discuss below
[Figure]
Discussion and implications of ﬁndings.
Unsurprisingly, we found that CNNs outperform ViTs when trained from scratch on medical image datasets, which corroborates previous ﬁndings in the natural image domain (Dosovitskiy et al., 2020).
This trend appears consistently and ﬁts with the “insufﬁcient data to overcome the lack of inductive bias” argument.
Thus, the usefulness of randomly initialized ViTs appears to be limited in the medical imaging domain.
When initialized with supervised I MAGE N ET pretrained weights, the gap between CNN and ViT performance disappears on medical tasks.
The beneﬁts of supervised I MAGE N ET pretraining of
CNNs is well-known, but it was unexpected that ViTs would beneﬁt so strongly.
To the best of our knowledge, we are the ﬁrst to conﬁrm that supervised I MAGE N ET pretraining is so effective for ViTs in the medical domain.
This suggests that further improvements could be gained via transfer learning from other domains more closely related to the task, as is the case for CNNs (Azizpour et al., 2016).
The beneﬁts of I MAGE N ET pretraining was not only observed for the ﬁne-tuned classiﬁcation tasks, but also self-supervised learning, k -nn classiﬁcation and segmentation.
The best overall performance on medical images is obtained using ViTs with in-domain selfsupervision, where small improvements over CNNs and other initialization methods were recorded.
Our k -NN evaluation showed that ViT features learned this way are even strong enough to outperform supervised learning with random initialization.
Interestingly, we did not observe as strong of an advantage for self-supervised ViTs over I MAGE N ET pretraining as was previously reported in the natural image domain, e.g. in Caron et al.
We suspect this is due to the limited size of medical datasets, suggesting a tantalizing opportunity to apply self-supervision on larger and easier to obtain unlabeled medical image datasets, where greater beneﬁts may appear.
Our other ablation studies showed that ViT performance scales with capacity in a similar manner to
CNNs; that ViT performance can be improved by reducing the patch size; and that ViTs attend to information across the whole image even in the earliest layers.
Although individually unsurprising, it is worthwhile to conﬁrm these ﬁndings for medical images.
It should be noted however, that the memory demands of DeiTs increase quadratically with the image and patch size which might limit their application in large medical images.
However, newer ViT architectures (e.g. Liu et al.
(2021) and Yang et al.
(2021)) mitigate this issue – while demonstrating increased predictive performance.
Interesting properties of ViTs.
Our results indicate that switching from CNNs to ViTs can be done without compromising performance, but are there any other advantages to switching to ViTs?
There are a number of important differences between CNNs and ViTs.
We brieﬂy discuss some of these differences and why they may be interesting for medical image analysis.
Lack of inductive bias – ViTs do away with convolutional layers.
Our experiments indicate their implicit locality bias is only necessary when training from scratch on medical datasets.
For larger datasets, evidence suggests removing this bias improves performance (Dosovitskiy et al., 2020).
Global + local features – ViTs, unlike CNNs, can combine information from distant and local regions of the image, even in early layers (see Appendix D).
This information can be propagated [Figure] efﬁciently to later layers due to better utilization of residual connections that are unpolluted by pooling layers.
This may prove beneﬁcial for medical modalities which rely on local features.
It may be necessary to use transfer learning to beneﬁt from local features, as they require huge amounts of data to learn (Raghu et al., 2021).
Interpretability – transformers’ self-attention mechanism provides, for free, new insight into how the model makes decisions.
CNNs do not naturally lend themselves well to visualizing saliency.
Popular CNN explainability methods such as class activation maps (Zhou et al., 2016) and GradCAM (Selvaraju et al., 2017) provide coarse visualizations because of pooled layers.
Transformer tokens give a ﬁner picture of attention, and the self-attention maps explicitly model interactions between every region in the image.
In Figure 4, we show examples from each dataset along with
Grad-CAM visualizations of R ES N ET -50 and the top-50% self-attention of 16 × 16 D EI T-S CLS token heads.
The ViTs provide a clear, localized picture of attention, e.g. attention at the boundary of the skin lesion in ISIC, on hemorrhages and exudates in APTOS, the membrane of the lymphatic nodule in PatchCamelyon, the opacity in the Chest X-ray, and a dense region in CBIS-DDSM.
C ONCLUSION
To answer the question posed in the title: vanilla transformers can reliably replace CNNs on medical image classiﬁcation with little effort.
More precisely, ViTs reach the same level of performance as CNNs in an array of medical classiﬁcation and segmentation tasks, but they require transfer learning to do so.
But since I MAGE N ET pretraining is the standard approach for CNNs, this does not incur any additional cost in practice.
The best overall performance on medical imaging tasks is achieved using in-domain self-supervised pretraining, where ViTs show a small advantage over CNNs.
As the data size grows, this advantage is expected to grow as well, as observed by Caron et al.
(2021) and our own experiments.
Additionally, ViTs have a number of properties that make them attractive: they scale similarly to CNNs (or better), their lack of inductive bias, global attention and skip connections may improve performance, and their self-attention mechanism provides a clearer picture of saliency.
From a practitioner’s point of view, these beneﬁts are compelling enough to explore the use of ViTs in the medical domain.
Finally, modern CNNs have been studied extensively for 15 years, while the ﬁrst pure vision transformer appeared one year ago – the potential for ViTs to improve is considerable.
A PPENDIX
The Appendix contains several details and additional experimental results that did not ﬁt within the page limit, organized by the order in which they are referenced in the main article.
A k -NN EVALUATION OF D EI T-S AND R ES N ET
The k -NN evaluation protocol is a technique to measure the usefulness of learned representations for some target task; it is typically applied before ﬁne-tuning.
For CNNs, it works by freezing the pretrained encoder, passing test examples through the network, and performing a k -NN classiﬁcation using the cosine similarity between embeddings of the test images and the k nearest training images ( k = 200 in our case).
For ViTs, the principle is the same except the output of the CLS token is used instead of the CNN’s penultimate layer.
We applied the k -NN evaluation protocol to test whether out-of-domain I MAGE N ET pretraining or in-domain self-supervised DINO pretraining yields useful representations when no supervised ﬁnetuning has been done.
The results appear in Table 3.
Pretraining with I MAGE N ET yields surprisingly useful features for both CNNs and ViTs, considering that the embeddings were learned on outof-domain data.
We observe even more impressive results with in-domain self-supervised DINO pretraining.
ViTs appear to beneﬁt more from self-supervision than CNNs – in fact, ViTs trained with DINO perform similar or better than the same model trained from scratch with full supervision!
The results indicate that representations from pretrained vision transformers, both in-domain and out-of domain, can serve as a solid starting point for transfer learning or k -shot methods.
[Table]
B D EIT AND R ES N ET WITH DIFFERENT CAPACITIES
To understand the practical implications of ViT model size for medical image classiﬁcation, we compare different model capacities of CNNs and ViTs.
Speciﬁcally, we test R ES N ET -18, R ES N ET -
[Table]
and R ES N ET -152 against D EI T-T, D EI T-S, and
D EI T-B.
These models were chosen because they are easily comparable in the number of parameters, memory requirements, and compute.
They are also some of simplest, most commonly used, and versatile architectures of their respective type.
All models were initialized with I MAGE N ET pretrained weights, and trained following the settings outlined in Section 3.
The results appear in
Table 4.
Both CNNs and ViTs seem to beneﬁt similarly from increased model capacity, within the margin of error.
For most of the datasets, switching from a small ViT to a bigger model results in minor improvements, with the exception of ISIC 2019 where model size appears to be an important factor for optimal performance.
[Table]
model types with similar capacity perform approximately on par in most of the datasets, indicating that ViTs scale similarly to CNNs.
The numbers in this table correspond to Figure 3.
C E FFECT OF P ATCH S IZE
ViTs have an additional parameter, patch size, that directly inﬂuences the memory and computational requirement.
We compare the default 16 × 16 patches, used in this work, with 8 × 8 patches, as it has been noted in previous works that reducing the patch size leads to increased performance (Caron et al., 2021).
The results appear in Table 5.
We observe that smaller patch sizes result in marginally better performance, though to a lesser degree than has been observed in the natural image domain.
[Table]
D M EAN A TTENDED D ISTANCE OF V I T S
We compute the attention distance , the average distance in the image across which information is integrated in ViTs.
As reported previously (e.g. Dosovitskiy et al.
(2020)), in early layers, some heads attend to local regions and some to the entire image.
Deeper into the network, attention becomes more global – though more quickly for some datasets than for others
[Figure]