V2M: Visual 2-Dimensional Mamba for Image Representation Learning

## Abstract

Abstract Mamba has garnered widespread attention due to its flexible design and efficient hardware performance to process 1D sequences based on the state space model (SSM).
Recent studies have attempted to apply Mamba to the visual domain by flattening 2D images into patches and then regarding them as a 1D sequence.
To compensate for the 2D structure information loss (e.g., local similarity) of the original image, most existing methods focus on designing different orders to sequentially process the tokens, which could only alleviate this issue to some extent.
In this paper, we propose a Visual 2-Dimensional Mamba (V2M) model as a complete solution, which directly processes image tokens in the 2D space.
We first generalize SSM to the 2-dimensional space which generates the next state considering two adjacent states on both dimensions (e.g., columns and rows).
We then construct our V2M based on the 2-dimensional SSM formulation and incorporate Mamba to achieve hardware-efficient parallel processing.
The proposed V2M effectively incorporates the 2D locality prior yet inherits the efficiency and input-dependent scalability of Mamba.
Extensive experimental results on ImageNet classification and downstream visual tasks including object detection and instance segmentation on COCO and semantic segmentation on ADE20K demonstrate the effectiveness of our V2M compared with other visual backbones. 1 1 1 Code link: https://github.com/wangck20/V2M .

## 1 Introduction

State Space Models (SSM) have recently received sustained attention in natural language processing (NLP), achieving excellent performance on various tasks because of its efficiency in handling long sequences (Gu et al., ; ; ).
Mamba (Gu & Dao, ) further equip the conventional SSMs with the ability to process time-varying input and exquisitely designed hardware acceleration algorithms, demonstrating the potential to compete with transformers (Kim et al., ).

Although convolutional neural networks (CNNs) (He et al., ; Liu et al., ) and visual transformers (ViTs) (Dosovitskiy et al., ; Liu et al., ) have consistently dominated computer vision, recent works applying Mamba to visual data modeling emerges as promising alternatives (Zhu et al., ; Liu et al., ; Huang et al., ).
To adapt 2D visual data to Mamba with 1D sequential processing, most existing methods follow ViTs to patchify images into tokens and design various ordering strategies to incorporate 2D structural prior.
For example, Vision Mamba (Vim) (Zhu et al., ) directly flattens images by rows and then employs a bidirectional modeling strategy to enhance the latent representation of the model.
LocalMamba (Huang et al., ) attempts to preserve the local invariance of images through local scanning and optimal scanning direction search.
However, these works are still constrained in the framework of 1-dimensional Mamba, which can only approximate the 2D structural modeling of images and inevitably disrupts the coherence and locality of visual data.

Figure: Figure 1: Motivation of the proposed V2M method. Previous vision Mambas processed image tokens with 1D SSM, whereas we extend SSM to a 2D form for more suitable image representation learning by introducing the prior of enhancing the relevance of adjacent regions for modeling.
Refer to caption: x1.png

In this paper, we propose a Visual 2-Dimensional Mamba (V2M) model as a complete solution for 2D visual modeling.
We directly processes image tokens in the 2D space instead of transforming them into 1D sequences to adapt to the 1D Mamba, as illustrated in Figure 1.
Specifically, we first generalize the state space model to the 2-dimensional space which generates the next state considering two adjacent states on both dimensions (e.g., columns and rows).
We then construct our V2M based on the 2-dimensional SSM formulation and incorporate Mamba to achieve hardware-efficient parallel processing.
In addition, considering the non-sequential nature of image tokens, we construct the 2-dimensional state equations starting from all four corners of the image.
Our proposed V2M approach essentially preserves the local similarity and coherence of image representations, and also inherits the efficiency and
input-dependent scalability of Mamba.
We conduct extensive experiments on the ImageNet-1K dataset for classification, similar to previous settings of convolution networks and vision transformers.
Furthermore, we transfer the pretrained models to downstream tasks such as object detection, instance segmentation, and semantic segmentation to evaluate the transferability of V2M.
We report both the ImageNet classification accuracy and the corresponding metrics of the downstream tasks for comparison.
Experimental results demonstrate the consistent improvements of our proposed method compared with the Vim baseline, verifying the superiority of V2M (e.g. $+0.4\%$ for Vim on ImageNet-1K).

## 2 Related Work

Generic Vision Backbones.
Convolutional Neural Networks (CNNs) and Vision Transformers (ViTs) are the most widely adopted generic backbones in computer vision.
Among them, CNNs leverage shift-invariance for feature extraction for each layer, setting a series of new benchmarks on numerous visual tasks (He et al., ; Liu et al., ; Krizhevsky et al., ; LeCun et al., ; Szegedy et al., ; Simonyan & Zisserman, ; Huang et al., ).
Nevertheless, the introduction of ViT has changed the landscape of computer vision, shifting the dominance of CNNs as the backbone to a new paradigm (Dosovitskiy et al., ; Liu et al., ; Touvron et al., ; Carion et al., ; Caron et al., ).
Specifically, the plain ViT (Chen et al., ) divides an image into multiple patches and linearly embeds these patches as the input sequence for the Transformer model with positional encoding.
Swin Transformer (Liu et al., ) introduces hierarchical feature maps and a self-attention mechanism within local windows, which enables the model to effectively handle image features at different scales.
Transformer-based models are also widely employed in multimodal tasks due to their versatility in handling data from different modalities.

In addition to CNNs and ViTs, a plethora of sophisticated visual architectures based on Mamba (Gu & Dao, ), have recently emerged in the domain of computer vision (Zhu et al., ; Liu et al., ; Huang et al., ; Liu et al., ; Ruan & Xiang, ; Yang et al., ; Pei et al., ).
For example, Vision Mamba (Vim) (Zhu et al., ) directly transfers Mamba from NLP to the visual domain, which does not incorporate significant considerations specifically for visual data.
VMamba (Liu et al., ) addresses these deficiencies through further exploration, utilizing one-dimensional scans in four distinct directions to simulate the modeling of two-dimensional images, and adopting a hierarchical structure for feature extraction. LocalMamba (Huang et al., ) further divides the input image into multiple local windows, thereby retaining the local similarity of the image to a certain extent.
Additionally, PlainMamba (Yang et al., ) is designed with a non-hierarchical structure to enhance feature fusion and modality fusion, whereas EfficientVMamba (Pei et al., ) conducts a streamlined and detailed study of the lightweight adaptation of VMamba.
These endeavors harness the high computational and memory efficiency of Mamba, yet predominantly involve the naive flattening of images input for the original 1-dimensional Mamba formulation, which compromises the local similarity and coherence of images and thus results in suboptimality.
On the contrary, we employ 2-dimensional state space equations to model image features, and thereby effectively incorporates the 2D locality prior, establishing a more rational application of Mamba within the realm of computer vision.

State Space Models.
State space models (SSMs) represent the output of the system (observation) as a function of the state of the system and typically assume that the evolution of the system state over time follows a Hidden Markov Model (HMM).
Considering the proficiency in the modeling of extended sequences, 1-dimensional SSMs have garnered extensive and conspicuous attention within the domain of natural language processing (Gu et al., ; ; ; Gupta et al., ; Gu & Dao, ).
For example, S4 (Gu et al., ) transforms the SSM into a large 1D convolution to achieve efficient parallelization and adopts HiPPO matrices as initialization to facilitate memory storage and prevent gradient explosion or vanishing.
S4D (Gu et al., ) systematically investigates the parameterization and initialization methods for diagonal state space models.
Mamba (Gu & Dao, ) further improves the traditional state space model, enabling it to selectively extract information based on the input content, combined with a hardware-aware parallel algorithm that achieves efficient processing.

Several studies have attempted to extend the formulation of the state space models into a 2-dimensional structure (Kung et al., ; Eising, ; Kurek, ; Hinamoto, ).
They typically increase the state dimension in the state space model and perform corresponding state iterative calculations along each dimension.
In computer vision, Baron *et al.* (Baron et al., ) integrates a 2-dimensional state space model into CNN and ViT architectures, implementing it as an additional layer for the processing of features.
Nonetheless, this approach neglects the input-dependent characteristic of SSM, which hampers the capacity of the model for adaptively extracting features from the input images.
In contrast, our proposed V2M generalizes the SSM to the 2-dimensional space which generates the next state considering two adjacent states on both dimensions, while also preserving the input-relevant feature of the SSM.

Visual Representation Learning.
Visual representation learning aims to acquire features that are capable of encapsulating and portraying the essence of visual data, which are widely harnessed in a diverse array of visual tasks including image classification, object detection, image segmentation, pose estimation, and video analysis.
Typical settings include pre-training on large datasets (such as ImageNet) and then fine-tuning on downstream tasks and specific datasets (He et al., ; Grill et al., ; He et al., ; Zhou et al., ; Chen & He, ; Khosla et al., ; He et al., ; Liu et al., ), which can be divided into supervised settings (Khosla et al., ; He et al., ; Liu et al., ) and unsupervised patterns (He et al., ; Grill et al., ; He et al., ; Zhou et al., ; Chen & He, ).
Specifically, supervised representation learning utilizes the ground truth labels of samples as the supervision while unsupervised representation learning constructs a label-free pretext task to conduct the training procedure, such as contrastive learning (He et al., ; Grill et al., ; Chen & He, ) and mask image modeling (He et al., ; Zhou et al., ).
In this paper, we follow the setting of supervised representation learning.
We pretrain the model with ground truth labels on the image classification task and then transfer to object detection, instance segmentation, as well as instance segmentation for verification of transferability.

## 3 Proposed Approach

In this section, we first present general preliminaries of 1-dimensional state space models, Mamba, and 2-dimensional state space models.
Then we elaborate on the implementations of the 2-dimensional state space models on visual representation learning.
Lastly, we provide an overview of the proposed V2M framework.

### 3.1 State Space Models for Vision

The state space models (SSMs) serve as the foundation for Mamba, leveraging hidden state vectors to map 1-dimensional inputs into their corresponding output manifestations, formulated as follows:

$$ $\displaystyle h^{{}^{\prime}}(t)=\mathbf{A}h(t)+\mathbf{B}x(t),$ $\displaystyle y(t)=\mathbf{C}h(t),$ (1) $$

where $x(t),y(t)\in\mathbb{R}$ denote the input and output data, $h(t)\in\mathbb{R}^{N}$ represents the latent state, $\mathbf{A}\in\mathbb{R}^{N\times N}$ denotes the evolution parameter, $\mathbf{B}\in\mathbb{R}^{N\times 1}$ and $\mathbf{C}\in\mathbb{R}^{1\times N}$ are the projection parameters, respectively.

The aforementioned form of the state space model is committed to handling continuous data, necessitating the discretization of the equations in practical applications. A conventional method of discretization is the zero-order hold (ZOH) technique, which translates continuous-time parameters $(\mathbf{A},\mathbf{B})$ into discrete forms $(\mathbf{\overline{A}},\mathbf{\overline{B}})$ by incorporating a timescale parameter $\mathbf{\Delta}\in\mathbb{R}>0$:

$$ $\displaystyle\mathbf{\overline{A}}=exp(\mathbf{\Delta A}),$ $\displaystyle\mathbf{\overline{B}}=(\mathbf{\Delta A})^{-1}(exp(\mathbf{\Delta A })-\mathbf{I})\cdot\mathbf{\Delta B}.$ (2) $$

Under such circumstances, the discretized state space models can be formulated as:

$$ $\displaystyle h_{t}=\mathbf{\overline{A}}h_{t-1}+\mathbf{\overline{B}}x_{t},$ $\displaystyle y_{t}=\mathbf{C}h_{t}.$ (3) $$

The aforementioned form constitutes the recurrent representation of the state space models, which is more amenable for inference because of the linear-time complexity. In the training phase, an equivalent convolutional form that is conducive to parallel processing with a pre-calculated convolutional kernel is indispensable:

$$ $\displaystyle\mathbf{\overline{K}}=(\mathbf{C}\mathbf{\overline{B}},\mathbf{C} \mathbf{\overline{A}}\mathbf{\overline{B}},...,\mathbf{C}\mathbf{\overline{A}} ^{L-1}\mathbf{\overline{B}}),$ $\displaystyle\mathbf{y}=\mathbf{x}\otimes\mathbf{\overline{K}},$ (4) $$

where $\otimes$ denotes the convolutional operation.

Conventional state space models exhibit constrained proficiency in the extraction of features from input data, an outcome attributed to the time-invariant nature of the parameters within the equations. Mamba rectifies this limitation by integrating time-varying parameters. Concretely, Mamba formulates the projection parameters $(\mathbf{B}\in\mathbb{R}^{B\times L\times N},\mathbf{C}\in\mathbb{R}^{B\times L
\times N})$ and timescale parameter $\mathbf{\Delta}\in\mathbb{R}^{B\times L\times D}$ in a manner that is intricately linked to the input data $\mathbf{x}\in\mathbb{R}^{B\times L\times D}$:

$$ $\displaystyle\mathbf{B}=s_{B}(\mathbf{x}),~{}~{}~{}\mathbf{C}=s_{C}(\mathbf{x} ),~{}~{}~{}\mathbf{\Delta}=\tau_{\Delta}(s_{\Delta}(\mathbf{x})+\mathbf{p}),$ (5) $$

where $s_{\{B,C,\Delta\}}(\cdot)$ denotes the linear projection of the input data, $\tau_{\Delta}(\cdot)$ is the activation function, and $\mathbf{p}$ represents the learnable parameter. Subsequent discretization follows equation 2.

However, the input-specific parameters render Mamba incapable of directly transitioning into a convolutional format, thereby precluding parallelization of the training process. Consequently, Mamba mitigates the decline in training velocity through parallel scanning tailored to hardware capabilities.

Mamba and 1D state space models are more congruent with the modeling of 1D input data. Conversely, for image samples, the extension of the state space model into a 2D format constitutes a more intrinsic method for feature extraction. A classical high-dimensional state space model is Roesser’s state space model (Kung et al., ), which employs multiple hidden states to describe the high-dimensional features. The 2D form is illustrated as follows:

$$ $\displaystyle\begin{bmatrix}h^{1}_{i,j+1}\\ h^{2}_{i+1,j}\end{bmatrix}$ $\displaystyle=\begin{bmatrix}\mathbf{A}_{1}\mathbf{A}_{2}\\ \mathbf{A}_{3}\mathbf{A}_{4}\end{bmatrix}\begin{bmatrix}h^{1}_{i,j}\\ h^{2}_{i,j}\end{bmatrix}+\begin{bmatrix}\mathbf{B}_{1}\\ \mathbf{B}_{2}\end{bmatrix}x_{i,j},$ $\displaystyle y_{i,j}$ $\displaystyle=[\mathbf{C}_{1},\mathbf{C}_{2}]\begin{bmatrix}h^{1}_{i,j}\\ h^{2}_{i,j}\end{bmatrix},$ $\displaystyle h_{i,j}=\begin{bmatrix}h^{1}_{i,j}\\ h^{2}_{i,j}\end{bmatrix},$ (6) $$

where the hidden state $h_{i,j}\in\mathbb{R}^{2N}$ is composed of two sub-states from each axis (horizontal component $h^{1}_{i,j}$ and vertical component $h^{2}_{i,j}$), the dimensions of the evolution parameters $\mathbf{A}_{i}$ and the projection parameters $(\mathbf{B}_{i},\mathbf{C}_{i})$ as well as the discretization process are equal to those in the 1D SSM. In this paper, we employ an input-contextualized discrete 2D SSM to establish the foundational architecture for image backbones.

Figure: Figure 2: Comparisons between Mamba, Vim, and V2M: (a) Mamba encodes the text information after tokenization with blocks based on 1-dimensional state space models. (b) Vision Mamba (Vim) employs a straightforward flattening of the input image and then processes it through 1-dimensional state space model blocks similar to Mamba. (c) Our proposed V2M framework involves 2-dimensional state space models into the encoding blocks without the flattening operations, which is more appropriate for processing images.
Refer to caption: x2.png

### 3.2 2-Dimensional SSM

The form of the aforementioned 2D state space model is time-invariant, i.e., the evolution parameters and projection parameters are independent of the input.
However, the transmutation of the 2D SSM into the input-dependent form precludes the direct application of convolutional training methodologies, and concurrently presents difficulties in leveraging hardware acceleration through parallel scanning processes.
Accordingly, we endeavor to convert the time-varying 2D SSM into an equivalent or near-equivalent 1D SSM, subsequently redeploying the parallel scanning methodology of 1D Mamba for the purpose of expediting the training process.

The state transition equation in equation 6 can be expanded as follows:

$$ $h^{1}_{i,j+1}=\mathbf{A}_{1}h^{1}_{i,j}+\mathbf{A}_{2}h^{2}_{i,j}+\mathbf{B}_{ 1}x_{i,j},$ (7) $$

$$ $h^{2}_{i+1,j}=\mathbf{A}_{3}h^{1}_{i,j}+\mathbf{A}_{4}h^{2}_{i,j}+\mathbf{B}_{ 2}x_{i,j}.$ (8) $$

We expound upon the state transition of the horizontal component as an illustrative example, which can be decomposed into the following two parts:

$$ $h^{1}_{i,j+1}=\mathbf{A}_{1}h^{1}_{i,j}+\mathbf{B}_{1}x_{i,j},$ (9) $$

$$ $h^{1}_{i,j+1}=\mathbf{A}_{2}h^{2}_{i,j}.$ (10) $$

We observe that equation 9 concurs with the format of the 1D SSM, thus we may seamlessly redeploy the encoding blocks from the 1D Mamba for computation with the integration of parallel scanning.
The trainable parameters are $\mathbf{A}_{1}$ and $\mathbf{B}_{1}$, respectively.
Specifically, for the input samples $\mathbf{x}\in\mathbb{R}^{B\times H\times W\times D}$, the computational procedure is conducted sequentially by rows.
Therefore, we process each row of features from the input sample as an independent sample for SSM computation, that is, fusing the $B$ and $H$ dimensions of the input sample, denoted as $\mathbf{x}^{{}^{\prime}}\in\mathbb{R}^{(BH)\times W\times D}$.

Under such circumstances, we utilize the state component values computed via equation 9 as new input values $x^{{}^{\prime}}_{i,j}$ at the corresponding positions and subsequently recombine them with equation 10:

$$ $h^{1}_{i,j+1}=\mathbf{A}_{2}h^{2}_{i,j}+\mathbf{I}x^{{}^{\prime}}_{i,j}.$ (11) $$

In equation 11, the precise computation of $h^{2}_{i,j}$ is intertwined with both state components simultaneously. We thus employ a judicious simplification, restricting the computation of $h^{2}_{i,j}$ to engage solely with the vertical aspect of the state, illustrated as follows:

$$ $h^{2}_{i+1,j}=\mathbf{A}_{2}h^{2}_{i,j}+\mathbf{I}x^{{}^{\prime}}_{i,j}.$ (12) $$

Evidently, this form can also correspond to the 1D SSM, requiring only that each column of the input sample be treated as an individual feature, with $B$ and $W$ dimensions integrated for subsequent computation, denoted as $\mathbf{x}^{{}^{\prime}}\in\mathbb{R}^{(BW)\times H\times D}$. Note that the projection parameter is set to $\mathbf{I}$, and the trainable parameters are confined to $\mathbf{A}_{2}$. Then $h^{1}_{i,j+1}$ can be directly derived from $h^{2}_{i,j}$ without additional computational processes.

Remark. The aforementioned simplified computational process may result in a less precise state calculation outcome, yet the state variables at each position remain coherent with the 2D SSM, encompassing both the horizontal and vertical state components simultaneously.
Furthermore, the computation of each state component (i.e., the state transition equation) is intricately intertwined with the two preceding state components.
Therefore, in practical applications, we are only required to separately endow the row and column features of the input sample with distinct sets of learnable parameters ( $(\mathbf{A}_{1},\mathbf{B}_{1})$ and $(\mathbf{A}_{2},\mathbf{I})$, respectively) and subsequently conduct individual 1D SSM processing on them.
Additionally, the specific implementation of equation 8 is analogous, differing only in the order of SSM computation for the sample columns and then the sample rows, as well as the differing learnable parameters.
Specially, we present a comparison diagram in Figure 2 to highlight the differences with previous backbones.

Figure: Figure 3: An overall framework of the proposed V2M approach. We employ the 2D SSM starting from four directions (upper left, upper right, lower right, and lower left) to conduct feature encoding. Within each V2M block, we initiate the process by calculating the 2D-SSM, subsequently leveraging a MLP for feature mapping. The output features from the final V2M block are utilized for classification or downstream heads
Refer to caption: x3.png

### 3.3 Visual 2D Mamba

The aforementioned implementation of the 2D SSM constitutes our Visual 2D Mamba (V2M) block.
We provide an overall framework of our V2M, illustrated in Figure 3.
Given an input batch $\mathbf{x}^{{}^{\prime}}\in\mathbb{R}^{B\times h\times w\times c}$, we first utilize the patch embedder to transform it into two-dimensional image patches and project to latent dimension, denoted as $\mathbf{x}\in\mathbb{R}^{B\times H\times W\times D}$, where $H/h=W/w$ represents the patch size.
Considering the non-temporal nature of image data, we commence the 2D SSM computation from the four corners of the image as the starting positions.
Concerning the implementation, we perform four-directional rotations of the embedding of the sample $\mathbf{x}$ around the central point, with respective angles of 0, $\pi/2$, $\pi$, and $3\pi/2$:

$$ $\displaystyle\mathbf{x}_{i}=rot_{(\pi/2)\cdot i}(\mathbf{x}),~{}~{}~{}~{}i=0,1 ,2,3,$ (13) $$

where $\mathbf{x}_{i}$ denotes the $i$th rotted embedding and $rot_{\theta}(\cdot)$ denotes rotating the sample by an angle of $\theta$.

Subsequently, we concatenate the rotated embeddings in accordance with the batch dimension as the input to the V2D blocks:

$$ $\mathbf{z}=concat([\mathbf{x}_{0},\mathbf{x}_{1},\mathbf{x}_{2},\mathbf{x}_{3} ],~{}~{}dim=0).$ (14) $$

The concatenated embeddings will undergo a comprehensive feature extraction process through K V2M blocks.
Within each V2M block, the features that have been modeled by the SSM will be aggregated in alignment with their original positions prior to rotation, presented as follows: (Note that we ignore the non-linear activation and the skip connection for brevity.)

$$ $\displaystyle\mathbf{z}^{k}=SSM_{2D}(\mathbf{z}^{k}),~{}~{}\mathbf{z}^{k}=[ \mathbf{z}_{i}^{k}],$ $\displaystyle\mathbf{z}^{k+1}=Linear(sum[rot_{2\pi-(\pi/2)\cdot i}(\mathbf{z}_ {i}^{k})]),$ $\displaystyle i=0,1,2,3$ (15) $$

where $\mathbf{z}^{k+1}$ will undergo a similar rotation operation as equation 13 before being passed to the subsequent V2M block.
We proceed to extract the class token from the output of the final V2M block, which will be passed through a classifier for subsequent supervision.

Efficiency Analysis. We concurrently implement hardware parallel scanning algorithms akin to those adopted in Mamba, alongside expedited computations facilitated by high-speed SRAM, and a recomputation method for backpropagation, all of which are orchestrated to uphold the efficiency of the model in computation and memory storage.
In addition, despite the four disparate orientations of 2D SSM in the proposed V2M framework, the concatenation along the batch dimension preserves the parallelistic characteristics of the computational process, which fully leverages the computational advantages of the GPU to forestall severe deterioration in processing speed.

## 4 Experiments

In this section, we conducted a series of comprehensive experiments to assess the efficacy of our V2M framework.
The network was initially pretrained using V2M on the ImageNet-1K dataset (Russakovsky et al., ), after which its performance was assessed across various downstream tasks.
Moreover, we provided in-depth ablation studies to delve into the intricacies of the effectiveness of V2M.

**Table 1: Classification results of different visual backbones on ImageNet. (${\dagger}$ denotes our reproduced performances under the default settings. * denotes using the pyramid architecture for V2M.)**
| Method | Backbone | Image | Params (M). | FLOPs (G) | Top-1 |
| --- | --- | --- | --- | --- | --- |
| Size | Acc (%) |  |  |  |  |
| ResNet-18 (He et al., ) | ConvNet | 224^2 | 12 | - | 69.8 |
| DeiT-T (Touvron et al., ) | Transformer | 224^2 | 6 | 1.3 | 72.2 |
| PlainMamba-L1 (Yang et al., ) | SSM | 224^2 | 7 | 3.0 | 77.9 |
| EffVMamba-T (Pei et al., ) | SSM | 224^2 | 6 | 0.8 | 76.5 |
| EffVMamba-S (Pei et al., ) | SSM | 224^2 | 11 | 1.3 | 78.7 |
| Vim-T${\dagger}$ (Zhu et al., ) | SSM | 224^2 | 7 | 1.5 | 75.8 |
| V2M-T (ours) | SSM | 224^2 | 7 | 1.9 | 76.2 |
| LocalVim-T (Huang et al., ) | SSM | 224^2 | 8 | 1.5 | 76.2 |
| V2M-T + local window (ours) | SSM | 224^2 | 8 | 1.8 | 76.4 |
| ResNet-50 (He et al., ) | ConvNet | 224^2 | 25 | - | 77.2 |
| RegNetY-4G (Radosavovic et al., ) | ConvNet | 224^2 | 21 | 4.0 | 80.0 |
| DeiT-S (Touvron et al., ) | Transformer | 224^2 | 22 | 4.6 | 79.9 |
| Swin-T (Liu et al., ) | Transformer | 224^2 | 29 | 4.5 | 81.2 |
| PlainMamba-L2 (Yang et al., ) | SSM | 224^2 | 25 | 8.1 | 81.6 |
| EffVMamba-B (Pei et al., ) | SSM | 224^2 | 33 | 4.0 | 81.8 |
| Vim-S${\dagger}$ (Zhu et al., ) | SSM | 224^2 | 26 | 5.1 | 80.3 |
| V2M-S (ours) | SSM | 224^2 | 26 | 5.9 | 80.5 |
| LocalVim-S (Huang et al., ) | SSM | 224^2 | 28 | 4.8 | 81.2 |
| V2M-S + local window (ours) | SSM | 224^2 | 28 | 5.4 | 81.3 |
| VMamba-T (Liu et al., ) | SSM | 224^2 | 30 | 4.9 | 82.6 |
| V2M-S* (ours) | SSM | 224^2 | 30 | 5.4 | 82.9 |
| ResNet-101 (He et al., ) | ConvNet | 224^2 | 45 | - | 78.3 |
| ResNet-152 (He et al., ) | ConvNet | 224^2 | 60 | - | 78.6 |
| RegNetY-8G (Radosavovic et al., ) | ConvNet | 224^2 | 39 | 8.0 | 81.7 |
| RegNetY-16G (Radosavovic et al., ) | ConvNet | 224^2 | 84 | 16.0 | 82.9 |
| ViT-B/16 (Dosovitskiy et al., ) | Transformer | 384^2 | 86 | 55.4 | 77.9 |
| DeiT-B (Touvron et al., ) | Transformer | 224^2 | 86 | 17.5 | 81.8 |
| Swin-S (Liu et al., ) | Transformer | 224^2 | 50 | 8.7 | 83.2 |
| Swin-B (Liu et al., ) | Transformer | 224^2 | 88 | 15.4 | 83.5 |
| PlainMamba-L3 (Yang et al., ) | SSM | 224^2 | 50 | 14.4 | 82.3 |
| VMamba-S (Liu et al., ) | SSM | 224^2 | 50 | 8.7 | 83.6 |
| V2M-B* (ours) | SSM | 224^2 | 50 | 9.6 | 83.8 |

### 4.1 Experimental Setup

Datasets.
We pretrain our V2M model on the training dataset of ImageNet-1K (Russakovsky et al., ), which comprises 1,280,000 samples across 1,000 distinct categories. The classification performance is verified on the validation set, which consists of 50,000 images.
Furthermore, we leverage the COCO (Lin et al., ) dataset for object detection and instance segmentation, encompassing 118K images for training, 5K for validation, and 40K for testing.
The semantic segmentation task is conducted on ADE20K (Zhou et al., ), which comprises 150 detailed semantic categories, offering 20K, 2K, and 3K images for training, validation, and testing, respectively.

Implementation Details.
We adopted Vision Mamba (Vim) (Zhu et al., ), LocalMamba (Huang et al., ), and VMamba (Liu et al., ) as our baselines for experimental comparisons, thus maintaining consistency in our experimental configuration.
We provided multiple V2M models according to their parameters, including V2M-T (Tiny), V2M-S (Small), and V2M-B (Base).
Additionally, the network architecture encompassed both the plain and pyramid configurations, with the latter denoted by *.
For image classification, we incorporated a suite of data augmentation techniques, including random cropping, random horizontal flipping, label smoothing regularization, mixup (Zhang et al., ), and random erasing.
We adopted AdamW as the optimizer with a momentum of 0.9 and a weight decay of 0.05.
We set the batch size to 512 and the training epoch to 300.
We set the learning rate to 5e^-4 with a cosine schedule.
The input images of both training and validation were cropped to 224$\times$224.
We reported both top-1 accuracy and FLOPs for evaluation.

As for object detection and instance segmentation, we adopted the MMDetection (Chen et al., ) codebase for evaluation.
We employed Mask R-CNN detector (He et al., ) with both 1x and 3x training schedules, similar to VMamba.
We reported AP with different settings for comparison.
Additionally, we equipped the pretrained models with UPerNet (Xiao et al., ) for semantic segmentation using the MMSegmentaion (Contributors, ) codebase.
We applied AdamW with a learning rate of 6e^-5 and a weight decay of 0.01.
We adopted a learning schedule of 160k reported mIoU for comparison.
All our experiments were performed on 8 RTX 3090 GPUs.

**Table 2: Object detection and instance segmentation results on COCO. (* denotes using the pyramid architecture for V2M.)**
| Method | Detector | Params (M). | $\mathbf{AP}^{b}$ | $\mathbf{AP}_{50}^{b}$ | $\mathbf{AP}_{75}^{b}$ | $\mathbf{AP}^{m}$ | $\mathbf{AP}_{50}^{m}$ | $\mathbf{AP}_{75}^{m}$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| ResNet-50 (He et al., ) | MaskRCNN@1x | 44 | 38.2 | 58.8 | 41.4 | 34.7 | 55.7 | 37.2 |
| ResNet-101 (He et al., ) | MaskRCNN@1x | 63 | 38.2 | 58.8 | 41.4 | 34.7 | 55.7 | 37.2 |
| ConvNeXt-T (Liu et al., ) | MaskRCNN@1x | 48 | 44.2 | 66.6 | 48.3 | 40.1 | 63.3 | 42.8 |
| ConvNeXt-S (Liu et al., ) | MaskRCNN@1x | 70 | 45.4 | 67.9 | 50.0 | 41.8 | 65.2 | 45.1 |
| ConvNeXt-T (Liu et al., ) | MaskRCNN@3x | 48 | 46.2 | 67.9 | 50.8 | 41.7 | 65.0 | 44.9 |
| ConvNeXt-S (Liu et al., ) | MaskRCNN@3x | 70 | 47.9 | 70.0 | 52.7 | 42.9 | 66.9 | 46.2 |
| Swin-T (Liu et al., ) | MaskRCNN@1x | 48 | 42.7 | 65.2 | 46.8 | 39.3 | 62.2 | 42.2 |
| Swin-S (Liu et al., ) | MaskRCNN@1x | 69 | 44.8 | 66.6 | 48.9 | 40.9 | 63.2 | 44.2 |
| Swin-T (Liu et al., ) | MaskRCNN@3x | 48 | 46.0 | 68.1 | 50.3 | 41.6 | 65.1 | 44.9 |
| Swin-S (Liu et al., ) | MaskRCNN@3x | 69 | 48.2 | 69.8 | 52.8 | 43.2 | 67.0 | 46.1 |
| VMamba-T (Liu et al., ) | MaskRCNN@1x | 50 | 47.3 | 69.3 | 52.0 | 42.7 | 66.4 | 45.9 |
| V2M-S* (ours) | MaskRCNN@1x | 50 | 47.6 | 69.4 | 52.2 | 42.9 | 66.5 | 46.3 |
| VMamba-S (Liu et al., ) | MaskRCNN@1x | 70 | 48.7 | 70.0 | 53.4 | 43.7 | 67.3 | 47.0 |
| V2M-B* (ours) | MaskRCNN@1x | 70 | 48.9 | 70.2 | 53.6 | 43.8 | 67.5 | 47.1 |
| VMamba-T (Liu et al., ) | MaskRCNN@3x | 50 | 48.8 | 70.4 | 53.5 | 43.7 | 67.4 | 47.0 |
| V2M-S* (ours) | MaskRCNN@3x | 50 | 49.0 | 70.6 | 53.5 | 43.8 | 67.5 | 47.2 |
| VMamba-S (Liu et al., ) | MaskRCNN@3x | 70 | 49.9 | 70.9 | 54.7 | 44.2 | 68.2 | 47.7 |
| V2M-B* (ours) | MaskRCNN@3x | 70 | 50.0 | 70.9 | 54.8 | 44.3 | 68.4 | 47.8 |

**Table 3: Semantic segmentation results on ADE20K. (* denotes using the pyramid architecture for V2M.)**
| Method | Segmentor | Image size | Params (M). | mIoU (SS) | mIoU (MS) |
| --- | --- | --- | --- | --- | --- |
| Swin-T (Liu et al., ) | UperNet@160k | 512^2 | 60 | 44.4 | 45.8 |
| Swin-S (Liu et al., ) | UperNet@160k | 512^2 | 81 | 47.6 | 49.5 |
| Vim-T (Zhu et al., ) | UperNet@160k | 512^2 | 13 | 41.0 | - |
| V2m-T (ours) | UperNet@160k | 512^2 | 13 | 41.4 | 42.0 |
| Vim-S (Zhu et al., ) | UperNet@160k | 512^2 | 46 | 44.9 | - |
| V2M-S (ours) | UperNet@160k | 512^2 | 46 | 45.1 | 46.1 |
| LocalVim-T (Huang et al., ) | UperNet@160k | 512^2 | 36 | 43.4 | 44.4 |
| V2M-T + local window (ours) | UperNet@160k | 512^2 | 36 | 43.5 | 44.6 |
| LocalVim-S (Huang et al., ) | UperNet@160k | 512^2 | 58 | 46.4 | 47.5 |
| V2M-S + local window (ours) | UperNet@160k | 512^2 | 58 | 46.6 | 47.7 |
| VMamba-T (Liu et al., ) | UperNet@160k | 512^2 | 62 | 47.9 | 48.8 |
| V2M-S* (ours) | UperNet@160k | 512^2 | 62 | 48.2 | 49.0 |
| VMamba-S (Liu et al., ) | UperNet@160k | 512^2 | 82 | 50.6 | 51.2 |
| V2M-B* (ours) | UperNet@160k | 512^2 | 82 | 50.8 | 51.3 |

### 4.2 Main Results

Image Classification on ImageNet.
We provided the evaluation results of V2M on ImageNet classification in Table 1.
The comparative methods include convolutional networks (ConvNets), transformer-based models, and the SSM baselines.
We observe that SSM-based methods exhibit comparative superiority over ConvNets and transformer-based models under comparable parameters.
For example, our proposed V2M-T achieved a 6.4% increase in Top-1 accuracy compared to ResNet-18 and a 4.0% improvement over DeiT-T.
Furthermore, V2M-S* outperforms ResNet-50, RegNetY-4G, and DeiT-S, with respective increases of 5.7%, 2.9%, and 3.0% in Top-1 accuracy.
In addition to ConvNets and transformer-based models, V2M also demonstrates a certain performance improvement across both the tiny and small model parameter configurations in comparison to the adopted baselines, registering a 0.4%/0.2%/0.3% enhancement in the tiny model and a 0.2%/0.1%/0.2% improvement in the small model compared with Vim/LocalMamba/VMamba, respectively.
The experimental results verify the effectiveness of the proposed V2M framework, which indicates that the adaptive modeling process of image data through 2D SSM is more suitable compared to the straightforward flattening of images with the 1D SSM modeling.

Transfer Learning on Object Detection and Instance Segmentation.
We evaluated the transferability of V2M on the COCO dataset, including object detection and instance segmentation tasks.
The experimental results are presented in Table 2.
Consistent with the image classification on ImageNet, the SSM-based models demonstrate superior performances on object detection and instance segmentation tasks compared to transformer-based counterparts with equivalent model parameters.
Moreover, compared with the corresponding baselines, the proposed V2M method exhibits certain advantages in both object detection and instance segmentation.
Specifically, V2M-S* outperforms VMamba by 0.3 box AP and 0.2 mask AP under 1x schedule, which demonstrates the superior generalization ability of the representations learned by V2M.

Transfer Learning on Semantic Segmentation.
We provided the semantic segmentation performances on ADE20K in Table 3.
The comparative backbones include ResNet-50, ResNet-101, ConvNeXts, and Swin Transformers.
We observe that V2M achieves consistent improvements over these backbones under approximate parameters.
Furthermore, V2M-S* surpasses the VMamba-T baseline by 0.3 mIoU (SS), demonstrating the effectiveness of the proposed framework.

Figure: Figure 4: Arrangements of the class token. (a) Obtaining a class token through feature mean pooling. (b) Adding an additional row and column of class tokens at the edge of the image and adopting the corner token for subsequent classification. (c) Adding an additional row and column of class tokens at the middle of the image and adopting the center token for subsequent classification. We adopt (c) as the default setting in the main experiments.
Refer to caption: x4.png

### 4.3 Experimental Analysis

Effect of Modeling Directions of the 2D SSM.
Our proposed V2M method undergoes 2D SSM modeling in four directions, and thus we assess the impact of different directions on model performance.
Specifically, we individually attempted modeling in 1, 2, 3, and 4 different directions, and subsequently tested the classification performance on ImageNet.
The comparison results are illustrated in Table 4.
We discern that the performance of V2M on the classification task ascends incrementally with the increase of modeling directions, culminating in optimal outcomes upon the integration of modeling across all four directions.

**Table 4: Effect of 2D SSM directions with V2M-T. (UL=Upper Left, UR=Upper Right, LL=Lower Left, LR=Lower Right)**
| Direction | Top-1 Acc | Top-5 Acc |
| --- | --- | --- |
| UL | 75.2 | 92.3 |
| UL + LR | 75.9 | 92.9 |
| UL + LR + UR | 75.9 | 93.0 |
| UL + LR + UR + LF | 76.2 | 93.1 |

Arrangements of Classification.
Diverging from the Vim baseline, our proposed V2M performs 2D SSM modeling by independently correlating the features of rows and columns without the flattening of the entire image patch.
Consequently, we are unable to merely append a solitary class token for subsequent classification tasks.
To address this limitation, we present three viable solutions, including obtaining the class token through feature mean pooling, augmenting the image with an additional row and column of class tokens at the edge, and incorporating a row and column of class tokens at the middle of the image, shown in Figure 4.
We observe that the third scheme yields a commendable performance and is thus designated as the default setting.

## 5 Conclusion

In this paper, we have presented a visual 2-dimensional mamba (V2M) framework for effective image representation learning.
We have employed the 2D SSM for the modeling process of images, thereby preserving the inherent prior of local invariance within the images.
We have concurrently performed 2D SSM modeling in four directions to enhance the representational capacity of the model considering the non-temporal nature of the input images.
We have conducted a series of experiments that encompass classification and various downstream tasks, including object detection, instance segmentation, and semantic segmentation, to assess the effectiveness of our framework.
Experimental results have demonstrated the superior classification as well as transferability performances of the proposed V2M framework.

Limitations.
The implementation of 2D SSM modeling in four directions within each V2M block necessitates a compromise in speed.
Future research will concentrate on both software and hardware algorithm optimizations to mitigate this drawback.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: - Writing: Use "Eq.~\eqref{xxx}" for equation references. The background formulation part can be shortened, e.g., Eqs. (1-5), as these have been widely described in previous vision Mamba papers.

- V2M seems better at spatial information modeling compared to ViM. If the positional embedding in V2M is removed, how significant is the performance drop? The performance drop could be minimal, right?

- How is the classification head designed?

- It would be interesting to see the formulation and results of V3M for video modeling.
soundness: 3
strengths: - Developing high-performance, high-dimensional state-space models is an important topic, and the paper is an interesting attempt in this direction.

- The results in the paper are generally good, demonstrating strong performance on ImageNet classification, COCO detection, and segmentation compared to strong baselines such as ViM, LocalMamba, and VMamba.
summary: The paper proposes a 2D state-space model (SSM) for visual representation learning, named V2M. Unlike previous Mamba-based visual representation methods that perform 1D sequence learning, V2M operates in 2D space, with each hidden state conditioned on its top and left tokens. Baron et al. have presented a 2D SSM but neglect the input-dependent characteristic of SSM, which is crucial for boosting performance.
weaknesses: - One of the essential advantages of Mamba is its high efficiency in processing long sequences. The paper does not show V2M's performance on long sequences.

- V2M's runtime is not reported. Given that hardware-efficient implementation is an important part of V2M, comparing its runtime with 1D Mamba is crucial.

---

confidence: 5
contribution: 3
initial_rating: 8
presentation: 2
questions: Please refer to the weakness part.
soundness: 3
strengths: - The consideration of 2d spatial relationship in the visual domain is intuitive and sounds reasonable.
- Different from other hierarchical vision mamba methods, this paper handles the visual perception with a plain, non-hierarchical architecture, maintaining the ability in multimodality applications.
- Compared with the baseline method, this paper presents significant improvements in both classification and dense prediction tasks.
- This paper aims at a good question. Scanning strategy with SSM methods is a key problem because the 1d original scanning does not treat all tokens equally.
summary: The paper presents Visual 2-Dimensional Mamba (V2M), a novel framework for image representation learning that adapts the state space model (SSM), specifically the Mamba model, to the two-dimensional structure of image data. This paper extends the 1d SSM to a 2d form to fundamentally accommodate the structure of image data. As a non-hierarchical vision mamba method, it achieves better results than its baseline Vim.
weaknesses: - The character corners in Figure 1 require further explanation.
- Note that function names in formulas are usually typeset properly in Roman font, e.g., `rot`, `concat`, `SSM`, `Linear`, `sum` in Eq. 13~15.
- The paper references preprints and arXiv versions of significant works, such as Mamba (COLM), Vision Mamba (ICML), and VMamba (NeurIPS). The authors should update these citations to their final published versions to reflect the current state of the literature.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: Could you add more efficiency experiments and ablation studies to this paper and provide a more thorough discussion based on those findings? 

Additionally, could you offer a more detailed analysis and explanation of the simplifications mentioned in Section 3.2?
soundness: 3
strengths: This paper innovatively introduces a 2D state space to address the spatial coherence issues faced by SSM models in the visual domain, achieving promising results.
summary: The paper proposes the Visual 2-Dimensional Mamba (V2M) framework, an innovative model for image representation learning. Unlike traditional Mamba models that process image data in a 1D sequence, V2M adapts a 2D state space model to retain spatial structure, preserving local relationships between pixels. This 2D adaptation allows V2M to process image patches as 2D entities directly, thereby enhancing locality and coherence in visual representation.
weaknesses: 1.Novelty

The novelty is moderate, as there has already been substantial research, such as VMamba[1], exploring how sequence models can preserve 2D structural information in visual tasks.

2.Implementation

While this paper introduces a theory of 2D state space to address coherence issues in visual tasks,  its implementation relies on excessive simplification, lacking a detailed explanation of the rationale behind this approach and its potential implications.

3.Performance

The performance demonstrated in this paper shows only minor gains over early baselines like Vim[2] and VMamba[1], without advantages over more recent models like EfficientVMamba[3].

4.Experiment

It is foreseeable that the proposed method will impact speed and memory usage; however, the paper lacks sufficient efficiency tests, memory usage tests, and ablation studies, resulting in limited experimental support.

5.Presentation

The expression of the formulas is somewhat confusing, such as the approach used to split Equation 7 into Equations 9 and 10, and the unclear logical progression from Equation 11 to Equation 12.

[1]VMamba: Visual State Space Model. Proceedings of the 38th Conference on Neural Information Processing Systems.

[2]Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model. Proceedings of the 41st International Conference on Machine Learning.

[3]EfficientVMamba: Atrous Selective Scan for Light Weight Visual Mamba. arXiv preprint arXiv:2403.09977.

---

confidence: 5
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. The method part (Sec. 3.2) is not clear enough. From equation 11 and 12, it seems that the horizontal component and the vertical component are the same after simplification, which contradicts formula 6. It is better if the authors can give a clearer and comprehensive explanation about the simplified computational process.
2. Please add the ablation study for 2D SSM block itself.
soundness: 3
strengths: 1. The proposed time-varying 2D SSM is based on Roesser’s state space model, which intuitively depicts the spatial information of the image. 

2. This paper proposes a simplified computational process for Roesser’s state space model, resulting in two consecutive 1D SSM block, thus achieving hardware-efficient parallel processing.

3. The proposed V2M model adopts four scanning directions and compares three types of classification tokens: mean pooling, edge class token and middle class token.
summary: The submission introduces a Visual 2-Dimensional Mamba (V2M) model, aimed at preserving the inherent prior of 2D image data. The V2M model incorporates a 2D State Space Model (SSM) that directly processes image tokens in 2D state space by considering two adjacent states in both dimensions. The authors convert the time-varying 2D SSM into a near-equivalent 1D Mamba format for efficient parallel computation. Experimental results show that V2M outperforms existing visual backbones on ImageNet classification and various downstream tasks.
weaknesses: 1. Although the theoretical analysis of 2D SSM is based on Roesser’s state space model, it is simplified into two 1D SSM processes for parallelism. I am not sure whether this simplification retains the ability of Roesser’s SSM to capture the spatial dependencies within the image.

2. The performance improvement is tiny especially in the downstream tasks. It is not sure whether the improvement is brought by the proposed 2D SSM. From the ablation study in Table 4, modeling in four directions seems to offer greater benefits than the 2D SSM itself. 

3. The FLOP of V2M is much higher than that of baseline models, for example, 1.5G for Vim-T and 1.9G for V2M-T. Specific operating speeds, such as throughput, should be provided. Please gradually add each module to the baseline model, and give the corresponding number of parameters, FLOPs, throughput and performance.

Meta-review: This paper proposes a Mamba-based backbone for vision tasks, using a 2D SSM formulation to effectively incorporate 2D locality priors into vision data. Reviewers generally find the motivation for this approach intuitive and interesting; however, they raised significant concerns regarding the method's complexity and performance. After the rebuttal and discussion phases, some major issues remain unaddressed. 

The paper received mixed reviews, with scores of 6, 5, 3, and 8, resulting in an average score of 5.5. The Area Chair carefully reviewed the paper and agreed with Reviewer 7MsE that the limited performance and efficiency could significantly diminish the significance of this method as a Mamba-based vision backbone. Furthermore, the novelty seems limited due to the existence of several prior works that utilize 2D scans. In light of these factors, the Area Chair recommends rejecting the paper.

Decision: Reject