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.

本文介绍了一种名为 Mamba-ND 的通用设计，将 Mamba 架构扩展到任意多维数据，并通过与其他替代方案的比较实证表明 Mamba-ND 在多个多维基准测试中表现出与最先进方法相竞争的性能。