CoAtNet

2021-12-23  本文已影响0人  Valar_Morghulis

CoAtNet: Marrying Convolution and Attention for All Data Sizes

NeurIPS 2021 Poster

这篇论文在NeurIPS2019中被抽中分配给两个独立的委员会,其中一个委员会给了拒绝,另一个委员会给了接收。

https://openreview.net/forum?id=dUk5Foj5CLf

https://openreview.net/forum?id=JXAyJeYqUkZ

NeurIPS有着悠久的实验历史。2014年,NeurIPS进行了一项实验,两个独立委员会对10%的提交材料进行了审查,以量化审查过程中的随机性。今年,我们重复了这个实验的一个变体,以了解审查过程的质量是如何随着时间的推移而变化的。这篇论文是实验的一部分,因此被分配给两个委员会(由评审员、一位区域主席和一位高级区域主席组成),由他们做出独立的决定。如果两个委员会提出了相同的建议,则遵循这一建议。如果一个委员会建议接受,则该文件被接受(少数情况除外,其中其他委员会确定了我们认为的致命缺陷,例如关键结果中的错误)。

评论:这篇文章得到了轻微的负面评价——主要关注点是缺乏新颖性,以及一些模棱两可的说法。在反驳和讨论阶段之后,评审员承认ImageNet结果的质量,但仍有一些问题没有回答。因此,我建议拒绝,并鼓励作者考虑评审员的意见,以改进手稿并重新提交到其他地点。

评论:反驳解决了所有评审员的担忧,所有评审员都建议接受。AC同意这一建议

摘要:Transformer在计算机视觉领域吸引了越来越多的兴趣,但它们仍然落后于最先进的卷积网络。在这项工作中,我们表明,虽然Transformer往往具有更大的模型容量,但由于缺乏正确的归纳偏置,其泛化能力可能比卷积网络差。为了有效地结合两种架构的优势,我们提出了CoAtNets(发音为“coat”网),这是一个基于两个关键观点构建的混合模型家族:(1)深度卷积和自注意可以通过简单的相对注意自然地统一起来;(2) 以一种有原则的方式垂直堆叠卷积层和注意层在提高泛化、容量和效率方面出人意料地有效。实验表明,我们的CoAtNet在不同数据集的不同资源约束下实现了最先进的性能:在没有额外数据的情况下,CoAtNet达到86.0%的ImageNet top-1准确度;当使用ImageNet-21K中的13M图像进行预训练时,我们的CoAtNet达到了88.56%的top-1准确度,与JFT-300M中的300M图像匹配ViT巨大预训练,同时使用的数据减少了23倍;值得注意的是,当我们使用JFT-3B进一步放大CoAtNet时,它在ImageNet上达到了90.88%的top-1准确度,建立了一个新的最先进的结果。

1导言

自从AlexNet[1]取得突破以来,卷积神经网络(卷积网络)一直是计算机视觉的主要模型结构[2,3,4,5]。与此同时,随着自然语言处理领域中Transformer[6]等自注意模型的成功[7,8],许多以前的工作试图将注意力的力量引入计算机视觉[9,10,11,12]。最近,Vision Transformer(ViT)[13]表明,使用几乎只有1层的vanilla Transformer层,仅在ImageNet-1K[14]上就可以获得合理的性能。更重要的是,当对大规模弱标记JFT-300M数据集[15]进行预训练时,ViT获得了与最先进的(SOTA)卷积网络相当的结果,表明Transformer模型在规模上可能比卷积网络具有更高的容量。

虽然ViT通过大量JFT 300M训练图像显示了令人印象深刻的效果,但在低数据区,其性能仍落后于卷积网络。例如,在没有额外的JFT-300M预训练的情况下,ViT的ImageNet准确度仍然显著低于具有可比模型尺寸的卷积网络[5](见表13)。随后的工作使用特殊的正则化和更强的数据增强来改进普通ViT[16,17,18],但在相同的数据量和计算量下,这些ViT变体在ImageNet分类上都无法超过SOTA卷积模型[19,20]。这表明,普通Transformer层可能缺少卷积网络所具有的某些理想归纳偏置,因此需要大量数据和计算资源进行补偿。毫不奇怪,最近的许多工作一直试图将卷积网络的归纳偏置纳入Transformer模型,为注意层施加局部感受野[21,22],或用隐式或显式卷积运算增加注意层和FFN层[23,24,25]。然而,这些方法要么是临时的,要么集中于注入特定的属性,缺乏对卷积和注意力在组合时各自作用的系统理解。

在这项工作中,我们从机器学习的两个基本方面——泛化和模型容量——系统地研究了卷积和注意的混合问题。我们的研究表明,由于卷积层具有较强的归纳偏置先验知识,卷积层具有更好的泛化能力和更快的收敛速度,而注意层具有更高的模型容量,可以从更大的数据集中获益。卷积层和注意层相结合可以获得更好的泛化能力和容量;然而,这里的一个关键挑战是如何有效地将它们结合起来,以便在准确性和效率之间实现更好的权衡。在本文中,我们研究了两个关键的观点:首先,我们观察到常用的深度卷积可以有效地合并到具有简单相对注意的注意层中;第二,简单地以适当的方式叠加卷积层和注意层,对于实现更好的泛化和容量可能会非常有效。基于这些见解,我们提出了一种简单而有效的网络架构CoAtNet,它兼有卷积网络和Transformer的优点。

我们的CoAtNet在不同数据大小的可比资源限制下实现SOTA性能。具体来说,在低数据条件下,由于有利的归纳偏置,CoAtNet继承了卷积网络的强大泛化特性。此外,由于数据丰富,CoAtNet不仅具有Transformer模型优越的可扩展性,而且实现了更快的收敛,从而提高了效率。当仅使用ImageNet-1K进行训练时,CoAtNet达到86.0%的top-1准确度,在类似计算资源和训练条件下与现有技术NFNet[20]匹配。此外,当在ImageNet-21K上预训练约10万张图像时,CoAtNet在ImageNet-1K上微调时达到88.56%的top-1准确度,与在JFT-300M上预训练的ViT(一个23×更大的数据集)匹配。最后,当使用JFT-3B进行预训练时,与ViT相比,CoAtNet表现出更好的效率,并将ImageNet-1K top-1准确度提高到90.88%,同时使用ViT-G/14设置的现有技术的1.5倍计算量[26]。

2型

在本节中,我们将重点讨论如何“最佳”组合卷积和Transformer的问题。大致来说,我们将问题分解为两部分:

1.如何在一个基本的计算块中结合卷积和自注意?

2.如何将不同类型的计算块垂直堆叠在一起,形成一个完整的网络?

随着我们逐渐揭示我们的设计选择,分解的原理将变得更加清晰。

2.1融合卷积和自注意

对于卷积,我们主要关注MBConv块[27],它采用深度卷积[28]来捕获空间交互。选择此选项的一个关键原因是,Transformer和MBConv中的FFN模块均采用“反向瓶颈”设计,该设计首先将输入的通道大小扩展4倍,然后将4倍宽的隐藏状态投影回原始通道大小,以实现残差连接。

除了倒转瓶颈的相似性外,我们还注意到深度卷积和自注意都可以表示为预定义感受野中每维值的加权和。具体地说,卷积依赖于一个固定的核从一个局部感受野收集信息

何席,易∈ rd分别是位置i处的输入和输出,L(i)表示i的局部邻域,例如,在图像处理中以i为中心的3x3网格。

相比之下,自注意允许感受野成为整个空间位置,并基于对之间的重新标准化成对相似性计算权重(xi,xj):

其中G表示全局空间。在讨论如何最好地结合它们之前,比较它们的相对优势和劣势是值得的,这有助于找出我们希望保留的优良特性。

•首先,深度卷积核wi−j是一个与输入无关的静态值参数,而注意权重Ai,j动态地依赖于输入的表示。因此,自注意更容易捕捉不同空间位置之间的复杂关系交互,这是我们在处理高级概念时最渴望的特性。然而,这种灵活性带来了更容易过度拟合的风险,尤其是在数据有限的情况下。

第二对,给出任意位置的卷积−j只关心它们之间的相对移动,即− j、 而不是i或j的具体值。这一特性通常被称为平移等变性,已经发现平移等变性可以改善有限数据集下的泛化[29]。由于使用绝对位置嵌入,标准Transformer(ViT)缺少此特性。这部分解释了为什么当数据集不是很大时,卷积网络通常比Transformers更好。

•最后,感受野的大小是自注意和卷积之间最关键的区别之一。一般来说,更大的感受野提供更多的上下文信息,这可能导致更高的模型容量。因此,全局感受野一直是在视觉中运用自注意的关键动机。然而,大的感受野需要更多的计算。在全局注意的情况下,复杂性是二次w.r.t.空间大小,这是应用自注意模型的基本权衡

表1

鉴于上述比较,理想模型应能够结合表1中的3个理想特性。与方程n中的深度卷积形式类似。(1) 和Eqn中的自注意。(2) ,可以实现这一点的简单想法是,在Softmax归一化之后或之前,简单地将全局静态卷积核与自适应注意矩阵相加,即。,

有趣的是,虽然这个想法似乎过于简化,但预归一化版本y pre对应于相对自注意的一种特殊变体[30,31]。在这种情况下,注意力权重Ai,j由wi共同决定−平移等变的j和输入自适应的x>ixj,根据它们的相对大小可以享受这两种效果。重要的是,请注意,为了在不增加参数数量的情况下启用全局卷积核,我们重新加载了wi的表示法−j作为标量(即w∈ rO(| G |)而不是方程n中的向量。(1). w的标量公式的另一个优点是检索wi−通过计算两两点积关注度,所有(i,j)的j明显包含在内,因此产生最小的额外成本(见附录A.1)。考虑到这些好处,我们将在Eqn中使用带有预归一化相对注意变量的Transformer块。(3) 作为提出的CoAtNet模型的关键组成部分。

2.2垂直布局设计

在找出一种巧妙的组合卷积和注意力的方法之后,我们接下来考虑如何利用它来堆叠整个网络。

正如我们上面所讨论的,全局上下文具有二次复杂性w.r.t.空间大小。因此,如果我们直接应用等式n中的相对注意。(3) 对于原始图像输入,由于任何普通尺寸的图像中都有大量像素,因此计算速度会非常慢。因此,要构建一个实际可行的网络,我们主要有三种选择:

(A) 在特征图达到可管理水平后,执行一些下采样以减小空间大小并使用全局相对注意。

(B) 加强局部注意,这将注意中的全局感受野G限制为局部野L,就像卷积[22,21]。

(C) 将二次型Softmax注意力替换为仅具有线性复杂度w.r.t.空间大小的线性注意力变量[12,32,33]。

我们简单地尝试了选项(C),但没有得到一个合理的好结果。对于选项(B),我们发现实现局部注意涉及许多需要大量内存访问的非平凡形状格式化操作。在我们选择的加速器(TPU)上,这样的操作非常缓慢[34],这不仅违背了加速全局注意力的最初目的,而且损害了模型的容量。因此,由于最近的一些研究工作已经研究了这个变量[22,21],我们将重点关注选项(A),并在我们的实证研究(第4节)中将我们的结果与他们的结果进行比较。

对于选项(A),可通过(1)具有积极步幅(例如,步幅16x16)的卷积杆(如ViT)或(2)具有渐进池的多级网络(如卷积网络)实现下采样。有了这些选择,我们得到了5个变体的搜索空间,并在对照实验中对它们进行了比较。

•当使用ViT杆时,我们会相对注意地直接堆叠LTransformer块,我们称之为VITERL。

•当使用多级布局时,我们模拟卷积网络构建一个由5个阶段(S0、S1、S2、S3和S4)组成的网络,空间分辨率从S0逐渐降低到S4。在每个阶段开始时,我们总是将空间大小减少2倍,并增加通道数量(有关详细的下采样实施,请参见附录A.1)。

第一阶段S0是一个简单的2层卷积干,S1总是使用带有挤压激发(SE)的MBConv块,因为空间尺寸太大,无法引起全局注意。从S2到S4,我们考虑了MbCONV或Transformer块,约束条件是在Transformer级之前必须出现卷积阶段。该约束基于一个先验知识,即卷积在处理早期更常见的局部模式时更好。这导致4种变型具有越来越多的Transformer级,C-C-C-C、C-C-C-T、C-C-T-T和C-T-T-T,其中C和T分别表示卷积和Transformer。

为了系统地研究设计选择,我们考虑了两个基本方面:泛化能力和模型容量:对于泛化,我们感兴趣的是训练损失和评估准确度之间的差距。如果两个模型具有相同的训练损失,那么评估准确度较高的模型具有更好的泛化能力,因为它可以更好地泛化到未见的评估数据集。当训练数据大小有限时,泛化能力对数据效率尤为重要。对于模型容量,我们衡量适应大型训练数据集的能力。当训练数据丰富且不存在过拟合问题时,经过合理的训练步骤,容量较大的模型将获得更好的最终性能。请注意,由于简单地增加模型尺寸可以提高模型容量,为了进行有意义的比较,我们确保5个变体的模型尺寸具有可比性。

为了比较泛化和模型容量,我们在ImageNet-1K(1.3M)和JFT(>300M)数据集上分别训练了300和3个epoch的混合模型的不同变体,这两个epoch都没有任何正则化或增强。图1总结了两个数据集的训练损失和评估准确度。

•从ImageNet-1K的结果来看,一个关键的观察结果是,就泛化能力(即,训练和评估指标之间的差距)而言,我们有

特别是,Vitel明显比变异体差很多,我们推测这与它的激进下采样干缺乏适当的低水平信息处理有关。在多阶段变量中,总体趋势是模型的卷积阶段越多,泛化差距越小。

•对于模型容量,根据JFT比较,训练结束时的训练和评估指标表明以下排名:

C-C-T-T≈ C-T-T-T>Vitell>C-C-C-T>C-C-C-C。

重要的是,这表明简单地拥有更多的Transformer块并不一定意味着更高的视觉处理能力。一方面,虽然最初情况更糟,但Vitel最终赶上了具有更多MBConv级的两种变体,这表明Transformer块的容量优势。另一方面,C-C-T-T和C-T-T-T都明显优于Vitel,这表明具有激进步幅的ViT干细胞可能丢失了太多信息,从而限制了模型容量。更有趣的是,C-C-T-T≈ C-T-T-T表明,对于处理低级信息,静态局部操作(如卷积)可以与自适应全局注意机制一样有效,同时大大节省计算和内存使用。

最后,为了在C-C-T-T和C-T-T-T之间做出决定,我们进行了另一个可迁移性测试3——我们在ImageNet-1K上对上述两个JFT预训练模型进行了30个epoch的微调,并比较了它们的迁移性能。从表2可以看出,尽管训练前表现相同,C-C-T-T-T的传递准确度明显优于C-T-T-T。

表2

考虑到通用性、模型容量、可迁移性和效率,我们采用了CoAtNet的C-C-T-T多级布局。更多模型详情见附录A.1。

3相关工作

卷积网络构建块。

卷积网络(卷积网络s)已经成为许多计算机视觉任务的主要神经结构。传统上,常规卷积(如ResNet块[3])在大规模网络中很流行;相比之下,深度卷积[28]由于其较低的计算成本和较小的参数大小而在移动平台中流行[27]。最近的工作表明,基于深度卷积的改进反向残差瓶颈(MBConv[27,35])可以实现高准确度和更高的效率[5,19]。如第2节所述,由于MBConv和Transformer块之间的强连接,本文主要使用MBConv作为卷积构建块。

自注意和Transformer。

Transformer具有自注意的关键成分,已被广泛用于神经语言处理和语音理解。作为一项早期工作,独立的自注意网络[34]表明,单独的自注意可以很好地用于不同的视觉任务,尽管有一些实际困难。最近,ViT[13]将一个普通的Transformer应用于ImageNet分类,并在大规模JFT数据集上进行预训练后获得了令人印象深刻的结果。然而,当训练数据有限时,ViT仍在很大程度上落后于最先进的卷积网络。此后,许多最近的工作都集中于改进vision Transformer,以提高数据效率和模型效率。为了更全面地回顾视觉Transformer,我们请读者参考专门的调查[36,37]。

相对注意力。

在相对关注这个总称下,文献中出现了各种变体[30、38、39、34、40、31]。一般来说,我们可以将它们分为两类:(a)输入依赖型,其中额外的相对注意分数是输入状态f(xi,xj,i)的函数− j) ,和(b)输入独立版本f(i− j) 。CoAtNet中的变量属于独立于输入的版本,与T5中使用的变量类似[31],但与T5不同,我们既不跨层共享相对注意参数,也不使用bucketing机制。作为输入独立性的一个好处,获得f(i− j) 对于所有(i,j)对,计算上比TPU上的输入相关版本便宜得多。此外,在推断时,只需计算一次并缓存以备将来使用。最近的一项工作[22]也利用了这种独立于输入的参数化,但它将感受野限制在局部窗口。

结合卷积和自注意。

将卷积和自注意相结合用于视觉识别的想法并不新鲜。一种常见的方法是使用显式自注意或非局部模块增强卷积网络主干[9、10、11、12],或使用标准自注意[11]或更灵活的线性注意和卷积混合[41]替换某些卷积层。虽然自注意通常会提高准确性,但它们通常会带来额外的计算成本,因此通常被视为卷积网络的附加模块,类似于挤压和激励[42]模块。相比之下,在ViT和ResNet ViT取得成功[13]之后,另一个流行的研究方向是从Transformer主干开始,并尝试将显式卷积或卷积的某些理想特性纳入Transformer主干[25、24、23、22、21、43、44]。

虽然我们的工作也属于这一类,但我们表明,我们的相对注意力实例化是深度卷积和基于内容的注意力的自然混合,额外成本最小。更重要的是,从通用性和模型容量的角度出发,我们采用系统的方法进行垂直布局设计,并说明不同网络阶段如何以及为什么喜欢不同类型的层。因此,与仅使用现成卷积网络作为stem层的模型(如ResNet ViT[13])相比,CoAtNet还可以在总体尺寸增加时缩放卷积阶段(S2)。另一方面,与采用局部注意的模型相比[22,21],CoAtNet始终对S3和S4充分注意,以确保模型容量,因为S3占据了大部分计算和参数。

4个实验

在本节中,我们将CoAtNet与以前在类似设置下的结果进行比较。为完整起见,此处未提及的所有超参数均包含在附录A.2中。

4.1实验设置

CoAtNet模型族。

为了与不同尺寸的现有模型进行比较,我们还设计了一系列CoAtNet模型,如表3所示。总的来说,我们总是将S1到S4的通道数量增加一倍,同时确保stemS0的宽度小于或等于S1的宽度。此外,为了简单起见,当增加网络深度时,我们只缩放S2和S3中的块数。

表3

评估协议。

我们的实验集中在图像分类上。为了评估模型在不同数据大小下的性能,我们使用了三个越来越大的数据集,即ImageNet-1K(1.28M图像)、ImageNet-21K(12.7M图像)和JFT(300M图像)。在之前的工作之后,我们首先在三个数据集中的每一个上以224分辨率分别对300、90和14个epoch的模型进行预训练。然后,我们在ImageNet-1K上以所需的分辨率对预先训练好的模型进行30次微调,并获得相应的评估准确度。一个例外是分辨率为224的ImageNet-1K性能,可在预训练结束时直接获得。请注意,与使用Transformer块的其他模型类似,在不进行微调的情况下,以更高的分辨率直接评估在ImageNet-1K上预先训练的模型通常会导致性能下降。因此,每当输入分辨率发生变化时,总是采用微调。

数据增强和正则化。

在这项工作中,我们只考虑两种广泛使用的数据增强,即Rangous[ 45 ]和MixUp〔46〕,以及三种常用技术,包括随机深度(47)、标记平滑(48)和权重衰减[49 ],以使模型正规化。直观地说,增广和正则化方法的特定超参数取决于模型大小和数据规模,其中强正则化通常适用于较大的模型和较小的数据集。

在一般原则下,当前范式下的一个复杂问题是如何调整预训练和微调的正则化,因为数据大小可能会发生变化。具体来说,我们有一个有趣的观察结果,即如果在训练前完全禁用某种类型的增强功能,那么在微调过程中简单地启用它很可能会损害性能,而不是提高性能。我们推测这可能与数据分布的改变有关。因此,对于建议模型的某些运行,我们在对两个较大的数据集ImageNet21-K和JFT进行预训练时,故意采用RandAugment和小程度的随机深度。尽管这种正则化会损害预训练度量,但这允许在微调期间进行更通用的正则化和增强,从而提高下游性能。

4.2主要结果

ImageNet-1K

仅使用ImageNet-1K数据集的实验结果如表4所示。在类似条件下,提出的CoAtNet模型不仅优于ViT变体,而且还匹配最佳的仅卷积结构,即EfficientNet-V2和NFNET。此外,我们还将图2中分辨率为224x224的所有结果可视化。正如我们所看到的,CoAtNet的扩展性比以前的带有注意模块的模型要好得多。

表4 图2

ImageNet-21K

如表4和图3所示,当ImageNet-21K用于预训练时,CoAtNet的优势变得更加明显,大大优于所有以前的模型。值得注意的是,最佳CoAtNet变体达到了88.56%的top-1准确度,与88.55%的ViTH/14性能相匹配,这需要在一个23倍大的专有弱标记数据集(JFT)上预训练2.3倍大的ViT模型2.2倍多的步骤。这标志着数据效率和计算效率的显著提高。

图3

JFT

最后,在表5中,我们使用JFT300M和JFT-3B进一步评估了大规模数据体制下的CoAtNet。令人鼓舞的是,我们的CoAtNet-4几乎可以与NFNet-F4+设置的JFT-300M的最佳性能相媲美,同时在TPU训练时间和参数计数方面效率提高了2倍。当我们扩大模型以消耗与NFNet-F4+类似的训练资源时,CoAtNet-5在top-1准确度上达到89.77%,在可比设置下优于先前的结果。

表5

此外,随着我们进一步将训练资源推向ViT-G/14所使用的水平,并利用同样的JFT-3B数据集,其规模甚至更大[26],计算量减少了4倍以上,CoAtNet-6能够与ViT-G/14的性能相匹配,计算量减少了1.5倍,CoAtNet-7的最高准确度达到了89.77%,即90.88%,实现新的最先进的性能。

4.3消融研究

在本节中,我们将详细介绍CoAtNet的设计选择。

首先,我们将卷积和注意结合到一个计算单元中,研究了相对注意的重要性。具体来说,我们比较了两种模型,一种是相对注意,另一种是不相对注意,分别在ImageNet-1K单独和ImageNet-21K迁移设置下进行比较。从表6中可以看出,当仅使用ImageNet-1K时,相对注意明显优于标准注意,表明具有更好的泛化能力。此外,在ImageNet-21K迁移设置下,尽管相对注意变量的预训练性能非常接近,但其迁移准确度仍大大提高。这表明相对注意在视觉处理中的主要优势不是更高的能力,而是更好的泛化能力。

其次,由于S2与MBConv块和S3与相对Transformer块占据了CoAtNet的大部分计算,因此需要提出的问题是如何在S2(MBConv)和S3(Transformer)之间分割计算以实现良好的性能。在实践中,它归结为决定每个阶段的区块数量,我们称之为“布局”设计。为此,我们比较了表7中试验的几种不同布局。

•如果我们将S2和S3中的块总数保持不变,并改变每个阶段的块数,我们观察到V0是V1和V2之间的最佳点。基本上,S3中有更多的Transformer块通常会导致更好的性能,直到S2中的MBConv块的数量太少而无法很好地推广。

•为了进一步评估迁移设置中是否也存在最佳点,通常认为更高的容量更重要,我们进一步比较ImageNet21K迁移到ImageNet-1K设置下的V0和V1。有趣的是,尽管V1和V0在ImageNet-21K预训练期间具有相同的性能,但V1的迁移准确度明显落后于V0。这再次表明卷积在实现良好的可迁移性和泛化方面的重要性。

最后,我们研究了模型细节的两种选择,即每个注意头的维度(默认为32)以及MBConv块中使用的归一化类型(默认为BatchNorm)。从表8中,我们可以看到,将注意力头大小从32增加到64会略微影响性能,尽管它实际上显著提高了TPU的速度。在实践中,这将是一种质量-速度权衡。另一方面,BatchNorm和LayerNorm具有几乎相同的性能,而BatchNorm在TPU上的速度要快10-20%,这取决于每个核心批大小。

5结论

在本文中,我们系统地研究了卷积和Transformer的性质,从而找到了一种将它们组合成一个新的模型族CoAtNet的原则性方法。大量的实验表明,CoAtNet具有良好的通用性(如卷积网络s)和优越的模型容量(如Transformers),在不同的数据大小和计算预算下实现了最先进的性能。

注意,本文目前主要关注用于模型开发的ImageNet分类。然而,我们相信我们的方法适用于更广泛的应用,如目标检测和语义分割。我们将把它们留给以后的工作。

附录

A.1模型细节

首先,图4示出了CoAtNet的概述。

图4

二维相对注意

在Eqn中实现相对注意的预归一化版本。3对于尺寸为[H×W]的2D图像,我们为每个头部创建一个尺寸为[(2H)的可训练参数P−1) ×(2W−1) ],因为最大距离为2H− 1和2W− 分别为1。然后,对于两个空间位置(i,j)和(i0,j0),相应的相对偏差为Pi−i 0+H,j−j 0+W在基于1的索引下。为了实现,我们需要从[(2H)索引H2W2元素−1) ×(2W−1) 矩阵。在TPU上,我们分别利用沿高度轴和宽度轴的两个EINSUM,以复杂度O(HW(H+W))索引相对偏差,该相对偏差严格包含在O(H2W2D)注意复杂度中。在GPU上,使用gather可以更高效地进行索引,它只需要访问内存。注意,在推断时,从[(2H)索引H2W2元素− 1) ×(2W− 1) ]矩阵可以预先计算和缓存,以进一步提高吞吐量。

当在更大的分辨率上进行微调时,我们只需使用双线性插值来增加大小[(2H− 1) ×(2W− 1) ]至所需尺寸[(2H0− 1) ×(2W0− 1) ]对于任何H0>H和W0>W。

预激活

为了促进模型架构的同质性,我们始终为MBConv和Transformer块使用预激活结构[50],即。,

十、← x+模(范数(x)),

其中,模块表示MBConv、自注意或FFN模块,而Norm对应于MBConv的BATCHTNORM和自注意和FFN的LayerNorm。我们已经尝试在MBConv块中使用LayerForm,它在加速器(TPU)上实现了相同的性能,但速度明显较慢。一般来说,我们建议您选择设备上速度更快的。本着同样的精神,在MBConv块和Transformer块中,高斯误差线性单位(GELUs)[51]用作激活函数。

下采样

对于从S1到S4的每个级内的第一块,对残差分支和标识分支独立地执行下采样。具体而言,对于Transformer块,stride 2的标准最大池直接应用于自注意模块的两个分支的输入状态,类似于漏斗Transformer[52]。同时,对恒等分支应用通道投影以扩大隐藏大小。因此,下采样自注意模块可以表示为

十、← 项目(池(x))+注意(池(常模(x)))。(4)

对于MBConv块,通过对归一化输入使用跨步2卷积来实现残差分支中的下采样,即。,

十、← 项目(池(x))+Transformer(深度Transformer(Transformer(标准(x),步幅=2)))。(5)

这与标准MBConv不同,标准MBConv通过对反向瓶颈隐藏状态应用Stripe2深度卷积进行下采样。我们后来发现,使用stride-2深度卷积有帮助,但在模型较小时速度较慢,但在模型缩放时速度较慢,如表9所示。因此,如果没有另外提及,正文中报告的数字使用Eqn中的下采样实现。(5). 在实践中,这可能是另一个质量-速度权衡,可以调整为较小的模型。

分类头

与在ViT中添加额外的<cls>token来执行分类不同,我们将全局平均池应用于最后阶段的输出,以获得简单的表示。

A.2超参数

表10:主要实验中使用的超参数。斜线符号“/”用于分隔用于各种CoAtNet型号尺寸的不同超参数:对于微调稍大的CoAtNet-3,使用2,20的RandAugment。†:表14中对CoAtNet-4的PT-RA变体应用2,5的随机扩增。

表10

表10总结了第4节中主要实验所用的超参数。

表11总结了JFT实验中使用的CoAtNet-5的模型尺寸。与表3中的标准CoAtNet模型不同,我们将CoAtNet-5的每个注意头的大小设置为64而不是32,因为这实现了更好的速度性能权衡,如第4.3节所述。

表11

对于CoAtNet-6和CoAtNet-7,为了减少内存消耗,我们将S2的2/3 MBConv块移动到S3中,并将其隐藏维度加倍。虽然此修改不会改变FLOPs的复杂性,但这将使这些MBConv块的激活相关内存使用量减少一半,这使我们能够构建更大的模型。通过此调整,S3成为混合块类型和隐藏尺寸的阶段。此外,我们将注意力头的大小进一步增加到128,以提高速度性能权衡。表12总结了具体尺寸。基本上,CoAtNet-6和CoAtNet-7的深度相同,但宽度不同。

表12

A.3完全比较

表13 表14
上一篇下一篇

猜你喜欢

热点阅读