Mamba-ND: Selective State Space Modeling for Multi-Dimensional Data

要約

近年、Transformers は、テキストおよび画像やビデオなどのさまざまな多次元データのシーケンス モデリングのための事実上のアーキテクチャになりました。
ただし、Transformer でセルフ アテンション レイヤーを使用すると、法外な計算とメモリの複雑さが発生し、それが二次関数的にスケールされます。
シーケンスの長さ。
状態空間モデルに基づく最近のアーキテクチャ Mamba は、テキスト シーケンスのモデリングで同等のパフォーマンスを達成し、シーケンスの長さに応じて線形にスケーリングすることが示されています。
この研究では、Mamba アーキテクチャを任意の多次元データに拡張する一般化された設計である Mamba-ND を紹介します。
私たちの設計では、行優先の順序に従って、さまざまな次元にわたって入力データを交互に解明します。
双方向 LSTM や S4ND などの以前の多次元拡張機能に基づいて、Mamba-ND と他のいくつかの代替案を体系的に比較します。
我々は、Mamba-ND が ImageNet-1K 分類、HMDB-51 行動認識、ERA5 天気予報などのさまざまな多次元ベンチマークにおいて、最先端のベンチマークに匹敵するパフォーマンスを実証していることを実証的に示しています。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google