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

最近，状态空间模型（SSMs）中，以Mamba为主要例子，已经展现出在具有线性复杂度的长程依赖建模方面的巨大潜力。接着，作者依次介绍了Vision Mamba和其后续架构，并展示了它们在视觉任务上的出色表现。将Mamba应用于视觉任务的关键步骤是以顺序方式构建2D视觉特征。为了通过1D选择性扫描在2D图像空间内有效地组织和构建视觉特征，我们提出了一种新颖的多头扫描（MHS）模块。从前一层提取的嵌入被投影到多个低维子空间中。随后，在每个子空间中，沿着不同的扫描路径进行选择性扫描。从多头扫描过程中获得的子嵌入最终被集成并投影回高维空间。此外，我们还加入了扫描路线注意力（SRA）机制，以增强模块识别复杂结构的能力。为了验证我们模块的有效性，我们将VM-UNet中的2D-Selective-Scan（SS2D）块专门替换为我们提出的模块，并在不使用任何预训练权重的情况下从头开始训练我们的模型。结果表明，在减少原始VM-UNet的参数的同时，性能有了显著的提高。本研究的代码可在https://github.com/PixDeep/MHS-VM上公开获取。

本论文旨在解决长期依赖建模的问题，并将Mamba应用于视觉任务中。同时，论文提出了一种新的Multi-Head Scan（MHS）模块，以有效地在2D图像空间内进行1D选择扫描，来构建视觉特征。

论文的关键思路是将先前层中提取的嵌入投影到多个低维子空间中，然后在每个子空间中沿着不同的扫描路径执行选择性扫描。最终，将从多头扫描过程中获得的子嵌入集成并投影回高维空间。同时，论文还引入了Scan Route Attention（SRA）机制来增强模块识别复杂结构的能力。

论文使用提出的MHS模块替换了VM-UNet中的2D-Selective-Scan（SS2D）块，并从头开始训练模型，结果表明在减少原始VM-UNet参数的同时，性能得到了显著提高。论文提供了公开可用的代码，同时还使用了多个数据集进行实验，值得进一步研究。

在最近的相关研究中，还有一些关于State Space Models（SSMs）和Mamba的研究。例如，Mamba在长期依赖建模方面表现出色。此外，还有一些关于视觉任务中使用SSMs的研究，例如对图像序列建模和视频分割的研究。