0. はじめに
MetaAI は、「すべてをセグメント化」できる視覚的な基本的な大規模モデル SAM を提案して以来、優れたセグメンテーション効果を提供し、大規模な視覚モデルを探索するための新しい方向性を提供してきました。SAM の効果は非常に優れていますが、SAM のバックボーンは ViT を使用しているため、推論中に多くのビデオ メモリを消費し、推論速度が遅く、ハードウェア要件が高いため、プロジェクト アプリケーションには大きな制限があります。FastSAM は検出ヘッドを交換することでより優れたリアルタイム セグメンテーションを実現できますが、モバイル アプリケーションに導入できるように SAM の「重い」デコーダを「軽い」デコーダに置き換えることはまだ非現実的です。そのため、「 Faster Segment Anything: Towards Lightweight 」 SAM for Mobile Applications は、 SAM の ViT-H デコーダを蒸留するための「分離蒸留」ソリューションを提案しており、その結果得られる軽量エンコーダは SAM のデコーダと「シームレスに互換性」を持つことができます。関連するコードもGithubでオープンソース化されています
図 1: セグメント何でもモデルの概要
1. 主な貢献
SAM (Space-Aware Multi-Head tention) のデフォルトの画像エンコーダが ViT-H (Vision Transformer-H) に基づいていることを考慮すると、 MobileSAMを取得する直接的な方法は、Kirillov et al. [2023] の公式プロセスに従うことです。 、ViT-L (Vision Transformer-L) または小型の ViT-B (Vision Transformer-B) などの小型の画像エンコーダーを使用して、新しい SAM を再トレーニングします。
- SAM アーキテクチャ システムの継続: 軽量 ViT デコーダを使用して潜在特徴を生成し、プロンプトを使用してデコーダをガイドして目的のマスクを生成します。
- この記事では、SAM の ViT-H デコーダを蒸留するための「分離蒸留」スキームを提案し、その結果得られる軽量エンコーダは SAM のデコーダと「シームレスに互換性」を持つことができます。
- 私たちの方法には、シンプルで効果的であると同時に、低コスト (単一 GPU で 1 日未満) で再現可能であるという利点があります。結果として得られる MobileSAM は、エンコーダ パラメータを 100 分の 1 に削減し、合計パラメータを 60 分の 1 に削減します。推論速度の点では、MobileSAM は画像の処理にわずか 10 ミリ秒 (エンコーダーで 8 ミリ秒、デコーダーで 2 ミリ秒) しかかからず、これは FastSAM の 4 倍高速であるため、MobileSAM はモバイル アプリケーションに非常に適しています。
2. モバイル SAM の背景とプロジェクトの目標
2.1 SAM（Segment Anything Model、「何でも切る」モデル）の背景
まずはSAMの構造と仕組みをまとめてみましょう。SAM は、ViT (Vision Transformer) ベースの画像エンコーダーとヒントガイド付きマスク デコーダーで構成されます。画像エンコーダは画像を入力として受け取り、埋め込みを生成し、それがマスク デコーダに供給されます。マスク デコーダは、背景からオブジェクトを切り取るためのヒント (ポイントやボックスなど) に基づいてマスクを生成します。さらに、SAM では、曖昧さを解決するために同じキューに対して複数のマスクを生成できるため、貴重な柔軟性が得られます。これを念頭に置いて、SAM のパイプラインを維持し、最初に ViT ベースのエンコーダーを使用して画像埋め込みを生成し、次にヒントガイド付きデコーダーを使用して必要なマスクを生成します。このプロセスは「何でもカット」に最適で、下流タスクの「すべてをカット」に使用できます。
2.2 プロジェクトの目標
このプロジェクトの目標は、軽量な方法で満足のいくパフォーマンスを実現し、元の SAM よりも高速なモバイル指向 SAM (MobileSAM) を生成することです。オリジナルの SAM のヒントガイド付きマスク デコーダのパラメータは 4M 未満であるため、軽量であると考えられます。公開デモンストレーションで示されているように、マスク デコーダは軽量であるため、エンコーダによって処理される画像埋め込みがある限り、SAM はリソースに制約のあるデバイス上でも実行できます。ただし、元の SAM のデフォルトの画像エンコーダは ViT-H に基づいており、6 億を超えるパラメータがあり、非常に重く、SAM プロセス全体がモバイル デバイスと互換性がなくなります。したがって、モバイル対応 SAM を取得する鍵は、重い画像エンコーダを軽量画像エンコーダに置き換えることです。軽量画像エンコーダは、元の SAM のすべての機能と特徴も自動的に保持します。次に、このプロジェクトの目標を達成するために提案されたアプローチについて詳しく説明します。
3. モバイルSAMが提案する手法
3.1 カップリングの改良
私たちのプロジェクトの目標を達成する簡単な方法は、 Kirillov et al. [2023] の公式プロセスに従って、より小さな画像エンコーダーを使用して新しい SAM を再トレーニングすることです。Kirillov et al. [2023] が説明しているように、ViT-H 画像エンコーダを使用して SAM をトレーニングするには、256 個の A100 GPU で 68 時間の実行時間が必要です。ViT-H を ViT-L または ViT-B に置き換えると、必要な GPU は 128 個に減りますが、コミュニティの多くの研究者にとって、結果を再現または改善することは依然として大きな負担となります。彼らのアプローチに従って、私たちはより小さな画像エンコーダーを採用し、彼らが提供したセグメンテーション データセット (11-T) を使用して新しい SAM を再トレーニングすることができます。データセットで提供されるマスクは、事前にトレーニングされた SAM (ViT 画像エンコーダーを使用) によって提供されることに注意することが重要です。本質的に、この再トレーニング プロセスは知識の蒸留 (Hinton et al.、2015) であり、ViT-H ベースの SAM からより小型の画像エンコーダーを備えた SAM に知識を転送します (図 2 左を参照)。
図 2: SAM の結合知識抽出。左側のグラフは完全に結合された改良を表し、右側のグラフは半結合された改良を表します。
3.2 セミカップリングからデカップリングおよびリファインへ
元の SAM からより小型の画像エンコーダを備えた SAM への KD (知識蒸留) を実行する場合、主な問題は画像エンコーダと結合されたデコーダの結合の最適化にあります。直感的には、画像エンコーダの最適化は画像デコーダの品質に依存し、逆も同様です。SAM の両方のモジュールが悪い状態にある場合、両方を良好な状態にトレーニングすることはより困難になります。分割統治アルゴリズム (Zhang et al., 2022c) にヒントを得て、KD タスクを 2 つのサブタスク (画像エンコーダーの洗練とマスク デコーダーの微調整) に分割することを提案します。具体的には、まず画像エンコーダで KD を実行し、ViT-H からの知識をより小さなエンコーダに転送します。オリジナルの SAM のマスク デコーダはすでに軽量であるため、そのアーキテクチャを維持する予定です。これにより、既製の組み合わせデコーダーをゼロからトレーニングするのではなく、微調整に直接使用できるという利点が得られます。結合リファインメントの最適化問題を軽減するための簡単なアプローチは、コピーおよび凍結されたマスク デコーダーを使用して画像エンコーダーを最適化することです (図 2 右を参照)。フリーズ操作は、画像エンコーダの品質の低下によるマスク デコーダの品質の低下を防ぐのに役立ちます。画像エンコーダの最適化がまだマスク デコーダから完全に分離されていないため、この改良を半結合と呼びます。実際には、キューの選択がランダムであるため、マスク デコーダーが可変になり、最適化の難易度が高くなるため、この最適化は依然として困難であることがわかりました。したがって、結合デコーダに依存せずに、元の SAM の ViT-H から直接小型画像エンコーダをリファインすることを提案します。これは分離リファインメントと呼ばれます (図 3 を参照)。画像埋め込みのリファインメントを実行するもう 1 つの利点は、Kirillov et al. [2023] (Lin et al., 2017) のようにマスク予測を行うために焦点損失を使用する代わりに、単純な MSE (平均二乗誤差) 損失を採用できることです。サイコロの損失 (Milletari et al., 2016)。
図 3: SAM の分離された知識の抽出
3.3 マスクデコーダの微調整の必要性について
半結合蒸留とは異なり、分離蒸留によってトレーニングされた軽量のエンコーダーには、凍結したデコーダーとの位置ずれの問題が発生する可能性があります。**経験に基づいて、この現象は存在しないことがわかりました。これは、生徒のエンコーダーによって生成された潜在特徴が、元の教師のエンコーダーによって生成された潜在特徴に非常に近いため、デコーダーと組み合わせて微調整する必要がないためです。もちろん、組み合わせをさらに微調整すると、パフォーマンスがさらに向上する可能性があります。