EdgeSAM: Prompt-In-the-Loop Distillation for SAM

TL;DR

EdgeSAM tackles the challenge of running SAM on edge devices by distilling SAM's ViT-based image encoder into a CNN backbone and incorporating a prompt-aware, prompt-in-the-loop distillation regime along with a lightweight granularity priors module. The approach preserves the SAM decoder and enables interactive segmentation with box and point prompts while achieving substantial on-device speedups (up to 37x on desktop and over 30 FPS on iPhone 14) and competitive or superior mIoUs on COCO/LVIS compared with MobileSAM and EfficientSAM. Key contributions include encoder-only KD analysis showing its insufficiency, the novel prompt-in-the-loop KD strategy, the granularity priors RPN, and a data-efficient training pipeline using 1% of SA-1B. The work demonstrates practical, real-time, on-device interactive segmentation with strong zero-shot transfer, enabling new on-device applications and workflows for edge AI.

Abstract

This paper presents EdgeSAM, an accelerated variant of the Segment Anything Model (SAM), optimized for efficient execution on edge devices with minimal compromise in performance. Our approach involves distilling the original ViT-based SAM image encoder into a purely CNN-based architecture, better suited for edge devices. We carefully benchmark various distillation strategies and demonstrate that task-agnostic encoder distillation fails to capture the full knowledge embodied in SAM. To overcome this bottleneck, we include both the prompt encoder and mask decoder in the distillation process, with box and point prompts in the loop, so that the distilled model can accurately capture the intricate dynamics between user input and mask generation. To mitigate dataset bias issues stemming from point prompt distillation, we incorporate a lightweight module within the encoder. As a result, EdgeSAM achieves a 37-fold speed increase compared to the original SAM, and it also outperforms MobileSAM/EfficientSAM, being over 7 times as fast when deployed on edge devices while enhancing the mIoUs on COCO and LVIS by 2.3/1.5 and 3.1/1.6, respectively. It is also the first SAM variant that can run at over 30 FPS on an iPhone 14. Code and demo are available at https://www.mmlab-ntu.com/project/edgesam.

Figures (5)

• Figure 1: Figure (a) shows the overview of EdgeSAM. We first apply encoder-only knowledge distillation between the output feature of the image encoder of SAM and EdgeSAM. Then, we adopt prompt-in-the-loop knowledge distillation, which interactively samples point prompts from the wrongly segmented areas. Note that the initial prompt can also be in a box. In (b), we show its throughput and the mIoU performance on the COCO and LVIS datasets with box prompts compared with SAM, MobileSAM, and EfficientSAM.
• Figure 2: Figures (a-c) show the behaviors of the SAM with different prompts. When prompted with ambiguous prompts, e.g., a point, SAM yields suboptimal results. More informative prompts solve the problem. In (d), we demonstrate the sampling process of prompt-in-the-loop KD. In particular, during distillation, we randomly sample new positive/negative point prompts from the false negative/false positive areas in a loop, so that the student dynamically focuses on those regions.
• Figure 3: Qualitative results of EdgeSAM with point and box prompts. The green and red stars indicate the positive and negative points respectively.
• Figure 4: We compare the visual results of applying prompt-in-the-loop knowledge distillation or not to demonstrate its effectiveness. The baseline follows the same network architecture of EdgeSAM but adopts MobileSAM's encoder-only distillation method.
• Figure 5: Qualitative results of EdgeSAM and SAM under point, box, and everything mode.