arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2406.05992eess.IVcs.CV

MHS-VM:面向 Vision Mamba 的并行子空间多头扫描

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

Hangzhou Dianzi University(杭州电子科技大学)

Zhongping Ji

AI总结:

本文提出多头扫描模块 MHS 和扫描路径注意力机制 SRA,通过在多个低维子空间沿不同路径执行选择性扫描来改进 Vision Mamba 的二维特征构建,替换 VM-UNet 中的 SS2D 后以更少参数实现显著性能提升。

摘要:

近年来,以 Mamba 为代表的状态空间模型(SSMs)在具有线性复杂度的长程依赖建模方面展现出巨大潜力。随后,Vision Mamba 及其后续架构相继被提出,并在视觉任务上表现良好。将 Mamba 应用于视觉任务的关键步骤是以序列方式构建二维视觉特征。为了通过一维选择性扫描在二维图像空间内有效组织和构建视觉特征,我们提出了一种新颖的多头扫描(Multi-Head Scan,MHS)模块。从前一层提取的嵌入被投影到多个低维子空间中。随后,在每个子空间内,沿不同的扫描路径执行选择性扫描。由多头扫描过程得到的子嵌入随后被整合,并最终投影回高维空间。此外,我们引入了一种扫描路径注意力(Scan Route Attention,SRA)机制,以增强该模块识别复杂结构的能力。为验证所提模块的有效性,我们仅将 VM-UNet 中的二维选择性扫描(SS2D)模块替换为所提出的模块,并在不使用任何预训练权重的情况下从头训练模型。结果表明,在减少原始 VM-UNet 参数量的同时,性能得到了显著提升。本研究的代码已公开于 https://github.com/PixDeep/MHS-VM。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑