[ICLR 2021] An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale

최근 Transformer와 Multispectral Pedestrain Detection을 접목한 논문을 리뷰하게 되었습니다. (해당 논문에 대해서는 김지원 연구원이 간단하에 X-Review로 작성해주셨기에 궁금하신 분들은 읽어보시는 것을 추천드립니다.) 사실 저는 Transformer 에 대해 말로만 들어봤지 제대로 알지는 못하여.. 이를 이해하기 위해서는 추가적인 공부가 필요했었습니다.

따라서 그 때의 공부의 연장선으로 Vision-Transformer에 대한 논문을 다시 읽어보고 리뷰를 해보려고 합니다. 제가 이번에 리뷰할 논문은 Vision에 처음으로 Transformer를 적용한 ViT 입니다. (다행스럽게도 본 논문이 아직 X-Review로는 작성되지는 않아 조금 편하게 리뷰를 작성할 수 있었습니다)

---

Background

Inductive bias?

최근 논문에 많이 나오는 용어인데 정확한 뜻에 대해서는 아직 모르고 있어서 찾아보았습니다. Bias와 Variance는 거의 짝지어 다니는 단어인데요, 많은 분들이 아시다시피 Bias는 ‘편향’, Variance는 ‘분산’ 이라고도 합니다. 쉽게말해 Bias는 모델의 예측값과 실제값이 얼마나 떨어져있는가를, Variance는 예측값으로부터 얼머나 퍼져있는냐는 의미합니다.

위키피디아에서의 정의를 빌리자면, Induvtive Bias란 학습 시 만나지 않은 상황에 대하여 정확한 예측을 하기 위해 사용하는 추가적인 가정을 의미합니다. 편향의 관점으로 다시 정의를 내리자면 모델이 학습 과정에서 본 적이 없는 분포의 데이터를 입력 받았을 때, 해당 데이터에 대한 판단을 내리기 위해 가지는 “학습과정에서 습득된 Bias(편향)” 을 Inductive Bias 라고 합니다.

[ViT] An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale

Open Review :: Paper / Review

Introduction

부제: 왜 Transformer를 Vision에 적용?

“Attention is All You need” 논문의 등장으로 Vision의 CNN과 같이 Self-attention 기반의 Transformer는 NLP에서의 Standard가 되었습니다.

• Transformer의 계산 효율성(Efficiency) 과 확장성(Scalability) 덕분에 100B를 가지는 파라미터 수도 학습이 가능하다
• 심지어 데이터셋이 클수록 모델을 키워도 되고, 성능이 Saturate될 가능성이 보이지 않다. → 즉, 더 키울수록 성능 역시 높아질 것이다

그러나 Vision에서는 여전히 Convolution 기반의 아키텍처가 지배적이었습니다. Transformer에 영감을 받아 Self-Attention을 적용하는 시도는 있었으나, Hardware Accelerators에 비효율적이었다고 합니다.

따라서 본 논문에서는 기존 방법들과는 다르게 NLP에서의 Transformer를 최대한 그대로 가져다 쓸 수 있도록 최소한으로 수정하여 Vision에 접목시켰습니다. 자세한 내용에 대해서는 뒷 단원에서 서술하도록 하겠습니다.

다만 Mid-Size 의 ImageNet으로 데이터셋을 학습할 경우, 기존 SOTA인 ResNet보다 낮은 성능을 보입니다. 이는 Transformer가 Locality와 Translation Equivatriance 같은 CNN에 내재된 일부 Inductive Bias가 부족하기 때문이라고 합니다. 떠ㅏ라서 JFT-300M 같이 ImageNet보다 더 큰 데이터로 사전학습 후 작은 데이터로 Fine Tunning해야 원하는 결과를 얻는다고 합니다.

즉, 충분히 많은 양의 데이터로 학습을 해야, Inductive Bias를 능가한다!!를 주장합니다.

(그나저나 이제는 ImageNet은 중간 사이즈의 데이터라고 논문에서 언급을 하는군요 ..)

Method

Vision Transformer (VIT)

위 그림이 바로 ViT의 개요입니다.

먼저 이미지를 고정된 크기의 작은 Patches로 분할하고, 각 Patch를 Linear Embedding의 시퀀스로 변환 후 위치 임베딩을 추가. 이 때, Classification을 수행하기 위해 학습 가능한 “Classification Token”을 시퀀스에 추가하여 트랜스포머 인코더의 인풋으로 전달한다.

본격적인 ViT에 대해 알아보기에 앞서, Transformer 가 어떻게 구성되있는지 알고있으면 더욱 이해가 용이할 것 같아 Transformer에 대해 예시와 그림을 함께 서술해주는 참고 사이트를 첨부합니다. Transformer에 대하여 [The Illustrated Transformer – NLP in Korean]

ViT에서의 Transformer 입력은 가능한 기존 NLP에서의 입력을 그대로 사용하도록 되었습니다

Input Image To Patches

먼저 HxWxC의 이미지를 flatten한 2D Patches의 sequence로 재구성합니다.

• 입력 이미지: $(H,W)$
• 패치: $(P, P)$
• 패치 개수(N): $N = HW/P^2$

Patch Embedding

이후, 분할한 2D Patch를 다시 1차원으로 Flatten(크기: $NP^2C$)을 하게 됩니다.

왜냐하면 Transformer의 내부 모든 layer에서는 latent vector의 size가 $D$로 통일되어 있기 때문에, 2차원 patch들을 다시 1차원으로 flatten한 후, trainable한 linear projection을 거쳐 $D$차원 벡터로 매핑시킵니다. 이 Projection의 결과를 Patch Embeddings 이라고 합니다.

Extra Learnable [Class] Token

BERT의 [CLS] 토큰과 같이 임베딩된 패치들 맨 앞에 하나의 학습가능한 class 토큰 임베딩 패치$(z_0^0 = x_{class})$를 추가합니다. 이 [class] 임베딩 패치는 여러 Encoder 층을 거쳐 최종 output으로 나왔을 때, 이미지에 대한 1차원 Representation vector로서의 역할을 합니다. (pre-training, fine-tuning 모두 이 image representation vector 위에 classification head가 부착됩니다.)

Position Embedding

또한 각 패치 별 위치 정보를 유지하기 위해 Trainable한 Position Embedding이 각 패치 임베딩에 추가됩니다. (이 때 이미지를 위한 2D 임베딩이 1D에 비해 그닥 유의미한 성능 향상을 가져오지 않았기에 1D position embedding을 사용했다고 합니다.)

이렇게 구성된 최종 임베딩 벡터들의 시퀀스가 Transformer Encoder의 Input이 됩니다.

Transformer Encoder

트랜스포머 인코더는 Multiheaded Self-Attention (MSA) Layer와 MLP 블록이 교차되어 구성됩니다. 모든 Block 앞에 Layer-Normalization(LN)이 적용되며, 모든 블록 뒤에는 Residual Connection이 적용됩니다. 또한 MLP는 GELU를 activation으로 사용하는 2개의 layer를 포함합니다.

아래는 트랜스포머 일련의 과정을 수식으로 작성한 것입니다.

FINE-TUNING AND HIGHER RESOLUTION

논문에서는 ViT를 large dataset들에 pre-train하고 더 작은규모의 downstream task 데이터셋에 fine-tune합니다. 이 과정에서 pre-trained prediction head는 제거되고, zero-initialized D×K의 feed forward layer가 부착됩니다. (K는 downstream classification task의 class 개수)

최근연구들에서 pre-training시보다 더 높은 해상도로 fine-tune 하는것이 종종 benefical 하다는 사실이 밝혀졌습니다. 본 논문에서도 fine-tuning시 사전학습 데이터보다 higher resolution에서 진행되었으며, 이러한 고해상도 image를 전달할 때 patch size (P×P)는 유지하고, 트랜스포머 유효 input 시퀀스 길이인 N을 늘리는 방향으로 조절하였다고 합니다.

ViT는 어떠한 길이의 시퀀스도 처리할 수 있으나 이 경우 pre-trained position embeddings는 의미를 잃게됩니다. 따라서 이를 극복하기 위해 사전학습된 위치임베딩에 원본이미지에서의 위치를 기준으로 하는 2D interpolation이 적용됩니다.

이와 같이 해상도 조정 및 patch extraction이 image의 2차원 구조에 대한 inductive bias가 ViT에 수동적으로 주입되는 유일한 과정이라고 해석할 수 있습니다.

Experiment

ResNet, ViT, 그리고 하이브리드 모델의 representation learning capabilities를 평가합니다.

1. 각 모델의 data requirements를 이해하기 위해, 다양한 크기의 데이터셋으로 사전학습을 진행하고 많은 Task 에 대해서도 평가합니다.
2. 모델의 사전학습 계산 비용 측면에서 ViT가 다른 모델들보다 더 적은 pre-training Cost로 SOTA를 달성하며
3. 끝으로 self-supervised를 이용한 실험을 통해 self-supervised된 ViT의 가능성을 확인합니다.

실험 설정

데이터셋

학습데이터: 1. ImageNet-1k (1.3M) / 2. ImageNet-21k(14M) / 3. JFT-18k(303M)

모델 설정

• ViT-L/16 : “Large” 매개변수 크기 및 16×16 Input Patch 사용
  • Sequence Lenght : Patch의 크기 제곱의 반비례 → 패치가 작아지면, 계산량 증가
• ResNet : Batch Normalization → Group Normalization & Weight Standardization 사용
  • Sequence Lenght : Patch의 크기 제곱의 반비례 → 패치가 작아지면, 계산량 증가
• Hybrid : Intermediate Feature Map을 입력으로 사용 → Patch Size 1×1
  1. ResNet50 – 4 stage 사용
  2. ResNet50 – 3 stage 사용(단, Layer 수 유지) → 4배 더 긴 Sequence 길이 ( ∵ Feature map 크기 증가되므로 )

COMPARISON TO STATE OF THE ART

    JFT로 사전 학습 후, 각 데이터 셋에 Transfer Learning 진행한 결과입니다.

먼저 ViT variants 모델들 중 가장 큰 size의 ViT-H/14와 ViT-L/16을 SOTA인 CNN들과 비교합니다.

하나는 대규모의 ResNet으로 지도방식의 전이학습을 수행시킨 Big Transfer(BiT)모델이고,
두 번째는 큰 EfficientNet으로 준지도학습을이용해 학습시킨 Noisy Student입니다.

본 연구가 진행되던 시점을 기준으로 ImageNet에 대해서는 Noisy Student가 SOTA이고, 다른 데이터셋들에 대해서는 BiT-L이 SOTA였다고 합니다.

Table2를 보면 동일한 데이터셋에 대해 사전학습되었을 때, 모든 downstream task에서 ViT-L/16의 성능이 BiT-L을 능가하면서 동시에 사전훈련에 드는 계산상의 cost는 확연하게 낮음을 확인할 수 있습니다. 특히 더 큰 모델인 ViT-H/14은 모든 실험 데이터셋에 대해 더욱 향상된 성능을 보여줍니다.

놀라운 사실은 이 거대한 모델이 기존 SOTA 모델과 비교할 때 더 적은 pre-train 비용이 필요합니다.
(그러나 사전학습의 효율성이 데이터 구조뿐 아니라 training schedule이나 optimizer, weight decay 등의 다른 파라미터들에도 영향을 받는다고도 합니다).

끝으로 ViT-L/16을 JFT보다 작은 데이터셋인 ImageNet-21k에 pre-train 하였을 때에도 대부분에 데이터 셋에 대해 상당히 좋은 성능을 보이는 것을 확인할 수 있습니다.

Pre-Training Data Requirements

사전 학습 시 데이터셋의 크기에 관한 내용입니다.

• 사전학습에서의 데이터셋 크기
  • Pre-Train : ImageNet-1k(1.3 M) / ImageNet-21k(14 M) / JFT-18k(303M) → Fine-Tune ImageNet-1k

그 결과, 사전학습 시 데이터 셋이 클수록 ViT 좋고, 사전학습시 데이터 셋이 작으면 좋지 않았음을 알 수 있습니다.

그 이유는 앞서 언급했던 것처럼  Convolutaional Inductive Bias가 없기 때문에, 더 많은 데이터 필요합니다.

Scaling Study

Transformer, Hybrid, BiT의 모델 크기 실험입니다. 즉, Transfer Learnig 실험 입니다.

JFT-300M 데이터셋으로부터의 Transfer Learnig 성능을 평가하는데, 여기서는 각 모델의 performance vs pre-training cost를 평가합니다. (ResNet 7개와 ViT 6개, 하이브리드모델 5개의 모델). 상단 Figure 5는 그 결과를 나타낸 그래프인데요, 이를 통해 다음과 같은 사실을 알아낼 수 있었습니다.

첫번째로 모든 비전트랜스포머 모델이 ResNet 모델들을 성능/계산코스트 trade-off에서 압도한다는 것입니다. 동일한 성능을 달성하기 위해 드는 계산 비용이 ViT가 2 ~ 4배는 더 적습니다.

두번째로 하이브리드 모델이 비교적 작은 계산 영역에서는 ViT의 성능을 뛰어넘는다는 것입니다. 이 차이는 모델 사이즈가 커짐에 따라 vanishing 되긴 하지만.. 어느 크기의 ViT에도 convoluational local feature processing이 도움을 줄 수 있다는 것을 기대할 수 있기에 놀라운 결과라고 합니다.

마지막 세 번째로 ViT는 실험을 시도해본 여러 range 내에서 saturated 되지 않았기 때문에,
추후 연구에서 더 큰 모델을 가지는 feature scaling efforts를 기대할 수 있다고 하네요.

Inspecting Vision Transformer

• Left: Filters of the initial linear embedding of RGB values of ViT-L/32
  • 각 Filter의 기능이 저차원의 CNN Filter 기능과 유사
• Center: Similarity of position embeddings of ViT-L/32
  • position embedding 간의 유사성 확인
  • position embedding 유사도 내에서 이미지 내부의 거리개념을 인코딩하는 방법을 배운다는 것을 보여준다. 
  • 그 결과 가까운 패치 간의 유사도가 높음: 패치 간 공간 정보 학습이 잘됨
    즉, 가까운 패치들은 유사한 position embedding을 갖고,
    row-column 구조(같은 행/열에 있는 patch는 유사한 임베딩을 갖는다.)를 가진다
• Right: Size of attended area by head and network depth
  • self-attention은 ViT가 이미지 전체의 정보를 통합하여 사용할 수 있게 한다.
  • 정보가 attention weights에 의해 통합되는 image space내의 평균거리를 계산하였다
  • 그 결과 attention distance = receptive field 서로 유사한 개념이라고 볼 수 있다.
  • 낮은 layer의 slef-attention head는 CNN 처럼 “Localization:” 효과를 보인다는 것을 확인할 수 있음

Conclusion

Transformer를 Image Classification에 직접 적용하는 방법에 대해 알아보았습니다.

Vision에 self-attention을 사용하는 이전의 연구들과는 다르게 본 논문에서는 모델을 설계할 때, 이미지를 patch들의 시퀀스로 보고 NLP에서의 표준 트랜스포머의 encoder 부분을 거의 그대로 차용하였습니다. 간단하면서도 확장성 좋은 방법론이지만 큰 데이터 셋에 대한 사전학습이 있을 때 좋은 성능이 나온다는 것도 확인하였습니다. ViT는 이미지 분류 데이터셋에 대해 SOTA를 달성하면서도 사전학습 cost가 상대적으로 적다는 장점을 보입니다.

그러나 여전히 해결해야할 과제가 남아있습니다. 다른 Task 로의 적용, Self-supervised Pre-trained 의 성능 향상 등이 있습니다.

추가적으로 앞으로 사전학습 모델의 활용은 계속 커지겠지만, JFT-300M은 구글 사내에서만 접근 가능한 데이터이기에 사전학습된 모델은 거의 정형화되어 있다고 봐도 무방할 것 같습니다. 또한 학습에 필요한 많은 리소스는 극복해야할 분명한 과제임이 확실해보입니다.

---

결국 Transformer를 이해하는 게 우선이었습니다. 다만 이 논문을 이후로 Transformer를 접목시킨 방법론이 무수히 많이 나오고 있는 것으로 보아 계속 팔로업해야할 방법론임을 분명한 것같습니다.