[논문 리뷰] Tokens-to-Token ViT : Training Vision Transformers from Scratch on ImageNet

Abstract / Introduction

ViT의 출현 이후 image classification을 비롯한 많은 vision task 분야에서 SOTA를 달성하고 있습니다. 과거에 CNN이 지배하던 vision 분야에서 ViT가 점점 더 우위를 차지하고 있는데요. 특히 large & big dataset에 대해서 좋은 성능을 보이고 있습니다. 

 

그러나 반대로 midsize 이하의 dataset 혹은 data scarce regime에서는 약한 모습을 보이는데요. 보통 이를 약한 inductive bias와 많은 parameter 때문이라고 분석합니다. 유사하게 이 논문에서도 이유를 다음과 같이 분석하네요. 

 

1) Simple tokenization

기존 ViT에서는 patch 단위로 image를 나눌 때 overlap이 없게 나눕니다. 그런데 이런 단순한 patch화가 edge, line과 같은 simple local structure 정보를 얻는데 방해가 된다는거죠. 보통 CNN에서는 이런 simple structure를 shallow layer에서 학습한다고 알려져 있습니다.

 

2) Redundant attention backbone

그림 1. Redundancy of ViT

그림 1을 보시면 CNN(ResNet50)의 feature map의 경우 local structure를 capture한 feature를 확인할 수 있고 zero feature들이 거의 없죠. 그러나 ViT의 feature들을 확인해보면 object의 global한 feature는 잘 잡아낸 것을 볼 수 있지만 local structure에 대한 feature은 찾기 힘듭니다. 게다가 빨간색으로 표시된 것과 같이 의미 없는 zero feature들이 많이 존재하고 있습니다. 

 

이런 단점들을 보완하기 위해서 이 논문에서는 T2T-ViT 구조를 제안합니다. 그림 1에서와 같이 ViT 기반의 structure임에도 local structure를 잘 캡쳐하고 zero feature들이 사라진 것을 확인할 수 있는데요. 자세한 방법을 다음 절에서 알아보겠습니다.

Method

그림 2. T2T ViT architecture

그림 2는 T2T ViT의 전체 구조도입니다. 이전 T2T block에서 출력된 token들을 re-structurization 후 soft split이란 과정을 거쳐 concat한 후 다음 T2T block으로 넣어주는 것을 확인할 수 있는데요. 각 step별로 자세히 알아보겠습니다.

 

1. Re-structuization

우선 각 token간의 local 정보를 반영하기 위해 token들을 원래 위치로 복원시켜주는 단계를 거칩니다. 예를 들면 그림 1에서 1번 token은 처음에 input image를 patch화 할 때 좌측 상단 첫 번째 patch에 해당하므로 그 위치에 놓아주는 거죠. 이렇게 re-struct한 token들은 그림 2의 $I_i$와 같이 공간 정보를 갖는 $d$-dimension 2D feature map이 됩니다. ($d$ = embedding dimension)

 

2. Soft-Split

Restructurization 과정을 거친 token들에서 이제 local information을 뽑아내야하는데요. $k\times k$ kernel을 overlap sliding하며 각 sliding마다 kernel에 속하는 token들을 concat하여 새로운 token을 만듭니다. 그림을 보면 이해가 쉬운데요. 빨간색 kernel로 sliding될 때 해당 kernel에 속해있는 1,2,4,5번 token이 concat 후 flatten 되어 새로운 token이 만들어집니다. 이런식으로 sliding을 진행하면 새롭게 생성되는 token의 개수는 다음 식 1과 같이 표현되겠죠. Convolution 연산시 생성되는 feature map size를 계산하는 방법과 유사합니다. 인접한 local token들을 하나의 token으로 병합해주면서 지역적인 정보를 더욱 강하게 만드는 방법이라고 볼 수 있겠습니다.

식 1. 새로 생성되는 token의 개수 $l_o$

그리고 새로운 token 각각의 length는 kernel 내에 기존 token $k^2$개가 들어가 있으니 $k^2d$가 되겠죠. ($d$:embedding dimension) 결과적으로 다음 T2T block의 input 크기는 $T_O\in R^{l_o\times dk^2} $ 가 되겠구요.

 

3. T2T transformer

앞서 설명한대로 restructurization과 soft split이 이루어진 결과를 input으로 하여 동작하는 transformer가 바로 T2T transformer입니다. 동작을 식으로 나타내면 다음 식 2와 같겠네요. 한 module에서 $T_i$라는 input이 multi head self attention, MLP layer를 거친 후 위에서 설명한 reconstruct와 soft split을 거쳐 다음 module로 들어갑니다.

식 2. T2T module의 동작

이 과정에서 token의 개수는 줄어들지만 각 token의 dimension이 크게 늘어나게 됩니다. 이를 보완하기 위해서 첫 T2T module의 input으로 각 token의 dimension을 매우 작게 설정합니다(32 or 64). 이렇게 T2T module을 전부 통과하고 나면 일정한 크기의 token들이 출력으로 나오게 됩니다.

그림 3. 전체 구조도

그러고 나면 이 new token들을 backbone transformer의 input으로 사용할 수 있는데요. 그러니까 T2T module의 의의는 patch를 분할하는 과정에서 1. local structure와 2. adjacent patch와의 연관성을 반영하는 역할을 한다고 볼 수 있겠습니다. 정말 간단히 말하면 더 좋은 patch embedding을 만든다고 볼 수 있겠네요. 

 

그리고 ViT backbone에 대해서도 vanilla ViT를 쓰는 것이 아니라 여러 변형을 거친 ViT를 사용합니다. DenseNet, WRN, SENet, ResNeXt, GhostNet 등의 기존 CNN 기반 network들의 idea를 따와서 변형한 architecture를 사용하는데요. 이 논문들의 내용을 모두 짚어보기엔 무리가 있으니 Appendix에 소개된 적용 방법과 해당 논문들을 비교하시며 읽어보시면 쉽게 이해가 되실 것 같습니다. 정말 직관적으로 아이디어를 반영해서 어렵지 않습니다. 그런데 웃긴건 성능표를 보면 SENet의 idea 이외에는 오히려 성능에 해가 되네요..

Result

표 1. ImageNet 성능

표 2. Transfer learning 성능

표 1을 보면 redundancy를 제거했기 때문에 훨씬 적은 parameter 수로도 기존 ViT 기반 model들보다 좋은 성능을 냈음을 확인할 수 있네요. 표 2에서 확인할 수 있듯이 역시 transfer learning에서도 더 좋은 성능을 보입니다. 그리고 위에서 말씀드린 CNN 기반 model들에서 가져온 다양한 idea들에 대한 성능 분석도 논문에 존재하니 확인하시면 좋겠습니다. 앞서 말씀드린것처럼 성능에 오히려 해가 된다고 합니다.

Conclusion

개인적으로 이 논문의 의의는 이 두가지라고 생각합니다.

1. Patch embedding (tokenize) 과정 개선 

2. 성능 향상은 없었지만 CNN based model에서의 idea를 ViT에 결합하려고 시도

특히 2번처럼 기존 idea를 새로운 framework에 적용 시도하는게 성능과 상관 없이 꽤 흥미롭네요. 

 

이상 

Yuan, Li, et al. "Tokens-to-token vit: Training vision transformers from scratch on imagenet." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2021.

논문 리뷰였습니다. 감사합니다.