ViTAE: Vision Transformer Advanced by Exploring Intrinsic Inductive Bias

NeurIPS 2021  ·  Yufei Xu, Qiming Zhang, Jing Zhang, DaCheng Tao ·

Transformers have shown great potential in various computer vision tasks owing to their strong capability in modeling long-range dependency using the self-attention mechanism. Nevertheless, vision transformers treat an image as 1D sequence of visual tokens, lacking an intrinsic inductive bias (IB) in modeling local visual structures and dealing with scale variance. Alternatively, they require large-scale training data and longer training schedules to learn the IB implicitly. In this paper, we propose a novel Vision Transformer Advanced by Exploring intrinsic IB from convolutions, ie, ViTAE. Technically, ViTAE has several spatial pyramid reduction modules to downsample and embed the input image into tokens with rich multi-scale context by using multiple convolutions with different dilation rates. In this way, it acquires an intrinsic scale invariance IB and is able to learn robust feature representation for objects at various scales. Moreover, in each transformer layer, ViTAE has a convolution block in parallel to the multi-head self-attention module, whose features are fused and fed into the feed-forward network. Consequently, it has the intrinsic locality IB and is able to learn local features and global dependencies collaboratively. Experiments on ImageNet as well as downstream tasks prove the superiority of ViTAE over the baseline transformer and concurrent works. Source code and pretrained models will be available at GitHub.

PDF Abstract NeurIPS 2021 PDF NeurIPS 2021 Abstract
Task Dataset Model Metric Name Metric Value Global Rank Result Benchmark
Video Object Segmentation DAVIS 2016 ViTAE-T-Stage Jaccard (Mean) 89.2 # 11
F-Score 90.4 # 15
J&F 89.8 # 13
Video Object Segmentation DAVIS 2017 ViTAE-T-Stage Jaccard (Mean) 79.4 # 2
J&F 82.5 # 1
F-Score 85.5 # 2
Image Classification ImageNet ViTAE-6M Top 1 Accuracy 77.9% # 792
Number of params 6.5M # 444
GFLOPs 4 # 191
Image Classification ImageNet ViTAE-T Top 1 Accuracy 75.3% # 880
GFLOPs 3.0 # 174
Image Classification ImageNet ViTAE-T-Stage Top 1 Accuracy 76.8% # 828
Number of params 4.8M # 394
GFLOPs 4.6 # 215
Image Classification ImageNet ViTAE-S-Stage Top 1 Accuracy 82.2% # 510
Number of params 19.2M # 533
GFLOPs 12.0 # 314
Image Classification ImageNet ViTAE-B-Stage Top 1 Accuracy 83.6% # 378
Number of params 48.5M # 716
GFLOPs 27.6 # 387
Image Classification ImageNet ViTAE-13M Top 1 Accuracy 81% # 614
Number of params 13.2M # 507
GFLOPs 6.8 # 246

Methods