PaperlayerView exact revision on arXiv

Source-linked AI summary

All Tokens Matter: Token Labeling for Training Better Vision Transformers

Zihang Jiang, Qibin Hou, Li Yuan, Daquan Zhou, Yujun Shi, Xiaojie Jin, Anran Wang, Jiashi Feng

arXiv:2104.10858v3cs.CV

TL;DR

Vision transformers often rely on a global class token while neglecting local information in patch tokens. Token labeling adds machine-generated, location-specific supervision for all patch tokens alongside class-token training. Experiments report improved performance across model sizes and downstream dense-prediction tasks, including 84.4% Top-1 accuracy with 26M parameters and 86.4% with 150M parameters on ImageNet-1K.

  • Problem

    Recent vision transformers commonly aggregate global information through class tokens while neglecting patch tokens that encode rich local information.

  • Method

    Token labeling uses a machine-generated K-dimensional score map to assign each patch token an individual location-specific label and train with dense auxiliary supervision.

  • Results

    84.4% Top-1 accuracy is achieved with 26M parameters and 86.4% with 150M parameters on ImageNet-1K; token labeling also benefits downstream dense prediction.

  • Takeaways & Limitations

    Token labeling provides a training technique that improves vision-transformer performance while adding negligible training computation, and pretrained models benefit on dense-prediction tasks.

  • Takeaways & Limitations

    Token labeling requires a pretrained model to serve as the machine annotator, although annotation can be performed in advance.

Abstract

from arXiv · show

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.

1 Introduction

Vision transformers capture global dependencies but typically underuse patch tokens carrying local information. Token labeling adds dense, location-specific supervision for patch tokens alongside class-token training, improving ImageNet performance and downstream dense prediction.

  • Recent vision transformers commonly use class tokens for global image classification while neglecting patch tokens that encode local image information.
  • Token labeling supervises both class and patch tokens using a machine-generated K-dimensional score map with location-specific labels.Each patch token receives supervision indicating target-object presence within its corresponding image patch.
  • Token labeling is presented as a new training objective that uses dense supervision to improve vision transformers’ object grounding and recognition with negligible computation overhead.
  • 85.4% top-1 accuracy is achieved on ImageNet with LV-ViT using 56M parameters, outperforming transformer models with no more than 100M parameters.At 150M parameters, performance increases to 86.4%.
  • Pretrained models using token labeling are also beneficial for downstream dense-prediction tasks such as semantic segmentation.

2 Related Work

Vision transformers have been extended across vision tasks and improved through local dependencies and pyramid structures. This work instead studies patch tokens and token labeling, using them to strengthen transformer baselines across model sizes.

  • Transformers use self-attention to build global dependencies and have been applied to classification, detection, segmentation, generation, video, and 3D point-cloud processing.
  • ViT achieved state-of-the-art ImageNet classification but required very large pretraining datasets and substantial computation for competitive performance.
  • Prior vision-transformer research introduced local dependencies through modified embedding or transformer blocks and reduced computation with pyramid structures.
  • Unlike architecture-focused approaches, this work studies patch tokens carrying rich local image information and uses token labeling to improve transformer baselines across model sizes.

3 Token Labeling Method

The method augments vision-transformer training by supervising patch tokens alongside the class token, using location-specific token labels and a combined loss. It also applies token-level mixing after patch embedding to preserve clean token content during augmentation.

  • 3.1 Revisiting Vision Transformer: A vision transformer tokenizes image patches, appends a learnable class token, and processes the sequence through transformer blocks.The final output contains one class token and N patch tokens.
  • 3.1 Revisiting Vision Transformer: Standard classification loss uses only the final class token and the image-level label, neglecting information encoded in individual patch tokens.The class token is sent to a linear layer for classification-score prediction.
  • 3.2 Token Labeling: Token labeling adds a K × N score map so each output patch token receives an aligned, location-specific supervision signal.The score map is generated by a machine annotator, and cross-entropy is computed between each patch token and its corresponding label.
  • 3.2 Token Labeling: The total objective combines class-token classification loss with the average token-labeling loss, weighted by hyper-parameter β set empirically to 0.5.This combines image-level and dense patch-level supervision during training.
  • 3.2 Token Labeling: Token labeling uses pretrained-model score maps generated in advance, requiring only cropping and interpolation during training, with negligible additional computation.The dense supervision also benefits pretrained models on downstream dense-prediction tasks such as semantic segmentation.
  • 3.3 Token Labeling with MixToken: MixToken performs mixing after patch embedding and applies the same binary mask to tokens and token labels, avoiding mixed regions within individual patches.Unlike CutMix on raw images, this preserves clean content for each resulting token.

4 Experiments

Experiments show that token labeling improves ViT performance across model sizes, augmentation settings, annotator models, and transformer variants, while transferring well to semantic segmentation.

  • Ablation Analysis: 83.3% ImageNet Top-1 accuracy is obtained when MixToken is combined with token labeling, exceeding the CutMix comparison by the reported ablation margin.MixToken is 0.1% higher than CutMix with original labels and 0.2% higher with ReLabel supervision.
  • Ablation Analysis: 83.1% ImageNet Top-1 accuracy is reached with four augmentations, while removing MixUp raises performance to 83.3%.The authors attribute the MixUp effect to excessive label noise when combined with MixToken; removing CutOut lowers performance by 0.3%, while RandAug contributes 0.5%.
  • Ablation Analysis: 100% token participation gives the best performance at negligible computation cost; using only 20% reduces accuracy by 0.5% for LV-ViT-S and 0.4% for LV-ViT-M.Performance improves consistently as more patch tokens contribute to the token-labeling loss.
  • Ablation Analysis: Token labeling consistently improves all 7 evaluated DeiT, T2T-ViT, and LV-ViT models and remains effective with different annotator networks.With NFNet-F6 supervision, LV-ViT-S reaches 83.3%; EfficientNet-L2 supervision further raises it to 83.5%.
  • Comparison to Other Methods: 83.3% ImageNet Top-1 accuracy is achieved with 26M parameters, 3.4% above DeiT-S, without extra training computations for supervision.Dense score maps are generated in advance, avoiding the additional training computation required by knowledge distillation.
  • Semantic Segmentation on ADE20K: 51.8 mIoU is achieved on ADE20K with LV-ViT-L, while LV-ViT-M reaches 50.6 mIoU using 77M parameters without ImageNet-22K pretraining.Pretrained models using token labeling also outperform those without it in FCN and UperNet semantic segmentation experiments.

5 Conclusions and Discussion

The paper introduces token labeling to improve vision transformers and reports strong ImageNet results at 26M and 150M parameters. Its main limitation is the requirement for a pretrained machine annotator.

  • Token labeling improves vision transformer performance while supporting 84.4% Top-1 accuracy with 26M parameters and 86.4% with 150M parameters on ImageNet-1K.
  • The method is analyzed across different annotators and variants of patch-based vision transformers for effectiveness and robustness.
  • Token labeling requires a pretrained model to serve as the machine annotator.Annotation can be performed in advance, avoiding extra computational cost during training.
  • Precomputed machine annotation distinguishes token labeling from knowledge distillation methods that rely on online teaching.

A More Experiment Details

The experiments specify default and large-resolution fine-tuning settings, including augmentation constraints and substantial GPU requirements.

  • Large-resolution fine-tuning uses batch size 512, learning rate 5e-6, weight decay 1e-8, and 30 epochs.
  • Standard training requires one machine node with 8 NVIDIA V100 GPUs, while fine-tuning the large model at 448 × 448 requires four such nodes.
  • Token labeling and ReLabel are trained without MixUp augmentation.

B.1 Training Technique Analysis

The training analysis combines architectural modifications, token labeling, and higher-resolution fine-tuning to improve the DeiT-Small baseline toward LV-ViT-S performance.

  • Training Technique Analysis: 26M learnable parameters improve from 79.9 to 84.4 (+4.5) through token labeling and other proposed training techniques.
  • Training Technique Analysis: Adding convolutional layers and narrowing patch-embedding kernels consistently improves performance over the original single-layer embedding, but further convolutional expansion is not adopted.
  • Training Technique Analysis: The patch embedding uses an initial stride-2, kernel-7 convolution, two stride-1, kernel-3 convolutions, and final stride-8, kernel-8 tokenization.
  • Training Technique Analysis: Patch-embedding ablations vary convolution count, kernel size, and stride in a 16-layer ViT with embedding size 384 and MLP expansion ratio 3.
  • Training Technique Analysis: Smaller residual scaling factors improve performance and accelerate convergence by preserving more information in the main branch.
  • Training Technique Analysis: Residual-scaling ablations use a 16-layer vision transformer with four-layer convolutional patch embedding, evaluating self-attention and feed-forward functions.
  • Training Technique Analysis: Fine-tuning at 384 × 384 improves LV-ViT-S performance by 1.1%.

B.2 Beyond Vision Transformers: Performance on MLP-Based and CNN-Based Models

Token labeling also improves representative CNN-based and MLP-based models, indicating that its location-specific dense supervision extends beyond vision transformers.

  • Token labeling improves performance over strong baselines for both MLP-based and CNN-based models.The reported models include Mixer-MLP and ResNeSt.
  • Results for ResNeSt and Mixer-MLP show consistent improvement across the different evaluated models.

B.3 Comparison with CaiT

The paper compares LV-ViT with CaiT and reports identical results while using fewer training techniques, a smaller model, and less computation.

  • LV-ViT achieves identical results to the state-of-the-art CaiT model with fewer training techniques, a smaller model size, and less computation.The comparison is presented in Table 12.

C Visualization

The paper visualizes DeiT-base and LV-ViT-S, including their top-2 predicted classes, and reports better target-object localization and higher-confidence classification for LV-ViT-S.

  • Visual comparisons between DeiT-base and LV-ViT-S examine how the models locate target objects.The comparisons are shown in Figure 6.
  • LV-ViT-S performs better in locating target objects and yields better classification performance with high confidence.
  • Figure 7 compares the top-2 classes predicted by DeiT-base and LV-ViT-S on images containing at least two classes.
Loading 2104.10858v3…