#  Multiclass MRI Brain Tumor Segmentation using 3D Attention-based U-Net

Maryann M. Gitonga
<EMAIL_ADDRESS>
Strathmore University, Nairobi, Kenya

###### Abstract

This paper proposes a 3D attention-based U-Net architecture for multi-region
segmentation of brain tumors using a single stacked multi-modal volume created
by combining three non-native MRI volumes. The attention mechanism added to
the decoder side of the U-Net helps to improve segmentation accuracy by de-
emphasizing healthy tissues and accentuating malignant tissues, resulting in
better generalization power and reduced computational resources. The method is
trained and evaluated on the BraTS 2021 Task 1 dataset, and demonstrates
improvement of accuracy over other approaches. My findings suggest that the
proposed approach has potential to enhance brain tumor segmentation using
multi-modal MRI data, contributing to better understanding and diagnosis of
brain diseases. This work highlights the importance of combining multiple
imaging modalities and incorporating attention mechanisms for improved
accuracy in brain tumor segmentation.

Keywords: attention mechanism, U-Net, MRI, NIfTI.

##  1 Introduction

Glioma is a common malignant brain tumor that originates from glial cells in
the brain and spinal cord. Gliomas are aggressive, and the median survival
time for glioma patients is about 12 months [1]. Early detection of these
tumors is critical, and MRI is a primary tool used for this purpose. MRI
provides high spatial resolution anatomical information and different
sequences such as T1-weighted, T2-weighted, T1-weighted contrast-enhanced, and
T2 Fluid Attenuated Inversion Recovery, which highlight different tumor
characteristics [2]. Accurate annotation and segmentation of tumor borders are
essential for tumor diagnosis. However, manual segmentation is costly, time-
consuming, and prone to human error, especially in cases where tumors have
varying intensities and shapes in different sub-regions [3].

##  2 Literature Review

Consequently, deep learning techniques have transformed brain tumor
segmentation from feature-driven to data-driven. Two types of deep learning
algorithms, Convolutional Neural Network(CNN)-based and Fully Convolutional
Network(FCN)-based, are used in brain tumor segmentation. Havaei et al.
proposed a multi-path CNN network, InputCascadeCNN, which uses variable size
convolution kernels to extract context features and has both local and global
routes [4]. The model attained a Dice coefficient of 0.81 for the complete
segmentation on the BraTS 2013 dataset. Myronenko proposed a 3D CNN-based
approach using a shared decoder and a variational auto-encoder branch for
regularization. The encoder of the network extracts features of the images and
the decoder reconstructs the dense segmentation masks for the scan. The
variational auto-encoder branch reconstructs the input image into itself and
is used only during training to regularize the shared decoder. The model
attained an average Dice coefficient of 0.82 on the BraTS 2018 dataset [5].
Bukhari et al. proposed the 3D U-Net which has a contracting path for
capturing context information and an expanding path for ensuring precise
location, which greatly improves the performance of medical picture
segmentation task. This state-of-the-art model attained a Dice coefficient of
approximately 0.92 on the BraTS 2021 dataset[6]. Lin et al. incorporated a
feature pyramid module into the U-Net architecture to combine multi-scale
semantic and location information. The solution had shortened distance between
the output layers and deep features by bottom-up path aggregation to try and
reduce the noise in the segmentation. The efficient feature pyramid was used
to improve mask prediction using fewer resources to complete the feature
pyramid effect. The model attained a Dice coefficient of 0.80 on the BraTS
2017 and BraTS 2018 datasets[7]. Jun et al. introduced a nn Unet architecture
which included an encoder and decoder composed of convolutions, normalization,
and skip connections, with deep supervision added to all but the two lowest
resolutions in the decoder. It attained an average Dice coefficient of 0.90 on
the BraTS 2021 dataset [8].

However, these techniques still take up significant time and resources in
training and evaluation to get good results. Additionally, small-scale tumors
are difficult to accurately segment due to decreased picture dimension during
downsampling. In this paper, I test the 3D attention-based U-Net network [9]
with the Dice Coefficient and Tversky Loss Function as metrics, to improve the
segmentation accuracy. [10] applied the same network but applied Hausdorff
Distance and augmented the dataset using the Positive Mining technique. I
apply the proposed method to the BraTS (Brain Tumor Segmentation) 2021 Dataset
provided by Medical Image Computing and Computer-Assisted Intervention
(MICCAI). To provide a richer spatial information in the input as well as
enable one-time segmentation, the three modalities out of the four are
combined into one. This is because the native modality (T1) highlights the
healthy anatomy of the brain and not the tumor regions [11]. This results to a
4D input of dimensions $3M\times L\times W\times S$, where $M$ is the
modality, $L$ is the length of the scan, $W$ is the width of the scan and $S$
is the number of slices in each volume. This allows the model to put more
focus on the regions of interest (region showing a potential tumor).

##  3 Methodology

###  3.1 3D Attention U-Net

The proposed architecture (cf. Figure 1) uses the U-Net architecture [12],
which employs a contracting path to down-sample image dimensions and an
expanding path to up-sample while retaining spatial information through skip
connections. Attention modules are utilized at the skip connections to
highlight relevant activations and suppress those in irrelevant regions during
training [9]. The soft attention module, which is differentiable and essential
for back-propagation, is used, consisting of two sub-modules: the channel
attention module and spatial attention module. The former selects important
feature maps, while the latter identifies important regions within the feature
maps, and both are used to take full advantage of the architecture. 3D
attention gates are introduced to generate 3D channel and spatial attention by
utilizing 3D inter-channel and inter-spatial feature relationships.

Figure 1: 3D Attention U-Net Architecture

The attention mechanism (cf. Figure 2) consists of an attention gate that
takes in two input vectors, x and g. g is acquired from the lower part of the
network and represents better features and x comes from early layers and
represent better spatial information. To align the dimensions of the two
vectors, x undergoes a strided 3D convolution while g undergoes a 3D
convolution with number of filters = $F_{g}$. The two vectors are summed
element-wise, with aligned weights becoming larger while unaligned weights
become relatively smaller. The resultant vector undergoes a ReLU activation
layer and a 1x1x1 convolution that reduces the dimensions to 1xHxWxD. A
sigmoid layer is applied to scale the vector between 0 and 1, generating
attention coefficients that indicate the relevance of the features. The
attention coefficients are upsampled to the original dimensions of vector x
and multiplied element-wise with vector x. The resultant scaled vector is
passed along in the skip connection. This mechanism helps to increase the
sensitivity of the network to small but important details and reduce the
impact of noisy or irrelevant features. The overall network is able to learn
better and more discriminative features which improves on its accuracy and
efficiency.

Figure 2: Visual representation of 3D Attention mechanism

##  4 Experiments

###  4.1 Dataset and Pre-processing

The BraTS 2021 Dataset, which consists of 1400 cases of multi-parametric MRI
(mpMRI) scans with expert neuro-radiologists’ ground truth annotations, was
used for this project. The dataset provides mpMRI scans in NIfTI format and
includes native (T1), post-contrast T1-weighted (T1CE), T2-weighted (T2), and
T2 Fluid Attenuated Inversion Recovery (T2-FLAIR) volumes, along with manually
annotated GD-enhancing tumor, peritumoral edematous/invaded tissue, necrotic
tumor core, and normal tissue.

The scans in each folder were transformed by scaling and translating the
features using the MinMax Scaler to shrink the features within the range of 0
to 1 [13]. The combined MRI scan was generated by merging the 3 volumes of
each brain scan to form a 4D array of 3(modalities) x length x width x number
of slices. This provides richer spatial information for one-time segmentation.
The fourth volume (native modality) was left out because the scan highlights
the healthy tissues of the brain [11] which does not majorly contribute to the
segmentation of the tumor regions. The combined scan and corresponding mask
were then cropped to remove useless blank regions, reducing bias and focusing
on the important parts of the volume. The pre-processed combined scan and mask
were saved as numpy arrays, with the mask features converted to class values
(labels) 0, 1, 2, and 3. Masks with a segmented region less than 1% were
excluded to retain only significant feature representation for the segmented
regions. The resulting dataset after pre-processing had about 1200 cases of
tumor without any additional in-house data. The dataset was divided into three
sets: the train, test, and validation dataset in the ratio 6:2:2 respectively.

###  4.2 Implementation Details

During the training of a segmentation model, several hyper-parameters were
investigated, including batch size, learning rate, epochs, activation
function, dropout rate, metric functions, and loss functions. To prevent
overfitting, the dropout approach and batch normalization were utilized for
model regularization. The dropout rate was distributed across the layers
between the range of 0.1 to 0.3 in both the encoder and decoder modules. To
fit the data into memory, a batch size of 2 was utilized, which was the
maximum allowable limit based on the GPU specifications obtained for this
research. This small batch size also offered a regularizing effect, resulting
in lower generalization error [14]. The Adam optimizer with a learning rate of
$1.0\times 10^{-4}$ was utilized for weight updates [15]. A pixel-wise softmax
activation function was employed in the last layer of the model. The Dice
Coefficient (Equation 1) was utilized as an evaluation metric for both
training and testing phases. It calculates the ratio between the intersection
and the union of the segmented and ground truth regions, focusing only on the
segmentation classes and not the background class. The pixels are classified
as True Positive ($TP$), False Negative ($FN$) and False Positive ($FP$).

$D=\frac{2TP}{2TP+FN+FP}$ (1)

The Tversky Loss [16] (Equation 3), based on the Tversky Index (Equation 2)
where $\alpha=0.7,\beta=0.3$ [17]), was used for both training and testing
phases.

$TI=\frac{TP}{TP+\alpha FN+\beta FP}$ (2) $TL=1-TI$ (3)

This loss function is a generalized approach to handle class imbalance issues,
resulting in a better balance between precision and recall during model
training. ReLU was used in the first trial of training and evaluation and in
the second trial, the activation function was changed to LeakyReLU to prevent
the dying ReLU problem. [18].

##  5 Results

The network was implemented in Tensorflow and trained it on NVIDIA Tesla V100
32GB GPU. Results for Dice Coefficient and Tversky Loss metrics evaluation on
the validation and testing datasets are presented in Table 1. The developed
model achieved promising results in brain tumor segmentation, with the best
performance attained during the second trial at the 127th epoch. Table 2 shows
the dice coefficients attained by other models from different studies in
comparison to the developed model. The use of Dice Coefficient and Tversky
Loss metrics evaluation on the validation and testing datasets demonstrated
the model’s effectiveness. The visualization of the testing dataset is as
shown in Figure 3. The model’s ability to accurately delineate the tumor and
its sub-regions from the input stack of three volumes (T2-FLAIR, T1CE, and T2)
helps to create an effective treatment plan based on the nature of the tumor
sub-regions observed from the error-proof segmentations obtained.

Trial |  Dice Coefficient |  Tversky Loss
---|---|---

Validation Dataset (BraTS 2021)
Trial 1 (epoch = 75) 0.9430 0.0570 Trial 2 (epoch = 127) 0.9562 0.0438 Testing
Dataset (BraTS 2021)
Trial 2 0.9864 0.0136

Table 1: Dice Coefficient and Tversky Loss metrics evaluation on the
validation and testing datasets.

Model |  Dataset |  Dice Coefficient
---|---|---

InputCascadeCNN [4] |  BraTS 2013 |  0.81
---|---|---
Encoder-Decoder with a variational auto-encoder branch [5] |  BraTS 2018 |  0.82
3D U-Net [6] |  BraTS 2021 |  0.92
Path aggregation U-Net [7] |  BraTS 2017, BraTS 2018 |  0.80
3D Attention U-Net [19] |  BraTS 2019 |  0.86
Nn U-Net [8] |  BraTS 2021 |  0.90
3D Attention U-Net* |  BraTS 2021 |  0.98

Table 2: Dice Coefficient metric comparison with models from other studies.
The developed model is marked with an *. Figure 3:  Predictions and Ground
truths visualized from a combined scan generated from 3 modalities: T2-FLAIR,
T1 and T1CE from the validation set from BraTS 2021. The annotation can be
interpreted as necrosis, edema (invaded tissue) and enhancing tumor.

This method has significant implications for early brain tumor detection,
which is crucial for effective treatment and ultimately saving lives. With
tumors being one of the leading causes of mortality worldwide, the model’s
output is critical in forecasting the tumor’s aggressiveness and the patient’s
survival early enough, allowing for the best chance for successful treatment.
The developed solution helps facilitate accurate and effective medical
diagnostics by optimizing computational resources consumed on irrelevant areas
on the images and facilitating better generalization of the network used for
the task.

##  6 Conclusion

In this paper, the model implements a 3D attention mechanism in a 3D U-Net
network to improve model sensitivity and accuracy to foreground pixels without
requiring significant computation overhead by progressively suppressing
features responses in irrelevant background regions. The model performs
segmentation on a stack of 3 modalities of the MRI scan, in their original
format (NIfTI) to attain richer feature representation during segmentation.
This work clearly exhibits the significance of the 3D attention mechanism in
multi-class segmentation with limited computational resources. The
significance of stacking the modalities into one array is also demonstrated:
providing better feature representation in the input and facilitation of one-
time segmentation of the multi-modal scans. This solution in its entirety
contributes to the development of accurate, extensive delineation tools for
brain tumors, allowing the physicians to develop effective treatment plans for
the patients based on the nature of the tumor sub-regions observed from error-
proof segmentations obtained.

## References

  * [1] Ramin Ranjbarzadeh et al.  “Brain tumor segmentation based on deep learning and an attention mechanism using MRI multi-modalities brain images”  In _Scientific Reports_ 11.1 Springer, 2021, pp. 1–17
  * [2] Francisco Javier Diaz-Pernas, Mario Martinez-Zarzuela, Miriam Antón-Rodriguez and David González-Ortega  “A deep learning approach for brain tumor classification and segmentation using a multiscale convolutional neural network”  In _Healthcare_ 9.2, 2021, pp. 153  MDPI
  * [3] Spyridon Bakas et al.  “The Brain Tumor Segmentation Challenge (2022 Continuous Updates & Generalizability Assessment)”  Zenodo, 2022  DOI: 10.5281/zenodo.6362180
  * [4] Mohammad Havaei et al.  “Brain tumor segmentation with deep neural networks”  In _Medical image analysis_ 35 Elsevier, 2017, pp. 18–31
  * [5] Andriy Myronenko  “3D MRI brain tumor segmentation using autoencoder regularization”  In _Brainlesion: Glioma, Multiple Sclerosis, Stroke and Traumatic Brain Injuries: 4th International Workshop, BrainLes 2018, Held in Conjunction with MICCAI 2018, Granada, Spain, September 16, 2018, Revised Selected Papers, Part II 4_ , 2019, pp. 311–320  Springer
  * [6] Syed Talha Bukhari and Hassan Mohy-ud-Din  “E1D3 U-Net for Brain Tumor Segmentation: Submission to the RSNA-ASNR-MICCAI BraTS 2021 Challenge”  In _arXiv e-prints_ , 2021, pp. arXiv–2110
  * [7] Fengming Lin et al.  “Path aggregation U-Net model for brain tumor segmentation”  In _Multimedia Tools and Applications_ 80 Springer, 2021, pp. 22951–22964
  * [8] Jun Ma and Jianan Chen  “NnUNet with Region-based Training and Loss Ensembles for Brain Tumor Segmentation”  In _Brainlesion: Glioma, Multiple Sclerosis, Stroke and Traumatic Brain Injuries: 7th International Workshop, BrainLes 2021, Held in Conjunction with MICCAI 2021, Virtual Event, September 27, 2021, Revised Selected Papers, Part I_ , 2022, pp. 421–430  Springer
  * [9] Ozan Oktay et al.  “Attention u-net: Learning where to look for the pancreas”  In _arXiv preprint arXiv:1804.03999_ , 2018
  * [10] Har Shwinder Singh  “Brain Tumor Segmentation Using Attention Activated U-Net with Positive Mining”  In _Brainlesion: Glioma, Multiple Sclerosis, Stroke and Traumatic Brain Injuries: 7th International Workshop, BrainLes 2021, Held in Conjunction with MICCAI 2021, Virtual Event, September 27, 2021, Revised Selected Papers, Part I_ , 2022, pp. 431–440  Springer
  * [11] Jana Vasković MD  “Normal brain mri”  In _Kenhub_ Kenhub, 2023  URL: https://www.kenhub.com/en/library/anatomy/normal-brain-mri
  * [12] Olaf Ronneberger, Philipp Fischer and Thomas Brox  “U-net: Convolutional networks for biomedical image segmentation”  In _Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18_ , 2015, pp. 234–241  Springer
  * [13] Hera Shaheen, Shikha Agarwal and Prabhat Ranjan  “MinMaxScaler binary PSO for feature selection”  In _First International Conference on Sustainable Technologies for Computational Intelligence: Proceedings of ICTSCI 2019_ , 2020, pp. 705–716  Springer
  * [14] Jason Brownlee  “How to control the stability of training neural networks with the batch size”  In _Machine Learning Mastery_ 20, 2019
  * [15] Peng Shi et al.  “An improved U-net image segmentation method and its application for metallic grain size statistics”  In _Materials_ 15.13 MDPI, 2022, pp. 4417
  * [16] Shruti Jadon  “A survey of loss functions for semantic segmentation”  In _2020 IEEE conference on computational intelligence in bioinformatics and computational biology (CIBCB)_ , 2020, pp. 1–7  IEEE
  * [17] Seyed Sadegh Mohseni Salehi, Deniz Erdogmus and Ali Gholipour  “Tversky loss function for image segmentation using 3D fully convolutional deep networks”  In _Machine Learning in Medical Imaging: 8th International Workshop, MLMI 2017, Held in Conjunction with MICCAI 2017, Quebec City, QC, Canada, September 10, 2017, Proceedings 8_ , 2017, pp. 379–387  Springer
  * [18] Jin Xu et al.  “Reluplex made more practical: Leaky ReLU”  In _2020 IEEE Symposium on Computers and communications (ISCC)_ , 2020, pp. 1–7  IEEE
  * [19] Mobarakol Islam et al.  “Brain tumor segmentation and survival prediction using 3D attention UNet”  In _Brainlesion: Glioma, Multiple Sclerosis, Stroke and Traumatic Brain Injuries: 5th International Workshop, BrainLes 2019, Held in Conjunction with MICCAI 2019, Shenzhen, China, October 17, 2019, Revised Selected Papers, Part I 5_ , 2020, pp. 262–272  Springer


# All Tokens Matter: Token Labeling for Training Better Vision Transformers

Zihang Jiang1 &Qibin Hou1 &Li Yuan1 &Daquan Zhou1 &Yujun Shi1 Xiaojie Jin2
&Anran Wang2 &Jiashi Feng1 1National University of Singapore &2ByteDance
<EMAIL_ADDRESS>
<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>Work done as
an intern at ByteDance AI Lab.

###### 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 all the training
details are publicly available at
https://github.com/zihangJiang/TokenLabeling.

##  1 Introduction

Transformers [39] have achieved great performance for almost all the natural
language processing (NLP) tasks over the past years [4, 15, 25]. 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 [7, 16, 26, 36, 40, 46],
especially the strong benefits of the self-attention mechanism in building
long-range dependencies between pairs of input tokens.

\begin{overpic}[width=433.62pt]{figures/TeaserComp.png} \end{overpic} Figure
1: Comparison between the proposed LV-ViT and other recent works based on
vision transformers, including T2T-ViT [46], ConViT [13], BoTNet [31], DeepViT
[59], DeiT [36], ViT [16], Swin Transformer [26], LambdaNet [1], CvT [43],
CrossViT [7], PVT [40], CaiT [37]. Note that we only show models whose model
sizes are under 100M. As can be seen, our LV-ViT achieves the best results
using the least amount of learnable parameters. The default test resolution is
$224\times 224$ unless specified after @.

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 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 [14], 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 [8, 45], as well as more challenging tasks such as image
classification [10, 16], object detection [5, 12, 55, 61], segmentation [8,
33, 41] and image generation [29]. Some works also extend transformers for
video and 3D point cloud processing [50, 53, 60].

Figure 2: Pipeline of training vision transformers with token labeling. Other
than utilizing the class token (pink rectangle), we also take advantage of all
the output patch tokens (orange rounded rectangle) by assigning each patch
token an individual location-specific prediction generated by a machine
annotator [3] as supervision (see the part in the red dash rectangle). Our
proposed token labeling method can be treated as an auxiliary objective to
provide each patch token the local details that aid vision transformers to
more accurately locate and recognize the target objects. Note that the
traditional vision transformer training does not include the red dash
rectangle part.

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 [19] 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 [22, 47] to improve model
performance.

Some recent works [7, 17, 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 [21, 26, 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 [16] and then describe the proposed training objective— _token
labeling_.

###  3.1 Revisiting Vision Transformer

A typical vision transformer [16] 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 $[X^{cls},X^{1},...,X^{N}]$, where
$N$ is the total number of patch tokens, and $X^{cls}$ and $X^{1},...,X^{N}$
correspond to the class token and the patch tokens, respectively. The
classification loss for image $I$ can be written as

$L_{cls}=H(X^{cls},y^{cls}),$ (1)

where $H(\cdot,\cdot)$ is the softmax cross-entropy loss and $y^{cls}$ 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 $y^{cls}$ but also a $K\times N$ matrix
or called a $K$-dimensional score map as represented by $[y^{1},...,y^{N}]$,
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 $X^{1},...,X^{N}$ and the corresponding labels $[y^{1},...,y^{N}]$, the
token labeling objective can be defined as

$L_{tl}=\frac{1}{N}\sum_{i=1}^{N}H(X^{i},y^{i}).$ (2)

Recall that $H$ is the cross-entropy loss. Therefore, the total loss function
can be written as

$\displaystyle L_{total}$ $\displaystyle=H(X^{cls},y^{cls})+\beta\cdot
L_{tl},$ (3)
$\displaystyle=H(X^{cls},y^{cls})+\beta\cdot\frac{1}{N}\sum_{i=1}^{N}H(X^{i},y^{i}),$
(4)

where $\beta$ 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.

Figure 3: Comparison between CutMix [48] (Left) and our proposed MixToken
(Right). CutMix is operated on the input images. This results in patches
containing mixed regions from the two images (see the patches enclosed by red
bounding boxes). Differently, MixToken targets at mixing tokens after patch
embedding. This enables each token after patch embedding to have clean content
as shown in the right part of this figure. The detailed advantage of MixToken
can be found in Sec. 4.2.

###  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 $I_{1},I_{2}$ and their
corresponding token labels $Y_{1}=[y_{1}^{1},...,y^{N}_{1}]$ as well as
$Y_{2}=[y^{1}_{2},...,y^{N}_{2}]$, we first feed the two images into the patch
embedding module to tokenize each as a sequence of tokens, resulting in
$T_{1}=[t^{1}_{1},...,t^{N}_{1}]$ and $T_{2}=[t^{1}_{2},...,t^{N}_{2}]$. Then,
we produce a new sequence of tokens by applying MixToken using a binary mask
$M$ as follows:

$\hat{T}=T_{1}\odot M+T_{2}\odot(1-M),$ (5)

where $\odot$ 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$:

$\hat{Y}=Y_{1}\odot M+Y_{2}\odot(1-M).$ (6)

The label for the class token can be written as

$\hat{y^{cls}}=\bar{M}y_{1}^{cls}+(1-\bar{M})y_{2}^{cls},$ (7)

where $\bar{M}$ is the average of all element values of $M$.

##  4 Experiments

###  4.1 Experiment Setup

We evaluate our method on the ImageNet [14] dataset. All experiments are built
and conducted upon PyTorch [30] 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 [11], 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 [28] with a linear
learning rate scaling strategy $lr=10^{-3}\times\frac{batch\\_size}{640}$ and
$5\times 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 [24]
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\times 576$, and
the dimension of the corresponding output score map for each image is
$L\in\mathbb{R}^{18\times 18\times 1000}$. During training, the target labels
for the tokens are generated by applying RoIAlign [18] 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.

###  4.2 Ablation Analysis

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 [16] 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.

Table 1: Performance of the proposed LV-ViT with different model sizes. Here, ‘depth’ denotes the number of transformer blocks used in different models. By default, the test resolution is set to $224\times 224$ except the last one which is $288\times 288$. Name | Depth | Embed dim. | MLP Ratio | #Heads | #Parameters | Resolution | Top-1 Acc. (%)
---|---|---|---|---|---|---|---
LV-ViT-T | 12 | 240 | 3.0 | 4 | 8.5M | 224x224 | 79.1
LV-ViT-S | 16 | 384 | 3.0 | 6 | 26M | 224x224 | 83.3
LV-ViT-M | 20 | 512 | 3.0 | 8 | 56M | 224x224 | 84.1
LV-ViT-L | 24 | 768 | 3.0 | 12 | 150M | 288x288 | 85.3

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 3, 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.
---|---|---|---|---
✓ | ✗ | ✓ | ✓ | 83.3
✗ | ✗ | ✓ | ✓ | 81.3
✓ | ✓ | ✓ | ✓ | 83.1
✓ | ✗ | ✗ | ✓ | 83.0
✓ | ✗ | ✓ | ✗ | 82.8

Data Augmentation: Here, we study the compatibility of MixToken with other
augmentation techniques, such as MixUp [52], CutOut [57] and RandAug [11]. 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.

Figure 4: Left: LV-ViT ImageNet Top-1 Accuracy w.r.t. the token participation
rate while applying token labeling. Token participation rate indicates the
percentage of patch tokens involved in computing the token labeling loss. This
experiment reflects that all tokens matter for vision transformers. Right: LV-
ViT-S ImageNet Top-1 Accuracy w.r.t. different annotator models. The point
size indicates the parameter number of the annotator model. Clearly, our token
labeling objective is robust to different annotator models.

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.

|  |
---|---|---
Figure 5: Performance of the proposed token labeling objective on three
different vision transformers: DeiT [36] (Left), T2T-ViT [46] (Middle), and
LV-ViT (Right). Our method has a consistent improvement on all 7 different ViT
models.

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.

Table 4: Top-1 accuracy comparison with other methods on ImageNet [14] and ImageNet Real [2]. All models are trained without external data. With the same computation and parameter constraint, our model consistently outperforms other CNN-based and transformer-based counterparts. The results of CNNs and ViT are referenced from [37]. | Network | Params | FLOPs | Train size | Test size | Top-1(%) | Real Top-1 (%)
---|---|---|---|---|---|---|---
CNNs |  EfficientNet-B5 [34] |  030M |  009.9B | $456$ | $456$ | 83.6 | 88.3
EfficientNet-B7 [34] |  066M |  037.0B | $600$ | $600$ | 84.3 | _
Fix-EfficientNet-B8 [34, 38] |  087M |  089.5B | $672$ | $800$ | 85.7 | 90.0
NFNet-F3 [3] | 255M | 114.8B | $320$ | $416$ | 85.7 | 89.4
NFNet-F4 [3] | 316M | 215.3B | $384$ | $512$ | 85.9 | 89.4
NFNet-F5 [3] | 377M | 289.8B | $416$ | $544$ | 86.0 | 89.2
Transformers |  ViT-B/16 [16] |  086M |  055.4B | $224$ | $384$ | 77.9 | 83.6
ViT-L/16 [16] | 307M | 190.7B | $224$ | $384$ | 76.5 | 82.2
T2T-ViT-14 [46] |  022M |  005.2B | $224$ | $224$ | 81.5 | _
T2T-ViT-14$\uparrow$384 [46] |  022M |  017.1B | $224$ | $384$ | 83.3 | _
CrossViT [7] |  045M |  056.6B | $224$ | $480$ | 84.1 | _
Swin-B [26] |  088M |  047.0B | $224$ | $384$ | 84.2 | _
TNT-B [17] |  066M |  014.1B | $224$ | $224$ | 82.8 | _
DeepViT-S [59] |  027M |  006.2B | $224$ | $224$ | 82.3 | _
DeepViT-L [59] |  055M |  012.5B | $224$ | $224$ | 83.1 | _
DeiT-S [36] |  022M |  004.6B | $224$ | $224$ | 79.9 | 85.7
DeiT-B [36] |  086M |  017.5B | $224$ | $224$ | 81.8 | 86.7
DeiT-B$\uparrow$384 [36] |  086M |  055.4B | $224$ | $384$ | 83.1 | 87.7
BoTNet-S1-128 [31] | 79.1M |  019.3B | 256 | 256 | 84.2 | -
BoTNet-S1-128$\uparrow$384 [31] | 79.1M |  045.8B | 256 | 384 | 84.7 | -
CaiT-S36$\uparrow$384 [37] |  068M |  048.0B | $224$ | $384$ | 85.4 | 89.8
CaiT-M36 [37] | 271M |  053.7B | $224$ | $224$ | 85.1 | 89.3
|  CaiT-M36$\uparrow$448 [37] | 271M | 247.8B | $224$ | $448$ | 86.3 | 90.2
Ours LV-ViT | LV-ViT-S |  026M |  006.6B | $224$ | $224$ | 83.3 | 88.1
LV-ViT-S$\uparrow$384 |  026M |  022.2B | $224$ | 384 | 84.4 | 88.9
LV-ViT-M |  056M |  016.0B | $224$ | $224$ | 84.1 | 88.4
LV-ViT-M$\uparrow$384 |  056M |  042.2B | $224$ | $384$ | 85.4 | 89.5
LV-ViT-L | 150M |  059.0B | $288$ | $288$ | 85.3 | 89.3
LV-ViT-L$\uparrow$448 | 150M | 157.2B | $288$ | $448$ | 85.9 | 89.7
LV-ViT-L$\uparrow$448 | 150M | 157.2B | $448$ | $448$ | 86.2 | 89.9
|  LV-ViT-L$\uparrow$512 | 151M | 214.8B | $448$ | $512$ | 86.4 | 90.1

###  4.3 Comparison to Other Methods

We compare our proposed model LV-ViT with other state-of-the-art methods in
Table 4. For small-sized models, when the test resolution is set to $224\times
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 medium-sized
models, when the test resolution is set to $384\times 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 $512\times 512$ achieves
an 86.4% top-1 accuracy, which is better than CaiT-M36 [37] but with far fewer
parameters and FLOPs.

###  4.4 Semantic Segmentation on ADE20K

It has been shown in [20] 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 [26], 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 [27] and UperNet [44] as our
segmentation frameworks and use the mmseg toolbox to implement. During
training, following [26], 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\times 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 5, 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 6. 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.

Table 5: Transfer performance of the proposed LV-ViT in semantic segmentation. We take two classic methods, FCN and UperNet, as segmentation architectures and show both single-scale (SS) and multi-scale (MS) results on the validation set. Method | Token Labeling | Model Size | mIoU (SS) | P. Acc. (SS) | mIoU (MS) | P. Acc. (MS)
---|---|---|---|---|---|---
LV-ViT-S + FCN | ✗ | 30M | 46.1 | 81.9 | 47.3 | 82.6
LV-ViT-S + FCN | ✓ | 30M | 47.2 | 82.4 | 48.4 | 83.0
LV-ViT-S + UperNet | ✗ | 44M | 46.5 | 82.1 | 47.6 | 82.7
LV-ViT-S + UperNet | ✓ | 44M | 47.9 | 82.6 | 48.6 | 83.1
Table 6: Comparison with previous work on ADE20K validation set. As far as we know, our LV-ViT-L + UperNet achieves the best result on ADE20K with only ImageNet-1K as training data in pretraining. †Pretrained on ImageNet-22K. | Backbone | Segmentation Architecture | Model Size | mIoU (MS) | Pixel Acc. (MS)
---|---|---|---|---|---
CNNs | ResNet-269 |  PSPNet [54] | - | 44.9 | 81.7
ResNet-101 |  UperNet [44] | 86M | 44.9 | -
ResNet-101 |  Strip Pooling [23] | - | 45.6 | 82.1
ResNeSt200 |  DeepLabV3+ [9] | 88M | 48.4 | -
Transformers | DeiT-S | UperNet | 52M | 44.0 | -
ViT-Large† |  SETR [56] | 308M | 50.3 | 83.5
Swin-T [26] | UperNet | 60M | 46.1 | -
Swin-S [26] | UperNet | 81M | 49.3 | -
Swin-B [26] | UperNet | 121M | 49.7 | -
Swin-B† [26] | UperNet | 121M | 51.6 | -
LV-ViT | LV-ViT-S | FCN | 30M | 48.4 | 83.0
LV-ViT-S | UperNet | 44M | 48.6 | 83.1
LV-ViT-M | UperNet | 77M | 50.6 | 83.5
LV-ViT-L | UperNet | 209M | 51.8 | 84.1

##  5 Conclusions and Discussion

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.

## References

  * [1] Irwan Bello.  Lambdanetworks: Modeling long-range interactions without attention.  arXiv preprint arXiv:2102.08602, 2021.
  * [2] Lucas Beyer, Olivier J Hénaff, Alexander Kolesnikov, Xiaohua Zhai, and Aäron van den Oord.  Are we done with imagenet?  arXiv preprint arXiv:2006.07159, 2020.
  * [3] Andrew Brock, Soham De, Samuel L Smith, and Karen Simonyan.  High-performance large-scale image recognition without normalization.  arXiv preprint arXiv:2102.06171, 2021.
  * [4] Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  arXiv preprint arXiv:2005.14165, 2020.
  * [5] Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko.  End-to-end object detection with transformers.  arXiv preprint arXiv:2005.12872, 2020.
  * [6] Hila Chefer, Shir Gur, and Lior Wolf.  Transformer interpretability beyond attention visualization.  arXiv preprint arXiv:2012.09838, 2020.
  * [7] Chun-Fu Chen, Quanfu Fan, and Rameswar Panda.  Crossvit: Cross-attention multi-scale vision transformer for image classification.  arXiv preprint arXiv:2103.14899, 2021.
  * [8] Hanting Chen, Yunhe Wang, Tianyu Guo, Chang Xu, Yiping Deng, Zhenhua Liu, Siwei Ma, Chunjing Xu, Chao Xu, and Wen Gao.  Pre-trained image processing transformer.  arXiv preprint arXiv:2012.00364, 2020.
  * [9] Liang-Chieh Chen, Yukun Zhu, George Papandreou, Florian Schroff, and Hartwig Adam.  Encoder-decoder with atrous separable convolution for semantic image segmentation.  In Proceedings of the European conference on computer vision (ECCV), pages 801–818, 2018.
  * [10] Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and Ilya Sutskever.  Generative pretraining from pixels.  In International Conference on Machine Learning, pages 1691–1703. PMLR, 2020.
  * [11] Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le.  Randaugment: Practical automated data augmentation with a reduced search space.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, pages 702–703, 2020.
  * [12] Zhigang Dai, Bolun Cai, Yugeng Lin, and Junying Chen.  Up-detr: Unsupervised pre-training for object detection with transformers.  arXiv preprint arXiv:2011.09094, 2020.
  * [13] Stéphane d’Ascoli, Hugo Touvron, Matthew Leavitt, Ari Morcos, Giulio Biroli, and Levent Sagun.  Convit: Improving vision transformers with soft convolutional inductive biases.  arXiv preprint arXiv:2103.10697, 2021.
  * [14] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In 2009 IEEE conference on computer vision and pattern recognition, pages 248–255. Ieee, 2009.
  * [15] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  arXiv preprint arXiv:1810.04805, 2018.
  * [16] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  arXiv preprint arXiv:2010.11929, 2020.
  * [17] Kai Han, An Xiao, Enhua Wu, Jianyuan Guo, Chunjing Xu, and Yunhe Wang.  Transformer in transformer.  arXiv preprint arXiv:2103.00112, 2021.
  * [18] Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick.  Mask r-cnn.  In Proceedings of the IEEE international conference on computer vision, pages 2961–2969, 2017.
  * [19] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016.
  * [20] Tong He, Zhi Zhang, Hang Zhang, Zhongyue Zhang, Junyuan Xie, and Mu Li.  Bag of tricks for image classification with convolutional neural networks.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 558–567, 2019.
  * [21] Byeongho Heo, Sangdoo Yun, Dongyoon Han, Sanghyuk Chun, Junsuk Choe, and Seong Joon Oh.  Rethinking spatial dimensions of vision transformers.  arXiv preprint arXiv:2103.16302, 2021.
  * [22] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.  Distilling the knowledge in a neural network.  arXiv preprint arXiv:1503.02531, 2015.
  * [23] Qibin Hou, Li Zhang, Ming-Ming Cheng, and Jiashi Feng.  Strip pooling: Rethinking spatial pooling for scene parsing.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 4003–4012, 2020.
  * [24] Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger.  Deep networks with stochastic depth.  In European conference on computer vision, pages 646–661. Springer, 2016.
  * [25] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov.  Roberta: A robustly optimized bert pretraining approach.  arXiv preprint arXiv:1907.11692, 2019.
  * [26] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  arXiv preprint arXiv:2103.14030, 2021.
  * [27] Jonathan Long, Evan Shelhamer, and Trevor Darrell.  Fully convolutional networks for semantic segmentation.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 3431–3440, 2015.
  * [28] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  arXiv preprint arXiv:1711.05101, 2017.
  * [29] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Łukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran.  Image transformer.  arXiv preprint arXiv:1802.05751, 2018.
  * [30] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al.  Pytorch: An imperative style, high-performance deep learning library.  In Advances in neural information processing systems, pages 8026–8037, 2019.
  * [31] Aravind Srinivas, Tsung-Yi Lin, Niki Parmar, Jonathon Shlens, Pieter Abbeel, and Ashish Vaswani.  Bottleneck transformers for visual recognition.  arXiv preprint arXiv:2101.11605, 2021.
  * [32] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov.  Dropout: a simple way to prevent neural networks from overfitting.  The journal of machine learning research, 15(1):1929–1958, 2014\.
  * [33] Zhiqing Sun, Shengcao Cao, Yiming Yang, and Kris Kitani.  Rethinking transformer-based set prediction for object detection.  arXiv preprint arXiv:2011.10881, 2020.
  * [34] Mingxing Tan and Quoc V Le.  Efficientnet: Rethinking model scaling for convolutional neural networks.  arXiv preprint arXiv:1905.11946, 2019.
  * [35] Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, et al.  Mlp-mixer: An all-mlp architecture for vision.  arXiv preprint arXiv:2105.01601, 2021.
  * [36] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou.  Training data-efficient image transformers & distillation through attention.  arXiv preprint arXiv:2012.12877, 2020.
  * [37] Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou.  Going deeper with image transformers.  arXiv preprint arXiv:2103.17239, 2021.
  * [38] Hugo Touvron, Andrea Vedaldi, Matthijs Douze, and Hervé Jégou.  Fixing the train-test resolution discrepancy.  arXiv preprint arXiv:1906.06423, 2019.
  * [39] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30:5998–6008, 2017.
  * [40] Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao.  Pyramid vision transformer: A versatile backbone for dense prediction without convolutions.  arXiv preprint arXiv:2102.12122, 2021.
  * [41] Yuqing Wang, Zhaoliang Xu, Xinlong Wang, Chunhua Shen, Baoshan Cheng, Hao Shen, and Huaxia Xia.  End-to-end video instance segmentation with transformers.  arXiv preprint arXiv:2011.14503, 2020.
  * [42] Ross Wightman.  Pytorch image models.  https://github.com/rwightman/pytorch-image-models, 2019.
  * [43] Haiping Wu, Bin Xiao, Noel Codella, Mengchen Liu, Xiyang Dai, Lu Yuan, and Lei Zhang.  Cvt: Introducing convolutions to vision transformers.  arXiv preprint arXiv:2103.15808, 2021.
  * [44] Tete Xiao, Yingcheng Liu, Bolei Zhou, Yuning Jiang, and Jian Sun.  Unified perceptual parsing for scene understanding.  In Proceedings of the European Conference on Computer Vision (ECCV), pages 418–434, 2018.
  * [45] Fuzhi Yang, Huan Yang, Jianlong Fu, Hongtao Lu, and Baining Guo.  Learning texture transformer network for image super-resolution.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 5791–5800, 2020.
  * [46] Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, Francis EH Tay, Jiashi Feng, and Shuicheng Yan.  Tokens-to-token vit: Training vision transformers from scratch on imagenet.  arXiv preprint arXiv:2101.11986, 2021.
  * [47] Li Yuan, Francis EH Tay, Guilin Li, Tao Wang, and Jiashi Feng.  Revisiting knowledge distillation via label smoothing regularization.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 3903–3911, 2020.
  * [48] Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo.  Cutmix: Regularization strategy to train strong classifiers with localizable features.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 6023–6032, 2019.
  * [49] Sangdoo Yun, Seong Joon Oh, Byeongho Heo, Dongyoon Han, Junsuk Choe, and Sanghyuk Chun.  Re-labeling imagenet: from single to multi-labels, from global to localized labels.  arXiv preprint arXiv:2101.05022, 2021.
  * [50] Yanhong Zeng, Jianlong Fu, and Hongyang Chao.  Learning joint spatial-temporal transformations for video inpainting.  In European Conference on Computer Vision, pages 528–543. Springer, 2020.
  * [51] Hang Zhang, Chongruo Wu, Zhongyue Zhang, Yi Zhu, Zhi Zhang, Haibin Lin, Yue Sun, Tong He, Jonas Muller, R. Manmatha, Mu Li, and Alexander Smola.  Resnest: Split-attention networks.  arXiv preprint arXiv:2004.08955, 2020.
  * [52] Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz.  mixup: Beyond empirical risk minimization.  arXiv preprint arXiv:1710.09412, 2017.
  * [53] Hengshuang Zhao, Li Jiang, Jiaya Jia, Philip Torr, and Vladlen Koltun.  Point transformer.  arXiv preprint arXiv:2012.09164, 2020.
  * [54] Hengshuang Zhao, Jianping Shi, Xiaojuan Qi, Xiaogang Wang, and Jiaya Jia.  Pyramid scene parsing network.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2881–2890, 2017.
  * [55] Minghang Zheng, Peng Gao, Xiaogang Wang, Hongsheng Li, and Hao Dong.  End-to-end object detection with adaptive clustering transformer.  arXiv preprint arXiv:2011.09315, 2020.
  * [56] Sixiao Zheng, Jiachen Lu, Hengshuang Zhao, Xiatian Zhu, Zekun Luo, Yabiao Wang, Yanwei Fu, Jianfeng Feng, Tao Xiang, Philip HS Torr, et al.  Rethinking semantic segmentation from a sequence-to-sequence perspective with transformers.  arXiv preprint arXiv:2012.15840, 2020.
  * [57] Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, and Yi Yang.  Random erasing data augmentation.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 13001–13008, 2020.
  * [58] Bolei Zhou, Hang Zhao, Xavier Puig, Tete Xiao, Sanja Fidler, Adela Barriuso, and Antonio Torralba.  Semantic understanding of scenes through the ade20k dataset.  International Journal of Computer Vision, 127(3):302–321, 2019\.
  * [59] Daquan Zhou, Bingyi Kang, Xiaojie Jin, Linjie Yang, Xiaochen Lian, Qibin Hou, and Jiashi Feng.  Deepvit: Towards deeper vision transformer.  arXiv preprint arXiv:2103.11886, 2021.
  * [60] Luowei Zhou, Yingbo Zhou, Jason J Corso, Richard Socher, and Caiming Xiong.  End-to-end dense video captioning with masked transformer.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 8739–8748, 2018.
  * [61] Xizhou Zhu, Weijie Su, Lewei Lu, Bin Li, Xiaogang Wang, and Jifeng Dai.  Deformable detr: Deformable transformers for end-to-end object detection.  arXiv preprint arXiv:2010.04159, 2020.

##  Appendix A More Experiment Details

We show the default hyper-parameters for our ImageNet classification
experiments in Table 7. In addition, for fine-tuning on larger image
resolution, we set batch size to 512, learning rate to 5e-6, weight decay to
1e-8 and fine-tune 30 epochs. Other hyper-parameters are set the same as
default. During training, a machine node with 8 NVIDIA V100 GPUs (32G memory)
is required. When fine-tuning our large model with image resolution of
$448\times 448$, we need 4 machine nodes with the same GPU settings as above.

Table 7: Default hyper-parameters for our experiments. Note that we do not use the MixUp augmentation method when ReLabel or token labeling is used. Supervision | Standard | ReLabel | Token labeling
---|---|---|---
Epoch | 300 | 300 | 300
Batch size | 1024 | 1024 | 1024
LR |  1e-3 $\cdot\frac{\text{batch\\_size}}{\text{1024}}$ |  1e-3$\cdot\frac{\text{batch\\_size}}{\text{1024}}$ |  1e-3$\cdot\frac{\text{batch\\_size}}{\text{640}}$
LR decay | cosine | cosine | cosine
Weight decay | 0.05 | 0.05 | 0.05
Warmup epochs | 5 | 5 | 5
Dropout | 0 | 0 | 0
Stoch. Depth | 0.1 | 0.1 | 0.1
MixUp alpha | 0.8 | - | -
Erasing prob. | 0.25 | 0.25 | 0.25
RandAug | 9/0.5 | 9/0.5 | 9/0.5

##  Appendix B More Experiments

###  B.1 Training Technique Analysis

We present a summary of our modification and proposed token labeling method to
improve vision transformer models in Table 8. We take the DeiT-Small [36]
model as our baseline and show the performance increment as more training
techniqeus are added. In this subsection, we will ablate the proposed
modifications and evaluate the effectiveness of them.

Table 8: Ablation path from the DeiT-Small [36] baseline to our LV-ViT-S. All experiments expect for larger input resolution can be finished within 3 days using a single server node with 8 V100 GPUs. Clearly, with only 26M learnable parameters, the performance can be boosted from 79.9 to 84.4 (+4.5) using the proposed Token Labeling and other proposed training techniques. Training techniques | #Param. | Top-1 Acc. (%)
---|---|---
Baseline (DeiT-Small [36]) | 22M | 79.9
\+ More transformers ($12\rightarrow 16$) | 28M |  81.2 (+1.2)
\+ Less MLP expansion ratio ($4\rightarrow 3$) | 25M |  81.1 (+1.1)
\+ More convs for patch embedding | 26M |  82.2 (+2.3)
\+ Enhanced residual connection | 26M |  82.4 (+2.5)
\+ Token labeling with MixToken | 26M |  83.3 (+3.4)
\+ Input resolution ($224\rightarrow 384$) | 26M |  84.4 (+4.5)

Explicit inductive bias for patch embedding: Ablation analysis of patch
embedding is presented in Table 9. The baseline is set to the same as the
setting as presented in the third row of Table 8. Clearly, by adding more
convolutional layers and narrow the kernel size in the patch embedding, we can
see a consistent increase in the performance comparing to the original single-
layer patch embedding. However, when further increasing the number of
convolutional layer in patch embedding to 6, we do not observe any performance
gain. This indicates that using 4-layer convolutions in patch embedding is
enough. Meanwhile, if we use a larger stride to reduce the size of the feature
map, we can largely reduce the computation cost, but the performance also
drops. Thus, we only apply a convolution of stride 2 and kernel size 7 at the
beginning of the patch embedding module, followed by two convolutional layers
with stride 1 and kernel size 3. The feature map is finally tokenized to a
sequence of tokens using a convolutional layer of stride 8 and kernel size 8
(see the fifth line in Table 9).

Table 9: Ablation on patch embedding. Baseline is set as 16 layer ViT with embedding size 384 and MLP expansion ratio of 3. All convolutional layers except the last block have 64 filters. #Convs indicatie the total number of convolutions for patch embedding, while the kernel size and stride correspond to each layer are shown as a list in the table. #Convs | Kerenl size | Stride | Params | Top-1 Acc. (%)
---|---|---|---|---
1 | [16] | [16] | 25M | 81.1
2 | [7,8] | [2,8] | 25M | 81.4
3 | [7,3,8] | [2,2,4] | 25M | 81.4
3 | [7,3,8] | [2,1,8] | 26M | 81.9
4 | [7,3,3,8] | [2,1,1,8] | 26M | 82.2
6 | [7,3,3,3,3,8] | [2,1,1,1,1,8] | 26M | 82.2

Enhanced residual connection: We found that introducing a residual scaling
factor can also bring benefit as shown in Table 10. We found that using
smaller scaling factor can lead to better performance and faster convergence.
Part of the reason is that more information can be preserved in the main
branch, leading to less information loss and better performance.

Table 10: Ablation on enhancing residual connection by applying a scaling factor. Baseline is a 16-layer vision transformer with 4-layer convolutional patch embedding. Here, function $F$ represents either self-attention (SA) or feed forward (FF). Forward Function | #Parameters | Top-1 Acc. (%)
---|---|---
$X\longleftarrow X+F(X)$ | 26M | 82.2
$X\longleftarrow X+F(X)/2$ | 26M | 82.4
$X\longleftarrow X+F(X)/3$ | 26M | 82.4

Larger input resolution: To adapt our model to larger input image, we
interpolate the positional encoding and fine-tune the model on larger image
resolution for a few epochs. Token labeling objective as well as MixToken are
also used during fine-tuning. As can be seen from Table 8, fine-tuning on
larger input resolution of $384\times 384$ can improve the performance by
$1.1\%$ for our LV-ViT-S model.

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

We further explore the performance of token labeling on other CNN-based and
MLP-based models. Results are shown in Table 11. Besides our re-implementation
with more data augmentation and regularization technique, we also list the
results from the original papers. It shows that for both MLP-based and CNN-
based models, token labeling objective can still improve the performance over
the strong baselines by providing the location-specific dense supervision.

Table 11: Performance of the proposed token labeling objective on representative CNN-based (ResNeSt) and MLP-based (Mixer-MLP) models. Our method has a consistent improvement on all different models. Here † indicates results reported in original papers. Model |  Mixer-S/16 [35] |  Mixer-B/16 [35] |  Mixer-L/16 [35] |  ResNeSt-50 [51]
---|---|---|---|---
Token Labeling | ✗ | ✗ | ✓ | ✗ | ✗ | ✓ | ✗ | ✗ | ✓ | ✗ | ✗ | ✓
Parameters | 18M | 18M | 18M | 59M | 59M | 59M | 207M | 207M | 207M | 27M | 27M | 27M
Top-1 Acc. (%) |  73.8† | 75.6 | 76.1 |  76.4† | 78.3 | 79.5 |  71.6† | 77.7 | 80.1 |  81.1† | 80.9 | 81.5

###  B.3 Comparison with CaiT

CaiT [37] is currently the best transformer-based model. We list the
comparison of training hyper-parameters and model configuration with CaiT in
Table 12. It can be seen that using less training techniques, computations,
and smaller model size, our LV-ViT achieves identical result to the state-of-
the-art CaiT model.

Table 12: Comparison with CaiT [37]. Our model exploits less training techniques, model size, and computations but achieve identical result to CaiT. Settings | LV-ViT (Ours) |  CaiT [37]
---|---|---
Transformer Blocks | 20 | 36
#Head in Self-attention | 8 | 12
MLP Expansion Ratio | 3 | 4
Embedding Dimension | 512 | 384
Stochastic Depth [24] | 0.2 (Linear) | 0.2 (Fixed)
Rand Augmentation [11] | ✓ | ✓
CutMix Augmentation [48] |  | ✓
MixUp Augmentation [52] |  | ✓
LayerScaling [37] |  | ✓
Class Attention [37] |  | ✓
Knowledge Distillation |  | ✓
Enhanced Residuals (Ours) | ✓ |
MixToken (Ours) | ✓ |
Token Labeling (Ours) | ✓ |
Test Resolution | $384\times 384$ | $384\times 384$
Model Size | 56M | 69M
Computations | 42B | 48B
Training Epoch | 300 | 400
ImageNet Top-1 Acc. | 85.4 | 85.4

##  Appendix C Visualization

We apply the method proposed in [6] to visualize both DeiT-base and our LV-
ViT-S. Results are shown in Figure 6 and Figure 7. In Figure 6, we can observe
that our LV-ViT-S model performs better in locating the target objects and
hence yields better classification performance with high confidence. In Figure
7, we visualize the top-2 classes predicted by the two models. Noted that we
follow [6] to select images with at least 2 classes existing. It can be seen
that our LV-ViT-S trained with token labeling can accurately locate both
classes while the DeiT-base sometimes fails in locating the entire target
object for a certain class. This demonstrates that our token labeling
objective does help in improving models’ visual grounding capability because
of the location-specific token-level information.

\begin{overpic}[width=346.89731pt]{figures/Top1.pdf} \end{overpic} Figure 6:
Visual comparisons between DeiT-base and LV-ViT-S.
\begin{overpic}[width=303.53267pt]{figures/Top2.pdf} \end{overpic} Figure 7:
Visual comparisons between DeiT-base and LV-ViT-S for the top-2 predicted
classes.