[Paper Review] A Comparative Study of CNN, ResNet, and Vision Transformers for Multi-Classification of Chest Diseases

Ananya Jain, Aviral Bhardwaj+2 more|arXiv (Cornell University)|May 31, 2024

AI in cancer detection4 citations

TL;DR

This study compares CNN, ResNet, and Vision Transformer (ViT) models for multi-classification of 14 chest diseases using the NIH Chest X-ray dataset. Pre-trained ViT models outperform CNN and ResNet, achieving 94% test accuracy when fine-tuned on ImageNet-21k, demonstrating the superiority of attention-based feature learning in medical imaging when pre-trained on large-scale data.

ABSTRACT

Large language models, notably utilizing Transformer architectures, have emerged as powerful tools due to their scalability and ability to process large amounts of data. Dosovitskiy et al. expanded this architecture to introduce Vision Transformers (ViT), extending its applicability to image processing tasks. Motivated by this advancement, we fine-tuned two variants of ViT models, one pre-trained on ImageNet and another trained from scratch, using the NIH Chest X-ray dataset containing over 100,000 frontal-view X-ray images. Our study evaluates the performance of these models in the multi-label classification of 14 distinct diseases, while using Convolutional Neural Networks (CNNs) and ResNet architectures as baseline models for comparison. Through rigorous assessment based on accuracy metrics, we identify that the pre-trained ViT model surpasses CNNs and ResNet in this multilabel classification task, highlighting its potential for accurate diagnosis of various lung conditions from chest X-ray images.

Motivation & Objective

• To evaluate and compare the performance of CNN, ResNet, and Vision Transformer (ViT) architectures in multi-label classification of 14 chest diseases from X-ray images.
• To assess the impact of pre-training on large-scale datasets like ImageNet-21k on ViT performance in medical imaging.
• To investigate the generalization capability and feature learning efficiency of ViTs versus traditional CNN-based models in low-resource medical imaging scenarios.
• To analyze attention maps and model interpretability in ViTs for clinical relevance and decision transparency.
• To identify limitations such as data imbalance and model optimization gaps in multi-class chest X-ray classification.

Proposed method

• Fine-tuned two ViT variants: one pre-trained on ImageNet and another trained from scratch using the NIH Chest X-ray dataset of 112,120 frontal-view X-ray images.
• Trained and evaluated a standard CNN and ResNet-50 as baseline models using the same dataset and training protocol.
• Employed transfer learning with a pre-trained InceptionResNetV2 model to generate attention maps by extracting attention weights from the last multi-head attention layer of the ViT encoder.
• Used 224×224 pixel RGB images as input, applied data augmentation, and trained models using cross-entropy loss with sigmoid activation for multi-label classification.
• Evaluated models using accuracy, AUC-ROC, and loss curves across training and validation sets to compare performance and generalization.
• Visualized attention maps by averaging attention weights across heads and patches, then upscaling and overlaying them on original images to interpret model predictions.

Experimental results

Research questions

• RQ1 Does a pre-trained Vision Transformer outperform CNN and ResNet in multi-label classification of chest X-ray images?
• RQ2 How does pre-training on large-scale datasets like ImageNet-21k affect ViT performance in medical image classification?
• RQ3 What is the role of attention mechanisms in identifying clinically relevant regions in chest X-rays compared to convolutional feature maps?
• RQ4 How do data imbalance and limited examples for rare diseases affect model performance across architectures?
• RQ5 Can ViT models achieve comparable or better interpretability through attention visualization than traditional CNNs?

Key findings

• The pre-trained Vision Transformer (ViT-v2) achieved a test accuracy of 94% when fine-tuned on the NIH Chest X-ray dataset, significantly outperforming both CNN and ResNet models.
• The ViT model pre-trained on ImageNet-21k demonstrated superior performance compared to the ViT trained from scratch and all baseline models, highlighting the importance of large-scale pre-training.
• AUC-ROC analysis revealed that ViT models performed poorly on the 'No Finding' class (AUC = 0.50), indicating difficulty distinguishing normal from abnormal cases, while showing moderate performance on diseases like Atelectasis (AUC = 0.51) and Pleural Thickening (AUC = 0.53).
• The ViT model faced challenges with rare diseases such as Cardiomegaly (AUC = 0.48) and Hernia (AUC = 0.43), primarily due to data imbalance in the dataset.
• Attention maps generated from the ViT model highlighted clinically relevant regions in X-ray images, such as the heart and lung fields, confirming the model’s ability to focus on diagnostically significant areas.
• Despite strong performance on average, the ViT model underperformed on the 'No Finding' class, suggesting a need for better data sampling or class-balancing techniques in future work.