Source-linked AI summary
MedViT: A Robust Vision Transformer for Generalized Medical Image Classification
Omid Nejati Manzari, Hamid Ahmadabadi, Hossein Kashiani, Shahriar B. Shokouhi, Ahmad Ayatollahi
TL;DR
Medical diagnosis systems need reliable classification despite adversarial attacks, while CNNs and Transformers each present limitations in locality, global modeling, data requirements, or computational complexity. MedViT combines efficient convolutional and Transformer components with feature-space augmentation, achieving state-of-the-art accuracy and robustness across standardized MedMNIST-2D datasets.
Problem
Medical image classifiers must remain reliable against adversarial attacks, while CNN locality and Transformer complexity limit robust, generalized medical-image analysis.
Method
MedViT is a hybrid CNN-Transformer architecture using efficient convolutional attention, local modeling components, and Patch Momentum Changer augmentation to support global features and smoother decision boundaries.
Results
MedViT achieves state-of-the-art accuracy and robustness on standardized large-scale 2D biomedical datasets, outperforming prior methods and improving on AutoML baselines in reported AUC and ACC comparisons.
Takeaways & Limitations
The results support MedViT as an efficient generalized medical-image classification design combining local representations, global features, robustness, and cross-domain generalization.
Abstract
from arXiv · showhide
Convolutional Neural Networks (CNNs) have advanced existing medical systems for automatic disease diagnosis. However, there are still concerns about the reliability of deep medical diagnosis systems against the potential threats of adversarial attacks since inaccurate diagnosis could lead to disastrous consequences in the safety realm. In this study, we propose a highly robust yet efficient CNN-Transformer hybrid model which is equipped with the locality of CNNs as well as the global connectivity of vision Transformers. To mitigate the high quadratic complexity of the self-attention mechanism while jointly attending to information in various representation subspaces, we construct our attention mechanism by means of an efficient convolution operation. Moreover, to alleviate the fragility of our Transformer model against adversarial attacks, we attempt to learn smoother decision boundaries. To this end, we augment the shape information of an image in the high-level feature space by permuting the feature mean and variance within mini-batches. With less computational complexity, our proposed hybrid model demonstrates its high robustness and generalization ability compared to the state-of-the-art studies on a large-scale collection of standardized MedMNIST-2D datasets.
1. Introduction
Medical image classification benefits from CNNs but faces challenges in modeling long-range dependencies, handling distribution variation, and resisting adversarial attacks. MedViT addresses these issues with an efficient hybrid architecture and feature-space augmentation designed to support generalization and robustness.
- CNNs provide accurate medical-image representations, but their locality bias limits long-range dependency modeling when anatomical texture, shape, and size vary.
- Transformer self-attention models long-range dependencies but has quadratic complexity and requires substantial training data, limiting practical medical-image use.
- Scanner, imaging-protocol, and epidemiological differences can substantially shift data and label distributions across medical sites.
- MedViT combines hierarchical patch embedding with convolution and Transformer blocks to model local and long-term information efficiently across CT, X-ray, ultrasound, and OCT domains.Its multi-head convolutional attention reduces the computational burden of conventional attention and improves inference speed.
- The local feed-forward network adds depth-wise convolution to reshaped 2D feature maps, encoding dependencies between nearby pixels throughout the Transformer stages.
- Feature mean-and-variance interpolation encourages reliance on global structure features and expands representation-space regions to learn smoother decision boundaries.The proposed augmentation is intended to improve both adversarial robustness and generalization performance.
2. Related Works
Prior work spans convolutional, Transformer, hybrid, and robustness-oriented approaches, each addressing different efficiency, locality, global-context, or adversarial-robustness concerns. MedViT extends this line by designing a robust hybrid architecture rather than merely comparing CNNs and ViTs.
- CNN research emphasizes discriminative feature extraction, while Vision Transformers use image patches and self-attention for visual recognition.
- Efficient vision models explore pyramidal designs, training strategies, and efficient self-attention to reduce computational demands.
- Hybrid models combine convolutional locality with Transformer-style global modeling through designs such as BoTNet, CMT, CvT, LeViT, MobileViT, and Mobile-Former.
- Adversarial-robustness research for CNNs uses architectural changes, augmentation, training strategies, pruning, quantization, activation functions, and pooling improvements.
- Earlier studies report greater adversarial robustness for Transformers than CNNs, but existing CNN robustness methods lacked evidence of improving ViT robustness.
- MedViT goes beyond direct CNN–ViT robustness comparisons by introducing a robust hybrid architecture family and an augmentation technique targeting Transformer fragility.
3. Method
MedViT combines hierarchical convolutional and Transformer stages to capture local dependencies, global context, and multi-frequency information efficiently. Its method also introduces feature-level augmentation to improve robustness and generalization.
- MedViT architecture: MedViT uses a hierarchical pyramid combining patch embedding, Transformer blocks, and stacked convolutional blocks across stages.Spatial resolution decreases by a total 32× ratio through [4×, 2×, 2×, and 2×] reductions, while channels double after convolution blocks.
- Locally Feed-Forward Network: The Local Feed-Forward Network reshapes token sequences into 2D feature maps and applies convolution to encode dependencies between nearby pixels.The resulting feature map is converted back into a token sequence for subsequent Transformer processing.
- Local Transformer Block: The Local Transformer Block combines Efficient Self-Attention and multi-head convolutional attention to capture global and local information with reduced computation.Efficient Self-Attention downsamples spatial dimensions before attention, while convolutional attention jointly attends across multiple representation subspaces.
- Transformer Augmentation Block: Patch Momentum Changer augments training by combining normalized features from one sample with feature moments from another at token-level feature representations.This asymmetric feature-space composition is intended to improve robustness and generalization when predicting medical image datasets.
4. Experiments
Experiments use the standardized MedMNIST-2D collection, spanning diverse modalities, datasets, and task types, with consistent training settings and ACC/AUC evaluation.
- Datasets: MedMNIST-2D includes 12 pre-processed datasets covering CT, X-ray, ultrasound, and OCT images across multiple classification task types.The collection ranges from at least 100 to more than 100,000 samples per dataset.
- Datasets: The benchmark includes datasets with binary, multi-class, multi-label, ordinal, and regression tasks, including PneumoniaMNIST and RetinaMNIST.PneumoniaMNIST has two classes, whereas RetinaMNIST uses five-level diabetic-retinopathy severity grading.
- Datasets: ChestMNIST contains 112,120 frontal X-ray images from 32,717 patients and represents 14 disease classes.
- Experimental setup: Experiments follow MedMNISTv2 training settings, train MedViT variants for 100 epochs, and evaluate performance using ACC and AUC.Average ACC and AUC are also reported across all 12 MedMNIST-2D datasets.
- Evaluation metrics: ACC evaluates thresholded discrete labels, whereas AUC evaluates continuous prediction scores and is less sensitive to class discrepancy.
5. Evaluation Results
MedViT performs strongly across MedMNIST-2D and comparable model-size benchmarks while retaining competitive complexity, localization, and adversarial robustness. Ablations further identify the Efficient Convolution Block and early-stage PMC placement as effective design choices.
- Results on Each Dataset: MedViT outperforms previous state-of-the-art methods across MedMNIST-2D, with gains over AutoML baselines on RetinaMNIST and TissueMNIST.MedViT-S improves RetinaMNIST by 2.3% AUC and 3.0% ACC, and TissueMNIST by 1.1% AUC and 2.8% ACC.
- Comparison with State-of-the-art Models: MedViT-S reaches 73.1% Top-1 accuracy, exceeding ResNet-50 by 5.1%, EfficientNet-B4 by 2.6%, and ConvNext-T by 0.5%.
- Average performance: MedViT-T, MedViT-S, and MedViT-L obtain average AUCs of 93.6%, 94.2%, and 93.5%, with average ACCs of 84%, 85.1%, and 84.2%, respectively.MedViT-S outperforms baseline ResNets and AutoML methods on both average metrics.
- Efficiency: MedViT achieves favorable accuracy-complexity trade-offs, with complexity on par with baseline ResNets and lower complexity than some stronger-performing comparators.MedViT-L outperforms ConvNext-B by 0.8% while using approximately half as many parameters.
- Qualitative analysis: Grad-CAM shows MedViT attends more precisely to relevant regions, less to background, and better handles scale variation across several datasets.The model localizes retinal lesions, focal bacterial infection, and multifocal pneumonia lesions in the reported examples.
- Adversarial robustness: MedViT-T* improves TissueMNIST robustness over ResNet-18 by 38.4% under FGSM and 30.2% under PGD.The reported attack settings use FGSM and PGD with adversarial noise magnitude ε = 8/255.
- Ablation Study: The Efficient Convolution Block exceeds the ConvNext block by 0.6% in clean accuracy and 4.4% in robust accuracy with lower model complexity.This ablation compares alternative blocks while keeping the other architecture components unchanged.
- Ablation Study: PMC performs best after the first stage of the four-stage MedViT-T, leading the paper to adopt first-stage augmentation.The ablation experiments are conducted on MedViT-T using TissueMNIST.
6. Conclusion
MedViT is presented as a hybrid CNN-Transformer architecture for medical image classification, combining local and global representations with a patch moment changer augmentation. The reported experiments show state-of-the-art accuracy and robustness on large-scale 2D biomedical datasets, while Table 8 evaluates PMC placement under FGSM attack.
- MedViT combines local and global representations through a hybrid CNN-Transformer architecture for medical image classification.
- Table 8 reports clean accuracy and adversarial robustness under FGSM attack on TissueMNIST across different PMC stages.
- Patch moment changer augmentation adds diversity and affinity to the training data.
- MedViT achieves state-of-the-art accuracy and robustness on a standard large-scale collection of 2D biomedical datasets.
CRediT authorship contribution statement
The authors’ contributions span conceptualization, methodology, software, validation, writing, supervision, reviewing, editing, resources, and final layout modification.
- Omid Nejati Manzari contributed conceptualization, software, original draft preparation, validation, and resources.
- Hamid Ahmadabadi contributed methodology, writing, reviewing, and editing.
- Hossein Kashiani contributed writing, reviewing, editing, and final layout modification.
- Shahriar B. Shokouhi and Ahmad Ayatollahi contributed supervision, review, and editing.