고정 헤더 영역
상세 컨텐츠
본문 제목
[5주차 / 황민아 / 논문리뷰] An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
본문
1. Introduction
Self-attention 기반의 모델들, 그 중에서도 Transformer는 NLP에서 유의미한 성과를 거뒀다. 보편적인 방법은 많은 양의 text corpus에 대해 pre-train을 진행한 후 더 작은, task-specific한 데이터셋에 fine-tuning을 하는 방법이다. Transformer의 연산 효율성과 확장 가능성 덕분에 100B 이상의 parameter를 갖는 큰 모델들을 학습시키는 것이 가능해졌으며, 모델과 데이터셋이 지속적으로 증가함에도 모델 성능의 포화가 이루어지지 않고 있다.
하지만 CV 분야에서는 convolution 기반의 모델이 아직까지는 우세하다. NLP의 성공에 따라 CNN 기반의 구조와 self-attention을 결합하거나 convolution을 attention으로 대체하려는 시도들이 있었지만 제대로 성과를 내지 못하였다.
이에 따라 본 논문에서는 Transformer를 약간의 변형을 통해 직접 이미지에 적용해보고자 하였다. 이미지를 patch로 나누고, 각 patch에 대한 linear embedding sequence를 입력으로 제공하였다. 이미지 patch를 NLP에서의 token으로 취급하여 image classification을 수행하게끔 학습시킨다.
2. Related Work
- Transformer, BERT
3. Method
3.1. Vision Transformer
NLP Transformer 구조는 scalable하며(확장 가능하며) 구현도 효과적이기에 transformer를 최대한 차용하여 모델을 설계하였다.
모델의 전체적인 구조는 다음과 같다.
- 이미지를 고정된 크기의 patch로 쪼갠다.
- 각각의 patch를 linearly embed하고, positional embedding을 추가하여 sequence of vectors를 만든다.
- Classification을 수행하기 위해 학습가능한 “classification token”을 sequence에 추가해주고, 이를 Transformer Encoder에 입력으로 넣는다.
여기서 classification token은 BERT의 class token과 유사하다. 학습가능한 임베딩 x_{class}으로, transformer encoder를 통과한 후 출력 단계에서는 이미지를 대표하는 y의 상태를 갖게된다.
Transformer encoder는 multiheaded self-attention layer와 MLP 블록으로 구성되며 그 식은 (2), (3)과 같다. Layernorm은 모든 블록 이전에 residual connection은 모든 블록 이후에 적용된다. MLP에서는 활성화 함수로 GELU를 사용한다.
Inductive Bias
Inductive bias란 보지 못한 데이터에 대해서도 귀납적 추론이 가능하게 하는 알고리즘이 가지고 있는 가정의 집합이다. ViT는 CNN에 비해 image-specific inductive bias가 훨씬 적다.CNN의 경우 모든 층이 locality와 translational equivariance를 갖는 반면 ViT에서는 MLP 층만이 local&translationally equivariant하고, self-attention 층은 global하기 때문이다.
Hybrid Architecture
Raw image patches 대신에 CNN의 feature map을 사용하여 input sequence를 만들 수도 있다. 이 경우엔 patch embedding projection E가 CNN feature map에서 추출된 patch에 적용된다. 특별한 경우에는 1x1 사이즈의 patch를 가질 수 있으며, 이는 feature map의 공간적 차원을 단순히 flattening함으로써 input sequence를 얻을 수 있음을 의미한다.
3.2. Fine-tuning and Higher Resolution
일반적으로 ViT는 대량의 데이터셋으로 pre-train시키고 더 작은(downstream) task에 맞게 fine-tune을 하여 사용한다. 이를 위해 pre-trained prediction head(=MLP head)를 제거하고 0으로 초기화된 D x K의 feedforward 층을 부착한다. 이 때 K는 downstream task의 class 개수이다.
Fine-tuning 시에는 pre-training보다 높은 해상도의 이미지를 사용하는 것이 효과적이다. 높은 해상도를 사용하는 경우에도 patch의 size는 같게 유지하는데, 이는 sequence의 길이 증가로 이어진다. ViT는 메모리의 한계까지 sequence 길이를 늘릴 수 있지만 pre-trained된 positional embedding은 의미를 잃게 된다. 따라서 pre-trained position embedding에 대해 원본 이미지에서의 위치에 기반하여 2D interpolation을 진행한다.
4. Experiments
SOTA 모델들과 비교해보았을 때 pre-trained된 ViT 모델들이 ResNet 기반의 모델들보다 더 좋은 성능을 보이는 것을 확인할 수 있다.
크기가 가장 작은 Imagenet부터, ImageNet-21k, 가장 큰 JFT-300M까지, 데이터셋의 크기에 따라 ViT에 어떠한 영향이 있는지 확인한 것이다. Pre-training 데이터셋의 크기가 작은 경우 좋은 성능을 보이지 못하였다. 이는 ViT의 inductive bias가 상대적으로 적기 때문이다.
ViT는 ResNet보다 성능/계산량 trade-off 측면에서 더 우세하였다.
좌측: Linear Projection 과정에서 학습된 embedding filter 중 top principal components를 나타낸 것이다. 이는 CNN에서의 low-level representation을 위해 사용되는 함수(필터)와 유사한 모양을 띈다.
중앙: 각 patch 간 position embedding의 코사인 유사도를 나타낸다. Patch가 가까울수록 position embedding이 유사하며, 같은 row/column에 있는 patch들은 더 유사한 embedding을 갖는다. 즉, 기존의 1차원 positional embedding은 2차원의 위상 배치를 표현할 수 있으므로 2차원 이미지를 표현하기에 충분하다.
우측: Network depth와 mean attention distance를 나타낸 것이다. Attention distance는 CNN에서의 reception field와 유사한데, ViT는 self-attention을 통해 가장 낮은 층에서조차도 전체 이미지에 대한 정보를 통합할 수 있다.
댓글 영역