Transformer가 CV에서도 적용될 수 있을까? 라는 질문에서부터 시작된 것으로 CNN에 의존하지 않고 이미지 패치를 순차적인 입력으로 처리하는 Transformer만으로도 이미지 분류 task에서 우수한 성능을 보일 수 있음을 강조한 방법론이다.
물론 기존 일부 연구들은 CV 분야에서 CNN에 self attention을 결합하거나 컨볼루션을 대체하려고 시도했지만 계산 규모가 너무 커지거나 컨볼루션 연산만큼 효율적으로 최적화된 구현체가 부족해 결과적으로 ResNet 계열이 여전히 SOTA를 달성했다.
본 논문에서는 Transformer에 최소한의 수정만 가하여 이미지에 직접 적용한다,
① 이미지를 여러 개의 패치로 분할하고
② 각 패치를 선형 임베딩하여
③ Transformer에 순차 입력으로 제공한다.
이미지 패치 자체를 NLP에서 단어 토큰 처리하듯 동일하게 취급한 것이다.
ImageNet 정도의 크기의 중간 규모 데이터셋에서 강한 regularization을 사용하지 않고 ViT를 학습하면 비슷한 크기의 ResNet과 비교했을 때 몇 퍼센트 포인트 낮은 정확도를 보인다. 바로 inductive biases 때문이라고 논문에서 이야기한다. Transformer가 CNN의 위치 불변성(translation equivariance)과 지역성(locality) 같은 inductive bias를 갖지 않아 데이터가 충분하지 않을 때 일반화 성능이 떨어진다.
하지만 학습 데이터 규모를 1400만 ~ 3억 장으로 늘리면 상황이 달라진다. 대규모 학습이 inductive bias보다 더 중요하게 작용하는 것이다. 본 논문에서는 ViT를 공개 ImageNet-21k나 자체 JFT-300M 데이터로 충분히 사전학습한 뒤 소규모 데이터셋에 전이 학습하면 여러 이미지 인식 벤치마크에서 최첨단 성능을 따라잡거나 능가한다고 한다.
💡inductive biases https://medium.com/@sanjithkumar986/inductive-bias-in-machine-learning-f360ea678a15 inductive biases는 귀납 편향으로 모델이 학습을 시작할 때 데이터에 대해 이미 가지고 있는 가정 (prior assumptions) 을 의미한다. 예를 들어 '이미지 속 패턴은 위치가 바뀌어도 같은 의미를 가진다.', '먼 픽셀들보다 인접한 픽셀들끼리 관꼐가 더 중요하다.' 와 같은 가정이 존재한다. 이런 가정들은 모델이 무작위의 모든 함수를 탐색하는 대신 돔 더 제한된 함수 공간에서만 해답을 찾도록 유도하므로 데이터가 적어도 효과적으로 학습할 수 있게 도와준다. CNN은 같은 필터를 이미지 전체에 슬라이딩하며 적용하기 때문에 다음과 같은 inductive bias를 가지게 된다. ① 위치 불변성(translation equivariance): 어떤 피처가 A 위치에 있던, B 위치에 있던 똑같이 인식할 수 있는 귀납편 https://stats.stackexchange.com/questions/208936/what-is-translation-invariance-in-computer-vision-and-convolutional-neural-netwo ② 지역성(locality): 컨볼루션 연산을 할 때 작은 영역만 놓고 연산을 하기 때문에 국소적인 패턴을 잘 포착하도록 구조화되어 있는 귀납편향
하지만 ViT의 경우엔 왜 데이터가 적으며 성능이 낮을까? ViT는 모든 패치가 모든 패치에 attention 연산을 하는 구조이기 때문에 어떤 거리에 있는 정보가 중요한지, 픽셀 간 관계가 공간적으로 어떻게 배치되었는지 학습 과정에서 알아내야 한다. 즉, CNN이 가지고 있는 inductive biases 가 없어서 데이터가 충분하지 않으면 일반화 성능이 떨어진다는 의미이다.
모델의 설계 같은 경우는 의도적으로 원본 Transformer를 거의 변경하지 않고 간단하게 구성함으로써 Transformer 아키텍처의 확장성을 입증한다.
표준 Transformer는 1D 시퀀스 형태의 토큰 임베딩을 입력으로 받는다. 본 논문에서는 2D 이미지를 처리하기 위해 해상도가 (H, W)인 이미지를 패치 크기 (P,P) 로 분할해 N개의 2D 패치로 만들고 각 패치를 펼쳐 P^2*C의 벡터로 변환한 후 linear projection을 거쳐 차원 D의 패치 임베딩을 얻는다.
이미지 크기2D 패치들패치의 개수
https://gaussian37.github.io/dl-concept-vit/
E는 linear projection 을 의미하고 x_{class}는 BERT의 [class] 토큰과 유사한 학습 가능한 임베딩으로 시퀀스 맨 앞에 추가된다. E_{pos}는 1D 위치 임베딩을 의미하고 패치 임베딩에 위치 정보를 더하기 위해 사용된다. 본 논문에서는 더 복잡한 2D 위치 인코딩도 사용해보았지만 성능 향상이 크지 않아서 표준 학습이 가능한 1D 위치 임베딩을 채택했다고 한다.
패치 임베딩 시퀀스 z_0은 Transformer 인코더의 입력이 되며 인코더는 MSA (multi headed self attention) 과 MLP 블록을 교대로 쌓은 구조이다. 각 블록 전에는 LayerNorm(LN)을, 후에는 residual connection을 적용하고 MLP는 두 개의 linear layer 사이에 GELU 비선형성 함수를 삽입한다.
인코더의 최종 출력 z_L^0 의 [class] 위치에 해당하는 벡터가 이미지 표현 y로 사용되며 pretraining과 finetuning 모두 해당 벡터 위 classification head를 붙여 학습한다. pretraining 시에는 hidden layer 하나를 둔 MLP. finetuning 시에는 단일 linear layer를 classifier로 사용한다.
CNN의 경우엔 전 layer에 걸쳐 locality, 2D 인접 구조, translation equivariance 라는 inductive biases를 내장하지만 ViT에서는 MLP만 locality와 translation equivariance 를 지니고 self attention은 global로 동작하여 CNN에 비해 inductive biases가 훨씬 적다. ViT의 위치 임베딩은 초기화된 시점에는 2D 좌표 정보를 전혀 담고 있지 않으며 오직 패치를 분할할 때 (→ 어떤 패치가 옆에 붙어있는지)와 해상도를 변경할 때 (→ pretraining 한 모델을 다른 해상도로 finetuning할 때 패치의 개수가 달라지므로 원래 학습된 위치 임베딩을 2D 보간 방식으로 확장 및 조정) 만 2D 구조를 알려주고, 초기화된 위치 임베딩은 2D 위치 정보를 담지 않아 모든 공간 관계를 학습으로부터 얻게 된다.
즉, 초기화된 E_{pos} 의 경우 순전히 각 패치 인덱스에 대응하는 임의의 벡터일 뿐 이 벡터들끼리의 attention 계산만으로 어떤 패치가 위, 아래, 왼쪽, 오른쪽에 있는지는 학습 과정에서 스스로 알아내야만 한다는 것이다.
또한 raw 이미지 패치 대신 CNN으로부터 얻은 feature map을 패치로 잘라 임베딩하여 입력 시퀀스로 사용할 수도 있다. 이를 Hybrid ViT라고 하며 실험에서 비교하는 방식으로도 사용한다.
일반적으로 ViT는 대규모 데이터셋에서 사전학습한 뒤 더 작은 데이터셋에 finetuning한다. finetuning 시 사용한 classification head를 제거하고 클래스 수 K에 맞춘 D * K 크기의 linear layer을 0으로 초기화한 head를 붙인다. 또한 pretraining 해상도보다 높은 해상도로 finetuning을 하면 성능이 향상되는 경우가 많다고 한다. 때문에 여기선 pretraining 된 위치 임베딩이 더 이상 적절하지 않기 때문에 2D 보간을 통해 위치 임베딩을 해상도에 맞게 조정한다.
ViT 실험은 다음과 같다.
Base와 Large는 BERT의 구성 그대로이다.
사용한 데이터는 다음과 같다.
ImageNet (ILSVRC-2012): 1,000개 클래스, 130만 장
ImageNet-21k: 2만1천개 클래스, 1,400만 장
JFT-300M: 1만8천개 클래스, 3억3백만 장
SOTA 모델들과 비교해보았을 때 ViT-L/16(JFT-300M 사전학습)은 같은 데이터로 학습된 BiT-L을 모든 과제에서 능가하면서 학습 비용은 훨씬 적었다. 또한 대형 ViT-H/14는 ImageNet, CIFAR-100, VTAB에서 성능이 더 향상되었고 적은 계산량으로 사전학습이 가능한 것을 볼 수 있다.
TPUv3-core-days: pretraining에 소요된 계산량 / 정확도의 경우 3번의 finetuning을 거친 후의 평균과 표준 편차
또한 그림 2를 살펴보았을 때 ViT-H/14는 Natural·Structured 그룹에서 BiT-R152x4 및 기존 기법들보다 우수했으며, Specialized 그룹에서는 상위 모델 간 차이가 크지 않은 것을 볼 수 있다.
그리고 ViT가 CNN보다 inductive bias가 적으므로 데이터셋의 규모가 얼마나 중요한지 실험해보았을 때..
ImageNet만으론 ViT-Large가 Base보다도 못 미쳤고 ImageNet-21k에선 비슷해졌으며 JFT-300M에선 대형 모델의 이점이 뚜렷하게 보이는 것을 알 수 있다. BiT 모델 분포(그림 3)와 비교하면 작은 데이터에선 BiT가 우수하지만, 대형 데이터셋에선 ViT가 앞선다.
데이터셋의 부분집합 실험으로 JFT-300M에서 무작위 9M, 30M, 90M 샘플로 사전학습한 뒤 few-shot 정확도를 측정(그림 4) 해보았을 때 작은 데이터에선 ViT가 과적합(overfit) 경향이 강해 ResNet보다 성능이 낮았으나 90M 이상에선 ViT가 더 우수한 결과를 보였다.
이를 통해 작은 데이터에서 CNN의 inductive bias가, 큰 데이터에서는 데이터로부터 직접 학습하는 것이 유리하다는 것을 알 수 있다.
확장성 연구의 일부로 JFT-300M 데이터에서 전이 성능을 측정하며 사전학습 계산 비용 대비 성능을 비교하였을 때 계산 대비 성능에서 ViT가 ResNet을 압도했다. 작은 사전학습 계산량에서는 하이브리드가 약간 우세하나 점점 차이가 사라지고 있는 것을 볼 수 있다. 여기서 하이브리드 모델은 raw 이미지 패치를 사용하는 것이 아니라 먼저 raw 이미지를 CNN에 넣어 feature map을 뽑아내고 이를 활용해 입력 시퀀스를 구성하는 모델이다.
또한 ViT가 내부적으로 이미지 데이터를 어떻게 처리하는지 이해하기 위해 내부 표현을 분석해보았을 때
가장 먼저 flatten한 패치 임베딩 행렬에 PCA를 진행한 상위 성분들을 보면 (그림 7의 left) 패치 내 세부 구조를 표현할 수 있는 basis 함수와 유사한 것을 볼 수 있다.
또한 위치 임베딩의 벡터 유사도로부터 (그림 7의 center) 가까운 패치는 임베딩이 유사한 거리 정보, 같은 행과 열 패치끼리 유사한 행/열 구조를 볼 수 있다. 이를 통해 2D 위치 정보를 임베딩이 직접 학습함을 보여준다.
마지막으로 self attention head 별로 공간 상 평균 attention 도달 거리를 계산했을 때 (그림 7의 right) 일부 head는 낮은 layer에서 global attention을 수행해 전역 정보 통합을 활용하고 다른 head는 local attention 패턴을 유지해 CNN의 초기 layer와 유사한 기능을 수행한다. layer 깊이가 늘어날수록 attention 거리가 증가하며 전체적으로 attention은 분류에 의미있는 시각 영역을 선택하는 것을 알 수 있다. (그림 6)