[論文レビュー] SAM-Lightening: A Lightweight Segment Anything Model with Dilated Flash Attention to Achieve 30 times Acceleration

Yanfei Songa, Bangzheng Pua+5 more|arXiv (Cornell University)|Mar 14, 2024

Data Visualization and Analytics被引用数 5

ひとこと要約

SAM-Lightening は SAM のエンコーダを Dilated Flash Attention と動的層別蒸留で再設計し、推論を約30×高速化し、メモリを大幅に削減しつつ、セグメンテーション品質を維持します。

ABSTRACT

Segment Anything Model (SAM) has garnered significant attention in segmentation tasks due to their zero-shot generalization ability. However, a broader application of SAMs to real-world practice has been restricted by their low inference speed and high computational memory demands, which mainly stem from the attention mechanism. Existing work concentrated on optimizing the encoder, yet has not adequately addressed the inefficiency of the attention mechanism itself, even when distilled to a smaller model, which thus leaves space for further improvement. In response, we introduce SAM-Lightening, a variant of SAM, that features a re-engineered attention mechanism, termed Dilated Flash Attention. It not only facilitates higher parallelism, enhancing processing efficiency but also retains compatibility with the existing FlashAttention. Correspondingly, we propose a progressive distillation to enable an efficient knowledge transfer from the vanilla SAM without costly training from scratch. Experiments on COCO and LVIS reveal that SAM-Lightening significantly outperforms the state-of-the-art methods in both run-time efficiency and segmentation accuracy. Specifically, it can achieve an inference speed of 7 milliseconds (ms) per image, for images of size 1024*1024 pixels, which is 30.1 times faster than the vanilla SAM and 2.1 times than the state-of-the-art. Moreover, it takes only 244MB memory, which is 3.5\% of the vanilla SAM. The code and weights are available at https://anonymous.4open.science/r/SAM-LIGHTENING-BC25/.

研究の動機と目的

• 現実世界のデプロイメントにおける Segment Anything Model (SAM) の計算ボトルネックに対処する。
• 推論を高速化しメモリを削減する Dilated Flash Attention を用いた効率的な画像エンコーダを導入する。
• 訓練なしで vanilla SAM から知識を転移する Dynamic Layer-Wise Distillation (DLD) を提案する。
• SAM-Lightening が COCO および LVIS において、効率性を大幅に高めつつ競争力のあるセグメンテーション性能を維持することを示す。

提案手法

• Vanilla self-attention を置換する Dilated Flash Attention メカニズムを設計し、より高い並列性と疎化されたセグメントを実現する。
• Dynamic Layer-Wise Distillation (DLD) を適用して、SAM から軽量エンコーダへ知識を段階的に転移する。
• 出力近傍の深い層に焦点を当てたディコプルド特徴蒸留を用いて、表現を教師と整合させる。
• 軽量エンコーダと整合するようデコーダを微調整し、プロンプトベースのセグメン tasks に対応する。
• 1% SA-1B データで訓練し、蒸留を加速するため SAM エンコーダ出力を保存し、標準ベンチマークで評価する。

実験結果

リサーチクエスチョン

• RQ1 再設計されたアテンション機構（Dilated Flash Attention）は、精度を犠牲にせず SAM のエンコーダを加速できるか。
• RQ2 Dynamic Layer-Wise Distillation は SAM から軽量エンコーダへ知識を効果的に転移できるか。
• RQ3 SAM-Lightening に対して、COCO および LVIS での速度、メモリ、およびセグメンテーション性能のトレードオフはどうなるか。
• RQ4 SAM-Lightening は Box、1P、3P の各プロンプトや Anything モードで最先端の SAM バリアントとどう比較されるか。

主な発見

Model	Enc. ms	Dec. ms	Tot. ms	S.U.	Mem.
SAM-ViT-H	216.1	3.8	219.9	1.0×	5.7GB
SAMFast	23.2	3.8	27.0	8.5×	4.1GB
FastSAM	20.7	3.4	24.1	9.1×	2.6GB
EfficientSAM	22.3	3.8	26.1	8.3×	309MB
MobileSAM	8.1	3.8	11.9	18.5×	309MB
SAM-Lightening	3.5	3.4	6.9	30.1×	224MB

• SAM-Lightening は 1024×1024 入力で 7 ms/画像を達成し、vanilla SAM の 30.1×、最先端と比較して 2.1× 高速。
• メモリ使用量は 224 MB に削減され、vanilla SAM の約 3.5% となる。
• 推論待機時間とメモリ効率は、プロンプトとデータセット全体で競合する軽量 SAM バリアントを上回る。
• SAM-Lightening は Box および Point プロンプト下で COCO および LVIS のセグメンテーション性能を vanilla SAM に近い水準で維持する。
• Dynamic Layer-Wise Distillation は、出力近くのより深い特徴層に焦点を当てつつ、段階的な層別重み付けで効率的な知識転移を可能にする。
• 点/ボックスプロンプトでのデコーダ微調整により、軽量エンコーダと凍結されたデコーダを整合させる。