TinySAM-Med3D: A Lightweight Segment Anything Model for Volumetric Medical Imaging with Mixture of Experts

Tianyuan Song,Guixia Kang,Yiqing Shen
DOI: https://doi.org/10.1007/978-3-031-66535-6_15
2024-01-01
Abstract:Segment anything models (SAMs) demonstrate exceptional zero-shot segmentation capabilities on natural images when provided appropriate prompts. However, directly applying SAMs to medical images presents challenges due to the complexity and diversity of medical data compared to natural images. Some research aims to integrate medical knowledge into SAM by fine-tuning the model. However, the encoders are based on the ViT architecture, which incurs high computational costs when applied directly to 3D medical data, limiting real-time performance on resource-constrained devices. To address these limitations, we introduce TinySAM-Med3D, an efficient SAM tailored for 3D medical image segmentation. TinySAM-Med3D builds on SAM-Med3D by distilling the encoder to a lightweight TinyViT and substituting the multilayer perceptron with a Mixture of Experts (MoEs) to preserve performance while significantly reducing computational and memory costs. This enables real-time segmentation on resource-constrained devices without sacrificing performance. Evaluations on the abdominal CT dataset of Total-Segmentator reveal that TinySAM-Med3D attains a 0.8440 Dice score while using only 33.87% of SAM-Med3D parameters. It also accelerates inference by 3.36x over SAM-Med3D. Thus, TinySAM-Med3D facilitates deploying SAMs for fast 3D medical image segmentation. Our code and models are available at https://github.com/songty21110133/TinySAMMed3D.
What problem does this paper attempt to address?