Visual Mamba が登場: 2.8 倍高速、メモリが 87% 削減

「万能トランスフォーマー」として知られるマンバは、発売から2か月も経たないうちに高性能ビジュアルバージョンが登場した。

木曜日、華中科技大学、Horizo​​n Robotics、智遠人工知能研究所などの研究者らがVision Mamba（Vim）を提案した。

• 論文アドレス: https://arxiv.org/pdf/2401.09417.pdf
• プロジェクトアドレス: https://github.com/hustvl/Vim
• 論文タイトル: Vision Mamba: 双方向状態空間モデルによる効率的な視覚表現学習

効果は何ですか？ ImageNet 分類タスク、COCO オブジェクト検出タスク、ADE20k セマンティック セグメンテーション タスクでは、Vim は DeiT などの成熟したビジュアル トランスフォーマーよりも高いパフォーマンスを実現し、計算効率とメモリ効率も大幅に向上します。たとえば、解像度 1248×1248 の画像に対してバッチ推論と特徴抽出を実行する場合、 Vim は DeiT よりも 2.8 倍高速で、GPU メモリを 86.8% 節約します。結果は、Vim が高解像度画像に対して Transformer スタイルの理解を実行する際の計算およびメモリの制限を克服でき、ビジョンベース モデルの次世代バックボーンになる大きな可能性を秘めていることを示しています。

次に、論文の内容を見てみましょう。

Mamba の導入により、状態空間モデル (SSM) に対する研究者の関心が高まっています。コンテキストの長さの増加に伴って計算の複雑さが 2 乗的に増加する Transformer の自己注意メカニズムとは異なり、SSM は長距離の依存関係を捉えるのに適しているため人気が高まっています。

この期間中、線形状態空間レイヤー (LSSL)、構造化状態空間シーケンス モデル (S4)、対角状態空間 (DSS)、S4D などの SSM ベースの方法が、特に長距離依存関係のモデリングにおいて、さまざまなシーケンス データを処理する研究者によって提案されました。

Mamba は、時間によって変化するパラメータを SSM に組み込み、効率的なトレーニングと推論を実現するためのハードウェア対応アルゴリズムを提案します。 Mamba の優れたスケーリング パフォーマンスは、言語モデリングにおける Transformer の有望な代替手段であることを示唆しています。

しかし、これまでのところ、研究者は視覚タスク用の一般的な純粋な SSM ベースのバックボーン ネットワークをまだ調査していません。

言語モデリングにおける Mamba の成功に触発され、研究者たちは、この成功を言語から視覚に応用できるかどうか、つまり、高度な SSM 手法を使用して汎用的で効率的な視覚バックボーンを設計できるかどうかを考え始めました。ただし、Mamba の独自のアーキテクチャにより、一方向モデリングと位置認識の欠如という 2 つの課題に対処する必要があります。

これらの問題に対処するために、研究者らは、データ依存型のグローバル視覚コンテキストモデリングのための双方向 SSM と位置認識視覚認識のための位置埋め込みを組み合わせた Vision Mamba (Vim) ブロックを提案しました。

視覚タスク用の他の SSM ベースのモデルと比較すると、Vim は純粋な SSM ベースのアプローチであり、画像を順次モデル化します。 Transformer ベースの DeiT と比較して、Vim は ImageNet 分類で優れたパフォーマンスを実現します。さらに、Vim は GPU メモリと高解像度画像の推論時間の点でより効率的です。

方法の紹介

Vision Mamba (Vim) の目標は、高度な状態空間モデル (SSM)、つまり Mamba をコンピューター ビジョンに導入することです。

Vim の概要を図 2 に示します。標準 Mamba は 1 次元シーケンス用に設計されています。視覚タスクを処理するには、まず2D 画像を展開された 2D パッチに変換する必要があります。ここで、(H, W) は入力画像のサイズ、C はチャンネル数、P は画像パッチのサイズです。次に、x_p をサイズ D のベクトルに線形投影し、位置埋め込みを追加して次の式を取得する必要があります。

Vimブロック

オリジナルの Mamba ブロックは 1 次元シーケンス用に設計されており、空間知覚理解を必要とする視覚タスクには適していません。 Vim ブロックは、視覚タスク用の双方向シーケンス モデリングを統合します。Vim ブロックは、上の図 2 に示されています。

Vim ブロックの操作アルゴリズムは次のとおりです。

アーキテクチャの詳細

アーキテクチャのハイパーパラメータは次のとおりです。

• L: ブロック数
• D: 隠れた状態の次元
• E: 拡張された状態次元
• N: SSM次元

ViT と DeiT に続いて、本研究ではまずカーネル サイズが 16×16 の投影層を採用し、重複しないパッチ埋め込みの 1 次元シーケンスを取得します。次に、L Vim ブロックを積み重ねるだけです。デフォルトでは、ブロック数 L は 24 に設定され、SSM 次元 N は 16 に設定されています。 DeiT シリーズ モデルのサイズと一致させるために、この研究では、極小サイズのバリアントの隠し状態次元 D を 192 に設定し、拡張状態次元 E を 384 に設定します。小型変異体の場合、研究では D を 384、E を 768 に設定しました。

実験

この研究では、ImageNet-1K データセットで Vim をベンチマークしました。

画像分類

表 1 は、Vim と ConvNet ベース、Transformer ベース、および SSM ベースのバックボーン ネットワークを比較しています。 ConvNet ベースの ResNet と比較すると、Vim は優れたパフォーマンスを示します。たとえば、パラメータがほぼ同じ場合、Vim-Small はトップ 1 の精度 80.3% を達成し、これは ResNet50 よりも 4.1 パーセントポイント高くなります。従来の自己注意ベースの ViT と比較すると、Vim はパラメータ数と分類精度の両方でかなりの利点があります。高度に最適化された ViT バリアントである DeiT と比較すると、VimTiny は DeiT-Tiny よりも 0.9 ポイント高く、Vim-Small は DeiT よりも 0.5 ポイント高くなります。 SSM ベースの S4ND-ViTB と比較すると、Vim は 3 分の 1 のパラメータで同様のトップ 1 精度を実現します。

図 1 (b) と (c) は、小さな Vim と DeiT の FPS と GPU メモリを比較しています。画像の解像度が高くなるにつれて、Vim は速度とメモリの面でより効率的になります。具体的には、画像サイズが 512 の場合、Vim は DeiT と同様の FPS とメモリを実現します。画像サイズが 1248 に増加すると、Vim は DeiT よりも 2.8 倍高速になり、GPU メモリを 86.8% 節約します。 Vim はシーケンス長を線形にスケーリングできるという大きな利点があり、高解像度のダウンストリーム ビジョン アプリケーションや長いシーケンスのマルチモーダル アプリケーションに最適です。

セマンティックセグメンテーション

表 2 に示すように、Vim はさまざまなスケールで一貫して DeiT よりも優れています。Vim-Ti は DeiT-Ti よりも 1.0 mIoU 高く、Vim-S は DeiT-S よりも 0.9 mIoU 高くなっています。 ResNet-101 バックボーン ネットワークと比較すると、Vim-S はほぼ 2 分の 1 のパラメータで同じセグメンテーション パフォーマンスを実現します。

下流のタスク (セグメンテーション、検出、インスタンスのセグメンテーションなど) に対するアプローチの効率をさらに評価するために、バックボーン ネットワークを一般的に使用される Feature Pyramid Network (FPN) モジュールと組み合わせ、その FPS と GPU メモリをベンチマークします。

図 3 および 4 に示すように、この研究ではバックボーン ネットワークに重い FPN を接続しましたが、効率曲線は純粋なバックボーン ネットワークの比較結果 (図 1) と同様でした。

物体検出とインスタンスセグメンテーション

表 3 は、Cascade Mask R-CNN フレームワークを使用して Vim-Ti と DeiT-Ti を比較したものです。 Vim-Ti は DeiT-Ti 1.3 ボックス AP および 1.1 マスク AP を上回ります。

下の図は可視化結果です。この論文の手法では、画像内で非常に大きなオブジェクトを捉えることができますが、これは DeiT-Ti らの手法では不可能です。

詳細については、原文論文を参照してください。