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

Yuan, Li, Yunpeng Chen외 7명|arXiv (Cornell University)|2021. 01. 28.

Multimodal Machine Learning Applications인용 수 32

한 줄 요약

이 논문은 Tokens-to-Token 모듈을 사용하여 이미지를 점진적으로 토큰화하고 깊고 좁은 백본을 갖춘 T2T-ViT를 제안한다. ImageNet에서 scratch로 학습하면서 더 적은 매개변수와 FLOPs로 ViT 및 강력한 CNN 베이스라인보다 더 높은 정확도에 도달한다.

ABSTRACT

Transformers, which are popular for language modeling, have been explored for solving vision tasks recently, e.g., the Vision Transformer (ViT) for image classification. The ViT model splits each image into a sequence of tokens with fixed length and then applies multiple Transformer layers to model their global relation for classification. However, ViT achieves inferior performance to CNNs when trained from scratch on a midsize dataset like ImageNet. We find it is because: 1) the simple tokenization of input images fails to model the important local structure such as edges and lines among neighboring pixels, leading to low training sample efficiency; 2) the redundant attention backbone design of ViT leads to limited feature richness for fixed computation budgets and limited training samples. To overcome such limitations, we propose a new Tokens-To-Token Vision Transformer (T2T-ViT), which incorporates 1) a layer-wise Tokens-to-Token (T2T) transformation to progressively structurize the image to tokens by recursively aggregating neighboring Tokens into one Token (Tokens-to-Token), such that local structure represented by surrounding tokens can be modeled and tokens length can be reduced; 2) an efficient backbone with a deep-narrow structure for vision transformer motivated by CNN architecture design after empirical study. Notably, T2T-ViT reduces the parameter count and MACs of vanilla ViT by half, while achieving more than 3.0\% improvement when trained from scratch on ImageNet. It also outperforms ResNets and achieves comparable performance with MobileNets by directly training on ImageNet. For example, T2T-ViT with comparable size to ResNet50 (21.5M parameters) can achieve 83.3\% top1 accuracy in image resolution 384$ imes$384 on ImageNet. (Code: https://github.com/yitu-opensource/T2T-ViT)

연구 동기 및 목표

• ImageNet과 같은 중간 규모 데이터에서 처음부터 학습할 때 순수한 트랜스포머 架構가 CNN보다 성능이 낮은 이유를 제시한다.
• 로컬 이미지 구조를 포착하고 토큰 길이를 반복적으로 축소하기 위한 Tokens-to-Token (T2T) 모듈을 제안한다.
• 피처의 풍부함을 개선하고 중복을 줄이기 위해 효율적인 깊고 좁은 ViT 백본을 설계한다.
• 라이크한 규모의 사전 학습 없이도 T2T-ViT가 ImageNet에서 CNN과 대등한 크기의 성능을 달성할 수 있음을 보여준다.
• ViT 백본에 CNN에서 영감을 받은 설계가 이점을 줄 수 있음을 보인다.

제안 방법

• 층별 Tokens-to-Token (T2T) 모듈을 도입하여 재구성화(Restructurization)와 소프트 분할(Soft Split)을 번갈아 수행하며 로컬 구조가 내재된 토큰으로 점진적으로 변환한다.
• 작은 숨김 차원과 더 많은 층을 갖춘 깊고 좁은 ViT 백본을 사용하여 매개변수와 MACs를 줄이면서 성능을 유지한다.
• 메모리와 연산을 관리하기 위해 T2T 모듈 내에서 Transformer 및 Performer 레이어를 실험한다.
• 비교 대상로 ImageNet에서 유사한 모델 크기로 ViT, ResNets, MobileNets와 비교한다.
• T2T 모듈과 깊고 좁은 아키텍처의 영향력을 정량화하기 위한 변형 실험을 수행하고 CIFAR-10/100으로의 전이도 탐구한다.

실험 결과

연구 질문

• RQ1 progressive tokens-to-token 모듈이 ImageNet에서 처음부터 학습된 ViT의 naive 토큰화보다 로컬 이미지 구조를 더 잘 포착할 수 있는가?
• RQ2 CNN에서 영감을 받은 깊고 좁은 백본이 표준 ViT에 비해 비전 트랜스포머의 중복을 줄이고 피처 풍부함을 향상시키는가?
• RQ3 ImageNet에서 scratch로 학습될 때 T2T-ViT가 비슷한 매개변수 수와 계산 예산에서 ResNets 및 MobileNets에 비해 어떤 성능을 보이는가?
• RQ4 다른 T2T 모듈 변형(Transformer vs Performer)이 성능과 효율성에 미치는 영향은?
• RQ5 사전 학습된 T2T-ViT 모델이 CIFAR-10/100과 같은 다운스트림 데이터셋으로 효과적으로 전이될 수 있는가?

주요 결과

모델	Top1-정확도 (%)	매개변수 (M)	MACs (G)
ViT-S/16 [12]	78.1	48.6	10.1
DeiT-small [36]	79.9	22.1	4.6
DeiT-small-Distilled [36]	81.2	22.1	4.7
T2T-ViT-14	81.5	21.5	4.8
T2T-ViT-14↑384	83.3	21.5	17.1
ViT-B/16 [12]	79.8	86.4	17.6
ViT-L/16 [12]	81.1	304.3	63.6
T2T-ViT-24	82.3	64.1	13.8
T2T-ViT t-14	81.7	21.5	6.1
T2T-ViT t-24	82.6	64.1	15.0

• 21.5M 매개변수와 4.8G MACs를 가진 T2T-ViT가 scratch로 ImageNet(224x224)에서 81.5% top-1 정확도를 달성하여 ViT-S/16을 능가하고 유사한 규모의 ResNets보다 더 나은 정확도를 보인다.
• 384x384 입력에서 T2T-ViT-14↑는 83.3% top-1 정확도를 달성해 해상도가 높아질 때 강력한 이득을 보여준다.
• ResNet50(25.5M 매개변수, 4.3G MACs)과 비교하면 T2T-ViT-14가 81.5% 정확도(또는 t-variant에서 6.1G MACs)를 달성하여 유사하거나 더 높은 계산으로 더 나은 정확도를 보여준다.
• T2T-ViT-24는 64.1M 매개변수와 13.8G MACs로 82.3% top-1에 이르러 더 큰 규모에서 경쟁력 있는 성능을 보여준다.
• Lite T2T-ViT 모델(T2T-ViT-7/12 등)은 MAC가 더 높음에도 MobileNets에 비해 경쟁력 있는 정확도를 제공하며, 증류를 통해 아주 작은 모델의 성능이 더 향상된다.
• 사전 학습된 T2T-ViT를 CIFAR-10/100으로 전이시키면 ViT 베이스라인 대비 경쟁력 있는 이득을 얻어 전이 가능성을 보여준다.