西交大 & 上海 AI Lab 提出 HRVMAMBA 用于高分辨率视觉状态空间模型,助力图像分类、人体姿态估计和语义分割等！

未来先知

发布于 2024-11-04 13:48:40

350

发布于 2024-11-04 13:48:40

文章被收录于专栏：未来先知

最近，由于其线性计算复杂度与 Token 长度成正比且具有全局感受野，具有高效硬件 Aware 的State Space Models（SSMs，如Mamba）在计算机视觉任务中展示出巨大的潜力。然而，Mamba在密集预测任务（包括人体姿态估计和语义分割）上的性能受到三个关键挑战的限制：归纳偏差不足、远程遗忘和低分辨率输出表示。为了解决这些问题，作者引入了Dynamic Visual State Space（DVSS）块，该块利用多尺度卷积核在不同尺度上提取局部特征并增强归纳偏差，同时采用变形卷积来缓解远程遗忘问题，同时根据输入和任务特定信息实现自适应空间聚合。通过利用HRNet（Wang等人，2020年）提出的多分辨率并行设计，作者引入了基于DVSS块的高分辨率视觉状态空间模型（HRVMamba），该模型在整个过程中保持高分辨率表示，同时促进有效多尺度特征学习。大量实验证明了HRVMamba在密集预测任务上的出色表现，在无花哨的条件下与现有基准模型实现竞争结果。代码可从https://github.com/zhanghao5201/HRVMamba获取。

1 Introduction

卷积神经网络（CNNs）和视觉 Transformer （ViTs）在图像分类、人体姿态估计和语义分割等任务上取得了显著进展。尽管CNN在局部特征提取方面表现出色，具有线性计算复杂度，但它们缺乏全局上下文建模能力。尽管ViTs通过自注意力机制捕获全局感受野，但面临平方复杂度，尤其是对于大输入缺乏归纳偏差。Mamba（Gu和Dao，2023年）引入了S6结构，提高了状态空间模型（SSMs）在长程特征提取方面的效率。通过使用输入相关的状态空间参数，Mamba实现了线性复杂度的更好的上下文建模。这导致了许多后续的视觉Mamba模型，如ViM，VMamba，LocalVMamba，GroupMamba 。

然而，视觉Mamba模型在密集预测任务上并未达到最佳性能，包括人体姿态估计（Xu等人，2024；Zhang等人，2024）和语义分割，主要原因有三个。

首先，与ViT类似，视觉Mamba将图像分割成一系列的patch（token）序列，并使用双向或四向扫描机制来遍历这些token，构建全局感受野。虽然这种方法对于处理长序列有效，但它破坏了图像的天然2D空间依赖关系，并且缺乏对有效局部表示学习的归纳偏差。

其次，Mamba的token处理导致之前隐藏状态的衰减，从而导致长期遗忘。因此，它可能失去与 Query patch相关的高层，任务特定的特征，而专注于低层边缘特征，如图1，第21列所示。最后，当前的视觉Mamba模型通常生成单尺度，低分辨率特征，这导致大量信息损失，并使其难以捕捉密集预测任务所需的细粒度细节和多尺度变化。

为了克服这些限制，作者提出了动态视觉状态空间（DVSS）块，该块基于V Mamba中提出的视觉状态空间（VSS）块。DVSS块利用多尺度卷积 Kernel 提取不同尺度的局部特征，从而增强各种视觉特征尺度的归纳偏差。此外，它还集成了可变形卷积v4（DCNv4）（Xiong等人，2024年），使其能够根据输入和任务特定信息进行动态高层空间聚合。这通过增强块之间的高层语义关系来减轻Mamba的远程遗忘问题，而不是主要关注低层特征。例如，如图1所示，右肩附近左肩和胸部的特征被突出显示（行1，列3），与右肩相连的 Head 特征被强调（行1，列5），同时双手和胸部的突出特征也被显示。

作者进一步采用HRNet（Wang等人，2020年）的多分辨率并行设计，将DVSS块嵌入并行多分辨率分支中，构建了高分辨率视觉状态空间模型（HRVMamba）。HRVMamba保持并增强了高分辨率表示，保留了更精细的细节，并使用多分辨率分支建模多尺度变化，使其非常适合密集预测任务。

这项研究的贡献如下：

作者提出了DVSS模块，它集成多尺度卷积 Kernel 和变形卷积，以减轻VSS模块中缺乏归纳偏见和长时遗忘问题的不足。
作者提出了一种基于DVSS模块的HRV Mamba模型，作为第一个应用于多分辨率分支结构的Mamba基础模型，旨在保留细粒度细节并捕捉多尺度变化，特别是在密集预测任务中。

HRV Mamba在图像分类、人体姿态估计和语义分割等任务中展现出良好的性能。实验结果显示，HRV Mamba在现有的CNN、ViT和SSM基准模型中取得了具有竞争力的结果。

相关研究 卷积神经网络（CNNs）和视觉 Transformer （ViTs）。

CNNs 一直是计算机视觉的基础，从早期的模型如 AlexNet和 ResNet（He 等，2016）演化为更近期的架构，如 ConvNeXt，SCGNet，FlashInternImage（Xiong 等，2024）和 FMGNet（Zhang 等，2024）。这些模型在局部特征提取方面表现出色，在图像分类、语义分割和人姿态估计等任务上实现了显著的性能。ViTs 引入了自然语言处理（NLP）中的自注意力机制，将图像分割成 patch 以捕捉全局依赖性，构成了大型视觉语言模型的基础（Zhang 等，2024; Ying 等，2024; Liu 等，2024）。

包括 DeiT 的蒸馏策略、Swin Transformer 的层次结构（Liu 等，2021）和 SwiftFormer 的有效注意力机制等各种方法已经出现，以扩大 ViTs 在视觉任务中的应用范围。最近，结合 CNNs 和 Transformer 优势的混合架构引起了关注。这些模型利用 CNNs 的局部特征提取的归纳偏见，同时结合 ViTs 的全局注意力能力，在 Backbone 网络研究方面标志着一个重要方向。

状态空间模型（SSMs）。 状态空间模型（SSMs）是一个数学框架，用于模拟具有线性计算复杂度的动态系统，使其在处理长序列时具有高效性。像S4（Gu等人，2021年）、S5（Smith等人，2022年）和H3（Fu等人，2022年）这样的模型通过结构优化、并行扫描和硬件改进等优化措施，提高了SSMs的性能。Mamba（Gu和Dao，2023年）引入了输入特定的参数化和并行扫描（S6），将SSMs定位为Transformers的有力的替代方案。自那时以来，SSMs已在视觉任务中广泛应用，其中S4ND 是首个将视觉数据处理为连续信号的模型。

基于Mamba，ViM 和VMamba（Liu等人，2024年）解决了Mamba的定向性，采用双向或四向扫描。LocalVMamba 通过窗口扫描捕获局部细节，而PlainMamba（Yang等人，2024年）则对2D扫描进行细化，以用于顺序处理。MambaVision 将SSMs与Transformers集成，而GroupMamba（Shaker等人，2024年）则通过蒸馏方法提高训练稳定性。然而，这些视觉Mamba模型通常会产生单尺度、低分辨率的特征，限制了它们在密集预测任务中捕捉细粒度细节和多尺度变化的能力。

高分辨率网络用于密集预测

高分辨率网络（High-Resolution network）首次在 HRNet（Wang 等人，2020 年）中引入，在人类姿态估计和语义分割等任务上表现出强大的性能。其多分辨率并行分支通过多尺度融合模块通过交换多分辨率特征来结合不同尺度的信息。在此基础上，HRFormer（Yuan 等人，2021 年）将局部窗口自注意力机制（Liu 等人，2021 年）与高分辨率结构相结合，在密集预测任务上取得了卓越的结果。

进一步的改进，包括 Lite-HRNet（Yu 等人，2021 年）、Dite-HRNet（Li 等人，2022 年）和 HFHRNet（Zhang 等人，2024 年），通过深度卷积和动态卷积等技术引入轻量级卷积神经网络，使得高分辨率网络可以在移动设备上更高效地部署。然而，Mamba 是否能在高分辨率结构中发挥最佳性能以及如何缓解 Mamba 在密集预测任务中的归纳偏差和长期遗忘仍需进一步研究。

3 Preliminaries

状态空间模型（SSMs）将一维函数或序列通过一个隐状态映射到输出序列，这些模型基于连续线性时不变（LTI）系统。

为了将深度模型集成并适应实际数据，必须应用离散化，将 SSMs 的连续微分方程通过零阶保持方法转换为离散函数。具体而言，通过一个离散时间步长，SSMs 被离散化为以下形式：

其中, 是系统的演化矩阵，和是投影矩阵。.

在Mamba（Gu & Dao，2023）中，引入了选择性状态空间模型（S6），以提高强上下文信息的提取。S6允许B、C和Δ作为输入x_t的函数而变化，而在S4（Gu et al., 2021）中，A、B、C和Δ是输入无关的，这限制了模型从输入序列中提取关键信息的能力。形式上，给定一个输入序列，其中B、L和C分别表示批量大小、序列长度和特征维度，输入相关参数B、C和Δ的计算如下：

其中，是一个可学习的参数，与 S4 中的参数相同。

4 High-Resolution Visual State Space Model

Multi-resolution Parallel VLambda

当前的视觉Mamba模型通常生成单尺度、低分辨率的特征，导致大量信息损失，并且难以捕捉到用于密集预测任务所需的细粒度特征和多尺度变化。为了解决这个问题，作者采用了HRNet的多分辨率并行设计（Wang等人，2020），使用并行分支开发了高分辨率视觉状态空间模型（HRVMamba）。作者将在图3（a）中展示HRVMamba的整体架构。对于输入图像，HRVMamba首先使用两个3x3卷积进行下采样，步长为2，将特征分辨率降低到。

然后，它经过四个阶段，其中后期阶段的分支包括前一个阶段的分辨率和额外的一个低分辨率，最终四个分支的特征维度为，，和。以前的研究（Yun & Ro，2024；Ma等人，2024）表明，在早期阶段对较大特征图进行卷积操作更有效，因此作者采用了HRNet的第一阶段类似的设计。其余阶段使用作者提出的动态视觉状态空间（DVSS）块（图3（c））作为基本单元。多尺度融合方法，遵循HRNet，包括一系列上采样和下采样块，将不同并行分支的特征进行合并。

Dynamic Visual State Space (DVSS) Block

作者提出了DVSS块，该块基于虚拟机管理 Agent 中的VSS块（图3（b））（刘等人，2024）。如图3（c）所示，DVSS块包括2D-选择性扫描与可变形卷积（DSS2D）块、多尺度宽度（MultiDW）块以及一个 FFN （FFN）作为特征提取单元。

基于可变形卷积的2D选择性扫描（DSS2D）块如先前的研究（Shi等人，2024年）中所述，在输入序列的SSM中，第个 Token 对第个 Token 的贡献可以表示为：

通常，学习的是负数，导致方程6中的随着序列距离的增加显著减小。这导致每个新 Token 与前一个隐状态的衰减一致，从而产生长期遗忘问题。因此，SSM可能失去与 Query Patch 和特定任务相关的较高层次、任务特定特征，而专注于较低层次的边缘特征，如图1中第二列所示。

为了减轻长程遗忘问题，作者用变形卷积 v4（DCNv4）（Xiong 等人，2024）替换 SS2D 块中的深度卷积，并构建了 DSS2D 块。给定输入，对于每个参考点，个点（对于 DCNv4，）的 DCNv4 操作定义如下：

表示空间聚合组的数量，设置为 4。标量表示第个采样点在第个组中的动态调制权重。是第个网格采样位置，即，而是其动态偏移。

另一方面，DCNv4通过增强像注意力机制这样的高级语义关系，而不是主要关注低级特征，来缓解Mamba的远程遗忘问题。这可能有助于确保在远程衰减的情况下，相邻 Token 之间的语义关系仍然受到影响。另一方面，与较大的卷积核和自注意力机制相比，DCNv4在提高远程特征建模的计算效率方面取得了很大提高。

多尺度深度卷积（多DW）块。 视觉Mamba使用双向或四向扫描将图像处理成token序列。然而，这种方法破坏了2D空间关系，并且缺乏局部特征所需的归纳偏差。

为了解决这个问题，作者引入了多DW块，它采用多尺度卷积核来捕捉不同尺度下的局部特征，从而增强不同尺度特征的归纳偏差。

具体而言，如公式9、10和11所示，输入特征首先在通道维度上分成组，其中设置为4。第组的特征经过一个大小为的深度卷积。所得特征然后进行拼接，接着进行乱序操作，以促进不同组之间的特征交互。

在此处，GELU表示激活函数。

HRVMamba Architecture Instantiation.

作者在表1中说明了HRVMamba架构配置。在第i阶段，表示块数，表示SSM扩展比，表示MLP扩展比，表示模块数。此外，作者还设计了两种规模的HRVMamba，即HRVMamba-S和HRVMamba-B。表2呈现了HRVMamba变体的详细信息。

Human Pose Estimation

训练设置作者在COCO数据集（Lin等人，2014）上评估HRVMamba在人体姿态估计任务上的性能，该数据集包括超过200,000张图像和250,000个带有17个关键点的 Token 人体实例。作者的实验在COCO train 2017数据集上进行，该数据集包括57,000张图像和150,000个人体实例。作者的模型在val 2017和test-dev 2017集上进行评估，分别包含5,000和20,000张图像。

在进行训练和评估时，作者遵循MMPOSE（贡献者，2020）的实现。批量大小设置为256，使用AdamW优化器，学习率设置为5e-4，β设置为（0.9，0.999），权重衰减设置为0.01。对于HRVMamba-B，没有使用预训练技术，而对于HRVMamba-S，作者应用了与HRFormer类似的ImageNet预训练（Deng等人，2009）。

结果。表3在COCO验证集上展示了结果。HRVMamba始终优于其他CNN模型、ViT模型和最新的状态最先进（SOTA）SSM方法。当输入大小为时，HRVMamba-S实现了74.6 AP，超过了FlashInternImage-B（74.1 AP）的同时仅使用五分之一的工作量。HRVMamba-B实现了76.4 AP，超过了像ViT-S、VMamba-B、MambaVision-B和GroupMamba-B这样的SOTA SSM方法。

在类似的计算复杂性下，HRVMamba-B相对于GroupMamba-B提高了3.2 AP和2.7 AR。此外，HRVMamba-B比ViTpose-B提高了0.6 AP和0.3 AR，参数减少50%，浮点运算减少20%。当输入大小为时，HRVMamba-S比HRFormer-S提高了0.8 AP；HRVMamba-B在不使用ImageNet预训练的情况下，比HRFormer-B提高了0.6 AP。

作者还在表4中提供了在COCO测试开发集上的比较。作者的HRVMamba-S在仅使用其十分之一的参数时，实现了75.3的AP，比ViTPose-B提高了0.2。它与VMamba-B的性能相当，但FLOPs仅为VMamba-B的五分之一。此外，HRVMamba-B在AP和AR方面分别比HRFormer-B提高了0.5和0.4，且无需在ImageNet上进行预训练，实现了最先进的性能。

Semantic Segmentation

作者采用Xiao等人（2018年）提出的UPerNet作为所有测试模型的基础框架。所有模型都在ImageNet-1K数据集（Deng等人，2009年）上进行预训练。

城市景观数据集（Cordts等人，2016年）是为了理解城市场景而设计的，其中使用了19个类别进行语义分割。经过精细标注的5000张图像被分为2975张训练图像、500张验证图像和1525张测试图像。作者设定了初始学习率为，权重衰减为0.01，裁剪尺寸为1024x512，批量大小为16，训练迭代次数为80K。如表5所示，HRVMamba-B在单尺度测试中比HRFormer-B提高了2.1 mIoU，在多尺度测试中提高了2.5 mIoU。同时，它还以更少的参数超过了SSM模型如GroupMamba-B和MambaVision-B。

PASCAL-Context 数据集（Mottaghi等人，2014年）包括59个语义类别和1个背景类别，共有4,998张训练图像和5,105张测试图像。作者将初始学习率设置为，权重衰减为0.01，裁剪大小为480x480，批量大小为16，训练迭代次数为80K。如表5所示，HRVMamba-B分别相较于HRFormer-B和MambaVision-B实现了0.9 mIoU和2.2 mIoU的提升。值得注意的是，LocalVMamba-S在具有变输入大小的测试中表现尤其糟糕，仅实现了12.2 mIoU。

ImageNet Classification

训练设置作者在ImageNet-1K数据集（Deng等人，2009年）上进行比较，该数据集包括1000个类别中的1280万训练图像和50000张验证图像。HRVMamba在运行在8个A100 GPU（每个GPU 80GB）上的Swin Transformer（Liu等人，2021年）训练框架。

结果。表6对比了HRVMamba与几种代表性的CNN、ViT和SSM方法。HRVMamba在各种等向量架构（Touvron等人，2021；Li等人，2024）、层次架构（Liu等人，2021，2024b）和高分辨率架构（Yuan等人，2021）上均表现出竞争力。具体而言，HRVMamba-B仅使用VideoMamba-M的30%的计算量，实现了Top-1准确率83.7，而VideoMamba-M的Top-1准确率为83.8。尽管MambaVision-B在整体准确率上达到84.2，采用了更先进的LAMB优化器，而其他模型没有使用该优化器。

此外，MambaVision-B在32A100 GPU上进行训练，而作者的模型仅使用8 A100 GPU，并依赖先进的训练技术，如蒸馏（Shaker等人，2024）。重要的是，如表3所示，MambaVision-B在密集预测任务中表现不佳，得分为73.4 AP，而HRVMamba-B得分为76.4 AP。值得注意的是，HRVMamba在所有高分辨率架构中实现了最佳性能，其中HRVMamba-S与HRFormer-S相比提高了0.6个点，而HRVMamba-B在相等的计算量下超过了HRFormer-B，提高了0.5个点。

Ablation Experiments

多分辨率并行架构。表3中的结果表明，利用多分辨率并行架构的HRVMamba在姿态估计任务上取得了SOTA性能。特别是，与其他SOTA SSM模型（如V Mamba（Liu等人，2024b）、V Mamba-B（Liu等人，2024b）、MambaVision-B（Hatamizadeh和Kautz，2024）以及GroupMamba-B（Shaker等人，2024））的比较凸显了多分辨率并行架构在密集预测任务中的优势。

如图7所示，DSS2D块相对于SS2D块（第2行 vs 第1行）提高了0.4个点，表明引入DCN增强了Mamba的空间特征提取。具体而言，如图1所示，DSS2D在早期阶段（S2）关注与 Query 块相关的较高层次特征，而SS2D则针对低层次的边缘特征。在后期阶段（S3），DSS2D突出了与人类相关的细节，而SS2D则倾向于捕获无关的背景信息。作者认为DCNv4增强了块之间较高层次语义关系特征，使得它们在长时间衰减（长期遗忘问题）后仍能相互影响。

多尺度深度卷积块（Multi-scale Depthwise Block）

HRFormer在FFN中引入了深度卷积以增强模型的归纳偏差。然而，作者在表7中的实验结果表明，将MultiDW块嵌入到FFN中（第3行与第2行）甚至可能降低DCN带来的性能提升。相反，当MultiDW块作为独立模块使用时，如图3所示，它将AP提高到74.2。然而，用卷积替换多尺度卷积核并未导致任何性能提升（第5行与第3行）。这表明多尺度卷积核在捕捉不同尺度下的局部特征方面是有效的，从而增强了特征的归纳偏差。