[논문 리뷰] All Tokens Matter: Token Labeling for Training Better Vision Transformers

Zihang Jiang, Qibin Hou외 6명|arXiv (Cornell University)|2021. 04. 22.

Advanced Neural Network Applications참고 문헌 56인용 수 140

한 줄 요약

이 논문은 비전 트랜스포머를 위한 조밀한 패치-토큰 감독Objective인 토큰 라벨링(token labeling)을 도입하여 ViT 변형 전반에 걸쳐 일관된 이득을 제공하고, 더 작은 모델에서 높은 정확도를 가능하게 한다. 26M 매개변수에서 ImageNet의 Top-1이 84.4%, 150M에서 86.4%를 보고하며, 다운스트림의 의미론적 분할 성능도 향상된다.

ABSTRACT

In this paper, we present token labeling -- a new training objective for training high-performance vision transformers (ViTs). Different from the standard training objective of ViTs that computes the classification loss on an additional trainable class token, our proposed one takes advantage of all the image patch tokens to compute the training loss in a dense manner. Specifically, token labeling reformulates the image classification problem into multiple token-level recognition problems and assigns each patch token with an individual location-specific supervision generated by a machine annotator. Experiments show that token labeling can clearly and consistently improve the performance of various ViT models across a wide spectrum. For a vision transformer with 26M learnable parameters serving as an example, with token labeling, the model can achieve 84.4% Top-1 accuracy on ImageNet. The result can be further increased to 86.4% by slightly scaling the model size up to 150M, delivering the minimal-sized model among previous models (250M+) reaching 86%. We also show that token labeling can clearly improve the generalization capability of the pre-trained models on downstream tasks with dense prediction, such as semantic segmentation. Our code and all the training details will be made publicly available at https://github.com/zihangJiang/TokenLabeling.

연구 동기 및 목표

• ViTs가 모든 이미지 패치 토큰에 대해 조밀하고 위치 특이적인 감독을 활용함으로써 훈련 방식을 재고한다.
• 이미지 분류를 다수의 토큰 수준 인식 문제로 전환하여 객체 바인딩 및 인식을 강화한다.
• 토큰 라벨링의 일반성과 효과를 ViT 변형과 모델 규모 전반에서 입증한다.

제안 방법

• 조밀한 감독을 위한 K차원 점수 맵을 사용하는 토큰 라벨링 Objective를 정의한다. N개의 패치 토큰에서.
• 손실 합성식: H(X^{cls}, y^{cls}) + beta * (1/N) * sum_i H(X^i, y^i)로 이미지 분류 토큰 손실과 조밀한 토큰 라벨링 손실을 결합한다.
• 사전 학습된 기계 주석자(pretrained machine annotator)를 사용하여 패치별 감독 신호를 생성한다.
• Patch 임베딩 이후 토큰 공간을 혼합하는 토큰 공간 증강인 MixToken을 도입하고, 레이블도 이에 맞춰 혼합한다.
• LV-ViT 변형에서 지역 정보를 더 잘 포착하기 위해 패치 임베딩 모듈을 작은 합성곱으로 교체한다.

실험 결과

연구 질문

• RQ1 조밀하고 위치 특이적인 토큰 감독이 전통적인 클래스 토큰 학습을 넘어 비전 트랜스포머를 향상시킬 수 있는가?
• RQ2 토큰 라벨링이 ViT 변형 전반의 기존 증강 및 모델 크기와 어떻게 상호작용하는가?
• RQ3 토큰 라벨링이 주석자 품질과 패치 임베딩 선택에 대해 강건한가, 그리고 다운스트림의 밀도 예측 작업에 도움이 되는가?
• RQ4 다중 ViT 아키텍처에 토큰 라벨링을 적용했을 때 ImageNet 및 ADE20K에서의 성능 향상은 어느 정도인가?

주요 결과

• 토큰 라벨링이 적용된 LV-ViT는 150M 매개변수에서 ImageNet의 Top-1 84.4%를 달성하고, 더 큰 규모에서 86.4%를 달성하며, 유사 예산의 다른 트랜스포머 기반 모델보다 우수하다.
• 토큰 라벨링은 일관되게 7개의 ViT 변형(DeiT, T2T-ViT, LV-ViT 등)을 향상시키며, 더 큰 모델에서 더 큰 이점을 보인다.
• MixToken은 토큰 기반 트랜스포머에서 CutMix보다 우수하며, LV-ViT-S에서 토큰 라벨링과 MixToken을 사용해 83.3% Top-1을 달성하는 데 기여한다.
• ADE20K에서의 다운스트림 의미론적 분할도 토큰 라벨링의 조밀한 감독으로 이익을 얻으며, LV-ViT-L + UperNet과 같은 비-ImageNet-22K 선훈련 설정에서 최첨단 성능에 도달한다.
• 이 접근은 토큰 수준 스코어를 생성하는 사전 학습된 주석자가 필요하지만, 주석이 사전에 계산되므로 추가 훈련 비용은 거의 발생하지 않는다.