Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet, 2021
https://arxiv.org/abs/2101.11986
ViT는 이미지 분류를 위한 트랜스포머로, 이미지를 고정된 길이의 토큰 시퀀스로 잘라서 트랜스포머 레이어에 넣고 그들의 글로벌한 관계를 학습하고자 한다. 하지만 ViT는 ImageNet과 같은 중간 사이즈에서 학습하는 경우 CNN보다 좋지못한 성능을 가진다. 우리는 이것이 1) 인풋 이미지를 토큰화하는 방법이 중요한 local 구조인 edge나 line을 잡아내지 못하여 학습 효율성이 떨어지고 2) ViT의 불필요하게 있는 attention 연산들이 고정된 계산 연산과 학습 데이터 내에서 feature의 다양성을 뽑아내는 것을 제한하기 때문으로 보았다. 이러한 문제를 해결하기 위해 Tokens-to-Token ViT(T2T-ViT)을 제안한다. T2T는 1) layerwise Tokens-to-Token 트래스포메이션으로 주변의 토큰들을 하나의 토큰으로 반복적으로 aggregate하면서 이미지를 토큰들로 구조화 한다. 2) CNN 아키텍쳐에 영감을 받아 deep-narrow 구조를 통해 효율적인 backbone을 만들었다. T2T는 ViT의 파라미터와 MACs*을 반으로 줄였고, ImageNet을 스크래치부터 학습했을 때의 성능은 3.0% 높았다. ResNet보다 성능이 좋았으며, Mobilenet과 비교할만한 성능을 냈다.
*MACs: Multi-Adds, FLOP과 같다고 함. https://stackoverflow.com/questions/51077386/how-to-count-multiply-adds-operations
ViT는 이미지를 14 x 14 또는 16 x 16개의 패치로 나누어 토큰으로 만든 뒤 이후 언어모델 트랜스포머처럼 진행을 한다. 이러한 ViT는 작은 사이즈의 CNN(ResNet)에 비해 성능이 떨어지는데, 1) 인풋 이미지를 단순히 토큰화 하는 방식이 로컬한 구조를 뽑지 못하고 2) ViT의 어텐션 구조가 비전 태스크를 위해 잘 디자인 되지 못하여 불필요한 연산 들이 있어 feature가 적고 모델 학습을 하는데 어려움을 겪는다고 생각했다.
이러한 가설이 맞는지 확인하기 위하여, ViT-L/16과 ResNet50의 학습된 벡터를 시각화하였다. ResNet은 첫번재 레이어에서 중간 레이어까지 local structure(line, edge)를 잘 잡아내는 것으로 보였다(초록색 박스). 반면에 ViT는 글로벌한 관계는 잘 잡아냈지만, 구조 정보는 잘 잡아내지 못했다. 이는 고정된 길이로 단순 토큰화시켰기 때문으로 보인다. 또한 우리는 ViT의 많은 채널에서 0을 가짐을 볼 수 있었다(빨간색 박스). 이는 ViT이 ResNet과 달리 리치한 피쳐를 뽑아내지 못함을 알 수 있다.
T2T-ViT는 두개의 요소로 구성되어 있다.
이미지의 local structure information을 뽑고 토큰의 길이를 점차 줄이는 모듈
T2T module에서 온 토큰들의 글로벌한 attention을 잡아냄
T2T 모듈은 1) Re-structurization 2) Soft Split(SS) 단계로 나누어져있다.