为了解决模型中的方向敏感性问题,我们引入了交叉扫描模块(CSM)。该模块能够遍历图像空间域,将任意视觉图像转化成有序序列。
VMamba降低注意力机制复杂度的概念来源于“具有选择性的扫描状态空间序列模型”(Selective Scan Space State Sequential Model ,S6)。S6原来应用于自然语言处理领域,与传统的注意力机制不同,S6使得在序列中的每一个元素能够与任意一个之前扫描过的样本交互。这样使得注意力机制的平方复杂度降低为线性。
然而,由于视觉数据的非因果特性(NLP有时序性),直接将S6方法应用于分块并展平的视觉图像上会导致全局感受野的损失。这是由于没有经过扫描的图像片间的注意力关联没有被估计。论文中将这一现象称为“方向敏感”问题。
为解决这一问题,提出“交叉扫描模块”(Cross-Scan Module,CSM)。CSM采用“四向”扫描策略,从图像的四角开始,曲折行进至对角。这种策略保证了在特征图种的每个元素都能够融合其他位置和方向的元素。因此,这种策略可以使模型在拥有全局感知野的同时,具有线性的计算复杂度。
论文在多种视觉任务上开展了详尽的实验来验证所提出VMama的有效性。如图2所示,在ImageNet-1K数据集上,相比于Resnet、ViT以及Swin模型,VMamba具有更强或者至少相当的性能。论文还针对下游稠密任务进行了实验,例如:VMambaTiny/Small/Base (分别有 22/44/75 M 参数)在COCO数据集上,使用MaskRCNN检测器达到了 46.5%/48.2%/48.5%的 mAP,并且在ADE20K数据集上使用UperNet达到了 47.3%/49.5%/50.0%的mIoU。这些结果表明了所提出的VMamba是一个性能很强的基础模型。进一步,当输入图像大小越来越大时,尽管ViT取得了更佳的性能,ViT的FLOP数目的增长速度显著高于CNN。而论文所提出的VMamba与ViT相比,在达到相当性能的同时,其FLOP数目增长更少,为近似线性增长。