ウェブサイト検索

視覚認識タスクにおける Mamba メカニズムの必要性の評価 - MambaOut

導入

トランスは、BERT、GPT シリーズ、ViT などのパワーアップ モデルのバックボーンです。ただし、そのアテンション メカニズムは 2 次の複雑さがあるため、長いシーケンスでは困難になります。これに取り組むために、線形複雑性を備えたさまざまなトークン ミキサーが開発されました。

最近、RNN ベースのモデルは、長いシーケンスに対する効率的なトレーニングと推論で注目を集めており、大規模な言語モデルのバックボーンとして有望であることが示されています。

これらの機能に触発されて、研究者は視覚認識タスクで Mamba を使用することを検討し、Vision Mamba、VMamba、LocalMamba、PlainMamba などのモデルにつながりました。それにもかかわらず、状態空間モデルや SSM ベースの視覚モデルは、最先端の畳み込みモデルや注意ベースのモデルと比較してパフォーマンスが劣ることが実験で明らかになりました。

この最近の論文は、新しいビジュアル Mamba モデルの設計には焦点を当てていません。代わりに、視覚認識タスクに Mamba は必要か?という重要な研究課題を調査します。

マンバとは何ですか？

Mamba は、カーネギー メロン大学とプリンストン大学の研究者によって開発された深層学習アーキテクチャであり、特に長いシーケンスにおける変換モデルの制限に対処するように設計されています。構造化状態空間シーケンス (S4) モデルを使用し、連続時間モデル、再帰モデル、畳み込みモデルの長所を組み合わせて、長い依存関係や不規則にサンプリングされたデータを効率的に処理します。

最近、研究者らは、ビジョン トランスフォーマー (ViT) の使用方法と同様に、コンピューター ビジョン タスクに Mamba を適応させました。 Vision Mamba (ViM) は、双方向状態空間モデル (SSM) を利用することで効率を向上させ、特に高解像度画像に対する従来の Transformer の高い計算要求に対処します。

マンバの建築

Mamba は、入力に基づいてパラメーターを適応させる独自の選択メカニズムを導入することで S4 モデルを強化し、シーケンス内の関連情報に焦点を当てることができます。この時間変化するフレームワークにより、計算効率が向上します。

また、Mamba は、GPU などの最新のハードウェアで効率的に計算を行うためのハードウェア対応アルゴリズムを採用し、パフォーマンスとメモリ使用量を最適化します。このアーキテクチャは SSM 設計を MLP ブロックと統合しており、言語、音声、ゲノミクスなどのさまざまなデータ タイプに適しています。

マンバの亜種

  • MambaByte: 生のバイト シーケンスを処理し、トークン化とそれに関連するバイアスを排除するトークンフリーの言語モデル。
  • Mamba Mixture of Experts (MOE): Mixture of Experts を Mamba と統合し、Mamba 層と MOE 層を交互に配置することで効率とスケーラビリティを強化します。
  • ビジョンマンバ (ViM): ViM は、ビジュアル シーケンス エンコーディングに双方向 Mamba ブロックを使用して、SSM をビジュアル データ処理に適応させます。これにより、計算量が削減され、ImageNet 分類、COCO オブジェクト検出、ADE20k セマンティック セグメンテーションなどのタスクのパフォーマンスが向上します。
  • Jamba: AI21 Labs によって開発された Jamba は、520 億のパラメーターと 256,000 のトークンのコンテキスト ウィンドウを備えたハイブリッド トランスフォーマーと Mamba SSM アーキテクチャです。

デモ

モデルの作業を開始する前に、リポジトリのクローンを作成し、必要なパッケージをいくつかインストールします。

!pip install timm==0.6.11
!git clone https://github.com/yuweihao/MambaOut.git
!pip install gradio

さらに、ステップを実行し、MambaOut で推論を実行するノートブックにアクセスするために使用できるリンクを追加しました。

cd /MambaOut

以下のセルは、gradio Web アプリを実行するのに役立ちます。

!python gradio_demo/app.py

RNN のようなモデルと因果的注意

以下の図は、因果的注意のメカニズムと RNN のようなモデルをメモリの観点から説明しています。ここで、 _xi_ は i 番目のステップの入力トークンを表します。

(a) 因果的注意: 以前のすべてのトークンのキー (k) と値 (v) をメモリとして保存します。メモリは、現在のトークンのキーと値を継続的に追加することで更新され、ロスレスになります。ただし、古いメモリと現在のトークンを統合する計算の複雑さは、シーケンスが長くなるにつれて増加します。したがって、注意力は短いシーケンスではうまく機能しますが、長いシーケンスでは困難になります。

(b) RNN のようなモデル: 以前のトークンを圧縮して、メモリとして機能する固定サイズの隠れ状態 (h) にします。この固定サイズは、RNN メモリが本質的に損失が多く、アテンション モデルの損失のないメモリ容量に匹敵できないことを意味します。それにもかかわらず、古いメモリと現在の入力をマージする複雑さはシーケンスの長さに関係なく一定のままであるため、RNN のようなモデルは長いシーケンスの処理に優れています。

Mamba は、その反復的な特性により、因果関係のあるトークンの混合を必要とするタスクに特に適しています。具体的には、Mamba は次の特性を持つタスクに優れています。

1.タスクには長いシーケンスの処理が含まれます。 2.タスクには因果関係のあるトークンの混合が必要です。

次の疑問は、視覚認識タスクには非常に長いシーケンスがあるのかということです。

ImageNet での画像分類の場合、一般的な入力画像サイズは 224x224 で、パッチ サイズが 16x16 のトークンが 196 個になります。この数値は、長時間シーケンスのタスクのしきい値よりもはるかに小さいため、ImageNet 分類は長時間シーケンスのタスクとは見なされません。

COCO でのオブジェクト検出とインスタンス セグメンテーション (画像サイズ 800x1280)、および ADE20K でのセマンティック セグメンテーション (ADE20K は、セマンティック セグメンテーション タスクに広く使用されているデータセットで、150 のセマンティック カテゴリで構成されています。このデータセットには、トレーニングに 20,000 枚の画像が含まれています)セットと検証セット内の 2,000 枚の画像）、画像サイズは512x2048、パッチ サイズ 16x16 でトークンの数は約 4,000 です。 4,000 トークンは小さなシーケンスのしきい値を超え、基本しきい値に近いため、COCO 検出と ADE20K セグメンテーションは両方とも長いシーケンスのタスクとみなされます。

MambaOutのフレームワーク

MambaOutの全体的な枠組み

図 (a) は視覚認識のための MambaOut の全体的なフレームワークを表しています: MambaOut は視覚認識用に設計されており、ResNet と同様の階層アーキテクチャに従っています。これは 4 つのステージで構成され、それぞれが異なるチャネル寸法 (_Di_ で示されます) を持ちます。この階層構造により、モデルは複数の抽象化レベルで視覚情報を処理できるようになり、画像内の複雑なパターンを認識する能力が強化されます。

(b) ゲート付き CNN ブロックのアーキテクチャ: Gated CNN ブロックは、MambaOut フレームワーク内のコンポーネントです。状態空間モデル (SSM) が含まれていないという点で Mamba ブロックとは異なります。どちらのブロックも、情報フローを制御するためにゲート機構を備えた畳み込みニューラル ネットワーク (CNN) を使用しますが、Gated CNN ブロックに SSM が存在しないことは、SSM を組み込んだ Mamba ブロックほど長いシーケンスや時間依存関係を処理する能力がないことを意味します。これらの目的のために。

Gated CNN と Mamba ブロックの主な違いは、状態空間モデル (SSM) の存在にあります。

MambaOut では、ConvNeXt と同様に、7x7 カーネル サイズの深さ方向の畳み込みが Gated CNN のトークン ミキサーとして使用されます。 ResNet と同様に、MambaOut は、図に示すように、各段階で Gated CNN ブロックを積み重ねることにより、4 段階のフレームワークを使用して構築されます。

次に進む前に、視覚認識に Mamba を導入する必要性に関する仮説をここに示します。 仮説 1: このタスクは特性 1 または特性 2 を満たさないため、ImageNet での画像分類に SSM を導入する必要はありません。 仮説 2: 特性 2 は満たされていないものの、これらのタスクは特性 1 と一致しているため、視覚的な検出とセグメンテーションに関する SSM の可能性をさらに調査する価値はあります。

トレーニング

ImageNet での画像分類

  • ImageNet は画像分類のベンチマークとして使用され、130 万枚のトレーニング画像と 50,000 枚の検証画像が使用されます。
  • トレーニングは、さまざまなデータ拡張手法や正則化手法を含む、蒸留することなく DeiT スキームに従います。
  • AdamW オプティマイザーは、学習率スケーリング ルール lr=バッチサイズ/1024 * 10^-3 でトレーニングに使用され、バッチ サイズ 4096 で学習率 0.004 となります。
  • MambaOut モデルは、PyTorch および timm ライブラリを使用して実装され、TPU v3 でトレーニングされます。

ImageNet での画像分類の結果

  • SSM が組み込まれていない MambaOut モデルは、ImageNet 上のすべてのモデル サイズにわたってビジュアル Mamba モデルよりも一貫して優れたパフォーマンスを発揮します。
  • たとえば、MambaOut-Small モデルは 84.1% のトップ 1 精度を達成し、MAC の 79% のみを必要としながら、LocalVMamba-S を 0.4% 上回ります。
  • これらの結果は仮説 1 を裏付けており、ImageNet 上の画像分類に SSM を導入する必要はないことを示唆しています。
  • Visual Mamba モデルは現在、ImageNet 上の最先端の畳み込みモデルやアテンション モデルに大きく遅れをとっています。
  • たとえば、CAFormer-M36 は、同等のサイズのすべてのビジュアル Mamba モデルよりも 1% 以上の精度で優れています。
  • 仮説 1 に挑戦することを目的とした将来の研究では、ImageNet で最先端のパフォーマンスを達成するために、畳み込みと SSM のトークン ミキサーを使用したビジュアル Mamba モデルを開発する必要があるかもしれません。

COCO でのオブジェクト検出とインスタンスのセグメンテーション

  • COCO 2017 は、オブジェクト検出とインスタンス セグメンテーションのベンチマークとして使用されます。
  • MambaOut は、Mask R-CNN 内のバックボーンとして利用され、ImageNet で事前にトレーニングされた重みで初期化されます。
  • トレーニングは 12 エポックの標準 1x スケジュールに従い、トレーニング画像のサイズは、短辺が 800 ピクセル、長辺が 1333 ピクセルを超えないようにサイズ変更されます。
  • AdamW オプティマイザーは、学習率 0.0001、合計バッチ サイズ 16 で使用されます。
  • 実装は PyTorch および mmdetection ライブラリを使用して行われ、FP16 の精度を利用してトレーニング コストを節約します。
  • 実験は 4 つの NVIDIA 4090 GPU で行われます。

COCO でのオブジェクト検出とインスタンスのセグメンテーションの結果

  • MambaOut は、COCO でのオブジェクト検出とインスタンス セグメンテーションにおいて一部のビジュアル Mamba モデルを上回るパフォーマンスを発揮しますが、VMamba や LocalVMamba などの最先端のビジュアル Mamba には依然として遅れをとっています。
  • たとえば、Mask R-CNN のバックボーンとしての MambaOut-Tiny は、VMamba-T を 1.4 APb および 1.1 APm 上回っています。
  • このパフォーマンスの違いは、長いシーケンスのビジュアル タスクに Mamba を統合する利点を強調し、仮説 2 を裏付けています。
  • ただし、Visual Mamba は、TransNeXt のような最先端のコンボリューション アテンション ハイブリッド モデルと比較すると、依然としてパフォーマンスに大きな差があります。 Visual Mamba は、視覚検出タスクにおいて他の最先端のモデルを上回るパフォーマンスを発揮することで、その有効性を実証する必要があります。

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

  • ADE20K はセマンティック セグメンテーション タスクのベンチマークとして使用され、トレーニング セットに 20,000 枚の画像、検証セットに 2,000 枚の画像を含む 150 のセマンティック カテゴリで構成されます。
  • Mamba は、ImageNet の事前トレーニングされた重みから初期化され、UperNet のバックボーンとして利用されます。
  • トレーニングは、AdamW オプティマイザーを使用して、学習率 0.0001、バッチ サイズ 16、160,000 反復で実行されます。
  • 実装は PyTorch および mmsegmentation ライブラリを使用して行われ、実験は 4 つの NVIDIA 4090 GPU で実行され、FP16 精度を利用してトレーニング速度が向上しました。

ADE20K でのセマンティック セグメンテーションの結果

  • COCO でのオブジェクト検出と同様に、ADE20K でのセマンティック セグメンテーションのパフォーマンス傾向は、MambaOut が一部のビジュアル Mamba モデルよりも優れたパフォーマンスを発揮できるものの、最先端の Mamba モデルの結果には匹敵しないことを示しています。
  • たとえば、LocalVMamba-T は単一スケール (SS) 評価とマルチスケール (MS) 評価の両方で MambaOut-Tiny を 0.5 mIoU 上回り、仮説 2 を経験的にさらに裏付けています。
  • さらに、ビジュアル Mamba モデルは、SG-Former や TransNeXt など、畳み込みとアテンション メカニズムを統合するより高度なハイブリッド モデルと比較して、顕著なパフォーマンスの低下を示し続けています。
  • Visual Mamba は、ビジュアル セグメンテーション タスクでより強力なパフォーマンスを達成することで、長期シーケンス モデリングにおける強みをさらに実証する必要があります。

結論

Mamba メカニズムは、長いシーケンスと自己回帰特性を持つタスクに最適です。 Mamba は、長いシーケンスの特性と一致する、視覚的な検出およびセグメンテーション タスクの可能性を示しています。 MambaOut モデルは、ImageNet 上のすべてのビジュアル Mamba モデルを上回っていますが、依然として最先端のビジュアル Mamba モデルには遅れています。

ただし、計算リソースの制限のため、この文書では視覚的なタスクに対する Mamba の概念を検証することに焦点を当てます。今後の研究では、Mamba と RNN の概念、さらには大規模言語モデル (LLM) と大規模マルチモーダル モデル (LMM) のための RNN と Transformer の統合がさらに探求され、これらの分野で新たな進歩がもたらされる可能性があります。

参考文献

  • 元の研究論文:- MambaOut: ビジョンに Mamba は本当に必要ですか?
  • Vision Mamba: 双方向状態空間モデルによる効率的な視覚表現学習
  • Mamba (深層学習アーキテクチャ)