CoAtNet Marrying Convolution and Attention for All Data Sizes 论文学习

Abstract

Transformers 在计算机视觉任务上得到了越来越多的关注，但仍落后于 SOTA 卷积网络。本文作者证明，尽管 Transformers 具有更强的模型能力，但因为缺乏归纳偏置特性，它的泛化性要落后于 CNN。为了有效地结合二者的长处，作者提出了 CoAtNets，它的构建主要基于两个关键想法：（1）我们可以通过简单的相对注意力将深度卷积和自注意力自然统一；（2）在提升泛化性、能力和效率方面，按照一定的原则，垂直摆放卷积层和注意力层会非常有效。实验证明，CoAtNets 在多个数据集上，根据不同的资源要求，可以取得 SOTA 的效果。例如，CoAtNet 在 ImageNet 上取得了 $86.0\%$ top-1 准确率，无需额外的数据，如果使用了 JFT 数据，则可达到 $89.77\%$，超越目前所有的 CNN 和 Transformers。值得注意的是，当我们用ImageNet-21K 的 1300 万张图像来预训练时，CoAtNet 得到了$88.56\%$的 top-1 准确率，可媲美 ViT-huge 预训练模型，而它用了3亿张图像。

1. Introduction

自从 AlexNet，CNN 已经成为了计算机视觉领域最主要的模型结构。同时，随着自注意力模型如 Transformers 在 NLP 领域取得成功，许多工作都尝试将注意力机制引入计算机视觉。最近，ViT 证明我们只需标准的 Transformer 层，就能在 ImageNet-1K 上取得不错的效果。而且，当我们在弱标注的 JFT-300M 数据集上进行预训练时，ViT 能取得和 SOTA 卷积网络相匹敌的结果，表明 Transformers 模型可以比卷积网络的能力更强。

尽管 ViT 在 JFT-300M 训练数据上表现非凡，但在数据较少的数据集上其表现就会差于 CNN。例如，如果不用 JFT-300M 数据集做预训练，ViT 在 ImageNet 上的准确率要明显低于同等大小的卷积网络。此后的一些工作试图通过特殊的正则化和更强的数据增广来优化 ViT，但是在给定数据和计算量的前提下，它们都没能超过 SOTA 卷积模型。这表明 Transformer 层可能缺乏归纳偏置特性，而它存在于卷积网络中，因此需要巨量的数据和计算资源来弥补该不足。最近一些工作尝试将卷积网络的归纳偏置融合到 Transformer 模型中，在注意力层中强加上局部感受野，或通过隐式或显式的卷积操作来增强注意力和 FFN 层。但是，这些方法都关注在某个特性的注入上，缺乏对卷积和注意力结合时各自角色的理解。

本文作者从机器学习的两个基本面，系统地研究了卷积和注意力融合的问题 — 泛化性和模型能力。研究显示得益于归纳偏置先验，卷积网络的泛化性更强，收敛速度更快，而注意力层则拥有更强的模型能力，对大数据集更有益。将卷积和注意力层结合，可以取得更好的泛化性和模型能力；但是，问题是如何将二者有效结合，实现准确率和效率的平衡。本文作者探讨了两点：首先，作者发现利用简单的相对注意力，深度卷积可以有效地融合入注意力层。其次，以适当的方式直接堆叠卷积层和注意力层，效果可能惊人地好，泛化性和能力都更高。于是乎，作者提出了一个简单而有效的网络结构，叫做 CoAtNet，具备卷积网络和 Transformers 的优势。

在资源要求一样的前提下，CoAtNet 对于不同的数据集大小都取得了 SOTA 表现。对于规模不大的数据集，CoAtNets 继承了卷积网络中极强的泛化性。而当数据充足时，CoAtNet 不仅具备 Transformer 模型的优越能力，也可实现更快的收敛速度，提高效率。当我们只用 ImageNet-1K 训练时，CoAtNet 取得了 $86.0\%$的 top-1 准确率，可与 NFNet（卷积网络的变体） 取得的最佳记录相媲美。当我们用 ImageNet-21K 训练时（差不多有1000万张图片），在 ImageNet-1K 基础上微调，CoAtNet 取得了 $88.56\%$ 的 top-1 准确率，可与 JFT-300M 上预训练的 ViT-Huge 模型相近，而该数据集的规模要大23倍。最后，当我们用 JFT 预训练，CoAtNet 展现了更高的效率，将 ImageNet-1K 上的 top-1 准确率提升到了 $89.77\%$，而计算量更低。

2. Model

本章节，作者重点探讨如何最完美地将卷积和 transformer 结合。作者将该问题拆分为下面两部分：

1. 如何将卷积和自注意力结合在一个计算模块内？
2. 如何将不同类型的计算模块垂直堆叠在一起，形成一个完整的网络？

2.1 Merging Convolution and Self-Attention

对于卷积，作者主要关注在 MBConv 模块，它利用深度卷积来获取空间联系。之所以这么做，是因为 Transformers 的 FFN 模块和 MBConv 都使用了 “inverted bottleneck” 设计，首先将输入的通道大小扩展4倍，然后将之映射回原来的通道大小，这样可以使用残差连接。

除了都使用了 inverted bottleneck，作者也注意到，深度卷积和自注意力都可以表示为一个预先定义好的感受野内的数值加权和。卷积需要一个固定的卷积核，从局部感受野中获取信息：

$$y_i=\sum _{j\in \mathcal{L}(i)}{w}_{i-j}\odot x_j\text{(depthwise convolution)}\text{\hspace{1em}(1)}$$

其中$x_i,y_i\in {\mathbb{R}}^D$分别是$i$位置的输入和输出，$\mathcal{L}(i)$表示$i$的相邻区域，即中心点为$i$的$3\times 3$网格。

为了比较，自注意力的感受野涵盖所有的空间位置，根据一对点$(x_i,x_j)$的归一化后的相似度来计算权值：

$$y_i=\sum _{j\in \mathcal{G}}\underset{A_{i,j}}{\frac{\exp (x_i^T{x}_j)}{{\sum }_{k\in \mathcal{G}}\exp (x_i^T{x}_k)}}{x}_j\text{(self-attention)}\text{\hspace{1em}(2)}$$

其中$\mathcal{G}$表示全局位置空间。在研究如何将它们最优地结合之前，我们应该分析一下它们各自的优劣，找出我们想要保留下来的优秀特性。

• 首先，深度卷积核$w_{i-j}$是一个与输入无关的静态参数，而注意力权重$A_{i,j}$动态地依赖于输入的表征。因此，自注意力可以很容易地获取不同空间位置的相互关系，当我们处理高级别任务时这是非常必须的性质。但是，这种灵活性也带来了容易过拟合的风险，尤其是当数据有限时。
• 其次，给定一对空间点$(i,j)$，相应的卷积权重$w_{i-j}$只关心二者的相对偏移，即$i-j$，而不是$i$或$j$的具体数值。这通常是指平移等价性，可以在有限的数据集上提升泛化性。因为使用了绝对位置 embedding，标准 ViT 缺少该特性。这就部分解释了为何当数据集有限时，卷积网络通常要比 Transformer 的表现好。
• 最后，感受野大小也是自注意力和卷积的最根本区别。一般而言，较大的感受野能提供更多的语义信息，模型能力就更强。因此，人们将自注意力应用在视觉领域的关键原因就是，它能提供全局感受野。但是，较大的感受野需要非常多的计算量。以全局感受野为例，复杂度相对于空间大小是指数的，这就限制了它的应用范围。

根据上面的比较，最理想的模型应该具备表1中的三点特性。根据等式1和等式2中的相似性，最直接的想法就是将全局静态卷积核和自适应的注意力矩阵相加，在 Softmax 归一化之前和之后：

$$y_i^{post}=\sum _{j\in \mathcal{G}}(\frac{\exp (x_i^T{x}_j)}{{\sum }_{k\in \mathcal{G}}\exp (x_i^T{x}_k)}+w_{i-j})x_j\text{or}{y}_i^{pre}=\sum _{j\in \mathcal{G}}\frac{\exp (x_i^T{x}_j+w_{i-j})}{{\sum }_{k\in \mathcal{G}}\exp (x_i^T{x}_k+w_{i-j})}{x}_j.\text{\hspace{1em}(3)}$$

尽管该想法看着很简单，提前归一化的版本$y^{pre}$对应着相对自注意力的特殊版本。我们通过具有平移等价性的$w_{i-j}$和输入自适应的$x_i^T{x}_j$来共同决定注意力权重$A_{i,j}$，根据它们相对的重要性来施加影响。注意，为了确保全局卷积核不会造成参数量爆炸，作者将$w_{i-j}$表示为一个标量（即$w\in {\mathbb{R}}^{O(|\mathcal{G}|)}$），而非等式1中的向量。将$w$表示为标量的另一个好处就是，通过计算 pairwise 点积注意力，其自然包含了所有$(i,j)$的$w_{i-j}$，这样计算开支是最低的（参考 Appendix A.1）。CoAtNet 模型的关键组成包括 Transformer 模块和等式3中的提前归一化了的相对注意力。

2.2 Vertical Layout Design

找到一个简洁地将卷积和注意力结合的方式之后，我们下一步就该思考如何利用它们，堆出一个网络。

如上所述，全局关系相对于空间大小呈指数复杂度。因此，如果我们直接在原始输入图像上使用等式3的相对注意力，计算会非常慢，因为任何一张图片的像素点个数都是非常惊人的。所以，要想构建出一个可行的网络，我们主要有三个选项：

• (A) 先进行下采样降低空间大小，然后当特征图达到可控的范围后，再使用全局相对注意力。
• (B) 强行加上局部注意力，将注意力中的全局感受野$\mathcal{G}$约束在一个局部范围$\mathcal{L}$内，就像卷积一样。
• © 将指数增长的 Softmax 注意力替换为线性注意力变体，它对于空间大小只有线性复杂度。

作者简单地用©进行了实验，却没有得到合理的效果。对于(B)选项，作者发现实现局部注意力会带来很多严重的形状变换操作，造成大量的内存读取。在 TPU 加速器上，该操作会非常慢，不仅无法实现全局注意力的加速，也会伤害模型的能力。所以，作者采用选项(A)，在实验研究（第4部分）中比较了它们的效果。

对于选项(A)，下采样操作可以是(1)一个步长很大（比如步长是$16\times 16$）的卷积操作（如 ViT 中的），或者是(2)一个使用逐步池化的多阶段网络（如卷积网络中的）。基于这些选项，作者设计了包含5个变体的搜索空间，对它们做了实验比较。

• 当我们用 ViT 时，利用相对注意力直接堆叠 $L$ 个 Transformer 模块，记做${\text{VIT}}_{\text{REL}}$。
• 当使用多阶段，我们模拟ConvNets，构建一个有5个阶段的网络（$\text{S0,S1,S2,S3 \& S4}$），从$\text{S0}$到$\text{S4}$其空间分辨率是逐渐下降的。在每个阶段开始位置，空间大小缩小2倍，增加通道数。第一个阶段$\text{S0}$是一个简单的2层CNN，$\text{S1}$使用 MBConv 模块和 squeeze-excitation 操作，其空间大小对于全局注意力来说太大了。从$\text{S2}$到$\text{S4}$，我们可以用 MBConv 或者 Transformer 模块，要求是卷积阶段必须在 Transformer 阶段之前。这个要求是基于一个先验得到的，卷积更擅长处理早期阶段中很常见的局部模式。这样随着 Transformer 阶段的增加，我们就有4种变体，C-C-C-C, C-C-C-T, C-C-T-T 和 C-T-T-T，其中C和T分别表示卷积和Transformer。

为了系统地研究这些设计选项，作者主要从两个基本面出发，即泛化能力和模型能力。对于泛化性，本文研究了训练损失和测试精度之间的误差。如果两个模型有着相同的训练损失，那么测试精度更高的模型就应该有着更强的泛化性，因为它对没见过的数据泛化更好。当数据有限时，泛化能力就非常重要。对于模型能力，作者评估了模型拟合大型数据集上的能力。当训练数据充足时，过拟合就不再是个问题，能力强的模型越有可能取得更好的效果。一般而言，模型越大，能力就越强，所以为了比较公平些，本文保证5个变体模型的大小是差不多的。

为了比较泛化性和模型能力，作者在 ImageNet-1K（130万张图片）和 JFT（$>3$亿张图片）数据集上训练不同模型的版本，分别训了300个 epochs 和3个epochs，都没用正则或增广。图1展示了在这两个数据集上得到的训练损失和测试精度。

• 从 ImageNet 上的结果可以看出，对于泛化能力（即训练和测试的差异）而言，我们有
  $$\text{C-C-C-C}\approx \text{C-C-C-T}\ge \text{C-C-T-T}>\text{C-T-T-T}\gg {\text{VIT}}_{\text{REL}}$$

${\text{VIT}}_{\text{REL}}$要比其它变体差好多，作者猜测这可能是因为在它激进的下采样过程中，缺失了低层级的信息。在多阶段变体中，整体趋势是，模型中卷积越多，泛化性越好。

• 就模型能力而言，作者参考 JFT 比较，从训练和测试角度来看，可得到如下排名：

$$\text{C-C-T-T}\approx \text{C-T-T-T}>{\text{VIT}}_{\text{REL}}>\text{C-C-C-T}>\text{C-C-C-C}$$

图1. 不同数据集规模，模型泛化性和能力的比较。为了公平比较，所有模型的参数量和计算量相近。

这说明，简单地增加 Transformer 模块并不一定会带来更强的模型能力。一方面，尽管初始化的很糟糕，${\text{VIT}}_{\text{REL}}$最终追赶上了另两个带有 MBConv 的变体，表明 Transformer 模块的能力很强。另一方面，$\text{C-C-T-T}$和$\text{C-T-T-T}$明显优于${\text{VIT}}_{\text{REL}}$，表明 ViT 如果步长过大，会丢失很多信息，不利于模型的能力。$\text{C-C-T-T}\approx \text{C-T-T-T}$表明，对于低层级信息而言，静态局部操作如卷积的能力可与自适应全局注意力机制相当，但能大幅度降低计算量和内存占用。

最后，为了在$\text{C-C-T-T}$和$\text{C-T-T-T}$之间做决定，作者进行了可迁移性测试，作者在 ImageNet-1K 上对两个在 JFT 数据集上预训练的模型进行微调，训练了30个epochs，比较它们的迁移效果。从表2可以看到，$\text{C-C-T-T}$的效果要好于$\text{C-T-T-T}$，尽管它们的预训练表现差不多。

根据泛化性、模型能力、可迁移性和效率，作者在 CoAtNet 中采用了$\text{C-C-T-T}$多阶段方案。附录 A.1 中包含了更多细节信息。

3. Related Work

卷积网络模块。CNN是目前计算机视觉任务的主流方法。传统卷积如 ResNet 模块，在大型卷积网络中很流行。而深度卷积则因为计算量和参数量少，在移动平台更受欢迎。最近一些工作表明，基于深度卷积而构建的 improved inverted residual bottlenecks (MBConv) 可以取得更高的准确率和效率。如第二部分所述，因为 MBConv 和 Transformer 模块具有很强关联性，本文用 MBConv 作为CNN构建模块。

自注意力和Transformers。自注意力是 Transformer 的核心组成，在 NLP 和语音理解领域应用范围很广。[33] 表示，自注意力网络在不同的视觉任务上可以表现得不错，尽管有一些实践难度。最近ViT [13] 在 ImageNet 分类任务上使用了一个标准 Transformer，在大规模 JFT 数据集上预训练后取得了惊人的效果。但是在数据有限的时候，ViT 仍然大幅度落后于 CNN。最近的一些工作就关注在如何提升 Transformer 的数据使用效率和模型效率。

相对注意力。一般而言，我们可以将相对注意力划分为两类：(a) 依赖输入的版本，相对注意力得分是关于输入状态的函数，$f(x_i,x_j,i-j)$。(b) 不依赖输入的版本$f(i-j)$。CoAtNet 变体属于第二类，与T5[30]类似，但又不同，本文没有在各层之间共享相对注意力参数，也没有用 bucketing 机制。不依赖输入的一个好处就是，要想获取所有$(i,j)$的$f(i-j)$，其在 TPU 上的计算成本要比依赖输入的低很多。此外，推理时，它只需要计算一次，然后就可缓存起来给以后使用。[22] 也用了不依赖输入的参数化方案，但是它将感受野约束为一个局部窗口。

结合卷积和自注意力。将卷积和自注意力结合的思想并不新。一个常见的办法就是用自注意力或non-local模块来增强卷积主干网络，也有人用标准的自注意力层或更灵活的线性注意力和卷积的混合体来替代部分卷积层。尽管自注意力通常能提升准确率，但它们通常增加计算量，因此经常用作为卷积网络的插件使用，就像 squeeze-and-excitation 模块。为了比较，在 ViT 和 ResNet-ViT 取得成功后，另一条科研路线采用 Transformer 作为主干，然后在里面直接加入卷积或一些卷积特性。

本文属于这一类方法，作者证明，相对注意力是深度卷积和注意力的自然融合，代价是最小的。从泛化性和模型能力角度出发，作者在垂直结构设计上采取系统分析的方法，阐述了如何以及为什么不同的网络阶段应采用不同类别的层。所以，和那些使用现成的卷积网络的方法相比（如 ResNet-ViT），当整体规模增大时，CoAtNet 也可扩大卷积阶段（S2）。另一方面，和利用局部注意力的方法相比，在S3&S4 中使用了全注意力，确保模型的能力，因为 S3 占据了计算量和参数量的大头。

Appendix

A.1 Model Details

首先，CoAtNe 整体如图4所示。

图4：CoAtNet 概览

$$y_i^{pre}=\sum _{j\in \mathcal{G}}\frac{\exp (x_i^T{x}_j+w_{i-j})}{{\sum }_{k\in \mathcal{G}}\exp (x_i^T{x}_k+w_{i-j})}{x}_j.\text{\hspace{1em}(3)}$$

2D 相对注意力。为了在大小是 $[H\times W]$ 的 2D 图像上实现等式3中的相对注意力 pre-norm 版本，对于每个 head，作者都创建了一个可训练的参数 $\mathbf{P}$，大小是$[(2H-1)\times (2W-1)]$，因为最大距离分别是$2H-1$和$2W-1$。然后，对于两个空间点$(i,j)$和$(i^{\prime },j^{\prime })$，在1-based索引下，其对应的相对偏置是 $P_{i-i^{\prime }+H,j-j^{\prime }+W}$。实现时，我们需要从$[(2H-1)\times (2W-1)]$矩阵中对$H^2{W}^2$个元素进行索引。在TPU上，作者使用了2个 einsums 操作，分别沿着x轴和y轴来索引相对偏置，复杂度是$O(HW(H+W))$，这肯定要小于$O(H^2{W}^{2}D)$注意力复杂度。在GPU上，索引可以通过$\text{gather}$操作更高效率地实现，只需内存读取。在推理时，对$[(2H-1)\times (2W-1)]$矩阵中$H^2{W}^2$个元素索引可以提前计算好缓存起来，进一步提升了吞吐量。

作者在更大分辨率的数据集上进行微调，只使用了双线性插值，将$[(2H-1)\times (2W-1)]$大小增加到$[(2H^{\prime }-1)\times (2W^{\prime }-1)]$，其中$H^{\prime }>H,W^{\prime }>W$。

Pre-Activation. 为了促进模型结构的一致性，作者对 MBConv 和 Transformer 模块都使用了 pre-activation 结构，即

$$\mathbf{x}\leftarrow \mathbf{x}+\text{Module}(\text{Norm}(\text{x}))$$

其中$\text{Module}$表示 MBConv、自注意力或 FFN 模块，$\text{Norm}$对应着 MBConv 中的 batch norm，或自注意力和 FFN 中的 layer norm，在取得相同精度同时，极大地降低TPU上的延迟。在 MBConv 和 Transformer 模块中，使用了 Gaussian Error Linear Units(GELUs) 作为激活函数。

下采样. 对于从S1到S4每个阶段中的第一个模块，在残差分支和恒等分支上，下采样是独立进行的。对于 Transformer 模块，将步长为2的最大池化操作，直接用在自注意力模块的每个分支的输入上，与 Funnel Transformer 相似。同样，为了增大其大小，对于恒等分支使用了通道映射。因此，下采样自注意力模块可以表示为：

$$\mathbf{x}\leftarrow \text{Proj}(\text{Pool}(\mathbf{x}))+\text{Attention}(\text{Pool}(\text{Norm}(\text{x})))$$

对于 MBConv 模块，残差分支上的下采样是通过一个步长为2的卷积实现的，作用到归一化输入上，即：

$$\mathbf{x}\leftarrow \text{Proj}(\text{Pool}(\mathbf{x}))+\text{Conv}(\text{DepthConv}(\text{Conv}(\text{Norm}(\text{x}),\text{stride}=2)))$$

这和标准的 MBConv 是不同的，MBConv 的下采样是通过步长为2的 depthwise conv 实现，作用在 inverted bottleneck 隐藏状态上。随后作者发现，使用步长为2的 depthwise conv 是有用的，但是当模型较小时速度会慢一点，而当模型增大时并不明显，如表9所示。因此，如果没有另外说明，余下论文都将使用等式5中的下采样实现。在实际操作中，对于小模型而言，我们可以权衡速度和准确度，再作出选择。

分类head。作者对最后一个stage的输出，采用全局平均池化来得到表征，而没有像 ViT 那样增加一个额外的 \text{} token 来分类。