이 논문에서는 convolution과 relative attention(의 변형)를 적절히 섞고 JFT에서 pre-training해 ImageNet-1K에서 90.88%로 SOTA를 차지하고 있는 CoAtNet(코트넷)에 대해 소개한다.
이 논문에서는 depthwise convolution을 이용한 MBConv와 Transformer을 합치는 방법을 고안한다. MBConv를 고른 이유는 Transformer의 FFN에서도 channel 개수가 늘어났다 줄어드는 inverted bottleneck 구조를 사용하기 때문이다. (4x 늘어났다 4x 줄어든다.)
익히 알듯 Convolution은 Translation Equivariance를 가지고 있는 것이 장점이 되고, Self-Attention은 Input-adaptive weighting이 되고 global receptive field가 장점이다. 논문에서는 둘을 합쳐본 아래와 같은 수식 2개를 제시한다.
흥미롭게도 $y^{pre}i$는 relative self-attention의 변형에 대응된다. 또 수식의 $w{i-j}$는 파라미터 개수를 줄이고 계산 비용을 줄이기 위해 스칼라로 쓴다고 한다. 이 논문에서는 위의 둘 중 후자인 pre-normalization relative attention varient를 사용한다고 한다.
self-attention의 time complexity는 $O(N^2)$이므로 ImageNet의 대형 이미지(224*224)에 그대로 쓸 수는 없다. 따라서 1) 다운샘플링 후 attention 2) local attention 사용 3) linear attention 사용 이 3 가지 방법이 있는데, 이 중 1)을 채택한다. 3)의 경우 간단히 해봤지만 좋은 결과가 나오지 않았고 2)의 경우 non-trivial shape formatting operation의 특성상 TPU에서 매우 느리게 돌아가서 그렇다고 한다.
다운샘플링은 크게 2가지 방법으로 시행한다.
ViT에서처럼 16*16 patch로 나눈다.
이렇게 ViT Steam이 쓰이면 이후 $L$개의 Transformer block을 바로 이어주고, 이를 $ViT_{REL}$로 표기한다.
multi-stage ConvNet으로 downsampling한다.
5개의 단계로 이루어진 ConvNet을 사용한다. 각각을 $S0, S1, S2, S3 \&S4$로 표기한다. 단계를 거치며 점점 해상도를 줄여나간다. 매 단계별 크기를 2x 줄이며 채널은 늘린다. 더 자세한 내용은 Appendix A.1과 후술할 내용 참조.
S0은 간단한 2-layer convolutional Stem이고, S1는 squeeze-excitation 포함한 MBConv이다. 이후 Transformer을 넣느냐 Convolution을 넣느냐에 따라 CCCC, CCCT, CCTT, CTTT를 후보군으로 두고 설정하고 일반화 능력, model capacity, transferability에 대해 test를 진행하고 CCTT를 최종 선정한다.
CCTT까지 선정하는 과정에서, $ViT_{REL}$의 일반화 능력은 예상대로 다른 모델에 비해 매우 떨어졌으며 T가 많다고 model capacity가 무작정 올라가는 건 아니었다.