Is our attention misplaced?—comparing feature importance of deep learning models to radiologist annotations

Introduction

Background

In 2012, a convolutional neural network (CNN) won the ImageNet Large Scale Visual Recognition Challenge (1) in which 1.2 million images are classified into 1,000 categories. The winning entry, AlexNet (2), achieved state-of-the-art performance which exceeded the second-place entry by over 10% in classification error. Since then, there has been a rapid growth in interest in using deep learning, particularly CNNs, in image and computer vision related tasks, quickly becoming the gold standard in many applications such as segmentation (3).

Unfortunately, deep learning typically requires large datasets in order to be successful and thus transfer learning started gaining substantial popularity to reduce training times and data requirements. Transfer learning by re-training previously trained models relies on the context of the original training being similar to the new task. Transfer learning from common, miscellaneous objects has become commonplace, but in certain contexts, such as medical imaging, it may not be appropriate as demonstrated by Raghu et al. (4).

The transformer architecture was first proposed in 2017 by Vaswani et al. (5). Transformers originally gained popularity as a deep learning method for large language models (LLMs) used for natural language processing and were eventually adapted for computer vision tasks with vision transformer (ViT) in 2019 (6). Since then, multiple variations and developments have been made to adapt this architecture to differing applications.

Rationale and knowledge gap

Transformers currently show competitive performance and a similar rise in popularity as the gold standard CNNs, while also uniquely proposing the use of a novel attention mechanism. Many researchers are now looking at ViTs as a replacement to the gold standard CNNs; however, there is little work done to evaluate if the transformer architecture, specifically the new attention mechanism, is more receptive to salient information of the input images. This issue is becoming increasingly critical as machine learning and artificial intelligence continue to expand its presence in healthcare where explainability can be equally critical as the diagnosis itself. For example, the cause and location of the diagnosis is important for treatment planning. Is the attention that ViTs received misplaced?

Objective

While machine learning shows promise to improve clinical decision-making and alleviate overburden, there are numerous barriers to adoption centered on explainability. Here we examine the attention mechanism and compare training methods that may lead to more explainable models than the current standard of transfer learning. Specifically, this study aimed to compare levels of agreement between the novel attention mechanism introduced by ViTs and important pixels of a CNN to radiologist annotations. This was meant to determine the receptiveness of the two neural network types to medical information in hopes of fostering more insight when developing machine learning tools with more explainability and better hopes for clinical adoption.

Methods

This study tested models using two publicly available datasets. The main dataset was the VinDr-CXR dataset (7) consisting of 18,000 chest X-ray images as it is the larger dataset where the bulk of the original development was done. The second dataset was the Mini-DDSM dataset (8) consisting of 10,000 mammograms. Figure 1 shows sample images from both datasets. Preprocessing of the VinDr-CXR dataset included resizing the images to 224×224, normalizing, and flipping contrast if necessary, as some images were provided with inverted contrast. Preprocessing of the Mini-DDSM dataset included cropping, resizing to 224×224, and normalizing. Resizing was done by first zero padding the images such that they were even in length and width, to prevent geometric distortions, then resampled. The dataset was labelled with various pathological conditions, but for the purpose of this study binary classification was used. All samples that were not labelled a healthy control subject were pooled for a single pathological label. The VinDr-CXR dataset was split into 6,143 pre-training, 8,854 downstream training, 3,000 test samples to create a balanced, downstream training portion. After preprocessing, the mini-DDSM dataset was also split into a pre-training, downstream training, and test sets. The data labelled as benign was reserved as the pre-training set, while the data labelled as normal or malignant was split into a training set of 3,109 and a test set of 1,536. Radiologists’ annotations were provided as co-ordinates (VinDr-CXR) or segmentation masks (Mini-DDSM) and converted into binary masks where 1 denotes a region of interest and 0 otherwise. The study was conducted in accordance with the Declaration of Helsinki and its subsequent amendments. No IRB approval and consent forms are needed as this study used publicly available data.

Figure 1 Sample images from the two datasets used for this study. Top row: samples from the VinDr-CXR dataset (6). From left to right: training sample labelled as a healthy subject, training sample labelled as a pathological case (interstitial lung disease, lung nodule/mass, lung tumor), a test sample labelled as a healthy subject, and a test sample labelled as pathological case (calcification). Bottom row: samples from the Mini-DDSM dataset (7). From left to right: a training sample labelled as a healthy subject, a training sample labelled as a malignant case, a test sample labelled as a healthy subject, and a test sample labelled as a malignant case.

CNN approach

To represent the CNN approach, ResNet50 models were trained. CNNs are a specific subset of neural networks that rely on convolutions to pool and learn features and are well suited to imaging and computer vision related tasks. There are three basic components to a CNN, convolutional layers, pooling layers, and fully connected layers (9).

Convolutional layers aim to learn the feature representation of the inputs. Each convolution layer is composed of several convolution kernels which are used to compute different feature maps (10). The feature maps are passed into a non-linear activation function for output to the next layer. They have parameters of stride and padding which describe the amount of overlap between kernels and the zero padding around borders respectively.

Pooling layers combine the previous layers outputs and essentially act to downsize the data (11). This serves two main purposes, to reduce the in-plane dimensionality and subsequently the number of learnable parameters, speeding up training and reducing underfitting. The most popular method is max pooling which takes a patch of feature inputs and outputs the maximum and is commonly implemented with a patch size of 2×2 and stride of 2 (9).

A flattening layer converts the outputs of the convolutional and pooling layers is implemented which converts the feature maps into a 1D-vector. Then fully connected layers typically follow which take all neurons in the previous layer and connects them to every neuron in the current layer to generate global semantic information (10). These fully connected layers are then connected to a non-linear activation function such as the rectified linear unit (ReLu) for output.

Unique to the ResNet architectures is the concept of residual connections which allow for deeper networks to be trained for better performance. Residual connections are implemented through skip connections and address the problem of vanishing gradients. ResNet50 gets its name as a 50-layer neural network created from blocks of 3 convolutional layers separated by skip connections. These blocks are preceded by a convolutional layer and pooling layer and followed by another pooling layer and fully connected layer.

As per common practice highlighted by Raghu et al. (4), in this study, training was done via transfer learning with pre-trained weights from ImageNet over 100 epochs. Pre-trained weights were frozen and treated as a feature extractor to be fed into a new classification layer. Feature extraction uses the previously trained model to generate features into a new “classification head”. Transfer learning, done through feature extraction, removes the classification layers from the previously trained network’s architecture. Only the weights of the hidden layers are reused. The outputs from the hidden layers are then used to train a new classification layer (or layers) in the desired task. It is worth noting that the previously trained weights are not adjusted. The previous weights are not trained to perform the new task. Instead, they are meant to create abstract features that the new layers can use to learn the new task.

ViT approach

To represent the transformer approach, base-size ViT models (6) were trained. The transformer architecture is composed of three key components: an encoder, a decoder, and the novel attention mechanism. Encoders are neural networks which aim to take input x and map it to an internal representation h through a mapping function f described by equation 1. A decoding neural network then takes the output h and creates a reconstruction r described by equation 2. By restricting the network sizes and reducing the input into a smaller latent vector, the encoder is forced to learn a representation of the input, which is expected to be only the salient information. Likewise, the decoder is forced to reconstruct based on a smaller input where the output is expected to be the original input.

b=f(x)[1]

r=g(b)[2]

Transformer models contain attention layers which conceptually inform the model to pay attention to certain portions of the input. The attention layers are described as mapping a query and a set of key-value pairs to an output. The query, keys, values, and output are all vectors. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key (5). This gives models the ability to distinguish context from other pertinent portions of the input, hence the prevalence and initial applications for natural language processing.

The ViT model itself is very similar to the original transformer architecture described by Vaswani et al. (5) and simply adapted to allow for 2D images as input. The 2D images are reshaped into a sequence of flattened 2D patches. This flattened sequence becomes the input to the transformer. Standard 1D learnable position embeddings are added to the patch embeddings to convey positional information.

In this study, ViT Models were first trained with a masked autoencoder in a generative pre-training step for 100 epochs followed by downstream training for another 100 epochs for classification. Generative pre-training was done as outlined by He et al. (12). In this generative pre-training stage, 75% of the training image is randomly masked before being given to the encoder. The decoder is tasked with generating a recreation of the original image. Once training is complete, the decoder is removed and the weights of the encoder are used for further training for the classification task. Of the healthy control training set, 6,143 were separated and used for pre-training. This left a balanced training set of 8,854 samples for downstream classification.

Implementation and attention masks

ResNet50 models were implemented with PyTorch and all transformer models were implemented with Huggingface and PyTorch. Relevant features from the ResNet50 models were extracted with the occlusion method included in Captum (13). The attention mechanisms of the ViTs were analyzed through the rollout method (14) to directly examine the attention mechanism. The choice of methods both produce maps that are considered important features for model decisions and therefore used as proxy for explanations and agreement calculations.

Occlusion is an iterative method of extracting the most relevant features for a correct classification. Different portions of the input are occluded, replaced by a grey square, and its effect on the output of the classifier is recorded. Occlusion generates two maps of features that correlate to either a correct or incorrect classification. These maps are combined in order to generate a single map of highly relevant features.

Attention rollout is implemented as described by Abnar et al. (14) by recursively multiplying the attention weight matrices in all the layers below it. Attention rollout is a means of investigating how information moves from the input layer to the embeddings in later layers. This method treats the network as a directed graph. Connected nodes represent input tokens and hidden embeddings. Each edge between nodes represents the attention from the nodes in each layer to those in the previous layer, and the weights of the edges are the attention weights. There are multiple paths between two nodes, so to compute the total information propagated, the values between all possible paths between two nodes are taken as a sum. Figure 2 depicts the training method and evaluations for both model types.

Figure 2 A visual outline of the training and evaluation process for the models. Top left: the training stage for the ResNet50 models. The models are pre-trained on ImageNet data and given the downstream training data with new classification layers through transfer learning. Top right: the training stage for the ViT models. A pre-training stage using a subset of the data is done rather than using pre-trained weights. The pretraining consists of a masked autoencoder. After pre-training, the decoder is removed and replaced with classification layers. The downstream training data is then used to train the model for classification. Bottom left: evaluation stage for the ResNet50 models. The test set is used to measure performance. The occlusion method is used to generate a map of important features. Bottom right: evaluation stage for the ViT models. The rollout method is used to generate a map of important features. CNN, convolutional neural network; ResNet, residual neural network; ViT, vision transformer.

Both outputs, the occlusion and the attention rollout, had a threshold applied with a cut-off of 0.5 to create binary attention masks. These attention mechanisms were compared to the radiologist annotations provided in the dataset through pixel-wise logical XNOR and averaged across the test set to measure the level of agreement.

Results

Aggregated model results are shown in Table 1. After training, the VinDr-CXR ResNet50 model had a test accuracy of 70.4% and the ViT model had a test accuracy of 77.4%. The Mini-DDSM ResNet50 model had a test accuracy of 90.26% and the ViT model had a test accuracy of 95.53%. There was a notable performance difference between the two datasets. The models trained on the Mini-DDSM dataset had approximately 20% higher accuracy than their counterparts trained on the VinDr-CXR dataset. The CNN models had higher sensitivity relative to their specificity.

The agreement was higher for the transformer models at 94.72% and 96.96% compared to 88.07% and 94.85% for the CNN models. Sample attention masks are shown in Figure 3. It is worth noting that the CNN models had both lower test accuracy and lower agreement. Likewise, the transformer models with higher test accuracy had higher agreement with radiologist annotations as well.

Figure 3 Sample outputs from both the VinDr-CXR and Mini-DDSM test sets. From left to right: test sample, occlusion output of highly relevant features from the ResNet model, rollout output from the ViT model, and the radiologist annotations. ResNet, residual neural network; ViT, vision transformer.

Discussion

Key findings

Overall, both the ResNet50 model and the ViT model had high levels of agreement with the radiologist annotations at 88.07% and 94.85% and 94.72% and 96.96% respectively; however, the ViTs had better performance in both test accuracy and agreement.

Strengths and limitations

A noted limitation is the subjective nature of radiologist annotations as a reference standard. It has been demonstrated that there is notable reader variability between radiologists (15,16). This is compounded by the annotations in the VinDr-CXR dataset which are provided as bounding boxes rather than finely segmented regions of interest (ROIs). This results in extra pixels being included in the ROI that would otherwise not be relevant to the radiologist’s attention (and hence their decision surrounding diagnosis).

Furthermore, in this study, only X-ray based images were included as the dataset (chest radiographs and mammography scans). The features of interest in radiographs are different from other imaging modalities such as positron emission tomography (PET), single photon emission computed tomography (SPECT), ultrasound or magnetic resonance imaging (MRI). These modality-based differences warrant further work to examine the generalizability of these results. Also, chest radiographs and mammograms are always performed with multiple views [lateral, frontal and rear views for radiographs, and medial-lateral oblique (MLO) and cranio-caudal for mammography]. This notable methodological difference was not differentiated in the study.

Comparison with similar research

Other work, such as Takahashi et al. (17), focuses on reviewing the literature and use traditional accuracy measurements as a performance indicator for comparison. However, this work aims to add another performance indicator, agreement with subject matter expertise. Deep learning models are often considered black box systems where it is difficult to understand how a decision is made. This adds an extra barrier to clinical adoption because explainability can be equally critical as the diagnosis itself. Therefore, other performance metrics may be of interest such as comparisons with subject matter expertise.

Explanation of findings

These findings may be due to multiple factors, one of note is the difference in training methods between the two models. However, as this study is a comparison between two architectures, training them as per their respective norms was deemed appropriate. Many transformers are trained with pre-training stages, with GPT being the most famous example (18) and at present, many CNN based applications are trained through transfer learning. The ResNet50 models were trained through transfer learning, using weights from a model originally trained to classify miscellaneous objects and re-trained for classifying the new dataset. On the other hand, the ViT models were also trained with a two-step process, but using weights learned from a subset of the data which were further trained for the classification task. Both methods have merits and advantages.

Transfer learning, based on models trained for the classification of miscellaneous objects, has the advantage of easily accessible data allowing for much larger training sets. The generative pre-training stage used to train the ViT has a comparatively smaller dataset, however the data is much more aligned with the downstream task and creates a more balanced downstream training set in contrast to the typically imbalanced data often found in medical datasets.

The higher sensitivity of the CNNs may be due to the pooling layers whereas attention matrices are calculated mostly through matrix multiplication. The pooling layers downsize the input and thus, may be responsible for the loss of specificity. On the other hand, there is no pooling done by the transformer models.

Implications and actions needed

Future work should look to expand the datasets and methods studied. There are a number of mechanisms proposed to extract relevant features from neural network models and many differing neural network model architectures, but only one method and one architecture was used here, for example, attention flow (14) for extracting information from transformers and saliency maps and integrated gradients (13) for CNNs.

Conclusions

This study aimed to compare model receptiveness to salient information between CNNs and ViTs in hopes to foster more insight for development of machine learning tools with more explainability and better odds for clinical adoption. Overall, in this study we demonstrated that ViTs warrant the surge in interest as the transformer model approach results in much better performance with an accuracy of 77.40% (9.9% increase) between the VinDr-CXR dataset models and 95.53% (5.8% increase) between the Mini-DDSM models. The transformer models also showed more agreement with the radiologist annotations with 6.65% more agreement between the VinDr-CXR dataset models and 2.11% between the Mini-DDSM models.

ViTs, and the generative pre-training methods, have many potential applications in medical imaging. Generative pre-training shows better results than transfer learning, further opening up more possibilities in training and model development. Of interest is the possibility of incorporating subject matter expertise in the form of agreement as a regularization tool implemented with the loss functions during training to push model training towards subject matter expert decisions.

Our attention has not been misplaced, rather, ViTs warrant more attention and future work. The novel attention mechanism and generative training methods show promise for fostering more insight for development of machine learning tools with better clinical adoption.

Acknowledgments

None.

Footnote

Peer Review File: Available at https://jmai.amegroups.com/article/view/10.21037/jmai-2025-125/prf

Funding: This work was supported in part by the Natural Sciences and Engineering Research Council (NSERC) of Canada (grant No. RGPIN-2023-05698 to M.D.N., grant No. RGPIN-2022-05438 to T.E.D.).

Conflicts of Interest: All authors have completed the ICMJE uniform disclosure form (available at https://jmai.amegroups.com/article/view/10.21037/jmai-2025-125/coif). M.D.N. is the CEO and co-founder of TBIFinder, Inc., a data analytics company focused on MRI-based brain injury detection. This entity has nothing to do with the current manuscript. M.D.N. has been paid by GE HealthCare Canada to provide lectures at a number of their symposia. T.E.D. is the President and Co-Founder of Lunar Medical Inc. The company was founded to solve the challenges of providing medical support for human spaceflight. This entity has nothing to do with the current manuscript. The other author has no conflicts of interest to declare.

Ethical Statement: The authors are accountable for all aspects of the work in ensuring that questions related to the accuracy or integrity of any part of the work are appropriately investigated and resolved. The study was conducted in accordance with the Declaration of Helsinki and its subsequent amendments. No IRB approval and consent forms are needed as this study used publicly available data.

Open Access Statement: This is an Open Access article distributed in accordance with the Creative Commons Attribution-NonCommercial-NoDerivs 4.0 International License (CC BY-NC-ND 4.0), which permits the non-commercial replication and distribution of the article with the strict proviso that no changes or edits are made and the original work is properly cited (including links to both the formal publication through the relevant DOI and the license). See: https://creativecommons.org/licenses/by-nc-nd/4.0/.

References

1. Russakovsky O, Deng J, Su H, Krause J, et al. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision 2015;115:211-52.
2. Krizhevsky A, Sutskever I, Hinton GE. Imagenet classification with deep convolutional neural networks. CACM 2017;60:84-90.
3. Ronneberger O, Fischer P, Brox T. U-net: Convolutional networks for biomedical image segmentation. In: International Conference on Medical image computing and computer-assisted intervention. Cham: Springer international publishing; 2015:234-41.
4. Raghu M, Zhang C, Kleinberg J, et al. Transfusion: Understanding transfer learning for medical imaging. 33rd Conference on Neural Information Processing Systems (NeurIPS 2019); Vancouver, Canada. Available online: https://openreview.net/pdf?id=B1e_PNSxLS
5. Vaswani A, Shazeer N, Parmar N, et al. Attention is all you need. 31st Conference on Neural Information Processing Systems (NIPS 2017); Long Beach, CA, USA. Available online: https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf
6. Dosovitskiy A, Beyer L, Kolesnikov A, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929. 2020 Oct 22.
7. Nguyen HQ, Lam K, Le LT, et al. VinDr-CXR: An open dataset of chest X-rays with radiologist’s annotations. Scientific Data 2022;9:429. [Crossref] [PubMed]
8. Lekamlage CD, Afzal F, Westerberg E, et al. Mini-DDSM: Mammography-based automatic age estimation. In: Proceedings of the 2020 3rd International Conference on Digital Medicine and Image Processing. 2020:1-6.
9. Yamashita R, Nishio M, Do RKG, et al. Convolutional neural networks: an overview and application in radiology. Insights Imaging 2018;9:611-29. [Crossref] [PubMed]
10. Gu J, Wang Z, Kuen J, et al. Recent advances in convolutional neural networks. Pattern Recognition 2018;77:354-77.
11. Albawi S, Mohammed TA, Al-Zawi S. Understanding of a convolutional neural network 2017 International Conference on Engineering and Technology (ICET); 21-23 August 2017; Antalya, Turkey. IEEE; 2017:1-6.
12. He K, Chen X, Xie S, et al. Masked autoencoders are scalable vision learners. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022:16000-9.
13. Kokhlikyan N, Miglani V, Martin M, et al. Captum: A unified and generic model interpretability library for pytorch. arXiv preprint arXiv:2009.07896. 2020 Sep 16.
14. Abnar S, Zuidema W. Quantifying attention flow in transformers. arXiv preprint arXiv:2005.00928. 2020 May 2.
15. Armato SG 3rd, McNitt-Gray MF, Reeves AP, et al. The Lung Image Database Consortium (LIDC): an evaluation of radiologist variability in the identification of lung nodules on CT scans. Acad Radiol 2007;14:1409-21. [Crossref] [PubMed]
16. McErlean A, Panicek DM, Zabor EC, et al. Intra- and interobserver variability in CT measurements in oncology. Radiology 2013;269:451-9. [Crossref] [PubMed]
17. Takahashi S, Sakaguchi Y, Kouno N, et al. Comparison of Vision Transformers and Convolutional Neural Networks in Medical Image Analysis: A Systematic Review. J Med Syst 2024;48:84. [Crossref] [PubMed]
18. Radford A, Narasimhan K, Salimans T, et al. Improving language understanding by generative pre-training. [Preprint]. 2018.