[Paper Review] All Tokens Matter: Token Labeling for Training Better Vision Transformers

Zihang Jiang, Qibin Hou+6 more|arXiv (Cornell University)|Apr 22, 2021

Advanced Neural Network Applications56 references140 citations

TL;DR

The paper introduces token labeling, a dense patch-token supervision objective for vision transformers, yielding consistent gains across ViT variants and enabling high accuracy with smaller models. It reports 84.4% Top-1 on ImageNet with 26M parameters and 86.4% with 150M, plus improved downstream semantic segmentation performance.

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.

Motivation & Objective

• Reconsider how ViTs are trained by leveraging dense, location-specific supervision for every image patch token.
• Convert image classification into multiple token-level recognition problems to enhance object grounding and recognition.
• demonstrate the generality and effectiveness of token labeling across ViT variants and model scales.

Proposed method

• Define a token labeling objective that uses a K-dimensional score map for dense supervision on N patch tokens.
• Combine image-class token loss with a dense token labeling loss via L_total = H(X^{cls}, y^{cls}) + beta * (1/N) * sum_i H(X^i, y^i).
• Generate token-level labels using a pretrained machine annotator to produce per-patch supervision.
• Introduce MixToken, a token-space augmentation that blends tokens after patch embedding, with labels blended accordingly.
• Replace the patch embedding module with a small convolution to better capture local information in LV-ViT variants.

Experimental results

Research questions

• RQ1 Can dense, location-specific token supervision improve vision transformers beyond traditional class-token training?
• RQ2 How does token labeling interact with existing augmentations and model sizes across ViT variants?
• RQ3 Is token labeling robust to different annotator quality and patch-embedding choices, and does it help downstream dense-prediction tasks?
• RQ4 What are the performance gains in ImageNet and ADE20K when applying token labeling across multiple ViT architectures?

Key findings

• LV-ViT with token labeling achieves 84.4% Top-1 on ImageNet with 150M parameters and 86.4% at larger scales, outperforming several transformer-based baselines under similar budgets.
• Token labeling consistently improves seven ViT variants (DeiT, T2T-ViT, LV-ViT, etc.) with larger gains for bigger models.
• MixToken outperforms CutMix for token-based transformers and helps achieve 83.3% Top-1 on ImageNet with LV-ViT-S using token labeling and MixToken.
• Dense supervision from token labeling also benefits downstream semantic segmentation on ADE20K, achieving state-of-the-art results among non-ImageNet-22K pretraining setups (e.g., 51.8 mIoU with LV-ViT-L + UperNet).
• The approach requires a pretrained annotator to generate token-level scores, but incurs negligible extra training cost since annotation is precomputed.