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 this https URL

提出了一种基于 State Space Models 的 Multi-Head Scan (MHS) 模块，通过 1D selective scan 在 2D 图像空间内构建视觉特征，结合了 Scan Route Attention (SRA) 机制来提高模块的复杂结构辨识能力，并在实验证明了该方法在性能上的显著改善与参数减少。