[Paper Review] Not All Patches are What You Need: Expediting Vision Transformers via Token Reorganizations

Youwei Liang, Chongjian Ge+4 more|arXiv (Cornell University)|Feb 16, 2022

Advanced Neural Network ApplicationsComputer Science95 citations

TL;DR

EViT reorganizes tokens in Vision Transformers by identifying attentive tokens and fusing inattentive ones during training to speed up inference without adding parameters, improving efficiency or accuracy under the same cost.

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 https://github.com/youweiliang/evit

Motivation & Objective

• Motivate acceleration of Vision Transformers (ViTs) by identifying token-level redundancy in MHSA.
• Propose a training-time token reorganization to preserve attentive tokens and fuse inattentive ones.
• Show that EViT reduces computation (MHSA and FFN) at inference without extra parameters.
• Demonstrate improved accuracy at the same compute budget by allowing more tokens (higher resolution) input.
• Explore the effect of using an oracle to guide token relevance and compare with existing acceleration methods.

Proposed method

• Compute the average attentiveness of the class token to each image token across MHSA heads.
• Keep the top-k attentive tokens and fuse the inattentive tokens into a single fused token.
• Fuse inattentive tokens via a weighted average using their attentiveness as weights (x_fused = sum_{i in N} a_i x_i).
• Incorporate token reorganization into ViT training at selected layers, with a cosine schedule for the keep rate.
• Optionally train with an oracle ViT to identify important tokens and initialize EViT with oracle weights.
• Demonstrate higher-resolution training by feeding more tokens at the same compute cost and validate with ImageNet experiments.

Experimental results

Research questions

• RQ1 Can token reorganization during ViT training reduce inference cost while maintaining accuracy?
• RQ2 Does fusing inattentive tokens preserve more information and stabilize training compared to simple token removal?
• RQ3 How does EViT perform under fixed compute and when given higher input resolutions?
• RQ4 What is the impact of using an oracle ViT to guide token selection on accuracy and efficiency?

Key findings

• EViT can speed up DeiT-S inference by about 50% with only ~0.3% accuracy loss at ImageNet.
• EViT achieves higher throughput at the same MACs and can maintain or improve accuracy when using higher-resolution inputs (e.g., DeiT-S gains 1% top-1 at the same compute).
• Inattentive token fusion helps preserve information and improves training stability and accuracy over token pruning alone.
• Training with an oracle further improves accuracy (e.g., DeiT-S from 79.8% to 80.7% in an oracle setup) while maintaining or reducing compute.
• Compared with DynamicViT, EViT delivers better accuracy at the same compute with fewer parameters and shows continued gains with longer training.
• EViT can be applied to different ViT variants (DeiT and LV-ViT) and yields favorable accuracy-throughput trade-offs across settings.