PaperlayerView exact revision on arXiv

Source-linked AI summary

CoAtNet: Marrying Convolution and Attention for All Data Sizes

Zihang Dai, Hanxiao Liu, Quoc V. Le, Mingxing Tan

arXiv:2106.04803v2cs.CVcs.LG

TL;DR

Vision Transformers can generalize worse than ConvNets in limited-data settings, while attention offers greater model capacity. CoAtNet combines their strengths through relative attention and convolution-first stacking, achieving state-of-the-art ImageNet results across data scales. The paper’s evaluation focuses on ImageNet classification, leaving broader applications for future work.

  • Problem

    Vanilla Transformers can have weaker low-data generalization than ConvNets because they lack desirable convolutional inductive biases.

  • Method

    CoAtNet unifies depthwise convolution and self-attention through relative attention and vertically stacks convolution and Transformer stages.

  • Results

    CoAtNet achieves state-of-the-art performance across data sizes, including 86.0% ImageNet-1K top-1 accuracy without extra data, 88.56% after ImageNet-21K pre-training, and 90.88% with JFT-3B.

  • Takeaways & Limitations

    The hybrid design combines ConvNet-like generalization with Transformer-like model capacity under different data sizes and computation budgets.

  • Takeaways & Limitations

    The paper focuses on ImageNet classification and leaves applications such as object detection and semantic segmentation for future work.

Abstract

from arXiv · show

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.

1 Introduction

CoAtNet addresses the weaker low-data generalization of vanilla Transformers by combining convolutional inductive bias with Transformer capacity. Its hybrid design achieves strong ImageNet performance across data regimes and resource constraints.

  • Motivation: ViT performance falls behind comparable ConvNets in low-data ImageNet settings, despite strong results with enormous JFT-300M pre-training.The paper attributes this gap to missing inductive biases in vanilla Transformer layers.
  • Contribution: CoAtNet studies generalization and model capacity to combine convolution’s faster-converging inductive bias with attention’s scalability on larger datasets.The design merges depthwise convolution and self-attention through relative attention and vertically stacks their layers.
  • Results: 86.0% top-1 accuracy is achieved on ImageNet-1K without extra data, matching NFNet under similar computation and training conditions.
  • Results: 88.56% top-1 accuracy is achieved after ImageNet-21K pre-training, matching ViT-Huge trained on JFT-300M while using a 23× smaller dataset.
  • Results: 90.88% ImageNet-1K top-1 accuracy is achieved with JFT-3B pre-training while using 1.5x less computation than ViT-G/14.

2 Model

The model design unifies convolution and attention within blocks, then arranges them in convolution-first multi-stage networks. Controlled comparisons examine their generalization, capacity, efficiency, and transferability.

  • Design decomposition: The design problem is decomposed into combining convolution and self-attention within a block and vertically stacking blocks into a complete network.
  • Merging Convolution and Self-Attention: Depthwise convolution and self-attention both compute per-dimension weighted sums over receptive fields, differing mainly in fixed local versus adaptive global weighting.Convolution uses a local neighborhood, while self-attention can use the entire spatial field and pairwise similarity.
  • Merging Convolution and Self-Attention: Convolution provides static translation-equivalent weights that support limited-data generalization, whereas attention provides adaptive interactions but risks overfitting.
  • Merging Convolution and Self-Attention: CoAtNet uses pre-normalization relative attention to combine translation-equivalent convolution weights with input-adaptive attention at minimal additional cost.
  • Vertical Layout Design: Global relative attention is applied after spatial down-sampling because its quadratic complexity is excessive on raw image resolutions; local or linear alternatives have practical drawbacks.
  • Vertical Layout Design: The multi-stage search space places convolution stages before Transformer stages, producing C-C-C-C, C-C-C-T, C-C-T-T, and C-T-T-T variants.Stages progressively halve spatial resolution and increase channels, with MBConv used early and either MBConv or Transformer blocks later.
  • Vertical Layout Design: More convolution stages generally reduce the ImageNet generalization gap, while JFT capacity ranks C-C-T-T ≈C-T-T-T > VITREL > C-C-C-T > C-C-C-C.
  • Vertical Layout Design: C-C-T-T transfers better than C-T-T-T despite identical pre-training performance, supporting the selected convolution-first layout.

3 Related Work

Prior vision work explored convolutional and attention-based architectures, including relative attention and hybrid combinations. CoAtNet distinguishes itself by treating the mixture and stage layout systematically while controlling computation and preserving model capacity.

  • Convolutional network building blocks: ConvNets use regular or depthwise convolutional blocks, with depthwise convolutions favored for lower computational cost and smaller parameter size.
  • Self-attention and Transformers: Vision Transformers achieve impressive results after large-scale JFT pre-training, motivating attention-based approaches for vision recognition.
  • Relative attention: Relative attention includes input-dependent and input-independent variants; CoAtNet uses the latter, which is cheaper to compute and can be cached during inference.
  • Combining convolution and self-attention: Earlier hybrid models commonly add attention to ConvNet backbones, replace selected convolutions, or mix linear attention with convolution, often increasing computational cost.
  • Combining convolution and self-attention: CoAtNet combines depthwise convolution and content-based attention with minimum additional cost, then assigns layer types to stages according to generalization and capacity considerations.
  • Combining convolution and self-attention: Unlike local-attention approaches, CoAtNet uses full attention in S3 and S4 to maintain capacity, while also scaling the convolutional S2 stage.

4 Experiments

CoAtNet is evaluated across increasing data scales and resource constraints, with experiments covering model scaling, ImageNet performance, JFT comparisons, and ablations of its hybrid design.

  • Experiment Setting: CoAtNet models scale by doubling channels across stages while increasing depth mainly in stages S2 and S3.Table 3 specifies block counts, hidden dimensions, kernel sizes, attention-head dimensions, and expansion rates.
  • Experiment Setting: The evaluation uses ImageNet-1K, ImageNet-21K, and JFT, with pre-training followed by ImageNet-1K fine-tuning when resolution changes.Models are pre-trained at 224 resolution and fine-tuned for 30 epochs at desired resolutions.
  • Main Results: 88.56% top-1 accuracy is achieved after ImageNet-21K pre-training, matching ViT-H/14 at 88.55% while using a 23x smaller pre-training dataset.The comparison also reports a 2.3x larger ViT model and 2.2x more training steps.
  • Main Results: On JFT-300M, CoAtNet-4 nearly matches NFNet-F4+ while using 2x less TPU training time and parameters, and CoAtNet-5 reaches 89.77% top-1 accuracy.With JFT-3B, CoAtNet-6 matches ViT-G/14 at 90.45%, while CoAtNet-7 reaches 90.88% top-1 accuracy.
  • Ablation Studies: Relative attention improves generalization and transfer accuracy, while the layout study identifies a trade-off between Transformer capacity and sufficient convolution for generalization.Increasing attention-head size can improve TPU speed despite slightly hurting performance; BatchNorm is 10–20% faster than LayerNorm with similar performance.

5 Conclusion

CoAtNet combines convolutional and Transformer components into a model family designed to retain ConvNet generalization and Transformer capacity. The paper reports state-of-the-art performance across data sizes and computation budgets, while focusing its development on ImageNet classification.

  • CoAtNet combines convolutional and Transformer properties to achieve good generalization and superior model capacity.
  • CoAtNet achieves state-of-the-art performance under different data sizes and computation budgets.
  • The paper focuses on ImageNet classification for model development.Object detection and semantic segmentation are identified as broader applications for future work.

A.1 Model Details

The CoAtNet implementation uses relative attention, pre-activation blocks, and stage-wise down-sampling, with several choices tuned for efficiency and model simplicity. Its classification head uses global average pooling rather than an additional class token.

  • CoAtNet uses trainable 2D relative-bias parameters whose indices encode spatial offsets between image locations.The parameter matrix has size [(2H−1)×(2W−1)] for images of size [H×W].
  • Relative-bias indexing has O(HW(H + W)) complexity on TPU and can use more efficient gather-based memory access on GPUs.Inference-time indexing can be pre-computed.
  • Both MBConv and Transformer blocks use a pre-activation structure with module-specific normalization.BatchNorm is used for MBConv, while LayerNorm is used for self-attention and FFN modules.
  • The first block of each stage performs down-sampling independently on residual and identity branches.Transformer blocks use stride-2 max pooling and channel projection, while MBConv uses stride-2 convolution on normalized inputs.
  • The reported main-text implementation uses first-convolution down-sampling, reflecting a quality-speed trade-off for smaller models.Stride-2 depthwise convolution can help smaller models but is slower, with less benefit as models scale.
  • CoAtNet replaces ViT’s additional <cls> token with global average pooling for classification.

A.2 Hyper-Parameters

The appendix summarizes hyper-parameters and model sizes for the main, JFT, and scaled CoAtNet experiments. Larger variants modify block placement, hidden dimensions, and attention-head sizes to improve speed, memory use, or scalability.

  • The main experiments’ hyper-parameters are summarized in Table 10, including model-size-specific settings.RandAugment settings differ for selected finetuning and PT-RA variants.
  • CoAtNet-5 uses attention heads of size 64 rather than 32 for a better speed-performance trade-off in the JFT experiment.
  • CoAtNet-5 model sizes are summarized in Table 11.
  • CoAtNet-6 and CoAtNet-7 move two-thirds of S2 MBConv blocks into S3 and double S3’s hidden dimension without changing FLOPs.The adjustment reduces activation-related memory usage by half and enables a larger model.
  • CoAtNet-6 and CoAtNet-7 share depth but differ in width, with attention head size increased to 128.Their specific model sizes are summarized in Table 12.

A.3 Complete Comparison

The appendix provides complete comparison tables for CoAtNet under ImageNet-1K-only training and ImageNet-21K pre-training followed by ImageNet-1K finetuning. The latter table also distinguishes RandAugment and extended pre-training variants.

  • Table 13 gives the complete comparison under the ImageNet-1K-only setting.
  • Table 14 gives the complete comparison for ImageNet-21K pre-training followed by ImageNet-1K finetuning.It defines PT-RA as RandAugment during 21K pre-training and E150 as 150 pre-training epochs.
Loading 2106.04803v2…