近年来，状态空间模型（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所示, 一张大小为
