Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

CoMViT: An Efficient Vision Backbone for Supervised Classification in Medical Imaging

Aon Safdar, Mohamed Saadeldin

TL;DR

Vision Transformers often struggle with small medical datasets and limited compute. CoMViT introduces a compact ViT backbone with a convolutional tokenizer, diagonal masking, learnable temperature, and learnable sequence pooling, achieving strong accuracy with substantially fewer parameters. Evaluated on 12 MedMNIST 2D datasets, it matches or surpasses larger CNNs and ViTs while using only a fraction of the parameters and FLOPs, and Grad-CAM analyses show interpretable localization. This work demonstrates that principled architectural redesign emphasizing locality and efficient pooling can deliver robust, deployable ViTs for low-resource medical imaging tasks.

Abstract

Vision Transformers (ViTs) have demonstrated strong potential in medical imaging; however, their high computational demands and tendency to overfit on small datasets limit their applicability in real-world clinical scenarios. In this paper, we present CoMViT, a compact and generalizable Vision Transformer architecture optimized for resource-constrained medical image analysis. CoMViT integrates a convolutional tokenizer, diagonal masking, dynamic temperature scaling, and pooling-based sequence aggregation to improve performance and generalization. Through systematic architectural optimization, CoMViT achieves robust performance across twelve MedMNIST datasets while maintaining a lightweight design with only ~4.5M parameters. It matches or outperforms deeper CNN and ViT variants, offering up to 5-20x parameter reduction without sacrificing accuracy. Qualitative Grad-CAM analyses show that CoMViT consistently attends to clinically relevant regions despite its compact size. These results highlight the potential of principled ViT redesign for developing efficient and interpretable models in low-resource medical imaging settings.

CoMViT: An Efficient Vision Backbone for Supervised Classification in Medical Imaging

TL;DR

Vision Transformers often struggle with small medical datasets and limited compute. CoMViT introduces a compact ViT backbone with a convolutional tokenizer, diagonal masking, learnable temperature, and learnable sequence pooling, achieving strong accuracy with substantially fewer parameters. Evaluated on 12 MedMNIST 2D datasets, it matches or surpasses larger CNNs and ViTs while using only a fraction of the parameters and FLOPs, and Grad-CAM analyses show interpretable localization. This work demonstrates that principled architectural redesign emphasizing locality and efficient pooling can deliver robust, deployable ViTs for low-resource medical imaging tasks.

Abstract

Vision Transformers (ViTs) have demonstrated strong potential in medical imaging; however, their high computational demands and tendency to overfit on small datasets limit their applicability in real-world clinical scenarios. In this paper, we present CoMViT, a compact and generalizable Vision Transformer architecture optimized for resource-constrained medical image analysis. CoMViT integrates a convolutional tokenizer, diagonal masking, dynamic temperature scaling, and pooling-based sequence aggregation to improve performance and generalization. Through systematic architectural optimization, CoMViT achieves robust performance across twelve MedMNIST datasets while maintaining a lightweight design with only ~4.5M parameters. It matches or outperforms deeper CNN and ViT variants, offering up to 5-20x parameter reduction without sacrificing accuracy. Qualitative Grad-CAM analyses show that CoMViT consistently attends to clinically relevant regions despite its compact size. These results highlight the potential of principled ViT redesign for developing efficient and interpretable models in low-resource medical imaging settings.

Paper Structure

This paper contains 13 sections, 2 equations, 4 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Method
  4. Experiments and Results
  5. Datasets
  6. Model Configuration
  7. Experiment Protocol
  8. Hyperparameters.
  9. Benchmarks and Metrics.
  10. Results and Discussion
  11. Conclusion
  12. Acknowledgments.
  13. Disclosure of Interests.

Figures (4)

  • Figure 1: CoMViT Architecture Overview
  • Figure 2: Overview of the MedMNIST Dataset Family. Modality, Classification Task, and split (Train / Val / Test) are mentioned. Best viewed in color and zoomed in.
  • Figure 3: Grad-CAM visualizations across six MedMNIST datasets highlighting the regions of interest that contribute most to the model's predictions. CoMViT accurately localizes relevant regions across diverse modalities.
  • Figure 4: Model Size vs. Accuracy on TissueMNIST. CoMViT achieves strong accuracy with minimal parameters, illustrating excellent efficiency.