这篇文章的主要内容可以概括如下:
背景与动机:
-
近年来,状态空间模型(SSM)在长序列建模中展现出巨大潜力,尤其是Mamba模型在硬件感知设计上的高效性。
-
然而,现有的SSM模型在处理视觉数据时面临挑战,因为视觉数据对位置敏感且需要全局上下文。
研究目标:
-
提出一种新的视觉骨干网络,称为Vision Mamba(Vim),基于双向Mamba块,旨在高效地进行视觉表示学习。
方法与创新:
-
Vim通过双向状态空间模型(SSM)进行数据依赖的全局视觉上下文建模,并通过位置嵌入实现位置感知的视觉理解。
-
与传统的视觉Transformer(如DeiT)相比,Vim在计算复杂度和内存使用上具有显著优势,且无需依赖自注意力机制。
实验与结果:
-
在ImageNet分类、COCO目标检测和ADE20k语义分割任务中,Vim的表现优于DeiT,并且在处理高分辨率图像时,Vim的计算速度更快,GPU内存使用更少。
-
例如,在处理分辨率为1248×1248的图像时,Vim比DeiT快2.8倍,并节省了86.8%的GPU内存。
结论与未来工作:
-
Vim展示了作为下一代视觉基础模型骨干的潜力,尤其在高分辨率图像处理和多模态应用中表现出色。
-
未来的工作将探索Vim在无监督学习和多模态任务中的应用。
总结:
-
本文提出了一种基于双向SSM的高效视觉骨干网络Vim,通过实验验证了其在多个视觉任务中的优越性能和效率。
这里是自己的论文阅读记录,感兴趣的话可以参考一下,如果需要阅读原文的话可以看这里,如下所示:
官方项目在这里,如下所示:
摘要
最近,具有高效硬件感知设计的状态空间模型(SSM),即Mamba深度学习模型,在长序列建模方面展现了巨大的潜力。同时,构建纯粹基于SSM的高效且通用的视觉骨干网络是一个吸引人的方向。然而,由于视觉数据对位置的敏感性以及视觉理解对全局上下文的需求,SSM在表示视觉数据方面面临挑战。在本文中,我们证明了视觉表示学习对自注意力机制的依赖并非必要,并提出了一种新的通用视觉骨干网络,采用双向Mamba块(Vim),通过位置嵌入标记图像序列,并使用双向状态空间模型压缩视觉表示。在ImageNet分类、COCO目标检测和ADE20k语义分割任务中,Vim相较于DeiT等成熟的视觉Transformer模型表现更优,同时在计算和内存效率方面显著提升。例如,在处理分辨率为1248×1248的图像进行批量推理以提取特征时,Vim比DeiT快2.8倍,并节省了86.8%的GPU内存。结果表明,Vim能够克服在高分辨率图像上进行Transformer风格理解的计算和内存限制,并具有成为下一代视觉基础模型骨干的巨大潜力。
图1:DeiT(Touvron等,2021a)与我们的Vim模型的性能和效率比较。结果显示,Vim在ImageNet分类和下游检测与分割任务中均优于DeiT,并且在处理高分辨率图像时,Vim在计算和内存效率方面比DeiT更高效。例如,在处理分辨率为1248×1248的图像进行批量推理以提取特征时,Vim比DeiT快2.8倍,并节省了86.8%的GPU内存,即每张图像6084个令牌
1 引言
最近的研究进展引发了人们对状态空间模型(SSM)的浓厚兴趣。起源于经典的卡尔曼滤波模型[19],现代SSM在捕捉长程依赖性方面表现出色,并受益于并行训练。一些基于SSM的方法,如线性状态空间层(LSSL)[7]、结构化状态空间序列模型(S4)[7]、对角状态空间(DSS)[6]和S4D[7],被提出用于处理跨多种任务和模态的序列数据,特别是在建模长程依赖性方面表现出色。它们在处理长序列时效率高,因为卷积计算和近线性计算。2-D SSM(Baron等,2023)、SGConvNeXt(Li等,)和ConvSSM(Smith等,)将SSM与CNN或Transformer架构结合,以处理2-D数据。最近的工作,Mamba(Gu & Dao,2023)将时变参数引入SSM,并提出了一种硬件感知算法,以实现非常高效的训练和推理。Mamba的卓越扩展性能表明,它是语言建模中Transformer的有力替代方案。然而,尚未探索用于处理视觉数据(如图像和视频)的通用纯SSM骨干网络。
视觉Transformer(ViTs)在视觉表示学习中取得了巨大成功,在大规模自监督预训练和下游任务中表现出色。与卷积神经网络相比,ViT的核心优势在于它能够通过自注意力为每个图像块提供数据/块依赖的全局上下文。这与卷积网络使用相同参数(即卷积滤波器)处理所有位置不同。另一个优势是通过将图像视为补丁序列,无需2D归纳偏差进行模态无关建模,这使其成为多模态应用的首选架构(Bavishi等,2023;Li等,2023;Liu等,2023)。同时,Transformer中的自注意力机制在处理长程视觉依赖性时,如处理高分辨率图像,在速度和内存使用方面提出了挑战。
受Mamba在语言建模中成功的启发,我们希望将这种成功从语言领域转移到视觉领域,即设计一种通用且高效的视觉骨干网络,采用先进的SSM方法。然而,Mamba面临两个挑战,即单向建模和缺乏位置感知。为了解决这些挑战,我们提出了Vision Mamba(Vim)模型,该模型结合了双向SSM进行数据依赖的全局视觉上下文建模,并通过位置嵌入实现位置感知的视觉识别。我们首先将输入图像分割成补丁,并将其线性投影为向量输入到Vim中。图像补丁在Vim块中被视为序列数据,通过提出的双向选择性状态空间高效压缩视觉表示。此外,Vim块中的位置嵌入提供了空间信息的感知,使Vim在密集预测任务中更加稳健。在当前阶段,我们在ImageNet数据集上使用监督图像分类任务训练Vim模型,然后使用预训练的Vim作为骨干网络,进行下游密集预测任务(即语义分割、目标检测和实例分割)的顺序视觉表示学习。与Transformer一样,Vim可以在大规模无监督视觉数据上进行预训练,以获得更好的视觉表示。得益于Mamba的更高效率,Vim的大规模预训练可以在更低的计算成本下实现。
与其他基于SSM的视觉任务模型相比,Vim是一种纯SSM方法,并以序列方式建模图像,这使其更有望成为通用且高效的骨干网络。得益于双向压缩建模和位置感知,Vim是第一个处理密集预测任务的纯SSM模型。与最具说服力的Transformer模型(即DeiT)相比,Vim在ImageNet分类上表现更优。此外,Vim在高分辨率图像的GPU内存和推理时间方面更加高效。在内存和速度方面的效率使Vim能够直接进行顺序视觉表示学习,而无需依赖2D先验(如ViTDet中用于高分辨率视觉理解任务的2D局部窗口),同时实现比DeiT更高的准确性。
我们的主要贡献可以总结如下:
-
我们提出了Vision Mamba(Vim),它结合了双向SSM进行数据依赖的全局视觉上下文建模,并通过位置嵌入实现位置感知的视觉理解。
-
无需注意力机制,所提出的Vim具有与ViT相同的建模能力,同时仅具有次二次时间计算和线性内存复杂度。具体来说,在处理分辨率为1248×1248的图像进行批量推理以提取特征时,Vim比DeiT快2.8倍,并节省了86.8%的GPU内存。
-
我们在ImageNet分类和密集预测下游任务上进行了广泛的实验。结果表明,Vim相较于成熟的、高度优化的普通视觉Transformer(即DeiT)表现更优。
2 相关工作
通用视觉骨干网络架构。 在早期,卷积网络(ConvNet)LeCun等(1998)是计算机视觉的实际标准网络设计。许多卷积神经网络架构(Krizhevsky等,2012;Szegedy等,2015;Simonyan & Zisserman,2014;He等,2016;Tan & Le,2019;Wang等,;Huang等,2017;Xie等,2017;Tan & Le,2021;Radosavovic等,2020)被提出作为各种视觉应用的视觉骨干网络。开创性的工作,视觉Transformer(ViT)Dosovitskiy等(2020)改变了这一格局。它将图像视为扁平化的2D补丁序列,并直接应用纯Transformer架构。ViT在图像分类上的惊人结果及其扩展能力鼓励了许多后续工作(Touvron等,;Tolstikhin等,2021;Touvron等,2022;Fang等,2022)。一些工作专注于通过将2D卷积先验引入ViT来设计混合架构(Wu等,2021;Dai等,2021;d'Ascoli等,2021;Dong等,2022)。PVT(Wang等,2021)提出了金字塔结构的Transformer。Swin Transformer(Liu等,2021)在移位窗口内应用自注意力。另一类工作专注于通过更高级的设置改进传统的2D卷积网络(Wang等,;Liu等,)。ConvNeXt(Liu等,)回顾了设计空间,并提出了纯卷积网络,可以像ViT及其变体一样扩展。RepLKNet(Ding等,2022)提出扩大现有卷积网络的内核大小以带来改进。
尽管这些主流的后续工作通过引入2D先验在ImageNet(Deng等,2009)和各种下游任务(Lin等,2014;Zhou等,2019)上展示了卓越的性能和更好的效率,但随着大规模视觉预训练(Bao等,2022;Fang等,2023;Caron等,2021)和多模态应用(Radford等,2021;Li等,;Lu等,2023;Bavishi等,2023;Jia等,2021)的兴起,原始的Transformer风格模型重新回到了计算机视觉的中心舞台。更大的建模能力、统一的多模态表示、对自监督学习的友好性等优势使其成为首选架构。然而,由于Transformer的二次复杂度,视觉令牌的数量受到限制。有许多工作(Choromanski等,2021;Wang等,;Kitaev等,2020;Child等,2019;Ding等,2023;Qin等,2023;Sun等,2023)试图解决这一长期存在的突出挑战,但很少有工作专注于视觉应用。最近,LongViT(Wang等,)构建了一种高效的Transformer架构,用于计算病理学应用,通过扩张注意力实现。LongViT的线性计算复杂度使其能够编码极长的视觉序列。在本工作中,我们受到Mamba(Gu & Dao,2023)的启发,探索构建一种纯SSM模型作为通用视觉骨干网络,而不使用注意力,同时保留ViT的顺序、模态无关建模优势。
用于长序列建模的状态空间模型。 Gu等()提出了结构化状态空间序列(S4)模型,这是一种CNN或Transformer的替代方案,用于建模长程依赖性。线性缩放序列长度的有前途的特性吸引了进一步的探索。(Wang等,2022)提出了双向门控SSM,以在没有注意力的情况下复制BERT(Devlin等,2018)的结果。(Smith等,)通过引入MIMO SSM和高效的并行扫描到S4层,提出了新的S5层。(Fu等,2023)设计了一种新的SSM层H3,几乎填补了SSM和Transformer注意力在语言建模中的性能差距。(Mehta等,2023)通过引入更多门控单元在S4上构建了门控状态空间层,以提高表达能力。最近,(Gu & Dao,2023)提出了数据依赖的SSM层,并构建了一种通用语言模型骨干Mamba,在各种规模的大规模真实数据上表现优于Transformer,并享受线性缩放序列长度。在本工作中,我们探索将Mamba的成功转移到视觉领域,即构建一种纯粹基于SSM的通用视觉骨干网络,而不使用注意力。
用于视觉应用的状态空间模型。 (Islam & Bertasius,2022)使用1D S4处理视频分类中的长程时间依赖性。(Nguyen等,2022)进一步扩展1D S4以处理包括2D图像和3D视频的多维数据。(Islam等,2023)结合S4和自注意力的优势,构建了TranS4mer模型,在电影场景检测中实现了最先进的性能。(Wang等,)引入了一种新的选择性机制到S4,大大提高了S4在长视频理解中的性能,同时内存占用更低。(Yan等,2023)用更具扩展性的SSM骨干取代注意力机制,生成高分辨率图像并在可接受的计算下处理细粒度表示。(Ma等,2024)提出了U-Mamba,一种混合CNN-SSM架构,用于处理生物医学图像分割中的长程依赖性。上述工作(Xing等,2024;Ma等,2024;Yan等,2023;Wang等,;Islam等,2023;Nguyen等,2022;Islam & Bertasius,2022)要么将SSM应用于特定的视觉应用,要么通过将SSM与卷积或注意力结合构建混合架构。与它们不同,我们构建了一种纯SSM模型,可以作为通用视觉骨干网络。值得注意的是,VMamba(Liu等,2024),与我们方法同时进行的工作,通过将Mamba与多向扫描和分层网络架构结合,在视觉识别中展示了令人印象深刻的结果。相比之下,Vim主要集中在视觉序列学习上,并为多模态数据提供统一的表示。
3 方法
Vision Mamba(Vim)的目标是将先进的状态空间模型(SSM),即Mamba(Gu & Dao,2023),引入计算机视觉。本节首先描述SSM的预备知识,然后概述Vim。接下来,我们详细说明Vim块如何处理输入令牌序列,并继续说明Vim的架构细节。最后,我们对所提出的Vim的效率进行分析。
预备知识
最后,模型通过全局卷积计算输出。
Vim块
原始的Mamba块设计用于1-D序列,不适用于需要空间感知理解的视觉任务。在本节中,我们介绍Vim块,该块结合了双向序列建模以处理视觉任务。Vim块如图2所示。
架构细节
总之,我们架构的超参数如下:L表示块的数量,D表示隐藏状态维度,E表示扩展状态维度,M表示SSM维度。遵循ViT(Dosovitskiy等,2020)和DeiT(Touvron等,),我们首先使用16×16内核大小的投影层来获得1-D非重叠补丁嵌入。随后,我们直接堆叠L个Vim块。默认情况下,我们将块数L设置为24,SSM维度N设置为16。为了与DeiT系列的模型大小对齐,我们将隐藏状态维度D设置为192,扩展状态维度E设置为384,用于微小尺寸变体。对于小型尺寸变体,我们将D设置为384,E设置为768。
效率分析
传统的基于SSM的方法利用快速傅里叶变换来加速卷积操作,如公式(3)所示。对于数据依赖的方法,如Mamba,算法1中的第11行的SSM操作不再等同于卷积。为了解决这个问题,Mamba和所提出的Vim选择了一种现代硬件友好的方式来确保效率。该优化的关键思想是避免现代硬件加速器(GPU)的IO绑定和内存绑定。
内存效率。 为了避免内存不足问题并在处理长序列时实现更低的内存使用,Vim选择了与Mamba相同的重新计算方法。对于大小为(B,M,E,N)的中间状态以计算梯度,Vim在网络反向传播时重新计算它们。对于激活函数的输出和卷积等中间激活,Vim也重新计算它们以优化GPU内存需求,因为激活值占用大量内存但重新计算速度快。
图2:所提出的Vim模型的概述。我们首先将输入图像分割成补丁,然后将它们投影为补丁令牌。最后,我们将令牌序列发送到所提出的Vim编码器。为了进行ImageNet分类,我们将一个额外的可学习分类令牌连接到补丁令牌序列中。与用于文本序列建模的Mamba不同,Vim编码器以正向和反向两个方向处理令牌序列
其中自注意力是序列长度M的二次方,而SSM是序列长度MM的线性方(N是固定参数,默认设置为16)。计算效率使Vim能够扩展到具有大序列长度的千兆像素应用。
4 实验
图像分类设置
表1:在ImageNet1K验证集上与不同骨干网络的比较。†表示该模型使用我们的长序列设置进行了微调
我们在ImageNet-1K数据集[9]上对Vim进行了基准测试,该数据集包含1.28M训练图像和50K验证图像,分为1,000个类别。所有模型都在训练集上进行训练,并在验证集上报告top-1准确率。为了公平比较,我们的训练设置主要遵循DeiT[24]。具体来说,我们应用随机裁剪、随机水平翻转、标签平滑正则化、mixup和随机擦除作为数据增强。在训练输入图像时,我们使用AdamW[16],动量为0.9,总批量大小为1024,权重衰减为0.05来优化模型。我们使用余弦调度、初始学习率和EMA训练Vim模型300个epoch。在测试时,我们对验证集应用中心裁剪以裁剪出图像。实验在8个A800 GPU上进行。
长序列微调 为了充分利用Vim高效的长序列建模能力,我们在预训练后继续使用长序列设置对Vim进行30个epoch的微调。具体来说,我们将补丁提取步幅设置为8,同时保持补丁大小不变,恒定学习率为,权重衰减为。
结果 表1比较了Vim与基于卷积网络、Transformer和SSM的骨干网络。与基于卷积网络的ResNet[7]相比,Vim展示了优越的性能。例如,当参数大致相同时,Vim-Small的top-1准确率达到80.3,比ResNet50高出4.1个百分点。与传统的基于自注意力的ViT[11]相比,Vim在参数数量和分类准确率方面都有显著优势。与高度优化的ViT变体DeiT[24]相比,Vim在不同规模上表现更优,参数数量相当:Vim-Tiny比DeiT-Tiny高出3.9个百分点,Vim-Small比DeiT-Small高出0.5个百分点,Vim-Base比DeiT-Base高出0.1个百分点。与基于SSM的S4ND-ViT-B[20]相比,Vim在相似的top-1准确率下参数减少了3倍。经过长序列微调后,、和都取得了更高的结果。其中,甚至取得了与DeiT-B相似的结果。结果表明,Vim可以轻松适应更长的序列建模,并提取更强的视觉表示。
图1(b)和(c)比较了微小尺寸Vim和DeiT的FPS和GPU内存。随着图像分辨率的增加,Vim在速度和内存方面表现出更好的效率。具体来说,当图像大小为512×512时,Vim的FPS和内存与DeiT相似。当图像大小增加到1248×1248时,Vim比DeiT快2.8倍,并节省了86.8%的GPU内存。Vim在序列长度上线性缩放的显著优势使其为高分辨率下游视觉应用和长序列多模态应用做好了准备。
语义分割
设置 我们在ADE20K[26]上进行了语义分割实验,并使用UperNet[22]作为分割框架。我们在附录B中提供了详细设置。结果 如表2所示,Vim在不同规模上始终优于DeiT:Vim-Ti比DeiT-Ti高出1.8 mIoU,Vim-S比DeiT-S高出0.9 mIoU。与ResNet-101骨干网络相比,我们的Vim-S以近2倍的参数数量实现了相同的分割性能。
表3:使用Cascade Mask R-CNN(Cai & Vasconcelos, 2019)框架在COCO验证集上的目标检测和实例分割结果
图3:在常用的下游框架中,DeiT-Ti(Touvron等,2021a)与我们的Vim-Ti的FPS比较。我们进行批量推理,并在包含骨干网络和FPN的架构上基准测试对数缩放的FPS。在较小的分辨率(即512×512)下,Vim的性能与DeiT相当。随着输入图像分辨率的增加,Vim的FPS更高
为了进一步评估下游任务(即分割、检测和实例分割)的效率,我们将骨干网络与常用的特征金字塔网络(FPN)模块结合,并基准测试其FPS和GPU内存。如图3和图4所示,效率曲线展示了与纯骨干网络(图1)相似的比较结果,尽管我们在骨干网络上附加了沉重的FPN。卓越的线性缩放性能归功于我们提出的高效骨干网络Vim,它为基础学习千兆像素级别的视觉表示提供了基础,无需多阶段编码(例如,航空图像、医学图像和计算病理学)。
目标检测和实例分割
设置 我们在COCO 2017数据集[27]上进行了目标检测和实例分割实验,并使用ViTDet[28]作为基本框架。我们在附录B中提供了详细设置。
结果 表3比较了使用Cascade Mask R-CNN框架[29]的Vim-Ti和DeiT-Ti。Vim-Ti在box AP和mask AP上分别比DeiT-Ti高出1.3和1.1。对于中型和大型对象,Vim-Ti在和上分别比DeiT-Ti高出1.6和1.4,展示了比DeiT更好的长程上下文学习(图5)。
我们强调,准确性的优势并非微不足道,因为DeiT配备了窗口注意力,而Vim以纯序列建模方式工作。具体来说,为了在高分辨率图像(即1024×1024)上进行表示学习,我们遵循ViTDet[30]并使用2D窗口注意力修改DeiT骨干网络,这注入了2D先验并打破了Transformer的顺序建模特性。得益于第3.5节中说明的效率,图1和图4,我们可以直接将Vim应用于1024×10241024×1024输入图像,并学习目标检测和实例分割的顺序视觉表示,而无需在骨干网络中使用2D先验。
图4:在常用的下游框架中,DeiT-Ti(Touvron等,2021a)与我们的Vim-Ti的GPU内存效率比较。我们进行批量推理,并在包含骨干网络和FPN的架构上基准测试GPU内存。在较小的分辨率(即512×512)下,Vim所需的GPU内存与DeiT相当。随着输入图像分辨率的增加,Vim将显著减少GPU内存的使用
消融研究
双向SSM 我们消融了Vim的关键双向设计,使用ImageNet-1K分类和ADE20K上的Segmenter[31]语义分割框架。为了充分评估ImageNet上的学习表示能力,我们使用仅2层的简单Segmenter头在语义分割上进行迁移学习。我们研究了以下双向策略。None:我们直接采用Mamba块以仅正向处理视觉序列。Bidirectional Sequence:在训练期间,我们随机翻转视觉序列。这类似于数据增强。Bidirectional Block:我们配对堆叠块。每对的第一个块以正向处理视觉序列,每对的第二个块以反向处理视觉序列。Bidirectional SSM:我们为每个块添加额外的SSM以反向处理视觉序列。Bidirectional SSM + Conv1d:在Bidirectional SSM的基础上,我们在反向SSM之前添加反向Conv1d(图2)。
表4:双向设计的消融研究。为了确保公平比较,我们在每个实验中不使用类别令牌。Vim的默认设置用蓝色标记
如表4所示,直接采用Mamba块在分类中取得了良好的性能。然而,不自然的单向方式在下游密集预测中提出了挑战。具体来说,使用Bidirectional Block的初步双向策略在分类上取得了7个百分点的top-1准确率下降。然而,它在语义分割上比普通的单向Mamba块高出1.3 mIoU。通过添加额外的反向SSM和Conv1d,我们实现了更高的分类准确率(73.9 top-1 acc vs. 73.2 top-1 acc)和卓越的分割优势(35.9 mIoU vs. 32.3 mIoU)。我们将Bidirectional SSM + Conv1d策略作为Vim块的默认设置。
表5:分类设计的消融研究。Vim的默认设置用蓝色标记
分类设计 我们消融了Vim的分类设计,在ImageNet-1K分类上进行基准测试。我们研究了以下分类策略。Mean pool:我们对最后一个Vim块的输出特征进行均值池化,并在此池化特征上进行分类。Max pool:我们首先在视觉序列的每个令牌上适应分类头,然后对序列进行最大池化以获得分类预测结果。Head class token:遵循DeiT[25],我们在视觉序列的头部连接类别令牌并进行分类。Double class token:在head class token策略的基础上,我们在视觉序列的尾部添加额外的类别令牌。Middle class token:我们在视觉序列的中间添加类别令牌,然后对最终的中间类别令牌进行分类。
如表5所示,实验表明,middle class token策略可以充分利用SSM的递归特性和ImageNet中的中心对象先验,展示了最佳的top-1准确率76.1。
5 结论和未来工作
我们提出了Vision Mamba(Vim),探索将最新的高效状态空间模型Mamba作为通用视觉骨干网络。与之前使用混合架构或等效全局2D卷积核的视觉任务状态空间模型不同,Vim以序列建模方式学习视觉表示,不引入图像特定的归纳偏差。得益于所提出的双向状态空间建模,Vim实现了数据依赖的全局视觉上下文,并具有与Transformer相同的建模能力,同时计算复杂度更低。得益于Mamba的硬件感知设计,Vim在处理高分辨率图像时的推理速度和内存使用显著优于ViT。标准计算机视觉基准上的实验结果验证了Vim的建模能力和高效率,表明Vim具有成为下一代视觉骨干网络的巨大潜力。
在未来的工作中,结合位置嵌入的双向SSM建模的Vim适用于无监督任务,如掩码图像建模预训练,以及与Mamba类似的架构,使其适用于CLIP风格的多模态预训练。基于预训练的Vim权重,探索Vim在分析高分辨率医学图像、遥感图像和长视频等下游任务中的有用性非常直接。
影响声明
我们提高了通用视觉骨干网络的效率。任何与提高效率相关的工作通常带来的社会后果或影响也适用于此处,因为此类工作必然提高了高分辨率输入图像的视觉骨干网络的实用性。
致谢
本工作部分由国家科技重大专项资助,项目编号2023YFF0905400,以及国家自然科学基金资助,项目编号62276108。
我们感谢程天恒、方宇鑫、杨书生、江波和姚景峰对初稿的有益反馈。
附录A 可视化
附加设置
语义分割设置 我们在ADE20K Zhou等(2019)数据集上进行了语义分割实验。ADE20K包含150个细粒度语义类别,训练集、验证集和测试集分别包含20K、2K和3K张图像。我们选择UperNet Xiao等()作为基础框架。在训练中,我们使用AdamW,权重衰减为0.01,总批量大小为16来优化模型。采用的训练计划使用初始学习率为,线性学习率衰减,线性预热1,500次迭代,总训练160K次迭代。数据增强遵循常见设置,包括随机水平翻转、在比例范围[0.5,2.0]内随机重新缩放和随机光度失真。在评估时,我们将图像缩放为短边512。
目标检测和实例分割设置 我们在COCO 2017数据集Lin等(2014)上进行了目标检测和实例分割实验。COCO 2017数据集包含118K张训练图像、5K张验证图像和20K张测试图像。我们使用Cascade Mask R-CNN Cai & Vasconcelos(2019)作为基础框架。对于基于ViT的骨干网络,我们应用额外的配置(例如,交错窗口和全局注意力)以处理高分辨率图像,遵循ViTDet Li等()。对于基于SSM的Vim,我们直接使用它,无需任何修改。其他训练和评估设置相同。在训练中,我们使用AdamW,权重衰减为0.1,总批量大小为64来优化模型。采用的训练计划使用初始学习率为,线性学习率衰减,总训练380K次迭代。数据增强使用大规模抖动数据增强Ghiasi等(2021)到1024×10241024×1024输入图像。在评估时,我们将图像缩放为短边1024。
分层架构的扩展比较
为了进一步与分层架构进行比较,我们提出了另一种变体Hier-Vim,通过将SwinTransformer中的移位局部窗口注意力替换为所提出的全局双向SSM。我们在表6中详细说明了配置。
ImageNet上的分类 遵循标准训练和验证协议Liu等(2021;2024),我们在表7中比较了Hier-Vim与流行的分层架构在微小、小型和基础模型大小上的表现。结果表明,Hier-Vim在微小尺寸上比Swin Transformer高出1.3%,在小型尺寸上高出0.2%,在基础尺寸上高出0.4%,展示了与成熟且高度优化的现代分层架构的竞争力。
表6:不同Hier-Vim变体的详细配置。我们提供了4个阶段的通道数和块数。
表7:与分层架构在ImageNet-1K验证集上的比较。
图5:使用Cascade Mask R-CNN Cai & Vasconcelos(2019)框架,DeiT-Ti Touvron等()和我们的Vim-Ti的可视化比较。得益于SSM的长程上下文学习,我们可以捕捉图像中的非常大的对象,而DeiT-Ti无法感知。