MedMambaLite: Hardware-Aware Mamba for Medical Image Classification

AI-powered medical devices have driven the need for real-time, on-device inference in healthcare domains 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, which integrates a Mamba structure for efficient feature extraction in medical imaging. We make the model lighter and faster in training and inference by reducing the redundancies in the architecture and integrating modifications. We then distill its knowledge into a smaller student model by reducing the embedding dimensions. The optimized model achieves $94.5 \%$ overall accuracy on 10 MedMNIST datasets. It also reduces parameters $22.8 \times$ compared to MedMamba. Deployment of MedMambaLite on an NVIDIA Jetson Orin Nano achieves 35.6 mJ energy per inference. This outperforms MedMamba by $63 \%$ improvement in energy per inference, demonstrating its suitability for edge medical applications.