CoAtNet

CoAtNet hybridises depthwise convolution and relative self-attention in a single backbone, motivated by the observation that convolutions excel at generalization (strong local inductive bias) while attention excels at capacity (data-dependent global mixing). The network is a four-stage pyramid preceded by a two-layer convolutional stem; the early stages use MBConv blocks (squeeze-and-excitation, expansion ratio 4) and the later stages use relative-attention transformer blocks operating on $\sqrt{HW}$-length token sequences.

The key building block of the transformer stages is relative self-attention:

where $r_{i-j}$ is a learned bias indexed by the relative spatial offset between tokens. This recovers the translation equivariance that convolutions provide for free while still permitting global, data-dependent mixing. CoAtNet stacks two MBConv stages (C-stages) followed by two transformer stages (T-stages) so that strong local features are extracted first and long-range dependencies are modelled at lower resolutions where attention is cheap. The CoAtNet-0 through CoAtNet-7 variants simply scale stage depths and channel widths along this skeleton.