[논문] 최근 AI의 이미지 인식에서 화제인 "Vision Transformer"에 대한 해설

※ 일본 블로그 내용을 번역한 것으로 오역이나 직역이 있을 수 있으며, 내용의 오류 지적해주시면 감사하겠습니다.

 

1. 개요

---

현재 AI계에서 화제가 되고 있는 "Vision Transformer"에 대해 다뤄보려고 한다. 이번에 제안된 Vison Transformer(이하 ViT)이라는 모델에 의해 더 이상 레이어를 겹치는 형태의 모델은 없어질지도 모른다. 

ViT는 Transformer을 거의 그대로 이미지 분류 업무에 이용한 것으로, ImangeNet/ImageNet-ReaL/CIFAR-100/VTAB  SoTA모델과 거의 비슷한 정도 혹은 그 이상을 성능을 달성하였다.  심지어 SoTA모델들(BiT/NoisyStudent)와 비교해서 계산 비용은 15분의 1까지 떨어졌다 (SoTA를 향상시킨 모델은 많아 봤자 4분의 1 혹은 5분의 1밖에 계산 비용을 줄이는데 그쳤지만 말이다). 

이번에 다룰 논문은 ICLR2021의 OpenReview에서 등록되어 아직 리뷰중에 있지만, NLP에서 Transformer가 RNN을 구축한 것 처럼, 이미지 인식에서도  Transformer가 CNN을 구축하지 않을까하고 주목하고 있는 상태이다. 그렇다면 본격적으로 기대를 받고 있는 새로운 모델 ViT에 대해 살펴보자.

논문 : "An Image is Worth 16x16 Words' Transformers for Image Recognition at Scale(2021)"  

코드 : PyTorch(비공식)

약칭	이름
ViT	Vision Transformer
BiT	Big Transfer
LN	Layer Normalization
SA	Self Attention

 

 

2. 바쁜 사람들을 위해

---

- 완전히 Convolution과 안녕하고 SoTA를 달성했다

- Vision Transformer에서 중요한 것은 세 가지이다.

① 이미지 패치(patch)를 단어와 같이 다루었다.

② 아키텍처는 Transformer의 엔코더 부분이다.

③ 거대한 데이터 세트인 JFT-300M으로 사전학습했다.

- SoTA보다 뛰어난 성능을 약 15분의 1의 계산 비용만으로 얻을 수 있었다.

- 사전학습 데이터 세트와 모델을 더욱 크게 하여 성능을 향상시킬 수 있는 여지가 있다.

 

 

3. Vision Transformer의 해설

---

Vision Transformer(=ViT)의 중요한 부분은 세 가지이다.

1) 입력 이미지

2) 아키텍처

3) 사전학습과 파인 튜닝

이것에 대해 각각 자세히 살펴보자.

 

1.1 입력 이미지

 제일 먼저 입력 이미지에 대한 내용이다. ViT는 Transformer을 베이스를 한 모델이지만, 이러한 ViT에 이미지를 어떻게 입력한다는 것일까?

 먼저 기본적으로 Transformer에서는 각 단어가 벡터로 표현되어 있는 문장을 한 번에 입력한다.

이러한 Transformer에 이미지를 넣어 다루고 싶은 경우 어떻게 하면 좋을까?

답은 간단하다. 아래의 이미지와 같이 이미지를 패치로 나눠서 각 배치를 단어처럼 다루는것이다.

 각 패치를 단어처럼 다루기 때문에 실제는 패치를 벡터에서 Flatten한다. 아래의 그림에서 왼쪽에서 부터 2번째의 패치에 대해 Flatten를 실행했을 때의 모습을 보여주는 것이다. 이처럼 모든 패치에 대해 동등하게 벡터와 Flatten처리하여 엔코더에 입력한다. 

 마지막에는 이것을 기호를 사용하여 나타낸다. 원래 이미지가 아래와 같을 때,

Flatten한 뒤부터 ViT에 입력하는 것은 아래와 같이 벡터로 처리하여 입력하는 것이다. 

여기서 N은 패치 수이고, P는 패치의 크기를 의미한다. 여기서 패치는 정사각형이므로, 즉 N은 다음과 같이 표현할 수 있다.

위 이미지의 원본 이미지의 사이즈가 256라고 한다면, N은 4, P는 128이 된다. 다음은 아키텍처에 대해 살펴보자.

 

1.2 아키텍처

1.2.1 아키텍처의 설명

아키텍처는 거의 Transformer의 엔코더부분과 동일한다(즉 BERT이다). 엔코더의 구조는 다음과 같다. 참고를 위해 오리지널의 Transformer의 엔코더도 오른쪽에 같이 넣어두었다.

 여기서 Norm은 Layer Normalization(=LM)을 의미한다. 위 그림을 보고 알 수 있듯이, 거의 오리지널의 엔코더와 동일하다. 오리지널 엔코더와 다른점은 두 가지이다.

① Pre-Norm : Norm이 Muti-Head Attention / MLP의 전에 위치해있다.

② GELU : MLP는 2단으로 활성화 함수로 GELU를 채용하고 있다(오리지널은 ReLU를 사용하지만 말이다).

 위 그림에서 ViT Encoder의 입력을 보면 "Embedded" Patches라고 되어 있다. 즉 제 1 블록에서의 ViT Encoder의 입력은 앞서 설명한 Patch 벡터들이다. 이것은 오리지널 Transformer에서도 마찬가지이다. 여기서 삽입하는 것은 아래의 두 가지이다.

 

(ViT는 Self-Attetion뿐이므로, 즉 패치의 위치 정보가 없으므로) 위치 엔코딩

 구체적으로는 아래와 같이 사용한다

- E에 각 패치을 길이 (P제곱・C)로 부터  D의 벡터에 삽입

- 각 패치에 위치 엔코딩  Epos를 계산

 여기서 혹시 위치 엔코딩의 Epos의 식에서 같이 N이 아닌 (N+1)로 되어 있는 것을 알아차린 사람들이 있을지도 모른다.이것은 E에 입력한 후에 입력의 맨 부분에 [CLS] 토큰을 연결하기 때문이다. 이것은 BERT와 완전히 동일하다. 다시 말해, [CLS] 토큰에 위치 엔코딩을 더하기 때문에 "패치 + [CLS] = N + 1"이 되는 것이다.  

 한편 위치 엔코딩도 학습 가능하다. 이미지 분류에서는 [CLS] 토큰의 출력을 식별기(MLP)에 넣는 것으로 최종적인 예측이 나온다. 여기까지의 내용을 한꺼번에 정리하자면 아래의 그림이 된다.

1.2.2 수식으로 정리

 여기까지의 설명을 수식으로 나타내보려고 한다. 수식이라고 해도 여기까지의 내용을 이해했다면 큰 문제 없을 것이다.

1.2.2.1 패치의 입력과 위치 엔코딩

 xkp 는 k번째의 패치를 표시하고 있다 (수식 입력이 안되서 xkp로 작성한 부분 양해 부탁드립니다) . 위의 식에서 볼 수 있듯, 각 배치는 xkp를 E에 곱해 [cls] 토큰을 연결한 후 위치 엔코딩 Epos를 더한다. 한편 E의 대신에 ResNet에 각 패치를 넣어는 것도 좋아보인다. 이 경우, 패치는 Flatten되지 않고 ResNet에 입력하여 그 출력에 대해 Flatten을 실시한다. 논문 속에서는 패치의 최초의 입력에 ResNet을 이용한 방법을 하이브리드라고 부르고 있다. 

1.2.2.2 엔코더

식(2)는 멀티 헤드 Attention을 표시하는 것이고 식(3)은 MLP를 나타내는 것이다. 제대로 스킵 결합도 실시한다. 

1.2.2.3 MLP 헤드

 마지막으로 식(4)이다. z0L는 최종 단계의 출력에 있어서 이전부터 0번째의 벡터 표현이므로 즉 [cls]코튼의 최종 출력이다. 이것을 LM에 넣어, y를 얻어낸다. MLP 헤드 자체의 식은 논문에 나오지는 않지만, 후에 이 y를 MLP에 넣는 것으로 최종적인 예측까지 나오는 것 같다. 이것으로 아키텍처에 대한 설명의 끝이다.

 

1.3 사전학습과 파인 튜닝

거대한 사전 학습 > 파인 튜닝

 여기서는 ViT의 3번째 중요한 특징인 사전학습과 파인 튜닝에 대해 설명할 것이다. ViT에는 지금까지의 이미지 인식 모델과 동일하게 '거대한 데이터 세트로 사전학습 + 파인 튜닝'이라는 단계를 거친다. 파인 튜닝을 할 때에는 ViT의 MLP헤드를 대체한다. 이 외에 세 부분정도 아이디어를 더하고 있다.

- 사전학습시에 해상도(e.g 224)보다도 파인튜닝시의 해상도를 높게(e.g384) 한다.

- 패치의 크기는 사전학습과 파인 튜닝 둘 다 일정하다( 즉, 파인튜닝시에는 해상도가 높으므로 패치의 수가 증가한다).

- 사전학습에서 학습된 위치 엔코딩을 파인 튜닝에서 부족한 부분에 삽입하여 보완한다.

  이제 ViT가 얼마나 대단한지를 실제로 살펴보자.

 

 

2. Vision Transformer의 실험결과

---

2.1 실험 whrjs

ViT의 실험은 ImageNet 소규모 데이터 세트를 사용하였다.

1) 데이터 세트

(1) 사전학습

① ILSVRC-2012 ImageNet ; 소규모, 클래스 1000개의 합계 130만장

② ImageNet-21K ; 중간규모, 클래스 21000개의 합계 1400만장

③ JFT-300M ; 대규모, 클래스 18000개의 합계 3억장.  Google에 있는 프라이베이트 데이터 세트

(2) 파인 튜닝

① ImageNet

② ImageNet-ReaL ; 보다 정확히 어노테이션된 평가용 라벨

③ CIFAR-10/100

④ Oxford-IIT Pets

⑤ VTAB ; 모델의 전이학습을 평가하는 19개의 이미지 데이터 세트군

    - Natural ; 일반적인 이미지 데이터 세트. 7개가 여기에 소속되어 있다.

    - Specialized ; 의료 데이터 등이 전문적인 이미지 데이터 세트. 4개가 여기에 소속되어 있다.

    - Structured ; 인공 3D 공간 등의 이미지 데이터 세트

2) 모델

 모델은 Base(B)/ Large(L) / Huge(H) 세 개를 이용하였다. 층의 수 등도 BERT를 따랐다. 그래고 Huge의 파라미터 수 6억은 정말 괴물이다.

	사전학습	파인 튜닝
옵티마이저	Adam	momentum이 딸린 SGD
배치 사이즈	4096	512
이미지 사이즈	224	384

- ImageNet은 파인튜닝에서만 이미지 사이즈 512나 518을 사용했다.

- ViT-H에 있어서만 가중치를 0.9999의 이동 평균으로 되어있다. 

- 모델은 'ViT-(1)/(2)' 라는 이름으로 표시되어 (1)에는 모델 사이트 B/L/H가 있고, (2)에는 패치의 크기가 16이나 14등으로 되어 있다. ViT-L/16이라면 ViT-Lage이며 입력 이미지의 패치 한 개의 크기가 16인 모델이라는 의미이다.

실험은 크게 세 가지로 나눌 수 있다.

- SoTA와의 비교

- 사전학습의 고찰(데이터 사이즈나 계산 비용)

- ViT의 고찰

 

2.2 SoTA과의 비교

 먼저 ViT-H/14와 ViT-L/16를 JFT-300M으로 사전학습시킨 경우의 결과를 SoTA모델의 BiT-L과 NoisyStudent와 비교하고 있다. SoTA모델들은 물론 CNN이다. 위 표를 살펴보면 일목요연하게 정리되어 있듯, 거의 모든 Vision Transformer가 SoTA를 갱신하고 있는 것을 알 수 있다. CIFAR-100에 있어서도 1% 갱신하고 있다. 단지, ImageNet에 있어서 Noisy Student의 수준에 미치지 못하고 잇지만, 보다 정확하게 어노테이션된 ReaL에서는 ViT의 쪽이 더욱 좋은 성능을 보여주었다.

 또한 ViT는 계산 시간도 대폭 감소시켰다. 위의 표의 TPUv3-days는 JFT-300M에의 사전학습에 걸린 시간을 표시하고 있다. BiT과 NoisyStudent에서는 약 1만 TPUv3-일이 걸렸지만, ViT-H/14에서는 약 4분의 1에서 5분의 1정도 시간이 절약되어 ViT-L/16에 있어서는 BiT-L보다 약간 더 좋을 성능을 약 15분의 1일 정도만에 달성하였다.

 더욱이 VTAB의 Natural/Specialized/Structed 각각의 성능을 비교하고 있다. 비교에는 ViT-H/14와 BiT-L의 다른 VTAP에 있어서 전부터 SoTA였던 VIVI과 S4L의 네 개를 사용하고 있다. 위 그림을 살펴보자면 VTAB 전체/ Specialized/Structed에서 ViT-H/14 가 제일 좋은 성능을 내고 있는 것을 알 수 있다. 

 Natural에 있어서는 BiT-L이 약간 더 높지만, 논문에서는 이것을 오차의 범위 내라고 설명하고 있다. 이렇듯 ViT는 폭넓게 전이학습에 굉장히 효과적임을 알 수 있다.

 

2.3 사전학습의 고찰

 방금 JFT-300M이라는 거대한 데이터 세트로 사전학습한 결과, 훌륭한 ViT로 SoTa를 달성하는 것이 가능하다는 것을 살펴보았다. 여기서는 사전학습에 사용한 데이터 세트의 사이즈가 어느정도 중요한 것인가 그리고 위에서 언급한 Hybrid(=입력 패치를 ResNet에 삽입한 ViT)와 비용대비 효율을 비교해볼 것이다. 

2.3.1 사전학습에서의 데이터 세트 사이즈

먼저 사전학습의 데이터 세트로써 ImageNet(작은 사이즈) / ImageNet-21k(중간 사이즈) / JFT-300 (거대 사이즈) 각각의 이용한 경우의 ImageNet에 있어서 파인 튜닝의 성능을살펴보자.

비교로써 BiT에의 결과도 표시하고 있다. 여기서 말할 수 있는 것은 ViT는 작은 데이터 세트에서의 사전 학습에서는 별로 좋은 성능을 보여주지 않지만, 거대한 데이터 세트에서는 진가를 발휘한다는 것이다. ViT의 사전학습에는 거대한 데이터 세트가 필요한 듯하다.

 계속해서 사전학습의 JFT-300M을 작게한 경우에도 실험해보았다. 여기서는 JFT-300M중9M/30M/90M/300M를 사용한 경우의 결과는 아래와 같다. 

 ViT-b란 ViT-B의 폭(즉, 숨겨진 레이어의 차원)을 반으로 줄인 것으로 보인다. 여기서에서도 아까와 같은 경향을 볼 수 있다. 90M이상에서 ViT가 진가를 발휘하고 있다는 것을 알 수 있다. 한편 9M에서는 ViT는 좋지 않다는 것을 알 수 있다. 여기까지 살펴보고 말할 수 있는 것은 Convolution이 가지고 있는 '국소성'과 같은 바이어스는 데이터 수가 적은 경우에 유효하지만 데이터수가 큰 경우에는 불필요하다는 것이다. 

2.3.2 Hybrid과의 비용대비 효과 비교

 Hybrid와의 비용대비 효과 비교를 살펴보자. 다양한 크기의 ResNet(5개) / ViT(6개) / Hybrid(4개)를 JFT-300M으로 사전학습시켰다. 이 때의 계산량과 전이학습에서의 정밀도를 표로 나타낸 것이 아래의 그림이다. 전이학습의 결과는 Average-5(왼쪽 그림)과 ImageNet(오른쪽 그림)이다. 여기서 Average-5가 어떤 5개인가 설명하는 부분이 논문에는 존재하지 않지만, ImageNet/CIFAR10/CIFAR100/Pets/Flowers 5개의 데이터 세트에 대한 평균치라고 생각된다.

 이 그림으로부터 알 수 있는 것은 세 가지이다.

① ViT은 BiT보다 비용대비 효과가 좋다(저비용으로도 같은 성능을 달성한다). 

② 비용이 한정되어 있는 경우는 Hybrid가 제일 효과적이다(높은 비용을 사용할 수 있는 경우 Hybrid도 ViT도 차이가 거의 없어진다). 

③ ViT에 대한 더 업그레이된 성능향상을 기대할 수 있다.

 

2.4 ViT의 고찰

ViT에 대해서 이해하기 위해 더욱 상세한 분석을 시시하였다. 여기서는 다음의 4가지 분석이 실시되었다.

① ViT의 내장 레이어

② ViT의 위치 엔코딩

③ Attention의 적용 범위

④ 자기교사학습

2.4.1 ViT의 내장 레이어

 ViT에서는 처음으로 내장 층을 이용하고 있다. 이 내장 층으로 학습한 것을 가시화한 것이 위 그림이다. 일체 Convolution 사용하지 않았음에서 가로선이나 세로선 등 CNN의 낮은 레이어에서 학습하는 것과 같이 확실히 학습할 수 있다는 것을 알 수 있다. 

2.4.2 ViT의 위치 엔코딩

 ViT에서는 잠재 층이후에 학습가능한 위치 엔코딩을 계산한다. 위치 엔코딩이들이 얼마나 비슷한지를 나타내는 것이 위 그림이다. 한 번에 알아보기 힘드므로 예를 통해 이 그림을 어떻게 보는지 설명하도록 하겠다.

 원래 위 그림은 ViT-L/32를 이용했을 때의 위치 엔토딩을 나타내고 있으며, 위 그림에서도 1장의 이미지에 대해서 패치수가 전부 49개 (7x7)라는 것을 알 수 있다(즉 원래 이미지의 크기는 32x7 = 224이므로 244x244라고 추정하였다). 여기서 위 그림 안의 위치(1,1) (즉 제일 왼쪽 첫번째)에 있는 1매스를 살펴보자. 이 그림만을 따로 빼내자면 아래와 같다.

 이것은 위치(1,1)의 위치 엔코딩과 다른 위치들의 위치 엔코딩들의 유사도를 전부표시한 것이 된다. 위 그림을 아래와 같이 편집하면 더욱 알기 쉬울 것이라고 생각한다. 

 이것은 위에서 설명한 것과 같이 위치(1,1)와 다른 위치의 엔코딩과의 유사도를 표시한 것으로, 위 그림의 (1,1)은 자기자신이므로 물론 1(노란색)로 표시되어 있다. 한편 위치(7,7)은 멀기 때문에 확실히 유사도가 낮다(진한 녹색). 이 동일한 행/열의 위치 엔코딩끼리는 비슷한 경향은 어느 위치에서도 볼 수 있다. 이를 통해 바로 같은 행/열의 위치 엔코딩끼리 비슷한 값이 되도록 학습 되고 있는 것을 나타내고 있다.

2.4.3 Attention의 적용 범위

  ViT는 Attension을 Pull에 사용한다. Convolution과 비교하여 Attention의 장점 중 하나로 "전체를 조감할 수 있다"를 꼽을 수 있다. 실제로 ViT는 Attention을 이용해 이미지 전체를 조감하고 있는 것일까? 이것을 조사하기 위해 먼저 논문에서는 'Attension거리' 라는 것을 정의하고 있다. Attension거리란 Attention스코어에 의한 Attention 적용 위치의 가중화로, 감각적으로는 말하자면 어떤 위치로 부터 먼 위치만 Attention 스코어가 크고 스코어가 큰 경우는 Attention 거리가 커지며, 반대로 가까운 위치라 Attention 스코어가 편중되어 있는 경우 Attention거리도 줄어든다.  즉 Attention거리는 CNN으로 말하자면, "Receptive Field의 크기"가 된다. 

 이 Attention거리를 각 층의 헤드마다 표시한 것이 아래의 그림이다. 가로 축이 각층, 세로 축이 Attention 거리로 점이 헤드를 나타내고 있다.  

 왼쪽 그림이 ViT-L/16이고 오른쪽은 Hybrid를 나타낸 것이다. 이 그림으로 부터 알 수 있는 것은 아래의 세 가지이다.

① 낮은 레이어에서는 Attention거리가 작은 헤드가 일정 수 존재한다. 즉, 낮은 레이어에서는 CNN와 같은 국소적인 특징량을 추출하고 있다.

② 낮은 레이어에서도 Attention 거리가 커다란 헤드가 일정 수 존재하다. 즉, 낮은 레이어에서는 전역적인 특징량도 추출하고 있다.

③ Hybrid의 낮은 레이어에서는 Attention 거리의 작은 헤드가 급감하는 것으로, 낮은 레이어에서의 Attention거리가 작은 헤드들은 CNN과 같은 역할을 달성하고 있다는 것을 알 수 있다. 

2.4.4 자기교사학습

 마지막으로 논문에서는 JFT-300M의 라벨을 Pull에 이용한 사전 학습대신에 BERT를 모방해 Masked Patch Prediction에 의한 사전학습도 실시하였다. 이것에 의해 Masked Patch Prediction으로 사전학습한 ViT-16/B가 ImageNet을 스크래치에서 학습시킨 모델보다도 2%정도 높은 성능을 나타내는 것처럼 보이지만, JFT-300M의 라벨을 Pull에 사용한 사전학습보다는 4%낮은 결과를 보여준 듯하다. 단지, BERT을 모방했을 뿐인 Masked Patch Prediction에서 성능향상이 확인되었으므로 보다 적당한 자기교사학습에 의한 사전학습 업무를 찾는 것으로 앞으로 JFT-300M과 같은 라벨이 붙은 비용이 상당한 데이터 세트를 반드시 필요로 하지 않고 높은 성능을 얻을 수 있는 가능성이 있다고 할 수 있다.

---

참고자료

qiita.com/omiita/items/0049ade809c4817670d7