MedMambaLite: Hardware-Aware Mamba for Medical Image Classification

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Aalishah, Romina, Navardi, Mozhgan, Mohsenin, Tinoosh
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866915433361178624
author Aalishah, Romina
Navardi, Mozhgan
Mohsenin, Tinoosh
author_facet Aalishah, Romina
Navardi, Mozhgan
Mohsenin, Tinoosh
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2508_05049
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle MedMambaLite: Hardware-Aware Mamba for Medical Image Classification
Aalishah, Romina
Navardi, Mozhgan
Mohsenin, Tinoosh
Image and Video Processing
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.
title MedMambaLite: Hardware-Aware Mamba for Medical Image Classification
topic Image and Video Processing
url https://arxiv.org/abs/2508.05049