视觉Mamba:基于双向状态空间模型的高效视觉表征学习

举报
AI浩 发表于 2024/08/27 19:30:23 2024/08/27
【摘要】 摘要https://arxiv.org/pdf/2401.09417v1.pdf最近,具有高效硬件感知设计的状态空间模型(SSMs),例如Mamba,在长序列建模方面展现出了巨大潜力。纯粹基于SSMs构建高效和通用的视觉骨干网络是一个吸引人的方向。然而,由于视觉数据的空间敏感性和视觉理解的全局上下文需求,用SSMs表示视觉数据是一项挑战。本文表明,视觉表示学习对自注意力的依赖不是必需的,并...

摘要

https://arxiv.org/pdf/2401.09417v1.pdf
最近,具有高效硬件感知设计的状态空间模型(SSMs),例如Mamba,在长序列建模方面展现出了巨大潜力。纯粹基于SSMs构建高效和通用的视觉骨干网络是一个吸引人的方向。然而,由于视觉数据的空间敏感性和视觉理解的全局上下文需求,用SSMs表示视觉数据是一项挑战。本文表明,视觉表示学习对自注意力的依赖不是必需的,并提出了一个新的通用视觉骨干网络,该网络使用双向Mamba块(Vim),通过位置嵌入标记图像序列,并使用双向状态空间模型压缩视觉表示。在ImageNet分类、COCO目标检测和ADE20k语义分割任务上,Vim与DeiT等完善的视觉变换器相比性能更高,同时还显著提高了计算和内存效率。例如,Vim比DeiT快2.8倍,在执行批量推理以提取分辨率为1248×1248的图像特征时节省了86.8%的GPU内存。结果表明,Vim能够克服在执行高分辨率图像的Transformer风格理解时的计算和内存限制,并且有潜力成为下一代视觉基础模型的骨干网络。
在这里插入图片描述

1. 引言

最近的研究进展引发了对状态空间模型(SSM)的极大兴趣。源自经典的状态空间模型[30],现代的SSM在捕捉长距离依赖方面表现出色,并受益于并行训练。一些基于SSM的方法,如线性状态空间层(LSSL)[22]、结构化状态空间序列模型(S4)[21]、对角线状态空间(DSS)[24]和S4D[23],被提出用于处理各种任务和模式的序列数据,特别是在建模长距离依赖方面。由于卷积计算和近线性计算,它们在处理长序列时效率很高。2-D SSM [2]、SGConvNeXt [37]和ConvSSM [52]将SSM与CNN或Transformer架构相结合,用于处理2-D数据。最近的工作,Mamba [20],将时间变化参数纳入SSM,并提出了一个硬件感知算法,以实现高效的训练和推理。Mamba的优良扩展性能表明,它是语言建模中Transformer的有前途的替代品。然而,针对视觉任务的纯SSM基础骨干尚未被探索。

视觉Transformer(ViTs)在视觉表示学习方面取得了巨大成功,在大规模自监督预训练和下游任务的高性能方面表现出色。与卷积神经网络相比,其核心优势在于ViT可以通过自注意力为每个图像块提供与数据/补丁相关的全局上下文。这不同于使用相同参数(即卷积滤波器)的所有位置的卷积网络。另一个优势是模态无关的建模,将图像视为一系列没有2D归纳偏置的补丁,使其成为多模态应用的优选架构[3, 36, 40]。同时,Transformer中的自注意力机制在处理长距离视觉依赖关系时,例如处理高分辨率图像时,在速度和内存使用方面带来了挑战。

受到Mamba在语言建模方面取得成功的启发,我们可以将这一成功从语言转移到视觉,即设计一个使用先进SSM方法的通用且高效的视觉骨干网络,这非常吸引人。然而,Mamba面临两个挑战:单向建模和缺乏位置意识。为了解决这些挑战,我们提出了Vision Mamba(Vim)块,该块结合了双向SSM来进行数据依赖的全局视觉上下文建模,并使用位置嵌入进行位置感知的视觉识别。我们首先将输入图像分割成补丁,并将其线性投影为向量传递给Vim。在Vim块中,图像补丁被视为序列数据,通过提出的双向选择性状态空间有效地压缩视觉表示。此外,Vim块中的位置嵌入提供了对空间信息的感知,使Vim在密集预测任务中更加鲁棒。在目前阶段,我们在ImageNet数据集上使用监督图像分类任务训练Vim模型,然后使用预训练的Vim作为骨干来执行下游密集预测任务的视觉表示学习,例如语义分割、目标检测和实例分割。与Transformers一样,Vim可以在大规模无监督视觉数据上进行预训练以获得更好的视觉表示。由于Mamba的效率更高,Vim的大规模预训练可以以更低的计算成本实现。

与针对视觉任务的其他基于SSM的模型相比,Vim是一种纯SSM方法,以序列方式对图像进行建模,这对于通用且高效的骨干网络来说更有前途。由于具有位置意识的双向压缩建模,Vim是第一个用于处理密集预测任务的纯SSM基础模型。与最令人信服的基于Transformer的模型(即DeiT [60])相比,Vim在ImageNet分类方面取得了优越的性能。此外,Vim在GPU内存和推理时间方面更加高效,对于高分辨率图像而言,这使得Vim能够直接进行顺序视觉表示学习,而无需依赖于2D先验(例如ViTDet [38]中的2D局部窗口)来完成高分辨率视觉理解任务,同时实现了比DeiT更高的准确性。

我们的主要贡献可以总结如下:

• 我们提出了Vision Mamba(Vim),该模型结合了双向SSM进行数据依赖的全局视觉上下文建模,并使用位置嵌入进行位置感知的视觉理解。

• 无需注意力机制,提出的Vim具有与ViT相同的建模能力,同时具有亚二次时间复杂度和线性内存复杂度。具体来说,我们的Vim比DeiT快2.8倍,在执行批量推理以提取分辨率为1248×1248的图像特征时,可节省86.8%的GPU内存。

• 我们对ImageNet分类和下游密集预测任务进行了大量实验。结果表明,与完善的优化普通视觉Transformer(即DeiT)相比,Vim实现了优越的性能。

• 由于Mamba的高效硬件感知设计,Vim在处理高分辨率计算机视觉任务(例如视频分割、航空图像分析、医学图像分割和计算病理学)时比基于自注意力的DeiT更加高效。

2、相关研究

通用视觉骨干网络架构。在早期阶段,ConvNet [34]成为计算机视觉的默认标准网络设计。许多卷积神经架构[25, 26, 33, 50, 51, 56–58, 63, 72]已被提出作为各种视觉应用的视觉骨干网络。具有开创性的工作,Vision Transformer(ViT) [14]改变了这一领域。它将图像视为一系列平坦的2D补丁,并直接应用纯Transformer架构。ViT在图像分类及其扩展能力方面的出色结果激发了许多后续工作 [16, 59, 61, 62]。一方面的工作专注于通过将2D卷积先验引入ViT来设计混合架构 [9, 13, 15, 69]。PVT [66]提出了一个金字塔结构Transformer。Swin Transformer [42]在移位窗口内应用自注意力。另一方面的研究工作侧重于通过更先进的设置改进传统的2D ConvNets [41, 67]。ConvNeXt [43]回顾了设计空间并提出了纯ConvNets,这些ConvNets可以像ViT及其变种一样可扩展。RepLKNet [12]提出了将现有ConvNets的核大小扩大,以带来改进。

尽管这些主导的后续工作通过引入2D先验在ImageNet [10]和各种下游任务[39, 74]上展示了优越的性能和更高的效率,但随着大规模视觉预训练[1, 5, 17]和多模态应用[3, 29, 35, 36, 40, 49]的兴起,原始的Transformer风格模型重新回到了计算机视觉的中心舞台。较大的建模能力、统一的多元表示、对自监督学习的友好性等的优势,使其成为首选架构。然而,由于Transformer的二次复杂性,视觉令牌的数量有限。有许多工作[7, 8, 11, 32, 48, 55, 65]来解决这一长期存在的突出挑战,但其中很少关注视觉应用。最近,LongViT [68]通过膨胀注意力为计算病理学应用构建了一个高效的Transformer架构。LongViT的线性计算复杂性允许它编码极长的视觉序列。在这项工作中,我们受到Mamba [20]的启发,探索构建一个基于纯SSM的模型作为通用视觉骨干网,而无需使用注意力,同时保留了ViT的顺序、模态无关建模优点。

对于长序列建模的状态空间模型。 [21] 提出了一种结构化状态空间序列(S4)模型,这是CNN或Transformer之外的一种新颖替代方案,用于建模长距离依赖关系。线性序列长度扩展的有前途的特性吸引了进一步的探索。 [53] 通过引入MIMO SSM和高效的并行扫描,提出了一种新的S5层。 [18] 设计了一个新的SSM层H3,几乎填补了SSM和Transformer注意力在语言建模方面的性能差距。 [46] 通过引入更多的门控单元来提高表达性,在S4上构建了门控状态空间层。最近,[20] 提出了一种数据依赖的SSM层,并构建了一个通用的语言模型骨干网Mamba,它在大型真实数据上对各种规模的Transformer具有更好的性能,并享有线性序列长度扩展。在本工作中,我们探索将Mamba的成功转移到视觉领域,即构建一个基于SSM的通用视觉骨干网,而无需使用注意力。

视觉应用中的状态空间模型。 [27] 使用1D S4处理视频分类中的长时序依赖关系。 [47] 将1D S4进一步扩展到处理包括2D图像和3D视频在内的多维数据。 [28] 结合S4和自注意力的优点构建了TranS4mer模型,在电影场景检测方面达到了最先进的性能。 [64] 向S4引入了一种新颖的选择性机制,在长格式视频理解方面大幅提高了S4的性能,同时大大降低了内存占用。 [73] 用更可扩展的基于SSM的骨干网取代了注意力机制,以生成高分辨率图像并在可负担的计算下处理精细的表示。 [45] 提出了U-Mamba,这是一种混合CNN-SSM架构,用于处理生物医学图像分割中的长时序依赖关系。上述工作要么将SSM应用于特定的视觉应用,要么通过将SSM与卷积或注意力相结合来构建混合架构。与它们不同,我们构建了一个基于纯SSM的模型,它可以作为一个通用的视觉骨干网。

3、方法

Vision Mamba(Vim)的目标是将先进的SSM(即Mamba [20])引入计算机视觉。本节首先介绍了SSM的预备知识。随后概述了Vim。接着详细介绍了Vim块如何处理输入令牌序列,并继续说明了Vim的架构细节。最后对所提Vim的效率进行了分析。

3.1、预备知识

基于SSM的模型,即结构化状态空间序列模型(S4)和Mamba,受到了连续系统的启发,该系统通过隐藏状态h(t)将1-D函数或序列x(t)映射到y(t)。这个系统使用ARN×N{A} \in \mathbb{R}^{\mathrm{N} \times \mathrm{N}}作为演化参数,使用BRN×1{B} \in \mathbb{R}^{\mathbb{N} \times 1}CR1×N\mathbf{C} \in \mathbb{R}^{1 \times \mathrm{N}}​作为投影参数。

h(t)=Ah(t)+Bx(t),y(t)=Ch(t).\begin{aligned} h^{\prime}(t) & =\mathbf{A} h(t)+\mathbf{B} x(t), \\ y(t) & =\mathbf{C} h(t) . \end{aligned}

S4和Mamba是连续系统的离散版本,包括时间尺度参数Δ\Delta,用于将连续参数A,BA,B转换为离散参数A,BA,B。常用的转换方法是零阶保持(ZOH),其定义如下:

A=exp(ΔA),B=(ΔA)1(exp(ΔA)I)ΔB.\begin{array}{l} \overline{\mathbf{A}}=\exp (\boldsymbol{\Delta} \mathbf{A}), \\ \overline{\mathbf{B}}=(\boldsymbol{\Delta} \mathbf{A})^{-1}(\exp (\boldsymbol{\Delta} \mathbf{A})-\mathbf{I}) \cdot \boldsymbol{\Delta} \mathbf{B} . \end{array}

在离散化A,BA,B之后,使用步长Δ\Delta的方程(1)的离散化形式可以重写为:

ht=Aht1+Bxt,yt=Cht.\begin{aligned} h_{t} & =\overline{\mathbf{A}} h_{t-1}+\overline{\mathbf{B}} x_{t}, \\ y_{t} & =\mathbf{C} h_{t} . \end{aligned}

最后,模型通过全局卷积计算输出。

K=(CB,CAB,,CAM1B),y=xK,\begin{aligned} \overline{\mathbf{K}} & =\left(\mathbf{C} \overline{\mathbf{B}}, \mathbf{C} \overline{\mathbf{A B}}, \ldots, \mathbf{C} \overline{\mathbf{A}}^{\mathrm{M}-1} \overline{\mathbf{B}}\right), \\ \mathbf{y} & =\mathbf{x} * \overline{\mathbf{K}}, \end{aligned}

其中M是输入序列x的长度,而KRM\overline{\mathbf{K}} \in \mathbb{R}^{M}是结构化的卷积核。

3.2、视觉Mamba

图2展示了所提出的Vim的概述。标准Mamba是为1-D序列设计的。为了处理视觉任务,我们首先将2-D图像tRH×W×Ct \in \mathbb{R}^{\mathrm{H} \times \mathrm{W} \times \mathrm{C}}转换为平坦的2-D补丁xpRJ×(P2C)x_p∈RJ×(P2⋅C)​,其中(H,W)是输入图像的大小,C是通道数,P是图像补丁的大小。接下来,我们将xpx_p​线性投影到大小为D的向量,并添加位置嵌入Epos E_{\text {pos }},如下所示:

T0=[tcls;tp1W;tp2W;;tpJW]+Epos,\mathbf{T}_{0}=\left[\mathbf{t}_{c l s} ; \mathbf{t}_{p}^{1} \mathbf{W} ; \mathbf{t}_{p}^{2} \mathbf{W} ; \cdots ; \mathbf{t}_{p}^{\mathrm{J}} \mathbf{W}\right]+\mathbf{E}_{p o s},

其中,tpj\mathbf{t}_{p}^{j} 是图像 t\mathbf{t} 的第 jj 个补丁,WR(P2C)×D\mathbf{W} \in \mathbb{R}^{\left(\mathrm{P}^{2} \cdot \mathrm{C}\right) \times \mathrm{D}} 是可学习的投影矩阵。受 ViT [14] 和 BERT [31] 的启发,我们也使用类别标记来表示整个补丁序列,记为 tcls\mathbf{t}_{c l s}。然后,我们将标记序列(T11\mathbf{T}_{1-1})发送到 Vim 编码器的第一层,并得到输出 T1\mathbf{T}_{1}。最后,我们规范化输出类别标记 TL0T_{L}^{0} 并将其输入到多层感知器(MLP)头部,以获得最终预测 p^\hat{p},如下所示:

Tl=Vim(T11)+T11,f=Norm(TL0),p^=MLP(f),\begin{aligned} \mathbf{T}_{l} & =\operatorname{Vim}\left(\mathbf{T}_{1-1}\right)+\mathbf{T}_{1-1}, \\ \mathbf{f} & =\operatorname{Norm}\left(\mathbf{T}_{\mathrm{L}}^{0}\right), \\ \hat{p} & =\operatorname{MLP}(\mathbf{f}), \end{aligned}

其中Vim是我们提出的视觉Mamba块,L是层数,Norm是归一化层。

3.3、Vim 块

原始的Mamba块是为1-D序列设计的,不适合需要空间感知理解的视觉任务。在本节中,我们介绍Vim块,它为视觉任务融入了双向序列建模。Vim块如图2所示。
在这里插入图片描述

具体来说,我们在算法1中展示了Vim块的操作。输入标记序列 T11\mathbf{T}_{1-1} 首先通过归一化层进行归一化。接下来,我们将归一化后的序列线性投影到具有维度大小 E 的 x\mathbf{x}z\mathbf{z}。然后,我们从前向和后向两个方向处理 x\mathrm{x}。对于每个方向,我们首先应用1-D卷积到 x\mathrm{x} 并得到 xo\mathrm{x}_{o}^{\prime}。然后,我们将 xo\mathbf{x}_{o}^{\prime} 线性投影到Bo\mathbf{B}_{o}, Co\mathbf{C}_{o}, Δo\boldsymbol{\Delta}_{o},分别。接着使用 Δo\boldsymbol{\Delta}_{o} 来转换 Ao,Bo\overline{\mathbf{A}}_{o}, \overline{\mathbf{B}}_{o},分别。最后,我们通过SSM计算 yforward \mathbf{y}_{\text {forward }}ybackward \mathbf{y}_{\text {backward }}yforward \mathbf{y}_{\text {forward }}ybackward \mathbf{y}_{\text {backward }} 然后由 z\mathbf{z} 进行门控并相加得到输出标记序列 T1\mathbf{T}_{1}
在这里插入图片描述

3.4、架构细节

总的来说,我们的架构的超参数如下:

L:块的数量, D:隐藏状态维度, E:扩展状态维度, N:SSM维度。

根据ViT [14]和DeiT [61],我们首先使用16×16大小的核大小投影层来获得非重叠补丁嵌入的1-D序列。接下来,我们直接堆叠L个Vim块。默认情况下,我们将块的数量L设置为24,SSM维度N设置为16。为了与DeiT系列模型的规模保持一致,我们将隐藏状态维度D设置为192,并将扩展状态维度E设置为384,用于tiny-size变体。对于small-size变体,我们将D设置为384,E设置为768。

3.5、效率分析

传统的基于SSM的方法利用快速傅里叶变换来加速卷积操作,如等式(4)所示。对于数据依赖的方法,如Mamba,算法1第11行的SSM操作不再等同于卷积。为了解决这个问题,Mamba和提出的Vim选择了一种现代硬件友好的方法来确保效率。这种优化的关键思想是避免现代硬件加速器(GPU)的IO约束和内存约束。

IO效率。高带宽内存(HBM)和SRAM是GPU的两个重要组件。其中,SRAM具有较大的带宽,HBM具有更大的内存大小。使用HBM的标准实现Vim的SSM操作需要数量为O(BMEN)的内存IO。受Mamba启发,Vim首先从慢速HBM读取O(BME+EN)字节的内存((Δo,Ao,Bo,Co)\left(\boldsymbol{\Delta}_{\mathbf{o}}, \mathbf{A}_{\mathbf{o}}, \mathbf{B}_{\mathbf{o}}, \mathbf{C}_{\mathbf{o}}\right)),然后将其存储在快速的SRAM中。接着,Vim在SRAM中获取大小为(B, M, E, N)的离散AoBo\overline{\mathbf{A}}_{\mathbf{o}},\overline{\mathbf{B}}_{\mathrm{o}}。最后,Vim在SRAM中执行SSM操作,并将大小为(B, M, E)的输出写回HBM。这种方法可以将IO从O(BMEN)减少到O(BME+EN)。

内存效率。为了避免内存溢出问题,并降低处理长序列时的内存使用量,Vim选择了与Mamba相同的重计算方法。对于大小为(B, M, E, N)的中间状态,以计算梯度,Vim在网络反向传播时对它们进行重计算。对于中间激活值,如激活函数的输出和卷积,Vim也进行重计算,以优化GPU内存需求,因为激活值需要大量内存,但重计算速度很快。

计算效率。Vim块中的SSM(算法1第11行)和Transformer中的自注意力都在提供全局上下文方面发挥着关键作用。给定一个视觉序列TR1×M×D\mathbf{T} \in R^{1 \times \mathrm{M} \times \mathrm{D}}和默认设置E=2D\mathrm{E}=2 \mathrm{D},全局自注意力和SSM的计算复杂度为:

Ω( self-attention )=4MD2+2M2D,Ω(SSM)=3M(2D)N+M(2D)N2,\begin{array}{l} \Omega(\text { self-attention })=4 \mathrm{MD}^{2}+2 \mathrm{M}^{2} \mathrm{D}, \\ \Omega(\mathrm{SSM})=3 \mathrm{M}(2 \mathrm{D}) \mathrm{N}+\mathrm{M}(2 \mathrm{D}) \mathrm{N}^{2}, \end{array}

其中,自注意力的计算复杂度与序列长度M的平方成正比,而SSM的计算复杂度与序列长度M成线性关系(N是一个固定参数,默认设置为16)。计算效率使Vim能够扩展到具有长序列长度的大型千兆像素应用。

4、实验

4.1、图像分类

设置。我们在ImageNet-1K数据集[10]上对Vim进行基准测试,该数据集包含1,280万张训练图像和5万张验证图像,共有1000个类别。所有模型都在训练集上进行训练,并在验证集上报告top-1准确率。为了公平比较,我们的训练设置主要遵循DeiT [61]。具体来说,我们应用随机裁剪、随机水平翻转、标签平滑正则化、混合和随机擦除作为数据增强。在2242224^2输入图像上进行训练时,我们使用AdamW [44]优化器,动量为0.9,总批量大小为1024,权重衰减为0.05来优化模型。我们使用余弦调度对Vim模型进行300轮训练,初始学习率为1×10(3)1×10^(-3),并使用EMA。在测试期间,我们对验证集进行中心裁剪,以裁剪出2242224^2张图像。实验在8个A800 GPU上执行。
在这里插入图片描述

结果。表1将Vim与基于卷积网络、基于Transformer和基于SSM的骨干网进行了比较。与基于卷积网络的ResNet相比,Vim展示了优越的性能。例如,当参数大致相似时,Vim-Small的top-1准确率达到了80.3,比ResNet50高出了4.1个百分点。与传统的基于自注意力的ViT相比,Vim在参数数量和分类准确率方面都大大优于它。与高度优化的ViT变体,即DeiT相比,Vim在参数数量相当的情况下,在不同尺度上都超过了它:VimTiny比DeiT-Tiny高0.9个百分点,Vim-Small比DeiT-Small高0.5个百分点。与基于SSM的S4ND-ViTB相比,Vim以少3倍的参数实现了相似的top-1准确率。

图1 (b)和©比较了小型Vim和DeiT的FPS和GPU内存。随着图像分辨率的增加,Vim在速度和内存方面表现出更好的效率。具体来说,当图像大小为512时,Vim的FPS和内存与DeiT相当。随着图像大小增加到1248,Vim比DeiT快2.8倍,并节省了86.8%的GPU内存。Vim在序列长度方面的明显优势使其适用于高分辨率下游视觉应用和长序列多模态应用。

4.2、语义分割

设置。我们在ADE20K数据集上进行了语义分割实验[74]。ADE20K包含150个精细的语义类别,训练、验证和测试图像分别为2万、2千和3千张。我们选择UperNet作为基础框架。在训练期间,我们使用AdamW优化器和权重衰减为0.01,总批量大小为16来优化模型。所采用的训练计划初始学习率为6×1056 \times 10^{-5},线性学习率衰减,线性热身1500次迭代,总训练16万次迭代。数据增强遵循常见设置,包括随机水平翻转、随机缩放(比例范围[0.5,2.0])和随机光学校正。在评估期间,我们将图像调整为较短的一侧为512。
在这里插入图片描述

结果。如表2所示,Vim在不同尺度上均优于DeiT:Vim-Ti比DeiT-Ti高出1.0 mIoU,Vim-S比DeiT-S高出0.9 mIoU。与ResNet-101骨干网相比,我们的Vim S在分割性能上几乎减少了2倍的参数数量。
在这里插入图片描述

为了进一步评估下游任务(即分割、检测和实例分割)的效率,我们将骨干网与常用的特征金字塔网络(FPN)模块相结合,并评估它们的FPS和GPU内存。如图3和图4所示,效率曲线展示了与纯骨干网(图1)相似的比较结果,尽管我们在骨干网上附加了一个沉重的FPN。出色的线性扩展性能归功于我们提出的效率型骨干网Vim,它为无需多阶段编码(例如航空图像、医学图像和计算病理学)的端到端学习千兆像素级视觉表示奠定了基础。
在这里插入图片描述

4.3、目标检测与实例分割

设置。我们在COCO 2017数据集上进行了目标检测和实例分割实验[39]。COCO 2017数据集包含11.8万张用于训练的图像、5千张用于验证的图像和2万张用于测试的图像。我们使用标准的Cascade Mask R-CNN [4]作为基础框架。对于基于ViT的骨干网络,我们应用额外的配置(例如交错的窗口和全局注意力)来处理高分辨率图像,遵循ViTDet [38]。对于基于SSM的Vim,我们无需任何修改直接使用。其他训练和评估设置保持不变。在训练期间,我们使用AdamW优化器和权重衰减为0.1,总批量大小为64来优化模型。所采用的训练计划初始学习率为1×1041 \times 10^{-4},线性学习率衰减,总训练迭代次数为38万次。数据增强使用大规模抖动数据增强[19]来处理1024x1024输入图像。在评估期间,我们将图像调整为较短的一侧为1024。
在这里插入图片描述

结果。表3比较了使用Cascade Mask R-CNN框架的Vim-Ti与DeiT-Ti [4]。Vim-Ti在box AP和mask AP上分别优于DeiT-Ti 1.3和1.1。对于中等大小和大型对象,Vim-Ti在box AP和mask AP上分别优于DeiT-Ti 1.6 APmbox / 1.3 APmmask 和1.4 AP1box / 1.8 AP1mask ,这表明Vim在长距离上下文学习方面优于DeiT(图5)。
在这里插入图片描述

我们想要强调的是,由于DeiT配备了窗口注意力,而Vim以纯序列建模方式工作,因此准确性的优越性是重要的。具体来说,为了在高清图像(即1024x1024)上进行表示学习,我们遵循ViTDet [38]的指导,并使用2D窗口注意力修改了DeiT骨干网络,从而引入了2D先验并打破了Transformer的序列建模性质。由于第3.5节、图1和图4中所述的效率,我们可以直接在1024x1024输入图像上应用Vim,并学习用于目标检测和实例分割的序列视觉表示,而无需骨干网络中的2D先验。

4.4、消融研究

我们在ImageNet-1K分类和ADE20K语义分割的Segmenter [54]上对Vim的关键双向设计进行了消融研究。为了充分评估在ImageNet上学到的表示能力,我们使用了一个只有2层的简单Segmenter头来执行语义分割的迁移学习。我们研究了以下这些双向策略:

  • 无。我们直接采用Mamba块仅沿前向方向处理视觉序列。
  • 双向序列。在训练期间,我们随机翻转视觉序列。这类似于数据增强。
  • 双向块。我们将堆叠的块进行配对。每对中的第一个块沿前向方向处理视觉序列,而每对中的第二个块沿后向方向处理。
  • 双向SSM。我们为每个块添加了一个额外的SSM,以沿后向方向处理视觉序列。
  • 双向SSM + Conv1d。基于双向SSM,我们在后向SSM之前进一步添加了一个后向Conv1d(图2)。

在这里插入图片描述

如表4所示,直接采用Mamba块在分类中取得了良好的性能。然而,不自然的单向方式在下游密集预测中提出了挑战。具体来说,使用双向块的初步策略在分类上的top-1准确率降低了7个点。然而,它在语义分割上的表现优于普通的单向Mamba块1.3 mIoU。通过添加额外的后向SSM和Conv1d,我们实现了相似的分类准确率(73.1 top-1 acc vs. 73.2 top-1 acc)和卓越的分割优势(34.8 mIoU vs. 32.3 mIoU)。我们在Vim块中默认使用双向SSM + Conv1d的策略。

5、结论和未来工作

我们提出了Vision Mamba(Vim)来探索最近的高效状态空间模型,即Mamba,作为通用的视觉骨干网。与之前用于视觉任务的混合架构或等效全局2D卷积核的状态空间模型不同,Vim以序列建模的方式学习视觉表示,并且不引入图像特定的归纳偏差。由于所提出的双向状态空间建模,Vim实现了数据相关的全局视觉上下文,并具有与Transformer相同的建模能力,同时具有较低的计算复杂性。得益于Mamba的硬件感知设计,Vim在处理高分辨率图像时的推理速度和内存使用量明显优于ViTs。在标准计算机视觉基准上的实验结果验证了Vim的建模能力和高效率,表明Vim具有成为下一代视觉骨干网的巨大潜力。

在未来的工作中,具有位置嵌入的双向SSM建模的Vim适合于无监督任务,例如掩模图像建模预训练和与Mamba类似的架构支持多模态任务,例如CLIP风格的预训练。基于预训练的Vim权重,探索Vim在分析高分辨率医学图像、遥感图像和长视频等下游任务中的用途是非常直接的。

【声明】本内容来自华为云开发者社区博主,不代表华为云及华为云开发者社区的观点和立场。转载时必须标注文章的来源(华为云社区)、文章链接、文章作者等基本信息,否则作者和本社区有权追究责任。如果您发现本社区中有涉嫌抄袭的内容,欢迎发送邮件进行举报,并提供相关证据,一经查实,本社区将立刻删除涉嫌侵权内容,举报邮箱: cloudbbs@huaweicloud.com
  • 点赞
  • 收藏
  • 关注作者

评论(0

0/1000
抱歉,系统识别当前为高风险访问,暂不支持该操作

全部回复

上滑加载中

设置昵称

在此一键设置昵称,即可参与社区互动!

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。