新的 Mamba 架构又双叒来了！上交大 & 华为提出 QuadMamba ：多个视觉任务 SOTA!

近年来，状态空间模型（State Space Models）中的Mamba取得了显著的进展，在性能上超过了主导的Transformer模型，尤其是在降低计算复杂度方面，从二次降到了一次。然而，由于视觉数据的独特特性，如图像中的空间局部性和邻接关系以及视觉 Token 的信息粒度变化很大，Mamba从语言任务转换到视觉任务时存在困难。现有的视觉Mamba方法要么将 Token 扁平化为在光栅扫描方式下的序列，这破坏了图像的局部邻接关系，要么手动将 Token 分区到窗口，这限制了它们的远程建模和泛化能力。

为了解决这些限制，作者提出了一种新的视觉Mamba模型，称为QuadMamba，它通过基于四叉树（quadtree）的图像分区和扫描有效地捕获了不同粒度局部依赖性。

具体来说，作者的轻量级四叉树扫描模块学习到在学习的窗口四角内保持空间区域2D局部性。模块根据每个 Token 的特征估计局部性分数，然后适当地将 Token 分区到窗口四角。作者还引入了一种全方位窗口平移方案，以捕获不同局部区域之间的更多完整和有用的特征。为了使离散化的四叉树分区端到端可训练，作者进一步设计了一种基于Gumbel-Softmax及其直接导数序列 Mask 策略。大量实验表明，QuadMamba在各种视觉任务中实现了最先进的性能，包括图像分类、目标检测、实例分割和语义分割。

代码已在https://github.com/VISION-SJTU/QuadMamba。

1 Introduction

结构化状态空间模型（SSMs）的架构在近年来得到了显著的普及。SSMs为序列建模提供了一种灵活的方法，在计算效率和模型灵活性之间实现了平衡。受到Mamba在语言任务中成功的影响，越来越多地使用SSMs进行各种视觉任务。这些应用范围从设计通用的 Backbone 模型 到推进图像分割 和合成（如[17]）等领域。这些进步突显了Mamba在视觉领域的适应性和潜力。

尽管在长序列建模方面，SSMs具有令人满意的线性复杂度，但直接将SSMs应用于视觉任务只能带来微小的改进，这相对于常见的CNN和视觉Transformer模型。在本文中，作者试图扩大Mamba模型在计算机视觉领域的适用性。作者观察到语言和视觉领域之间的差异可能会在将Mamba应用于后者时带来巨大的障碍。

这些挑战源于图像数据的两个自然特性：

1）图像数据具有严格的2D空间依赖性，这意味着将图像块平铺成序列可能会破坏高级理解。

2）自然视觉信号具有严重的空间冗余性--例如，无关的图像块不会影响物体的表示。为了解决这两个问题，作者开发了一种专用的扫描方法，为视觉Mamba构建1D Token 序列。通常，视觉Mamba模型需要将2D图像转换为1D序列进行处理。

如图1(a)所示，直接将空间数据平铺成1D Token 会破坏自然局部2D依赖关系。LocalMamba通过将图像分割成多个窗口来提高局部表示，如图1(b)所示。每个窗口在单独扫描后，在窗口之间进行遍历，确保相同2D语义区域内的 Token 被紧密地处理在一起。然而，手工制作的窗口分割缺乏处理不同物体尺度的灵活性，并且无法忽略信息较少的区域。

在本研究中，作者提出了一种名为QuadMamba的新的Mamba架构，该架构通过关注更多具有信息性的区域来提高局部表示，从而实现局部感知序列建模的改进。如图1(c)所示，QuadMamba的核心在于可学习的窗口划分，该划分能够自适应地以粗粒度到细粒度的方式建模局部依赖关系。作者提出在视觉Mamba模型中的多个轻量级预测模块，用于评估每个空间 Token 的局部邻接关系。得分最高的四分之一区域进一步以递归的方式分割为子四分之一区域进行细粒度扫描，而其他区域（可能包含信息性较弱的 Token ）保持粗粒度。这一过程导致了从二维图像特征中分得的不同粒度的窗口四分之一区域。

值得注意的是，直接从基于索引的二维窗口图像特征中进行采样是非可微分的，这使得学习窗口选择变得不可行。为了解决这个问题，作者采用了Gumbel-Softmax，从分区间得分映射中生成一个序列 Mask 。然后，作者使用全可微分运算符（即Hadamard乘法和逐元素求和），从序列 Mask 和局部窗口构造一维 Token 序列。这导致了端到端可训练的流水线，计算开销可以忽略不计。对于跨越两个相邻四分之一窗口的有用 Token ，作者在连续块中应用了全方位位移方案。在两个方向上移动二维图像特征允许四分之一窗口分区间在任意位置出现的目标建模更加灵活。

1 相关工作

Generic Vision Backbones

卷积神经网络（CNNs）[10; 30; 31]和视觉 Transformer （ViT）[7]是计算机视觉领域两种主导的后备网络。它们在广泛的计算机视觉任务中，包括但不局限于图像分类[29; 53; 55; 20; 23; 21; 74; 4; 12]，分割[44; 19]，目标检测[36; 79]，视频理解[28; 76]，生成[11]等方面，都证明了自己是通用的视觉 Backbone 网络。与CNNs的受限制的感知域不同，视觉 Transformer （ViT）[7; 42; 61]从语言任务中借用，在全局上下文建模方面具有优势。后来，为了更好地适应视觉领域，提出了许多视觉特定的修改，如引入层次特征[42; 61]，优化训练[58]，以及将CNN元素集成[5; 54]。因此，视觉 Transformer 在各种视觉应用中表现出色。然而，这以注意力操作的平方时间复杂度和内存复杂度为代价，尽管提出了[42; 61; 72; 57]的补救措施，但其可扩展性仍然受到阻碍。

近年来，状态空间模型（SSMs）作为一种强大的范式，在语言任务中建模顺序数据。先进的SSM模型与最先进的视觉 Transformer （ViT）架构相比，在性能上甚至更优，同时具有线性复杂度。它们在视觉任务上的初步成功，更重要的是，惊人的计算效率，暗示了SSM作为CNN和Transformer的有前途的通用后端替代品的潜力。

2 State Space Models

状态空间模型（SSMs）[16; 15; 18; 35] 是一类用于序列建模的完全循环架构。最近的研究 使得 SSMs 的性能达到了 Transformer  Level ，同时其复杂度呈线性增长。作为一项重大突破，Mamba [13] 革新了传统的 SSM，采用了输入相关的参数化方法，并支持可扩展的、面向硬件优化的计算，使其在涉及顺序 1D 数据的多种任务中，与高级 Transformer 模型相媲美或表现更优。

随着Mamba的成功，ViM [80]和VMamba [41]将Mamba的1D扫描转换为双向和四向的2D交叉扫描，以便处理图像。此后，SSMs已迅速应用于视觉任务（语义分割[51; 65; 46]，目标检测[26; 3]，图像修复[17; 52]，图像生成[9]，等）以及其他模态的数据（例如，视频[67; 32]，点云[40; 73]，图[2]，以及跨模态学习[60; 6]。

在将Mamba适应非一维数据时，一个基本考虑因素是设计一个路径，该路径遍历并映射所有图像块到一个SSM友好的1D序列。在这个方向上，初步的工作包括ViM中的双向ZigZag扫描[80]，VMamba中的4个方向交叉扫描[41]，以及PlainMamba和ZigMa中的蛇形扫描[66;22]，这些工作都是在高度和宽度轴所覆盖的空间域中进行的。其他工作[52;75;33]将扫描扩展到额外的通道[52;33]或时间[75;33]维度。然而，这些扫描策略在遍历块时忽视了空间局部性的重要性。LocalMamba[26]部分地缓解了这一固有弱点，它将块分成窗口并在每个窗口内进行遍历。

然而，由于整个图像域的单一致局部细粒度，由任意的窗口大小控制，很难确定最优粒度。LocalMamba选择DARTS [38]以分别搜索每个层的最佳窗口大小和最佳扫描方向，这使得方法变得更加复杂。另一方面，所有现有方法都涉及硬编码的扫描策略，这些策略可能是次优的。与所有这些方法不同，本文引入了一种可学习的四叉树结构来扫描具有不同局部细粒度的图像块。

3 Preliminaries

状态空间模型（SSMs）[16;15;18;35]本质上是一种线性时不变系统，它通过隐藏状态 （序列长度 和状态大小 ）将一维输入序列 循环映射到输出响应序列 (其中 和 分别为序列长度和状态大小)。从数学上讲, 这样的系统可以表示为以下常微分方程（ODEs）：

其中, 矩阵 包含了演化参数; 和 是投影矩阵。然而, 在实际中, 通过零阶保持 (ZOH) 规则[16], 方程 1 通常被转化为离散形式[18], 其中 的值在样本间隔 内保持不变。离散化的ODEs可以表示为:

在上述内容中, 和 是 和 的离散表示: , 并且 . 为实现高效计算, 方程 2 中的迭代计算可以并行执行, 同时进行全局卷积操作。

其中 表示卷积运算符， 表示SSM核。

选择状态空间模型（S6）。 传统状态空间模型（SSMs）具有输入无关的参数。为了改进这一点，提出了选择状态空间模型（S6）或“Mamba”，它们具有输入相关的参数，使得 和 变得可学习。为了弥补并行性困难，还进行了硬件感知的优化。在本工作中，作者特别研究了Mamba架构在视觉任务中的有效适应性。

早期的工作，如ViM [80]和V Mamba [41]，通过将2D图像转换为以光栅扫描方式的一维序列来探索直观的适应性。作者认为简单的光栅扫描并不是最优设计，因为它破坏了图像的局部邻域。在作者的工作中，提出了一种基于四叉树的新颖可学习扫描方案。

4 Method

General Architecture

QuadMamba 共享了与许多卷积神经网络（CNNs）[20; 64; 23]和视觉 Transformer （Vision Transformers） 相似的多尺度背身设计。如图2所示, 一张大小为 的图像首先被分割成大小为 的patch, 从而得到 个视觉tokens。一个线性层将这些视觉tokens映射到具有维度 的隐藏嵌入，然后被输入到作者提出的基于Quadtree的视觉状态空间（QuadVSS）块中。

与在语言建模中使用的Mamba结构不同, QuadVSS块遵循了 Transformer 块的流行结构 , 68], 如图2(b)所示。QuadMamba由四个阶段的QuadVSS块组成，其中阶段 有 个QuadVSS块。在每个阶段中, 一个降采样层将特征图的空间大小减半, 同时将它们的通道维度加倍。由于Mamba的线性复杂度, 作者可以在前两个阶段中堆叠更多的QuadVSS块, 这使得它们的局部特征保持和建模能力可以充分利用, 同时引入的计算开销最小。

Quadtree-based Visual State Space Block

如图2所示，作者的QuadVSS块采用了视觉Transformer的元架构[68]，由一个 Token 运算符、一个 FFN （FFN）和两个残差连接组成。Token 运算符包括一个位移模块、一个分区映射预测器、一个基于四叉树扫描器和一个Mamba Layer。在 Token 运算符内部，一个轻量级预测模块首先在特征 Token 上预测一个分区映射。四叉树策略然后通过递归地将二维图像空间分成四个象限或窗口来对它进行细分。根据粗粒度分区映射的得分，跳过较不具有信息量的细粒度子窗口。

因此，一个多尺度、多粒度的1D Token 序列被构建，在更具有信息量的区域捕获更多的局部性，同时保留其他区域的全球上下文建模。QuadVSS块的关键组件如下：

分区映射预测。图像特征 ，包含 个嵌入 Token，首先被投影为得分嵌入 :

φs 是一个轻量级的 Projector ，具有规范线性 GELU 层。为了更好地评估每个 Token 的局部性，作者利用每个象限内的本地嵌入和上下文信息。具体而言，作者首先在通道维度上对 xs 进行拆分，得到局部特征 xs^local 和全局特征 xs^global:

因此, 作者得到了聚合得分嵌入 , 并将其输入到线性 GELU 层 中, 用于预测分段得分：where Interpolate 是双线性插值运算符，将上下文向量插值到空间大小为 的位置。

从而, 作者得到了聚合得分嵌入 , 并将其输入到线性 GELU 层 中, 用于预测分段得分:

其中 表示在空间坐标 的 Token 的配分。基于四叉树（Quadtree）的窗口划分。在预测每个特征 Token 的划分得分 后，作者应用了一个快速的四叉树基窗口划分策略, 计算成本可以忽略不计。作者构建了双级窗口四象限，从粗粒度到细粒度捕捉空间局部性。

四叉树基策略将图像特征在粗粒度 Level 划分为 个子窗口, 在细粒度 Level 划分为 个子窗口。与Transformer方案不同, 它仅通过将 Query 特征设置为相同大小来保持图像特征的空间形状，而由QuadMamba学习的扫描生成的 Token 序列应该与具有原始空间大小的特征合并。

因此, 作者选择在粗粒度 Level 具有最高平均局部邻域得分的顶点 四象限, 并将其进一步划分为四个子窗口:

其中， 包含在粗粒度窗口和细粒度窗口中的样本，并发送到 SS2D 块进行序列建模。

为了考虑最具有信息量的 Token 跨越相邻的两个窗口四分之一，作者借用了Swin Transformer [42]中的移位窗口方案。不同之处在于，Swin Transformer在窗口内的每个 Token 内部都忽略了空间局部性，而Mamba窗口内的 Token 序列仍然具有方向性。

因此，在后续的VSS模块中，作者添加了额外的移位方向，如图4所示，与Swin Transformer中只有一个方向移位相比。

Model Configuration

值得注意的是，QuadMamba模型的容量可以通过调整输入特征维度和（Q）VSS层数量来定制。在本研究中，作者构建了四种具有不同容量的QuadMamba架构变体，分别为QuadMamba-Li/T/S/B：

值得注意的是，QuadMamba-S在目标检测上比EfficientVMamba-B（一种基于Mamba的 Backbone 网络，以其更高的效率而著称）提高了3.0%的mAP，在实例分割上提高了2.2%，使用相似的参数。此外，QuadMamba-S能够跟上并甚至超越LocalVMamba-T（一种局部Mamba Backbone 网络）的表现，同时避免了表2的复杂性测量所反映的架构和扫描搜索的麻烦。

这些结果表明，QuadMamba可以作为一种实用且功能强大的视觉 Backbone 网络，在计算复杂度、设计成本和性能之间实现平衡。

Semantic Segmentation on ADE20K

如图3所示，在可比的网络复杂度和效率下，QuadMamba相较于基于ConvNet的ResNet-50/101 [20]和ConvNeXt [43]，基于Transformer的DeiT [58]和Swin Transformer [42]，以及大多数基于Mamba的架构 [80, 49, 66]，实现了显著更高的分割精度。

例如，QuadMamba-S在 的mloU下报告了优越的分割精度，超过了Vim-S ( ）， LocalVim-S (46.4%), EfficientVMamba-B (46.5%)，以及PlainMamba-L2 (46.8%)，并且与 VMamba-T ( )具有竞争力的结果。与LocalMamba-S/B相比, QuadMamba-S/B虽然稍逊一筹, 但仍具有不产生额外网络搜索成本的优势。

值得注意的是，LocalMamba是由神经架构搜索（NAS）技术设计的，该技术依赖于数据，并且与其他数据模式和数据源缺乏灵活性。

Ablation Studies

作者从多个角度进行了消融实验来验证QuadMamba设计选择的正确性。除非另有说明，所有实验都使用了QuadMamba-T。

Mamba中局域性的影响。 作者在构建一维 Token 序列时，考虑了粗粒度和细粒度局部性建模对模型性能的影响。具体而言，作者将[80, 41]中的原始窗口无损平铺策略与三个尺度（即28x28、14x14、2x2）的窗口划分（即特征局部性的三个粒度 Level ）进行了比较。在实际中，作者将QuadMamba-T模型中的QuadVSS块替换为[41]中的简单VSS块。

为了排除填充操作的负面影响，作者在第一个模型阶段只划分了空间尺寸为56x56的特征。如表4所示，原始扫描策略导致与采用窗口扫描相比，目标检测和实例分割性能显著降低。局部窗口的规模也显著影响了模型性能，这表明在给定图像分辨率的情况下，太大的或太小的窗口可能是次优的。

基于四叉树的分区分辨率。作者探讨了在双层四叉树分区间策略中分区分辨率的选取。表格5中配置的分区分辨率在第一阶段两个模型中应用，特征分辨率为{56×56, 28×28, 14×14}。实验中作者推导出粗粒度窗口和细粒度窗口的最优分辨率为{1/2, 1/4}。这种手工配置的设置可能在未来的工作中被更灵活和可学习的设置所替代。

分层模型阶段的层模式 作者研究了在分层模型 Pipeline 中的层模式设计选择。从图5可以看出，层模式LP2在减少了0.2%的准确率上超过了层模式LP1。这可能是因为局部建模效果在较浅的阶段比在较深的阶段更明显，以及第三阶段的填充操作产生了负面影响。LP3在第一和第二阶段将QuadVSS模块以交错方式放置，实现了最佳性能，并作为作者的模型设计。

必要性多向窗口平移。与Swin Transformer [42]中的单向平移不同，图5显示了在添加互补平移方向时，准确率提高了0.2%。这可以预见，因为Transformer中的注意力是非因果的，而Mamba中的一维序列具有因果性质，因此它们对相对位置非常敏感。在处理信息区域跨越相邻窗口的情况时，多向平移操作也是必不可少的。图6进一步可视化了在不同层级的层次结构中学习的细粒度四象限中的平移，这些四象限在不同层适当地关注不同的空间细节。

每个阶段使用的（Quad）VSS模块数量 作者进行了实验来评估每个阶段不同数量（Quad）VSS模块的影响。表6呈现了遵循图5中设计规则LP3的四个配置，固定通道维度为96。作者发现，庞大的第二或第四阶段会导致与第三阶段设计相比性能下降，而将（Quad）VSS模块在第二和第三阶段之间分配得更加均匀，可以获得与有利的复杂性相当甚至更好的性能。这些证据可以作为未来模型设计的一条基本规则，尤其是在模型扩展时。

6 Conclusion

在本文中，作者提出了一种名为QuadMamba的视觉Mamba架构，作为一种通用且高效的后端，用于诸如图像分类和密集预测等视觉任务。

QuadMamba通过可学习的四叉树扫描有效地捕获了不同粒度的局部依赖关系，同时适应地保留了图像数据的固有局部性，且计算开销极低。

QuadMamba的有效性已经通过大量实验和消融研究得到证明，其性能优于流行的卷积神经网络（CNNs）和视觉 Transformer （ViTs）。

然而，QuadMamba的一个局限性是窗口分级的深度尚未探索，这可能特别适用于处理密集预测视觉任务和高分辨率数据，如遥感图像。细粒度分区的区域刚性且缺乏针对任意形状和大小的区域的灵活性，这留待作者未来研究。

作者希望作者的方法能激发进一步将Mamba应用于更多样化和复杂的视觉任务的研究。

参考文献

[0]. QuadMamba: Learning Quadtree-based Selective Scan for Visual State Space Model.