All Tokens Matter: Token Labeling for Training Better Vision Transformers

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.

References in corpus (16)

• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Distilling the Knowledge in a Neural Network
• Decoupled Weight Decay Regularization
• mixup: Beyond Empirical Risk Minimization
• Language Models are Few-Shot Learners
• MLP-Mixer: An all-MLP Architecture for Vision
• Transformer in Transformer
• DeepViT: Towards Deeper Vision Transformer
• High-Performance Large-Scale Image Recognition Without Normalization
• CvT: Introducing Convolutions to Vision Transformers
• Pre-Trained Image Processing Transformer
• End-to-End Object Detection with Adaptive Clustering Transformer
• Are we done with ImageNet?
• Rethinking Spatial Dimensions of Vision Transformers
• End-to-End Video Instance Segmentation with Transformers

Cited by in corpus (9)

• Transformers in Vision: A Survey
• Vision Transformer with Attentive Pooling for Robust Facial Expression Recognition
• A survey on efficient vision transformers: algorithms, techniques, and performance benchmarking
• A Battle of Network Structures: An Empirical Study of CNN, Transformer, and MLP
• VOLO: Vision Outlooker for Visual Recognition
• A Unified Pruning Framework for Vision Transformers
• ClipFormer: Key-Value Clipping of Transformers on Memristive Crossbars for Write Noise Mitigation
• DS-Net++: Dynamic Weight Slicing for Efficient Inference in CNNs and Transformers
• 3rd Place Solution for VisDA 2021 Challenge -- Universally Domain Adaptive Image Recognition