Ultra-Lightweight 3D Medical Image Segmentation via Knowledge Distillation in Label-Scarce Domains
Authors
Abstract
Mamba is a neural network architecture that reduces computational complexity, and has been shown to be effective for 3D medical image segmentation. However, existing models still remain computationally expensive. We propose compressing Mamba by distillation while achieving competitive performance. Specifically, we demonstrate compression of a 139M-parameter SegMamba-V2 teacher into a lightweight TinyUNet3D student with about 0.96M parameters, achieving about 145$\times$ compression with a performance drop of just over one Dice point. We further investigate the utility of knowledge distillation in a range of data-labeling scenarios under identical experimental settings, and show that its advantage is most prominent when labeled data is scarce. To demonstrate the generalizability of our approach, experiments are conducted on three datasets: BraTS 2023, MSD Spleen and MSD Hippocampus. The code for TinyUNet3D is publicly available at: https://github.com/2802990u/tinyunet3d-dist-kd.