V2M: Visual 2-Dimensional Mamba for Image Representation Learning

TL;DR

V2M introduces a novel 2D spatially-aware Mamba model for image representation learning, directly processing image tokens in 2D to better preserve local structure and improve performance on various visual tasks.

Contribution

The paper proposes a new 2D extension of the Mamba model using 2D state space modeling, enabling efficient and locality-aware image processing.

Findings

Outperforms other visual backbones on ImageNet classification.

Effective in downstream tasks like object detection and segmentation.

Maintains hardware efficiency and scalability.

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.