MedMambaLite: Hardware-Aware Mamba for Medical Image Classification

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 dimensions. The optimized model achieves 94.5% overall accuracy on 10 MedMNIST datasets. It also reduces parameters 22.8x compared to MedMamba. Deployment on an NVIDIA Jetson Orin Nano achieves 35.6 GOPS/J energy per inference. This outperforms MedMamba by 63% improvement in energy per inference.