부스트캠프 Ai tech/논문리뷰

[1] CoAtNet

소갱 2022. 2. 11. 12:10
CoAtNet : Marrying Convolution and Attention for All Data Sizes

Abstract

  • Computer Vision에서 Transformers의 관심은 커지고 있지만, SOTA convolutional networks 보다 뒤쳐져있다.
  • Transformers는 larger model capacity를 가지고 있지만, right inductive bias의 부족때문에 generalization은 Convolutional network보다 떨어짐
  • Transformers와 Convolutional network 두 아키텍처의 장점을 결합하기 위해서 CoAtNets을 제시했다.
  • 두 가지 key insight
    1. depthwise Convolution과 self-Attension은 간단한 relative attention을 통해 결합 가능하다.
    2. vertically하게 convolution layers와 attention layers를 쌓아서 generalization, capacity, efficiency 세 가지 측면에서 성능 개선
  • Experiment에서도 나오지만 ViT보다 23배 적은 데이터를 사용하여도 비슷한 성능을 가졌다.

generalization = 학습 Data와 Input이 달라져도 상관없게 하는 것

 

Introduction

  • AlexNet의 획기적인 발전 이후, ConvNets가 CV에서 dominat하였고, 최근에는 Transformer를 vision에 적용하려는 시도가 많았다.
  • Vision Transformers (ViT)는 ImageNet-1K에서 vanilla Transformers layer와 비슷한 성적을 가짐
  • 큰 Datasets에서 ViT는 State-of-the-art(SOTA) ConvNets와 비슷하지만 low data regime에서는 ConvNets에 비해 떨어지는 모습을 보였다.
  • 후에 regularization과 data augmentation을 통해서 ViT의 성능을 향상시켰지만, ConVNets을 뛰어넘기는 힘들었다.
  • 이는 Transformers가 desirable inductive biases가 부족하기 때문에 많은 data와 Computation이 필요하다는 것을 알 수 있다.
  • 최근 많은 연구들이 ConvNets의 inductive biases를 Transformers에 주입하는 방향으로 시도되고 있다.
  • 위를 포함한 여러 방법들은 한정적인 방법들이며, systematic understanding이 부족하다.
  • 그렇기 때문에 이번엔 Machine Learning의 gerneralization과 model capacity측면에서 convolution과 attention을 결합하는 방향으로 연구한다.
  • 이번 연구에선 convolutional layers가 strong inductive bias로 generalization 성능이 좋고, attention layer는 high capacity로 많은 데이터셋에서 이점을 가진다.
  • 따라서 우리는 accuracy와 efficiency의 좋은 trade-off를 위해 attention과 convolution의 효율적인 결합 방법을 생각했다.
  • 두 가지 key insight
    1. depthwise Convolution과 self-Attension은 간단한 relative attention을 통해 결합 가능하다.
    2. vertically하게 convolution layers와 attention layers를 쌓아서 generalization, capacity, efficiency 세 가지 측면에서 성능 개선

Model

convolution과 transformer를 결합하는 최적의 방법을 찾는다.
어떻게 하면 convolution과 self-attention을 하나의 블록으로 만들 수 있을까
어떻게 다른 두가지 타입을 쌓을 수 있을까

Merging Convolution and Self-Attention

  • MBConv block에 집중 (depthwise convolution을 사용) -> Key is "inverted bottleneck"
  • depthwise convolution과 self-attention은 predefined receptive field에서 per-dimension weighted sum으로 표현 가능한 공통점이 있다.

  • self-attention은 receptive field -> entire spatial location

  • depthwise convolution kernel의 w는 input에 독립적인 static value지만 attention weight의 A_(i,j)는 input에 representation에 따라 dynamic하다.
  • 따라서 self-attention에서 spatial positions들의 상호 연관관계를 더 잘 포착할 수 있다. 하지만 그렇기 때문에 overfitting의 risk가 크기도 하다.
  • Convolution weight은 상대적인 shift만을 고려하기 때문에 translation equivariance 성질(Conv성질을 기억하면 이해하기 쉽다. 특정 위치의 값이 바뀌면 출력도 따라서 바뀌는 성질) 을 가지며 이로 인해서 적은 dataset에서 generalization 성능이 좋다.
  • ViT는 absolution positional embedding을 사용하였기 때문에 위에서 언급한 translation equivariance 성질이 부족하다. <-> translation invariance. 이러한 것이 ConvNets가 대게 크지 않은 데이터셋에서 Transformers보다 나은 이유이다.

  • 또한, receptive field의 크기가 self-attention과 convolution의 중요한 차이중 하나이다. 일반적으로 larger receptive field가 model capacity를 키워주는 contextual information을 가져 온다. 그렇기 때문에 global receptive field가 self-attention을 CV에 기용하는데 중요한 역할이 된다.
  • 하지만 global receptive field는 spatial size에 대해 quadratic한 계산량 증가 문제가 있다.

  • 이상적인 모델은 위 표에 나와있는 properties가 모두 좋아야 한다.

  • 위의 식은 Attention layer의 softmax nomalize의 pre 또는 post에 global static convolution kernel을 더하는 아이디어이다.
  • pre, 즉 이전에 더하는 경우 서로가 연관이 되어 relative self-attention이며 attention weight은 W_(i,j)와 input-adaptive(x_i,x_j)에 의해 결정되기 때문에 두 개의 장점을 다 가진다.
  • parameter 수를 크게 늘리지 않으면서 global convolution kernel을 사용하기 위해서 논문에서는 W_(i,j)를 스칼라 값이라고 생각
  • 위의 이유로 인하여 pre-normalization relative attention variant를 Transformer block에 사용한다.

Vertical Layout Design

Convolution과 attention의 효율적인 layer stacking에 대한 고려
  • vertical로 바로 쌓을 경우 quadratic complexity로 인하여 속도가 매우 느릴것.
  • 3가지 방법을 고려
    (A) - down sampling으로 spatical size를 줄이고, feature map의 size가 충분히 작아졌을 때 global relative attention을 적용
    (B) - global receptive field를 local field로 변경
    (C) - quadratic softmax attention을 spatial size에 대해서 linear complexity를 가진 linear attention으로 변경

실험을 통했을 때 (C)는 좋은 성능을 내지 못하였으며, (B)는 global attention 사용을 불가능하게 할 뿐 아니라 model capacity를 낮추며 shape formatting 과정에서 많은 memory access로 속도 또한 좋지 않았다. 그리고 (A)의 down-sampling은 ViT-stem or CNN의 gradual pooling을 사용할 수 있었다.

  • CNN 형식의 down-sampling은 5-stages로 구성되어 있다. (S0,S1,S2,S3,S4)

S0 = simple 2-layer convolutional Stem

S1 = SE-module이 사용된 MBConv blocks

S2~S4 = MBConv or Transformer block 사용 가능

  • Convolution은 초기의 local patterns을 학습하는데 더 좋기에 convolution은 transformer보다 항상 먼저 위치해야함 C -> T
  • 그렇게 나온 조합은 {C-C-C-C}, {C-C-C-T}, {C-C-T-T}, {C-T-T-T}

dataset for 300 and 3 epochs

ImageNet-1K에서 ViT의 성능이 크게 떨어지는데, 이는 적절한 low-level information을 얻는데 실패하였을 것이다. 전체적으로 보았을 때 Convolution stage가 많을수록 generalization 성능이 좋은 것을 확인할 수 있었다. 실험을 통해 Trnasformers block의 수가 많다고 성능이 반드시 더 좋은것이 아님을 알게 되었고, CCTT와 CTTT의 성능이 비슷한것은 low-level information에 한해서 static local convolution operation, adaptive global attention의 성능이 비슷하다고 해석할 수 있다.

이후 C-C-T-T와 C-T-T-T 중 최종 모델 구조를 결정하기 위해서 Transferability test를 수행하였다.

  • C-C-T-T의 성능이 전체적으로 더 좋은 것을 확인할 수 있었다. 따라서 C-C-T-T의 구조를 CoAtNet의 multi-stage layout으로 최종결정 하였다.

Expriments

  • 위 표를 보면 CoatNet의 Top 1 accuracy 는 88.56%이고 이는 ViT를 취약 데이터를 기준으로 23배 더 학습시킨 결과와 비슷하다.
  • 또한 1K 데이터 셋에서도 높은 accuracy를 보이는 것을 표를 통해 확인할 수 있다.

  • Relative attention을 적용했을 때 더 높은 정확성을 보이는 것을 확인할 수 있었다.
  • Attention Head size를 늘렸을 때 32 -> 64 정확도가 낮아지는 것을 확인

  • BatchNorm이나 LayerNorm은 문제가 되지 않음

 

Conclusion

  • Convolution과 Transformer의 속성에 대한 체계적인 연구를 통해 둘을 최적의 방법으로 결합한 CoAtNet을 제시
  • ConvNet의 일반화 성능과 Transformer의 model capacity를 모두 갖추어 다양한 Dataset에서 CoAtNet이 SOTA를 달성함
  • future work로 object detection, semantic segmentation에서의 적용을 제시

 

Relative Attention : Attention을 Global static convolution kernel을 SoftMax 전에 넣을 지 후에 넣을 지 두가지 경우로 나누는 방법