Image Classification task (ImageNet)의 새로운 SOTA로 Transformer와 ConvNet이 결합된 모델이며 확실히 Computer vision 분야에서도 Transformer기반 모델이 대세가 되었음을 보여준다.
Abstract
컴퓨터 비전 분야에서 Transformers가 많은 관심을 끌고 있다. 하지만 아직도 SOTA, ConvNet기반 모델에 비해 성능이 좋지 못함. 그 이유는 Transformers가 ConvNet보다 바이어스를 올바르게 유도하지 못하기 때문이며, 이로 인해 일반화 능력이 떨어지기 때문이다. Transformers, ConvNet 두 개의 아키텍처를 효율적으로 조합하기 위해 CoAtNets를 제안한다.
CoAtNets은 다음 두 가지 Key insight로부터 설계 되었다.
depthwise convolution과 self-attention은 간단한 relative attention을 통해 자연스럽게 통합될 수 있다.
convolution layer, attention layer를 수직적으로 쌓는 방법은 일반화 성능, 용량, 효율성에 있어 매우 효과적이다.
다음과 같은 실험들은 CoAtNets가 다양한 데이터셋에서 SOTA급 성능을 달성한 것을 보여준다. (추가 데이터 없음)
ImageNet-21K에서 13M개의 이미지로 사전학습 했을 때, top-1 accuracy: 86.0%
JFT-300M에서 300M개의 이미지로 사전학습 했을 때, top-1 accuracy: 88.56%
JFT-3B를 같이 사용하여 CoAtNets를 확장 했을 때, top-1 accuracy: 90.88%(SOTA)
Introduction
AlexNet의 등장 이후로 CNN은 컴퓨터 비전 분야에서 지배적인 아키텍처이다. 한편 NLP에서는 Transformers와 같은 self-attention 모델이 성공해서 이전의 많은 실험들은 이러한 강력한 attention을 Computer Vision으로 가져오려는 시도를 많이 했다.
최근에는 ImageNet-1K에 한해 Vision Transformer(ViT) 모델이 기본적인 Transformer layer만 사용해 좋은 성능을 얻었다. 중요한 점은 대규모의 weakly labeled JFT-300M 데이터셋으로 사전학습 했을 때, ViT는 ConvNet과 유사한 성능을 달성한다는 점이다. 이는 Transformer 모델이 ConvNet에 비해 잠재적인 용량이 더 크다는 것을 의미한다.
ViT는 JFT 300M과 같은 대규모 데이터셋에서는 높은 성능을 보이지만 low data regime(규모가 작은 데이터)에서는 ConvNet보다 성능이 확실히 떨어진다. 예를 들어, 추가적인 JFT-300M 사전학습이 없을 때, ViT는 비슷한 크기의 ConvNet보다 ImageNet에서의 성능이 많이 안좋다. 바닐라 ViT의 성능을 향상시키기 위해 후속 작업에서 special regularization과 강한 data augmentation을 적용했다. 하지만 동일한 양의 데이터와 컴퓨팅 자원이 주어지면 이 중 어느것도 SOTA 성능을 내는 Convolution 모델을 능가할 수는 없었다. 이는 바닐라 Transformer layer가 바이어스 유도 능력이 ConvNet보다 떨어지고, 비슷한 성능을 내기 위해 많은 양의 데이터와 컴퓨팅 자원이 필요하다는 것을 의미함.
당연히도 최근에 많은 실험이 Transformer 모델에 ConvNet을 통합하려고 했다. attention layer를 위한 local receptive fields를 추가하거나 암시적, 명시적 convolution 연산으로 attention, FFN layer를 증강한다. 그러나 이러한 접근 방식은 ad-hoc(일반화 할 수 없는 해결책) 이거나 결합할 때 Convolution, Attention 각각의 역할에 대한 이해가 부족하다. (과거에는 무작정 결합하는 데에만 중점을 둠)
본 논문에서는 Convolution, Attention을 어떻게 결합할지 체계적으로 연구했다.
Model
Convolution과 Attention을 어떻게 최적의 방법으로 조합할 것인가? 이는 다음 두 개의 질문으로 나누어 생각할 수 있다.
Convolution과 Self-attention을 하나의 block으로 어떻게 합칠 것인가?
완전한 네트워크를 형성하기 위해 다양한 block을 어떻게 쌓을 것인가?
※ Merging Convolution and Self-Attention
Convolution convolution은 고정된 커널에 의존하여 local receptive field로부터 정보를 얻는다. 수식은 다음과 같다.
\( \mathcal L(i) \)는 \(i\)의 이웃 픽셀(\(i\)를 중심으로 하는 3*3 grid)을 뜻한다.
Self-Attention self-attention은 receptive field가 전체 공간 위치가 되도록 허용하고 \( (x_i, x_j) \)ㅅ 쌍의 유사도를 기반으로 weight를 계산한다. 수식은 다음과 같다.
\( \mathcal G \)는 global spatial space를 의미함.
※ Convolution VS Self-Attention
Input dependency · Convolution : input-independent parameter of static value. (입력에 따라 달라지지 않음) · Self-Attention: dynamically depends on the representation of the input. (입력에 따라 값이 달라짐)
즉, Self-attention이 서로 다른 위치의 상호작용을 포착하기 쉽다.
하지만 데이터가 제한적일 때 과적합 되기 쉽다.
Position · Convolution 임의의 위치 쌍 \(i, j\)가 주어지면 이들의 특정 값 보다는 상대적 이동에만 관심을 둔다. (ex: \(i-j\)) → 이러한 속성은 제한된 크기의 데이터셋에서 일반화 능력을 향상 시킨다 (= translation equivalence)
· Self-Attention positional embedding을 사용하기 때문에 ViT는 이러한 속성이 존재하지 않는다.
데이터 셋의 크기가 그렇게 크지 않을 때 ConvNet이 Transformer보다 성능이 좋은 것을 어느정도 설명할 수 있다.
The size of receptive field 일반적으로 receptive field가 클수록 더 많은 context 정보를 제공한다고 알려져 있다. 그래서 vision에서 self-attention을 사용하려는 시도를 했다. 하지만 receptive field가 클수록 많은 연산이 필요(용량이 커짐)하다. global attention의 복잡도는 2차원(제곱으로 증가)이며 이것이 self-attention 적용의 trade-off가 된다.Table 1
이상적인 모델은 Table 1의 3가지 속성을 결합할 수 있어야 한다. 이는 간단하게 softmax 정규화 이전 또는 이후에 convolution, kernel, adaptive attention matrix를 더하면 된다.
정규화 이후정규화 이전
정규화 이전의 경우, Attention weight \(A_{i,j}\)가 translation equivariance \(w_{i, j}\)와 input-adaptive \(x_i^\top x_j\)에 의해 결정된다. 이는 상대적인 크기에 따라 두 가지 효과를 모두 누릴 수 있다.
CoAtNet의 핵심 구성요소인 Transformer block (정규화 이전)을 사용하고 이를 Relative Attention이라 한다.
※ Vertical Layout Design
Global context의 문제점은 제곱으로 증가하는 연산량(2차원)이다. 이 때문에 relative attention을 직접 적용하게 되면 raw image가 입력으로 들어올 경우 수 많은 pixel로 인해 계산이 매우 느려질 것이다.
그래서 실제로 네트워크를 구성하기 위해 다음과 같이 3가지 옵션을 사용한다.
Down-sampling을 통해 크기를 줄여서 feature map을 관리 가능하도록 한 뒤 적용.
Convolution에서의 local field와 같이 global receptive field를 제한한다.
2번의 경우 TPU상에서 비효율적이기 때문에 연산 속도를 올린다는 목적을 달성할 수 없어서 사용하지 않는다.
3번의 경우 좋은 성능을 보이지 못했으므로 사용하지 않는다.
따라서 1번 방식만을 고려하여 실험을 진행한다.
다운 샘플링의 경우 크게 두 가지 방법이 존재한다.
ViT에서와 같이 높은 stride를 가진 convolution stem 해당 방식에서는 Relative attention을 적용한 Transformer 블록을 바로 쌓는다. 이를 \( \text{ViT}_\text{REL} \)
ConvNet에서와 같이 pooling과 함께 multi-stage network 구성 해당 방식에서는 5 stages(S0 ~ S4)로 구성된 ConvNet을 사용한다. (S0에서 S4로 갈수록 resolution 감소) 각 stage의 시작점에서 항상 공간 크기를 1/2로 줄이고 채널 수를 늘린다.
· S0: 간단한 2-layer convolutional stem · S1: global attention을 적용하기엔 사이즈가 너무 커서 squeeze-excitation(SE)를 사용한 MBConv block 사용. · S2 ~ S4: MBConv 또는 Transformer block고려, Convoltuion이 Transformer stage보다 먼저 나타나야 함.
Convolution이 초기 단계에서 로컬 패턴을 처리하기에 더 좋다는 사실에 근거해 다음과 같이 4가지 경우를 비교한다.
(C: Convolution / T: Transformers)
C-C-C-C
C-C-C-T
C-C-T-T
C-T-T-T
이러한 구조를 선택하기 위해 일반화 능력과 모델의 용량 두 가지를 고려한다.
일반화 일반화를 위해 training loss와 evaluation accuracy 사이의 갭에 초점을 둔다. 만약 두 개의 모델이 같은 training loss를 가질 때 evaluation accuracy가 높은 모델이 일반화 능력이 좋다고 볼 수 있다.
일반화 능력은 ImageNet-1K의 결과로 측정할 수 있으며 다음과 같다.
\( \text{ViT}_\text{REL} \)의 경우 일반화 성능이 확연히 떨어지며, Conv가 많을수록 일반화 성능이 좋아짐을 알 수 있다.
용량 (Capacity) 모델이 큰 training dataset에 얼마나 fit할 수 있는지 측정한다. training data가 방대하고 overfitting이 발생하지 않는다면 용량이 더 큰 모델이 더 좋은 성능을 가질 것이다.
모델의 용량은 JFT를 이용해 측정 가능하다. 결과는 다음과 같다. 여기서 중요한 점은 단순히 Transformer 블록을 많이 쌓을수록 높은 용량을 가지는 것은 아니라는 점이다.
최종적으로 C-C-T-T 와 C-T-T-T를 선택하여 JFT pre-trained 모델을 ImageNet-1K로 30 epoch transfer하는 과정을 거친다. 두 모델의 성능 비교는 다음과 같다.C-C-T-T가 더 좋은 성능을 보이는 것을 확인할 수 있다.
