图像分类论文翻译——CoAtNet：结合卷积和注意力 适用于所有数据大小

Posted 2021-07-30 AI浩

 

摘要

Transformers 在计算机视觉方面吸引了越来越多的兴趣，但它们仍然落后于最先进的卷积网络。在这项工作中，我们表明虽然 Transformer 往往具有更大的模型容量，但由于缺乏正确的归纳偏差，它们的泛化可能比卷积网络更差。为了有效地结合两种架构的优势，我们提出了 CoAtNets（发音为“coat”nets），这是一个基于两个关键见解构建的混合模型系列： (1)depthwise Convolution和self-Attention可以通过简单的相对注意力自然地统一起来； (2) 以有原则的方式垂直堆叠卷积层和注意力层在提高泛化、容量和效率方面非常有效。实验表明，我们的 CoAtNets 在各种数据集的不同资源限制下实现了最先进的性能。例如，CoAtNet 在没有额外数据的情况下实现了 86.0% 的 ImageNet top-1 准确率，在额外的 JFT 数据下达到了 89.77%，优于卷积网络和 Transformer 的现有技术。值得注意的是，当使用来自 ImageNet-21K 的 1300 万张图像进行预训练时，我们的 CoAtNet 达到了 88.56% 的 top-1 准确率，与使用来自 JFT 的 300M 图像预训练的 ViT-huge 相匹配，同时使用的数据减少了 23 倍。

1 简介

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

虽然 ViT 用巨大的 JFT 300M 训练图像显示了令人印象深刻的结果，但它的性能在低数据情况下仍然落后于 ConvNets。例如，在没有额外的 JFT-300M 预训练的情况下，ViT 的 ImageNet 精度仍然明显低于具有可比模型大小的 ConvNets [5]（见表 12）。随后的工作使用特殊的正则化和更强的数据增强来改进 vanilla ViT [16, 17, 18]，但鉴于相同数量的数据和计算，这些 ViT 变体都无法在 ImageNet 分类上优于仅 SOTA 卷积模型 [19， 20]。这表明普通的 Transformer 层可能缺乏 ConvNets 拥有的某些理想的归纳偏差，因此需要大量的数据和计算资源来补偿。毫不奇怪，最近的许多工作一直试图将 ConvNets 的归纳偏差纳入 Transformer 模型，通过为注意力层施加局部感受野 [21, 22] 或使用隐式或显式卷积操作来增强注意力和 FFN 层 [23, 24] , 25]。然而，这些方法要么是临时的，要么专注于注入特定的属性，缺乏对卷积和 结合时注意。

在这项工作中，我们从机器学习的两个基本方面——泛化和模型容量，系统地研究了卷积和注意力的混合问题。我们的研究表明，卷积层往往具有更好的泛化能力和更快的收敛速度，这要归功于它们强大的归纳偏差先验，而注意力层具有更高的模型容量，可以从更大的数据集中受益。结合卷积层和注意力层可以实现更好的泛化和容量；然而，这里的一个关键挑战是如何有效地组合它们以在准确性和效率之间实现更好的权衡。在本文中，我们研究了两个关键见解：首先，我们观察到常用的深度卷积可以有效地合并到具有简单相对注意力的注意力层中；其次，以适当的方式简单地堆叠卷积层和注意力层对于实现更好的泛化和容量可能会非常有效。基于这些见解，我们提出了一种名为 CoAtNet 的简单而有效的网络架构，它兼有 ConvNets 和 Transformers 的优势。

我们的 CoAtNet 在不同数据大小的可比资源限制下实现了 SOTA 性能。 具体来说，在低数据机制下，由于有利的归纳偏差，CoAtNet 继承了 ConvNets 的强大泛化特性。 此外，在大量数据的情况下，CoAtNet 不仅享有 Transformer 模型优越的可扩展性，还能实现更快的收敛，从而提高效率。 当仅使用 ImageNet-1K 进行训练时，CoAtNet 达到了 86.0% 的 top-1 准确率，匹配了 ConvNet 变体 NFNet 设置的最佳公共记录集。 此外，当在 ImageNet-21K 上用大约 1000 万张图像进行预训练时，CoAtNet 在 ImageNet-1K 上进行微调时达到 88.56% 的 top-1 准确率，与在 JFT-300M（一个 23 倍大的数据集）上预训练的 ViT-Huge 相匹配。 最后，当使用 JFT 进行预训练时，与 ViT 相比，CoAtNet 表现出更好的效率，以更少的计算量将 ImageNet-1K top-1 精度提高到 89.77%。

2 模型

在本节中，我们重点讨论如何“最优”结合卷积和变换器的问题。 粗略地说，我们将问题分解为两部分： 1. 如何在一个基本计算块内结合卷积和自注意力？ 2、如何将不同类型的计算块垂直堆叠在一起，形成一个完整的网络？ 随着我们逐渐揭示我们的设计选择，分解的基本原理将变得更加清晰。

2.1 合并卷积和自注意力

对于卷积，我们主要关注 MBConv 块 [26]，它采用深度卷积 [27] 来捕获空间交互。 这种选择的一个关键原因是 Transformer 和 MBConv 中的 FFN 模块都采用了“反向瓶颈”的设计，首先将输入的通道大小扩展了 4 倍，然后将 4 倍宽的隐藏状态投影回原始状态 通道大小以启用残差连接。 除了倒置瓶颈的相似性之外，我们还注意到深度卷积和自注意力都可以表示为预定义感受野中值的加权和。 具体来说，卷积依赖于一个固定的内核来从局部感受野收集信息

其中 ， 分别是位置 i 的输入和输出，(i) 表示 i 的局部邻域，例如图像处理中以 i 为中心的 3x3 网格。相比之下，self-attention 允许感受野是整个空间位置，并根据对 (; xj) 之间重新归一化的成对相似度计算权重：

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

       • 首先，depthwise 卷积核  是一个静态值的输入独立参数，而注意力权重 动态地取决于输入的表示。因此，它是 自注意力更容易捕捉不同空间位置之间复杂的关系交互，这是我们在处理高级概念时最想要的属性。然而，灵活性伴随着更容易过度拟合的风险，尤其是在数据有限的情况下。

        • 其次，注意给定任何位置对(i; j)，对应的卷积权重只关心它们之间的相对位移，即i-j，而不是i 或j 的具体值。这个属性通常被称为翻译等效性，已经发现它可以提高有限大小数据集下的泛化能力 [28]。由于使用绝对位置嵌入，标准 Transformer (ViT) 缺少此属性。这部分解释了为什么当数据集不是很大时，ConvNets 通常比 Transformers 更好。

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

 鉴于上述比较，理想模型应该能够结合表 1 中的 3 个理想属性。 (1) 和 Eqn 中的 self-attention。 （2），可以实现这一点的一个简单的想法是简单地将全局静态卷积核与自适应注意矩阵相加，无论是在 Softmax 归一化之后还是之前，即，

有趣的是，虽然这个想法似乎过于简化，但预规范化版本  对应于相对自我注意的特定变体 [29, 30]。 在这种情况下，注意力权重  由平移等方差的和输入自适应 x 共同决定，根据它们的相对大小可以同时享受这两种效果。 重要的是，请注意，为了在不增加参数数量的情况下启用全局卷积核，我们将  的符号重新加载为标量（即 w）而不是方程中的向量。 (1). w 的标量公式的另一个优点是，为所有 (i; j) 检索 显然是通过计算成对点积注意力来包含的，因此导致最小的额外成本（参见附录 A.1）。 考虑到好处，我们将使用 Transformer 块和 Eqn 中的预归一化相对注意变量。 (3) 作为所提出的 CoAtNet 模型的关键组件。

2.2 纵向布局设计

在找出将卷积和注意力结合起来的巧妙方法之后，我们接下来考虑如何利用它来堆叠整个网络。 正如我们上面讨论的，全局上下文具有二次复杂性 w.r.t.空间大小。因此，如果我们直接应用方程中的相对注意力。对于原始图像输入，由于在任何常见尺寸的图像中都有大量像素，因此计算会过慢。因此，要构建一个在实践中可行的网络，我们主要有三种选择：

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

(B) 强制局部注意力，就像在卷积中一样，将全局感受野 G 限制在局部场 上 [22, 21]。

(C) 用某些线性注意力变量替换二次 Softmax 注意力，该变量只有线性复杂度 w.r.t.空间大小 [12, 31, 32]。

我们对选项 (C) 进行了简短的试验，但没有得到相当好的结果。对于选项 (B)，我们发现实现局部注意力涉及许多需要密集内存访问的非平凡形状格式化操作。在我们选择的加速器（TPU）上，这种操作被证明是极其缓慢的[33]，这不仅违背了加速全局注意力的初衷，而且损害了模型容量。因此，由于最近的一些工作已经研究了这种变体 [22, 21]，我们将重点关注选项 (A)，并将我们的结果与我们的实证研究（第 4 节）中的结果进行比较。 对于选项 (A)，下采样可以通过 (1) 具有侵略性的卷积茎来实现 步幅（例如，步幅 16x16）如 ViT 或（2）多阶段网络，如 ConvNets 中的渐进池化。通过这些选择，我们推导出 5 个变体的搜索空间，并在对照实验中对它们进行比较。

• 当使用 ViT Stem 时，我们直接将 L 个 Transformer 块相对注意力堆叠起来，这 我们表示为 VITREL。

• 当使用多阶段布局时，我们模仿 ConvNets 构建 5 个阶段的网络（S0、 S1、S2、S3 和 S4)，空间分辨率从 S0 到 S4 逐渐降低。在每个阶段的开始，我们总是将空间大小减少 2 倍并增加通道数。第一阶段 S0 是一个简单的 2 层卷积 Stem，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 个时期的混合模型的不同变体，两者都没有任何正则化或增强。图 1 总结了两个数据集上的训练损失和评估准确度。

• 从 ImageNet-1K 结果来看，一个关键观察是，就泛化能力（即训练和评估指标之间的差距）而言，我们有 C-C-C-C ≈ C-C-C-T ≥ C-C-T-T > C-T-T-T >>：

特别是，VITREL 明显比其他变体差很多，我们推测 与在其激进的下采样 Stem 中缺乏适当的低级信息处理有关。 在多阶段变体中，总体趋势是模型的卷积阶段越多，泛化差距越小。

• 至于模型容量，从 JFT 比较来看，训练结束时的训练和评估指标都表明以下排名：

C-C-T-T ≈ C-T-T-T >  > C-C-C-T > C-C-C-C：

 重要的是，这表明仅仅拥有更多的 Transformer 块并不一定意味着更高的视觉处理能力。 一方面，虽然最初更糟，但  最终赶上了两个具有更多 MBConv 阶段的变体，表明 Transformer 块的容量优势。 另一方面，C-C-T-T 和 C-T-T-T 明显优于 ，这表明具有激进步幅的 ViT 词干可能丢失了太多信息，因此限制了模型容量。 更有趣的是，C-C-T-T ≈ C-T-T-T 的事实表明，为了处理低级信息，像卷积这样的静态局部操作可以与自适应全局注意力机制一样强大，同时大大节省计算和内存使用。

最后，为了在 C-C-T-T 和 C-T-T-T 之间做出决定，我们进行了另一个可转移性测试3——我们在 ImageNet-1K 上对上述两个 JFT 预训练模型进行了 30 次微调，并比较了它们的转移性能。 从表 2 中可以看出，尽管具有相同的预训练性能，但 C-C-T-T 的传输精度明显优于 C-T-T-T。

考虑到泛化性、模型容量、可转移性和效率，我们为 CoAtNet 调整了 C-C-T-T 多阶段布局。 更多模型细节包含在附录 A.1 中。

3 相关工作

卷积网络构建块。 卷积网络 (ConvNets) 一直是许多计算机视觉任务的主要神经架构。 传统上，常规卷积，例如 ResNet 块 [3]，在大规模 ConvNet 中很流行； 相比之下，深度卷积 [27] 因其较低的计算成本和较小的参数大小而在移动平台中很受欢迎 [26]。 最近的工作表明，基于深度卷积的改进的反向残差瓶颈（MBConv [26, 34]）可以实现高精度和更好的效率 [5, 19]。 如第 2 节所述，由于 MBConv 和 Transformer blocks 之间的强连接，本文主要采用 MBConv 作为卷积构建块。

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

Relative attention.。在相对关注的总称下，出现了各种变种 在文献中 [29, 37, 38, 33, 39, 30]。一般来说，我们可以将它们分为两类： (a) 依赖于输入的版本，其中额外的相对注意力分数是输入状态 f(xi; xj; i − j) 的函数，以及 (b) 独立于输入的版本 f(i − j)。 CoAtNet 中的变体属于与输入无关的版本，类似于 T5 [30] 中使用的变体，但与 T5 不同的是，我们既不共享跨层的相对注意力参数，也不使用分桶机制。作为输入独立性的一个好处，为所有 (i; j) 对获得 f(i − j) 在计算上比在 TPU 上依赖于输入的版本便宜得多。此外，在推理时，这只需要计算一次并缓存以备将来使用。最近的一项工作 [22] 也利用了这种与输入无关的参数化，但它将感受野限制为局部窗口。 结合卷积和自注意力。将卷积和自注意力相结合进行视觉识别的想法并不新鲜。一种常见的方法是使用显式自注意力或非局部模块 [9, 10, 11, 12] 来增强 ConvNet 主干，或者用标准自注意力 [11] 或更灵活的线性组合替换某些卷积层注意和卷积 [40]。虽然自注意力通常会提高准确性，但它们通常会带来额外的计算成本，因此通常被视为 ConvNet 的附加组件，类似于挤压和激励 [41] 模块。相比之下，在 ViT 和 ResNet-ViT [13] 取得成功之后，另一个流行的研究路线从 Transformer 主干开始，并试图将显式卷积或卷积的一些理想特性合并到 Transformer 主干中 [25, 24, 23, 22 , 21, 42, 43]。

      虽然我们的工作也属于这一类，但我们表明我们的相对注意力实例是深度卷积和基于内容的注意力的自然混合，附加成本最低。 更重要的是，从泛化和模型容量的角度出发，我们采用系统的方法进行垂直布局设计，并展示不同的网络阶段如何以及为什么喜欢不同类型的层。 因此，与简单地使用现成的 ConvNet 作为茎层的模型相比，例如 ResNet-ViT [13]，CoAtNet 在整体尺寸增加时也会缩放卷积阶段（S2）。 另一方面，与采用局部注意力[22, 21] 的模型相比，CoAtNet 始终对 S3 和 S4 使用完全注意力来确保模型容量，因为 S3 占据了大部分计算和参数。

4 实验

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

4.1 实验设置

        CoAtNet 模型族。为了与现有不同大小的模型进行比较，我们还设计了一系列 CoAtNet 模型，如表 3 所示。总体而言，我们始终将 S1 到 S4 的通道数量加倍，同时确保 Stem S0 的宽度更小或相等到S1。 此外，为简单起见，当增加网络深度时，我们只缩放 S2 和 S3 中的块数。

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

数据增强和正则化。在这项工作中，我们只考虑两种广泛使用的数据增强，即 RandAugment [44] 和 MixUp [45]，以及三种常用技术，包括随机深度 [46]、标签平滑 [47] 和权重衰减 [48]，来正则化该模型。直观上，增强和正则化方法的特定超参数取决于模型大小和数据规模，其中强正则化通常应用于较大的模型和较小的数据集。 在一般原则下，当前范式下的一个复杂问题是如何随着数据大小的变化调整预训练和微调的正则化。具体来说，我们有一个有趣的观察结果，如果某种类型的增强在预训练期间完全被禁用，那么在微调期间简单地打开它很可能会损害性能而不是改进。 我们推测这可能与数据分布偏移有关。因此，对于所提出模型的某些运行，我们在两个较大的数据集 ImageNet21-K 和 JFT 上进行预训练时，特意应用了 RandAugment 和小程度的随机深度。尽管这种正则化可能会损害预训练指标，但这允许在微调期间进行更通用的正则化和增强，从而提高下游性能。

4.2 主要结果

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

 

 

 ImageNet-21K 从表 4 和图 3 中我们可以看出，当使用 ImageNet-21K 进行预训练时，CoAtNet 的优势变得更加明显，明显优于之前所有的模型。 值得注意的是，最好的 CoAtNet 变体实现了 88.56% 的 top-1 准确率，与 88.55% 的 ViTH/14 性能相匹配，这需要在 23 倍大的专有弱标记数据集 (JFT) 上对 2.3 倍大的 ViT 模型进行预训练 2.2 x 更多步骤。 这标志着数据效率和计算效率的显着提高。 JFT 最后，在表 5 中，我们使用 JFT 在大规模数据机制下进一步评估了 CoAtNet。 令人鼓舞的是，我们的 CoAtNet-4 几乎可以与 NFNet-F4+ 设置的 JFT 相媲美，同时在 TPU 训练时间和参数数量方面效率提高 2 倍。 当我们扩大模型以消耗与 NFNet-F4+ 类似的训练资源时，CoAtNet 在 top-1 准确率上达到 89.77%，在可比设置下优于之前的结果。

4.3 消融研究

    在本节中，我们将消除我们对 CoAtNet 的设计选择。

    首先，我们研究了将卷积和注意力结合到单个计算单元中的相对注意力的重要性。 具体来说，我们比较了两种模型，一种具有相对注意力 另一个没有，在单独的 ImageNet-1K 和 ImageNet-21K 传输设置下。 从表 6 可以看出，当仅使用 ImageNet-1K 时，相对注意力明显优于标准注意力，表明具有更好的泛化能力。 此外，在 ImageNet-21K 传输设置下，相对注意力变量实现了明显更好的传输精度，尽管它们 非常接近的训练前表现。 这表明视觉处理中相对注意力的主要优势不是更高的容量，而是更好的概括。

 

其次，由于带有 MBConv 块的 S2 和带有相关 Transformer 块的 S3 占据了 CoAtNet 的大部分计算，所以要问的问题是如何拆分 S2（MBConv）和 S3（Transformer）之间的计算以获得良好的性能。 在实践中，它归结为决定每个阶段的块数，我们将其称为“布局”设计。 为此，我们比较了我们在表 7 中试验过的几种不同布局。

     • 如果我们保持 S2 和 S3 中的区块总数固定并改变每个阶段的数量，我们会观察到 V0 是 V1 和 V2 之间的最佳点。 基本上，在 S3 中拥有更多 Transformer 块通常会带来更好的性能，直到 S2 中 MBConv 块的数量太小而无法很好地泛化。

    • 为了进一步评估最佳点是否也适用于转移设置，其中更高的容量通常被认为更重要，我们进一步比较了在 ImageNet-21K 转移到 ImageNet-1K 设置下的 V0 和 V1。 有趣的是，尽管 V1 在 ImageNet-21K 预训练期间的性能略好于 V0，但 V1 的传输精度明显落后于 V0。 同样，这表明卷积在实现良好泛化方面的重要性。  

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

5 结论

         在本文中，我们系统地研究了卷积和 Transformer 的特性，这导致了一种将它们组合成名为 CoAtNet 的新模型系列的原则方法。 大量实验表明，CoAtNet 兼具 ConvNets 等良好的泛化能力和 Transformers 等卓越的模型容量，在不同的数据大小和计算预算下实现了最先进的性能。 请注意，本文目前专注于用于模型开发的 ImageNet 分类。 然而，我们相信我们的方法适用于更广泛的应用，如对象检测和语义分割。 我们将把它们留给以后的工作。