Appleとオレゴン州立大学がAutoFocusFormerを提案: 従来のグリッドを廃止し、適応型ダウンサンプリング画像セグメンテーションを使用

従来の RGB 画像はラスター形式で保存され、ピクセルは画像全体に均等に分散されます。ただし、この均一な分布は、画像の実際のコンテンツの密度分布とは大きく異なることがよくあります。特に、現在一般的に使用されているディープ ネットワークでは、エンコード部分で頻繁にダウンサンプリングを行った後、小さなオブジェクトはごくわずかなポイントしか占有せず、大きなオブジェクトは多くのポイントを占有します。下の図に示すように、背景の混雑した群衆はほんの数個のドットで表現されていますが、画像の下部にある多数のドットは情報量が非常に少ない地面で占められています。画像認識プロセスを、保存される特徴の数と計算能力の観点から考えると、地面の特徴が大量に保存され、計算能力の大部分がこれらの地面の特徴を計算するために使用されることが想像できます。ただし、本当に重要な集団の場合、ポイントの数が少ないため、割り当てられる特徴も少なくなり、計算に使用される計算能力も非常に小さくなります。

写真

さまざまな小さな物体の特徴を記述するのに十分なポイントがないため、ディープ ネットワークは小さな物体やカメラから遠く離れた物体ではうまく機能しません。この問題は画像認識において非常に重要です。なぜなら、どのような種類のカメラや他の視覚センサーであっても、「近くの物体は大きく、遠くの物体は小さくなる」ことは避けられないからです。近くの物体の場合、センサーはより多くのサンプリング ポイント (ピクセル) を返しますが、遠くの物体の場合、サンプリング ポイント (ピクセル) は少なくなります。したがって、遠くにある物体を識別するのは当然難しくなります。しかし、標準モデルではダウンサンプリングが頻繁に行われ、遠くにあるオブジェクトや小さなオブジェクトのポイント数がさらに少なくなるため、実際にはこの問題がさらに悪化します。

現在主流となっている解決策は、1) 画像解像度の向上、2) モデルパラメータの数の増加です。これら 2 つの方式により、複数回のダウンサンプリング後に画像にさらに多くのポイントを含めることができるようになり、パラメータの数が増えると、複数のオブジェクトの特徴情報など、各ポイントにさらに多くの情報を保存できるようになります。しかし、どちらのソリューションもネットワーク計算量が大幅に増加し、物体認識の速度が低下します。同時に、大きな物体の表面の均一で特徴のない部分で多くの無意味な計算が無駄になります。

画像認識の主流がTransformerに完全に移行した時代でも、画像認識ネットワークでは従来のグリッドと均一ダウンサンプリングが依然として広く使用されています。これは、単純な均一ダウンサンプリングは実装が容易（たとえばストライド畳み込みを使用）であり、グリッドベースの構造はGPU上で計算できるためと考えられます。しかし、グリッド構造は実際には計算負荷をより効率的に分散する方法を妨げ、小さくて遠くにある物体を識別するという前述の問題が、力ずくの方法でしか解決できないという結果になってしまいます。文献には変形可能な畳み込みや変形可能な注意に基づく優れたスキームが数多くありますが、それらは依然として画像のグリッド表現に限定されており、ダウンサンプリングやアップサンプリングの操作のために画像をグリッド形式に戻す必要がある場合が多くあります。

従来のグリッドによって引き起こされる問題を認識して、検討できる別の解決策は、より柔軟なポイント クラウド形式を使用して画像の特徴を保存することです。ポイント クラウド形式を使用すると、画像の各部分で異なる密度を使用できるため、大きな特徴のない表面上のサンプリング ポイントの数が減り、遠く離れた領域や小さなオブジェクトが集中している領域ではより多くのサンプリング ポイントが保持されます。しかし、ポイントクラウドの表現は主流のグリッドとはまったく異なるため、ポイントクラウドを使用する画像認識ソリューションでは、ポイントクラウド上の近傍をどのように分割するかなど、多くの問題に直面します。上記の機能を実現するための適応ダウンサンプリング方法を学ぶにはどうすればよいでしょうか?地域によって異なる近隣規模の問題にどう対処するか?デコード方法は？

Apple とオレゴン州立大学の研究者は、ポイント クラウドを使用した画像セグメンテーションの完全なソリューションを提案しました。このソリューションは、上記の問題の多くを解決し、画像セグメンテーションの損失関数、ローカル アテンション Transformer レイヤー、ポイント クラウド ベースの Transformer デコーダーに基づいてエンドツーエンドで学習できるダウンサンプリング アルゴリズムを提案しています。

写真

• 論文アドレス: https://arxiv.org/abs/2304.12406
• コードアドレス: https://github.com/apple/ml-autofocusformer

彼らは、ADE20K および Cityscapes データセット、特に難しく、遠くの小さなオブジェクトが多数含まれる Cityscapes データセットで、通常のラスター アルゴリズムよりも優れた結果を達成しました。彼らは、197M パラメータの大規模モデルを使用する SOTA アルゴリズム Mask2Former を、42.6M パラメータの小規模モデルのみで上回りました。 PapersWithCode ランキングでは、AutoFocusFormer-Base モデルが、インスタンス セグメンテーション ベンチマークで Mapillary Vista の追加トレーニング データを使用しないすべての Large 以上のモデルを上回っています。

写真

モデル

本論文で提案する AutoFocusFormer モデルは、Transformer をベースにした適応型モデルです。複数の適応型ダウンサンプリングにより、高解像度の画像を少数の特徴点群に素早く縮小できます。さらに、モデルはタスクの目的 (損失関数) に応じて画像のさまざまな領域のサンプリング密度を自動的に調整するため、生成された特徴点群は情報量の多い領域では密度が高くなり、重要度の低い領域では密度が低くなります。

これまでのいくつかの Transformer ベースの適応モデル (AdaViT、DynamicViT、A-ViT など) では、通常、グローバル アテンションを使用します。つまり、自己アテンションでは、すべてのトークンが他のすべてのトークンに注意を向けます。これによってもたらされる二次の複雑さにより、モデルを高解像度の画像セグメンテーション タスクに効率的に適用することができません。これらの方法は、微分可能なバイナリ マスクをトレーニングして「重要な」トークンと「重要でない」トークンを区別し、モデル推論でゼロ値を持つポイントを破棄します。ただし、これらの方法では勾配を伝播できないため、トレーニング中に連続的なダウンサンプリングを実行することはできません。さらに、これらの方法では、マスクのサイズを制限するための追加の損失関数も必要です。

AutoFocusFormer の設計では、これらの欠点をうまく回避しています。AFF はローカル アテンションを使用して、各トークンのアテンションを固定サイズの近傍に制限します。AFF は新しい近傍マージ方法を使用して、トレーニング中に適応型ダウンサンプリングを実際に学習します。AFF は、タスク自体の損失関数以外のガイダンスを必要としません。その結果、AutoFocusFormer は高解像度の入力で実行できる最初の適応型ダウンサンプリング モデルとなり、画像セグメンテーションなどのピクセルレベルの予測タスクに適したものになりました。

下の図は、AFF のモデル フレームワーク構造を示しています。 AFF は、2 つの畳み込み層で構成されるパッチ埋め込みから始まり、4 つの段階を経て、最終的に画像分類またはセグメンテーション デコーダーによる予測結果を出力します。各ステージは、バランスのとれたクラスタリング、ローカル アテンション ブロック、適応型ダウンサンプリングの 3 つのモジュールに分かれています。その中でも最も重要なのは、革新的なアダプティブ ダウンサンプリング モジュールです。

1. バランスのとれたクラスタリング

AFF はローカル アテンションを使用します。つまり、各ポイントは近傍の K 個の隣接ポイントにのみ注意を払うことができるため、高解像度の状況での計算の複雑さが軽減されます。したがって、グローバルアテンションと比較して、モデルでは最初に各ポイントの近傍を定義する必要があります。ポイント クラウドでは、従来の方法は K 最近傍法です。これは、ユークリッド距離に基づいて各ポイントに最も近い K ポイントを見つけ、それらのセットをこのトークンの近傍として定義します。ただし、KNN アルゴリズムは高解像度の画像では効率的ではありません。いくつかの効率的な KNN アルゴリズムにヒントを得て、著者らはまずポイント クラウドを等しいサイズの (小さな) クラスター (たとえば 8 個のポイント) に分割し、次に各ポイントの近傍をそれに最も近い R 個のクラスター (たとえば 6 個のクラスター) として定義します。クラスターのサイズは同じなので、各トークンには同じ数の隣接トークンがあり、ゼロパディングを使用して GPU メモリを無駄にする必要はありません。また、各トークンの近傍が重なり合うことで、画像全体にスムーズに情報を伝達できるため、Swin モデルのように近傍の位置を交互に切り替える必要がなく、近傍を見つける問題を効率的に解決できます。

各クラスター内のポイント数が等しくなるようにするために、著者らは論文の中で革新的な均一クラスタリング アルゴリズムを提案しました。 k-means や局所性に敏感なハッシュなどの従来のクラスタリング アルゴリズムでは、複数回の反復が必要であり、各クラスターのサイズが同じになることは保証されません。クラスターのサイズが不均等な場合、近隣のサイズも異なり、メモリが無駄になります。 AFF で使用される方法では反復処理は必要なく、クラスター サイズが同じであることが保証されます。この方法では、まず空間充填曲線 (ヒルベルト曲線など) を使用して 2D 画面上のすべてのトークンを 1D 配列に接続し、次にその配列をターゲットのクラスター数に直接分割します。これにより、各クラスター内のサンプル数が厳密に等しくなります。ただし、トークンを曲線で単純に接続すると、生成されたクラスターの円周がすべての方向に不均一になります。これに基づいて、著者らは、トークンと充填曲線間のバッファ構造として使用される空間充填アンカーの概念を提案しました。詳細な方法については、AFF の記事を参照してください。

2. 地元の注目

ポイント クラウドを同じサイズのクラスターに分割した後、モデルは各トークンの近傍を、それに最も近い R クラスターとして定義します。各近傍において、AFF は標準の QKV ローカル自己注意式を使用します。

写真

式中、P は位置埋め込み行列です。著者らは、単純な完全接続層を使用して、トークンとその近傍の間の相対位置情報 (x 座標と y 座標の差) をエンコードし、モデルが注意を計算するプロセスで位置情報を使用できるようにしています。

ただし、x 座標と y 座標の差のみから導出される位置エンコーディングは、スケール不変でも回転不変でもありません。グリッドベースのローカル アテンション モデルでは、近傍のサイズは常に同じであるため、これは問題になりません。ただし、ポイント クラウド ベースの近隣は、場所や画像によってスケールが異なる場合があります。この時点では、モデルは異なるサイズの同じオブジェクトや、回転後の同じオブジェクトに一般化することはできません。そのため、著者らは相対位置情報を拡張し、2つのトークン間の距離、コサイン値、サイン値を含めることにしました。

距離は回転不変ですが、コサイン値とサイン値はスケール不変であることに注意してください。したがって、この位置エンコーディングを通じて、ディープ モデルは使用する必要がある情報を自由に選択することを学習でき、柔軟性が向上します。

3. 適応ダウンサンプリング

各ステージの最後のモジュールとして、アダプティブ ダウンサンプリングはトークンの数を元の x% に削減する役割を果たします。従来のストライド 2 ダウンサンプリングでは数値が 1/4 に削減されますが、AFF のより柔軟なポイント クラウド構造では x% の値に制限はありません。実験では、著者らは 1/4 および 1/5 のダウンサンプリング レートでのモデルのトレーニング結果を示しました。

AFF の適応ダウンサンプリング モジュールは、1. 重要度スコアを計算する、2. 重要度に応じてマージ センターを選択する、3. 近傍をマージする、という 3 つのステップに分かれています。最後のステップを下の図に示します。

写真

最初のステップでは、モデルは各トークンの特徴ベクトルを完全接続層に入力し、それをシグモイド関数に渡して 0 から 1 の間の重要度スコアを取得します。次に、次のステップで、最も「重要度」の高い x% のトークンがマージ センターとして選択されます。最後に、モデルは PointConv レイヤーを使用して、選択された各トークンの近傍にあるすべての特徴ベクトルを新しいベクトルにマージします。

マージ プロセス中、PointConv レイヤー内の各隣接ポイントの畳み込み重みは、それぞれの「重要度スコア」によって調整されることに注意してください。トークンの重要度が 0 に近い場合、その特徴ベクトルは、マージされたベクトル特徴に大きな影響を与えません。この変調プロセスにより、重要度の勾配を選択されていないポイントにまで伝播させることができるため、モデルはタスク損失に応じてパラメータを自動的に調整し、比較的重要なポイントを保持することができます。これにより、画像の内容とタスクの目的に基づいた適応型ダウンサンプリングの効果が得られます。

最後に、この記事ではグリッド事前分布についても言及しています。著者らは、「重要度」が特徴のみに基づいて計算され、ポイントの 2D 位置が完全に無視されると、特徴のない領域 (大きなオブジェクトの表面など) で過剰なランダム サンプリングが発生すると示唆しています。このような均一な特徴の領域では、理想的には、モデルは自動的に従来の均一密度ダウンサンプリングにフォールバックできます。したがって、トークンを選択するときに、著者は「重要度」にグリッドの事前レイヤーを追加して、従来のグリッド サンプリングでそれらのポイントに優先順位を付けます。具体的な方法については論文を参照してください。

4. デコーダー（ポイントベースデコーダー）

従来の画像セグメンテーション デコーダーもグリッドに基づいています。この論文では、Mask2Former デコーダーをポイント クラウドに一般化し、3x3 畳み込みの代わりに PointConv を使用し、ポイント クラウドに変形可能な注意を実装するための補間方法を提供し、適応型ダウンサンプリング ポイント クラウド スキームを使用して画像セグメンテーションを学習するためのエンドツーエンドのネットワークを実現します。

実験

著者らは、画像分類および画像セグメンテーションタスクにおける AFF の実験結果を紹介します。

彼らは、古典的な ImageNet-1K データセットで AFF の画像分類機能をテストしました。ベースライン メソッドの Swin Transformer と比較すると、AFF はモデル サイズに関係なくパフォーマンスが向上します。特に、AFF の 1/5 ダウンサンプリング バージョンは、Swin Transformer の精度を上回るだけでなく、FLOP も約 30% 節約します。

著者らは、ADE20K データセットで AFF のセマンティック セグメンテーション機能を実験しました。同様に、モデルのサイズに関係なく、AFF の 1/5 ダウンサンプリング バージョンはベースライン パフォーマンスを上回るだけでなく、約 30% の FLOP を節約します。

最後に、著者らは Cityscapes データセットでインスタンスセグメンテーションとパノプティックセグメンテーションに関する実験を実施しました。実験結果によると、AFF の適応能力は、都市景観内のオブジェクトのサイズに大きな差がある屋外シーンに非常に適していることがわかりました。特に注目すべきは、AFF-Tiny が Swin-Base と同等のパフォーマンスを実現するのに対し、4260 万個のパラメータのみを使用する AFF-Small モデルは、197 万個のパラメータを使用する Swin-Large と同等のパフォーマンスを実現し、約 78% のパラメータ節約を実現していることです。

次の 2 つの例は、AFF の適応能力が Cityscapes 上の遠くにある小さな物体を識別するのにどのように役立つかを直感的に理解するのに役立ちます。例 1 では、4 回のダウンサンプリングを行った後でも、画像内で非常に小さい遠くの車に対して十分なサンプルが残っています。その結果、モデルはこれらのオブジェクトを正常にキャプチャできます。対照的に、グリッドベースのモデルでは遠くにあるこれらの車を見逃してしまいます。

写真

以下の例では、AFF はベースライン アルゴリズム Swin Transformer よりも左側のコーヒー ショップの頭部を多く認識します。 AFF は、中央の大きなグレーの領域に計算能力を浪費する代わりに、道路脇の混雑した屋外コーヒーショップに、より多くのサンプリング ポイントをより柔軟に割り当てます。

写真

要約する

AutoFocusFormer は、マルチレイヤー適応ダウンサンプリングを使用し、高解像度の入力で実行できる最初のモデルであり、画像セグメンテーションなどのピクセルレベルの予測タスクを実行するのに適しています。柔軟なデータ構造により、より多様なダウンサンプリング レート オプションが提供され、実験では、AFF が適切なポイントを正しくサンプリングできるため、1/5 のダウンサンプリング レート モデルは、計算量を大幅に節約しながらベースラインよりも優れた結果を達成できることが示されています。包括的な実験結果に基づいて、AFF は FLOP 数とパラメータの数を大幅に節約しながら、小さな物体の認識において大きな進歩を遂げました。