Source-linked AI summary
Mamba-ND: Selective State Space Modeling for Multi-Dimensional Data
Shufan Li, Harkanwar Singh, Aditya Grover
TL;DR
Transformers scale quadratically with sequence length, while Mamba's linear-complexity results had been shown for text, leaving its extension to multidimensional data open. Mamba-ND alternates row-major orderings across 1D state-space layers, and the paper reports stronger performance than Transformers across varied tasks with fewer parameters. The authors also note that many possible ordering choices remain unexplored.
Problem
Whether Mamba can be effectively extended beyond 1D text sequences to multidimensional inputs remains open, while Transformer self-attention has quadratic complexity in sequence length.
Source: §1 · p. 1Method
Mamba-ND extends Mamba to multidimensional data by stacking 1D state-space layers and alternating row-major sequence orderings.
Source: §1 · pp. 2–3Results
Mamba-ND outperforms Transformers across varied tasks with fewer parameters, and the conclusion characterizes its complexity as subquadratic.
Source: §7 · p. 15Takeaways & Limitations
Experiments with alternative designs show the importance of multidirectional choices over unidirectional and bidirectional baselines.
Source: §7 · p. 15The study primarily explores row-major orderings; the exponential design space leaves other orderings, such as diagonal or zig-zag patterns, unexplored.
Source: §11 · p. 23
Abstract
from arXiv · showhide
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.
1 Introduction
The paper asks whether Mamba, which has linear complexity on text sequences, can be extended effectively to multidimensional data. It proposes alternating row-major sequence orderings, finding this simple design performs strongly across tasks with fewer parameters than Transformers.
1 Introduction
- Mamba's multidimensional extension requires choosing a sequence ordering, because multidimensional inputs have no default order and one-way information flow may be suboptimal.The proposed alternatives include bidirectional processing, directional SSMs, and splitting channels into heads that process different directions.Source: §1 · p. 2
- Mamba-ND surpasses Transformer-based models across diverse tasks, with a lower parameter count and linear complexity in input sequence length.The reported tasks include image classification, action recognition, weather forecasting, and 3D segmentation.Source: §1 · pp. 2–3
- Alternating fixed row-major orderings across layers is one of the best-performing strategies on both 2D and 3D data.Mamba-ND stacks unmodified 1D-SSM layers as black boxes and changes sequence order between layers.Source: §1 · pp. 2–3
- The study compares possible design choices, including more complicated extensions, and finds that complexity alone does not ensure stronger performance.The introduction identifies extensive design exploration, while its result summary states that complicated designs do not necessarily perform better.Source: §1 · pp. 2–3
2 Related Works
Related work spans convolutional and Transformer approaches to images and video, alongside state-space models. Mamba-ND targets a general multidimensional framework, including settings where the third axis is spatial rather than temporal.
Modeling 2D Data
Modeling 3D Data
- Mamba-ND treats video as a generic 3D array of RGB pixels without assumptions about temporal structure, which also suits data whose third axis is spatial.The text gives weather forecasting as an example where the third dimension represents an additional spatial dimension.Source: §2 · pp. 3–4
- Unlike S4ND, which extends time-invariant S4 through a convolutional formulation, Mamba's selective, non-time-invariant state space model requires alternative multidimensional designs.Source: §2 · p. 4
- Transformers have also succeeded in climate forecasting and medical image segmentation, while facing scaling challenges similar to those in vision.Source: §2 · p. 4
3 Background
The background describes state-space models as continuous-time systems discretized for sequence modeling, and explains why Mamba's selective parameters complicate parallelization. Mamba-ND organizes 1D layers with alternating orderings to process multidimensional inputs.
3.1 State Space Models
- State-space models represent input data with an ordinary differential equation and are used to model long sequences.Source: §3.1 · p. 4
- The continuous input x(t) and output y(t) are discretized; the zero-order hold rule yields a difference equation for sequence processing.Source: §3.1 · p. 4
- Mamba makes Δ, B, and C input-dependent, so its parameters vary over time and prevent direct use of the time-invariant global-convolution approach; associative scan helps mitigate parallelization difficulty.Source: §3.1 · p. 4
3.2 Mamba Layers
- Mamba-ND patchifies multidimensional input while retaining its structure, processes it through blocks of 1D layers in alternating orderings, then reshapes it for task-specific heads.The illustrated 3D ordering is H+H-W+W-T+T-; the 2D ordering is H+H-W+W-.Source: Figure 2 · p. 5
- A Mamba layer combines a 1D convolution, an SSM kernel, and a residual connection; extending the SSM to multiple dimensions is nontrivial because Mamba is not time-invariant.Source: §3.2 · p. 5
4 Methodology
The methodology adapts Mamba to multidimensional inputs by varying axis scan order and organizing 1D SSM layers at layer or block level. Its final design alternates directional scans while retaining standard 1D-SSM layers.
4 Methodology
- Multidirectionality can be added inside SSM layers or by arranging unchanged Mamba layers in different scan directions.Source: §4 · pp. 5–6
4.1 Scan Orderings
- A scan ordering permutes axes, then flattens in row-major order, forward or reversed; 2D H/W traversals have four possibilities.Source: §4.1 · pp. 6–7
4.2 Adapting the Mamba Layer
- Bi-SSM uses forward/reverse kernels; ND-SSM adds kernels for more orderings, while multi-head splits feature dimensions among ordering-specific SSMs.Source: §4.2 · pp. 7–8
4.3 Arranging Mamba Layers
- Block-level options alternate directions across layers or group directions into bidirectional, quad-directional, or other blocks; each layer scans a flattened sequence with global receptive field.Source: §4.3 · p. 8
4.4 Scan Factorization
- For T×H×W inputs, factorization splits scans into T sequences of HW or TH of W; more parallel sequences increase memory use and training time.A future implementation may reduce this memory cost.Source: §4.4 · pp. 8–9
4.5 Final Design
- The final design uses standard 1D-SSM layers with H+H-W+W-T+T- orderings, then reshapes processed sequences for task-specific heads.Source: §4.5 · p. 9
5 Experiments
The experiments evaluate Mamba-ND across image, video, weather, and medical-imaging tasks, then compare alternative layer arrangements and scan factorizations. Results favor the alternating-directional design, while factorized scans incur performance and implementation costs.
5 Experiments
- The evaluation covers ImageNet-1K, HMDB-51, UCF-101, ERA5, and BTCV, using top-1 accuracy, RMSE and ACC, or DICE as appropriate.The setup also specifies dataset splits and training data periods for ImageNet, HMDB-51, and ERA5.Source: §5 · p. 9
5.1 Image Classification
- +3.8 Top-1 accuracy versus ViT-B, with Mamba-ND-S using 20.7% of ViT-B’s parameters on ImageNet-1K.Mamba-ND also achieves results on par with S4ND.Source: Table 1 · p. 10, §5.1 · p. 10
5.2 Video Action Recognition
- +1.3 accuracy versus Video Swin Transformer on Kinetics-400, using 44% of its parameters; Mamba-3D also outperforms prior models on Breakfast classification.The video experiments use ImageNet-pretrained weights; the Kinetics-400 results include 32-frame and 64-frame settings.Source: Table 3 · p. 11, §5.2 · p. 11
5.3 Global Weather Forecasting
- On ERA5 5.625° weather forecasting, Mamba-3D achieved +0.7 ACC versus Cli-ViT while reducing parameter count by 44.5%, for geopotential at the 500 hPa level.The table reports ACC and RMSE.Source: Table 4 · p. 12
5.4 3D Medical Image Segmentation
- +0.1 DICE versus UNETR with Mamba-3D-S at 64% fewer parameters; Mamba-3D-B gains +2.7 DICE at a similar parameter count.The reported BTCV comparison uses split0 results, with six validation and 24 training samples.Source: Table 5 · p. 12, §5.4 · p. 12
5.5 Meta Architectures
- +4.8 ImageNet and +26.9 HMDB-51 accuracy versus 1D-Mamba accompany the top-performing alternating layer design; block-level multidirectional processing instead degrades performance.The authors hypothesize that fixed layer counts reduce computation-graph depth in block-level designs.Source: §5.5 · p. 13, Table 6 · p. 13
5.6 Scan Factorization
- All tested scan factorizations perform worse; 2D+1D outperforms 2D+3D, while high 1D-factorization cost leads the authors not to explore it further.The current factorization implementation also adds memory and runtime overhead that the authors describe as external to the designs.Source: §5.5 · p. 13, §5.6 · p. 14
6 Discussion
The ERF comparison shows more uniform input sensitivity for the multidirectional model than for unidirectional or bidirectional baselines. The discussion links weaker block-level designs to reduced depth when layer counts are fixed.
6.1 Effective Receptive Field (ERF)
- The multidirectional model has a more uniform sensitivity pattern, unlike the unidirectional model’s sharp cutoff or the bidirectional model’s horizontal bias.The authors connect this pattern with the multidirectional model’s superior performance.Source: §6.1 · p. 14
6.2 Depths versus Widths
- Block-level designs other than H+H-W+W-T+T- trade depth for width; the authors’ results appear to support depth’s greater importance.They hypothesize that reduced effective depth helps explain why more complicated multidirectional designs may not improve performance.Source: §6.2 · p. 15
7 Conclusion
Mamba-ND extends Mamba to multi-dimensional inputs with performance that exceeds Transformers across several tasks, using fewer parameters and subquadratic complexity.
7 Conclusion
- Mamba-ND outperforms Transformers on varied tasks with fewer parameters and subquadratic complexity, and its multidirectional designs outperform unidirectional and bidirectional baselines.The conclusion attributes this to extending Mamba’s performance to multidimensional inputs.Source: §7 · p. 15
9 Additional Ablation Studies
These ablations examine directional ordering, temporal-layer initialization, positional-weight inflation, patch size, and scan ordering across image, video, and 3D segmentation settings. Alternating scan directions perform strongly, while tested ordering changes show no meaningful difference.
9.1 Additional ImageNet and HMDB-51 Results.
- Alternating Bi-SSM directions across height and width improve ImageNet-1K by +1.5 and HMDB-51 by +8.4.The default Bi-SSM uses both directions along height in each layer; the alternative alternates height and width directions across layers.Source: §9.1 · p. 20
- HMDB-51 models trained from scratch converge slowly or diverge, and the authors report no useful results.The reported HMDB-51 training setup uses ImageNet-pretrained weights.Source: §9.1 · p. 20
9.2 Video-Specific Ablations 9.3 ∆Initialization
- A ∆ scale of 1.0 is optimal for initializing temporal layers on HMDB-51, matching the prior conclusion from S4ND.The scale controls the initial temporal receptive field of newly initialized T+ and T− SSM layers.Source: §9.3 · p. 21, Table 9 · p. 21
9.4 Weight Inflation
- The two tested 2D-to-3D positional-embedding inflation policies show no significant performance difference.The alternatives either repeat and scale the 2D embedding across time or place it once at the temporal midpoint and initialize other embeddings to zero.Source: §9.4 · p. 21
9.5 Ablation on 3D Segmentation
- Alternating scan directions are the top-performing layer design on BTCV segmentation, while non-bidirectional 1D-SSM also performs strongly.The ablation compares layer designs using ImageNet-1K and HMDB-51 top-1 accuracy and BTCV Dice score.Source: §9.5 · p. 21, Table 10 · p. 22
9.6 Patch Size
- Lower patch size achieves better ImageNet-1K performance, aligning with the stated fixed-compute motivation for smaller patches or higher input resolution.The paper identifies global context at linear complexity as allowing smaller patches or higher-resolution inputs under fixed computation.Source: §9.6 · p. 21, Table 11 · p. 22
9.7 Ordering
- Two BTCV ordering changes showed no meaningful difference within a 48-choice design space that the authors could not exhaust.The authors hypothesize that residual connections mitigate the effect of different orderings.Source: §9.7 · p. 22
10 Implementation Details
The implementation details describe model sizes, parameter-count comparisons, and dataset-specific hyperparameters. The authors also note unexplored scan patterns and distinguish Mamba-ND’s general design focus from concurrent application-specific approaches.
10.1 Model Size
- Two Mamba layers have roughly the parameter count of one ViT block, and Fig. 7 compares computational cost as sequence length increases.The figure compares a 12-layer ViT-B with 12- and 24-layer Mamba-2D models; 142 denotes a 224 × 224 image with patch size 16.Source: §10.1 · p. 22, Figure 7 · p. 23
10.2 Additional Hyperparameters
- Hyperparameters vary by task: ImageNet-1K, video, BTCV, and ERA5 use different learning rates, while image and video ablations fix patch size at 16.The passage also specifies task-specific augmentation and regularization settings.Source: §10.2 · p. 23
11 Limitations
- The study primarily explores row-major orderings; diagonal and zig-zag patterns remain possible alternatives for future work.The authors cite the exponential number of possible orderings as leaving areas unexplored.Source: §11 · p. 23
12 Concurrent Works
- Unlike application-specific concurrent Mamba works, Mamba-ND explores a general recipe for adapting Mamba-Kernel to multidimensional data.The passage characterizes VisionMamba as a Bi-SSM variant and VMamba as a hierarchical ND-SSM variant, each with additional designs.Source: §12 · p. 24