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

Shufan Li, Harkanwar Singh, Aditya Grover
<EMAIL_ADDRESS>
University of California, Los Angeles

###### 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.

##  1 Introduction

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].

The design of flexible and scalable neural network architectures is
fundamental to the success of deep learning across diverse domains.
Convolutional neural networks [23] excel at handling continuous data such as
images, audio, and video. Recently, they have been surpassed by Transformers
[45], which process continuous data as a discrete sequence of patches [12].
Despite their superior performance on many tasks such as Image Classification
and Video Action Recognition, Transformer-based models struggle to scale to
larger patch sequence as they scale quadratically with respect to sequence
length. Most recently, a special kind of State Space Model (SSM) termed as
Mamba [15], has demonstrated stronger performance than Transformers while
maintaining a linear complexity. However, the impressive performance of Mamba
was shown on 1D text sequences. This leaves open the question whether Mamba
can be effectively extended to multi-dimensional data such as images, video,
or scientific datasets, which is the focus of this work.

Unlike convolution or self-attention operations, which can be computed in
parallel across the ND input data, Mamba requires a specific ordering of the
data. Determining such an order is not an easy task. Building on past work in
architecture design, we could consider many choices. One naive approach is to
flatten the data in row-major order. Intuitively, this is non-optimal because
in this setting information only flows in one direction, which is suboptimal
for modeling complex relationships over multi-dimensional data that do not
have any default ordering. Drawing inspiration from early works on Bi-
directional LSTM, another alternative is to process the sequence in two
directions at each layer and aggregate the results. We call this method Bi-
SSM. While this design, in principle, allows information exchange between two
arbitrary patches, two patches adjacent to each other spatially may have a
huge distance between them on the computation graph. A natural extension of
this method, which we call ND-SSM, is to process the input in $2D$ directions,
where $D$ is the dimension of the data, and aggregate the results.

Rather than picking orders in sequence, another possibility is to borrow
inspiration from multi-head self-attention in transformers, wherein we can
split the channels into multiple heads and let each head process an SSM in a
different direction. This design is similar to ND-SSM but differs in that it
has less computational burden within each layer.

At a high level, there is also the question of block-level design, which
specifies how Mamba layers are organized. For example, even using the vanilla
Mamba layer as a black box, one can still apply the bi-directional or
N-directional design at the block level by processing the same input through
multiple layers with different directions.

In this work, we conducted an extensive study on these possible design
choices. Surprisingly, we find that simply alternating between three fixed
row-major orderings is one of the best-performing strategies on both 2D and 3D
data. Armed with these findings, we propose Mamba-ND, a surprisingly simple
yet effective design to extend Mamba to multi-dimensional data. Mamba-ND does
not introduce any complicated changes to 1D-SSM layers. By stacking 1D-SSM
layers as black boxes and alternating the sequence order between each layer,
Mamba-ND was able to surpass Transformer-based models on various tasks,
including image classification, action recognition, and weather forecasting,
with a lower parameter count. It also maintains a linear complexity with
respect to input sequence length. We show these results in Figure 1.

In summary, our main contributions are as follows.

  * •

We proposed Mamba-ND , which extends SSM to higher dimension through simply
alternating sequence odering across layers.

  * •

Compared with Transformers, our model was able to achieve stronger performance
at lower parameter count. Compared with ViT, our model achieves +1.5 increase
in ImageNet-1K classification accuracy while reducing the number of parameters
by 20.7%. On HMDB-51 video classification task, Mamba-ND our model achieves
+0.9 increase compared with 3D Swin-Transformer while reducing the number of
parameters by 39.0%. On ERA5 5.625 degree 72 hour weather forecast tasks,
Mamba-ND achieves +0.7 gain in Anomaly Correlation Coefficient (ACC) while
reducing the number of parameters by 44.5% when compared with Cli-ViT.

  * •

We conducted extensive ablation studies on various more complicated approaches
that extend SSM to higher dimension. We find that complicated designs do not
necessarily translate into stronger performance. These results should help
future work on adapting 1D Sequence models to higher dimension.

##  2 Related Works

#### Image Modeling.

Convolutional Neural Networks (CNNs) [23, 17, 20, 39, 21, 38] have
historically been the state-of-the-art approach for image recognition tasks.
These networks learn to extract hierarchical features from images by using
interleaved convolution and pooling layers. While they have demonstrated
strong performance, the receptive field of convolution kernels only grows
linearly with the model depth. This limitation restricts the ability of CNNs
to effectively capture global context.

More recently, Vision Transformers (ViTs) [12] have emerged as a strong
candidate for vision tasks. Drawing inspiration from the field of natural
language processing, these models first divide an input image into a grid of
discrete patches and then process them as a 1D sequence. While they provide a
global receptive field through self-attention mechanisms at every layer, such
advantages come with a cost. In particular, they exhibit quadratic
computational and memory complexity with respect to input length, making them
challenging to scale to large images.

To address this limitation, several works have proposed hybrid architectures
that incorporate attention mechanisms into CNNs [47, 5, 24] or introduce
hierarchical designs into transformers [26, 25, 8, 10, 11, 29, 40, 46, 52].
Other approaches, such as PixelRNN [43], attempt to apply recurrent models,
which are inherently designed for linear complexity and global context, to
images.

There has also been considerable efforts to apply state space models to 2D
images, such as 2D-SSM[2] and S4ND [33]. Our work falls within the this
category, wherein we specifically explore ways to adapt selective state space
models to images.

#### Video Modeling.

Modeling videos has long presented a significant challenge. Earlier works [4,
41] primarily focused on extending 2D ConvNets into the 3D domain. More
recently, 3D Transformers [1, 28, 32] have demonstrated superior performance.
Some of these works add temporal attention on top of existing 2D Vision
Transformers, while others attempt to apply self-attention in a pure 3D
fashion.

The unique characteristics of the temporal dimension in video data have
inspired various video-specific high-performing designs. For example, some
works leverage extracted optical flow features [6, 14] or motion vectors [50,
51]. Others, like SlowFast [13], employ specific architectural designs to
account for the fact that pixel values change less in the temporal dimension
than in the spatial dimensions. In contrast to these approaches, our goal is
to devise a generic framework for modeling multi-dimensional data.
Consequently, Mamba-ND treats videos as simple 3D arrays of RGB pixels, making
no additional assumptions about the temporal structure. This is crucial
because in certain use cases such as weather forecasting, the third dimension
of the data represents an additional spatial dimension instead of a temporal
one.

S4ND [33] is a prior attempt to model videos using state space models (S4
[16]) and has achieved competitive performance. It directly extends the S4
formulation to higher dimensions and leverages the time invariance constraint
to parallelize computations as the outer product of three 1D convolutions. In
contrast to S4ND, we aim to extend selective state space models, i.e. Mamba,
to higher dimensions, which do not adhere to linear time invariance (LTI). As
a result, the convolutional trick is not applicable, necessitating alternative
designs.

#### Climate Modeling.

Climate modeling presents a significant challenge in the 3D domain. We
specifically focus on global weather forecasting. Typically, this task is
evaluated using the ERA5 reanalysis dataset [18]. Early works [48, 49]
employed CNN architectures. Pangu-Weather [3] was the first to introduce
Transformers to this domain. It adapted a 3D-Swin Transformer to handle
atmospheric measurements in 3D spaces. ClimaX [34] extended the Vision
Transformer to weather data and proposed an architecture called Cli-ViT. Using
this architecture, it established a strong foundational model, called ClimaX,
capable of various tasks such as forecasting and downscaling through large-
scale pretraining on simulated data.

Using state space models for climate modeling remains an unexplored area. We
are the first to apply SSMs in the 3D weather forecasting task.

##  3 Background

###  3.1 State Space Models

Recent state space models (SSMs) [16, 33] have shown superior performance on
long sequences. Formally, SSMs model the input data using the following
ordinary differential equation (ODE):

$\displaystyle h^{\prime}(t)$ $\displaystyle=Ah(t)+Bx(t)$ (1) $\displaystyle
y(t)$ $\displaystyle=Ch(t)+Dx(t)$ (2)

Here, $x(t)\in\mathbb{R}$ is a continuous input signal in the time domain, and
$y(t)\in\mathbb{R}$ is a continuous output signal in the time domain. In
modern SSMs, this ODE is approximated through discretization. One common
discretization method is the zero-order hold (ZOH) rule, which gives the
following difference equation:

$\displaystyle\bar{A}$ $\displaystyle=\exp(\Delta A)$ (3)
$\displaystyle\bar{B}$ $\displaystyle=(\Delta A)^{-1}(\exp(\Delta
A)-I)\cdot\Delta B$ (4) $\displaystyle h_{t}$
$\displaystyle=\bar{A}h_{t-1}+\bar{B}x_{t}$ (5) $\displaystyle y_{t}$
$\displaystyle=Ch_{t}$ (6)

Early works such as S4 [16] and S4ND [33] assume linear time invariance (LTI).
This constraint makes it possible to solve the above difference equation using
a global convolution kernel. Selective state space models, i.e., Mamba [15],
introduce time-varying parameters that do not follow the LTI assumption. In
particular, $\Delta$, $B$, and $C$ become functions of the input signal
$x(t)$. Consequently, from Equations 3, 4, 5 and 6, $\bar{A}$ and $\bar{B}$
also become dependent on time. This makes it hard to parallelize the
computation. A hardware-aware optimization that makes use of associative scan
is employed to mitigate this issue.

###  3.2 Mamba Layers

Mamba [15] proposed an implementation of selective state space model (sSSM)
layer (Figure 3 column 1). It consists of a 1D convolution, an SSM kernel, and
a residual connection. In each layer, the input sequence is first processed by
the convolution operation and then by the SSM kernel. The result is added back
to the input through the residual connection. While the Conv1D operation can
be easily extended to multiple dimensions with ConvND operations, it is non-
trivial to convert the SSM to multiple dimensions. S4ND [33] resolved this by
formulating higher-dimensional SSM as a time-invariant system of partial
differential equations (PDEs), which can be computed using global convolution.
However, since Mamba is not a time-invariant system, this approach is not
applicable. In this work, we discuss various alternative ways to extend SSM to
higher dimensions by flattening the input data into 1D sequences.

##  4 Methodology

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 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 tasks.

Mamba-ND explores several approaches that adapt Mamba to multidimensional
data. The key element of these designs is to devise a combination of sequence
orderings to flatten the multidimensional data into 1-D sequences.
Intuitively, some level of bidirectional or multidirectional design is
required to allow information exchange between two arbitrary data points in
multidimensional space. This can be achieved at two levels: layer level and
block level. As mentioned in Section 3.2, a Mamba layer consists of a 1D
convolution, an SSM kernel, and a residual connection. One example of layer-
level design is to pass the output of the convolution to two independent SSM
kernels and sum up the results (Figure 3, col 2). Contrary to layer-level
designs, block-level design keeps the internal design of a Mamba layer
unchanged. Instead, it tries to achieve bidirectionality on the block level.
For example, one block-level design is to alternate between each axis from one
layer to another and apply a bidirectional design for each axis (Figure 6, row
2). One possible benefit of layer-level design is that it keeps the optimized
fused kernel of a Mamba layer and the memory access pattern of the underlying
array unchanged.

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: Different ways of arranging Mamba layers. The first row visualizes
alternating-directional design. The second row visualizes bidirectional
design. The third row visualizes quad-directional design.

1

Figure 6: Visualization of various Scan-Factorization policies. Col 1: No
factorization, there is only 1 sequence. Col 2: Factorizing the 3D sequence
into $D$ 2D sequences, where $D$ is the length of a single dimension. Col 3:
Factorizing the 3D sequence into $D^{2}$ 1D sequences.

###  4.1 Scan Orderings

Consider the input N-dimensional data $X$ of shape $D_{1}\times
D_{2}\times\ldots\times D_{N}$, where $D_{i}$ is the length of data along the
$i$th dimension. Let $L=\prod_{i=1}^{N}D_{i}$ be the total sequence length.
There are a total of $L!$ possible ways of flattening $X$ into a 1D sequence.
However, we only consider a subset of these choices, which we call the scan
orderings. Formally, a scan ordering is obtained by permuting the order of
axes of the original data $X$, and flattening it into a 1D sequence either in
the forward or reverse direction. Since there are $N!$ permutations of $N$
axes and 2 possible directions, there are a total of $2N!$ scan orderings.

We denote a particular scan ordering $s$ as $(k_{1},k_{2},\ldots,k_{N})\pm$,
which represents the unique ordering obtained by first permuting the axis
order from $1,2,\ldots,N$ to $k_{1},k_{2},\ldots,k_{N}$, and then flattening
the sequence in row-major order. The symbols $+$ and $-$ indicate whether the
order of the final 1D sequence is in forward or reversed direction.

We focus on 2D data of shape $H\times W$ and 3D data of $T\times H\times W$.
The 2D data has 4 possible orderings: $(HW)+$, $(HW)-$, $(WH)+$, and $(WH)-$.
The 3D data has 12 possible orderings; examples include $(HWT)+$ and $(WHT)-$.
To provide a concrete example, the ordering $(WH)-$ refers to first permuting
the 2D data into a 2D array of $W\times H$, then flattening it into a 1D
sequence in row-major order, and finally reversing the order of this 1D
sequence.

For simplicity, we use $H$ to represent $WH$ and $W$ to represent $HW$ for 2D
data. Similarly, we use $H$ to represent $TWH$, $W$ to represent $THW$, and
$T$ for $HWT$ for 3D data. In this notation, $(WH)-$ becomes $H-$ and $(HWT)+$
becomes $T+$. We also use $L$ to represent $THW$ or $HW$, as this is the naive
way of flattening 3D and 2D data to a 1D sequence without changing the memory
layout. Notably, the last dimension will be traversed continuously. We
visualize 2D scan orderings in Figure 4.

###  4.2 Adapting the Mamba Layer

We explore three possible alterations to the standard Mamba layer design,
which are illustrated in Figure 3.

#### Bi-SSM

layer passes the output of the convolution layer to two independent SSM
kernels, one in the forward direction and another in the reversed direction.
This design draws inspiration from bidirectional LSTM [7], in which two LSTM
layers of opposite directions are stacked together in a similar fashion.

#### ND-SSM

layer extends Bi-SSM by incorporating additional SSMs to accommodate other
possible orderings. In the 2D case, there are four orderings $W+$, $W-$, $H+$,
$H-$.

#### Multi-head SSM

layer is a mimic of the multi-head design in the Transformer. It splits an
input sequence of dimension $D$ into $H$ sequences of dimension $D/H$, where
$H$ is the number of orderings. Each of the heads is then passed to separate
SSM kernels in respective orderings. In the 2D case, the orderings are $W+$,
$W-$, $H+$, $H-$.

###  4.3 Arranging Mamba Layers

In addition to making direct changes to the internal structure of Mamba, one
can also change the way in which the layers are organized to achieve multi-
directionality. We illustrate these variations in Figure 6.

#### Alternating-Directional: H+H-W+W-T+T-

keeps the sequential ordering of Mamba layers and changes the direction of SSM
in each layer in an alternating fashion. The ordering is H+H-W+W-T+T-

#### Bi-Directional: [H+H-][W+W-][T+T-]

adopts a design on the block level. In each block, the input is passed to two
Mamba layers at opposite directions. The ordering is [H+H-][W+W-][T+T-], where
each $[.]$ denotes a bidirectional block consisting of two layers. To avoid
confusion, we will explicitly refer to this method as [H+H-][W+W-][T+T-]. The
term Bi-Directional will mostly be used for the Bi-SSM layer mentioned in
Section 4.2.

#### Quad-Directional: [H+H-W+W-][T+T-]

builds on top of the [H+H-][W+W-][T+T-] design. It further groups the H and W
directions. This design is inspired by works in Video Recognition such as
[42], which factorize 3D convolution into a 2D operation on the spatial
dimensions and a 1D operation in the temporal domain.

There are more possible ways to organize multi-directional blocks, but they
generally follow a similar design. Crucially, it is important to note that
while each layer has a specific ordering, the SSM kernel operates on a single
flattened input sequence. This means that all these layers have a global
receptive field.

###  4.4 Scan Factorization

In order to mitigate the quadratic complexity of the Transformer, prior works
[19] factorize full 3D attention into three 1D attentions along each axis.
While SSMs already achieve linear complexity, the sequence length is still
quadratic in the length of a single dimension. In this work, we also explore
various ways of factorizing an SSM scan into multiple smaller scans. For an
input array of dimensions $T\times H\times W$, the standard approach is to
flatten it into a single sequence of length $THW$. Alternatively, we can
factorize it into $T$ sequences with length $HW$, or $TH$ sequences with
length $W$. We visualize these factorization techniques in Figure 6. Since SSM
only retains one copy of a state per sequence in the linear scan process,
increasing the number of sequences in parallel actually leads to an increase
in memory consumption and training time because more hidden states need to be
materialized in the GPU memory. However, we note that these sub-sequences need
not be computed in parallel. Hence a better implementation in the future may
reduce the memory cost of this design. Thus, we still investigate the
performance of this design.

###  4.5 Final Design

After extensive experiments, our final design uses the standard 1D-SSM layer
and an alternating-directional: (H+H-W+W-T+T-) at the block level. We find
that this simple design surprisingly outperforms more complicated ones. We
provide more details of our experiments in Section 5.

Our overall image is shown in Figure 2. Given multi-dimensional input data, we
first patchify it into a 1D sequence. During this process, we keep track of
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 SSM layers that
process the sequence in alternating orderings. In 2D space, we use the order
H+H-W+W-. In 3D space, we use the order H+H-W+W-T+T-. Finally, the sequence is
reshaped back to its original 3D structure and passed to task-specific heads
for downstream tasks.

##  5 Experiments

#### Settings

We aim to evaluate the effectiveness of Mamba-ND on various multi-dimensional
data tasks. Specifically, we use ImageNet-1K [9] for image classification,
HMDB-51 [22] for action recognition, and ERA5 5.625-degree weather forecasting
[18] for 3D tasks. ImageNet-1K is a large-scale dataset containing 1.2 million
images across 1000 classes, HMDB-51 is an action recognition dataset
comprising 7,030 video clips categorized into 51 action classes, and ERA5
consists of 3D atmospheric weather measurements, such as temperature and wind
speed, across 13 pressure levels. We use the standard train and validation
split for ImageNet, split1 of HMDB-51, and data from the years 1979-2016 for
ERA5 as the training set, data from 2017 as the validation set, and data from
2018 as the test set.

#### Metrics

We measure top-1 accuracy for image classification and action recognition
tasks. For climate forecasting, we report both the Residual Mean Squared Error
(RMSE) and the Anomaly Correlation Coefficient (ACC), which is calculated as:

$ACC=\frac{\sum(f-c)(a-c)}{\sqrt{\sum(f-c)^{2}\sum(a-c)^{2}}}$ (7)

Here, $f$ represents the forecast, $c$ is the climatology (long time average
measurement), and $a$ is the prediction target.

###  5.1 Image Classification

Following the approach of previous studies [12, 33, 17], Mamba-ND is trained
on the ImageNet-1K dataset for 300 epochs using the AdamW optimizer with
$\beta=(0.9,0.999)$ and a learning rate of $1e-3$. The results are presented
in Table 1. Our model demonstrates superior performance compared to
transformer-based models when operating under similar conditions, and it
achieves results on par with the state-of-the-art state-space model, S4ND
[12]. We compare our results with Hyena [37], ViT [12], and S4ND [33].
Notably, Mamba-ND-S shows a remarkable improvement of $+1.5$ in accuracy when
compared to ViT, while simultaneously reducing the parameter count by 20.7%.
This performance gap is consistent with prior research on 1D sequences [15],
where Mamba consistently outperforms transformers with fewer parameters.

Table 1: ImageNet 1K Classification Results. We report Top 1 Accuracy on the validation set. Mamba-ND-S shows a remarkable improvement of $+1.5$ in accuracy when compared to ViT-B while reducing the parameter count by 20.7% Model | Image Size | Params | Acc.$\uparrow$
---|---|---|---
ViT-B | 384 | 86M | 77.9
ViT-L | 384 | 307M | 76.5
S4ND-ViT-B | 224 | 86M | 80.4
Hynea-ViT-B | 224 | 88M | 78.5
Mamba-2D-T | 224 | 24M | 79.2
Mamba-2D-S | 224 | 63M | 79.4
Mamba-2D-B | 224 | 92M | 80.2

###  5.2 Video Action Recognition

Prior works [4, 28] demonstrate strong performance on video datasets by
adapting ImageNet-pretrained vision models to 3D tasks. Following their
strategies, we inflate Mamba-2D to Mamba-3D. Since we adopted an alternating
design, each layer only sees a 1D input sequence. This means the sizes of
weights in Mamba-2D and Mamba-3D layers are identical, and we can directly
load the Mamba-2D checkpoint. The only weight that have different shapes are
the patch embeddings. We adopted a temporal patch size of 2, so we duplicated
the ImageNet weights along the time dimension and divide the value by 2. To
address the change in ordering from H+H-W+W- to H+H-W+W-T+T-, we simply append
new layers for T+T- every other four layers.

We sample 32 frames at a frame interval of 2 from each video clip, which
amounts to around 2 seconds of video. We use the AdamW optimizer with
$\beta=(0.9,0.999)$ and a learning rate of $1e-3$. The learning rate on the
backbone is multiplied by a factor of 0.1. We train our model with a global
batch size of 64 for 50 epochs. We select Inception-I3D [4] and ConvNeXt-I3D
[27] as our convolutional baseline. For transformer-based models, we select
Video Swin Transformer [28].

We report the results in Table 2. Our method achieves a $+0.9$ accuracy
compared to Video Swin Transformer while using 39% fewer parameters. Compared
to the state-of-the-art ConvNet ConvNeXt, we achieve a $+5.5$ accuracy.

Table 2: HMDB-51 Video Classification Results. All models are initialized with ImageNet weights. *: Numbers from S4ND[33] paper. ${\dagger}$: Our reproduced numbers. Memory: Training memory measured in GB on a A100 GPU. All models are trained with a batch size of 16 per GPU, except S4ND, which has a batch size of 8 (OOM at 16). We also report the samples per second.  Model | Acc.$\uparrow$ | Params | Memory | Samples/s
---|---|---|---|---
ConvNeXt-I3D* | 58.1 | 29M | - | -
S4ND-ConvNeXt-3D* | 62.1 | 29M | - | -
Inception-I3D | 49.8 | 25M | - | -
ConvNeXt-I3D${\dagger}$ | 53.5 | 29M | 17GB | 7.5
S4ND-ConvNeXt-3D${\dagger}$ | 55.2 | 29M | 77GB | 6.3
Video-Swin-T | 53.0 | 30M | 35GB | 22.6
Video-Swin-S | 58.1 | 54M | 73GB | 39.2
Mamba-2D | 51.2 | 24M | 12GB | 39.2
Mamba-3D | 59.0 | 36M | 17GB | 19.6

###  5.3 Global Weather Forecast

For weather forecasting tasks, we use Cli-ViT [34, 35, 36] and Pangu-Weather
[3] as our baselines. Because these models were originally trained on
terabytes of high-resolution data, typically around 1 degree, for a long time,
we cannot directly compare them due to computational constraints. Instead, we
use the 5.625-degree version of the ERA data across 7 pressure levels. Our
baselines Cli-ViT are based on Vision Transformer, and Pangu-Weather is based
on Swin-Transformer. For Cli-ViT, we use the official implementation from
ClimaX [34]. For Pangu-Weather, we re-implement it in PyTorch. We train all
models from scratch for 50 epochs, using a learning rate of 0.0005. The
learning rate follows a cosine decay schedule with a warm-up period of 5
epochs. For Mamba-3D, we use a patch size of $2\times 2\times 2$.

We show the results in Table 3. Compared to Cli-ViT, Mamba-3D achieved a
$+0.7$ ACC while reducing the parameter count by 44.5%.

Table 3: ERA5 5.625 Degree Weather Forecast Results. We report the Anomaly Correlation Coefficient (ACC) and RMSE on geopotential at the 500 hPa pressure level. Compared to Cli-ViT, Mamba-3D achieved a $+0.7$ ACC while reducing the parameter count by 44.5%.  | Arch | Parms | RMSE$\downarrow$ | ACC$\uparrow$
---|---|---|---|---
Cli-ViT | ViT | 108M | 467 | 89.3
Pangu | Swin | 50M | 462 | 89.0
Mamba-3D | SSM | 50M | 433 | 90.1

###  5.4 Meta Architectures

We perform extensive ablation studies on various design choices mentioned in
Section 4. We show that the alternating-directional design is the simplest and
most effective one among a wide range of possible choices.

#### Layer Design

We use ImageNet-1K and HMDB-51 classification as our benchmarks. We perform an
ablation study on Mamba-2D and Mamba-3D. We adopted various designs mentioned
in Section 4.2. Results show that our final alternating-directional design
achieves stronger performance than all proposed layer-wise changes. We show
results in Table 4. Compared with the 1D-Mamba baseline, in which no special
design is adopted after naively flattening the multi-dimensional sequence into
a 1D sequence, we achieved a $+4.8$ accuracy on ImageNet and a $+26.9$
accuracy on HMDB-51.

Additionally, we observe that most of the proposed multi-directional designs
are able to outperform the naive and bi-directional baselines by considerable
margins, suggesting the importance of multi-directional design.

Table 4: Ablation Study on Layer Designs. We report top-1 accuracy on the ImageNet-1K validation set. The Alt-Directional design is the top-performing one.  |  | IN1K$\uparrow$ | HMDB-51 $\uparrow$
---|---|---|---
Alt-Directional | Block Level | 79.4 | 59.0
Multi-Head-SSM | Layer-Level | 77.6 | 51.5
ND-SSM | Layer-Level | 77.2 | 46.7
1D-SSM | - | 76.4 | 34.9
Bi-SSM | Layer-Level | 74.6 | 32.1

#### Layer arrangement

We perform comprehensive ablation studies on various organizations of SSM
layers mentioned in Section 4.3 at the block level. In addition to three
choices discussed in Section 4.3, we also experiment with a hex-directional
design [H+H-W+W-T+T-], in which inputs of a block are processed by multiple
layers at different ordering in parallel, and the results are summed together
in the end. We show the results in Table 5. Contrary to intuition and the
layer-level results, adding multi-directional design at the block level
degrades the performance. We hypothesize that this behavior is caused by the
reduced number of model depths. Since we keep the number of layers fixed
across all designs, incorporating some level of multi-directional parallel
processing will inevitably reduce the depths of the computation graph. We
provide more discussions in the discussion section.

Table 5: Ablation Study on Layer Arrangement. We report top-1 accuracy on the HMDB-51 dataset. H+H-W+W-T+T- is the top-performing design.  Ordering | HMDB-51$\uparrow$
---|---
H+H-W+W-T+T- | 59.0
[H+H-W+W-T+T-] | 38.3
[H+H-][W+W-][T+T-] | 49.8
[H+H-W+W-][T+T-] | 47.4

###  5.5 Scan Factorization

While factorizing the sequence in the current implementation of Mamba leads to
considerable memory and runtime overhead, these costs are external to the
designs themselves and may be patched in the future. For example, a scan of N
sequences of length $L$ can be considered as a long scan of one sequence of
$NL$ with $\bar{A}$ from Equation 5 set to zero in places where the scan
operation moves from one short sequence to another. Hence, we find it useful
to study the effects of various factorization techniques despite their
inferior runtime efficiency.

We show such results in Table 6. We also report the memory cost on a single
Nvidia A100 GPU. While all factorizations lead to worse performance, we find
that 2D+1D factorization outperforms the 2D+3D setup, suggesting there may be
merits of having certain layers dedicated to processing temporal
correspondence. Because the cost of 1D factorization is high, we choose not to
explore it further at this time.

Table 6: Ablation Study on Various Factorization Policies \+ indicates layer ordering. For example, the notation 2D+3D represents a 2D factorization layer and a 3D layer with no factorization arranged sequentially. In all studies, the total number of layers is fixed. We also show the training memory cost and the number of sequences. $B$ is the batch size, and $D$ is the length of a single dimension. Factorization | HMDB-51$\uparrow$ | Mem | Num.
---|---|---|---
1D + 1D + 1D | 44.5 | 80GB | $O(BD^{2})$
2D + 1D | 55.8 | 77GB | $O(BD^{2})$
2D + 3D | 51.9 | 18GB | $O(BD)$
3D | 59.0 | 17GB | $O(B)$

###  5.6 Video Specific Ablations

When initializing new SSM layers in T+T- direction, we can adjust the scale of
$\Delta$ to control the initial temporal receptive field. We show such results
in Table 7. We find that 1.0 is the optimal value, which coincides with the
prior conclusion from S4ND [33].

Table 7: Ablation Study on the Scaling Factor of $\Delta$ for HMDB51 Experiments. We report the performance on HMDB-51 under different initializations of temporal layers (T+ and T-). $\Delta=1.0$ is the optimal choice.  $\Delta$ | HMDB-51$\uparrow$
---|---
0.1 | 58.0
0.2 | 57.8
1.0 | 59.0
5.0 | 55.3

##  6 Discussion

Figure 7: Effective Receptive Field of Various Designs. Darkness indicates the
sensitivity of the central patch of the output to each pixel of the input
image, normalized to the range of $(0,1)$. All images are 224x224. We use
ImageNet-1K pretrained checkpoints for these visualizations.

###  6.1 Effective Receptive Field (ERF)

To validate the significance of multi-directional design, we visualize the
effective receptive field (ERF) [30] of Mamba-ND in Figure 7. ERF is computed
by setting the gradient of the center patch to one and backpropagating through
the network. In particular, we visualize the ERF of ImageNet-1K pretrained
weights of Mamba-ND against the Bi-directional and Uni-directional baseline
(Table 4, row 5 and Table 4, row 4). Sensitivity is measured by the whiteness
of a pixel in the visualization. The Uni-Directional model exhibits a sharp
cutoff and is insensitive to all patches after it in the flattened sequence.
This is clearly undesirable in data where causal relations do not exist. The
Bi-Directional model exhibits a global receptive field, but it is heavily
biased in the horizontal direction. This is also undesirable because images
can have vertical structures. The Multi-Directional model exhibits a more
uniform pattern in the sensitivity visualization, which explains their
superior performance.

###  6.2 Depths versus Widths

In Section 5.4, we showed that a complicated multi-directional design may not
necessarily lead to improved performance over the alternating-directional
baseline. We hypothesize that this result is related to the effective depth of
the model.

Each of the block-level designs can be represented by a directed acyclic graph
(DAG) of layers, where the direction of edges implies the computation order.
In this interpretation, designs other than H+H-W+W-T+T- are trading tree
depths for widths. Existing literature [44] shows that depth is more important
than width in deep neural networks. Our results seem to reaffirm this
conclusion.

##  7 Conclusion

Transformers have been the go-to choice in image, language, and video tasks in
recent years, just like ConvNets and RNNs in earlier years. Mamba [31]
presents itself as a competitive challenger to Transformers in 1D sequence
modeling. However, due to their intricate formulations discussed in Section 4,
it is not straightforward whether they will be equally competitive in multi-
dimensional data. In this work, we proposed Mamba-ND which successfully
extends the strong performance of Mamba to multi-dimensional space. It was
able to outperform Transformers on a variety of tasks with significantly fewer
parameters, while reducing the quadratic complexity to a linear one. Through
extensive experiments with alternative designs, we demonstrate the importance
of multi-directional design in dealing with multi-directional data and
highlight its advantages over uni-directional and bi-directional baselines. We
hope these results will facilitate future research in multi-dimensional
modeling.

##  8 Acknowledgement

We are grateful to Microsoft Research for supporting this research through
their Accelerate Foundation Models Research program.

## References

  * [1] H. Akbari, L. Yuan, R. Qian, W.-H. Chuang, S.-F. Chang, Y. Cui, and B. Gong.  Vatt: Transformers for multimodal self-supervised learning from raw video, audio and text.  Advances in Neural Information Processing Systems, 34:24206–24221, 2021.
  * [2] E. Baron, I. Zimerman, and L. Wolf.  2-d ssm: A general spatial layer for visual transformers.  arXiv preprint arXiv:2306.06635, 2023.
  * [3] K. Bi, L. Xie, H. Zhang, X. Chen, X. Gu, and Q. Tian.  Pangu-weather: A 3d high-resolution model for fast and accurate global weather forecast.  arXiv preprint arXiv:2211.02556, 2022.
  * [4] J. Carreira and A. Zisserman.  Quo vadis, action recognition? a new model and the kinetics dataset.  In proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 6299–6308, 2017.
  * [5] L. Chen, H. Zhang, J. Xiao, L. Nie, J. Shao, W. Liu, and T.-S. Chua.  Sca-cnn: Spatial and channel-wise attention in convolutional networks for image captioning.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 5659–5667, 2017.
  * [6] R. Christoph and F. A. Pinz.  Spatiotemporal residual networks for video action recognition.  Advances in neural information processing systems, 2, 2016.
  * [7] Z. Cui, R. Ke, Z. Pu, and Y. Wang.  Deep bidirectional and unidirectional lstm recurrent neural network for network-wide traffic speed prediction.  arXiv preprint arXiv:1801.02143, 2018.
  * [8] Z. Dai, H. Liu, Q. V. Le, and M. Tan.  Coatnet: Marrying convolution and attention for all data sizes.  Advances in neural information processing systems, 34:3965–3977, 2021.
  * [9] J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei.  ImageNet: A large-scale hierarchical image database.  In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 248–255. Ieee, June 2009.
  * [10] M. Ding, B. Xiao, N. Codella, P. Luo, J. Wang, and L. Yuan.  Davit: Dual attention vision transformers.  In European Conference on Computer Vision, pages 74–92. Springer, 2022.
  * [11] X. Dong, J. Bao, D. Chen, W. Zhang, N. Yu, L. Yuan, D. Chen, and B. Guo.  Cswin transformer: A general vision transformer backbone with cross-shaped windows.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12124–12134, 2022.
  * [12] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  arXiv preprint arXiv:2010.11929, 2020.
  * [13] C. Feichtenhofer, H. Fan, J. Malik, and K. He.  Slowfast networks for video recognition.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 6202–6211, 2019.
  * [14] C. Feichtenhofer, A. Pinz, and A. Zisserman.  Convolutional two-stream network fusion for video action recognition.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1933–1941, 2016.
  * [15] A. Gu and T. Dao.  Mamba: Linear-Time Sequence Modeling with Selective State Spaces, Dec. 2023.
  * [16] A. Gu, K. Goel, and C. Ré.  Efficiently modeling long sequences with structured state spaces.  In The International Conference on Learning Representations (ICLR), 2022.
  * [17] K. He, X. Zhang, S. Ren, and J. Sun.  Deep residual learning for image recognition.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016.
  * [18] H. Hersbach, B. Bell, P. Berrisford, G. Biavati, A. Horányi, J. Muñoz Sabater, J. Nicolas, C. Peubey, R. Radu, I. Rozum, et al.  Era5 hourly data on single levels from 1979 to present.  Copernicus climate change service (c3s) climate data store (cds), 10(10.24381), 2018.
  * [19] J. Ho, N. Kalchbrenner, D. Weissenborn, and T. Salimans.  Axial attention in multidimensional transformers.  arXiv preprint arXiv:1912.12180, 2019.
  * [20] G. Huang, Z. Liu, L. Van Der Maaten, and K. Q. Weinberger.  Densely connected convolutional networks.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4700–4708, 2017.
  * [21] A. Krizhevsky, I. Sutskever, and G. E. Hinton.  Imagenet classification with deep convolutional neural networks.  Advances in neural information processing systems, 25, 2012.
  * [22] H. Kuehne, H. Jhuang, E. Garrote, T. Poggio, and T. Serre.  Hmdb: a large video database for human motion recognition.  In 2011 International conference on computer vision, pages 2556–2563. IEEE, 2011.
  * [23] Y. LeCun, B. Boser, J. Denker, D. Henderson, R. Howard, W. Hubbard, and L. Jackel.  Handwritten Digit Recognition with a Back-Propagation Network.  In Advances in Neural Information Processing Systems, volume 2. Morgan-Kaufmann, 1989.
  * [24] J. Li, K. Jin, D. Zhou, N. Kubota, and Z. Ju.  Attention mechanism-based cnn for facial expression recognition.  Neurocomputing, 411:340–350, 2020.
  * [25] Y. Li, C.-Y. Wu, H. Fan, K. Mangalam, B. Xiong, J. Malik, and C. Feichtenhofer.  Mvitv2: Improved multiscale vision transformers for classification and detection.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 4804–4814, 2022.
  * [26] Z. Liu, Y. Lin, Y. Cao, H. Hu, Y. Wei, Z. Zhang, S. Lin, and B. Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 10012–10022, 2021.
  * [27] Z. Liu, H. Mao, C.-Y. Wu, C. Feichtenhofer, T. Darrell, and S. Xie.  A convnet for the 2020s.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 11976–11986, 2022.
  * [28] Z. Liu, J. Ning, Y. Cao, Y. Wei, Z. Zhang, S. Lin, and H. Hu.  Video swin transformer.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 3202–3211, 2022.
  * [29] J. Lu, R. Mottaghi, A. Kembhavi, et al.  Container: Context aggregation networks.  Advances in neural information processing systems, 34:19160–19171, 2021.
  * [30] W. Luo, Y. Li, R. Urtasun, and R. Zemel.  Understanding the effective receptive field in deep convolutional neural networks.  Advances in neural information processing systems, 29, 2016.
  * [31] K. Mangalam, H. Fan, Y. Li, C.-Y. Wu, B. Xiong, C. Feichtenhofer, and J. Malik.  Reversible Vision Transformers.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10830–10840, 2022.
  * [32] D. Neimark, O. Bar, M. Zohar, and D. Asselmann.  Video transformer network.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 3163–3172, 2021.
  * [33] E. Nguyen, K. Goel, A. Gu, G. W. Downs, P. Shah, T. Dao, S. A. Baccus, and C. Ré.  S4nd: Modeling images and videos as multidimensional signals using state spaces.  Advances in Neural Information Processing Systems, 35, 2022.
  * [34] T. Nguyen, J. Brandstetter, A. Kapoor, J. K. Gupta, and A. Grover.  Climax: A foundation model for weather and climate.  arXiv preprint arXiv:2301.10343, 2023.
  * [35] T. Nguyen, J. Jewik, H. Bansal, P. Sharma, and A. Grover.  Climatelearn: Benchmarking machine learning for weather and climate modeling.  arXiv preprint arXiv:2307.01909, 2023.
  * [36] T. Nguyen, R. Shah, H. Bansal, T. Arcomano, S. Madireddy, R. Maulik, V. Kotamarthi, I. Foster, and A. Grover.  Scaling transformer neural networks for skillful and reliable medium-range weather forecasting.  arXiv preprint arXiv:2312.03876, 2023.
  * [37] M. Poli, S. Massaroli, E. Nguyen, D. Y. Fu, T. Dao, S. Baccus, Y. Bengio, S. Ermon, and C. Ré.  Hyena hierarchy: Towards larger convolutional language models.  arXiv preprint arXiv:2302.10866, 2023.
  * [38] K. Simonyan and A. Zisserman.  Very deep convolutional networks for large-scale image recognition.  arXiv preprint arXiv:1409.1556, 2014.
  * [39] C. Szegedy, W. Liu, Y. Jia, P. Sermanet, S. Reed, D. Anguelov, D. Erhan, V. Vanhoucke, and A. Rabinovich.  Going deeper with convolutions.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1–9, 2015.
  * [40] Y. Tian, L. Xie, Z. Wang, L. Wei, X. Zhang, J. Jiao, Y. Wang, Q. Tian, and Q. Ye.  Integrally pre-trained transformer pyramid networks.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 18610–18620, 2023.
  * [41] D. Tran, L. Bourdev, R. Fergus, L. Torresani, and M. Paluri.  Learning spatiotemporal features with 3d convolutional networks.  In Proceedings of the IEEE international conference on computer vision, pages 4489–4497, 2015.
  * [42] D. Tran, H. Wang, L. Torresani, J. Ray, Y. LeCun, and M. Paluri.  A closer look at spatiotemporal convolutions for action recognition.  In Proceedings of the IEEE conference on Computer Vision and Pattern Recognition, pages 6450–6459, 2018.
  * [43] A. Van Den Oord, N. Kalchbrenner, and K. Kavukcuoglu.  Pixel recurrent neural networks.  In International conference on machine learning, pages 1747–1756. PMLR, 2016.
  * [44] G. Vardi, G. Yehudai, and O. Shamir.  Width is less important than depth in relu neural networks.  In Conference on Learning Theory, pages 1249–1281. PMLR, 2022.
  * [45] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin.  Attention is All you Need.  In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
  * [46] W. Wang, E. Xie, X. Li, D.-P. Fan, K. Song, D. Liang, T. Lu, P. Luo, and L. Shao.  Pyramid vision transformer: A versatile backbone for dense prediction without convolutions.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 568–578, 2021.
  * [47] X. Wang, R. Girshick, A. Gupta, and K. He.  Non-local neural networks.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 7794–7803, 2018.
  * [48] T. Weber, A. Corotan, B. Hutchinson, B. Kravitz, and R. Link.  Deep learning for creating surrogate models of precipitation in earth system models.  Atmospheric Chemistry and Physics, 20(4):2303–2317, 2020.
  * [49] J. A. Weyn, D. R. Durran, R. Caruana, and N. Cresswell-Clay.  Sub-seasonal forecasting with a large ensemble of deep-learning weather prediction models.  Journal of Advances in Modeling Earth Systems, 13(7):e2021MS002502, 2021.
  * [50] B. Zhang, L. Wang, Z. Wang, Y. Qiao, and H. Wang.  Real-time action recognition with enhanced motion vector cnns.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2718–2726, 2016.
  * [51] B. Zhang, L. Wang, Z. Wang, Y. Qiao, and H. Wang.  Real-time action recognition with deeply transferred motion vector cnns.  IEEE Transactions on Image Processing, 27(5):2326–2339, 2018.
  * [52] X. Zhang, Y. Tian, L. Xie, W. Huang, Q. Dai, Q. Ye, and Q. Tian.  Hivit: A simpler and more efficient design of hierarchical vision transformer.  In The Eleventh International Conference on Learning Representations, 2022.