V2M: Visual 2-Dimensional Mamba for Image Representation Learning

TL;DR

V2M introduces a Visual 2-Dimensional Mamba framework that extends state space models to 2D, enabling direct, four-directional processing of image tokens to capture local locality without flattening. By decomposing a 2D SSM into tractable row/column computations and leveraging hardware-aware parallelism, V2M achieves strong ImageNet classification performance and superior transferability to object detection, instance segmentation, and semantic segmentation compared with 1D Mamba baselines. Key contributions include the 2D SSM formulation, four-directional encoding from image corners, and an effective class-token arrangement, validated by extensive experiments and ablations. The approach delivers practical impact by offering a scalable, locality-preserving alternative to transformers for vision tasks, with competitive efficiency on large-scale benchmarks.

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.

Figures (4)

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