摘要
https://arxiv.org/pdf/2206.09959
我们提出了全局上下文视觉转换器(GC ViT),这是一种新颖的架构,旨在提高计算机视觉中的参数和计算利用率。我们的方法利用全局上下文自注意力模块与标准的局部自注意力相结合,以有效且高效的方式对长程和短程空间交互进行建模,而无需执行昂贵的操作,如计算注意力掩码或移动局部窗口。此外,我们还解决了ViT中归纳偏差缺失的问题,并提出在我们的架构中利用修改后的融合倒置残差块。我们提出的GC ViT在图像分类、目标检测和语义分割任务中均取得了最先进的结果。在ImageNet-1K分类数据集上,具有51M、90M和201M参数的GC ViT变体在
224
×
224
224 \times 224
224×224图像分辨率下,无需任何预训练,分别达到了
84.3
%
84.3\%
84.3%、
85.0
%
85.0\%
85.0%和
85.7
%
85.7\%
85.7%的Top-1准确率,因此大幅超越了同等规模的先前技术,如基于卷积神经网络的ConvNeXt和基于ViT的MaxViT和Swin Transformer。在MS COCO和ADE20K数据集上进行的目标检测、实例分割和语义分割等下游任务中,预训练的GC ViT主干网络始终优于先前的工作。具体而言,具有4尺度DINO检测头的GC ViT在MS COCO数据集上实现了
58.3
%
58.3\%
58.3%的边界框AP。代码可在https://github.com/NVlabs/GCViT获取。
1、引言
- 近年来,Transformer(Vaswani等人,2017)在自然语言处理(NLP)基准测试中取得了最先进的(SOTA)性能,并已成为各种任务中的实际模型。Transformer成功的一个关键要素是自注意力机制,它允许通过关注远程和附近的标记来捕获上下文表示(Yin等人,2021)。遵循这一趋势,Vision Transformer(ViT)(Dosovitskiy等人,2020)提出将图像块作为标记,在整体架构中使用,与原始Transformer的编码器相比略有不同。尽管卷积神经网络(CNN)在计算机视觉领域具有历史性的主导地位,但基于ViT的模型在各种计算机视觉任务中取得了SOTA或竞争性的性能。
本质上,与卷积神经网络(CNN)相比,ViT中的自注意力机制能够学习更一致的长距离和短距离信息(Raghu等人,2021)。然而,ViT的一体式架构和自注意力的二次计算复杂度阻碍了它们在高分辨率图像(Yang等人,2021a)中的快速应用,而在高分辨率图像中捕获多尺度长距离信息对于准确表示建模至关重要。
几项工作(Liu等人,2021;Dong等人,2022;Chu等人,2021a;Tu等人,2022),尤其是Swin Transformer(Liu等人,2021),通过提出多分辨率架构来解决短距离和长距离空间依赖之间的平衡问题,在该架构中,自注意力是在局部窗口内计算的。在这种范式中,使用诸如窗口移位之类的跨窗口连接来建模不同区域之间的相互作用。尽管取得了进展,但局部窗口的有限感受野挑战了自注意力捕获长距离信息的能力,而诸如移位之类的窗口连接方案仅覆盖每个窗口附近的一个小邻域。随后的一些工作,如Focal Transformer(Yang等人,2021b),试图通过设计模型复杂度更高的高度复杂的自注意力模块来解决这个问题。
在这项工作中,我们引入了全局上下文(GC)ViT网络来解决这些限制。具体来说,我们提出了一种由局部和全局自注意力模块组成的分层ViT架构。在每个阶段,我们使用一种新颖的融合倒残差块(我们称之为改进的Fused-MBConv块)来计算全局查询标记,该块涵盖了来自不同图像区域的全局上下文信息。虽然局部自注意力模块负责建模短距离信息,但全局查询标记在所有全局自注意力模块之间共享,以便与局部键和值表示进行交互。
我们提出的全局查询生成器和自注意力框架的设计直观且简单,可以使用主要的深度学习框架高效实现。因此,它消除了复杂且计算量大的操作,并确保了自注意力在应用于高分辨率图像时的有效性。此外,我们还提出了一种具有参数高效的融合MBConv层的新型下采样块,以解决ViT中归纳偏差不足的问题,并增强通道间依赖性的建模。
我们使用三个公开可用的数据集对各种计算机视觉任务进行了广泛的实验,以验证所提出的GC ViT的有效性。在使用ImageNet-1K数据集进行图像分类时,具有51M、90M和201M参数的GC ViT达到了新的SOTA基准,Top-1准确率分别为84.3%、85.0%和85.7%,且没有使用额外数据或预训练。
因此,GC ViT在性能上始终优于ConvNeXt(Liu等人,2022b)、Max ViT(Tu等人,2022)和Swin Transformer(Liu等人,2021)模型,有时优势还相当明显(见图1)。
使用具有Cascade Mask RCNN(He等人,2017)头的ImageNet-1K预训练GC ViT基础主干网络,我们的模型在MS COCO数据集上实现了 52.9 52.9 52.9的边界框mAP(mean Average Precision,平均精度均值)用于目标检测,以及 45.8 45.8 45.8的掩码mAP用于实例分割,且仅使用了单尺度推理。我们还使用了具有4尺度DINO检测头的ImageNet-21K GC ViT模型,实现了 58.3 % 58.3\% 58.3%的边界框AP(Average Precision,平均精度)。
此外,使用UPerNet(Xiao等人,2018)头,我们的模型在ADE20K数据集上通过仅使用单尺度推理方案实现了
49.2
49.2
49.2的mIoU(mean Intersection over Union,平均交并比)用于语义分割。与MS COCO和ADE20K数据集上类似大小的模型相比,具有不同学习能力的GC ViT的其他变体也展示了SOTA(state-of-the-art,最先进)结果。因此,GC ViT在各种下游任务中对高分辨率图像表现出极大的可扩展性,验证了所提出框架在捕获短距离和长距离信息方面的有效性。
我们的工作主要贡献总结如下:
- 我们引入了一种计算和参数优化的分层ViT,并对设计空间(如嵌入维度、头数、MLP比率)进行了重新参数化。
- 我们设计了一种高效的类似CNN的令牌生成器,该生成器以不同分辨率对空间特征进行编码,以用于全局查询表示。
- 我们提出了全局查询令牌,这些令牌能够以高效的方式有效地捕获上下文信息,并建模局部和全局交互。
- 我们引入了一个参数高效的降采样模块,该模块具有改进的Fused MB-Conv块,不仅集成了归纳偏置,还实现了通道间依赖性的建模。
- 我们展示了新的SOTA基准测试,包括:(1)在ImageNet-1K上,针对参数数量和浮点运算次数(FLOPs)的Pareto前沿的ImageNet分类;(2)在MS COCO和ADE20K上的下游任务,如检测、实例分割和语义分割。
2、GC ViT架构
架构。图2展示了GC ViT的架构。我们提出了一种分层框架,通过减小空间维度同时扩大嵌入维度(均以2为因子)来获得多个分辨率下的特征表示(称为阶段)。
首先,给定一个分辨率为 x ∈ R H × W × 3 \mathbf{x} \in \mathbb{R}^{H \times W \times 3} x∈RH×W×3的输入图像,我们通过应用一个步长为2且带有适当填充的 3 × 3 3 \times 3 3×3卷积层来获得重叠的块(patches)。然后,通过另一个步长为2的 3 × 3 3 \times 3 3×3卷积层将这些块投影到 C C C维的嵌入空间中。
每个GC ViT阶段都由交替的局部和全局自注意力模块组成,以提取空间特征。两者都像Swin Transformer(Liu et al., 2021)一样在局部窗口中操作,但全局自注意力可以访问由全局查询生成器提取的全局特征。查询生成器是一个类似于CNN的模块,它在每个阶段仅从整个图像中提取一次特征。在每个阶段之后,空间分辨率减小2倍,而通道数通过降采样块增加2倍。将得到的特征通过平均池化和线性层传递,以创建用于下游任务的嵌入。
GC ViT架构受益于下一节中描述的新型块,如下采样算子、全局查询生成器和全局自注意力模块。
下采样器:我们借鉴了CNN模型中空间特征收缩的思想,该思想在减少维度的同时引入了局部偏差和跨通道交互。我们利用了一个修改后的Fused-MBConv块,后面跟着一个核大小为3且步长为2的最大池化层作为下采样算子。我们工作中的Fused-MBConv块与EfficientNetV2(Tan & Le, 2021)中的块类似,但做了一些修改,如下所示:
x ^ = DWConv 3 × 3 ( x ) x ^ = GELU ( x ^ ) x ^ = SE ( x ^ ) x = Conv 1 × 1 ( x ^ ) + x \begin{array}{l} \hat{\mathbf{x}}=\operatorname{DWConv}_{3 \times 3}(\mathbf{x}) \\ \hat{\mathbf{x}}=\operatorname{GELU}(\hat{\mathbf{x}}) \\ \hat{\mathbf{x}}=\operatorname{SE}(\hat{\mathbf{x}}) \\ \mathbf{x}=\operatorname{Conv}_{1 \times 1}(\hat{\mathbf{x}})+\mathbf{x} \end{array} x^=DWConv3×3(x)x^=GELU(x^)x^=SE(x^)x=Conv1×1(x^)+x
其中,SE、GELU和DW-Conv_{3 \times 3}分别表示Squeeze-and-Excitation块(Hu et al., 2018)、高斯误差线性单元(Hendrycks & Gimpel, 2016)和 3 × 3 3 \times 3 3×3的深度卷积。在我们提出的架构中,Fused-MBConv块提供了所需的属性,如归纳偏置和跨通道依赖性的建模。这些特性在表8中得到了验证。
2.1. 全局自注意力
图3展示了我们的主要贡献背后的主要思想。局部自注意力只能在局部窗口内查询块,而全局自注意力可以在操作窗口内的同时查询不同的图像区域。在每个阶段,全局查询组件都是预先计算好的。全局自注意力利用提取的全局查询标记,并在所有块之间共享,以便与局部键和值表示进行交互。此外,GC ViT采用交替的局部和全局自注意力块来有效地捕获局部和全局空间信息。图S.1展示了局部自注意力和全局自注意力之间的差异。全局注意力查询
q
g
\mathrm{q}_{\mathrm{g}}
qg 的大小为
B
×
C
×
h
×
w
B \times C \times h \times w
B×C×h×w,其中
B
,
C
,
h
B, C, h
B,C,h 和
w
w
w 分别表示批量大小、嵌入维度、局部窗口的高度和宽度。此外,
q
g
\mathrm{q}_{\mathrm{g}}
qg 沿着批量维度重复,以补偿窗口总数和聚合批量大小
B
∗
=
B
×
N
∗
B^{*}=B \times N^{*}
B∗=B×N∗,其中
N
∗
N^{*}
N∗ 是局部窗口的数量。
q
g
\mathrm{q}_{\mathrm{g}}
qg 被进一步重塑为多个头。值和键在每个局部窗口内使用线性层计算。
由于分割的窗口仅包含局部信息,因此与全局查询标记中嵌入的丰富上下文信息的交互提供了一种扩大感受野并关注输入特征图中不同区域的有效方式。自注意力模块的计算方式如下:
Attention ( q g , k , v ) = Softmax ( q g k d + b ) v \operatorname{Attention}\left(\mathbf{q}_{\mathrm{g}}, k, v\right)=\operatorname{Softmax}\left(\frac{\mathbf{q}_{\mathrm{g}} k}{\sqrt{d}}+\mathbf{b}\right) \mathbf{v} Attention(qg,k,v)=Softmax(dqgk+b)v
其中
d
\mathbf{d}
d 是缩放因子,
b
\mathbf{b}
b 是一个可学习的相对位置偏置项。假设沿水平和垂直轴的位置变化在
[
−
p
+
1
,
p
−
1
]
[-p+1, p-1]
[−p+1,p−1] 范围内,
b
\mathbf{b}
b 从网格
b
^
∈
R
(
2
p
−
1
)
×
(
2
p
−
1
)
\hat{\mathbf{b}} \in \mathbb{R}^{(2 p-1) \times(2 p-1)}
b^∈R(2p−1)×(2p−1) 中采样。如第4节所示,相对位置偏置提高了性能,特别是对于密集预测下游任务。在算法1中,我们提供了在GC ViT中计算全局自注意力的类似PyTorch的伪代码。
2.2. 复杂度分析
给定每个阶段输入的特征图 x ∈ R H × W × C x \in \mathcal{R}^{H \times W \times C} x∈RH×W×C,其中窗口大小为 h × w h \times w h×w,GC ViT的计算复杂度如下:
O ( GC ViT ) = 2 H W ( 2 C 2 + h w C ) \mathcal{O}\left(\text{GC ViT}\right)=2 H W\left(2 C^{2}+h w C\right) O(GC ViT)=2HW(2C2+hwC)
全局查询标记生成器和其他组件的高效设计使得GC ViT能够在与Swin Transformer(Liu et al., 2021)相似的计算复杂度下,捕获长距离信息,并在分类以及检测和分割等下游任务中实现更高的准确性。
3. 实验
对于图像分类,我们在ImageNet-1K数据集(Deng et al., 2009)上训练和测试了我们的模型。为了进行公平比较,所有GC ViT变体均遵循先前工作(Liu et al., 2021; Yang et al., 2021b; Chu et al., 2021a)的训练配置进行训练。具体来说,所有模型均使用AdamW(Kingma&Ba,2014)优化器进行300个周期的训练,初始学习率为0.001,权重衰减为0.05,采用余弦衰减调度器,并分别设置20个预热和冷却周期。
对于目标检测和实例分割,我们使用DINO(He et al., 2017)和Mask-RCNN(He et al., 2017)头部在MS COCO(Lin et al., 2014)数据集上训练了我们的模型,采用
×
3
\times 3
×3学习率计划,初始学习率为0.0001,批量大小为16,权重衰减为0.05。按照(Liu et al., 2022b)的做法,我们使用Cascade Mask-RCNN与Tiny、Small和Base模型变体进行比较,但仅使用Mask-RCNN与Tiny变体进行比较。对于语义分割,我们使用ADE20K数据集(Zhou et al., 2017)和UPerNet(Xiao et al., 2018)分割头部。与之前的工作一样,我们为输入图像设置了
512
×
512
512 \times 512
512×512的随机裁剪大小。
3.1. 分类
我们在表1中展示了ImageNet-1K分类基准,并与不同模型大小的基于CNN和ViT的模型进行了比较。与ConvNeXt(Liu et al., 2022b)等其他已建立的基准相比,我们的模型取得了更好的性能。此外,如图1所示,在竞争对手模型上,GC ViT模型在FLOPs数量方面表现出更好或相当的计算效率。
3.2. 检测和实例分割
在表2中,我们展示了MS COCO数据集上的目标检测和实例分割基准。使用Mask-RCNN头部时,具有预训练GC ViT-T(47.9/43.2)主干的模型在框AP和掩码AP方面分别比具有预训练ConvNeXt-T(Liu et al., 2022b)(46.2/41.7)的模型高出+1.7和+1.5,比Swin-T(Liu et al., 2021)(46.0/41.6)的模型高出+1.9和+1.6。使用Cascade Mask-RCNN头部时,具有预训练GC ViT-T(51.6/44.6)和GC ViT-S(52.4/45.4)主干的模型在框AP和掩码AP方面分别比ConvNeXt-T(Liu et al., 2022b)(50.4/43.7)高出+1.2和+0.9,比ConvNeXt-S(Liu et al., 2022b)(51.9/45.0)高出+0.5和+0.4。此外,具有GC ViT-B(52.9/45.8)主干的模型在框AP和掩码AP方面分别比具有ConvNeXt-B(Liu et al., 2022b)(52.7/45.6)的模型高出+0.2和+0.2。
- 消融实验
组件级分析。如表5所示,我们研究了GC ViT模型中每个组件在分类、检测、实例分割和语义分割中的作用。为简便起见,我们以Swin Transformer作为基础模型,并逐步重新设计组件,以展示它们在提高性能方面的重要性。首先,我们移除了窗口平移机制,并预见到在所有任务中性能都有显著下降。将参数分布更改为我们的设计后,在精度、框AP、掩码AP和mIoU方面分别提高了+1.7%、+2.8%、+2.2%和+1.7%。这种重新参数化包括更改窗口大小、MLP比率、层数等。将GC ViT的基于CNN的stem添加到模型中,在精度、框AP、掩码AP和mIoU方面分别带来了+0.3%、+0.2%、+0.2%和+0.2%的额外提升。此外,使用我们提出的下采样器进一步将精度、框AP、掩码AP和mIoU分别提高了+0.4%、+0.1%、+0.1%和+0.3%。最后两项变化证明了卷积归纳偏差和捕获通道间依赖性在我们模型中的重要性。最后,利用我们提出的全局自注意力机制,在精度、框AP、掩码AP和mIoU方面分别提升了+0.9%、+0.8%、+0.6%和+1.2%。因此,这验证了所提出的全局自注意力机制的有效性,特别是对于高分辨率图像的下游任务(如语义分割),在这些任务中建模长距离空间依赖关系至关重要。
4.1. ImageNet-21K
在表6中,我们将在ImageNet-21K数据集上预训练并在ImageNet-1K数据集上微调后的GC ViT-L模型与同类方法进行了性能比较。GC ViTL在Top-1准确率上分别比Swin-L和CSwin-L高出+0.3%和+0.1%,同时与ConvNeXt-L模型的性能相当。因此,这验证了模型在大规模数据环境下的有效性。
4.2. 泛化能力
在表7中,我们评估了GC ViT在ImageNetV2数据集(Recht等人,2019)上的性能,以进一步衡量其鲁棒性。具体而言,我们采用了不同的采样策略,包括Matched Frequency和Threshold-0.7。这些基准测试展示了GC ViT在ImageNetV2数据集上的竞争性能,并验证了其在鲁棒性和泛化能力方面的有效性。
4.3. 下采样器设计
我们在表8中研究了各种下采样块的有效性。我们设计的最简单替代方案是一对卷积层和最大池化层。然而,这导致ImageNet Top-1准确率下降了-0.8%。Patch merging是Swin Transformers(Liu等人,2021)中引入的另一种变体。
然而,这导致准确率下降了-0.6%。最后,我们的下采样器由改进的Fused-MBConv块和带步长的卷积组成,并显示出最佳结果。前者组件的重要性通过SE操作来解释,该操作在保持参数数量和浮点运算次数较低的同时,增强了跨通道交互。我们得出结论,我们提出的下采样器对于实现高精度至关重要,因为它引入了卷积归纳偏差。
5. 可解释性
为了对提出的全局自注意力和查询标记的可解释性提供进一步的见解,我们在图5中展示了学习到的注意力和Grad CAM(Selvaraju等人,2017)映射的可视化。全局自注意力模块揭示的相关注意力分布与图像语义保持一致,因此为局部注意力模块提供了信息丰富的来源。此外,相应的Grad-CAM映射展示了精确的对象定位,并包含了最复杂的细节。
6、相关工作
卷积神经网络(ConvNet)。自深度学习兴起以来,卷积神经网络(CNN)(Krizhevsky等人,2012;Simonyan & Zisserman,2014;Howard等人,2017;He等人,2016;Szegedy等人,2016;Huang等人,2017;Hu等人,2018)凭借其最优性能在计算机视觉基准测试中占据主导地位。最近,ConvNeXt(Liu等人,2022b)对ResNet(He等人,2016)的架构进行了修改,并在分类、检测和分割任务上取得了具有竞争力的基准测试结果。
Transformer。ViT(Dosovitskiy等人,2020)首次作为CNN的替代方案被提出,其自注意力层具有扩大感受野的优势。然而,它缺乏CNN所期望的属性,如归纳偏差和平移不变性,并且需要大规模的训练数据集才能达到具有竞争力的性能。数据高效图像Transformer(DeiT)(Touvron等人,2021)引入了一种基于蒸馏的训练策略,显著提高了分类精度。
混合模型。LeViT(Graham等人,2021)提出了一种混合模型,该模型重新设计了多层感知器(MLP)和自注意力模块,这些模块经过高度优化以实现快速推理。交叉协方差图像转换器(XCiT)(Ali等人,2021)引入了一个转置自注意力模块,用于对特征通道之间的相互作用进行建模。卷积视觉转换器(CvT)(Wu等人,2021)在分层架构中引入了卷积令牌嵌入层和转换器块,以提高ViT的效率和准确性。条件位置编码视觉转换器(CPVT)(Chu等人,2021b)通过根据局部补丁令牌对位置编码进行条件化,在图像分类和对象检测等各种任务上展示了改进的性能。令牌到令牌视觉转换器(T2T-ViT)(Yuan等人,2021)提出了一种转换层,用于聚合相邻令牌并利用空间相关性建立图像先验。金字塔视觉转换器(PVT)(Wang等人,2021)提出了一种分层架构,在每个阶段的开始处进行补丁嵌入和空间维度缩减,以提高计算效率。独立地,Swin转换器(Liu等人,2021)也提出了一种分层架构,其中自注意力是在局部窗口内计算的,这些窗口会移动以实现区域交互。Twins转换器(Chu等人,2021a)提出了一种空间可分离的自注意力机制,该机制具有局部分组和全局子采样模块,以提高效率。
全局注意力。在计算机视觉中的EdgeViT(Pan等人,2022)和自然语言处理中的BigBird(Zaheer等人,2020)等其他研究工作中,已经在其各自的应用中提出了全局自注意力。然而,GC ViT中的全局注意力与这些方法在本质上是不同的。例如,EdgeViT采样代表性令牌,并且仅在减少特征大小的这些代表性令牌之间计算稀疏自注意力。相反,GC ViT在其各自的局部区域内,计算全局查询(而不仅仅是令牌)与局部键和值之间的自注意力,而不进行任何子采样。此外,在EdgeViT中,只有每个区域的子采样代表性令牌在自注意力模块中进行交互;然而,在GC ViT中,全局查询与整个局部区域进行交互。此外,BigBird结合了随机、窗口和全局注意力机制,这与GC ViT中提出的局部和全局自注意力方案不同。BigBird没有任何特定的机制来提取全局令牌,因为现有的令牌或额外的特殊令牌可以被指定为全局令牌。然而,GC ViT中的全局令牌是通过查询生成器从整个输入特征中提取上下文信息获得的。最后,BigBird使用一组全局令牌来关注整个输入序列。然而,在GC ViT中,全局查询令牌关注分区窗口中的局部键和值令牌,因为考虑到输入特征较大的尺寸,关注整个输入序列是不可行的。
7、结论
在本工作中,我们介绍了一种新颖的分层ViT,即GC ViT,它能够通过利用全局查询令牌有效地捕获全局上下文,并与局部区域进行交互。我们在各种任务上广泛验证了模型的有效性。所提出的GC ViT模型在ImageNet-1K数据集上,针对不同模型大小的图像分类任务,均达到了新的最优基准(SOTA),并且以显著优势超越了基于CNN和ViT的同类模型。此外,我们在高分辨率图像的检测和语义分割等下游任务上也取得了最优或具有竞争力的性能。