[평범한 학부생이 하는 논문 리뷰] An Image is Worth 16X16 Words: Transformers for Image Recognotion at Scale (ViT)

Transformer 아키텍쳐가 NLP에서 많이 쓰이지만 Vision 분야에서도 쓰인다는 것을 최근에 들었다. Transformer 논문을 최근에 리뷰했는데 이를 Vision 분야에서도 이용한 논문을 읽어봐야겠다는 생각이 들었다. 그 중에서 가장 대표적인 논문인 ViT를 읽게 되었다. 그래서 이번엔 ViT 논문에 대한 리뷰를 진행할 생각이다.

 

본 논문 리뷰는 원본 paper와 유튜브 거꾸로 읽는 AI 이야기의 발표 자료를 가지고 진행할 예정이다.

 

1. Abstract & Introduction

Transformer 아키텍처는 NLP분야에서는 표준이지만 Vision분야에서는 제한적으로 쓰이고 있다. 비전에서 attention는 컨볼루션 네트워크와 함께 적용되거나 전체 구조를 그대로 유지하면서 컨볼루션 네트워크의 특정 구성 요소를 대체하는 데 사용된다. 여러 다른 논문에서 이를 도전했지만 실패했지 이 논문에서는 Transformer 아키텍쳐만으로 이미지 분류에서 CNN구조의 성능을 뛰어 넘는 것을 보여준다. 데이터 수가 적으면 귀납적 편향이 부족해 성능이 떨어지지만, 대규모의 데이터셋으로 훈련을 시키고 적은 데이터 포인트의 작업에 적용하면 좋은 결과를 낸다고 한다.

 

다음으로 Vision Transformer에 대해 설명해보겠다.

 

2. Method

모델 설계는 가능한 원본 Transformer를 따랐다고 한다.

2.1 Vision Transformer

Vision Transformer의 전체 아키텍쳐

위의 이미지는 Vision Transformer의 전체적인 아키텍처를 나타낸 것이다. 이 아키텍처를 자세히 뜯어보면서 이해해보자.

이미지 $\mathbf{x} \in \mathbb{R}^{H\times W\times C}$가 input으로 주어지면 이미지 $\mathbf{x}$를 크기가 $P\times P$인 패치들로 나눈다. 패치 $N$개를 각각 $x^1_p \in \mathbb{R}^{P^2\cdot C}$로 flatten시킨다. 이후 $E \in \mathbb{R}^{(P^2\cdot C)}$로 linear projection 시킨다. 그 결과 $x^1_pE$는 위 이미지의 노란 박스의 분홍 부분을 나타낸다. 그리고 각 결과에 이미지의 어느 위치에 있는 패치인지를 나타내기 위해 learnable한 1-D 위치 임베딩을 더해준다. 여기서 분류 task를 위한 모델이기에 class를 나타내는 토큰을 앞에 추가한다. 이 과정의 결과를 수식으로 나타내면 아래와 같다.

$$\mathbf{z}_0 = [\mathbf{x}_{class}; \mathbf{x}^1_p\mathbf{E}; \mathbf{x}^2_p\mathbf{E}; \cdots ; \mathbf{x}^ N _p \mathbf{E}] + \mathbf{E}_{pos},\quad \mathbf{E} ∈ \mathbb{R}^{ (P^2 \cdot C)×D}, \mathbf{E}_{pos} ∈ \mathbb{R} ^{(N+1)×D}\quad (1)$$

위 이미지가 Transformer encoder의 구조를 나타낸다. 기존 Transformer의 구조와 다른 점은 Norm의 순서와 activation function으로 GELU를 사용했다는 점이다. ViT에서 사용된 Encoder는 위 이미지처럼 Layernorm을 하고 Multi-head attention과 MLP를 지나는 형태이다. 여기서 MLP는 GELU를 가진 2개의 레이어로 구성된다. 위에서 도출한 $\mathbf{z}_0$를 Encoder의 input으로 넣는다 그렇게 해서 나온 output을 수식으로 나타내면 아래와 같다.

$$\mathbf{z}^{\prime}_{\mathcal{l}} = MSA(LN(\mathbf{z}_{l-1})) + \mathbf{z}_{l-1},\quad l = 1, \cdots ,L\quad (2)$$

$$\mathbf{z}_{l} = MLP(LN(\mathbf{z}^{\prime}_l)) + \mathbf{z}^{\prime}_l, \quad l = 1,\cdots,L \quad (3)$$

Encoder를 총 L개를 두어 위 과정을 L번 반복한다. 그렇게 해서 L번째 Encoder에서 나온 output을 $\mathbf{z}_L$이라 한다. 최종 output에서 class를 나타내는 0번째 벡터 $\mathbf{z}_0$를 MLP헤드에 넣어 class의 확률을 구한다.

$$\mathbf{y} = LN(\mathbf{z}^0_L)\quad (4)$$

2.2 Fine-tuning and higher resolution

일반적으로 대규모 데이터 세트에 대해 ViT를 pre-training하고 (더 작은) 다운스트림 작업에 맞게 fine-tuning한다. 이를 위해 pre-trained된 prediction head를 제거하고 zero-initialized된 D × K 피드포워드 레이어를 연결한다(여기서 K는 다운스트림 클래스 수). Pre-training보다 더 높은 해상도로 미세 조정하는 것이 유리한 경우가 많다. 더 높은 해상도의 이미지를 공급할 때 패치 크기를 동일하게 유지하면 유효 시퀀스 길이가 더 길어진다. 이 때, pre-trained된 위치 임베딩은 더 이상 의미가 없을 수 있다. 따라서 원본 이미지에서의 위치에 따라 pre-trained된 위치 임베딩의 2D 보간을 수행한다. 이 해상도 조정과 패치 추출은 이미지의 2D 구조에 대한 귀납적 편향이 비전 트랜스포머에 수동으로 주입되는 유일한 지점이라고 한다.

3. Experiments

위 표와 같이 ViT는 Base, Large, Huge 3개의 구성으로 설정되고 실험이 진행된다. ViT-L/16은 입력 패치 크기가 16×16인 "Large" ViT를 의미한다.

JFT-300M에서 pre-trained된 더 작은 ViT-L/16 모델은 모든 작업에서 동일한 데이터 세트에서 pre-trained된 BiT-L보다 성능이 뛰어나면서도 훈련에 훨씬 적은 컴퓨팅 리소스를 필요로 한다. 더 큰 모델인 ViT-H/14는 특히 더 까다로운 데이터 세트인 ImageNet, CIFAR-100 및 VTAB에서 성능을 더욱 향상시킨다. 흥미롭게도 여전히 pre-training에 이전 sota 모델보다 훨씬 적은 컴퓨팅이 소요된다.

위 이미지들은 데이터 셋의 크기가 모델에 영향을 주는지에 관한 실험의 결과를 보여준다. 위 결과를 보면 데이터 셋의 크기가 작으면 기존의 CNN 계열 모델들을 이용한 BiT가 성능이 좋은 것을 볼 수 있다. 하지만 데이터 셋의 크기가 커지면서 ViT의 성능이 더 좋아지는 것을 볼 수 있다. 본 논문에 CNN은 inductive bias가 Vision Transformer보다 훨씬 많으므로 데이터 셋 크기가 작으면 이 점이 유리하게 작용한다고 설명한다. 그러나 데이터 셋의 크기가 커지면 그렇지 않은 듯 싶다. 

Inductive Bias에 대해 이 논문을 읽으면서 처음 들어봤는데 자세히 이해를 하지 못한 것 같다.

위 실험 결과를 통해 몇 가지를 관찰할 수 있다. 첫째, Vision Transformer는 성능/컴퓨팅 트레이드오프에서 ResNet을 압도한다. ViT는 동일한 성능을 달성하기 위해 약 2~4배 적은 컴퓨팅을 사용한다. 둘째, 적은 계산 예산에서는 하이브리드가 ViT보다 약간 더 나은 성능을 보이지만, 대규모 모델에서는 그 차이가 사라진다. 셋째, Vision Transformer는 시도한 범위 내에서 포화 상태가 되지 않는 것으로 나타나 향후 확장에 동기부여한다.

• Inspecting Vision Transformer

위 이미지에서 중간 이미지를 보면 그리드에서 행이나 열이 같으면 position embedding이 유사하게 나타나고 주변 그리드의 position embedding 또한 유사하게 나타나는 결과를 확인할 수 있다.

오른쪽 이미지를 보면 네트워크 깊이가 낮음에도 attention 거리가 먼 것들이 표현되는 것을 확인 할 수 있는데, 이 것은 얕은 레이어에서도 전체적인 이미지에 대한 정보를 가진다는 것을 의미한다. CNN계열은 앞 레이어들은 이미지의 전체적인 대략적인 특징을 포착하고 뒷 레이어들은 미세한 디테일들을 포착하는 것과는 대조적인 모습이다.

 

이외에도 다양한 실험을 진행했다고 한다. 논문에 설명되어있기 때문에 논문에서 참고하길 바란다.

4. Conclusion

기존 self-attention을 비전에 이용한 다른 모델들과 달리 Transformer encoder를 거의 그대로 이용해서 좋은 성능을 냈다는 것이 인상적이다. 게다가 기존 CNN계열 모델들보다 computational cost도 적다. 앞으로 vision transformer을 이용한 다른 논문들이 궁금해졌다.

 

많이 부족한 포스팅 읽어주셔서 감사합니다! 아직 많이 부족합니다. 많은 피드백 남겨주시면 감사하겠습니다~~