读论文：VMamba: Visual State Space Model

作者：我家小花儿 | 2024-08-01 18:02:57

踩

vmamba

（引用）

论文所在期刊：
发布时间：18 Jan 2024
所在级别：
影响因子：

论文笔记

注意机制在图像大小方面要求二次复杂度，导致在处理下游密集预测任务(如目标检测、语义分割等)时产生昂贵的计算开销
为了解决第一点的问题，现有的方法是通过限制计算窗口的大小或跨步来提高注意力的效率已经付出了大量的努力，尽管这是以限制接受域的规模为代价的
由于视觉数据的非因果性，直接将vmamba应用于贴片和扁平图像将不可避免地导致受限制的接受野，因为无法估计与未扫描贴片的关系——引入的交叉扫描模块(CSM)：该策略确保了特征映射中的每个元素都集成了来自不同方向的所有其他位置的信息，从而在不增加线性计算复杂度的情况下呈现全局接受场

2D选择性扫描：尽管S6具有独特的特性，但它会对输入数据进行因果处理，因此只能捕获数据扫描部分内的信息。这自然使S6与涉及时间数据的NLP任务对齐，但在适应非因果数据（如图像、图形、集合等）时会带来重大挑战。尽管图像具有非因果性质，但它们与文本的不同之处在于它们包含2D空间信息（例如局部纹理和全局结构），作者选择通过坚持选择性扫描方法来保留动态权重，不幸的是，这不允许我们遵循并集成卷积运算。
为了解决这个问题，作者提出了交叉扫描模块（CSM），如图2所示。选择将图像块沿着行和列展开为序列（扫描扩展），然后沿四个不同的方向进行沿着扫描：左上至右下、右下至左上、右上至左下和左下至右上。通过这种方式，任何像素（例如图2中的中心像素）都集成了来自不同方向上所有其他像素的信息。然后，将每个序列重塑为单个图像，所有序列合并为一个新的序列，如图3所示（扫描合并）。

整体架构
- VMamba通过使用茎模块将输入图像划分为补丁开始该过程，类似于ViTs，但没有进一步将补丁平坦化为1-D序列。这种修改保留了图像的2D结构，从而产生了尺寸为H/4 × W/4 × C1的特征图。
- 然后在特征图上堆叠几个VSS块，保持相同的维度，构成“阶段1”。
- VMamba中的分层表示是通过补丁合并操作在“阶段1”中对特征图进行下采样来构建的。随后，涉及更多的VSS块，产生输出分辨率为H/8 × W/8，并形成“阶段2”。重复此过程以分别创建分辨率为H/16 × W/16和H/32 × W/32的“阶段3”和“阶段4”。
- 所有这些阶段共同构建了类似于流行的CNN模型和一些ViT的分层表示。由此产生的架构可以作为一个通用的替代其他视觉模型在实际应用中具有类似的要求。
VSS模块：VSS块的结构如图4（b）所示。
- 输入经历初始线性嵌入层，输出分成两个信息流。
- 一个流通过一个3 × 3的深度方向卷积层，然后是一个Silu激活函数，然后进入核心SS2D模块。
- SS2D的输出经过层规范化层，然后添加到其他信息流的输出中，该信息流已经过Silu激活。这种组合产生VSS块的最终输出。
Vmamba的三种规模

基于ImageNet-1 K的图像分类
VMamba-T/S/B从头开始训练300个epoch（前20个epoch预热），使用1024的批量大小。训练过程采用AdamW优化器，beta设置为（0.9，0.999），动量为0.9，余弦衰减学习率调度器，初始学习率为1 × 10⁻³，权重衰减为0.05。还采用了其他技术，如标签平滑（0.1）和指数移动平均（EMA）。
COCO上的目标检测
训练框架建立在mmdetection库上，并且我们使用Mask-RCNN检测器坚持Swin中的超参数。具体来说，我们使用AdamW优化器，并对12和36个epoch的预训练分类模型（在ImageNet-1 K上）进行微调。对于VMamba-T/S/B，丢弃路径速率分别设置为0.2%/0.2%/0.2%2。学习率初始化为1 × 10⁻⁴，在第9和第11个历元时降低10倍。我们实现了多尺度训练和随机翻转，批量大小为16。
基于ADE20K的语义分割
在预训练模型的基础上构建一个UperHead。使用AdamW优化器，我们将学习率设置为6 × 10⁻⁵。微调过程总共跨越160k次迭代，批处理大小为16。默认输入分辨率为512×512，我们还提供了使用640 × 640输入和多尺度（MS）测试的实验结果。