All Tokens Matter: Token Labeling for Training Better Vision Transformers

TL;DR

This paper introduces token labeling, a new training objective for vision transformers that leverages all patch tokens with location-specific supervision, significantly improving performance on ImageNet and downstream dense prediction tasks.

Contribution

It proposes token labeling, a novel dense supervision method for ViTs, enhancing accuracy and generalization over traditional class token training.

Findings

Achieves 84.4% Top-1 accuracy on ImageNet with a 26M parameter ViT.

Scaling to 150M parameters increases accuracy to 86.4%.

Improves generalization on dense prediction tasks like semantic segmentation.

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.