[Paper Review] MedViT: A Robust Vision Transformer for Generalized Medical Image Classification

Omid Nejati Manzari, Hamid Ahmadabadi+3 more|arXiv (Cornell University)|Feb 19, 2023

COVID-19 diagnosis using AI9 citations

TL;DR

MedViT proposes a CNN-Transformer hybrid architecture with efficient convolutional blocks and a multi-frequency LTB, aiming for robust, generalized medical image classification against adversarial attacks on MedMNIST-2D datasets.

ABSTRACT

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.

Motivation & Objective

• Motivate robust medical image classification amid adversarial threats and domain shifts in medical imaging.
• Develop a hybrid CNN-Transformer architecture that combines local and global feature modeling with low computation.
• Enhance robustness by encouraging reliance on global structure features and introducing feature-level augmentation.
• Propose efficient blocks and multi-frequency fusion to balance accuracy and efficiency across medical imaging modalities.

Proposed method

• Introduce Efficient Convolution Block (ECB) combining Locally Feed Forward Network (LFFN) and Multi-Head Convolutional Attention (MHCA) for local feature learning.
• Develop Local Transformer Block (LTB) to fuse multi-frequency information using Efficient Self Attention (ESA) and MHCA for global-local fusion.
• Use a Transformer Augmentation Block (TAB) with Patch Momentum Changer (PMC) to augment training data at feature level and promote smoother decision boundaries.
• Adopt a patch embedding as a hierarchical pyramid with shared stage-wise convolution and transformer blocks to manage computational complexity.
• Optimize norm/activation with Batch Normalization and ReLU for deployment efficiency in medical contexts.
• Provide multiple MedViT variants (T/S/L) with detailed stage configurations and channel dimensions.

Experimental results

Research questions

• RQ1 Can MedViT achieve robust generalization across diverse medical imaging modalities and domain shifts?
• RQ2 Does the combination of ECB and LTB enhance modeling of local and global information efficiently?
• RQ3 Does the PMC augmentation improve robustness to adversarial perturbations and smooth decision boundaries?
• RQ4 How does MedViT perform relative to baselines on standardized MedMNIST-2D datasets in terms of accuracy and efficiency?

Key findings

• MedViT demonstrates robustness and generalization advantages over state-of-the-art studies on MedMNIST-2D datasets.
• The architecture achieves high efficiency by replacing standard self-attention with convolution-based attention and local feed-forward networks.
• PMC feature-level augmentation broadens interaction between tokens and improves robustness and generalization.
• LTB effectively mixes multi-frequency information, boosting performance with a lightweight design.
• The model variants (MedViT-T/S/L) provide scalable configurations with detailed stage-wise blocks and channel counts.