MHS-VM: Multi-Head Scanning in Parallel Subspaces for Vision Mamba

TL;DR

This paper introduces MHS-VM, a novel multi-head scanning module for vision models that improves feature organization in 2D images, enhancing performance and reducing parameters in VM-UNet.

Contribution

The paper proposes a multi-head scan module with scan route attention for better 2D feature construction in vision models, replacing the 2D-Selective-Scan block in VM-UNet.

Findings

Significant performance improvements on visual tasks.

Reduced model parameters compared to original VM-UNet.

Effective organization of 2D features via multi-head scanning.

Abstract

Recently, State Space Models (SSMs), with Mamba as a prime example, have shown great promise for long-range dependency modeling with linear complexity. Then, Vision Mamba and the subsequent architectures are presented successively, and they perform well on visual tasks. The crucial step of applying Mamba to visual tasks is to construct 2D visual features in sequential manners. To effectively organize and construct visual features within the 2D image space through 1D selective scan, we propose a novel Multi-Head Scan (MHS) module. The embeddings extracted from the preceding layer are projected into multiple lower-dimensional subspaces. Subsequently, within each subspace, the selective scan is performed along distinct scan routes. The resulting sub-embeddings, obtained from the multi-head scan process, are then integrated and ultimately projected back into the high-dimensional space. Moreover, we incorporate a Scan Route Attention (SRA) mechanism to enhance the module's capability to discern complex structures. To validate the efficacy of our module, we exclusively substitute the 2D-Selective-Scan (SS2D) block in VM-UNet with our proposed module, and we train our models from scratch without using any pre-trained weights. The results indicate a significant improvement in performance while reducing the parameters of the original VM-UNet. The code for this study is publicly available at https://github.com/PixDeep/MHS-VM.