All Tokens Matter: Token Labeling for Training Better Vision Transformers

1 Introduction . 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 . 2 Related Work . 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 . 3 Token Labeling Method . In this section , we first briefly review the structure of the vision transformer [ 15 ] and then describe the proposed training objective—token labeling . 3.1 Revisiting Vision Transformer . 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 . 3.2 Token Labeling . 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 . 3.3 Token Labeling with MixToken . 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 .

Summary: 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.