An Image Is Worth 16X16 Words:Transformers For Image Recognition At Scale
https://arxiv.org/pdf/2010.11929
0. Abstract
NLP에서 Transformer가 사실상 standard가 되었지만 computer vision에는 적용이 제한되어있다. 하지만 저자들은 image classification에서 CNN없이 image patch의 sequence를 적용한 pure transformer를 성공적으로 보여주었다. 많은 양의 data를 pre-train하고 중간이나 작은 사이즈의 image로 구성된 benchmark를 사용했을 때, Vision Transformer(ViT)는 상대적으로 작은 계산량을 필요로하면서 최신 CNN보다 우수한 결과를 보여주었다.
1. Introduction
self-attention을 기반으로 하는 아키텍처(Transformer)는 NLP분야에서 주요 모델이 되었다. transformer의 계산효율성, 확장성 덕에, NLP분야는 전례없는 크기의 모델을 학습 가능하게 했다. 하지만 computer vision에서는 여전히 convolutional 아키텍처가 우세하다. 저자들은 transformer에서 영감을 받아, 최소한의 수정을 거쳐 표준 transformer를 그대로 image에 사용하는 방안을 제시한다. image를 patch로 나누고 이 patch에 linear embedding의 sequence를 transformer 입력으로 사용하였다. image patch는 NLP의 token과 동일한 용도로 사용된다. 정규화없이 ImageNet같은 mid-size 데이터셋을 학습시킬 때, ViT는 ResNet보다 정확도가 약간 떨어진다. Transformer는 CNN 고유의 translation equivariance, locality 같은 inductive bias가 부족하기 때문에 불충분한 양의 데이터를 학습시킬 때 일반화가 잘 이루어지지 않는다. (translation equivariance : CNN의 필터는 입력 이미지의 특정 위치에서 패턴을 감지하기 때문에 이미지가 이동해도 동일하게 처리된다. locality : CNN의 필터는 이미지의 작은 부분(local patch)만 처리하며 전체 이미지를 한번에 보지 않는다. 이러한 특징 때문에 CNN은 적은 데이터로도 높은 성능을 발휘할 수 있다. 하지만 Trnasformer는 모든 패턴을 데이터로부터 직접 학습해야 하기 때문에 학습데이터양이 적으면 일반화 성능이 떨어진다.)
하지만 많은 데이터셋을 학습시키면 inductive bias를 능가한다. 충분한 양의 데이터로 pre-train했을 때 Vision Transformer(ViT)는 좋은 성능을 보인다.
2. Related Work
transformer는 기계번역에서 제안되었고, 많은 NLP분야에서 최신 방법으로 자리잡고 있다.
attention 아키텍처를 computer vision에 접목시키려 시도하고 있지만 hardware accelerator에서 효율적으로 구현하려면 복잡한 엔지니어링이 필요하다.
ViT와 관련된 모델로 입력이미지를 2*2 패치로 나누어 full self-attention을 하는 방법이 있다. 이는 ViT와 유사하지만 ViT는 많은 양의 pre-train을 통해 vanila transformer를 CNN의 최신 성능과 경쟁할 수 있게 한다. 또한 다른 방법은 오직 작은 해상도의 이미지만 사용하지만 ViT는 중간 해상도 이미지도 처리할 수 있다.
저자들은 표준 ImageNet dataset보다 큰 규모에서 image recognition을 탐구하며, 추가적인 데이터는 표준 benchmark에서 좋은 성능을 얻을 수 있다.
3. Method
3.1. Vision Transformer(ViT)
표준 transformer는 1차원의 token을 받는다. 2차원의 이미지를 다루기 위해 image를 flatten하여준다.
여기서 (H, W)는 이미지 해상도, C는 채널수, (P, P)는 이미지 패치 해상도, N = HW / P^2는 패치 갯수이자 transformer의 입력 sequence 길이이다. transformer는 모든 레이어에서 일정한 벡터 크기 D를 사용하기 때문에 패치를 flatten해야하며, linear project을 사용해 D차원으로 매핑한다. 이 projection의 출력을 patch embedding이라고 한다.
BERT의 [class] 토큰과 유사하게, 임베딩 sequence 앞에 학습가능한 토큰을 추가한다.(초기화)
이 임베딩의 상태는 transformer encoder의 출력으로 나타나고 pre-train, fine-tuning동안 classification head의 입력으로 사용된다. classification head는 pre-train 시에는 MLP로 구현하고 fine-tune 시에는 단순 선형 레이어로 구현된다. (학습)
위치정보를 주기 위해 patch embedding에 position embedding을 추가한다. 이 때 저자들은 2차원 위치 임베딩과 1차원 위치 임베딩의 큰 차이가 없어 1차원 임베딩을 사용한다고 언급했다. 이 sequence embedding의 결과는 encoder의 입력으로 사용된다. (식 1)
transformer encoder는 multiheaded self-attention(MSA)와 MLP로 구성되어 있다. 참고로 MLP는 2개의 layer와 GELU로 구성되어 있다. (식 2, 3)
Layernorm(LN)은 모든 블록 전에 적용하고(MSA, MLP로 들어가기 전에 데이터의 분포를 조정하여 과도하게 큰 값이나 작은 값으로 인해 학습에 영향이 가지 않도록 해주는 과정), residual connection은 모든 블록 후에 적용한다.(정보 손실 방지, gradient vanishing 문제 해결) (식 4)
Inductive bias
Vision Transformer는 CNN에 비해 inductive bias가 매우 적다. CNN에서는 locality, 2차원 neighborhood structure, translation equivariance는 각 layer에 내장되어있다. ViT에서는 MLP 레이어에만 local, translation equivariance가 있다. 2차원 neighborhood structure는 매우 드물게 사용된다.
Hybrid Arhitecture
raw image patch의 대안책으로 입력 sequence를 CNN map의 feature로 사용할 수 있다. 이 추출한 feature 패치들을 patch embedding projection E(식 1)에 적용한다. 이를 transformer의 입력으로 변환하여 학습한다. (많은 데이터가 없으면 일반화가 잘 되지 않는 Transformer의 성능을 개선하기 위해 CNN의 강점을 활용한 방법 제시)
3.2 Fine-tuning and higher resolution
일반적으로 ViT는 큰 데이터셋에서 pre-train하고 작은 데이터셋에서 downstram에 대해 fine-tune을 한다. 이를 위해 pre-train한 prediction head를 제거하고 0으로 초기화한 D * K feedforward layer(K는 downstream class수)를 추가한다. fine-tune할 때는 더 높은 해상도의 데이터셋을 사용하는 것이 성능이 더 좋기도 하다. 이 때 더 높은 해상도를 사용하면 패치 크기는 동일하지만 sequence 길이는 증가한다. 그렇게 되면 pre-train한 position embedding이 의미가 없어진다. 그래서 원본 이미지의 위치에 따라 pre-train한 position embedding의 2차원 보간을 수행한다. 이러한 해상도 조정과 패치 추출이 inductive bias가 수동적으로 주입되는 유일한 구간이다. (해상도를 조정할 때 inductive bias를 주입할 수 있으니 고해상도를 썼을 때 성능이 더 좋아질 수도 있다.)
4. Experiments
ResNet, ViT, hybrid 모델을 비교 실험했다. 각 모델의 데이터 요구사항을 이해하기 위해, 다양한 크기의 데이트셋에서 pre-train하고 benchmark작업을 수행하였다. pre-train 계산량을 고려했을 때 ViT가 cost가 낮으면서 최신 성능을 달성했다.
데이터셋의 필요조건을 확인하기 위해 실험해본 결과 Figure 3, 4에서 보시다시피 데이터셋이 작은 경우 CNN(inductive bias)이 유용하지만, 데이터셋이 클 경우 관련 패턴을 데이터셋에서 직접 학습하는 것이 더 효과적이다는 것을 확인할 수 있다.
저자들은 ViT가 classification과 관련있는 구역에 집중하는 것을 Figure 6과 같이 확인했다.
Figure 7의 왼쪽을 보면 embedding 필터의 상위 component를 보여준다. 이는 각 패치내의 fine structure에 대한 저차원 representation의 basis와 유사하다. (패치 임베딩 필터에 PCA를 적용하여 상위 component를 얻은 것인데, 이 상위 component가 basis function처럼 보인다. 즉 패치의 구조적 특성을 적절하게 요약한 저차원 representation이다.)
projection 후, position embedding은 patch representation에 추가된다. Figure 7 중앙을 보면 모델이 position embedding의 유사성에서 이미지 내 거리를 인코딩한다는 것을 보여준다. (가까운 patch는 더 유사한 position embedding을 가진다.) 또한, 행렬 구조를 나타낸다. (같은 행,열을 가지는 패치는 유사한 임베딩을 가진다.)
5. Conclusion
이미지 인식에 transformer를 그대로 적용하는 방법을 알아보았다. sequence 패치로 이미지를 해석하고, NLP에서 사용하는 표준 transformer encoder를 그대로 사용했다. 이를 통해 ViT는 image classification 데이터셋에서 낮은 학습비용으로 최신성능을 달성하였다.