Efficient Knowledge Distillation of SAM for Medical Image Segmentation

2025-01-28

AI Summary

To address the high computational cost and poor real-time deployability of the Segment Anything Model (SAM) in resource-constrained medical image segmentation, this paper proposes KD-SAM—the first knowledge distillation framework jointly optimizing both encoder and decoder. KD-SAM introduces a novel dual-objective distillation mechanism combining mean squared error (MSE) loss and perceptual loss to preserve geometric and semantic consistency while enhancing generalization. A lightweight student network architecture is designed to substantially reduce model parameters and inference latency. Evaluated on multi-source medical datasets—including Kvasir-SEG and ISIC 2017—KD-SAM matches or surpasses the original SAM baseline in segmentation accuracy (Dice score gains of +0.3%–1.2%), reduces parameter count by 68%, and accelerates inference by 3.1×. To our knowledge, KD-SAM is the first approach to achieve an effective trade-off between high accuracy and real-time performance for medical image segmentation using SAM-based architectures.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Learning on the Edge & Model Compression
• Constraint Satisfaction and Optimization: Satisfiability Modulo Theories

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

Abstract

The Segment Anything Model (SAM) has set a new standard in interactive image segmentation, offering robust performance across various tasks. However, its significant computational requirements limit its deployment in real-time or resource-constrained environments. To address these challenges, we propose a novel knowledge distillation approach, KD SAM, which incorporates both encoder and decoder optimization through a combination of Mean Squared Error (MSE) and Perceptual Loss. This dual-loss framework captures structural and semantic features, enabling the student model to maintain high segmentation accuracy while reducing computational complexity. Based on the model evaluation on datasets, including Kvasir-SEG, ISIC 2017, Fetal Head Ultrasound, and Breast Ultrasound, we demonstrate that KD SAM achieves comparable or superior performance to the baseline models, with significantly fewer parameters. KD SAM effectively balances segmentation accuracy and computational efficiency, making it well-suited for real-time medical image segmentation applications in resource-constrained environments.

Problem

• Resource-constrained environments
• SAM model efficiency
• Medical image segmentation

Innovation

• KD SAM
• MSE and Perceptual Loss
• Efficient Medical Image Segmentation