MedMambaLite: Hardware-Aware Mamba for Medical Image Classification
Romina Aalishah, Mozhgan Navardi, Tinoosh Mohsenin

TL;DR
MedMambaLite is a hardware-aware, lightweight medical image classification model that uses knowledge distillation to achieve high accuracy and energy efficiency on edge devices, enabling real-time medical diagnostics.
Contribution
This paper introduces MedMambaLite, a novel, optimized Mamba-based model for medical imaging that is smaller, faster, and more energy-efficient through architecture modifications and knowledge distillation.
Findings
Achieves 94.5% accuracy on MedMNIST datasets
Reduces model parameters by 22.8x compared to MedMamba
Improves energy efficiency by 63% on NVIDIA Jetson Orin Nano
Abstract
AI-powered medical devices have driven the need for real-time, on-device inference such as biomedical image classification. Deployment of deep learning models at the edge is now used for applications such as anomaly detection and classification in medical images. However, achieving this level of performance on edge devices remains challenging due to limitations in model size and computational capacity. To address this, we present MedMambaLite, a hardware-aware Mamba-based model optimized through knowledge distillation for medical image classification. We start with a powerful MedMamba model, integrating a Mamba structure for efficient feature extraction in medical imaging. We make the model lighter and faster in training and inference by modifying and reducing the redundancies in the architecture. We then distill its knowledge into a smaller student model by reducing the embedding…
Peer Reviews
No public reviews on file for this paper yet. If you reviewed it on a platform where reviews are public (OpenReview, ICLR, NeurIPS, ICML), you can paste yours below so the community can read it here.
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
