Not All Patches are What You Need: Expediting Vision Transformers via Token Reorganizations

Youwei Liang,Chongjian Ge,Zhan Tong,Yibing Song,Jue Wang,Pengtao Xie

DOI: https://doi.org/10.48550/arXiv.2202.07800

2022-04-14

Abstract:Vision Transformers (ViTs) take all the image patches as tokens and construct multi-head self-attention (MHSA) among them. Complete leverage of these image tokens brings redundant computations since not all the tokens are attentive in MHSA. Examples include that tokens containing semantically meaningless or distractive image backgrounds do not positively contribute to the ViT predictions. In this work, we propose to reorganize image tokens during the feed-forward process of ViT models, which is integrated into ViT during training. For each forward inference, we identify the attentive image tokens between MHSA and FFN (i.e., feed-forward network) modules, which is guided by the corresponding class token attention. Then, we reorganize image tokens by preserving attentive image tokens and fusing inattentive ones to expedite subsequent MHSA and FFN computations. To this end, our method EViT improves ViTs from two perspectives. First, under the same amount of input image tokens, our method reduces MHSA and FFN computation for efficient inference. For instance, the inference speed of DeiT-S is increased by 50% while its recognition accuracy is decreased by only 0.3% for ImageNet classification. Second, by maintaining the same computational cost, our method empowers ViTs to take more image tokens as input for recognition accuracy improvement, where the image tokens are from higher resolution images. An example is that we improve the recognition accuracy of DeiT-S by 1% for ImageNet classification at the same computational cost of a vanilla DeiT-S. Meanwhile, our method does not introduce more parameters to ViTs. Experiments on the standard benchmarks show the effectiveness of our method. The code is available at <a class="link-external link-https" href="https://github.com/youweiliang/evit" rel="external noopener nofollow">this https URL</a>

Computer Vision and Pattern Recognition,Machine Learning

What problem does this paper attempt to address?

The problem that this paper attempts to solve is the computational redundancy issue in Vision Transformers (ViTs) when processing images. Specifically, ViTs divide an image into multiple patches and perform multi - head self - attention (MHSA) calculations on these patches. However, not all patches contribute positively to the final prediction results. For example, patches containing semantically meaningless or interfering backgrounds do not contribute to ViT's prediction. Therefore, the paper proposes a method, that is, to accelerate the inference process of the ViT model by reorganizing the tokens in the image while maintaining or improving the recognition accuracy.

The main contributions of the paper are as follows:
1. **Reducing computational cost**: Under the same number of input image tokens, the inference efficiency is improved by reducing the amount of calculations in MHSA and FFN (feed - forward network). For example, for the DeiT - S model, this method can increase the inference speed by 50%, while the recognition accuracy of ImageNet classification only drops by 0.3%.
2. **Improving recognition accuracy**: While maintaining the same computational cost, by processing higher - resolution images, ViT can receive more image tokens, thereby improving the recognition accuracy. For example, under the same computational cost, this method can increase the recognition accuracy of DeiT - S by 1%.
3. **Not increasing the number of parameters**: This method does not introduce additional parameters, thus maintaining the simplicity of the model.

In conclusion, this paper aims to optimize the ViT model by reorganizing image tokens, reducing unnecessary calculations, improving the inference speed and recognition accuracy, while not increasing the complexity of the model.