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

TL;DR

Mamba-ND extends the selective state space model framework to multi-dimensional data by alternating simple row-major scan orderings across layers, achieving linear complexity with respect to input length. The approach avoids heavy architectural changes by stacking standard 1D Mamba layers and reordering data, enabling strong performance on image classification, video action recognition, weather forecasting, and 3D medical segmentation with fewer parameters than Transformers. Extensive ablation studies show the alternating-directional design as the most effective, while factorization and complex multi-directional schemes offer limited gains and can increase memory costs. The work demonstrates that a simple, directional layering strategy can deliver competitive results across diverse ND domains, highlighting the practical impact of SSM-based ND modeling for scalable, efficient sequence processing.

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.

Figures (7)

• Figure 1: Mamba-ND outperforms Transformers while significantly reducing the number of parameters. On ImageNet-1k, we compare against ViT vit. On HMDB-51kuehne2011hmdb, we compare against Video-Swin liu2022video. On ERA5, we compare against Cli-ViT. On BTCV, we compare against UNETRhatamizadeh2022unetr. nguyen2023climax.
• Figure 2: Mamba-ND Architecture. We visualize Mamba-3D as an example. Given 3D input, 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 1D Mamba layers that process the sequence in alternating orderings. In 3D space, we use the order H+H-W+W-T+T-. In 2D space, the sequence would be H+H-W+W-. Finally, the sequence is reshaped back to its original 3D structure and passed to task-specific heads for downstream processing.
• Figure 3: Variations of SSM Layer Design. Col 1 represents the standard 1D SSM layer. Col 2 represents Bi-SSM, which adds bidirectionality in a similar fashion as LSTM. Col 3 represents ND-SSM block, which extends Bi-SSM to more directions. Col 4 represents multi-head SSM block inspired by multi-head attention in Transformers.
• Figure 4: Visualization of 2D scan orderings. We visualize the set of possible all scan ordering on 2D data. Arrow indicates the scan order.
• Figure 5: Visualization of block level design and factorization policies.