All Tokens Matter: Token Labeling for Training Better Vision Transformers
Year: 2022

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 pretrained models on downstream tasks with dense prediction, such as semantic segmentation. Our code and model are publicly available at https://github.com/zihangJiang/TokenLabeling.

heading: 1 Introduction
text: Transformers [39] have achieved great performance for almost all the natural language processing (NLP) tasks over the past years [4, 14, 24]. Motivated by such success, recently, many researchers attempt to build transformer models for vision tasks, and their encouraging results have shown the great potential of transformer based models for image classification [6, 15, 25, 36, 40, 46], especially the strong benefits of the self-attention mechanism in building long-range dependencies between pairs of input tokens.
Despite the importance of gathering long-range dependencies, recent work on local data augmentation [57] has demonstrated that well modeling and leveraging local information for image classification would avoid biasing the model towards skewed and non-generalizable patterns and substantially ∗Work done as an intern at ByteDance AI Lab. †Corresponding author. Part of this work was done as a research fellow at NUS.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
improve the model performance. However, recent vision transformers normally utilize class tokens that aggregate global information to predict the output class while neglecting the role of other patch tokens that encode rich information on their respective local image patches.
In this paper, we present a new training objective for vision transformers, termed token labeling, that takes advantage of both the patch tokens and the class tokens. Our method takes a K-dimensional score map generated by a machine annotator as supervision to supervise all the tokens in a dense manner, where K is the number of categories for the target dataset. In this way, each patch token is explicitly associated with an individual location-specific supervision indicating the existence of the target objects inside the corresponding image patch, so as to improve the object grounding and recognition capabilities of vision transformers with negligible computation overhead. To the best of our knowledge, this is the first work demonstrating that dense supervision is beneficial to vision transformers in image classification.
According to our experiments, utilizing the proposed token labeling objective can clearly boost the performance of vision transformers. As shown in Figure 1, our model, named LV-ViT, with 56M parameters, yields 85.4% top-1 accuracy on ImageNet [13], behaving better than all the other transformer-based models having no more than 100M parameters. When the model size is scaled up to 150M, the result can be further improved to 86.4%. In addition, we have empirically found that the pretrained models with token labeling are also beneficial to downstream tasks with dense prediction, such as semantic segmentation.

heading: 2 Related Work
text: Transformers [39] refer to the models that entirely rely on the self-attention mechanism to build global dependencies, which are originally designed for natural language processing tasks. Due to their strong capability of capturing spatial information, transformers have also been successfully applied to a variety of vision problems, including low-level vision tasks like image enhancement [7, 45], as well as more challenging tasks such as image classification [9, 15], object detection [5, 11, 55, 61], segmentation [7, 33, 41] and image generation [28]. Some works also extend transformers for video and 3D point cloud processing [50, 53, 60].
Vision Transformer (ViT) is one of the earlier attempts that achieved state-of-the-art performance on ImageNet classification, using pure transformers as basic building blocks. However, ViTs need pretraining on very large datasets, such as ImageNet-22k and JFT-300M, and huge computation resources to achieve comparable performance to ResNet [18] with a similar model size trained on ImageNet. Later, DeiT [36] manages to tackle the data-inefficiency problem by simply adjusting the network architecture and adding an additional token along with the class token for Knowledge Distillation [21, 47] to improve model performance.
Some recent works [6, 16, 43, 46] also attempt to introduce the local dependency into vision transformers by modifying the patch embedding block or the transformer block or both, leading to significant performance gains. Moreover, there are also some works [20, 25, 40] adopting a pyramid structure to reduce the overall computation while maintaining the model’s ability to capture low-level features.
Unlike most aforementioned works that design new transformer blocks or transformer architectures, we attempt to improve vision transformers by studying the role of patch tokens that embed rich local information inside image patches. We show that by slightly tuning the structure of vision transformers and employing the proposed token labeling objective, we can achieve strong baselines for transformer models at different model size levels.

heading: 3 Token Labeling Method
text: In this section, we first briefly review the structure of the vision transformer [15] and then describe the proposed training objective—token labeling.

heading: 3.1 Revisiting Vision Transformer
text: A typical vision transformer [15] first decomposes a fixed-size input image into a sequence of small patches. Each small patch is mapped to a feature vector, or called a token, by projection with a linear layer. Then, all the tokens combined with an additional learnable class token for classification score prediction are sent into a stack of transformer blocks for feature encoding.
In loss computing, the class token from the output tokens of the last transformer block is usually selected and sent into a linear layer for the classification score prediction. Mathematically, given an image I , denote the output of the last transformer block as [Xcls, X1, ..., XN ], where N is the total number of patch tokens, and Xcls and X1, ..., XN correspond to the class token and the patch tokens, respectively. The classification loss for image I can be written as
Lcls = H(X cls, ycls), (1)
where H(·, ·) is the softmax cross-entropy loss and ycls is the class label.

heading: 3.2 Token Labeling
text: The above classification problem only adopts an image-level label as supervision whereas it neglects the rich information embedded in each image patch. In this subsection, we present a new training objective—token labeling—that takes advantage of the complementary information between the patch tokens and the class tokens.
Token Labeling: Different from the classification loss as formulated in Eqn. (1) that measures the distance between the single class token (representing the whole input image) and the corresponding image-level label, token labeling emphasizes the importance of all output tokens and advocates that each output token should be associated with an individual location-specific label. Therefore, in our method, the ground truth for an input image involves not only a single K-dimensional vector ycls but also a K ×N matrix or called a K-dimensional score map as represented by [y1, ..., yN ], where N is the number of the output patch tokens.
Specifically, we leverage a dense score map for each training image and use the cross-entropy loss between each output patch token and the corresponding aligned label in the dense score map as an auxiliary loss at the training phase. Figure 2 provides an intuitive interpretation. Given the output patch tokens X1, ..., XN and the corresponding labels [y1, ..., yN ], the token labeling objective can be defined as
Ltl = 1
N N∑ i=1 H(Xi, yi). (2)
Recall that H is the cross-entropy loss. Therefore, the total loss function can be written as
Ltotal = H(X cls, ycls) + β · Ltl, (3)
= H(Xcls, ycls) + β · 1 N N∑ i=1 H(Xi, yi), (4)
where β is a hyper-parameter to balance the two terms. In our experiment, we empirically set it to 0.5.
Advantages: Our token labeling offers the following advantages. First of all, unlike knowledge distillation methods that require a teacher model to generate supervision labels online, token labeling is a cheap operation. The dense score map can be generated by a pretrained model in advance (e.g., EfficientNet [34] or NFNet [3]). During training, we only need to crop the score map and perform interpolation to make it aligned with the cropped image in the spatial coordinate. Thus, the additional computations are negligible. Second, rather than utilizing a single label vector as supervision as done in most classification models and the ReLabel strategy [49], we also harness score maps to supervise the models in a dense manner and thereby the label for each patch token provides location-specific information, which can aid the training models to easily discover the target objects and improve the recognition accuracy. Last but not the least, as dense supervision is adopted in training, we found that the pretrained models with token labeling benefit downstream tasks with dense prediction, like semantic segmentation.

heading: 3.3 Token Labeling with MixToken
text: While training vision transformer, previous studies [36, 46] have shown that augmentation methods, like MixUp [52] and CutMix [48], can effectively boost the performance and robustness of the models. However, vision transformers rely on patch-based tokenization to map each input image to a sequence of tokens and our token labeling strategy also operates on patch-based token labels. If we apply CutMix directly on the raw image, some of the resulting patches may contain content from two images, leading to mixed regions within a small patch as shown in Figure 3. When performing token labeling, it is difficult to assign each output token a clean and correct label. Taking this situation into account, we rethink the CutMix augmentation method and present MixToken, which can be viewed as a modified version of CutMix operating on the tokens after patch embedding as illustrated in the right part of Figure 3.
To be specific, for two images denoted as I1, I2 and their corresponding token labels Y1 = [y11 , ..., y N 1 ] as well as Y2 = [y12 , ..., y N 2 ], we first feed the two images into the patch embedding module to tokenize each as a sequence of tokens, resulting in T1 = [t11, ..., t N 1 ] and T2 = [t 1 2, ..., t N 2 ]. Then, we produce a new sequence of tokens by applying MixToken using a binary mask M as follows:
T̂ = T1 M + T2 (1−M), (5)
where is element-wise multiplication. We use the same way to generate the mask M as in [48]. For the corresponding token labels, we also mix them using the same mask M :
Ŷ = Y1 M + Y2 (1−M). (6)
The label for the class token can be written as
ˆycls = M̄ycls1 + (1− M̄)ycls2 , (7)
where M̄ is the average of all element values of M .

heading: 4 Experiments

heading: 4.1 Experiment Setup
text: We evaluate our method on the ImageNet [13] dataset. All experiments are built and conducted upon PyTorch [29] and the timm [42] library. We follow the standard training schedule and train our models on the ImageNet dataset for 300 epochs. Besides normal augmentations like CutOut [57] and RandAug [10], we also explore the effect of applying MixUp [52] and CutMix [48] together with our proposed token labeling. Empirically, we have found that using MixUp together with token labeling brings no benefit to the performance, and thus we do not apply it in our experiments.
For optimization, by default, we use the AdamW optimizer [27] with a linear learning rate scaling strategy lr = 10−3 × batch_size640 and 5 × 10
−2 weight decay rate. For Dropout regularization, we observe that for small models, using Dropout hurts the performance. This has also been observed in a few other works related to training vision transformers [36, 37, 46]. As a result, we do not apply Dropout [32] and use Stochastic Depth [23] instead. More details on hyper-parameters and finetuning can be found in our supplementary materials.
We use the NFNet-F6 [3] trained on ImageNet with an 86.3% Top-1 accuracy as the machine annotator to generate dense score maps for the ImageNet dataset, yielding a 1000-dimensional score map for each image for training. The score map generation procedure is similar to [49], but we limit our experiment setting by training all models from scratch on ImageNet without extra data support, such as JFT-300M and ImageNet-22K. This is different from the original ReLabel paper [49], in which the EfficientNet-L2 model pretrained on JFT-300M is used. The input resolution for NFNet-F6 is 576× 576, and the dimension of the corresponding output score map for each image is L ∈ R18×18×1000. During training, the target labels for the tokens are generated by applying RoIAlign [17] on the corresponding score map. In practice, we only store the top-5 score maps for each position in half-precision to save space as storing the entire score maps for all the images results in 2TB storage. In our experiment, we only need 10GB of storage to store all the score maps.

heading: 4.2 Ablation Analysis
text: Model Settings: The default settings of the proposed LV-ViT are given in Table 1, where both token labeling and MixToken are used. A slight architecture modification to ViT [15] is that we replace the patch embedding module with a 4-layer convolution to better tokenize the input image and integrate local information. Detailed ablation about patch embedding can be found in our supplementary materials. As can be seen, our LV-ViT-T with only 8.5M parameters can already achieve a top-1 accuracy of 79.1% on ImageNet. Increasing the embedding dimension and network depth can further boost the performance. More experiments compared to other methods can be found in Sec. 4.3. In the following ablation experiments, we will set our LV-ViT-S as baseline and show the advantages of the proposed token labeling and MixToken methods.
MixToken: We use MixToken as a substitution for CutMix while applying token labeling. Our experiments show that MixToken performs better than CutMix for token-based transformer models. As shown in Table 2, when training with the original ImageNet labels, using MixToken is 0.1% higher than using CutMix. When using the ReLabel supervision, we can also see an advantage of 0.2% over the CutMix baseline. Combining with our token labeling, the performance can be further raised to 83.3%.
Table 2: Ablation on the proposed MixToken and token labeling augmentations. We also show results with either the ImageNet hard label and the ReLabel [49] as supervision.
Aug. Method Supervision Top-1 Acc.
MixToken Token labeling 83.3 MixToken ReLabel 83.0
CutMix ReLabel 82.8 Mixtoken ImageNet Label 82.5 CutMix ImageNet Label 82.4
Table 3: Ablation on different widely-used data augmentations. We have empirically found our proposed MixToken performs even better than the combination of MixUp and CutMix in vision transformers.
MixToken MixUp CutOut RandAug Top-1 Acc.
X 7 X X 83.3 7 7 X X 81.3 X X X X 83.1 X 7 7 X 83.0 X 7 X 7 82.8
Data Augmentation: Here, we study the compatibility of MixToken with other augmentation techniques, such as MixUp [52], CutOut [57] and RandAug [10]. The ablation results are shown in Table 3. We can see when all the four augmentation methods are used, a top-1 accuracy of 83.1% is achieved. Interestingly, when the MixUp augmentation is removed, the performance can be improved to 83.3%. This may be explained as, using MixToken and MixUp at the same time would bring too much noise in the label, and consequently cause confusion of the model. Moreover, the CutOut augmentation, which randomly erases some parts of the image, is also effective and removing it brings a performance drop of 0.3%. Similarly, the RandAug augmentation also contributes to the performance and using it brings an improvement of 0.5%.
All Tokens Matter: To show the importance of involving all tokens in our token labeling method, we attempt to randomly drop some tokens and use the remaining ones for computing the token labeling loss. The percentage of the remaining tokens is denoted as Token Participation Rate. As shown in Figure 4 (Left), we conduct experiments on two models: LV-ViT-S and LV-ViT-M. As can be seen, using only 20% of the tokens to compute the token labeling loss decreases the performance (−0.5% for LV-ViT-S and −0.4% for LV-ViT-M). Involving more tokens for loss computation consistently leads to better performance. Since involving all tokens brings negligible computation cost and gives the best performance, we always set the token participation rate as 100% in the following experiments.
Online Token Labeling: Unlike the online knowledge distillation method which generates labels by a teacher model online, our token labeling approach utilizes the dense label map generated in advance and directly applies the corresponding augmentation methods, such as random crop, on the label map to obtain token-level labels. To directly compare with the online knowledge distillation based method and validate the effectiveness of token-level supervision, we further conduct experiments on the online version of our token labeling method, which generates token-level labels online during training. Following DeiT [36], we use RegNetY-16GF [30] as the online teacher model. Results in terms of DeiT-S/LV-ViT-S Top-1 accuracy and training time for our token labeling, online knowledge distillation, and ReLabel [49] are listed in Table 4, with number of utilized tokens also included for clear comparison. As can be seen, for both online and offline cases, using token-level supervision can improve the overall performance with only negligible additional training cost. Meanwhile, compared to the vanilla training baseline, our proposed offline token labeling brings almost no additional training cost, and boosts the overall performance of LV-ViT-S by 0.9%, which well demonstrates its efficiency and effectiveness.
Robustness to Different Annotators: To evaluate the robustness of our token labeling method, we use different pretrained CNNs, including EfficientNet-B3,B4,B5,B6,B7,B8 [34], NFNet-F6 [3] and ResNest269E [51], as annotator models to provide dense supervision. Results are shown in the right part of Figure 4. We can see that, even if we use an annotator with relatively lower performance, such as EfficientNet-B3 whose Top-1 accuracy is 81.6%, it can still provide multi-label location-specific supervision and help improve the performance of our LV-ViT-S model. Meanwhile, annotator models with better performance can provide more accurate supervision, bringing even better performance, as stronger annotator models can generate better token-level labels. The largest annotator NFNet-F6 [3], which has the best performance of 86.3%, allows us to achieve the best result for LV-ViT-S, which is 83.3%. In addition, we also attempt to use a better model, EfficientNet-L2 pretrained on JFT-300M as described in [49] which has 88.2% Top-1 ImageNet accuracy, as our annotator. The performance of LV-ViT-S can be further improved to 83.5%. However, to fairly compare with the models without
extra training data, we only report results based on dense supervision produced by NFNet-F6 [3] that uses only ImageNet training data.
Robustness to Different ViT Variants: To further evaluate the robustness of our token labeling, we train different transformer-based networks, including DeiT [36], T2T-ViT [3] and our model LV-ViT, with the proposed training objective. Results are shown in Figure 5. It can be found that, all the models trained with token labeling consistently outperform their vanilla counterparts, demonstrating the robustness of token labeling with respect to different variants of patch-based vision transformers. Meanwhile, for different scales of the models, the improvement is also consistent. Interestingly, we observe larger improvements for larger models. These indicate that our proposed token labeling method is widely applicable to a large range of patch-based vision transformer variants.
Beyond Vision Transformers: We further explore the performance of token labeling on other CNN-based and MLP-based models. Results are shown in Table 5. Besides our re-implementation with more data augmentation and regularization techniques, we also provide the results from the original papers. It can be found that for both MLP-based and CNN-based models, our token labeling objective can also improve the performance over strong baselines by providing location-specific dense supervision.

heading: 4.3 Comparison to Other Methods
text: We compare our proposed model LV-ViT with other state-of-the-art methods in Table 6. For smallsized models, when the test resolution is set to 224×224, we achieve an 83.3% accuracy on ImageNet with only 26M parameters, which is 3.4% higher than the strong baseline DeiT-S [36]. For mediumsized models, when the test resolution is set to 384 × 384 we achieve the performance of 85.4%, the same as CaiT-S36 [37], but with much less computational cost and parameters. Note that both DeiT and CaiT use knowledge distillation to improve their models, which introduce much more computations in training. However, we do not require any extra computations in training and only have to compute and store the dense score maps in advance. For large-sized models, our LV-ViT-L with a test resolution of 448 × 448 achieves an 86.2% top-1 accuracy, which is comparable to CaiT-M36 [37] but with far fewer FLOPs and parameters.

heading: 4.4 Semantic Segmentation on ADE20K
text: It has been shown in [19] that different training techniques for pretrained models have different impacts on downstream tasks with dense prediction, like semantic segmentation. To demonstrate the advantage of the proposed token labeling objective on tasks with dense prediction, we apply our pretrained LV-ViT with token labeling to the semantic segmentation task.
Similar to previous work [25], we run experiments on the widely-used ADE20K [58] dataset. ADE20K contains 25K images in total, including 20K images for training, 2K images for validation and 3K images for test, and covering 150 different foreground categories. We take both FCN [26] and UperNet [44] as our segmentation frameworks and use the mmseg toolbox to implement. During training, following [25], we use the AdamW optimizer with an initial learning rate of 6e-5 and a weight decay of 0.01. We also use a linear learning schedule with a minimum learning rate of 5e-6. All models are trained on 8 GPUs and with a batch size of 16 (i.e., 2 images on each GPU). The input resolution is set to 512× 512. In inference, a multi-scale test with interpolation rates of [0.75, 1.0, 1.25, 1.5, 1.75] is used. As suggested by [58], we report results in terms of both mean intersection-over-union (mIoU) and the average pixel accuracy (Pixel Acc.).
In Table 7, we test the performance of token labeling on both FCN and UperNet frameworks. The FCN framework has a light convolutional head and can directly reflect the performance of the pretrained models in terms of transferable capability. As can be seen, pretrained models with token
labeling perform better than those without token labeling. This indicates token labeling is indeed beneficial to semantic segmentation.
We also compare our segmentation results with previous state-of-the-art segmentation methods in Table 8. Without pretraining on large-scale datasets such as ImageNet-22K, our LV-ViT-M with the UperNet segmentation architecture achieves an mIoU score of 50.6 with only 77M parameters. This result is much better than the previous CNN-based and transformer-based models. Furthermore, using our LV-ViT-L as the pretrained model yields a better result of 51.8 in terms of mIoU. As far as we know, this is the best result reported on ADE20K with no pretraining on ImageNet-22K or other large-scale datasets.

heading: 5 Conclusions and Discussion
text: In this paper, we introduce a new token labeling method to help improve the performance of vision transformers. We also analyze the effectiveness and robustness of our token labeling with respect to different annotators and different variants of patch-based vision transformers. By applying token labeling, our proposed LV-ViT achieves 84.4% Top-1 accuracy with only 26M parameters and 86.4% Top-1 accuracy with 150M parameters on ImageNet-1K benchmark.
Despite the effectiveness, token labeling has a limitation of requiring a pretrained model as the machine annotator. Fortunately, the machine annotating procedure can be done in advance to avoid introducing extra computational cost in training. This makes our method quite different from knowledge distillation methods that rely on online teaching. For users with limited machine resources on hand, our token labeling provides a promising training technique to improve the performance of vision transformers.

Review summaries:
This paper presents a novel approach for training vision transformer. Instead of using the image-level label as the supervision, the proposed approach uses pixel-level labels as supervision, and use the cutmix-like image augmentation (MixToken) to generate pixel-level labels. The proposed approach can be regarded as the combination of cutmix [46] and relabeling [47]. Experimental results implied that proposed scheme improves the training quality. 

This paper presents a new strategy to more effectively train vision transformers by introducing token-wise supervision. A token labeling loss, MixToken data augmentation and several modifications on ViT architectures are proposed. Extensive experiments are conducted to show the effectiveness of the proposed method. 

This paper introduces token labelling, a new training objective for vision transformers. The core idea is to supervise each of the patches of the last layer, instead of simply discarding them as is usually done. The proposed approach significantly improves performance of various architectures on ImageNet.