[論文レビュー] EdgeSAM: Prompt-In-the-Loop Distillation for SAM

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

Advanced Memory and Neural Computing参考文献 61被引用数 11

ひとこと要約

EdgeSAMはSAMをCNNベースのエンコーダへ蒸留し、ループ内プロンプト蒸留と軽量な粒度事前情報モジュールを組み合わせ、競争力のある精度を保ちながらリアルタイムのデバイス上対話的セグメンテーションを実現する。

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.

研究の動機と目的

• 大きな計算コストをかけずに、エッジデバイス上でSAM風の機能を有効にして、オンデバイスの対話型セグメンテーションを動機づける。
• エンコーダ蒸留戦略を調査し、タスク非特異的アプローチの限界を示す。
• 動的なプロンプト・イン・ザ・ループ機構を備えた、プロンプト認識型知識蒸留を提案する。
• データセット固有の粒度事前情報を扱う軽量な粒度 priors モジュールを組み込む。
• 標準的なセグメンテーションベンチマークにおけるゼロショット転送とリアルタイム性能を評価する。

提案手法

• SAMの ViT ベースの画像エンコーダを、エッジ展開向けの純粋なCNNベースのアーキテクチャへ蒸留する。
• エンコーダのみの知識蒸留を適用し、タスク固有の指針がない場合の限界を示す。
• 学生が誤る領域でプロンプトを反復的にサンプリングしてマスクデコーダを導く、動的なプロンプト・イン・ザ・ループ蒸留を導入する。
• SAMのマスクデコーダを保持し、デコーダ出力に焦点を当てた蒸留ターゲットで訓練する。
• 推論時にデータセット固有の粒度事前情報を組み込む軽量な粒度 priors モジュール（FPNを備えたRPN）を追加する。
• 段階的な訓練を実施: エンコーダのみのKD、プロンプト・イン・ザ・ループKD、オプションとして軽量RPNによるグラウンディング。

実験結果

リサーチクエスチョン

• RQ1 SAMの構成要素をCNNベースのエンコーダに蒸留して、セグメンテーション品質の大幅な低下を招くことなく、エッジデバイス上で効果的に展開できるだろうか？
• RQ2 蒸留時のタスク認識型プロンプト指導は、SAMのような密な予測タスクにおいてエンコーダのみの蒸留より優れているか？
• RQ3 プロンプトタイプと粒度 priors はゼロショット転送性とデバイス上の性能にどう影響するか？
• RQ4 エッジハードウェア上でスループットと精度のバランスを最も良く取るバックボーン/バックボーン適応は何か？
• RQ5 粒度 priors を活用する軽量なRPNは曖昧なプロンプトの性能を改善できるか？

主な発見

Model	FPS (2080 Ti)	FPS (iPhone 14)	MParam.	GFLOPs
SAM	4.3	N/A	641.1	2734.8
MobileSAM	111.7	5.2	9.8	38.2
EdgeSAM	185.9	72.3	9.6	22.1
EdgeSAM*	135.8	57.8	9.8	22.2

• EdgeSAMはiPhone 14で30 FPS超を実現し、リアルタイムのエッジデバイス性能を備えた最初のSAMバリアントである。
• エンコーダのみの蒸留だけでは元のSAMに対する性能ギャップが生じるため、プロンプト認識型蒸留を動機づける。
• プロンプト・イン・ザ・ループ蒸留は、追加のリファインメントプロンプトとともに、エンコーダのみKDよりマスク精度を改善する。
• 純粋なCNNベースのバックボーン（RepViT-M1 with FPN）は、エッジ展開において最高のスループットと精度のバランスを提供する。
• EdgeSAMはエッジデバイス上でSAMより最大40x、MobileSAMより最大14xの速度向上を達成し、COCOおよびLVISで競争力のあるまたはそれを上回るmIoU/IoUの向上を示す。
• GTボックスをプロンプトとして用いると、EdgeSAMはSAMとの差を縮め、いくつかの設定でMobileSAMを上回る。粒度 priors RPN はCOCOで中心点性能をさらに向上させる。