ViT - An Image is Worth 16x16 Words Transformers for Image Recognition at Scale [2020]

의의

NLP에서 사용되는 Transformer를 Vision에 적용
기존의 제한적인 Attention Mechanism에서 벗어나, CNN구조 대부분을 Transformer로 대체
대용량 데이터셋을 Pre-train 후 Small Image 데이터셋에서 전이학습(Transfer Learning)
- 훨씬 적은 계산 비용으로 우수한 결과를 뽑음
- 단, 많은 데이터를 사전 학습해야 된다는 제한사항 존재

요약

이전 Vision Task에서 Self-Attention적용의 한계
- Self-Attention을 적용하는 시도는 있었으나 비효율적이였음
- 기존의 Transformer를 최대한 그대로 적용하고자 함
Transformer의 장점
- Parameter의 수가 많아도 학습 가능
- 계산 효율성 및 확장성이 높음
- 데이터셋이 크면 클수록 모델을 키워도 되고, 성능이 포화(Saturate)될 징후가 없다
Transformer의 적용 방안
- 이미지를 Patch로 분할 후 Sequence로 입력 (NLP에서 단어가 입력되는 방식과 동일)
- Supervised Learning 방식으로 학습
Transformer의 특징
- ImageNet과 같은 Mid-sized 데이터셋으로 학습시, ResNet보다 낮은 성능을 보인다.
- JFT-300M 사전 학습 후, Transfer Learning 하면 CNN 구조보다 매우 좋은 성능을 달성한다.
- Transformer는 Inudective Biases가 없다(=Locality와 Translation Equivariance와 같은 CNN의 특성이 없다)
ViT 구조
- 구조 자체는 BERT와 굉장이 유사하다.
- BERT에서도 [Class] Token을 따로 두고 Classficiation task에 사용했는데, ViT에서도 이와 동일하다.
Embedding
- Transformer에서는 Token들로 이뤄진 1D Sequence를 Embedding으로 한번 전처리를 하고, 그 결과에 Sinusoidal Positional Embedding을 추가한 것을 input으로 사용한다.
- ViT에서는 Image를 일정 크기의 Patch로 자르고, 이 Patch들로 Sequence를 구성한다
- Positional Embedding을 patch embedding에 추가해서, 위치 정보를 담고 있도록 한다

Transformer의 계산 효율성과 Scalability를 비전에 활용

표준 Transformer를 최대한 변형 없이 직접적으로 이미지에 적용하였다. 이를 위해 이미지를 패치 조각으로 쪼갠다. 그리고 각각의 패치에 대해 선형적인 임베딩의 시퀀스를 제공하여 Transformer에 입력하였다. 이때 이미지의 패치는 마치 NLP에서의 토큰과 같이 다루게 된다. 모델은 이미지 분류 태스크에 대해 지도 학습 방식으로 학습하였다.

Transformer는 CNN계열의 모델은 모델링할 수 있는 지역성(Locality)과 입-출력 위치의 동일성(Translation Equivariance) 등 이미지 이해에 필수적인 바이어스(Inductive Bias)를 학습하지 못하고, 따라서 데이터가 충분하지 못할 때 일반화 성능이 떨어지는 것

하지만 모델을 더 많은 데이터셋에 대해 학습했을 때는 다른 결과가 나왔다. 대량의 이미지로 학습하는 것이 CNN이 학습하는 Bias의 힘보다 강했다. 충분히 큰 스케일에서 ViT를 사전학습한 결과, 더 적은 데이터셋을 가진 하위 태스크에 전이 학습하여 좋은 성능을 얻을 수 있었다.

Model

이미지는 고정된 크기의 Patch로 쪼개고, 각각을 선형적으로 임베딩한 후 위치임베딩을 더하여 결과 벡터를 일반적인 Transformer 인코더의 인풋으로 입력한다. 분류 과제를 수행하기 위해 추가적으로 학습되는 classification token을 만들어 Sequence에 더한다.

Vision Transformer(ViT)

[이미지 인풋]

일반적인 Transformer는 토큰 임베딩에 대한 1차원의 Sequence를 입력으로 받는다
2차원의 이미지를 다루기 위해 논문에서는 이미지를 Flatten된 2차원의 patch의 Sequence로 변환하였다
즉, $H * W * C$ -> $N * (P^2 * C)$로 변환된다
- $(H, W)$는 원본 이미지의 크기
- $C$는 채널의 개수
- $(P, P)$는 이미지 패치의 크기
- $N = HW/P^2$는 패치의 크기
Transformer는 모든 Layer에서 고정된 벡터 크기 D를 사용하기 때문에 Image Patch는 Flatten한 다음, D차원 벡터로 Linear Projection을 시킨다.
BERT의 [CLS]토큰과 비슷하게 임베딩 된 Patch의 Sequence에 $z0 = x_{class}$ 임베딩을 추가로 붙여넣는다.
이후, 이 패치에 대해 나온 인코더 아웃풋은 이미지 Representation으로 해석하여 분류에 사용한다

[위치 임베딩]

각각의 Patch Embedding에 Positional Embedding을 더하여 위치 정보를 활용할 수 있도록 한다.
학습 가능한 1차원의 Embedding을 사용한다
2차원의 정보를 유지하는 Positional Embedding도 활용해 보았으나, 유의미한 성능 향상은 없었다.

[Hybrid Architecture]

Image Patch를 그대로 사용하는 대신, CNN의 결과 나온 Feature Map을 인풋 Sequence로 사용할 수 있다.
Hybrid 모델에서는 Patch Embedding Projection을 CNN Feature Map에서 결과로 나온 Patch에 대해 적용한다.
특수한 케이스로 Patch는 1x1 크기를 가질 수 있는데, 이 경우는 Input Sequence를 단순히 Feature Map에 대한 차원으로 Flatten한 후, Transformer의 차원으로 Projection한 결과이다.
[CLS]에 해당하는 Input Embedding과 Positional Embedding은 기존모델과 동일하게 적용한다.

Fine-tuning과 높은 해상도 이미지 다루기

ViT는 대량의 데이터셋에 대해 사전 학습한 후, 더 작은 다운스트림 태스크에 Fine-tuning하는 방법을 취한다.Fine-tuning시에는 사전 학습된 Prediction head를 제거하고, 0으로 초기화된 $D * K$차원의 FCL를 연결한다.

이때 Fine-tuning단계에서는 더 높은 해상도에서 학습하는 것이 정확도 향상에 좋다는 것이 알려져 있다. 더 높은 해상도의 이미지를 처리해야 할 경우, Image Patch크기를 동일하게 유지함으로써 더 긴 Patch Sequence를 사용한다.

ViT는더 높은 하드웨어의 메모리가 허용하는 한, 임의의 길이의 Sequence를 처리할 수 있다. 단, 이 경우 사전학습된 Positional Embedding이 의미없어진다. 이 경우, Pre-train된 Positional Embedding에 원본 이미지에서의 위치에 따라 2D Interpolation을 수행한다.

이러한 해상도 조절과 패치 추출 방식은 ViT에서 이미지의 2차원 구조에 대한 Inductive Bias를 수동적으로 다루는 유일한 포인트이다.

실험결과

Vision Transformer 탐색하기

[Embedding Projection]

ViT는 Flatten Patch를 더 낮은 차원의 공간으로 매핑한다.
아래 그림은 학습된 임베딩 필터 중 중요한 몇 가지 구성 요소들을 시각화한 것이다.

이러한 구성요소는 각각의 Patch에 대해 저차원의 Representation을 만드는 기본 함수들을 나타내는 것이다.

[Positional Embedding]

Linear Projection이후, 각각의 Patch Representation에는 Positional Embedding이 더해지게된다.
모델은 이미지 내의 거리 개념을 인코딩하여 Positional Embedding에서 유사성이 나타난다는 것을 알 수 있다.
즉, 가까운 거리에 있는 Patch는 비슷한 Positional Embedding을 가지게된다.
더 나아가 행-렬 개념의 구조가 나타나는데, 같은 열이나 행에 위치한 임베딩이 비슷하게 나타날것이다.
또한, 더 큰 grid에 대해서는 sinusoidal 구조가 더 명확하게 나타나는데, 이는 위치 임베딩은 2차원의 이미지를 나타내는 법을 학습한다는 것을 의미한다. 사람이 디자인한 2차원 구조를 인식하는 Embedding이 성능 향상에 기여하지 못한 이유가 여기에 있을 것으로 보인다.

[Self-Attention]

Self-Attention은 가장 밑단에 있는 Layer에서부터도 ViT가 전체 이미지에 있는 정보를 통합하도록 돕는다.
이러한 능력을 어느 정도까지 활용할 수 있는지 조사하기위해 Attention Weight에 기반하여 이미지 공간 상에서 정보가 취합되는 평균거리를 구해보았다.(Attention Distance는 CNN의 Receptive Field와 비슷하게 해석가능)

실험 결과, Attention Head 중 일부는 가장 낮은 Layer에서부터 대부분의 이미지에 집중하고 있고, 이렇게 글로벌하게 정보를 통합하는 능력을 모델이 활용하는 것으로 보인다.
또 다른 Attention Head는 밑단 Layer에서 일관적으로 작은 거리의 Patch에 집중하는 모습을 보였는데, 이렇게 지역적인 Attention은 하이브리드 모델에서는 좀처럼 나타나지 않는다. 즉, 이 Attnetion head는 CNN의 밑단에서 일어나는 것과 비슷한 작용을 하는 것이라고 유추할 수 있다.
Attention이 일어나는 거리는 네트워크의 깊이가 깊어질수록 늘어난다는 것도 찾을 수 있다.
또한 전체적으로 모델은 의미적으로 분류 과제에 필요한 부분에 집중하는것을 찾을 수 있다.