[Paper Review] EdgeSAM: Prompt-In-the-Loop Distillation for SAM

Chong Zhou, Xiangtai Li+2 more|arXiv (Cornell University)|Dec 11, 2023

Advanced Memory and Neural ComputingEngineering61 references11 citations

TL;DR

EdgeSAM distills SAM into a CNN-based encoder with prompt-in-the-loop distillation and a lightweight granularity priors module, achieving real-time on-device interactive segmentation with competitive accuracy.

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.

Motivation & Objective

• Motivate on-device interactive segmentation by enabling SAM-like capabilities on edge devices without large computational costs.
• Investigate encoder distillation strategies and demonstrate limitations of task-agnostic approaches.
• Propose prompt-aware knowledge distillation with a dynamic prompt-in-the-loop mechanism.
• Incorporate a lightweight granularity priors module to handle dataset-specific prompt granularity.
• Evaluate zero-shot transfer and real-time performance across standard segmentation benchmarks.

Proposed method

• Distill SAM's ViT-based image encoder into a purely CNN-based architecture for edge deployment.
• Apply encoder-only knowledge distillation and show its limitations without task-specific guidance.
• Introduce a dynamic prompt-in-the-loop distillation that iteratively samples prompts in regions where the student errs to guide the mask decoder.
• Retain SAM's mask decoder and train it with distillation targets focused on the decoder outputs.
• Add a lightweight granularity priors module (RPN with FPN) to incorporate dataset-specific granularity priors during inference.
• Conduct staged training: encoder-only KD, prompt-in-the-loop KD, and optionally grounding with a lightweight RPN.

Experimental results

Research questions

• RQ1 Can SAM be effectively deployed on edge devices by distilling its components into a CNN-based encoder without substantial loss in segmentation quality?
• RQ2 Does task-aware prompt guidance during distillation outperform encoder-only distillation for dense prediction tasks like SAM?
• RQ3 How do prompt types and granularity priors influence zero-shot transferability and on-device performance?
• RQ4 What backbone/backbone adaptations best balance throughput and accuracy on edge hardware?
• RQ5 Can a lightweight RPN leveraging granularity priors improve performance for ambiguous prompts?

Key findings

• EdgeSAM runs over 30 FPS on an iPhone 14, the first SAM variant with real-time edge-device performance.
• Encoder-only distillation alone yields a performance gap versus the original SAM, motivating prompt-aware distillation.
• Prompt-in-the-loop distillation improves mask accuracy over encoder-only KD, especially with additional refinement prompts.
• A purely CNN-based backbone (RepViT-M1 with FPN) provides the best throughput-accuracy balance for edge deployment.
• EdgeSAM achieves up to 40x speed-up over SAM and up to 14x speed-up over MobileSAM on edge devices, with competitive or better mIoU/IoU gains on COCO and LVIS.
• With GT boxes as prompts, EdgeSAM closes the gap to SAM and even surpasses MobileSAM in several settings; a granularity-priors RPN further boosts center-point performance on COCO.