点击下方卡片，关注 「AI视界引擎」 公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
选择性状态空间模型（SSMs），如Mamba（Gu & Dao，2023），在捕捉一维序列数据中的长期依赖性方面表现出色，但其应用于2D视觉任务仍面临挑战。当前的视觉SSM通常将图像转换为一维序列，并采用各种扫描模式来包含局部空间依赖性。
然而，这些方法在有效捕捉复杂图像空间结构和由于扫描路径加长导致的计算成本增加方面存在局限性。
为解决这些问题，作者提出了一种新颖的方法Spatial-Mamba，它直接在状态空间中建立邻域连通性。
作者不再仅依赖顺序状态转换，而是引入了一个_结构感知状态融合_方程，该方程利用膨胀卷积来捕捉图像空间结构依赖性，显著增强了视觉上下文信息的 Stream 。
Spatial-Mamba分为三个阶段：单向扫描的初始状态计算，通过结构感知状态融合获取空间上下文，以及使用观察方程计算最终状态。作者的理论分析表明，Spatial-Mamba将原始Mamba和线性注意力统一在相同的矩阵乘法框架下，为作者提供了更深入的理解作者的方法。
实验结果显示，即使只有一个扫描，Spatial-Mamba也达到了或超过了基于SSM的最先进图像分类、检测和分割模型。
源代码和训练好的模型可以在 https://github.com/EdwardChasel/Spatial-Mamba。
unset
unset1 Introductionunset
unset
状态空间模型（SSMs）是分析具有隐藏状态的动态系统的强大工具，它们在控制理论、信号处理和经济学等领域已长期使用。最近，由于其使用专门参数化的矩阵，状态空间模型被引入深度学习（Gu等人，2021年），尤其是在自然语言处理（NLP）中（Gu和Dao，2023年）。
这是状态空间模型的一项重大进步，即引入了选择性机制和并行计算的硬件感知优化，如Mamba（Gu和Dao，2023年），该模型根据序列中每个元素的相关性选择保留或丢弃信息，高效地建模具有线性复杂度的远程依赖关系。
Mamba在NLP领域的显著成功激发了研究行人研究如何将SSMs应用于视觉任务。与一维序列不同，视觉数据通常具有2D空间结构。因此，在适应Mamba中的信息选择和传播机制的同时，保持图像内的空间依赖性至关重要。现有的视觉SSMs通常使用一些扫描策略将2D视觉数据从不同方向扫描为多个1D序列，然后使用原始Mamba处理这些扫描得到的1D序列。这些扫描策略可以大致分为三类：扫描扫描、连续扫描和局部扫描，如图1（a）-1（c）所示。
Vim (Zhu et al., 2024) 和 V Mamba (Liu et al., 2024) 分别采用双向扫描和大方向扫描，如图1(a)所示。这些策略旨在降低空间方向敏感性，并调整网络结构以适应视觉任务。杨等人（2024）认为大方向扫描忽视了空间连续性的重要性，他们引入了一种连续扫描顺序，如图1(b)所示，以更好地整合视觉数据的归纳偏差。
黄等人（2024）认为扫描整个图像可能无法有效地捕捉局部空间关系。相反，他们提出了一种名为 LocalMamba 的方法，通过使用几种局部扫描模式，如图1(c)所示，将图像划分为不同的窗口以捕捉局部依赖性。除了上述扫描模式之外，还有其他扫描方法，如希尔伯特扫描（He et al., 2024）和动态树扫描（Xiao et al., 2024），也被提出以将 SSMs 适应到视觉任务。
虽然现有的扫描策略部分解决了将空间结构与顺序SSMs对齐的问题，但在模型有效性和效率上仍有局限。一方面，定向扫描不可避免地改变了像素之间的空间关系，破坏了图像的固有空间上下文。例如，在扫描过程中（图1（a）），一个像素与其左或右邻居之间的距离为1，而与顶部或底部邻居之间的距离等于图像的宽度。这种失真可能会阻碍模型在原始视觉数据中理解空间关系。另一方面，固定扫描路径，如常用的四方向扫描（图1（a）和图1（b）），在捕捉图像中复杂且变化的空间关系方面效果不够，同时引入更多的扫描方向也会导致计算过度。因此，探索如何设计更有效且结构感知的SSMs对于视觉任务至关重要。
为了实现这一目标，本文做出重要尝试，并提出了一种名为空间-Mamba 的方法，该方法旨在捕获潜在状态空间中相邻特征的空间依赖性。空间-Mamba的处理流程包括三个阶段。首先，如图1(d)左所示，将视觉数据转换为顺序数据，使用单向扫描。基于原始SSMs的状态转换方程计算状态变量，然后 Reshape 为视觉格式。
其次，将状态变量通过一个结构感知的状态融合 （SASF）方程处理，该方程使用膨胀卷积重新加权并合并附近的状态变量，如图1(d)左所示。最后，这些结构感知的状态变量被输入到观测方程中，产生最终输出变量。SASF方程不仅使序列中的非顺序元素之间实现有效的 Shortcut ，还增强了模型捕捉空间关系的能力，从而使底层视觉结构得到更准确的表示。
此外，作者证明空间-Mamba、原始Mamba和线性注意都可以在同一个框架下用结构矩阵表示，这为作者提出的算法提供了更连贯的理解。作者在图像分类、检测和分割等基本视觉任务上验证了空间-Mamba的优势。结果表明，即使只有一个扫描，空间-Mamba也实现了或超过了使用不同扫描策略的最近状态最先进技术。
unset
unset2 相关工作unset
unset
Gu等人（2021年）首次将线性状态空间层（LSSL）引入到HiPPO框架（Gu等人，2020年）中，以有效处理长序列中的长程依赖性。Gu等人（2021年）通过将参数表示为对角线加低秩矩阵，显著提高了SSMs的效率。所谓S4模型引发了结构化SSMs（Smith等人，2022年；Fu等人，2022年；Gupta等人，2022年；Gu和Dao，2023年）的一波。Smith等人（2022年）通过在S4层引入并行扫描，提出了S5，同时保持了S4的计算效率。最近，Gu和Dao（2023年）开发了Mamba，将数据相关的选择机制引入到S4层，以硬件友好的方式简化了计算和架构，实现了具有线性复杂度的Transformer类似建模能力。基于此，Dao和Gu（2024年）提出了Mamba2，揭示了SSMs与注意力之间的特定结构矩阵的联系。该框架将参数矩阵简化为标量表示，在不牺牲效率的情况下，探索更大更具有表达力的状态空间。
视觉SSM（单向递归神经网络）虽然在大规模自然语言处理任务中表现出色，捕捉时间依赖关系，但在处理视觉数据的多维空间结构方面存在局限性。这给开发有效的视觉SSM带来了挑战。S4ND（Nguyen等人，2022年）是首个基于SSM的多维数据模型，它通过独立的1D SSM将每个维度分离。Baron等人（2023年）将S4ND泛化为离散的多轴系统，并提出了2D-SSM空间层，成功地将1D SSM扩展到2D SSM。
近年来，越来越多的视觉SSM倾向于设计多个扫描顺序或模式以保持空间一致性，包括双向（Liu等人，2024年）、四向（Zhu等人，2024年）、连续（Yang等人，2024年）、锯齿状、窗口和拓扑扫描。这些视觉SSM已在多模态基础模型、图像修复、医学图像分析和其他视觉任务中得到应用，展示了SSM在视觉数据理解方面的潜力。
unset
unset3 Preliminariesunset
unset
SSMs（状态空间模型）通常用于分析序列数据和建模连续线性时不变（LTI）系统（Williams & Lawrence，2007）。一个输入序列 通过一个状态变量 被转换为输出序列 。在这里， 表示时间索引， 表示状态变量的维度。这个动态系统可以由线性状态转移和观测方程（Kalman，1960）描述：，其中 是状态转移矩阵，、， 控制系统的动力学。状态转移和观测方程描述了系统随时间演变的过程以及状态变量与观测输出之间的关系。
为了有效地将连续时间SSMs集成到深度学习框架中，必须将连续时间模型离散化。一种常用的技术是零阶保持（ZOH）离散化（Gu & Dao，2023）。ZOH方法通过在每个离散时间间隔内保持输入常数来近似连续时间系统。具体而言，给定一个时间尺度，表示离散时间步之间的间隔，并定义和为离散参数，则ZOH规则应用为和。
然而，实际过程中的过程往往随时间变化，不能准确地由LTI系统描述。正如Mamba（Gu & Dao，2023年）所指出的，时间变化的系统可以更关注相关信息，并能够提供更准确、更真实的动态系统表示。在Mamba中，通过选择性函数，使SSM的参数具有上下文感知和自适应能力。这通过将参数修改为输入序列的简单函数来实现，从而得到输入相关的参数和。然后，可以根据相应的输入相关的离散参数和计算。
因此，离散状态转移和观测方程可以计算如下：
如图2(a)所示，上述过程的简化说明如下。
Formulation of Spatial-Mamba
Spatial-Mamba旨在捕获潜在状态空间中相邻特征的空间依赖性。为实现这一目标，与之前的方法（如Zhu等人，2024；Liu等人，2024；Huang等人，2024）不同，作者引入了一种新的结构感知状态融合（SASF）方程到原始的Mamba公式。
Spatial-Mamba的整个过程可以用三个方程来描述：状态转换方程、SASF和观测方程，它们分别表示为：
图2（b）展示了在所提出的Spatial-Mamba中SSM（Spatial-Self-Attention Mechanism）的流程。与图2（a）中的原始Mamba相比，作者可以看到原始状态变量xt直接受到其前状态xt-1的影响，而结构感知状态变量ht通过融合机制整合了额外的邻居状态变量xtρt（t），xtρ2（t）…，xtρK（t），其中K表示邻居集的大小。考虑时间和空间信息，融合后的状态变量ht获得了更丰富的上下文，从而提高了适应性，并使图像得到了更全面的解释。
提出的空间-Mamba模型可以分为三个步骤实现。如图1(d)所示，输入图像首先被 flatten 成一个1D序列 ，利用状态转换方程 计算状态 。计算出的状态随后被 reshape 回2D格式。为了使每个状态都能感知到其在2D空间中的相邻状态，作者引入了SASF方程 。对于状态变量 ，作者在邻域 使用权重 对邻域中的状态 进行线性加权，以便有效地将局部依赖信息集成到新的状态 中，从而获得更具有上下文信息表示。最后，通过观察方程 从丰富的状态 生成输出。这种SASF方法有助于模型在视觉学习中集成局部结构信息，同时保留原始Mamba的优势。
在实际应用中，作者简单地采用多尺度膨胀卷积（Yu，2015）来线性加权相邻状态变量，增强空间关系描述，并实现 Shortcut 。具体而言，作者使用三个大小为 的逐点卷积滤波器，膨胀因子分别为 来构建邻居集合 。将公式（2）中的 SASF 方程重写为：
在位置的膨胀因子的卷积权重为，表示为位置的卷积权重。表示位于位置的状态的相邻状态。表示图像的宽度。
为了直观理解SASF，作者在图3中可视化了原始状态变量和作者提出的结构感知状态变量。可以看到，原始状态变量x_t（图3（b））在区分前景和背景上存在困难。相比之下，通过SASF过程进行优化的结构感知状态变量h_t（图3（c）），有效地将这两个区域分离开来。
此外，图3（d）中的原始状态变量x_t仅在扫描方向上显示出水平衰减（从左上角最亮值逐渐变暗），而融合后的状态变量h_t在图3（e）中展示了水平、垂直和对角方向的衰减。这种改进源于其能够利用图像内的空间关系，从而实现更准确和上下文敏感的表示。
Network Architecture
空间-Mamba的整体架构如图4所示。它由四个连续阶段组成，类似于Swin-Transformer（刘等人，2021）的结构。作者在不同尺度上引入了空间-Mamba模型的三种变体：空间-Mamba-T（微小）、空间-Mamba-S（小型）和空间-Mamba-B（基础）。详细的配置在附录A中提供。
具体来说，一个输入图像首先经过一个重叠的茎层处理，生成一个2D特征图，维度为。这个特征图然后被输入到四个连续阶段。每个阶段都包含多个空间-Mamba块，然后是一个下采样层，下采样因子为2（除最后一个阶段外），从而产生层次特征。最后，一个 Head 层被用来处理这些特征，为特定的下游任务生成相应的图像表示。
空间-Mamba块构成了作者架构的基本构建单元，该单元包括一个结构感知的单共享多维卷积（SASM）和具有 Shortcut 的 FFN （FFN），如图4左下角的示例所示。
在Mamba块设计的基础上（Gu和Dao，2023），在图4右下角的结构感知单共享多维卷积（SASM）的实现是通过用3x3的逐点卷积替换原始的1D因果卷积，并用作者提出的SASF模块替换原始的S6模块，实现了状态空间中的局部邻居连接，其线性复杂度。此外，在空间-Mamba块和FFN之前，作者还使用了一个局部感知单元（LPU）（Guo等人，2022），用于在图像块内提取局部信息。
Connection with Original Mamba and Linear Attention
作者深入分析了线性注意力（Katharopoulos等人，2020年）、原始Mamba（Gu和Dao，2023年）以及作者提出的空间-Mamba之间的相似性和差异性，以更好地理解作者提出的方法的运行机制。详细的推导见附录B。
线性自注意力 是一种改进的自注意力（SA）机制，通过使用核函数φ将SA的计算复杂度降低到线性。对于输入序列， Query 、键和值通过不同的权重矩阵进行投影而计算得到。在自回归模型中，为了避免模型关注未来 Token ，第个 Query 受到前序键的约束，即。因此，如果核函数φ是一个恒等映射，单头线性自注意力（不进行归一化）的计算可以表示为：。令，则作者得到。线性自注意力可以重写为。这表明线性自注意力实际上是线性递归的一个特殊情况。隐藏状态变量由外积更新，最终输出通过将与 Query 相乘得到。如果作者定义和，线性自注意力可以表示为类似于SSM的形式：
是线性变换的线性变换，即。
Mamba 本质上定义在公式（1）中。通过将初始状态变量 设置为零，状态变量可以递归地推导为 。在这里， 表示从 到 的状态转移矩阵的乘积，其值为 当且仅当 。
观察方程中没有 的最终输出可以重写为：
基于等式（2）和等式（5），作者从提出的SASF方程中得到的结构感知状态变量可以表示为 。为了简洁表达，可以省略 ，那么Spatial-Mamba的最终输出可以改写为：
备注 。从上述分析中，作者可以看出线性注意、Mamba和Spatial-Mamba这三种范式都可以在统一的矩阵乘法框架内进行建模，具体为。
差异在于矩阵的结构。对于线性注意和Mamba，的形式为下三角矩阵，而Spatial-Mamba的是一个邻接矩阵。图5提供了这些矩阵的可视化。在线性注意中，较亮位置的垂直方向保持一致，这表明SA机制关注一个小图像 Token 集合。
另一方面，Mamba显示随着时间的衰减模式，这是其状态转换矩阵的影响。这种动态转换允许Mamba在不同图像 Token 之间切换焦点。与线性注意和Mamba不同，作者的Spatial-Mamba考虑在更广泛的时空邻域上对所有状态进行加权求和，从而实现对空间关系的更全面表示。
实验结果 在本节中，作者进行了一系列实验来比较 Spatial-Mamba 与顶级基准模型，包括基于卷积神经网络（CNNs）的模型 ，视觉 Transformer ，以及最近的视觉SSMs 。
遵循以前的工作（刘等人，2021年；2024年），作者训练了 Spatial-Mamba 的三种变体，即 Spatial-Mamba-T，Spatial-Mamba-S 和 Spatial-Mamba-B。
性能评估是在基本视觉任务上进行的，包括图像分类、目标检测和语义分割。
Image Classification on ImageNet-1K
设置。 首先，作者在ImageNet-1K（Deng等，2009）上评估了Spatial-Mamba在图像分类中的表示学习能力。作者采用了以前工作中使用的实验配置（Liu等人，2021；2024），这些详细在附录A中说明。作者将作者的方法与最先进的算法进行了比较，包括RegNetY。
结果 。表1呈现了Spatial-Mamba与最新方法的全面比较。值得注意的是，Spatial-Mamba-T在Top-1精度上达到83.5%，与基于CNN的ConvNext-T相比提高了1.4%，在参数数量和GFLOPs相当的情况下。与基于Transformer的方法相比，Spatial-Mamba-T超过了Swin-T的2.2%和NAT-T的0.3%。与基于SSM的方法相比，Spatial-Mamba-T超过了VMamba-T的1.0%和LocalVMamba-T的0.8%。对于其他变体，Spatial-Mamba也显示出优势。
具体来说，Spatial-Mamba-S和Spatial-Mamba-B分别达到了Top-1精度84.6%和85.3%，分别比NAT-S和NAT-B提高了1.6%和1.0%，比VMamba-S和VMamba-B分别提高了1.0%和1.4%。虽然由于架构差异，Spatial-Mamba-T比VMamba-T稍微慢一些，但Spatial-Mamba的Small和Base变体比VMamba对应版本更快。此外，它们都比CNN和Transformer基础方法显著快。
Object Detection and Instance Segmentation on COCO
设置如下。作者在目标检测和实例分割任务中评估Spatial-Mamba，使用COCO 2017数据集（Lin等人，2014年）和MMDetection库（Chen等人，2019年）。作者采用Mask R-CNN（He等人，2017年）作为检测器头，应用在ImageNet-1K上预训练的Spatial-Mamba-T/S/B作为 Backbone 。
遵循常见做法（Liu等人，2021年，2024年），作者对预训练模型进行微调，训练12个周期（计划）和36个周期（计划）的多尺度输入。在训练过程中，采用AdamW优化器，初始学习率为0.0001，批量大小为16。
结果. 在COCO上的详细结果报告在表2中。可以看出，所有版本的Spatial-Mamba在不同的时间表上都优于其竞争对手。对于1×时间表，Spatial-Mamba-T在box mAP上达到47.6，在mask mAP上达到42.9，分别比Swin-T/VAMaba-T高4.9/0.3，且在参数和FLOPs上分别更少。同样，Spatial-Mamba-S/B在相同配置下也展现出优于其他方法的优秀性能。此外，随着3×多尺度训练时间表的改进性能趋势保持不变。值得注意的是，Spatial-Mamba-S在box mAP上达到50.5，在mask mAP上达到44.6，分别比VMamba-S高0.6和0.4。
Semantic Segmentation on ADE20K
设置。 为评估Spatial-Mamba在语义分割任务上的性能，作者使用广泛使用的UPerNet分割器 和MMSegmenation工具包来训练作者的模型。与以前的工作一致，作者在ImageNet-1K上预训练作者的模型，并将其用作UPerNet在ADE20K数据集上的基础进行训练。这个过程包括160000次迭代，批量大小为16。使用AdamW作为优化器，权重衰减为0.01。学习率设置为6×10^-5，线性学习率衰减。所有输入图像都裁剪为512×512。
结果 。语义分割的结果汇总在表3中。空间-Mamba变体始终实现了出色的性能。例如，空间-Mamba-T在单尺度上实现了48.6的mIoU，多尺度上实现了49.4的mIoU。这比NAT-T提高了1.5 mIoU，比VAMba-T提高了0.6 mIoU。在单尺度输入时，这种优势在多尺度输入时得到保持，空间-Mamba-T比NAT-T高1.0 mIoU，比VAMba-T高0.6 mIoU。此外，空间-Mamba-B在多尺度上实现了最佳性能，多尺度mIoU为52.6。
Ablation Studies
在本节中，作者在ImageNet-1K分类任务上消融了Spatial-Mamba-T在空间-Mamba-T上的各种关键组件。请参阅表4。根据Swin-T（Liu等人，2021年）的配置，作者构建了基准模型Spatial-Mamba-T，但不含SASF模块。该基准模型使用4x4卷积，步长为4作为基本层，并合并2x2邻域块进行下采样。
邻域集。 首先，对邻域集Ω的调整表明，将大小从3x3的邻域增加到5x5会导致准确度的逐步提高，从82.3%提高到82.5%，尽管相应的吞吐量会降低。此外，使用扩展邻域集，其中因子d=1，3，5，可以将准确度提高至82.7%，同时将吞吐量降低至1158。这表明存在速度与更大邻域集之间的权衡。
局部增强。 作者用重叠卷积（见表4中的'重叠茎'）替换原始非重叠茎和下采样层，从而提高了0.2%的准确性。作者还引入了LPU（Guo等人，2022），这是一种在每块和FFN顶部放置的深度卷积，进一步提高了准确性0.4%。这些修改丰富了在SASF模块处理图像块之前可用的局部信息，使其能更好地捕捉结构依赖关系。
优化 。为了进一步提高模型效率，作者利用再参数化技术（Ding等人，2022）实现SASF模块并优化CUDA核函数。这使模型速度至少提高30%，并将吞吐量从1065提升至1438。最后，通过集成MESA（Du等人，2022）以减轻过拟合，又额外获得了0.2%的准确率提升。
unset
unset6 Conclusionunset
unset
在本文中，作者提出了Spatial-Mamba，一种专为视觉任务设计的全新状态空间模型。
Spatial-Mamba的关键在于作者提出的结构感知状态融合（SASF）模块，该模块有效地捕获了图像空间依赖性，从而提高了上下文建模能力。
作者在图像分类、检测和分割等基本视觉任务上进行了广泛的实验。结果表明，使用SASF，Spatial-Mamba仅用一次信号扫描就超越了最先进的状态空间模型，显示出其强大的视觉特征学习能力。
作者还深入分析了Spatial-Mamba与原始Mamba和线性注意力的关系，并将其统一起来，置于相同的矩阵乘法框架下，为视觉表示学习提供了更深入的理解。
[0]. Spatial-Mamba: Effective Visual State Space Models via Structure-aware State Fusion.
点击上方卡片，关注 「AI视界引擎」 公众号