[论文解读] CoAtNet: Marrying Convolution and Attention for All Data Sizes

Zihang Dai, Hanxiao Liu+2 more|arXiv (Cornell University)|Jun 9, 2021

Advanced Neural Network Applications参考文献 49被引用 735

一句话总结

CoAtNet 将深度卷积与相对自注意力结合在同一计算单元中，采用卷积与变换器块的分层堆叠以在不同数据条件下实现强泛化与高容量，达到在多种数据预算下的 ImageNet 最新结果。

ABSTRACT

Transformers have attracted increasing interests in computer vision, but they still fall behind state-of-the-art convolutional networks. In this work, we show that while Transformers tend to have larger model capacity, their generalization can be worse than convolutional networks due to the lack of the right inductive bias. To effectively combine the strengths from both architectures, we present CoAtNets(pronounced "coat" nets), a family of hybrid models built from two key insights: (1) depthwise Convolution and self-Attention can be naturally unified via simple relative attention; (2) vertically stacking convolution layers and attention layers in a principled way is surprisingly effective in improving generalization, capacity and efficiency. Experiments show that our CoAtNets achieve state-of-the-art performance under different resource constraints across various datasets: Without extra data, CoAtNet achieves 86.0% ImageNet top-1 accuracy; When pre-trained with 13M images from ImageNet-21K, our CoAtNet achieves 88.56% top-1 accuracy, matching ViT-huge pre-trained with 300M images from JFT-300M while using 23x less data; Notably, when we further scale up CoAtNet with JFT-3B, it achieves 90.88% top-1 accuracy on ImageNet, establishing a new state-of-the-art result.

研究动机与目标

• 说明将卷积神经网络与 Transformer 结合的原因，以同时利用归纳偏置和高容量。
• 提出一个统一块，将深度卷积与相对自注意力合并在一个 Transformer 风格的块中。
• 探索垂直网络布局，在卷积阶段与注意力阶段之间取得平衡，以实现泛化与容量的折中。
• 展示在不同数据域（1K、21K 和大规模 JFT）上的最先进性能。
• 提供消融分析以支持设计选择与布局决策。

提出的方法

• 采用一个混合块，在同一个 Transformer 风格块中将深度卷积与相对自注意力合并（预归一化的相对注意力，带全局/有噪声的静态核）。
• 用一个标量相对核 w 表示所有位置对 i-j，以实现高效全局受 receptive 字段且无需额外参数。
• 比较具有不同垂直布局的网络设计（卷积为主的阶段 vs 变换器为主的阶段），研究泛化与容量的权衡。
• 使用多阶段布局（S0 stem、S1 MBConv、S2 MBConv、S3 Transformer Rel、S4 Transformer Rel），在早期捕捉局部模式，在后期捕捉全局上下文。
• 评估大规模预训练（ImageNet-21K、JFT-300M/3B）并进行微调，以展示相对于 ViT 与 ConvNets 的数据与计算效率优势。

实验结果

研究问题

• RQ1 在不产生过高成本的情况下，如何在单一计算单元中融合卷积与自注意力？
• RQ2 在有限数据下哪些卷积与注意力块的垂直布局能获得最佳泛化，在充足数据下又能获得最佳容量？
• RQ3 相对注意力是否在泛化与迁移性方面优于视觉任务中的标准注意力？
• RQ4 混合的 ConvNet-Transformer 模型是否能够在 ImageNet-1K、ImageNet-21K 与 JFT 数据域上达到最先进的性能？

主要发现

• 预归一化的相对注意力块有效地将深度卷积和自注意力统一在一个计算单元中，保留两者的优点。
• 垂直堆叠中卷积阶段在 Transformer 阶段之前（如 C-C-T-T）能带来更好的泛化且容量具有竞争力，相较于更偏向 Transformer 的布局。
• 相对注意力在泛化方面优于标准注意力，尤其在迁移场景（ImageNet-21K 到 ImageNet-1K）。
• CoAtNet 在仅有 1K 数据的情况下，ImageNet-1K 的 top-1 达到 86.0%，在预训练于 ImageNet-21K 再微调后达到 88.56%，相当于在更大数据规模上训练的 ViT-Huge 的表现。
• 在大规模预训练（JFT-3B）下，CoAtNet 在 ImageNet 上达到 90.88% 的 top-1 精度，在可比计算下创下新的状态最优。
• 消融分析表明更多的卷积阶段能够提升泛化，且最佳布局在 S2 MBConv 与 S3 Transformer 块之间达到对迁移性和效率的平衡。