CoAtNet: Marrying Convolution and Attention for All Data Sizes

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.


Reviews and author responses:

 After reading the rebuttal, I still recommend accepting the paper due to its excellent performance and comprehensive experiments. But please note I still think the novelty of the paper is limited and the authors did not address my concerns.

 
> The authors mentioned the limitation is that they only focus on the
> classification. I think the authors should include more, for example, the
> models' actual latency on TPUs/GPUs.

-   There are some genuine difficulties to benchmark the TPU latency for all
    models, as any sparse operation (gather) or reshape operation would be super
    slow on TPU.
-   Given the difficulty, we only managed to compare the TPU-V3 latency of
    CoAtNet, NFNet and DeiT (ViT) (measured with per-core batch size = 32 and resolution 224):
    [Figure link](https://i.ibb.co/23zCKfR/speed.png").

    As we can see from the Figure, CoAtNet achieves comparable speed-quality
    trade to NFNet while being significantly better than DeiT (ViT).

 
> ... But I am curious about whether the relative attention itself can well
> exploit the information given by the inductive biases. It would be more clear
> to add another experiment replacing all convolutions in ResNet with the
> relative attention designed in this paper.

- As relative attention is global, directly replacing "local" convolution with
relative attention in ResNet would be too expensive/slow to run in reality.
- Meanwhile, from the ablation of relative attention in Table 6, we can see adding relative
attention upon standard attention does help to boost the performance.

--------------------------------------------------------------------------------

> the indexing of the relative self-attention in Eq (3) makes me a bit confused.
> It would be better to have a figure here demonstrating the entire process.

Thanks for the suggestion. We will include a diagram to illustrate the process.

 
> ... would be good to also compare with PVT

Sure, we will add PVT into experimental comparison (see [Figure link](https://i.ibb.co/hx574JH/imagnet1k.png)) as well as
discussing the difference in related work.

--------------------------------------------------------------------------------

> are the weights w_{i-j} the same or different for each head?

They are **different for each head**. We discussed this detail in the Appendix
(L474).

--------------------------------------------------------------------------------

> footnote 2: Says this omits q k v projections for now, but I don't see them
> re-introduced later. If these projections are used in the actual system, I
> think it would be good to either include, or mention that they are used
> possibly with multihead attention reference.

In our actual implementation, we always use the **multi-head** Q, K, V projections.
We will clarify that in the updated version.

--------------------------------------------------------------------------------

> ... but I think key sentences or equations of A.1 including downsampling
> placements would be better included the main text, ...

Thanks for the feedback. We will add a detailed explanation in the main text
clarify that.

--------------------------------------------------------------------------------

> Table 7: How does moving layers between S2/S3 affect FLOPs? Also text says V1
> has better imgnet21k pretraining accuracy but both v0 and v1 in the table show
> 53.0.

-   Firstly, the FLOPs are **roughly comparable** when moving a layer from S2 to
    S3.
    -   When D > W x H (in S3, WxH=49), the dominating complexity of a Conv/TFM
        layer is O(WHD^2) (linear projections), where D is the hidden dim, and W & H are the width
        and height.
    -   From S2 to S3, we always increase the D by 2x and reduce the W & H by 2x
        at the same time
    -   Hence, the complexity does not change.
-   For the 21K pre-training accuracy, V0 and V1 **should be comparable** (the exact
    numbers are V0 = 53.01, V1 = 53.02). We will correct the text and thanks for
    the catch.

--------------------------------------------------------------------------------

> 3.2 option (A) --- PVT and CvT also use a form of downsampling, .... I think
> this could be good to contrast.

-   Effectively, both PvT and CvT utilize **strided (separable) Convolution** to
    reduce the resolution, since reshaping [H x W x C] into [H/R x W/R x R^2C]
    and performing a linear projection is equivalent to stride-R Convolution
    with kernel RxR.
-   In comparison, CoAtNet simply uses **strided pooling**. This difference has
    appeared before in Computer Vision with ConvNet, where models have converged
    to stride pooling. That's also the reason that we use strided pooling.


 
> ... as relative position bias have been discussed in many previous works, it
> is hard to say this is the contribution of this paper.

We totally agree that relative attention **NOT** is our contribution. We only
hope to **point out our observation** that while there are many efforts trying to
put convolution priors into the attention, the "well established" relative
attention is a "natural" combination with great performance.

This paper explores how to combine convolution and self-attention for vision models. It points out that depthwise convolution and self-attention can be unified as relative attention. In addition, it proposes a family of hybrid models, named CoAtNets, which vertically stack convolution layers and self-attention layers. Experiments show that CoAtNets strike a balance between generalization and model capacity and these models achieve state-of-the-art performance in different data regimes.   This paper systematically analyzes how to hybrid convolution and attention. It studies this from aspects of model generalization ability and model capacity.  

It shows depthwise convolution and self-attention can be simply merged as relative self-attention module. This unified view is interesting but not so innovative.

In addition, CoAtNets are proposed by searching in 5 variants. Rules are designed to evaluate the generalization and capacity of models. These strategies are straightforward but inspiring. They may be useful for future network design.

Extensive experiments show CoAtNets achieve both generalization and capacity. Their SOTA performances are impressive.
 1. This paper regards relative self-attention as the hybrid of convolution and self-attention. This understanding is interesting. However, as relative position bias have been discussed in many previous works, it is hard to say this is the contribution of this paper.

This paper presents a system and experimental studies on combining convolutions with transfomers.  Multiple placements of a convnet->stacked transformer transition are evaluated, finding about three conv stages and two transformer stages are best in most cases.  A set of shift-parameterized weights are added to attention similarities in the transformers, which conceptually combines convolution and attention operations in these stages as well.  The system achieves high accuracy for FLOPs on ImageNet, in cases both with and without extra data and pretraining, and the effect of different parts of the system are evaluated with ablations in resource allotment.
  This is a good set of studies and explanations of the effect of conv->transformer stacked models.  The eventual model arrived at is similar to the hybrid ResNet50-ViT stride 16 model described in Visual Transformers [13], though with differences around the structure of the convolution stages, using depthwise conv blocks.  Section 3.2 provides a more comprehensive study around conv->transformer placement, evaluating different numbers of conv stages from stride 8 up to 32.

Combining convolutional (shift-dependent) weights in the attention weights is also an interesting approach, which the experiments show has positive effect.  I would have liked to see a comparison between the two variants mentioned in the description ("pre" and "post", only "pre" is used).

Another recent work that combines convnets and transformers is PVT https://arxiv.org/pdf/2102.12122.pdf, which I think would be interesting to compare with.  The hybrid ResNet50-ViT in [13] is also not mentioned, and I think would also be good to add to Table 5 and Figures 2 and 3.

This is a performant system that is explained well, including an interesting relative attention module that combines convlution-like weights directly into attention values.  Much of the general architecture ends up similar to recent other hybrid approaches, with additional studies validating these choices.



Additional comments and questions:

- would be good to also compare with PVT https://arxiv.org/pdf/2102.12122.pdf

- for multi-headed attention, are the weights w_{i-j} the same or different for each head?

- sec 3.1 footnote 2:  Says this omits q k v projections for now, but I don't see them re-introduced later.  If these projections are used in the actual system, I think it would be good to either include, or mention that they are used possibly with multihead attention reference.


- sec 3.2:  how is 2x downsampling performed for stages S2..5, for either C or T variants?  This is in appendix A.1, but I think key sentences or equations of A.1 including downsampling placements would be better included the main text, as they help explain questions that came up immediately for me when reading 3.2.


- Table 7:  How does moving layers between S2/S3 affect FLOPs?  Also text says V1 has better imnet21k pretraining accuracy but both v0 and v1 in the table show 53.0.


- 3.2 option (A) --- PVT and CvT also use a form of downsampling, within each attention layer on kv maps (in addition to striding between stages), which is different from the approach here that is only studying downsampling between stages.  I think this could be good to contrast.

 -

This paper proposes CoAtNet, a hybrid model that can take the advantage of convolution and self-attention. The efficacy of CoAtNet can be attributed to two components, relative self-attention and vertical layout design. The proposed relative self-attention is essentially a combination of self-attention and relative positional encoding. Based on that, the CoAtNet further applies some other tweaks including relacing the self-attention into convolutions in the early stages, adding SE modules etc.
Based on the above observations, CoAtNet achieves state-of-the-art performance on ImageNet.  Originality: 

- The novelty of this paper is somewhat limited. The relative attention proposed in this work is very similar to the positional encoding designed in LambdaNet [1]. Or it could be regarded as a 2D-version relative positional encoding proposed in HaloNet [2] and SASA [3]. It would be preferable if the author could concretely compare with the above modules in their paper.

- The vertical layout design can also be found in some existing works like LambdaNet [1] and AANet [4]. To me, it is not surprising that convolution can achieve a better cost-accuracy tradeoff in the early stages. 

---

Quality:

- The design of CoAtNet is backed by extensive and comprehensive experiments. I appreciate the efforts that the authors made to analyze the generalization and model capacity of convolution and self-attention (Figure 1). But I am curious about whether the relative attention itself can well exploit the information given by the inductive biases. It would be more clear to add another experiment replacing all $3\times 3$ convolutions in ResNet with the relative attention designed in this paper.


---


Significance: 

- In general, this paper should be impactful in its corresponding field due to the comprehensive experiments and strong performance. However, as it does not extend to other downstream tasks like object detection, the impact of this paper may be somewhat limited in other fields of computer vision.

---

Clarity: 

- This paper is well written and easy to understand. However, it may be hard to reproduce as there are lots of tweaks (e.g. SE, Stem, MBConv, Pre-activation, Relative Attention...) introduced in just one network. It would be preferable if the authors promise to release the code of the proposed method. Meanwhile, the indexing of the relative self-attention in Eq (3) makes me a bit confused. It would be better to have a figure here demonstrating the entire process.

---

References

[1] Bello I. Lambdanetworks: Modeling long-range interactions without attention. ICLR 2021

[2] Vaswani A, Ramachandran P, Srinivas A, et al. Scaling local self-attention for parameter efficient visual backbones. CVPR 2021

[3] Ramachandran P, Parmar N, Vaswani A, et al. Stand-alone self-attention in vision models. NeurIPS 2019

[4] Bello I, Zoph B, Vaswani A, et al. Attention augmented convolutional networks. ICCV 2019 Limitations:

- Some of my concerns are given in the main review. See above.

- It would be preferable to provide a fair comparison in terms of throughput with other networks listed in Table 4.

- It would be more clear to provide the pseudocode of the relative self-attention in the appendix.

This paper proposes an architecture CoAtNet that combines convolution and relative attention. The paper claims that CoAtNet can have the advantages of both operations and can achieve better or similar performance in all data sizes. Experiments show state-of-the-art classification performance on ImageNet-1K, ImageNet-21K, and the larger JFT dataset with the less computational resource.
  Pros:

+ The paper is clear and easy to understand.
+ The paper presents many good practices, which can be very useful to the community and shed light on the design of convolution and attention mixtures.
+ The experiments are comprehensive and as many details as possible are included in the paper. This makes other researchers easily reproduce the results.
+ The proposed CoAtNet achieves very good performance, outperforming previous SOTA methods such as Swin.

Cons:
- The paper has limited novelty. Using the mixture of convolution and attention has been attempted by many previous works such as BoTNet [11]. The contribution of the work is just to find a better combination, thus being incremental.
 The authors mentioned the limitation is that they only focus on the classification. I think the authors should include more, for example, the models' actual latency on TPUs/GPUs.

Meta-review: The rebuttal addressed all of the reviewers concerns, and all reviewers recommend acceptance. The AC agrees with this recommendation.

Decision: accept