轻量级Mamba模型实现医疗图像分类，能效比提升63%。
MedMambaLite: Hardware-Aware Mamba for Medical Image Classification
- 基于知识蒸馏优化Mamba结构，压缩模型冗余并降低嵌入维度
- 在10个MedMNIST数据集上达94.5%准确率，参数量减少22.8倍
- 部署于Jetson Orin Nano时每推理能耗仅35.6 GOPS/J，适合边缘医疗设备
基于人工智能的医疗设备推动了实时、本地化推理的需求，如生物医学图像分类。将深度学习模型部署至边缘设备用于异常检测和图像分类已成为现实应用。然而，受限于模型规模与计算能力，实现在边缘设备上的高性能仍具挑战。为此，我们提出MedMambaLite，一种面向硬件优化的Mamba基医疗图像分类模型，通过知识蒸馏实现高效压缩。以强大但复杂的MedMamba为基础，引入Mamba结构以高效提取医学影像特征；通过结构调整与冗余削减，显著减小模型体积并加快训练与推理速度；再通过降低嵌入维度将知识蒸馏至更小的学生模型。该优化模型在10个MedMNIST数据集上实现94.5%的整体准确率，相比MedMamba参数量减少22.8倍。部署于NVIDIA Jetson Orin Nano时，每推理能耗仅为35.6 GOPS/J，较MedMamba能效提升63%，充分验证其在边缘医疗场景中的适用性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。