人们对Transformers应用在在视觉领域越来越感兴趣,但是他们仍然落后于现在的最好的卷积神经网络。在本篇论文中,我们显示了尽管 Transformers试图有更大的模型性能,但是由于缺乏正确的归纳偏置,Transformers的泛化能力比卷积神经网络要差。为了有效的结合两个架构的强项,我们呈现了 CoAtNets模型(读作‘coat’ nets),根据两个关键构思而构建的混合模型家族: (1) 通过简单的思考,很自然的将深度卷积和自注意力统一起来。 (2) 按照某种方式垂直堆叠卷积层和注意力层在提高泛化能力、性能、高效上表现的出奇的有效。实验证明我们的 CoAtNets在不同数据集的不同资源约束下达到了目前最好的性能:没用额外的数据, CoAtNet在 ImageNet上的top-1精度达到了86.0%;当在具有1300万个图片的 ImageNet-21K上预训练时,我们的 CoAtNet达到了 88.56% top-1精度,达到了ViT-huge在具有3亿张图片的 JFT-300M上的预训练结果,尽管我们是在比 1 23 \frac{1}{23} 231还低的数据集上( 300 M 13 M = 23.0769230769 \frac{300M}{13M}=23.0769230769 13M300M=23.0769230769)。显然,当我们进一步的在CoAtNet使用更大规模的JFT-3B数据集时,它在 ImageNet上达到了90.88%的top-1精度,取得了一个新的目前最好的结果。
自从AlexNet取得了重大突破,卷积神经网络在计算机视觉上一直是主要的模型架构。同时,在自注意力模型像 Transformers在自然语言处理中的成功运用的推动下,很多以前的工作尝试将强有力的注意力引入计算机视觉中。最近,Vision Transformer (ViT)已经证明只需用很普通的 Transformer层就能在 ImageNet-1K上获得不错的性能。更重要的是,当在大规模弱标记的JFT-300M数据集上预训练时,ViT达到了可以和目前最好的卷积神经网络相媲美的结果,这表明Transformer模型比 ConvNets在大规模数据集上具有更高的潜在的性能。
尽管ViT在足够多的JFT 300M训练图片上表现出了令人印象深刻的结果,在小样本数据集上,它的性能仍然落后于ConvNets。例如,不使用额外的 JFT-300M数据集预训练, ViT在相同的模型大小下在 ImageNet上的精度仍然显著低于 ConvNets(卷积神经网络)(看表13)。随后的工作用特殊的正则化和更强的数据增强去改进普通的 ViT,然而至今这些 ViT的变种在 ImageNet分类上用相同数量的数据和计算量没有一个能胜过目前最好的只基于卷积神经网络的模型。这表明普通的Transformer层可能缺乏ConvNets拥有的某些理想的归纳偏置,因此需要大量数据和计算资源来补偿。并不惊讶,最近很多工作尝试将ConvNets的归纳偏置融合到 Transformer模型中,通过在注意力层使用局部感受野或者通过隐式或者显式卷积操作增强注意力和FFN(前馈神经网络)层。然而,当组合的时候,这些方法要么点对点要么聚焦于投射到一个特殊的属性上,缺乏对卷积和注意力各自角色的系统理解。
在本篇论文中,我们从机器学习当中的泛化性能和模型性能这两个基本角度系统的研究卷积和注意力混合的问题。我们的研究表明,卷积层倾向于拥有更好的泛化能力,更快的收敛速度这得益于他们强大的归纳偏置能力。然而注意力层拥有更高的模型性能这得益于更大的数据集。将卷积和注意力结合在一起能获得更好的泛化能力和性能。然而,在这里有一个关键的挑战是如何将他们有效的组合起来以在精度和有效性上达到更好的权衡。在本篇论文中,我们研究了两个关键的想法:第一,我们观察到用简单的相对注意力能有效的将公用的深度卷积融合到注意力层中;第二,用恰当的方式,简单的堆叠卷积和注意力层能取得令人惊讶的效果以达到更好的泛化能力和性能。基于这些观察,我们提出了一个简单的至今有效的网络架构并命名为 CoAtNet,它汲取了 ConvNets 和 Transformers的强项。
我们的CoAtNet在可比的资源有限的不同的数据大小下,达到了目前最好的性能。特别的,在小数据样本上, CoAtNet内在的很好的泛化属性得益于讨人喜欢的归纳偏置。而且,如果提供大量数据,CoAtNet不仅具有 Transformer模型的优异可扩展性,而且实现了更快的收敛,从而提高了效率。当只用 ImageNet-1K进行训练时,CoAtNet达到了 86.0%的top-1精度,在相同的计算资源和训练条件下与现有技术NFNet模型相匹敌。进一步的当我们在ImageNet-21K上用1000万张图片进行预训练,在 ImageNet-1K上微调,CoAtNet达到了88.56%的 top-1精度,与ViT-Huge在JFT-300M(是当前数据集的23倍)上的预训练结果相匹敌。最后,当使用JFT-300M进行预训练, CoAtNet展示出了比ViT更高效,把 ImageNet-1K上的top-1精度推到了90.88%,然而使用的计算量比现有的ViT-G/14却少1.5倍。
在这部分,我们聚焦在如何将卷积和transformer最佳的结合在一起。简单的说,我们将问题分为两部分:
1、在一个基本的计算块上怎样将卷积和注意力结合在一起?
2、怎样将不同类型的计算块垂直堆叠在一起来形成完整的网络?
随着我们逐渐揭示我们的设计选择,分解的基本原理将变得更加清晰。
对于卷积,我们主要关注在MBConv块,它使用深度卷积来捕获空间交互作用。这样选择的一个关键原因是在 Transformer 和 MBConv中的FFN模块都采用“倒瓶颈”设计,这种设计首先将输入通道扩大4倍,然后在原始通道后边投射四倍宽的隐藏状态并使用残差连接。
除了相似的倒瓶颈设计,我们还注意到,深度卷积和自我注意都可以表示为预定义感受野中每维值的加权和。特别的,卷积依赖一个固定的核从一个局部感受野中收集信息 y i = ∑ y ∈ L ( i ) w i − j ⊙ x j ( 深 度 卷 积 ) ( 1 ) y_i = \sum_{y\in\mathcal L(i)}w_{i-j}\odot x_j\quad(深度卷积)\quad\quad\quad\quad(1) yi=y∈L(i)∑wi−j⊙xj(深度卷积)(1)其中 x i , y i ∈ R D x_i, y_i\in \Bbb R^D xi,yi∈RD分别是在位置的输入和输出, L ( i ) \mathcal L(i) L(i)表示i的邻域,例如:图像处理中以i为中心的3x3网格。
相比之下,自注意力允许感受野是整个空间并且计算权重是基于在 ( x i , x j ) (x_i , x_j ) (xi,xj)对中的重新归一化后的成对的相似性。 y i = ∑ j ∈ G e x p ( x i T x j ) ∑ k ∈ G e x p ( x i T x k ) ⏟ A i , j x j ( 自 注 意 力 ) ( 2 ) y_i = \sum_{j\in\mathcal G}\frac{exp(x_i^Tx_j)}{\underbrace{\sum_{k\in\mathcal G}exp(x_i^Tx_k)}_{A_{i,j}}}x_j\quad(自注意力)\quad\quad(2) yi=j∈G∑Ai,j k∈G∑exp(xiTxk)exp(xiTxj)xj(自注意力)(2)其中 G \mathcal G G表示全局空间.在讨论如何最好地组合它们之前,比较一下他们的相对优势和弱势是值得的,这帮助我们指出我们需要保留的好的属性。
Table 1: 在卷积或者自注意力中发现的有价值的属性
根据上边的比较,一个理想的模型应该能结合在表一中的三种属性。与方程式1中的深度卷积和方程式2中的自注意力的形式相同,一个直接的想法能做到这一点就是要么在Softmax归一化之前或者之后让一个全局静态卷积核和自适应注意力矩阵相加例如: y i p o s t = ∑ j ∈ G ( e x p ( x I T x j ) ∑ k ∈ G e x p ( x i T x k ) + w i − j ) x j o r y i p r e = ∑ j ∈ G e x p ( x i T x j + w i − j ) ∑ j ∈ G e x p ( x i T x k + w i − k ) x j ( 3 ) y_i^{post} = \sum_{j\in\mathcal G}\left(\frac{exp(x_I^Tx_j)}{\sum_{k\in\mathcal G}exp(x_i^Tx_k)}+w_{i-j} \right)x_j \quad \mathrm{or} \quad y_i^{pre} = \sum_{j\in\mathcal G}\frac{exp(x_i^Tx_j+w_{i-j})}{\sum_{j\in\mathcal G}exp(x_i^Tx_k+w_{i-k})}x_j\quad(3) yipost=j∈G∑(∑k∈Gexp(xiTxk)exp(xITxj)+wi−j)xjoryipre=j∈G∑∑j∈Gexp(xiTxk+wi−k)exp(xiTxj+wi−j)xj(3)有意思的是,尽管这个想法看起来过于简单,预归一化版本相当于相对自注意力的一个特殊变体。在这种情况下,注意力权重 A i , j A_{i,j} Ai,j是由平移不变性的 w i − j w_{i-j} wi−j和输入自适应 x i T x j x_i^Tx_j xiTxj共同决定的,这能根据他们的相对重要性大小来决定他们的效果。重要的是,为了在不增加参数量的情况下使用全局卷积核,我们需要重新加载 w i − j w_{i-j} wi−j做为一个标量而不是方程式1中的向量。公式中标量w的另一个优势是为所有的(i, j)检索 w i − j w_{i-j} wi−j是清楚的包括通过计算成对点积注意力,因此导致少量的额外开销。考虑到这些好处,我们将会使用方程式三中的带有与归一化的相对注意力变种的 Transformer 块作为我们所提的 CoAtNet模型的一个关键部分。
将卷积和注意力结合在一起的思路构思好后,我们接下来考虑怎样将他们堆叠在一整个网络中。
正如我们在上边讨论的那样,全局上下文的复杂度和空间大小呈平方关系。因此,如果我们直接应用方程式3中的相对注意力到未加工过的图片输入上,由于任何常见尺寸的图片的像素都很大,所以计算速度会非常慢。因此,为了构造一个在实际中可行的网络,我们有三个主要选项:
(A)执行一些下采样来减小空间大小,在特征映射之后使用全局相对注意力来达到可处理水平。
(B)加强局部注意力,这将全局感受野 G \mathcal G G限制到局部感受野 L \mathcal L L中就像在卷积中的操作一样。
(C) 用某个线性注意力变体取代平方级的Softmax注意力,线性注意力拥有的复杂度和空间大小呈线性关系。
我们在选项C上做了个简单的实验,并没有得到一个不错的结果。对于选项B我们发现实现局部注意力涉及很多非平凡形状操作这需要密集的内存访问。在我们选择的加速器上 (TPU),这样的操作导致非常的慢,这不仅违背了加速全局注意力的目的,而且破坏了模型性能。因此,一些最近的工作已经研究了这个变体,我们将会聚焦在选项A上,在我们的实验中比较了我们的结果和他们的结果(在第四部分)。
对于选项A,下采样可以通过一个大步长的卷积(比如16x16)或者用一个分步池化的多阶段网络就像在卷积神经网络中的那样来实现。带着这些选项,我们推导了5个变体的搜索空间,并在对照实验中对它们进行了比较。
第一阶段S0是一个简单的两层卷积, S1应用带有 squeeze-excitation (SE)的MBConv块,因为对于全局注意力来说空间大小是太大了。从S2到S4我们考虑要么采用 MBConv要么采用Transformer块并且卷积必须在Transformer之前。做这种限制是基于以前的经验:卷积在处理局部图案时表现较好,而且在早期是一种通用做法。这导致了随着Transformer的增加出现了四个变体:C-C-C-C, C-C-C-T, C-C-T-T和C-T-T-T,其中C和T分别表示卷积和 Transformer的第一个大写字母。
为了系统研究设计选择,我们考虑两个基础因素泛化能力和模型性能:对于泛化能力,我们对训练损失和评估精度之间的差距感兴趣。如果两个模型有相同的训练损失,则具有较高评估精度的模型则拥有更好的泛化能力,因为它将在从未使用过的数据集上泛化的更好。当训练数据集是有限时,泛化能力对数据的有效性是尤为重要的。对于模型性能,我们在大的训练数据集上评估其性能。当训练数据是足够的时,过拟合就不是问题了,具有更高性能的模型在之后的推理中就会达到更好的最终性能。注意:因为简单的增加模型大小能造成更高的模型性能,为了进行一个有意义的竞争,我们确保五个变体的模型的大小是可比较的。
Figure 1: 在不同大小的数据集下对比模型的泛化能力和性能。为了公平,所有模型有相同的参数大小和计算开销。
为了比较泛化能力和模型性能,我们在ImageNet-1K (1.3M) 和 JFT数据集上分别对不同的混合模型变体进行300轮和3轮的训练,两个都没有做任何正则化或者数据增强操作。在两个数据集上的训练损失和评估精度看图1。
最后,为了在 C-C-T-T 和 C-T-T-T中做决定,我们安排了另一个迁移测试—我们在ImageNet-1K上微调两个JFT预训练模型30轮,对比他们的迁移性能。从表2上看,它显示出 C-C-T-T比 C-T-T-T有更好的迁移精度,尽管预训练性能相同。
考虑到泛化能力,模型性能,迁移能力,和高效性,我们为CoAtNet采用了C-C-T-T多阶段布局。更多的细节看附录 A.1。
卷积网络构建块 卷积网络 (ConvNets)是许多计算机视觉任务的主要神经架构。传统的,常规卷积像ResNet在大规模卷积网络中是很受欢迎的。相比之下,在移动平台下深度卷积是受欢迎的,因为它更低的计算代价和更小的参数量。最近的工作表明,一个改进的倒置残差瓶颈 (MBConv)是建立在深度卷积的基础之上的,即能达到高精度,又能达到高效率。正如在第二部分讨论的那样,由于MBConv 和 Transformer块之间有很强的联系,本篇论文主要使用MBConv作为卷积构建块。
自注意力和Transformers 在关键元素自注意力的作用下,Transformers被广泛应用到自然语言处理,和语音辨识。作为早期工作,独立的自注意力网络表明单独的自注意力网络在不同的视觉任务上可以做的很好,尽管有一些实际的困难。最近, ViT应用一个普通的 Transformer到 ImageNet分类任务上,在大规模 JFT数据集上预训练后达到了一个令人印象深刻的结果。然而, 当训练数据集有限时ViT仍然大幅落后于目前最好的卷积神经网络。基于此,很多最近的工作聚焦于提高vision Transformers的数据高效性和模型高效性。为了对vision Transformers有一个更广泛的回顾,我们建议读者参考专门的调查 [36, 37]。
相对注意力 在相对注意力的统称下,有很多个变种[30, 38, 39, 34, 40, 31]。一般而言,我们能将其分为两个类别:(a) 输入依赖版本它的额外相对注意力得分是一个输入状态函数 f ( x i , x j , i − j ) f(x_i , x_j , i-j) f(xi,xj,i−j)。 (b) 输入依赖版本是 f ( i − j ) f(i-j) f(i−j)。 CoAtNet的变体属于输入依赖版本和 T5中的一个类似,但是与T5不同的是,我们既没有在跨层之间共享相对注意力参数也没有使用桶机制。输入独立的一个好处是在所有的 ( i , j ) (i, j) (i,j)对中获取 f ( i − j ) f(i-j) f(i−j)比在TPU上的输入依赖版本的计算代价更低。除此之外,在推理时,这只需要去计算一次并缓存下未来要用到的。一个最近的工作也统一一个输入依赖的参数,但是它将感受野限制在了局部窗口中。
将卷积和自注意力相结合 将卷积和自注意力相结合的想法对视觉识别任务来说并不新鲜。一个通用的做法去增强卷积神经网络的骨架是: 用明确的自注意力或者非局部模块或者用标准的自注意力取代某些卷积层或者灵活的将注意力和卷积线性混合。尽管自注意力通常用来提高精度,但是它们经常会产生额外的计算代价,因此,经常被视作 ConvNets的附加组件,类似于 squeeze-and-excitation模块。相比之下,在 ViT 和 ResNet-ViT的成功运用之后,另一个流行的研究方向是Transformer骨干网络即尝试在Transformer骨干网络中加入明确的卷积或者一些卷积的有价值的属性。
尽管我们的工作也属于这个类别,但我们的相对注意力实例化是一个深度卷积和基于内容注意力的自然混合,具有较小的附加代价。更重要的是,从泛化能力和模型性能的角度出发,我们采用系统的方法进行垂直布局设计,并说明不同的网络阶段如何以及为什么使用不同类型的层。因此,与那些简单的使用现成的卷积网络比如 ResNet-ViT作为主干层相比当整体大小增加时CoAtNet也会有大规模的卷积阶段。在另一方面,与使用局部注意力的模型相比CoAtNet在 S3 和 S4阶段一直使用全注意力,来确保模型的性能,因为S3占据了大部分计算和参数。
在这个部分,我们在可比较的设置下对比了CoAtNet和以前的模型的结果。为了完整性,所有的超参数在这里并没有提到,但在附录A.2中可以查看。
Table 3: L表示模块的数量,D表示隐藏维度 (或者叫通道数量),对于所有的Conv 和 MBConv块,我们都用核大小为3的核。对于所有的Transformer块,我们按照[22]设置每个注意力头大小为32。倒瓶颈的膨胀率为4,SE的膨胀(收缩)率为0.25。
CoAtNet模型家族 为了和已存在的不同大小的模型进行比较,我们也设计了一个模型家族,表三中是对这个家族的概述。总的来说,从S1 to S4每过一个阶段通道数量翻倍,同时确保主干S0的宽度比S1的更小或者相等。同样,为了简化,随着网络深度的增加,我们只在S2和S3缩放块的数量。
评估协议我们的实验聚焦在图片分类,为了在不同数据集大小上评估模型的性能, 我们利用了三个越来越大的数据集,即ImageNet-1K (128万个图片), ImageNet-21K (1270万个图片) and JFT (3亿个图片)。按照以前的工作,我们首先在三个数据集上使用224的像素分别预训练 300, 90和14轮,然后,我们在ImageNet-1K上以期望的分辨率对预训练模型进行了30轮的微调,获得相应的评估精度。一个例外是ImageNet-1K在分辨率为224时可以直接获取预训练结果。注意,与使用Transformer块的其他模型类似,在不进行微调的情况下,以更高的分辨率直接在ImageNet-1K上评估预训练模型通常会导致性能下降。因此,不管输入分辨率如何变化,总是要使用微调的。
数据增强和正则化在本篇论文中,我们只考虑了两个广泛使用的数据增强方法,即随机数据增强和混合数据增强,和三个通用技术包括:随机深度,标签光滑,权重衰退,来正则化模型,直观上来说,明确的数据增强和正则化方法的超参数依赖于模型的大小和数据规模,其中强正则化通常应用在较大的模型和较小的数据集。
常规情况下,当前范式下的一个复杂问题是如何随着数据大小的变化调整预训练和微调的正则化。特别的,我们有一个有意思的观察,在预训练期间如果某一类型的增强是完全的被禁用,在微调过程中简单地打开它很可能会损害性能,而不是提高性能。我们推测这可能与数据分布转移有关。最终,在所提模型的某些运行时,我们故意在俩个较大的ImageNet21-K 和 JFT数据集上使用一个小幅度的随机增强和随机深度增强进行预训练。尽管这样的正则化可能会损害预训练指标,但这允许在微调期间使用更多的正则化和数据增强,进而提高了下游性能。
Table 4:在 ImageNet. 1K上的模型性能,注意训练也只在 ImageNet-1K上;21K+1K意思是训练在 ImageNet-21K,在 ImageNet-1K上微调; PT-RA意思是在21K上预训练时用随机增强,E150意思是在21K上预训练150轮,这比标准的90轮要长。更多结果看看附录A.3。
Figure 2:只设置成224x224在 ImageNet-1K上的精度-浮点量比例曲线。
ImageNet-1K 表四中展示的是只在 ImageNet-1K数据集上的实验结果。在相同条件下,我们所提的CoAtNet模型不仅胜过了ViT的变种,而且能匹敌最好的只基于卷积架构,例如EfficientNet-V2 和 NFNets。除此之外,我们在表2上也将在224x224分辨率下的所有结果可视化了。正如我们所看到的那样, 带有注意力模块的CoAtNet比以前的模型好得多。
Figure 3: 在ImageNet-21K上预训练然后迁移到ImageNet-1K上微调的精度-参数量的比例曲线。
ImageNet-21K 正如我们从表4和图三中看的那样,当在ImageNet-21K上进行预训练时。 CoAtNet的优势变得十分明显,基本上胜过了以前的所有的模型。值得注意的是,最好的CoAtNet的变体达到了88.56%的 top-1精度。与ViT-H/14的 88.55%相匹敌,这要求比 CoAtNet大2.3倍的ViT-H/14在比ImageNet-21K大23倍的若标签数据集JFT上预训练比在ImageNet-21K上的2.2倍还多的轮次。这标志着一个戏剧性的提升在数据高效性和计算高效性上。
Table 5: 大规模JFT数据集上的性能对比,TPUv3-core-days表示预训练时间,Top-1 Accuracy表示在 ImageNets上的微调精度。注意:最后三行在一个更大的数据集JFT-3B上进行预训练,其他的是在 JFT-300M上。CoAtNet-5/6/7的大小细节看附录A.2。在MBConv块中的下采样是通过步长为2的深度卷积完成的。ViT-G/14的计算消耗请看 [26]这篇论文的图一。
JFT 最后,在表5中,在大规模数据集JFT-300M 和 JFT-3B上进一步评估了CoAtNet。令人鼓舞人心的是,我们的CoAtNet-4通过设置 NFNet-F4+在JFT-300M上能匹敌以前最好的性能而且效率就在TPU上的训练时间和参数量而言是其两倍。当我们放大模型到相同的训练资源上如NFNet-F4+, CoAtNet-5达到了89.77%的 top-1精度,在可比较的设置上胜过了以前的结果。
而且,随着我们进一步将训练资源放到 ViT-G/14使用的水平,使用一个更大的 JFT-3B数据集,超过四倍小的计算量, CoAtNet-6是能与ViT-G/14的90.45%相匹配的,具有1.5倍少的计算量,CoAtNet-7达到了89.77% top-1精度为 90.88%,达到了新的目前最好的性能。
在这部分,我们将对我们设计的CoAtNet做消融实验。
Table 6: 相对注意力上的消融实验
首先,我们通过将卷积和注意力整合到一个计算单元中来研究相对注意力的重要性。特别的,我们对比两个模型,一个具有相对注意力,另一个没有,两个都在单独在ImageNet-1K和 ImageNet-21K上做迁移设置。正如我们从表六上看到的那样,当只是用ImageNet-1K时,相对注意力很明显胜过了标准注意力,具有更好的泛化能力。除此之外,在 ImageNet-21K迁移设置中,相对注意力变体达到了一个更好的迁移精度,尽管他们的预训练性能接近。这表明相对注意力在视觉处理中的主要优势是不仅具有更高的性能,而且具有更好的泛化能力。
Table 7: 在架构布局上的消融实验
第二,因为带有MBConv块的 S2和具有相对Transformer块的 S3占了CoAtNet的主要计算。一个问题是如何切分在S2 (MBConv)和and S3 (Transformer)的计算量来达到一个好的计算性能。在实践中,它由在每个阶段中块的个数决定,我们称之为“布局”设计。对于这个问题,我们对比了一些不同的布局。实验结果在表七。
最后,我们研究了两个所选的模型的细节,即在 MBConv块中的每个注意力头的维度和归一化类型。从表8中我们可以看出随着注意力头的数量从32增加到64稍微的有损性能,尽管它实际上大大提高了TPU的速度。实际上,这将是一种质量-速度权衡。另一方面,批量归一化 和 层归一化几乎拥有相同的性能,而TPU上的BatchNorm速度要快10-20%,这取决于每个核心的批量大小。
在本篇论文中,我们系统的研究了卷积和Transformers的属性,这就产生了一个准则,将它们组合成一个名为CoAtNet的新模型家族。拓展实验表明CoAtNet拥有像ConvNets这样好的泛化能力和像Transformers这样的有监督模型性能,在不同的数据大小和计算成本下达到了目前最好的性能。
注意,本文目前主要关注于模型开发的ImageNet分类。然而,我们相信我们的方法适用于更广泛的应用,如目标检测和语义分割。我们将把它们留给以后的工作。