首页
学习
活动
专区
圈层
工具
发布
首页
学习
活动
专区
圈层
工具
MCP广场
社区首页 >专栏 >A2Mamba:融合Transformer与Mamba的混合视觉架构

A2Mamba:融合Transformer与Mamba的混合视觉架构

作者头像
AI浩
发布2025-09-11 20:00:45
发布2025-09-11 20:00:45
1530
举报
文章被收录于专栏:AI智韵AI智韵

论文简要总结

近年来,基于Transformer和Mamba的架构在计算机视觉领域展现出强大的潜力。然而,现有方法通常只是简单堆叠这两类模块,缺乏深层次的交互机制。论文《A2Mamba: Attention-Augmented Mamba for Vision》提出了一种新型的混合视觉主干网络A2Mamba,通过多尺度注意力增强状态空间模型(MASS)实现了Transformer与Mamba的深度融合。

核心创新点包括:

  1. 多尺度注意力机制(AMA):结合滑动局部注意力(SLA)和空洞滑动注意力(DLA),自适应捕获局部细节与长程依赖关系
  2. 注意力增强状态空间模型(A2SSM):通过交叉注意力变体将多尺度注意力图与SSM隐藏状态结合,增强空间感知能力
  3. MM-Refine解码模块:为语义分割任务设计的渐进式多尺度细化模块,保留细节信息的同时提取丰富上下文

A2Mamba在多个视觉任务上取得突破性性能:

  • ImageNet-1K分类:A2Mamba-L达到86.1% top-1准确率
  • 目标检测:在COCO数据集上超过现有最佳方法2.6% AP
  • 语义分割:在ADE20K上达到54.1% mIoU
  • 高效率:在高分辨率输入下保持线性计算复杂度

该方法首次实现了注意力机制与状态空间模型的深度集成,为视觉主干网络设计提供了新方向。

论文翻译全文

摘要

最初为自然语言处理发明的Transformer和Mamba,已经激发了用于视觉识别的主干架构。最近的研究将局部注意力Transformer(Local Attention Transformers)与Mamba集成,以捕捉局部细节和全局上下文。尽管性能具有竞争力,但这些方法仅限于简单堆叠Transformer和Mamba层,而没有它们之间的任何交互机制。因此,Transformer层和Mamba层之间的深度集成仍然是一个悬而未决的问题。我们通过提出A2Mamba来解决这个问题,这是一个强大的Transformer-Mamba混合网络架构,其特点是采用了一种称为多尺度注意力增强状态空间模型(Multi-scale Attention-augmented State Space Model, MASS)的新令牌混合器(token mixer),其中多尺度注意力图被集成到一个注意力增强的SSM(A2SSM)中。A2SSM的一个关键步骤是执行一种交叉注意力的变体,即使用多尺度注意力图在空间上聚合SSM的隐藏状态,这增强了关于二维空间的空间依赖性,同时提高了SSM的动态建模能力。我们的A2Mamba在视觉识别任务中优于所有先前的基于ConvNet、Transformer和Mamba的架构。例如,A2Mamba-L在ImageNet-1K上达到了令人印象深刻的86.1% top-1准确率。在语义分割中,A2Mamba-B在mIoU上超过CAFormer-S36 2.5%,同时展现出更高的效率。在使用Cascade Mask R-CNN的目标检测和实例分割中,A2Mamba-S在 上超过MambaVision-B 1.2%/0.9%,同时参数减少40%。代码公开在 https://github.com/LMMMEng/A2Mamba。

https://arxiv.org/pdf/2507.16624

索引术语—视觉识别,视觉主干架构,Transformer,注意力,Mamba,状态空间模型

1 引言

视觉Transformer(ViTs)[1] 因其使用多头自注意力(MHSA)[2] 建模长距离依赖关系的能力,已成为各种视觉任务的事实选择。然而,MHSA的二次复杂度导致了高计算成本,特别是在需要高分辨率输入的密集预测任务中,如语义分割和目标检测。为此,后续努力提出了高效的注意力机制,如窗口注意力(window attention)[3]-[6]、空间缩减注意力(spatial reduction attention)[7]-[9]和空洞注意力(dilated attention)[10]-[12]以降低计算复杂度。最近,由于Mamba架构[13]能够以线性时间复杂度建模长距离依赖关系,许多努力致力于开发基于Mamba的架构用于视觉识别[14]-[20]。与通过下采样或混洗(shuffling)来减少序列长度的空间缩减注意力和空洞注意力不同,Mamba通过状态空间模型(SSMs)直接在原始序列上建模长距离依赖关系。这种架构使得在长序列处理过程中能够保留细粒度信息,非常有希望使视觉模型在密集预测任务中实现卓越性能[21]。

SSMs中的顺序扫描机制自然适合语言建模,因为词序很重要,而图像表现出具有非顺序像素依赖关系的复杂二维结构。因此,SSMs难以全面理解图像的空间结构。尽管一些努力[15], [16]利用了替代的扫描策略来部分克服这一限制,但顺序扫描带来的内在因果性仍然在一定程度上损害了潜在的空间依赖性。因此,Transformer-Mamba混合架构已成为视觉识别的一个有前景的方向。例如,MambaVision [22]通过在其更深的阶段堆叠MHSA和SSM块来构建视觉主干,使用MHSA来补充SSM。然而,尽管效率很高,其性能在各种视觉任务上仍然落后于先进的ViTs [23]-[26]。最近,一个通用的Transformer-Mamba混合架构,称为SegMAN编码器[27],采用统一的令牌混合器来结合滑动局部注意力(sliding local attention)[5]和SS2D [15],实现了有竞争力的性能,并与领先的ViTs相比具有有利的权衡。然而,由于这些努力代表了将Transformer和Mamba集成用于视觉任务的早期尝试,其令牌混合器中的基于注意力和SSM的模块只是简单堆叠。仍然缺乏有效的方法来实现Transformer和Mamba层之间更深的集成,从而产生一个强大的视觉主干,在效率和性能上都能够超越领先的ViTs。

在这项工作中,我们提出了一种新颖的混合令牌混合器,称为多尺度注意力增强状态空间模型(MASS),它利用了自注意力和SSM的优势。具体来说,我们首先引入了一种自适应多尺度注意力(AMA)机制,包含两个互补的路径:(1) 捕捉细粒度空间细节的常规滑动局部注意力(SLA);(2) 自适应调整空洞率以建模长距离依赖关系的空洞滑动注意力(DLA)。这种设计背后的动机是鼓励在多粒度上进行特征和上下文表示。该机制中的注意力矩阵具有多尺度的动态空间依赖性。其次,为了实现SSM和自注意力层之间更深的集成,SSM的隐藏状态通过一种交叉注意力的变体与前述的多尺度注意力矩阵进行交互。该设计旨在动态增强二维空间依赖性并缓解顺序扫描引入的因果性,从而提高SSM的空间感知和动态建模能力。总体而言,我们的MASS有效地将自适应多尺度表示和长距离依赖关系建模封装到一个混合令牌混合器中。

在这里插入图片描述
在这里插入图片描述

通过分层堆叠MASS令牌混合器和前馈网络(FFN)层,我们提出了一种通用的Transformer-Mamba混合视觉主干架构,称为A2Mamba。如图1所示,A2Mamba在各种视觉任务上表现出比先进的ConvNets、Transformers和基于Mamba的架构显著更好的性能。例如,我们仅具有约30M参数的A2Mamba-S模型实现了84.7%的惊人top-1准确率,分别超过了RMT-S [25]和TransNeXt-T [26] 0.6%和0.7%,同时具有更高的效率。此外,A2Mamba-S甚至以仅约三分之一的计算复杂度,在top-1准确率上超过了混合模型MambaVision-B [22] 0.5%。A2Mamba在密集预测任务中始终表现出优于其他基线的性能。例如,在使用UperNet [28]的语义分割任务中,A2Mamba-B在mIoU上分别超过了BiFormer-B [6]和UniFormer-B [26] 2.3%和3.3%。同时,在使用Cascade Mask R-CNN [29]的目标检测和实例分割任务中,A2Mamba-L在APb /APm上分别领先CAFormer-M36 [24]和MogaNet-L [30] 1.8%/1.6%和2.3%/2.0%。这些实验结果表明A2Mamba具有更强的全局建模和局部细节保留能力。

这项工作的一个初步版本已发表在CVPR 2025 [27]。在初步版本中,我们的贡献总结如下:

  1. 我们引入了一种新颖的视觉主干架构,称为SegMAN编码器,其特点是采用混合LASS混合器。LASS协同结合了局部注意力和状态空间模型,用于高效的局部细节编码和全局上下文建模。
  2. 我们提出了基于Mamba的多尺度上下文提取(MMSCopE),这是一种专为语义分割任务设计的新特征解码器。MMSCopE在随输入分辨率自适应缩放的多尺度特征图上操作,在细粒度细节保留和全尺度上下文建模方面均超越了先前的方法。
  3. 通过集成SegMAN编码器和MMSCopE,设计了一个强大的分割网络架构SegMAN。在语义分割任务上的大量实验证明了我们方法的卓越性能和具有竞争力的效率。

在这个扩展版本中,我们旨在进一步释放Transformer-Mamba混合架构在视觉识别方面的潜力。与我们的会议论文相比,本版本在以下方面进行了实质性改进:

  1. 我们提出了一种新的混合令牌混合器,称为MASS,它可以更深地集成自注意力和SSM,在单个混合器内实现强大的多尺度上下文建模和长距离依赖关系建模能力。注意,MASS令牌混合器是对会议论文中LASS令牌混合器的更强大替代。
  2. 基于MASS,我们提出了一个更强的视觉主干架构,称为A2Mamba,它为各种视觉识别任务编码更具区分性的特征表示。此外,我们利用MASS为语义分割构建了一个新的解码器,称为基于MASS的多尺度细化(MM-Refine)模块,将其与A2Mamba结合形成一个新的分割网络架构SegMAN-V2。
  3. 我们在更广泛的视觉识别任务上对我们架构进行了更广泛的实验验证,包括不同分辨率下的图像分类以及密集预测(包括语义分割、目标检测和实例分割)。广泛的结果表明,我们的方法在计算成本更低的情况下优于所有现有的基线。

2 相关工作

2.1 ConvNets

自从AlexNet [31]出现以来,ConvNets释放了深度学习的潜力,并逐渐成为视觉识别的主流架构。最初,ConvNet的设计侧重于使用小核(即3×3)构建深度网络,逐渐增加感受野,例如VGGNet [32]、ResNet [33]和DenseNet [34]。然而,现代ConvNet设计[35]-[38],以ConvNeXt [35]为例,已将焦点转向增加核大小以更快地扩大感受野,旨在实现与基于Transformer和Mamba的模型相当的性能。同时,门控机制已成功集成到现代ConvNets中以提升性能[30], [39], [40]。最近,OverLoCK [41]通过从生物自上而下神经注意力[42]中汲取灵感,重新发明了ConvNet架构,在各种视觉任务上显著优于先前的ConvNets。然而,同时跨网络层获得更多信息的多尺度表示和全局依赖性仍然具有挑战性,本文旨在探索一个更强大的解决方案。

2.2 视觉Transformer

ViT[1]的出现激发了在视觉领域探索多头自注意力(MHSA)的许多后续工作围绕MHSA构建视觉主干模型。然而,普通的MHSA存在二次复杂度,导致长序列建模中的高计算成本。为此,提出了各种高效的注意力机制来捕获长距离依赖关系,同时保持计算效率,例如窗口注意力[3]-[6]、空间缩减注意力[7], [8]和空洞注意力[10], [12]。为了进一步提高性能,BiFormer [6]引入了双层路由注意力(bilevel routing attention),以从粗到细的方式捕获局部范围依赖关系。最近,RMT [25]提出了曼哈顿注意力(Manhattan attention),它在注意力计算中注入了空间先验,以实现更准确的全局信息感知。尽管取得了显著成果,但这些工作中使用的高效注意力机制通常通过牺牲序列长度来逐步捕获长距离上下文。相比之下,本文旨在开发一种结合多尺度注意力和状态空间模型(SSM)[13]的混合架构,以在不减少序列长度的情况下建模细粒度多尺度线索和全局上下文,从而产生更强大的视觉架构。

2.3 视觉Mamba

受Mamba [13]在自然语言处理(NLP)任务中出色性能的启发,研究人员将其应用扩展到计算机视觉任务。作为Mamba的核心,状态空间模型(SSM)能够以线性时间复杂性建模长距离依赖关系,在视觉任务中展现出卓越性能。ViM [14]首先引入了双向SSM模块,并构建了类似于ViT [1]的朴素架构。VMamba [15]将扫描顺序扩展到包括四个方向,并提出了一个早期的基于SSM的分层架构。随后,提出了一系列代表性的基于Mamba的视觉主干模型[16]-[20]。例如,SpatialMamba[19]提出了一种结构化的SSM来增强对图像结构的空间感知。SparX-Mamba [20]专注于通过提出一种新的稀疏跳跃连接机制来改进基于Mamba的网络架构。这项工作采用多尺度自注意力来内在且动态地增强SSM的表征能力,从而进一步释放基于Mamba的模型在视觉任务中的潜力。

2.4 混合视觉主干架构

混合视觉模型已成为视觉识别中一个有前景的方向。先前,各种Transformer-ConvNet混合模型已被广泛研究,展现出卓越的性能[10], [24], [43]-[47]。混合视觉模型的主要优势在于能够利用两个子混合器的优势,例如ACmix [44]和MixFormer [45],它们并行使用深度卷积(DWConv)和移位窗口注意力(shifted window attention)。最近,TransNeXt [26]提出了一种中心凹自注意力(foveal self-attention)机制和ConvGLU,开发了一个强大的Transformer-ConvNet混合视觉主干架构,在各种视觉任务上展示了显著成果。自从引入Mamba以来,将Mamba集成到混合模型中已显示出有前景的性能。MambaVision [22]将Conv、SSM和MHSA集成到一个单一网络中,虽然展示了高效率,但其性能仍然落后于先进的视觉主干架构。我们的初步工作SegMAN [27]提出了一个有效的Transformer-Mamba混合视觉主干和一个伴随的基于Mamba的解码器,在语义分割任务中展示了相对于其他基线的显著性能改进。在这项工作中,我们通过引入一个新的、更强大的令牌混合器(称为多尺度注意力增强SSM)来进一步释放Transformer-Mamba混合视觉架构的潜力,该混合器更深度地集成了注意力与状态空间模型。

3 方法

在本节中,我们首先简要回顾我们初步工作[27]中的网络架构。然后,详细阐述一个具有显著性能改进的升级版本。

3.1 SegMAN回顾

我们早期的工作[27]代表了探索结合局部自注意力和状态空间模型以构建强大视觉主干架构的早期尝试,即SegMAN编码器。该令牌混合器由两个互补的堆叠模块组成:用于捕获局部细节的滑动局部注意力(SLA)[5]和用于建模长距离依赖关系的选择性扫描2D(SS2D)[15]。与先前使用空间下采样自注意力来减少序列长度以建模长距离依赖关系的工作不同,最近状态空间模型的线性时间复杂性使我们的SegMAN编码器能够在不牺牲序列长度的情况下建模全局信息,从而允许保留对密集预测至关重要的细粒度空间信息。在ImageNet-1K分类任务中,SegMAN编码器展示了卓越的性能,显著优于先前的ConvNets、Transformers和基于Mamba的架构,同时与先进的基于Transformer的架构(即RMT [25]和TransNeXt [26])持平。

另一方面,我们还为语义分割提出了一个基于Mamba的解码器,该解码器包含一个新颖的基于Mamba的多尺度上下文提取(MMSCopE)模块。在实践中,MMSCopE首先计算多个尺度的特征,然后将它们馈送到SS2D中。这种设计背后的动机是多尺度特征可以促进不同粒度下的上下文建模,从而获得更好的语义分割结果。通过集成提出的编码器和解码器,我们引入了一个新的分割网络架构,称为SegMAN,该架构在三个具有挑战性的数据集上进行了评估,包括ADE20K[48]、Cityscapes[49]和COCO-Stuff[50],以显著优势超过了先前的state-of-the-art分割网络架构,如SegNeXt [51]和VWFormer [52]。

3.2 A2Mamba的整体架构

在这项工作中,我们提出了一种新颖的混合视觉主干架构A2Mamba,它利用了Transformer和Mamba架构的优势。A2Mamba是SegMAN编码器的一个全面升级版本,在性能和效率上都提供了显著改进。如图2所示,A2Mamba是一个具有四个阶段的金字塔架构,与先前的工作[3], [8], [33], [53]相同。每个阶段的下采样因子分别为 和,而通道维度随深度增加。对于分类任务,最深阶段的输出被馈送到分类器以生成图像级预测。相比之下,分层特征用于密集预测任务,例如目标检测和语义分割。

在这里插入图片描述
在这里插入图片描述

A2Mamba的关键层是A2Mamba块,每个块主要由三个组件组成:一个增强位置信息的残差 深度卷积(DWConv),一个作为核心令牌混合器以捕获全尺度上下文信息的新颖的多尺度注意力增强状态空间模型(MASS),以及一个提升通道多样性的卷积前馈网络(ConvFFN)[8]。

在这里插入图片描述
在这里插入图片描述

3.3 MASS令牌混合器

自适应多尺度注意力(Adaptive Multi-scale Attention)。所提出的MASS通过将动态多尺度聚合与长距离传播集成来增强其上下文建模能力,同时使用门控机制[13], [30]来进一步消除上下文噪声。如图3(a)所示,给定一个输入特征图 ,其中C表示通道维度, 表示空间维度,我们首先将X在通道维度上均匀划分为 。然后, 使用标准的SLA [5]进行处理。具体来说,在 上的一个滑动窗口内计算多头自注意力(MHSA)[2],其中只有查询位于中心,生成一个注意力图 ,其中 是原始X上的注意力头数, 表示窗口大小。该注意力图通过注意力加权求和动态聚合 中的细粒度局部邻域,产生一个新的特征图 。同时, 使用空洞局部注意力(DLA)[11]进行处理,它通过类似于空洞卷积[54]的空洞机制来扩大感受野。为了在不同分辨率下一致地捕获长距离依赖关系,空洞率r自适应地确定如下:

这个公式背后的动机是使空洞滑动窗口具有与输入特征图相同的大小,而不管绝对分辨率如何。因此,基于注意力的上下文建模范围覆盖了整个输入空间。之后,生成的特征图 沿通道维度连接形成 。这种整合结合了来自标准SLA的细粒度局部细节和由DLA捕获的稀疏采样的长距离依赖关系,产生了输入依赖的多尺度表示。Y被馈送到一个注意力增强的状态空间模型中,下文将详细说明。在实践中,我们按照先前的工作[27],将四个阶段的窗口大小分别设置为[11, 9, 7, 7]。

注意力增强状态空间模型(Attention-augmented State Space Model)。与先前使用SSM或SS2D模块进一步编码基于注意力的输出以进行全局建模的Transformer-Mamba混合模型不同,我们提出了一种新颖的注意力增强状态空间模型(A2SSM),它有效地利用预先计算的自注意力图来提升SSM的空间感知和动态建模能力。如图3(b)所示,输入 被展平并投影为三个输入相关的序列:∆, B, 和 。然后,∆, 和一个可学习向量 通过选择性扫描操作用于生成一个因果隐藏状态图(HSM),其中第t个令牌动态地融合了所有先前位置的令牌。注意, 被省略了,因为先前的工作[15], [19], [20], [27]为了计算效率将其设置为1。

在普通SSM中,S和 进行逐元素相乘以实现全局上下文调制。然而,我们的A2SSM可以有效地与自注意力集成以执行更强大的全局建模。我们首先重塑 并沿通道维度将其拆分为 ,然后对其应用预先计算的注意力图 。具体来说, 和 被视为“值”分量,其多尺度邻域分别使用 和 进行动态聚合。这背后的动机是 和 在不同范围内捕获了无因果性的动态亲和力。特别是, 密集地捕获了每个令牌与其邻居之间的动态亲和力,而 稀疏地捕获了每个令牌与一组规则间隔的远处令牌之间的动态亲和力。同时, 和 由于其窗口形的空间范围而具有固有的归纳偏置。因此,通过考虑 和 ,得到的HSMs 不仅具有动态增强的、关于二维空间而非一维序列的空间相干性和依赖性,而且还抑制了SSM或SS2D中顺序扫描引入的因果性。此外,我们注意力图的归纳偏置有助于感知二维图像结构。因此,我们的A2SSM提高了普通SSM的空间感知和动态建模能力。接下来, 和 沿通道维度连接,然后与重塑后的C'逐元素相乘,以实现增强的全局上下文调制。剩余的操作遵循普通SSM,其中通过在学习权重向量D与输入Y集成后添加加权残差连接,生成A2SSM的最终输出。

与我们早期尝试[27](简单地堆叠局部注意力和SSM层)相比,本扩展版本中的MASS混合器更深度地集成了注意力机制与状态空间模型,从而产生了更强大的混合架构。总体而言,我们的MASS混合器可以形式化地表示为:

3.4 架构变体

为了使不同设备上更多的潜在应用成为可能,我们的A2Mamba有5种架构变体,包括Nano(N)、Tiny(T)、Small(S)、Base(B)和Large(L)。如表1所列,我们通过调整每个阶段的通道数和块数来控制模型大小。例如,A2Mamba-S有4个阶段,通道数分别为[64, 128, 320, 512],深度为 。四个阶段中的注意力头数分别为[2, 4, 10, 16]。四个阶段中使用的窗口大小分别为 。

表 1

A2Mamba模型变体的配置。

A2Mamba

通道数

块数

头数

窗口大小

Nano

[32, 64, 128, 192]

[2,2, 8,2]

[2, 2, 4, 8]

[11, 9, 7,7]

Tiny

[48, 96, 256, 448]

[2, 2, 10, 2]

[2, 4, 8, 16]

[11, 9, 7,7]

Small

[64, 128, 320, 512]

[2, 4, 12, 4]

[2, 4, 10, 16]

[11, 9, 7,7]

Base

[96, 192, 384, 512]

[4, 6, 12, 6]

[4, 8, 12, 16]

[11, 9, 7,7]

Large

[112, 224, 512, 720]

[4, 6, 12, 6]

[4, 8,16, 30]

[11,9, 7, 7]

3.5 用于改进语义分割的SegMAN-V2

概述。如同在我们的初步工作[27]中一样,除了主干架构(A2Mamba),我们还进一步提出了一个专门为语义分割量身定制的解码器。如图4所示,我们的解码器聚合了A2Mamba中多个抽象层次的特征(即从阶段1的低级特征到阶段4的高级特征),与先前的工作[51], [55]相同。具体来说,我们使用三个并行的1×1卷积层将阶段{2,3,4}中的特征图投影到较低的维度。然后,我们使用双线性插值对来自阶段3和4的投影特征图进行上采样,以匹配来自阶段2的投影特征图的空间维度。三个转换后的特征图被连接起来并通过另一个1×1卷积层,产生一个融合的特征图 。接下来,F被进一步编码,通过多个操作符,包括全局平均池化(GAP)以获得图像级全局上下文,恒等映射以保留原始信息和平滑训练,以及一个新的基于MASS的多尺度细化(MM-Refine)模块以捕获丰富的多尺度上下文信息。这些操作符的输出被连接起来,随后馈入一个线性层,接着是一个双线性插值层,产生一个特征图 。之后,我们执行低级增强以进一步细化空间细节[52], [56]。即,A2Mamba中阶段1的输出被线性投影到一个低维特征空间 ,该特征图与 连接,并馈入一个1×1卷积层,以融合低级空间细节和高级上下文信息。最后,融合的特征图被上采样以产生密集的分割预测。通过集成A2Mamba和这个解码器,我们获得了一个用于语义分割的升级网络架构,称为SegMAN-V2。

在这里插入图片描述
在这里插入图片描述

MM-Refine。为了将多尺度丰富的上下文信息封装到上述解码器中,在这项工作中,我们进一步提出了MM-Refine模块,它是[27]中MMSCopE模块的升级版本。如图4所示,我们通过使用更少的参数同时减少信息损失来改进MMSCopE [27]中的下采样操作。具体来说,在第一个分支中,F首先通过一个像素重排(pixel unshuffle)层实现无损下采样,然后馈入一个步长(stride)=2的 卷积以获得 。与MMSCopE直接使用像素重排后接1×1卷积将分辨率降低到 不同,我们的渐进式下采样方法可以更好地减轻信息损失。在第二个分支中,我们首先使用一个步长=2的 卷积获得一个中间特征 ,然后使用另一个步长=2的 卷积进一步降低分辨率以获得 。同时, 也被馈入一个像素重排层,后接一个1×1卷积,将其分辨率降低到 ,得到 。这背后的动机是高效地捕获不同尺度的多个区域聚合上下文,即 代表了多个粒度上的语义信息。与MMSCopE相比,MM-Refine的下采样方法更加渐进,并且使用更少的卷积层,从而具有更高的效率。最后, 沿通道维度连接,并馈入提出的MASS混合器,后接FFN和双线性上采样层。注意,由于特征分辨率较小,这里的MASS混合器采用全局自注意力,而不是第3.3节中使用的多尺度自注意力。由于 封装了多尺度信息,MASS可以为具有广泛大小范围的对象捕获丰富的上下文信息。

尽管我们谨慎地使用了渐进式下采样,

表 2

在ImageNet-1K上使用224×224输入的图像分类全面比较。#F和#P分别表示模型的FLOPs和参数数量。类型指模型类型,其中“C”、“T”、“M”和“H”分别指ConvNet、Transformer、Mamba和混合模型。

方法

类型

# P (M)

# F (G)

准确率 (%)

方法

类型

# P (M)

# F (G)

准确率 (%)

PVTv2-B0 [8]

T

4

0.6

70.5

Swin-S [3]

T

50

8.7

83.0

QuadMamba-Li [57]

M

5

0.8

74.2

ConvNeXt-S [53]

C

50

8.7

83.1

MSCAN-T [51]

C

4

0.9

75.9

MambaVision-S [22]

H

50

7.5

83.3

ConvNeXt-V2-A [35]

C

4

0.5

76.2

FocalNet-S [40]

C

50

8.7

83.5

EfficientVMamba-T [18]

M

6

0.8

76.5

InceptionNeXt-S [58]

C

49

8.4

83.5

UniRepLKNet-A [37]

C

4

0.6

77.0

PVTv2-B4 [8]

T

63

10.1

83.6

MSVMamba-N [59]

M

7

0.9

77.3

VMamba-S [15]

M

50

8.7

83.6

SegMAN-T Encoder [27]

H

4

0.7

76.2

NAT-S [5]

T

51

7.8

83.7

A2Mamba-N

H

4

0.8

78.7

LocalVMamba-S [16]

M

50

11.4

83.7

PVTv2-B1[8]

T

14

2.1

78.7

RDNet-S [60]

C

50

8.7

83.7

EffcientVMamba-S[18]

M

11

1.3

78.7

QuadMamba-B [57]

M

50

9.3

83.8

MSVMamba-M [59]

M

12

1.5

79.8

SLaK-S [36]

C

55

9.8

83.8

RegionViT-T [61]

T

14

2.4

80.4

UniFormer-B [47]

H

50

8.3

83.9

MPViT-XS [62]

T

11

2.9

80.9

PeLK-S[38]

C

50

10.7

83.9

ConvNeXt-V2-N [35]

C

16

2.5

81.2

UniRepLKNet-S [37]

C

56

9.1

83.9

BiFormer-T [6]

T

13

2.2

81.4

HorNet-S [39]

C

50

8.8

84.0

Conv2Former-N [35]

C

15

2.2

81.5

MSVMamba-S [59]

M

50

8.8

84.1

UniRepLKNet-N [37]

C

18

2.8

81.6

MambaOut-S [21]

C

48

9.0

84.1

NAT-M [5]

T

20

2.7

81.8

Conv2Former-S [63]

C

50

8.7

84.1

SMT-T [64]

H

12

2.4

82.2

InternImage-S [65]

C

50

8.0

84.2

RMT-T [25]

T

14

2.7

82.4

SparX-Mamba-S [20]

M

47

9.3

84.2

TransNeXt-M [26]

T

13

2.7

82.5

BiFormer-B [6]

T

57

9.8

84.3

A2Mamba-T

H

15

2.7

83.0

MogaNet-B [30]

C

44

9.9

84.3

Swin-T [3]

T

28

4.5

81.3

MLLA-S [23]

T

43

7.3

84.4

EfficientVMamba-B [18]

M

33

4.0

81.8

MaxViT-S [10]

H

69

11.7

84.5

PVTv2-B2 [8]

T

25

4.0

82.0

CAFormer-M36 [24]

H

57

12.8

84.5

ConvNeXt-T [53]

C

29

4.5

82.3

Spatial-Mamba-S [19]

M

43

7.1

84.6

FocalNet-T [40]

C

29

4.2

82.3

TransNeXt-S [26]

T

50

10.3

84.7

InceptionNeXt-T [58]

C

28

4.5

82.4

RMT-B [25]

T

54

10.4

85.0

QuadMamba-S [57]

M

31

5.5

82.4

SegMAN-B Encoder [27]

H

45

9.9

85.1

ConvNeXt-V2-T [35]

C

29

4.5

82.5

A2Mamba-B

H

51

10.7

85.7

SLaK-T [36]

C

30

5.0

82.5

Swin-B [3]

T

88

15.4

83.5

VMamba-T [15]

M

29

4.9

82.6

FocalNet-B [40]

C

89

15.4

83.7

PeLK-T [38]

C

29

5.6

82.6

PVTv2-B5 [8]

T

82

11.8

83.8

CSWin-T [4]

T

23

4.5

82.7

ConvNeXt-B [53]

C

89

15.4

83.8

LocalVMamba-T [16]

M

26

5.7

82.7

VMamba-B [15]

M

89

15.4

83.9

MambaVision-T2 [22]

H

35

5.1

82.7

SLaK-B[36]

C

95

17.1

84.0

MambaOut-T [21]

C

27

4.5

82.7

InceptionNeXt-B [58]

C

87

14.9

84.0

HorNet-T [39]

C

22

4.0

82.8

CSWin-B [4]

T

78

15.0

84.2

RDNet-T [60]

C

24

5.0

82.8

MambaVision-B [22]

H

98

15.0

84.2

UniFormer-S [47]

H

22

3.6

82.9

PeLK-B[38]

C

85

15.8

84.2

MPViT-S [62]

T

23

4.7

83.0

ConvNeXt-V2-B [35]

C

89

18.3

84.2

MSVMamba-T [59]

M

32

5.1

83.0

MPViT-B [62]

T

75

15.4

84.3

NAT-T [5]

T

28

4.3

83.2

NAT-B [5]

T

90

16.4

84.3

Conv2Former-T [63]

C

27

4.4

83.2

HorNet-S [39]

C

87

13.7

84.3

UniRepLKNet-T [37]

C

31

4.9

83.2

MSVMamba-B [59]

M

91

15.6

84.3

MogaNet-S [30]

C

25

5.0

83.4

RDNet-B [60]

C

87

16.3

84.3

CMT-S [43]

T

25

4.0

83.5

Conv2Former-B [63]

C

91

15.4

84.4

MLLA-T [23]

T

25

4.2

83.5

SparX-Mamba-B [20]

M

87

15.9

84.4

Spatial-Mamba-T [19]

M

27

4.5

83.5

MogaNet-L [30]

C

96

15.9

84.4

SparX-Mamba-T [20]

M

27

5.2

83.5

TransNeXt-B [26]

T

95

15.9

84.5

InternImage-T [65]

C

30

5.0

83.5

MaxViT-B [10]

H

120

18.4

84.7

CAFormer-S18 [24]

H

26

4.1

83.6

InternImage-B [65]

C

97

24.0

84.8

MaxViT-T [10]

H

31

5.6

83.7

MLLA-B[23]

T

96

16.0

84.9

SMT-S[64]

H

21

4.7

83.7

Spatial-Mamba-B [19]

M

81

16.2

84.9

BiFormer-S [6]

T

26

4.5

83.8

CAFormer-B36 [24]

H

95

16.8

85.3

TransNeXt-T [26]

T

28

5.7

84.0

RMT-L [25]

T

66

23.2

85.3

RMT-S [25]

T

27

4.8

84.1

SegMAN-L Encoder [27]

H

96

19.6

85.5

SegMAN-S Encoder [27]

H

26

4.1

84.0

A2Mamba-L

H

81

16.8

86.1

A2Mamba-S

H

31

5.4

84.7

某些重要的局部线索可能仍然会丢失。为了解决这个问题,我们引入了一个基于5×5空洞RepConv [37]的额外轻量级卷积快捷方式,以加强局部细节建模能力。最终的特征F′不仅拥有丰富的多尺度上下文信息,而且保留了局部细节,这两者对于高质量的语义分割都是不可或缺的。

4 实验

4.1 图像分类

设置。我们在ImageNet-1K数据集[66]上评估我们的方法,并遵循与先前工作[3], [23]相同的实验设置以确保公平比较。具体来说,我们使用AdamW优化器[67]训练所有模型300个周期。A2Mamba-N、-T、-S、-B和-L模型的随机深度率[68]分别设置为0.05、0.1、0.2、0.4和0.5。在224×224输入上预训练基础和大型模型后,我们进一步在384×384输入上对它们进行30个周期的微调,以评估高分辨率输入下的性能。所有实验在8个NVIDIA H800 GPU上运行。

结果。如表2所示,我们先前的工作SegMAN编码器已经实现了与state-of-the-art (SOTA)视觉主干模型相竞争的性能。然而,升级版本A2Mamba在所有先前的基于ConvNet、Transformer和Mamba的模型上都带来了显著的性能改进。具体来说,我们的A2Mamba-S模型实现了令人印象深刻的84.7% top-1准确率,分别超过了RMT-S [25]和TransNeXt-T [26] 0.6%和0.7%。此外,A2Mamba-B进一步将top-1准确率提高到85.7%,以大约一半的计算复杂度超过了MLLA-B_[23] 0.4%。值得注意的是,我们的A2Mamba-L实现了显著的86.1% top-1准确率,以更少的复杂度显著超过了CAFormer-B36 [24] 0.6%。如表3所列,在384×384输入上微调A2Mamba-B得到了86.4%的top-1准确率,这优于TransNeXt-B和RMT-L,同时仅使用大约一半的计算复杂度。此外,A2Mamba-L进一步将top-1准确率提高到86.7%,显著超过了其对应的模型。

表 3

使用384×384输入的图像分类性能比较。

方法

类型

# P (M)

# F (G)

准确率 (%)

Swin-B [3]

T

88

47

84.5

CSWin-B [4]

T

78

47

85.4

ConvNeXt-B [53]

C

89

45

85.1

ConvNeXt-L [53]

C

198

101

85.5

MaxViT-S [10]

H

69

36

85.2

MaxViT-B [10]

H

120

74

85.7

TransNeXt-S [46]

H

50

32

86.0

TransNeXt-B [46]

H

106

56

86.2

RMT-L [25]

T

95

59

85.5

A2Mamba-B

H

51

34

86.4

A2Mamba-L

H

95

54

86.7

4.2 目标检测和实例分割

设置。我们使用COCO 2017数据集[50]评估我们的A2Mamba网络架构在目标检测和实例分割任务上的性能。遵循Swin [3]的实验设置,我们采用Mask R-CNN [69]和Cascade Mask R-CNN [29]框架。我们的主干网络在ImageNet-1K上预训练,然后使用多尺度训练(3× + MS schedule)进行36个周期的微调。

结果。如表4和表5所示,我们的模型在目标检测和实例分割方面取得了令人印象深刻的性能。例如,使用Mask R-CNN框架,A2Mamba-S在 上以显著优势超过了UniFormer-S [47] 3.3%/1.9%,甚至以仅约一半的复杂度超过了CSWin-B 0.7%/0.4%。使用Cascade Mask R-CNN框架,我们的方法表现出更显著的性能提升。例如,A2Mamba-B在 上以显著优势超过了CAFormer-S36[24] 2.2%/1.6%,并且还以节省约三分之一的参数,在 上显著超过了MambaVision-B[22] 2.6%/1.9%。这种显著的性能提升有效地证明了我们的方法在建模多尺度和全局上下文方面的强大能力。

表 4

在COCO数据集上使用Mask R-CNN的主干架构比较。FLOPs以800×1280的图像分辨率计算。

主干网络

# P (M)

#F (G)

ConvNeXt-T [35]

48

262

46.2

41.7

FocalNet-T [40]

49

268

48.0

42.9

InternImage-T [65]

49

270

49.1

43.7

RDNet-T [60]

43

278

47.3

42.2

MogaNet-S [30]

45

272

48.5

43.1

VMamba-T [15]

50

271

48.8

43.7

MSVMamba-T [59]

52

275

48.7

43.4

Spatial-Mamba-T [19]

46

261

49.3

43.6

SparX-Mamba-T [20]

47

279

50.2

44.7

Swin-T [3]

48

267

46.0

41.6

PVTv2-B2 [8]

45

309

47.8

43.1

CSWin-T [4]

42

279

49.0

43.6

MPViT-S [62]

43

268

48.4

43.9

UniFormer-S [47]

41

269

48.2

43.4

NAT-T [5]

48

258

47.8

42.6

SMT-S [64]

40

265

49.0

43.4

RMT-S [25]

45

288

50.7

44.9

A2Mamba-S

49

283

51.5

45.3

ConvNeXt-S [53]

70

348

47.9

42.9

FocalNet-S [40]

72

365

49.3

43.8

InternImage-S [65]

69

340

49.7

44.5

MogaNet-B [30]

63

373

50.3

44.4

VMamba-S [15]

70

384

49.9

44.2

MSVMamba-S [59]

70

349

49.7

44.2

Spatial-Mamba-S[19]

63

315

50.5

44.6

SparX-Mamba-S [20]

63

339

51.0

45.2

Swin-S [3]

67

361

48.2

43.2

PVTv2-B3 [8]

69

354

48.4

43.2

CSWin-S[4]

65

397

50.0

44.5

UniFormer-B [47]

54

342

50.3

44.8

NAT-S [5]

69

399

48.4

43.2

SMT-B [64]

70

330

48.4

43.2

RMT-B [25]

52

328

49.8

44.0

A2Mamba-B

73

422

52.2

46.1

ConvNeXt-B [53]

108

486

48.5

43.5

FocalNet-B [40]

111

507

49.8

44.1

InternImage-B [65]

115

507

50.3

44.8

MogaNet-L[30]

102

495

50.5

44.5

SparX-Mamba-B [20]

103

498

51.8

45.8

Swin-B [3]

107

496

48.6

43.3

PVTv2-B5 [8]

102

557

48.4

42.9

CSWin-B [4]

97

526

50.8

44.9

MPViT-B [62]

95

503

49.5

44.5

A2Mamba-L

113

552

53.0

46.8

表 5

在COCO数据集上使用Cascade Mask R-CNN的主干架构比较。FLOPs以800×1280的图像分辨率计算。

主干网络

# P (M)

# F (G)

ConvNeXt-T [53]

86

741

50.4

43.7

HorNet-T [39]

80

730

51.7

44.8

RDNet-T [60]

81

757

51.6

44.6

PeLK-T [38]

86

770

51.4

44.6

UniRepLKNet-T [37]

89

749

51.8

44.9

MogaNet-S [30]

78

750

51.6

45.1

MambaVision-T [22]

86

740

51.1

44.3

Swin-T [3]

86

745

50.4

43.7

PVTv2-B2 [8]

83

788

51.1

-

CSWin-T [4]

80

757

52.5

45.3

UniFormer-S [47]

79

747

52.1

45.2

NAT-T [5]

85

737

51.4

44.5

SMT-S [64]

78

744

51.9

44.7

CAFormer-S18 [24]

69

733

51.5

44.6

RMT-S [25]

83

767

53.2

46.1

A2Mamba-S

87

762

54.0

46.6

ConvNeXt-S [53]

108

827

51.9

45.0

HorNet-S [39]

108

827

52.7

45.6

RDNet-S [60]

108

832

52.3

45.3

PeLK-S [38]

108

874

52.2

45.3

UniRepLKNet-S [37]

113

835

53.0

45.9

MogaNet-B [30]

101

851

52.6

46.0

MambaVision-S [22]

106

828

52.3

45.2

Swin-S [3]

107

838

51.8

44.7

CSWin-S [4]

92

820

53.7

46.4

UniFormer-B [47]

107

878

53.8

46.4

NAT-S [5]

108

809

52.0

44.9

CAFormer-S36 [24]

81

811

53.2

46.0

RMT-B [25]

111

906

54.5

47.2

A2Mamba-B

108

889

55.4

47.6

ConvNeXt-B [53]

146

964

52.7

45.6

HorNet-B [39]

144

969

53.3

46.1

RDNet-B [60]

144

971

52.3

45.3

PeLK-B [38]

147

1028

52.9

45.9

MogaNet-L [30]

149

974

53.3

46.1

MambaVision-B [22]

145

964

52.8

45.7

Swin-B [3]

145

982

51.9

45.0

CSWin-B [4]

135

1004

53.9

46.4

NAT-B [5]

147

931

52.5

45.2

CAFormer-M36 [24]

106

920

53.8

46.5

A2Mamba-L

151

1027

55.6

48.1

4.3 语义分割

设置。我们使用ADE20K数据集[48]和UperNet框架[28]评估我们的主干架构(A2Mamba变体)在语义分割上的性能,遵循与Swin[3]相同的训练协议。此外,我们在三个数据集上评估我们的分割网络架构(SegMANV2):ADE20K、Cityscapes[49]和COCO-Stuff [50],使用与SegFormer [55]相同的训练协议。为了公平比较,所有主干网络都使用ImageNet-1K预训练权重进行初始化。

结果。如表6所示,当使用相同的特征解码器来公平比较不同主干的性能时,我们的A2Mamba实现了领先的性能。例如,A2Mamba-S实现了51.6%的显著mIoU,分别以大约一半的参数数量显著超过了InternImage-B_[65] 0.8%和VMamba-B [15] 0.6%。这进一步证明了我们提出的A2Mamba在密集预测任务上的强大性能。另一方面,与其他语义分割模型相比,我们之前的模型SegMAN [27]已经显示出显著的性能优势。然而,SegMAN-V2在SegMAN的基础上进一步改进,实现了更显著的性能提升。例如,SegMAN-V2-S的参数仅为Segformer-B5 [55]的三分之一左右,但在ADE20K、Cityscapes和COCO-Stuff数据集上的mIoU分别高出1.0%、1.4%和1.3%。同时,我们的SegMAN-V2-B在三个数据集上分别显著提高了LRFormer-B [77] 2.5%、1.2%和1.8%。此外,我们的SegMAN-V2-L取得了显著的改进,在三个数据集上分别超过了VWFormer-B5[52] 2.1%、1.8%和1.5%。跨不同数据集和模型规模的持续性能增益验证了我们提出的SegMAN-V2的有效性,它可以通过其基于MASS的特征编码器和基于MM-Refine的特征解码器同时捕获全局上下文、局部细节和多尺度线索。

表 6

在ADE20K数据集上使用UperNet的各种视觉主干性能比较。FLOPs以512×2048的分辨率计算。

主干网络

# P (M)

# F (G)

mIoUss (%)

mIoUMS (%)

ConvNeXt-T [53]

60

939

46.0

46.7

SLaK-T [36]

65

936

47.6

-

InternImage-T [65]

59

944

47.9

48.1

PeLK-T [38]

62

970

48.1

-

MogaNet-S [30]

55

946

49.2

-

VMamba-T [15]

62

949

48.0

48.8

MSVMamba-T [59]

63

953

47.9

48.5

MambaVision-T [22]

55

945

46.0

-

SparX-Mamba-T [20]

50

954

50.0

50.8

Spatial-Mamba-T [19]

57

936

48.6

49.4

CSWin-T [4]

59

959

49.3

50.7

UniFormer-S [47]

52

1008

47.6

48.5

BiFormer-S [6]

55

1025

49.8

50.8

CAFormer-S18 [24]

54

1024

48.9

-

TransNeXt-T [26]

59

978

51.1

51.2

RMT-S [25]

56

970

49.8

-

A2Mamba-S

60

959

51.6

52.0

ConvNeXt-S [53]

82

1027

48.7

49.6

SLaK-S [36]

91

1028

49.4

-

InternImage-S [65]

80

1017

50.1

50.9

PeLK-S [38]

84

1077

49.7

-

UniRepLKNet-S [37]

86

1036

50.5

51.0

MogaNet-B [30]

74

1050

50.1

-

VMamba-S [15]

82

1038

50.6

51.2

MambaVision-S [22]

84

1135

48.2

-

SparX-Mamba-S [20]

73

1039

51.3

52.5

Spatial-Mamba-S [19]

81

992

50.6

51.4

Swin-S [3]

81

1038

47.6

49.5

CSWin-S [4]

65

1057

50.4

51.5

UniFormer-B [47]

80

1027

50.0

51.5

BiFormer-B [6]

88

1227

51.0

50.8

CAFormer-S36 [24]

67

1184

50.8

51.7

TransNeXt-S [26]

80

1197

52.2

-

RMT-B [25]

83

1089

52.0

52.3

A2Mamba-B

80

1111

53.3

53.9

SLaK-B[36]

109

1090

-

-

InternImage-B [65]

135

1172

50.2

51.3

PeLK-B[38]

128

1185

50.8

-

MogaNet-L [30]

126

1237

50.4

-

VMamba-B [15]

113

1176

50.9

-

MambaVision-B [22]

122

1170

51.0

51.6

SparX-Mamba-B [20]

126

1342

49.1

-

Spatial-Mamba-B [19]

115

1181

52.3

53.4

Swin-B [3]

127

1176

51.8

52.6

CSWin-B [4]

121

1188

48.1

49.7

NAT-B [5]

109

1222

51.1

52.2

MPViT-B [62]

123

1137

48.5

49.7

CAFormer-M36 [24]

105

1186

50.3

-

TransNeXt-B [26]

84

1346

51.7

-

RMT-L [25]

121

1268

53.0

53.4

A2Mamba-L

125

1324

52.8

53.4

126

1237

53.7

54.1

表 7

不同分割模型之间的语义分割性能比较。FLOPs以512×512(ADE20K和COCO-Stuff)和1024×2048(Cityscapes)分辨率计算。

方法

# P (M)

ADE20K

Cityscapes

COCO-Stuff

#F (G)

mIoU (%)

#F (G)

mIoU (%)

#F (G)

mIoU (%)

Segformer-B0 [55]

3.8

8.4

37.4

126

76.2

8.4

35.6

SegNeXt-T [51]

4.3

7.7

41.1

62

78.9

7.7

38.7

VWFormer-B0 [52]

3.7

5.8

38.9

112

77.2

5.8

36.2

EDAFormer-T [70]

4.9

5.8

42.3

152

78.7

5.8

40.3

CGRSeg-T [71]

9.4

4.8

42.5

66

78.3

4.8

40.4

SegMAN-T [27]

6.4

6.2

43.0

53

80.3

6.2

41.3

SegMAN-V2-N

6.6

7.4

44.4

67

81.0

7.4

41.9

ViT-CoMer-S [72]

61

296

46.5

-

-

-

-

OCRNet [73]

71

165

45.6

-

-

-

-

Segformer-B2 [55]

28

62

46.5

717

81.0

62

44.6

MaskFormer [74]

42

55

46.7

-

-

-

-

Mask2Former [75]

47

74

47.7

-

-

-

-

SegNeXt-B [51]

28

35

48.5

279

82.6

35

45.8

FeedFormer-B2 [76]

29

43

48.0

523

81.5

43

-

VWFormer-B2 [52]

27

47

48.1

415

81.7

47

45.2

EDAFormer-B [70]

29

32

49.0

606

81.6

32

45.9

CGRSeg-B [71]

36

17

47.3

200

80.2

17

45.2

LRFormer-S [77]

32

40

50.0

295

81.9

40

46.4

SegMAN-S [27]

29

25

51.3

218

83.2

25

47.5

SegMAN-V2-S

32

34

52.0

282

83.8

34

48.0

Segformer-B3 [55]

47

79

49.4

963

81.7

79

45.5

SegNeXt-L [51]

49

70

51.0

578

83.2

70

46.5

VWFormer-B3 [52]

47

63

50.3

637

82.4

63

46.8

LRFormer-B [77]

69

75

51.0

555

83.0

75

47.2

SegMAN-B [27]

52

58

52.6

479

83.8

58

48.4

SegMAN-V2-B

56

66

53.5

552

84.2

66

49.0

ViT-CoMer-B [72]

145

455

48.8

-

-

-

-

Segformer-B5 [55]

85

110

51.0

1150

82.4

110

46.7

VWFormer-B5 [52]

85

96

52.0

1140

82.8

96

48.0

LRFormer-L [77]

113

183

52.6

908

83.2

183

47.9

SegMAN-L [27]

92

97

53.2

796

84.2

97

48.8

SegMAN-V2-L

108

109

54.1

871

84.6

109

49.5

4.4 分析实验

速度比较和增加分辨率的影响。受VMamba [15]的启发,我们评估了不同视觉主干在各种输入分辨率下的推理速度和泛化能力。如表8所列,我们利用在ImageNet-1K上预训练的模型在一系列图像分辨率(包括224×224、512×512和1024×1024)上执行推理,并报告相应的GPU内存消耗(Mem.)和推理吞吐量(Thr.)。三种分辨率使用的批量大小分别为128、32和8。所有实验均在单个NVIDIA L40S GPU上进行。可以观察到,我们提出的A2Mamba与其他基线相比,实现了具有竞争力的效率和更强的泛化能力。例如,使用224×224输入时,A2Mamba-S在准确率上优于RMT-S,并实现了1.5倍更高的吞吐量。当分辨率增加到512×512时,A2Mamba-S在top-1准确率上以显著优势超过RMT-S 8.5%,同时保持近1.7倍的速度提升和更低的内存消耗。此外,当分辨率扩展到1024×1024时,A2Mamba-S在top-1准确率上以显著优势超过RMT-S 29.9%,同时消耗近一半的内存并以2倍的速度运行。此外,一个有趣的现象是,我们发现先进的视觉transformer,如BiFormer、RMT和TransNeXt,在分辨率增大时内存消耗显著增加,速度下降。这是因为尽管使用了高效的注意力机制,但在高分辨率下计算成本仍然显著增加。相比之下,我们的A2Mamba模型有效地避免了这种现象,得益于其线性时间模块(包括高效的自注意力和SSM),这使得计算和内存使用都高效,同时性能强大,使其成为复杂和高分辨率视觉识别任务更有前途的基础模型。

在这里插入图片描述
在这里插入图片描述

有效感受野分析。为了进一步了解A2Mamba相对于先前方法的优越性,我们可视化了有效感受野(ERFs)[78]。具体来说,我们使用从ImageNet-1K验证集中随机采样的500多张分辨率为224×224的图像生成可视化,同时确保所有比较的模型具有可比的复杂度。如图5所示,与使用具有四个并行方向扫描的SS2D的SegMAN编码器-S相比,我们的A2Mamba产生了更大的ERF,表明即使使用单次扫描,注意力增强的SSM也可以拥有更强的全局表示能力。此外,与强大的基于Transformer的模型(包括RMT和TransNeXt)相比,我们的A2Mamba不仅表现出更大的ERF,而且由于多尺度滑动注意力而表现出更强的局部敏感性。总体而言,我们的A2Mamba模型在所有强大的竞争对手中实现了最大的ERF,包括所有先前的基于ConvNet、Transformer和Mamba的模型。

表 8

随着输入分辨率增加,推理速度和泛化能力的比较。

方法

# P (M)

224×224

512×512

1024×1024

# F (G)

Mem. (MB)

Thr. (imgs/s)

Acc. (%)

# F (G)

Mem. (MB)

Thr. (imgs/s)

Acc. (%)

# F (G)

Mem. (MB)

Thr. (imgs/s)

Acc. (%)

ConvNeXt-T [53]

29

4.5

1507

1102

82.1

23.3

3865

212

78.0

93.0

3747

70

55.4

ConvNeXt-S [53]

50

8.7

3263

691

83.1

45.5

3965

133

80.4

182.0

3847

43

65.4

ConvNeXt-B [53]

89

15.4

3343

477

83.8

80.3

4921

117

80.6

321.0

4715

28

52.9

FocalNet-T [40]

29

4.5

7151

1102

82.1

23.5

9847

212

78.5

94.0

11065

52

62.2

FocalNet-S [40]

50

8.7

8679

691

83.5

45.7

12685

133

81.3

183.0

15267

33

67.7

FocalNet-B [40]

89

15.4

12155

477

83.8

80.6

15737

88

82.9

322.0

20858

22

72.3

MogaNet-S [30]

25

5.0

4803

766

83.8

25.9

5921

70

78.2

104.0

5967

17

57.0

MogaNet-B [30]

44

9.9

4961

373

84.3

51.7

6123

53

80.2

207.0

6053

13

44.8

MogaNet-L [30]

83

15.9

5159

282

84.7

82.9

7489

42

83.3

332.0

7215

22

73.7

VMamba-T [15]

29

4.9

1324

1118

82.6

25.6

5691

226

80.9

103.0

5699

56

57.4

VMamba-S [15]

50

8.7

544

869

83.6

45.5

7431

72

82.6

185.0

7489

16

74.8

VMamba-B [15]

89

15.4

439

512

83.9

80.2

11035

84

83.3

321.0

11527

21

61.9

Swin-T [3]

28

4.5

1324

1118

81.3

26.6

5777

213

79.0

153.0

5521

54

28.0

Swin-S [3]

50

8.7

4893

596

83.0

45.5

7489

89

81.3

342.0

7215

22

39.3

Swin-B [3]

88

15.4

4961

544

83.5

80.9

8767

87

83.0

342.0

7215

22

47.0

MPViT-XS [62]

11

2.9

3511

808

80.9

15.6

4243

212

78.0

62.0

4237

48

25.1

MPViT-S [62]

23

4.7

3599

380

83.0

40.8

3841

156

77.1

163.0

3841

39

51.6

MPViT-B [62]

75

16.4

5981

1740

84.3

71.7

4775

109

78.8

287.0

4773

27

64.8

NAT-M[5]

20

2.7

2747

1740

81.8

14.2

3191

330

70.7

57.0

3191

81

38.1

NAT-T [5]

28

4.3

2771

1287

83.2

40.8

3841

156

77.1

163.0

3841

39

39.3

NAT-S [5]

51

7.8

3265

823

83.7

66.9

11085

42

78.0

430.0

21761

7

47.0

NAT-B [5]

90

13.7

4087

644

84.3

105.1

18043

22

83.8

555.0

38633

3

69.6

BiFormer-T [6]

13

2.2

4567

1103

81.4

16.3

7591

135

71.3

117.0

14507

21

30.0

BiFormer-S [6]

26

4.5

4635

341

84.3

66.9

11085

42

78.0

430.0

21761

7

52.3

BiFormer-B [6]

57

9.8

6419

225

84.8

105.1

18043

22

83.8

555.0

38633

3

69.6

MLLA-T [23]

25

4.2

4429

580

83.5

18.2

6881

106

74.4

81.0

17217

16

34.2

MLLA-S [23]

43

7.3

4505

341

84.4

26.9

7035

81

74.6

122.0

10981

13

42.2

MLLA-B[23]

96

13.7

6427

1054

84.9

57.7

8781

42

80.7

258.0

13675

9

50.9

TransNeXt-M [26]

13

2.7

4345

769

82.5

14.3

5401

139

76.2

89.0

3499

30

45.7

TransNeXt-T [26]

28

5.7

5977

527

84.0

28.4

6551

51

82.4

213.0

5401

11

71.5

TransNeXt-S [26]

50

10.3

6909

341

84.7

52.3

6747

37

83.1

357.0

5401

10

72.7

TransNeXt-B [26]

106

18.4

7691

225

85.0

106.7

11853

14

84.8

463.0

18957

6

68.1

RMT-T [25]

14

2.7

869

512

82.4

18.2

3325

387

70.3

14.0

3499

91

45.7

RMT-S [25]

27

4.8

512

260

84.1

26.9

5375

122

74.6

122.0

10981

13

42.2

RMT-B [25]

54

10.4

260

176

85.0

57.7

8247

42

80.7

258.0

13675

9

50.9

RMT-L [25]

95

19.6

176

106

85.5

106.7

8329

29

84.9

463.0

18957

6

56.6

SegMAN-T Encoder [27]

4

0.7

2118

708

76.2

4.4

4141

445

74.4

18.0

3889

108

43.9

SegMAN-S Encoder [27]

26

4.1

4417

269

84.0

21.4

5005

139

82.4

89.0

3499

30

45.7

SegMAN-B Encoder [27]

45

9.9

6551

200

85.1

52.3

5935

51

83.0

213.0

5401

11

66.5

SegMAN-L Encoder [27]

96

16.8

6747

140

85.5

88.3

8637

60

84.3

357.0

5401

10

72.1

A2Mamba-N

4

0.8

3273

2486

78.7

4.4

4141

445

74.4

18.0

3889

108

43.9

A2Mamba-T

15

2.7

4417

762

83.0

14.0

5005

320

83.1

60.0

3499

48

66.5

A2Mamba-S

31

5.4

6551

258

84.7

22.8

5935

140

83.1

99.0

5401

32

72.1

A2Mamba-B

51

10.7

6747

200

85.7

44.0

8637

60

84.3

177.0

5401

14

74.8

A2Mamba-L

95

17.4

7691

140

86.1

71.6

9665

48

85.3

287.0

5401

11

75.4

4.5 消融研究

设置。我们在图像分类和语义分割任务上进行了全面的消融研究,以评估我们模型中各个组件的有效性。具体来说,我们按照第4.1节中概述的训练设置,在ImageNet-1K数据集上训练每个模型变体300个周期。随后,我们使用与SegFormer [55]相同的所有设置,在ADE20K数据集上对预训练模型进行160K次迭代的微调。除非另有说明,分割网络都建立在我们基于MM-Refine的解码器之上。FLOPs和吞吐量是在512×512图像分辨率、批量大小为32的情况下,使用主干在单个NVIDIA L40S GPU上测量的,遵循[27]的协议。

从LASS到MASS的路线图。我们提供了LASS混合器[27]向MASS混合器演变的详细过程。如表9所列,我们首先将A2Mamba-T模型中的所有MASS混合器替换为LASS混合器,得到我们的基线模型,其top-1准确率为82.2%,mIoU为48.2%。然后,我们将LASS中的Natten [5]替换为我们第3.3节讨论的自适应多尺度注意力(AMA),在top-1/mIoU上分别产生0.3%/0.5%的改进。这突出了自适应多尺度建模的重要性,特别是在语义分割任务中。接下来,我们将SS2D[15]替换为普通SSM[13],这导致性能显著下降,top-1准确率为81.4%,mIoU为47.3%。这表明仅使用单向扫描严重损害了模型捕获输入图像上下文信息的能力。然而,当我们将SSM替换为第3.3节讨论的提出的A2SSM时,性能在top-1/mIoU上分别大幅提高了1.3%/1.9%,证明了我们A2SSM强大的空间感知和动态能力。最后,我们向模型引入了一个门控机制[13], [30],这产生了我们MASS混合器的最终版本,与基线模型相比,实现了改进的性能和效率。

表 9

逐步将LASS [27]演变为我们提出的MASS的详细路线图。

模型

# P (M)

# F (G)

Thr. (imgs/s)

准确率 (%)

mIoU (%)

Baseline

13

14.5

176

82.2

48.2

Natten → AMA

13

14.5

172

82.5

48.7

SS2D → SSM

13

12.0

256

81.4

47.3

SSM →A2SSM

13

13.3

235

82.7

49.2

w Gate

15

14.0

220

82.9

49.7

自适应空洞率的影响。我们研究了AMA中空洞率(r)对模型性能的影响。基线模型是A2Mamba-T,它使用公式1中描述的自适应空洞率。首先,我们将空洞率设置为固定值,即3、5和7。如表10所示,使用固定的r对图像分类性能的影响可以忽略不计,但会导致语义分割性能显著下降。此外,我们还将双分支AMA修改为四分支版本,其中一个分支是常规滑动局部注意力,其余三个分支是空洞局部注意力,分别具有 。然而,这种修改并没有带来性能改进,反而降低了效率。这些结果表明,根据输入分辨率自适应调整空洞率可以在密集预测中捕获更有用的多尺度信息。

表 10

AMA中空洞率的调查研究。

模型

# P (M)

# F (G)

Thr. (imgs/s)

准确率 (%)

mIoU (%)

Baseline

15

14

220

83.0

49.7

Dilation=3

15

14

221

82.8(−0.2)

49.1(−0.6)

Dilation=5

15

14

221

83.0(+0.0)

49.3(−0.3)

Dilation=7

15

14

221

83.0(+0.0)

49.2(−0.5)

Dilation={3, 5, 7}

15

14

209

82.8(−0.2)

49.5(−0.2)

共享注意力图的影响。我们A2SSM的核心是使用具有共享多尺度注意力图的交叉注意力变体来有效增强SSM的空间感知和动态建模能力。为了验证这一点,我们以A2Mamba-T作为基线模型,并将交叉注意力操作替换为其他相关操作,包括空洞RepConv [37]和DCNv2 [9]。为了公平比较,我们对空洞RepConv和DCNv2使用与原始局部注意力窗口大小相同的核大小。注意,我们使用DCNv2的深度wise版本,因为原始版本会产生显著的计算成本。如表11所列,使用空洞RepConv或DCNv2都会导致性能和效率显著下降。这是因为这些操作符无法动态捕获令牌之间的多尺度关系,导致嵌入到SSM中时无法进行有效的空间结构感知和动态增强。

表 11

不同混合器对SSM的影响。

模型

# P (M)

# F (G)

Thr. (imgs/s)

准确率 (%)

mIoU (%)

Baseline

15

14.0

220

83.0

49.7

Dilated RepConv [37]

16

13.9

201

81.9

47.9

DCNv2 [79]

16

14.6

93

82.1

48.3

令牌混合器的比较。遵循我们的会议版本[27],我们将SegMAN-S编码器中的令牌混合器替换为其他视觉主干的令牌混合器,包括PVT[7]、MaxViT[10]、ACmix[44]和BiFormer [6],以公平比较不同的令牌混合器。如表12所示,我们的MASS令牌混合器在分类和分割任务上都取得了显著的性能改进,同时保持了具有竞争力的计算成本。性能提升可归因于我们方法的互补性,它建模了自适应多尺度线索和更鲁棒的全局上下文。

表 12

不同令牌混合器的比较。

令牌混合器

# P (M)

# F (G)

Thr. (imgs/s)

准确率 (%)

mIoU (%)

PVT [7]

30

22.0

169

82.8

49.1

MaxViT [10]

25

29.8

96

83.5

47.2

ACmix [44]

25

19.3

104

83.1

48.6

BiFormer [6]

25

30.5

97

82.9

48.8

LASS [27]

26

21.4

139

84.0

51.3

MASS

27

22.8

160

84.3

51.8

从SegMAN解码器到SegMAN-V2解码器的路线图。SegMAN-V2解码器是SegMAN解码器[27]的升级版本,旨在实现更细粒度的语义分割。为此,我们提供了一个详细的路线图来说明我们SegMAN-V2解码器的性能改进。所有实验均使用A2Mamba-T作为编码器,在ADE20K数据集上进行,遵循与SegFormer [55]相同的训练设置。分割网络的FLOPs和吞吐量使用512×512输入分辨率、批量大小为32,在单个NVIDIA L40S GPU上进行评估,遵循[27]的设置。如表13所列,我们首先将MMSCopE中的原始下采样修改为第3.5节中描述的更渐进的下采样,从而提高了性能和效率。接下来,我们引入一个基于空洞RepConv的局部嵌入来补充丢失的局部细节,我们的实验表明5×5(k=5)的核大小实现了最佳权衡。随后,我们用MASS混合器替换SS2D,导致进一步的显著性能改进。最后,我们采用低级增强,这在不明显影响效率的情况下带来了适度的性能提升。

表 13

逐步将SegMAN解码器演变为我们SegMAN-V2解码器的详细路线图。

模型

# P (M)

# F (G)

Thr. (imgs/s)

mIoU (%)

MMSCopE [27]

17

18.1

142

48.5

w Progressive Down.

16

17.1

150

48.8

w Local Embed. (k=3)

17

17.2

147

48.9

w Local Embed. (k=5)

17

17.2

143

49.1

w Local Embed. (k=7)

17

17.2

136

49.1

w MASS

18

17.1

140

49.5

w Low Level

18

17.6

137

49.7

5 结论

这项工作提出了A2Mamba,一个强大的Transformer-Mamba混合视觉主干架构,其特点是采用了一个统一的令牌混合器,称为多尺度注意力增强状态空间模型(MASS)。MASS模块自适应地提取多尺度上下文,同时存储中间注意力图以进一步增强后续SSM层的全局感知和动态建模能力。我们在各种视觉任务上评估了A2Mamba,包括图像分类和密集预测,并证明了其相对于现有强大的基于ConvNet、Transformer和Mamba的视觉主干架构的显著性能优势。

本文参与 腾讯云自媒体同步曝光计划,分享自微信公众号。
原始发表:2025-09-08,如有侵权请联系 cloudcommunity@tencent.com 删除

本文分享自 AI智韵 微信公众号,前往查看

如有侵权,请联系 cloudcommunity@tencent.com 删除。

本文参与 腾讯云自媒体同步曝光计划  ,欢迎热爱写作的你一起参与!

评论
登录后参与评论
0 条评论
热度
最新
推荐阅读
目录
  • 论文简要总结
  • 论文翻译全文
  • 摘要
  • 1 引言
  • 2 相关工作
    • 2.1 ConvNets
    • 2.2 视觉Transformer
    • 2.3 视觉Mamba
    • 2.4 混合视觉主干架构
  • 3 方法
    • 3.1 SegMAN回顾
    • 3.2 A2Mamba的整体架构
    • 3.3 MASS令牌混合器
    • 3.4 架构变体
    • 3.5 用于改进语义分割的SegMAN-V2
  • 4 实验
    • 4.1 图像分类
    • 4.2 目标检测和实例分割
    • 4.3 语义分割
    • 4.4 分析实验
    • 4.5 消融研究
  • 5 结论
领券
问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档