0. 论文信息
标题：QuadMamba: Learning Quadtree-based Selective Scan for Visual State Space Model
作者：Fei Xie, Weijia Zhang, Zhongdao Wang, Chao Ma
机构：上海交通大学、华为
原文链接：https://arxiv.org/abs/2410.06806
代码链接：https://github.com/VISIONSJTU/QuadMamba
1. 导读
状态空间模型的最新进展，特别是Mamba，已经证明了优于主流变压器模型的性能，特别是在将计算复杂度从二次降低到线性方面。然而，由于视觉数据的不同特征，例如图像内的空间位置和相邻性，以及跨视觉表征的信息粒度的大变化，使得在将Mamba从语言调整到视觉任务时出现困难。现有的vision Mamba方法要么以光栅扫描的方式将记号展平成序列，这破坏了图像的局部相邻性，要么将记号手动划分成窗口，这限制了它们的远程建模和概括能力。为了解决这些限制，我们提出了一种新的视觉Mamba模型，称为QuadMamba，它通过基于四叉树的图像分区和扫描来有效地捕捉不同粒度的局部依赖性。具体地，我们的轻量级的基于四叉树的扫描模块学习在学习的窗口象限内保持空间区域的2D局部性。在自适应地将记号划分成窗口象限之前，该模块根据它们的特征来估计每个记号的位置分数。一个全方位的窗口移动方案也被引入，以捕捉不同局部区域的更完整和更丰富的特征。为了使离散化的四叉树划分端到端可训练，我们进一步设计了基于Gumbel-Softmax及其直通梯度估计器的序列屏蔽策略。大量实验表明，QuadMamba在各种视觉任务中实现了最先进的性能，包括图像分类、对象检测、实例分割和语义分割。
2. 引言
结构化状态空间模型（SSM）架构近年来获得了广泛关注。SSM为序列建模提供了一种多功能方法，在计算效率和模型灵活性之间取得了平衡。受Mamba在语言任务中取得成功的启发，SSM在各种视觉任务中的应用也逐渐增多。这些应用从设计通用骨干模型到推动图像分割和合成等领域的发展。这些进展凸显了Mamba在视觉领域的适应性和潜力。
尽管SSM在长序列建模中具有吸引人的线性复杂度，但直接将SSM应用于视觉任务相较于流行的卷积神经网络（CNN）和视觉Transformer模型仅带来了有限的改进。在本文中，我们旨在拓展Mamba模型在计算机视觉中的应用。
我们观察到，语言和视觉领域之间的差异在将Mamba应用于后者时可能构成重大障碍。这些挑战源自图像数据的两个自然特性：1）视觉数据具有严格的二维空间依赖性，这意味着将图像块展平为序列可能会破坏高级理解。2）自然视觉信号具有大量空间冗余——例如，不相关的图像块不会影响对象的表示。
为了解决这两个问题，我们开发了一种针对视觉的特定扫描方法，以构建用于Vision Mamba的一维标记序列。
在本文中，我们介绍了一种新颖的Mamba架构，该架构通过关注更具信息量的区域来进行局部感知序列建模，从而学习改进局部表示。如图1（c）所示，QuadMamba的核心在于可学习的窗口划分，该划分能够自适应地学习以粗到细的方式建模局部依赖性。我们提议在视觉Mamba模型的多个层中使用一个轻量级预测模块，以评估每个空间标记的局部邻近性。得分最高的象限将以递归方式进一步划分为子象限进行精细扫描，而其他可能包含信息较少标记的象限则保持在较粗的粒度。此过程导致从二维图像特征中划分出不同粒度的窗口象限。
值得注意的是，直接根据索引从二维窗口图像特征中进行采样是不可微的，这使得窗口选择的学习变得难以处理。为了处理这个问题，我们采用Gumbel-Softmax从分区得分图中生成序列掩码。然后，我们使用完全可微的操作（即Hadamard乘积和逐元素求和）从序列掩码和局部窗口中构建一维标记序列。这些操作导致了一个端到端的可训练管道，且计算开销可忽略不计。对于跨越两个相邻四分之一窗口的信息标记，我们在连续块中应用全向移位方案。在两个方向上移位二维图像特征使得四分之一窗口划分能够灵活地建模出现在任意位置的对象。
3. 效果展示
通常，视觉Mamba模型需要将二维图像转换为一维序列进行处理。如图1（a）所示，直接将空间数据展平为一维标记的直观方法（例如Vim）会破坏自然的局部二维依赖性。LocalMamba通过将图像划分为多个窗口来改进局部表示，如图1（b）所示。在跨窗口遍历之前，分别对每个窗口进行扫描，确保同一二维语义区域内的标记被紧密地一起处理。然而，手工制作的窗口划分缺乏处理不同对象尺度的灵活性，并且无法忽略信息较少的区域。
4. 主要贡献
在ImageNet-1k和COCO2017上的大量实验表明，QuadMamba在图像分类、目标检测和分割任务上表现出色，相较于现有的CNN、Transformer和Mamba模型具有显著优势。例如，QuadMamba在ImageNet-1k上实现了78.2%的Top-1准确率，而其模型大小与PVT-Tiny（75.1%）和LocalViM（76.2%）相似。
5. 方法
QuadMamba与许多CNN和视觉Transformer具有相似的多尺度骨干设计。如图2所示，首先将图像I ∈ RHim×Wim×3划分为大小为4 × 4的图像块，得到N = H × W = ⌊Him/4⌋ × ⌊Wim/4⌋个视觉标记。一个线性层将这些视觉标记映射到维度为d的隐藏嵌入中，然后将其输入到我们提出的基于四叉树的视觉状态空间（QuadVSS）块中。与用于语言建模的Mamba结构不同，QuadVSS块遵循Transformer块的流行结构，如图2（b）所示。QuadMamba由四个阶段中组织的QuadVSS块的级联组成，阶段i（i ∈ {1, 2, 3, 4}）包含Si个QuadVSS块。在每个阶段中，下采样层将特征图的空间尺寸减半，同时将其通道维度加倍。得益于Mamba的线性复杂度，我们可以自由地在前两个阶段中堆叠更多的QuadVSS块，从而能够在引入最小计算开销的情况下充分利用其局部特征保留和建模能力。
6. 实验结果
表1展示了QuadMamba在准确性和效率方面的优越性。具体而言，QuadMamba-S在Top-1准确率上比RegNetT-8G高出2.5%，比DeiT-S高出2.6%，比Swin-T高出1.1%，同时其浮点运算次数（FLOPs）相当或更少。这一优势在比较具有相似参数数量或FLOPs的其他模型变体时也依然存在。与其他基于Mamba的视觉主干网络相比，QuadMamba在相似的网络复杂度下也表现出优越的性能。例如，QuadMamba-B（83.8%）的性能与VMamba-B（83.7%）、LocalVim-S（81.2%）和PlainMamba-L3（82.3%）相当或更优，但其参数数量和FLOPs却显著更少。这些结果证明了QuadMamba在性能和复杂度上的优越性，以及其作为强大且高效的视觉主干网络的潜力。此外，QuadMamba在性能上达到或超过了LocalMamba，同时完全避免了后者昂贵的架构和扫描策略搜索，这使得它成为更加实用且多功能的视觉主干网络选择。
在目标检测和实例分割任务中，QuadMamba在具有相似复杂度的模型和架构中（以网络参数数量和FLOPs衡量）表现出极高的效率。在参数少于或大约30M的小型主干网络中，QuadMamba的竞争对手寥寥无几。如表2所示，QuadMamba-T不仅在目标检测和实例分割任务上大幅优于ResNet18和PVT-T，还在目标检测上领先EfficientVMamba-S 3.0%的平均精度均值（mAP），在实例分割上领先2.1%的mAP。在更大的主干网络中，QuadMamba再次超越了所有基于卷积神经网络（ConvNet）、Transformer和Mamba的竞争对手。值得注意的是，QuadMamba-S在目标检测和实例分割任务上分别比以高效率为特点的基于Mamba的主干网络EfficientVMamba-B高出3.0%和2.2%的mAP，而参数数量相当。此外，QuadMamba-S能够保持甚至超越LocalVMamba-T的性能，同时避免了表2的复杂度测量所未反映的大量架构和扫描搜索麻烦。这些结果表明，QuadMamba在计算复杂度、设计成本和性能之间做出了务实的权衡，可以作为一个强大且多功能的视觉主干网络。
推荐课程：
面向自动驾驶领域的3D点云目标检测全栈学习路线！(单模态+多模态/数据+代码)
。
如表3所示，在相似的网络复杂度和效率下，QuadMamba的分割精度显著高于基于ConvNet的ResNet-50/101和ConvNeXt、基于Transformer的DeiT和Swin Transformer，以及大多数基于Mamba的架构。例如，QuadMamba-S以47.2%的平均交并比（mIoU）报告了比Vim-S（44.9%）、LocalVim-S（46.4%）、EfficientVMamba-B（46.5%）和PlainMamba-L2（46.8%）更高的分割精度，并且与VMamba-T（47.3%）的结果相当。与LocalMamba-S/B相比，QuadMamba-S/B虽然略有落后，但无需承担额外的网络搜索成本。值得注意的是，LocalMamba是通过神经架构搜索（NAS）技术设计的，该技术依赖于数据，且对其他数据模态和数据源缺乏灵活性。
7. 总结 & 未来工作
在本文中，我们提出了QuadMamba，这是一种视觉Mamba架构，可作为图像分类和密集预测等视觉任务的多功能且高效的主干网络。QuadMamba通过可学习的基于四叉树的扫描有效地捕获了不同粒度的局部依赖性，同时以可忽略的计算开销自适应地保留了图像数据中的固有局部性。通过广泛的实验和消融研究，已经证明了QuadMamba的有效性，其性能优于流行的卷积神经网络（CNN）和视觉Transformer。然而，QuadMamba的一个限制是尚未探索超过两级的窗口划分，这可能对于处理密集预测视觉任务和更高分辨率的数据（如遥感图像）特别相关。细粒度的划分区域是刚性的，并且在关注任意形状和大小的区域时缺乏灵活性，这留待我们未来的工作来解决。我们希望我们的方法能够激励进一步研究将Mamba应用于更多样化和更复杂的视觉任务。
对更多实验结果和文章细节感兴趣的读者，可以阅读一下论文原文~
本文仅做学术分享，如有侵权，请联系删文。
3D视觉交流群，成立啦！
目前我们已经建立了3D视觉方向多个社群，包括
2D计算机视觉
、
最前沿
、
工业3D视觉
、
SLAM
、
自动驾驶
、
三维重建
、
无人机
等方向，细分群包括：
工业3D视觉
：相机标定、立体匹配、三维点云、结构光、机械臂抓取、缺陷检测、6D位姿估计、相位偏折术、Halcon、摄影测量、阵列相机、光度立体视觉等。
SLAM
：视觉SLAM、激光SLAM、语义SLAM、滤波算法、多传感器融合、多传感器标定、动态SLAM、MOT SLAM、NeRF SLAM、机器人导航等。
自动驾驶：深度估计、Transformer、毫米波|激光雷达|视觉摄像头传感器、多传感器标定、多传感器融合、自动驾驶综合群等、3D目标检测、路径规划、轨迹预测、3D点云分割、模型部署、车道线检测、Occupancy、目标跟踪等。
三维重建
：3DGS、NeRF、多视图几何、OpenMVS、MVSNet、colmap、纹理贴图等
无人机
：四旋翼建模、无人机飞控等
2D计算机视觉
：图像分类/分割、目标/检测、医学影像、GAN、OCR、2D缺陷检测、遥感测绘、超分辨率、人脸检测、行为识别、模型量化剪枝、迁移学习、人体姿态估计等
最前沿
：具身智能、大模型、Mamba、扩散模型等
除了这些，还有
求职
、
硬件选型
、
视觉产品落地、产品、行业新闻
等交流群
添加小助理: dddvision，备注：
研究方向+学校/公司+昵称
（如
3D点云+清华+小草莓
）, 拉你入群。
▲长按扫码添加助理：cv3d008
3D视觉知识星球
