Multi-Tailed Vision Transformer for Efficient Inference

TL;DR

A Multi-Tailed Vision Transformer that adopts multiple tails to produce visual sequences of different lengths for the following Transformer encoder and can achieve a significant reduction on FLOPs with no degradation of the accuracy and outperform compared methods in both accuracy and FLOPs.

Abstract

Recently, Vision Transformer (ViT) has achieved promising performance in image recognition and gradually serves as a powerful backbone in various vision tasks. To satisfy the sequential input of Transformer, the tail of ViT first splits each image into a sequence of visual tokens with a fixed length. Then the following self-attention layers constructs the global relationship between tokens to produce useful representation for the downstream tasks. Empirically, representing the image with more tokens leads to better performance, yet the quadratic computational complexity of self-attention layer to the number of tokens could seriously influence the efficiency of ViT's inference. For computational reduction, a few pruning methods progressively prune uninformative tokens in the Transformer encoder, while leaving the number of tokens before the Transformer untouched. In fact, fewer tokens as the input for the Transformer encoder can directly reduce the following computational cost. In this spirit, we propose a Multi-Tailed Vision Transformer (MT-ViT) in the paper. MT-ViT adopts multiple tails to produce visual sequences of different lengths for the following Transformer encoder. A tail predictor is introduced to decide which tail is the most efficient for the image to produce accurate prediction. Both modules are optimized in an end-to-end fashion, with the Gumbel-Softmax trick. Experiments on ImageNet-1K demonstrate that MT-ViT can achieve a significant reduction on FLOPs with no degradation of the accuracy and outperform other compared methods in both accuracy and FLOPs.

Figures (5)

• Figure 1: The throughput and confidence of prediction from DeiT-S, with a different number of tokens, e.g., 7$\times$7, 10$\times$10 and 14$\times$14. The 'tick' and 'cross' sign denote the right and false prediction respectively.
• Figure 2: The framework of proposed method contains two main components: scale predictor $\pi_\theta$ and multi-tailed vision transformer (MT-ViT). The number of scale $K$ is set to 3. (a) The tail predictor is a CNN-based model that determines the appropriate tail for the image. (b) By using the multiple tails in MT-ViT, patches with different sizes are all projected into a $d$ dimension embedding. This makes it possible to share the Transformer encoders and MLP head.
• Figure 3: Ablation study on the hyper-parameter $\alpha$ and $\eta$ in FLOPs regularization term. $\alpha$ is fixed to 0.5 in the left figure and $\eta$ is fixed to 0.5 in the right figure.
• Figure 4: The Accuracy-FLOPs curve of DVT and MT-ViT based on DeiT-S.
• Figure 5: The visualized results of tail predictor in ImageNet-1K. The images in each row are from the class 'Soccer', 'Pineapple', 'Car' and 'Castle', respectively. The decision of each tail predictor illustrates how the predictor translates to instance difficulty.