[論文レビュー] Mamba-ND: Selective State Space Modeling for Multi-Dimensional Data

Shufan Li, Harkanwar Singh+1 more|arXiv (Cornell University)|Feb 8, 2024

Neural Networks and Applications被引用数 4

ひとこと要約

Mamba-ND は、Mamba 状態空間モデルを多次元データへ拡張するために、レイヤー間で行優先順序を交互に切り替えることで、画像、動画、気象予測タスクで最先端の性能を達成した。パラメータ数を著しく削減し、線形計算量を維持しており、ViT や Swin-Transformer などのトランスフォーマーを凌駕するとともに、モデルサイズを最大 44.5% 減少させた。

ABSTRACT

In recent years, Transformers have become the de-facto architecture for sequence modeling on text and a variety of multi-dimensional data, such as images and video. However, the use of self-attention layers in a Transformer incurs prohibitive compute and memory complexity that scales quadratically w.r.t. the sequence length. A recent architecture, Mamba, based on state space models has been shown to achieve comparable performance for modeling text sequences, while scaling linearly with the sequence length. In this work, we present Mamba-ND, a generalized design extending the Mamba architecture to arbitrary multi-dimensional data. Our design alternatively unravels the input data across different dimensions following row-major orderings. We provide a systematic comparison of Mamba-ND with several other alternatives, based on prior multi-dimensional extensions such as Bi-directional LSTMs and S4ND. Empirically, we show that Mamba-ND demonstrates performance competitive with the state-of-the-art on a variety of multi-dimensional benchmarks, including ImageNet-1K classification, HMDB-51 action recognition, and ERA5 weather forecasting.

研究の動機と目的

• 1D シーケンスを想定して設計された Mamba 状態空間モデルを、画像、動画、科学的データセットなどの多次元データへ拡張すること。
• Mamba における逐次処理の課題、特に多次元入力に適用する際に特定のデータ順序が必要な点を解決すること。
• 双方向、N 方向、マルチヘッド設計を含む、SSM の多次元拡張のさまざまな手法を評価・比較すること。
• 複雑な因子分解やアーキテクチャを導入せずに、1D SSM を高次元へ効果的かつ効率的に拡張する最適なアーキテクチャを同定すること。
• 1D SSM を高次元へ拡張する際に、単純なレイヤー間での順序切り替え戦略が、複雑な多方向設計を上回る性能と効率を達成できることを示すこと。

提案手法

• Mamba-ND は 1D Mamba レイヤーをブラックボックスとして適用し、連続するレイヤー間で交互に行優先順序で並べ替えた多次元データを処理する。
• モデルはレイヤー間でシーケンス順序を切り替える——例として行優先、列優先、対角線に近い順序——これにより、下流の SSM 機構を変更せずに多方向の情報伝達を可能にする。
• Mamba の選択的状態空間メカニズムを活用することで、シーケンス長に対して線形の計算量を維持し、自己注意機構の二次的計算量を回避する。
• 因子分解や特別な 2D/3D SSM レイヤーを必要とせず、代わりに各 1D SSM レイヤーの前に入力を再順序付けすることで実現する。
• 各ブロックが異なる順序を適用するブロックレベルの設計を採用し、複数の方向にわたる空間的・時間的依存関係を捉える。
• 入力次元に依存しない汎用性を持ち、データをパッチのシーケンスとして扱うことで、2D（画像）、3D（動画）、さらには 4D（気象データ）へも一般化可能である。

Figure 1: Mamba-ND outperforms Transformers while significantly reducing the number of parameters. On ImageNet-1k, we compare against ViT [ 12 ] . On HMDB-51 [ 22 ] , we compare against Video-Swin [ 28 ] . On ERA5, we compare against Cli-ViT [ 34 ] .

実験結果

リサーチクエスチョン

• RQ1 元々 1D テキストシーケンスを想定して設計された Mamba アーキテクチャを、画像や動画などの多次元データへ効果的に拡張できるか？
• RQ2 レイヤー間で単純な順序切り替え戦略を採用することで、より複雑な多方向またはマルチヘッド SSM 設計を上回る性能が得られるか？
• RQ3 Mamba-ND の有効な受容 field は、一方向および双方向ベースラインと比較して、空間的・時間的依存関係を捉える上でどのように異なるか？
• RQ4 因子分解、レイヤー順序、深さといったアーキテクチャ的選択が、性能およびパラメータ効率に与える影響は何か？
• RQ5 単にレイヤー間で入力シーケンスを再順序付けするという最小限の変更のみで、最先端のトランスフォーマーと同等の性能が達成できるか？

主な発見

• ImageNet-1K では、ViT よりもトップ-1 精度が 1.5% 高く、パラメータ数を 20.7% 減少させた。
• HMDB-51 動作認識ベンチマークでは、3D Swin-Transformer よりも 0.9% 高い精度を達成し、パラメータ数を 39.0% 減少させた。
• ERA5 72 時間気象予測タスクでは、Cli-ViT よりも異常相関係数（ACC）で 0.7 ポints 高く、パラメータ数を 44.5% 減少させた。
• より複雑な設計（2D+3D 因子分解、マルチヘッド SSM、双方向 SSM）に比べて、順序切り替え戦略が優れた性能を示した。
• 有効な受容 field の可視化から、Mamba-ND は一方向または双方向モデルと比較して、より均一でグローバルな感度パターンを示しており、優れた性能の背景にあることが明らかになった。
• アブレーションスタディの結果、より深いまたは広い多方向アーキテクチャよりも、単純な順序切り替え設計がより効果的であることが確認され、SSM において深さが幅よりも重要であることが示された。

Figure 2: Architecture of Mamba-ND: In this figure we visualize Mamba-3D. Given 3D input data, we patchify it into $L$ patches. During this process, we maintain the original 3D structure of the input. This sequence is then passed through $K$ Mamba-ND blocks, each of which consists of a chain of 1-D