SAKD-SAM: efficient breast tumor MRI segmentation via self-attention knowledge distillation and 3D sparse flash attention

Abstract

Segmenting breast tumors in magnetic resonance imaging (MRI) is a critical step in computer-aided diagnosis. Recently, the Segment Anything Model (SAM) has demonstrated remarkable potential in medical image segmentation. However, the immense parameter size and heavy computational burden of SAM's Transformer-based architecture severely restrict its clinical deployment, especially for 3D volumetric data. To address these challenges, we propose a lightweight breast tumor segmentation framework based on a novel Self-Attention Knowledge Distillation (SAKD) strategy. By transferring the rich semantic representations from a cumbersome teacher model to a lightweight student model equipped with a TinyViT encoder, we significantly reduce the model size. Furthermore, we design an innovative 3D Sparse Flash Attention mechanism coupled with an Adapter module to accelerate inference speed and preserve task-specific medical prior knowledge. Extensive experiments on multi-center datasets demonstrate that our method achieves a competitive Dice Similarity Coefficient (DSC) of 78.9% with 10 prompt points. Notably, the encoder's inference time is reduced to 4.4 ms and the memory requirement is decreased to 1.04 GB, achieving a 9.2× acceleration compared to SAM-Med3D. Our approach provides a highly robust and efficient solution for deploying large foundation models in resource-constrained clinical workstations.