点击下方卡片，关注
「AI视界引擎」
公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
状态空间模型（SSMs）近年来在计算机视觉领域引起了广泛关注。然而，由于图像数据的独特特性，将SSMs从自然语言处理领域调整到计算机视觉领域，并未超越最新的卷积神经网络（CNNs）和Vision Transformers（ViTs）。
现有的视觉SSMs主要依赖于手动设计的扫描方法，将图像块局部或全局地扩展为序列。这种方法破坏了图像原有的语义空间邻接关系，并缺乏灵活性，难以捕捉复杂的图像结构。
为了克服这一限制，作者提出了一种名为动态自适应扫描（DAS）的数据驱动方法，它可以自适应地分配扫描顺序和区域。这使得模型具有更大的灵活性，同时保持线性计算复杂度和全局建模能力。
基于DAS，作者进一步提出了视觉主干网DAMamba，该模型在图像分类、目标检测、实例分割和语义分割等视觉任务中显著优于当前最先进的Mamba模型，并且在某些情况下甚至超过了最新的CNNs和ViTs。
相关代码将会在
https://github.com/ltzovo/DAMamba发布。
- Introduction
近年来，为了克服传统卷积神经网络（CNN）（LeCun等人，1998年）在建模长程依赖关系方面的局限性，Transformer（Vaswani等人，2017年）被引入到计算机视觉领域，并在图像分类任务中取得了最先进的性能。然而，由于自注意力机制具有的二次计算复杂度，在高分辨率视觉下游任务如目标检测和图像分割中应用时受到了限制。为了解决这一问题，研究行人提出了各种稀疏注意力机制。这些机制通过在注意力计算中引入稀疏性来减少复杂度，但通常会以牺牲模型的整体建模能力为代价，从而限制了其在实际应用中的性能。
状态空间模型（SSMs）（Gu等，2021），由Mamba 代表，最近吸引了大量研究行人的关注。核心模块S6块根据序列中每个元素的相关性选择性地保留或丢弃信息。通过结合并行计算的选择机制和硬件感知优化，S6块不仅提供了全面的全局感受野，还实现了与序列长度线性相关的计算复杂度。这一特性使Mamba在自然语言处理任务中能够超越流行的Transformer模型。受Mamba的启发，一些研究工作，如Vim（Zhu等，2024）和VMamba（Liu等，2024），将其应用扩展到了计算机视觉领域。这些方法将二维图像分割为patches，并采用特定的扫描策略从不同方向将图像分解为多个一维序列。这成功地将Mamba模型整合到视觉任务中，取得了令人满意的性能，并展示了SSMs在计算机视觉领域的潜力。
与一维序列语言数据不同，视觉数据通常表现出二维空间结构。将 Mamba 模型适应于视觉任务的核心挑战之一在于设计适用于图像的适当扫描策略。扫描策略使得专为一维序列处理设计的序列Self-Attention模型（SSMs）能够有效应对图像的二维空间结构。目前，视觉 SSM 的扫描策略可以大致分为三种类型：扫视扫描、连续扫描和局部扫描。Vim 和 VMamba 采用扫视扫描策略，模拟从左到右、从上到下逐行扫描的方式，使 Mamba 能够适应图像的二维空间结构。然而，PlainMamba（Yang 等人，2024）认为扫视扫描忽略了图像内部的空间连续性，因此引入了连续扫描策略以确保相邻 patch 之间的相关性。与此同时，LocalMamba（Huang 等人，2024）提出了局部扫描策略，旨在捕捉图像内的局部空间关系。
尽管上述方法在实践中证明是有效的，但它们依赖于手动设计的扫描模式，这些模式与输入数据无关，可能并不最优。例如，在SSM计算中，扫掠和平滑连续扫描会导致空间上接近的Patch变得距离遥远，从而损失局部信息。另一方面，局部扫描能够捕捉局部空间关系，却限制了模型捕捉长范围依赖的能力。显然，需要一种更灵活的扫描策略，可以根据每个输入数据实例的特性动态调整扫描区域。例如，在处理狗的图像时，理想的扫描策略应当自适应地关注狗的身体，而过滤掉背景中的无关信息。然而，这种动态调整超出了现有手动设计扫描方法的能力范围。
为了解决上述问题，作者提出了一种灵活且高效的扫描策略，名为动态自适应扫描（DAS）。与传统的手工设计的扫描方法不同，DAS 在训练过程中能够动态地学习和调整扫描区域及其顺序，从而实现更智能、更精确的功能提取。具体而言，DAS 首先定义了一组可学习的位置，这些位置的初始值对应于每个块的原始位置。然后通过一个可学习的偏移预测网络（OPN），为每个块生成一组偏移值。通过将这些偏移值与原始块位置结合，计算出预测的块位置。使用双线性插值，这些预测位置可以通过梯度与特征图链接起来，在训练过程中实现偏移值的自适应优化。根据原始位置，预测的块从上到下、从左到右排列，动态地形成新的序列顺序，以适应输入数据。通过这种方式，DAS 能够更加关注关键区域，以更高的灵活性捕捉重要的特征和复杂的空间结构。
基于提出的DAS，作者开发了一种强大的视觉Mamba模型，称为DAMamba。DAMamba可以作为各种视觉任务的多功能视觉 Backbone 。例如，作者的DAMamba-T在图像分类中的准确率为83.8%，目标检测中的APb值为48.5，在实例分割中的APm值为43.4，以及语义分割中的mIoU值为50.3。这些结果在分类准确率上超过了之前的前沿视觉Mamba（VMamba）1.2%，在APb上提高了1.2，在APm上提高了0.7，在mIoU上提高了2.3。此外，DAMamba在这几类视觉任务中也优于一些最近的前沿ViTs和CNNs，展示了其卓越的性能和广泛应用性。
- Related Work
2.1. Vision State Space Models
尽管Transformer（Vaswani等，2017）在自然语言处理领域取得了显著的成功，但其二次复杂度在处理长序列结构时带来了挑战。为了应对这一问题，状态空间模型（Gu等，2021）（SSMs），例如Mamba（Gu & Dao，2023）逐渐成为Transformer的一种替代方案。在视觉任务中，标准的自注意力机制同样存在二次复杂度的问题，这对处理高分辨率图像构成了挑战。因此，Vim（Zhu等，2024）和VMamba（Liu等，2024）尝试将Mamba引入计算机视觉任务中。
然而，将图像输入到SSM模型中仍然是一个关键挑战。Vim和VMamba通过采用双向扫描和四向扫描策略，将图像块转换为一维序列来解决这个问题。在此基础上，后续研究引入了连续扫描（Yang等，2024）和局部四向扫描（Huang等，2024），以更好地与图像的二维结构相匹配。尽管Mamba模型在计算机视觉方面取得了显著成就，但现有的扫描方法主要依赖于手动设计，难以动态且灵活地适应输入的变化。这种限制阻碍了模型捕捉复杂二维结构的能力。因此，作者的目标是提出一种能够根据输入图像自适应和灵活调整扫描路径的视觉Mamba模型，进一步提升其在视觉任务中的性能。
2.2. Vision Transformers
Transformer（瓦斯尼等，2017）模型最初在2017年为自然语言处理（NLP）任务引入。凭借其强大的全局建模能力和出色的并行性能，Transformer很快在NLP领域获得了广泛认可。截至2020年底，视觉Transformer（多索夫斯基等，2021）（ViT）成功将Transformer模型扩展到大规模图像分类任务中，并取得了最先进的性能。
随后，DeiT（图弗罗恩等，2021）通过引入知识蒸馏（亨廷顿等，2015）和更高效的训练策略，改进了ViT，即使在相对较小的数据集如ImageNet1K（鲁萨科夫斯等，2015）上也能实现有效的训练。沿着这一发展轨迹，研究行人提出了多种减少高分辨率图像计算复杂度的分层Transformer模型，通过各种稀疏注意力机制实现。值得注意的例子包括Swin Transformer（刘等，2021）和PVT（王等，2021；2022）。后续研究（王等，2021；2022；夏等，2022；董等，2022；朱等，2023；贾0等，2023；张等，2024；施，2024）引入了多种稀疏注意力机制，以平衡全局建模能力和计算复杂性。然而，这些改进的稀疏注意力机制的全局建模能力仍然不及标准的自注意力机制。
2.3. Convolutional Neural Networks
卷积神经网络（CNN）（LeCun等，1998年）最初用于手写数字识别，但在2012年AlexNet（Krizhevsky等，2012年）的引入引发了“ImageNet时刻”之后，CNN的全部潜力才得以实现。这一突破推动了计算机视觉领域的快速发展，神经网络的复兴使得CNN成为计算机视觉任务的标准架构。在此期间，许多代表性CNN模型相继涌现，例如VGG（Simonyan & Zisserman，2014年）、GoogLeNet（Szegedy等，2015年）、ResNet（He等，2016年）、DenseNet（Huang等，2017年）、DCN（Zhang等，2019年；Wang等，2023年）和EfficientNet（Tan & Le，2019年）。
这些模型专注于不同的方面，包括准确性、效率和可扩展性，同时也促进了有价值的设计原则。近年来，受ViT启发，一些CNN模型（Liu等，2022年；Woo等，2023年；Ding等，2022年；Liu等，2023年）引入了大 Kernel 卷积来捕获长范围依赖关系，取得了与ViT相当的性能。与此同时，CNN被广泛整合到各种ViT和视觉Mamba中，以增强局部建模能力，形成了两者之间的互补协同效应。这些进展推动了视觉任务中模型设计的多样化和融合。
- Methodology
3.1. Preliminaries
状态空间模型（SSMs）（Gu等人，2021；Gu & Dao，2023）是一类常用在深度学习中的序列建模方法，能够通过中间隐状态
来表示动态系统。其核心方程如下：
为了在实践中实现连续时间模型，需要使用离散化技术。常用的零阶保持（Zero-Order Hold，ZOH）方法在每个时间间隔内保持输入不变，将连续时间参数
转换为离散形式如下：
其中
表示采样时间尺度。得到的离散化模型可以表示为：
该方法不仅支持高效的并行计算，还通过卷积操作直接生成序列输出：
其中，
是SSM核，
表示卷积操作。这种并行化显著提升了计算效率和可扩展性。
尽管传统的序列结构模型（如S4（Gu等人，2021））实现了线性的时间复杂度，但它们静态的参数化限制了其捕捉序列上下文的能力。为了克服这一局限，Mamba（Gu & Dao, 2023）引入了一种动态且依赖输入的参数化机制。与使用恒定转换参数
的传统模型不同，Mamba模型从输入序列
动态计算参数
、
和
。这使得模型能够进行更丰富和更注重序列建模的动态建模。
在 Mamba 中，连续时间参数通过基于输入序列
的选择性函数，自适应地调整为输入相关的参数。
并且可以根据输入计算得到依赖于输入的离散参数
和
。离散状态转换和观测方程如下：
Mamba模型的动态参数化不仅提高了序列建模性能，还在语言建模和视觉任务中展现了强大的竞争力。例如：Vim（Zhu等，2024）将双向Mamba块结合用于替代传统的Transformer块进行视觉建模。VMamba（Liu等，2024）通过引入二维选择性扫描构建了层次结构，类似于Swin Transformer（Liu等，2021）的设计。这些进展扩大了SSMs在基础视觉任务中的应用潜力，进一步推动了SSM模型的发展。
3.2. Dynamic Adaptive Scan
如下图3所示，作者提出了一种动态自适应扫描（Dynamic Adaptive Scan，DAS），该方法在特征图的重要区域指导下有效地建模了图像patches之间的关系。这些聚焦区域由多个可学习采样点集确定，这些采样点由偏移预测网络（Offset Prediction Network，OPN）从输入特征图中学习得到。在获得OPN预测的二维坐标后，作者使用双线性插值从特征图中采样特征，并将采样特征输入到SSM中进行特征聚合。此外，DAS的位置为SSM提供了更强的相对位置偏差，有助于SSM的学习。
具体地，作者首先将特征图输入到一个OPN中，以预测感兴趣patches相对于原始patches的二维坐标偏移量
：
然后，将这些偏移量加到原始Patch的位置上，以确定感兴趣Patch的采样位置：
通过建立图像块特征与偏移量之间的关系，所提出的方法可以自适应地学习。
其中，
表示双线性采样权重函数，
，
，并且
用于索引
中的所有位置，在这些位置上，只有最接近输入位置的四个网格点处
的值才非零。
在获得感兴趣的采样特征向量
后，作者按照其原始Patch从上到下、从左到右的顺序排列它们，然后将它们输入到SSM中进行特征提取。
3.3. Architecture Design of DAMamba
其中，
和
分别代表特征图中原_patch_和感兴趣_patch_的二维坐标集合。
和
的值范围在 -1 到 1 之间，(-1,-1) 表示左上角，(1,1) 表示右下角。
和
分别表示图像_patch_在高度和宽度方向上的坐标。
此外，DAMamba 架构集成了卷积位置编码以及卷积 FFN （ConvFFN）（Wang等，2022），进一步增强了提取局部特征的能力。在架构中，DAMamba 块灵活地堆叠在不同的阶段，以适应各种下游任务。最后，模型输出经过 BatchNorm 、全局平均池化和线性分类头，生成最终的图像分类任务特征表示。不同模型尺度的具体配置详见表1。
- Experiments
4.1. Image Classification on ImageNet-1K
实验设置：作者基于包含1,281,167张训练图像和50,000张验证图像的ImageNet-1K数据集（Russakovsky等，2015）进行了图像分类实验，这些图像覆盖了1,000个类别。作者实验方法的实现依赖于PyTorch（Paszke等，2019）和Timm（Wightman等，2019）库。为了确保公平比较，作者采用了DeiT中常用的实验设置（Touvron等，2021）。所用的优化器为AdamW（Loshchilov & Hutter，2017），采用余弦退化学习率调度，并在前20个周期进行线性 Warm up 。模型在分辨率
的图像上训练了300个周期。
为了数据增强和正则化，作者采用了诸如RandAugment（Cubuk等，2020）、重复增强（Hoffer等，2020）、Mixup（Zhang等，2017）、CutMix（Yun等，2019）、随机擦除（Zhong等，2020）、权重衰减、标签平滑（Szegedy等，2016）和随机深度（Huang等，2016）等技术。测试时，对验证图像应用中心裁剪以生成分辨率为
的输入图像。实验在16块RTX 3090 GPU上进行。值得注意的是，指数移动平均（EMA）（Polyak & Juditsky，1992）并未显著提升最终模型性能，在实验中未被采用。
结果：如表2所示，作者对比了提出的DAMamba与几种最先进的模型。提出的DAMamba在准确率上持续超越ViT、CNN和SSM模型。具体来说，DAMamba-B实现了85.2%的准确率，比当前最先进的SSM模型VMamba-B高出1.3%。与最新的先进CNN模型(ConvNeXt V2)和ViT模型(TransNext)相比，提出的DAMamba-T在准确率上有显著提升。即使将模型参数扩展到约50百万和30百万时，DAMamba依然保持了其优异的表现，分别实现了83.8%和84.8%的Top-1准确率。
4.2. Object Detection and Instance Segmentation on C0C02017
实验设置：作者在COCO 2017数据集上进行了目标检测和实例分割实验。COCO 2017数据集（林等，2014）包含大约118K张训练图像和5K张验证图像，并且常被用作目标检测和实例分割任务的标准基准。为了评估所提模型在下游视觉任务中的性能，作者将DAMamba作为主干网络嵌入到检测器中，以从图像中提取目标和实例特征。DAMamba被集成到了经典的Mask R-CNN（何等，2017）检测器中，并以在ImageNet-1K数据集上预训练的权重初始化，经过300个周期（epoch）的训练。对于目标检测和实例分割任务，作者分别对模型进行了12个周期（1×）和36个周期（3×）的训练。所有实验均基于MMDetection框架（陈等，2019）进行。
结果：DAMamba在COCO2017数据集上的目标检测和实例分割结果如表3所示。在边界框和 Mask 平均精度（AP b 和 AP m ）方面，DAMamba表现出色。采用一次微调策略时，DAMamba-T/S/B分别实现了48.5/49.8/50.6的目标检测mAP，分别超过VMambaT/S/B的1.2/1.1/1.4 mAP，Swin-T/S/B的5.8/5.0/3.7 mAP，以及ConvNeXt-T/S/B的4.3/4.4/3.6 mAP。在相同配置下，DAMamba-T/S/B的实例分割mAP也显著优于VMambaT/S/B、Swin-T/S/B和ConvNeXt-T/S/B。即使使用三次训练配置，DAMamba依然保持了较大的性能提升。这些结果表明，DAMamba有可能在密集预测下游任务中实现最佳性能。
4.3. Semantic Segmentation on ADE20K
实验设置：作者使用ADE20K数据集进行了语义分割实验，并在UperNet框架（Xiao等，2018）中对DAMamba和其他模型进行了对比分析。在UperNet框架中，主干网络使用预训练于ImageNet-1K数据集的权重初始化，其余部分则随机初始化。模型优化采用了AdamW优化器，批量大小为16。为了确保公平比较，在UperNet框架中所有模型都进行了
次迭代的训练。所有实验均使用MMSegmentation框架（贡献者，2020）进行。
结果：表4展示了DAMamba在UperNet（ Xiao et al., 2018）框架下的语义分割结果。实验表明，DAMamba-T、DAMamba-S和DAMamba-B分别实现了
分数为
、
和
，显著优于其他类型的模型。即使在使用多尺度测试时，作者DAMamba的表现提升也十分明显。这些结果进一步验证了DAMamba在下游任务中出色的一般化能力。
4.4. Ablation Study
为了验证作者方法的有效性，作者在ImageNet-1k上使用DAMamba-T进行了图像分类消融实验（参见表5）。与使用扫掠扫描的 Baseline 模型相比，作者提出的方法动态自适应扫描（DAS）在Top-1精度上提高了0.4%，并且只消耗了少量额外的FLOPs和参数开销。此外，对于以全局建模见长的视觉SSMs，作者观察到使用卷积位置编码（Convpos）（Chu等，2023；Islam等，2020；Li等，2021）和卷积 FFN （ConvFFN）（Wang等，2022）进行局部建模可以提高准确性。
- Conclusion
在本文中，作者提出了一种新颖的视觉状态空间模型，称为DAMamba。DAMamba显著增强了视觉状态空间模型中的灵活性，并提高了捕捉复杂图像结构的能力，同时保持了局部和全局上下文的理解能力。
具体地，作者提出了一种动态自适应扫描机制，该机制能够根据输入图像自适应地分配扫描顺序和区域。
在各种数据集和流行视觉任务上的 extensive 实验表明，所提出的 DAMamba 在图像分类、目标检测、实例分割和语义分割方面显著且一致地超越了当前最先进的视觉状态空间模型以及流行的ViT和CNN架构，建立了新的基准。
作者的发现强调了扫描机制在视觉状态空间模型中的重要性，并突显了状态空间模型作为视觉 Backbone 网的巨大潜力。
参考
[0]. DAMamba: Vision State Space Model with Dynamic Adaptive Scan .
点击上方卡片，关注
「AI视界引擎」
公众号