[논문 정리] An image is worth 16x16 words : Transformers for image recognition

논문정보

An image is worth 16x16 words

:Transformers for image recognition

 

An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

While the Transformer architecture has become the de-facto standard for natural language processing tasks, its applications to computer vision remain limited. In vision, attention is either applied in conjunction with convolutional networks, or used to rep

arxiv.org


논문정리

Abstract

자연어 처리 분야에서 트랜스포머 구조가 사실상의 표준이 되었지만 컴퓨터 비전에 응용하기에는 여전히 한계가 있었다.

비전 영역에서는, 어텐션을 컨볼루션 네트워크와 혼합하여 사용하거나 전체적인 구조는 유지하되 특정 요소를 대체하는 식으로 사용하였다.

본 연구진은 굳이 CNN에 의존하지 않고 순수 트랜스포머를 이미지 패치의 시퀀스에 적용하여도 이미지 분류 작업을 잘 수행할 수 있음을 보여주었다.

대용량의 데이터에 대해 사전훈련한 뒤 중/소형 크기의 벤치마크 이미지 데이터셋(ImageNet, CIFAR-100, VTAB, etc)에 전이학습했을 때 Vision Transformer(ViT)는 convnet의 SOTA에 견줄만한 결과를 보여주면서도 연산리소스는 적게 사용하였다.

 

Introduction

셀프 어텐션을 기반으로한 구조(특히 트랜스포머)는 자연어처리 분야에서 널리 쓰이는 모델이 되었다.

이는 대규모 텍스트 말뭉치(corpus)에 대해 사전훈련 한 후 더 작은 크기의 text-specific 데이터셋에 대해 파인튜닝하는 방식이다.

트랜스머의 연산 효율성과 확장성 덕에 과거엔 다룰 수 없었던 크기의 데이터로도 모델을 훈련할 수 있게 되었다.

모델과 데이터셋이 더 증가하여도 성능이 손실되지 않았다.

 

반면 컴퓨터비전 영역에서는 컨볼루션 구조가 여전히 우세하였다.

NLP가 성공하면서 많은 연구들이 셀프 어텐션을 CNN구조와 결합하기도 하고, 전체 컨볼루션 구조를 대체하기도 하였다.

후자의 경우, 이론적으로는 효율적이지만 독특한 어텐션 패턴 탓에 현대 하드웨어 가속기에 확장되지 못했다.

따라서,  대규모 이미지 인식 작업에 경우엔 클래식한 ResNet 구조가 여전히 SOTA!

 

본 연구진은 standard Transformer를 최대한 변형하지 않고 그대로 이미지에 적용하였다. 

먼저 이미지를 패치로 분할하고 이 패치들의 선형 임베딩 시퀀스를 트랜스포머의 입력으로 준비한다.

이미지 패치는 NLP 작업에서 토큰(단어)와 같은 개념

 

이미지넷과 같은 중형 데이터셋으로 훈련했을 때는, 비슷한 크기의 ResNet모델보다 해당 모델의 정확도가 몇퍼센트 정도 뒤쳐졌다. 

Why? 트랜스포머에는 translation equivalance, locality와 같이 CNN에 내제된 inductive bias가 없기 때문에 불충분한 데이터로 학습했을 때는 잘 일반화 되지 않기 때문.

Inductive bias란 모델이 처음보는 입력에 대한 출력을 예측하기 위해 사용하는 가정

위치 (i,j)가 주어졌을 때, 컨볼루션 가중치는 상대적인 이동 (i-j)만 고려한다.
이것을 translation equivalance라고도 하는데 해당 위치의 값이 바뀌면 출력도 바뀌는 성질이다.
이로 인해 컨볼루션의 경우 제한된 데이터셋에서 일반화 성능이 좋다.
ViT의 경우 absolution positional embedding을 이용하여 위 property가 부족하다.
(따라서, 데이터셋이 엄청 크지 않으면 ConvNet이 트랜스포머보다 낫다)

또한, 컨볼루션 연산시, 이미지 전체에서 커널이 이미지의 일부분만 보게 된다. 이 특정 영역만으로 컨볼루션 필터가 특징을 추출하는데 이를 locality라 한다.

반면, 트랜스포머는 attention구조를 사용하여 데이터 전체를 보고 attention할 위치를 정한다. 

inductive bias가 적다는 것은, 그만큼 이미지 전체에서 정보를 얻을 수 있다는 장점이 있지만 optimal parameter를 찾을 때 공간이 커져버리기 때문에 데이터가 충분하게 많지 않으면 학습이 잘 안되는 것이다.
따라서, ViT는 JFT-300M(3억)에 의해 사전학습된 가중치를 사용한다.

 

반면, 모델이 대용량 데이터셋(14M-300M장)을 훈련한다면 얘기가 달라진다.

이 대규모 학습이 inductive bias를 뛰어넘는 다는 것.

ViT는 충분한 크기의 데이터로 사전훈련한 후, 더 작은 데이터셋에 전이되었을 때 훌륭한 결과를 보였다.

ImageNet-21K이나 JFT-300M 데이터셋으로 훈련했을 때 ViT는 SOTA를 달성하거나 그에 준하는 성능을 보였다.

 

Related Work

skip

 

Method

기존의 Transformer를 가능한 그대로 따를 수 있도록 모델을 설계하였다.

 

<Vision Transformer - ViT>

Fig 1. 모델 개요

기존의 트랜스포머는 입력으로 1차원의 토큰 임베딩 시퀀스를 받았다.

2D 이미지를 다루기 위해서, 이미지 $\mathbf{x}\in \mathbb{R}^{H\times W\times C}$ (3차원)를 flatten하여 2D 패치의 시퀀스 $\mathbb{x}_p \in \mathbb{R}^{N\times (P^2 \cdot C)}$로 변환하였다.

$(H,W)$ : 원본 이미지 해상도
$C$ : 채널 수
$(P,P)$ : 이미지 패치 해상도
$N=HW/P^2$ : 패치 수 (트랜스포머에 입력으로 들어가는 시퀀스 길이 - 하이퍼파라미터)

 

식 1~4

트랜스포머의 모든 레이어에서 latent vector 크기를 D로 고정하였으며 따라서 패치를 flatten한 후 trainable linear projection을 통해 D차원이 되도록 매핑한다. (식 1)

이 projection을 거친 결과벡터를 패치 임베딩이라 칭한다.

 

BERT에서 [cls] 토큰과 유사하게, 임베딩된 패치들의 맨앞에 학습가능한 임베딩을 덧붙였다.

$\mathbf{z}^0_0 = \mathbf{x}_{class}$

트랜스포머 인코더의 최종 출력 $\mathbf{z}^0_L$에서 이미지 표현 $\mathbf{y}$로 여겨진다. (식 4)

이미지에 대한 1차원 representation vector로써의 역할 수행

 

트랜스포머 인코더의 출력 또한 클래스 토큰과 벡터로 이뤄져있는데, 여기서 클래스 토큰만 사용하여 위 아키텍쳐의 MLP Head를 구성하고 이 MLP head를 통하여 MLP하면 최종적으로 클래스를 분류함

 

pre-training, fine-tuning 동안 이 classification head가 $\mathbf{z}^0_L$에 붙어있는데 사전학습시에는 1개의 은닉층을 가진 MLP로 구현되고 fine-tuning시에는 단일 선형 레이어로 구현된다.

 

이미지의 위치 정보를 유지하기 위해 패치 임베딩에 1D의 위치 임베딩이 추가된다.

더 개선된 2D 위치 임베딩을 사용하여도 성능이 유의미하게 향상되지 않았다고 한다.

이제 이 결과 시퀀스가 트랜스포머 인코더에 입력으로 들어간다.

 

트랜스포머 인코더는 멀티헤드 셀프 어텐션과 MLP 블록이 교차되어 구성되어 있다.

LayerNorm(LN)이 각 블록이 시작 되기 전에 적용되고 각 블록 이후에 Residual connection을 진행

MLP는 GELU(Gaussian Error Linear Unit)을 활성함수로 사용하는 2개의 레이어를 포함한다.

 

Inductive bias

ViT가 CNN에 비해 이미지별 inductive bias가 훨씬 작다.

CNN은 locality, 2차원의 neighborhood structure, translation equivariance이 전체 모델의 각 레이어에 적용되지만 ViT의 경우 MLP 레이어만 local하고 translationally equivariant하고 셀프 어텐션 층은 global 하다.

2차원 neighborhood structure는 매우 드물게 사용된다. 모델이 시작할 때 이미지를 패치로 자르고 fine-tuning시에 다른 이미지에 대한 위치 임베딩 조정.

이 외에는 초기화 시에 위치 임베딩이 2D 위치에 대한 정보를 전달하지 않으며, 패치 간 모든 공간관계는 처음부터 학습해야 함

 

Hybrid Architecture

원시 이미지 패치를 사용하는 것 대신에, 입력 시퀀스를 CNN의 특징맵으로부터 형성한다.

하이브리드 모델에서는, 패치 임베딩 프로젝션 $\mathbf{E}$ (식1)이 CNN의 특징맵에서 나온 패치에 적용된다.

하나의 경우로 패치의 공간 사이즈가 1x1일 수 있는데, 이는 입력 시퀀스가 특징맵의 차원들을 flatten한 뒤 트랜스포머의 차원으로 프로젝션 된 것을 의미한다.

그 후 위에서 처럼 classification input embedding, position embedding 진행

 

<Fine-Tuning and Higher Resolution>

일반적으로 ViT는 대용량 데이터셋에 대해 사전훈련한 후, 더 작은 downstream task에 대해 파인튜닝 한다.

이를 위해, 사전훈련된 prediction head가 제거되고 0으로 초기화된 $D\times K$의 feed forward 레이어가 부착된다.

$K$ : downstream classification task에서 클래스 수

파인 튜닝 시에는 사전훈련 에 사용된 데이터보다 더 높은 해상도로 진행하는 것이 좋다.

이때, 패치 크기 ($P \times P)$는 유지하고, 트랜스포머의 입력 시퀀스 길이인 $N$을 늘렸다.

ViT는 임의의 시퀀스 길이를 처리할 수 있지만 이 경우엔 pretrained position embedding의 의미가 없어진다.

사전학습된 위치 임베딩이 파인튜닝에 적용될 때 원본이미지에서의 위치를 기준으로  2D interpolation을 적용

위의 해상도 조정과 패치 추출은 2D 이미지에 대한 inductive bias가 수동적으로 주입되는 유일한 순간이다.

2D 위치에 대한 inductive bias가 없는 트랜스포머에게 위치에 대한 개념을 주입시키는 것