モダンな機械学習および人工知能フレームワークでは、トランスフォーマーはGPTシリーズやBERTなどの自然言語処理、コンピュータービジョンタスクなどのさまざまなドメインで最も広く使用されているコンポーネントの1つです。ただし、トランスフォーマーのアテンションモジュールはシーケンスの長さに比例して計算コストが増加し、高い計算コストの課題をもたらします。さまざまなモデルは、カーネル化、履歴メモリ圧縮、トークン混合範囲の制限、低ランクアプローチなどのさまざまな戦略を使用してこの課題に対処してきました。最近、MambaやRWKVなどの再帰型ニューラルネットワークは、大規模言語モデルの結果が約束されているため、注目を集めています。

マンボは、再帰型ニューラルネットワークのようなトークンミキサーのアーキテクチャを持つモデルファミリーで、最近、トランスフォーマーのアテンションメカニズムの二次的な複雑さに対処するために導入され、後にビジョンタスクに適用されました。研究者は、MambaやSSM（State Space Model）を視覚認識タスクに組み込む方法をすでに探索しており、Mambaを使用して等方性ビジョンモデルを開発するVision Mambaはその例です。一方、LocalMambaは視覚マンボモデルを強化するためにローカル誘導バイアスを組み込み、VMambaフレームワークはベースマンボモデルを使用してResNetやAlexNetに似た階層モデルを構築します。ただし、ビジョン認識コンテキストタスクにマンボフレームワークは本当に必要なのでしょうか。マンボファミリーのモデルのパフォーマンスは、伝統的なアテンションベースおよび畳み込みモデルの比較で、ビジョンタスクで今まで不調な結果を出しています。

MambaOutは、Mambaファミリーのモデルの性質を調査し、Mambaはオートリグレッシブまたは長シーケンスの特性を持つタスクに適しているという仮説を立てています。ただし、ほとんどのビジョンタスクはこれらの特性の両方を備えていないため、MambaOutは次の2つの仮説を提案しています。まず、状態空間モデルは画像分類に必要ないということです。画像分類タスクは、長シーケンスまたはオートリグレッシブの特性のいずれにも対応していないからです。2つ目は、インスタンスセグメンテーションやセマンティックセグメンテーション、およびオブジェクト検出は長シーケンスの特性を示しますが、オートリグレッシブではないため、状態空間モデルはこれらのタスクに有益である可能性があるということです。実験結果は、MambaOutフレームワークの仮説を支持しており、MambaOutフレームワークは画像分類タスクで視覚マンボモデルを上回り、検出およびセグメンテーションタスクでは最先端のマンボモデルに匹敵することができません。

この記事では、MambaOutフレームワークを深く掘り下げ、メカニズム、方法、フレームワークのアーキテクチャ、最先端のフレームワークとの比較について探ります。では、始めましょう。

マンボアウト：マンボは本当にビジョンに必要ですか？

機械学習アプリケーションの進歩とともに、トランスフォーマーは、Vision Transformers、GPTシリーズ、BERTなどのモデルを支える主流のバックボーンとして登場しました。ただし、トランスフォーマーのトークンミキサーはシーケンスの長さに比例して二次的な複雑さをもたらします。これを解決するために、Linformer、Longformer、Performer、Dynamic Convolution、Big Birdなどのトークンミキサーが導入されています。ただし、最近、RNNのようなモデルは並列化可能なトレーニングと長シーケンスでの効率的なパフォーマンスを提供するため、注目を集めています。Mambaファミリーのモデルは、トランスフォーマーのトークンミキサーの二次的な複雑さに対処するために導入され、最近、視覚認識タスクに適用されています。

MambaOutは、Mambaファミリーのモデルの性質を調査し、Mambaはオートリグレッシブまたは長シーケンスの特性を持つタスクに適しているという仮説を立てています。ただし、ほとんどのビジョンタスクはこれらの特性の両方を備えていないため、MambaOutは次の2つの仮説を提案しています。まず、状態空間モデルは画像分類に必要ないということです。画像分類タスクは、長シーケンスまたはオートリグレッシブの特性のいずれにも対応していないからです。2つ目は、インスタンスセグメンテーションやセマンティックセグメンテーション、およびオブジェクト検出は長シーケンスの特性を示しますが、オートリグレッシブではないため、状態空間モデルはこれらのタスクに有益である可能性があるということです。実験結果は、MambaOutフレームワークの仮説を支持しており、MambaOutフレームワークは画像分類タスクで視覚マンボモデルを上回り、検出およびセグメンテーションタスクでは最先端のマンボモデルに匹敵することができません。

マンボに適したタスクは何か？

マンボフレームワークのトークンミキサーは、選択的な状態空間モデルで、4つの入力依存パラメーターを定義します。フレームワークの再帰的特性は、RNNのような状態空間モデルを因果的アテンションから区別します。隠れ状態は、歴史情報を保存する固定サイズのメモリとして見ることができます。固定サイズのメモリは、計算コストを一定に保つことを保証します。一方、因果的アテンションレイヤーは、前のトークンからのすべてのキーと値を保存し、新しいトークンが入力されるたびにキーと値を追加します。メモリサイズは、入力されるトークンの数に比例して増加します。

状態空間モデルのメモリは、因果的アテンションのメモリに比べ、損失のあるメモリです。したがって、マンボモデルは、短シーケンスを扱うタスクで因果的アテンションメカニズムに比べ、優位性を示すことができません。ただし、長シーケンスを扱うタスクでは、因果的アテンションアプローチは二次的な複雑さのために失敗します。このシナリオでは、マンボフレームワークはメモリを現在の入力と効率的に統合する能力を示し、長シーケンスをスムーズに処理することができます。

また、状態空間モデルの再帰的特性により、マンボモデルは、現在および過去のタイムステップからの情報にのみアクセスできます。このタイプのトークンミキサーは、因果モードと呼ばれます。因果モードは、オートリグレッシブ生成タスクに適しています。

視覚認識タスク、因果的トークンミキサーコード、そして非常に大きなシーケンス

視覚認識タスクは、モデルが画像全体を一度に認識できるため、トークンミキサーの制限を必要としません。因果モードは、オートリグレッシブタスクに適しています。一方、完全可視モードは、モデルがすべての入力を一度に認識できるため、理解タスクに適しています。

実験的検証と結果

MambaOutフレームワークの仮説を実験的に検証するために、Mambaブロックは、ゲート化畳み込みニューラルネットワークブロックに基づいて構築されています。MambaOutフレームワークは、Mambaブロックを積み重ね、状態空間モデルを削除することで構築されています。

MambaOutフレームワークの実験結果は、MambaOutフレームワークの仮説を支持しており、画像分類タスクで視覚マンボモデルを上回り、検出およびセグメンテーションタスクでは最先端のマンボモデルに匹敵することができません。

画像分類タスク

ImageNetは、画像分類タスクのベンチマークとして使用され、1000以上の一般的なクラス、130万以上のトレーニング画像、5万以上の検証画像で構成されています。実験で使用されるデータ増強には、ランダムリサイズクロップ、Mixup、カラージッター、ランダム消去、CutMix、Rand Augmentが含まれます。次の表は、Mambaファミリーのモデル、MambaOutモデル、および他のアテンションベースおよび畳み込みモデルのImageNetデータセットでのパフォーマンスをまとめたものです。MambaOutフレームワークは、状態空間モデルなしで構築されており、画像分類タスクで視覚マンボモデルを上回ります。

例えば、MambaOut-Smallモデルは、トップ1の精度スコアで84%以上を達成し、最も近いMambaの競合相手よりも0.4%高くなります。この結果は、画像分類タスクに状態空間モデルを導入する必要がないという最初の仮説を強く支持しています。

オブジェクト検出およびインスタンスセグメンテーションタスク

COCOは、オブジェクト検出およびインスタンスセグメンテーションタスクのベンチマークとして使用されます。MambaOutフレームワークは、一部の視覚マンボモデルを上回りますが、最先端の視覚マンボモデル、LocalVMambaやVMambaには及ばないことがわかります。MambaOutと最先端の視覚モデルの間のパフォーマンスの差は、長シーケンスの視覚タスクにおけるMambaファミリーのモデルの潜在性を強調しています。

最終的な考え

マンボファミリーのモデルは、オートリグレッシブまたは長シーケンスの特性を持つタスクに適しています。MambaOutフレームワークは、画像分類タスクにマンボは必要ないという仮説を立てています。検出およびセグメンテーションタスクはオートリグレッシブではないものの、長シーケンスの特性を示します。MambaOutフレームワークは、Mambaブロックを積み重ね、状態空間モデルを削除することで構築されています。実験結果は、MambaOutフレームワークの仮説を支持しており、画像分類タスクで視覚マンボモデルを上回り、検出およびセグメンテーションタスクでは最先端のマンボモデルに匹敵することができません。