[Paper Review] Rethinking Spatial Dimensions of Vision Transformers

Byeongho Heo, Sangdoo Yun+4 more|arXiv (Cornell University)|Mar 30, 2021

Advanced Neural Network Applications39 references67 citations

TL;DR

PiT introduces a pooling-based pooling layer to ViT, creating a ResNet-style spatial dimension reduction for Vision Transformers, yielding better accuracy and generalization than ViT across ImageNet, detection, and robustness benchmarks.

ABSTRACT

Vision Transformer (ViT) extends the application range of transformers from language processing to computer vision tasks as being an alternative architecture against the existing convolutional neural networks (CNN). Since the transformer-based architecture has been innovative for computer vision modeling, the design convention towards an effective architecture has been less studied yet. From the successful design principles of CNN, we investigate the role of spatial dimension conversion and its effectiveness on transformer-based architecture. We particularly attend to the dimension reduction principle of CNNs; as the depth increases, a conventional CNN increases channel dimension and decreases spatial dimensions. We empirically show that such a spatial dimension reduction is beneficial to a transformer architecture as well, and propose a novel Pooling-based Vision Transformer (PiT) upon the original ViT model. We show that PiT achieves the improved model capability and generalization performance against ViT. Throughout the extensive experiments, we further show PiT outperforms the baseline on several tasks such as image classification, object detection, and robustness evaluation. Source codes and ImageNet models are available at https://github.com/naver-ai/pit

Motivation & Objective

• Investigate whether CNN-style spatial dimension reduction (downsampling) benefits transformer-based vision models.
• Extend ViT with a pooling-based mechanism to enable spatial downsampling without sacrificing token-based processing.
• Evaluate PiT across image classification, object detection, and robustness benchmarks to assess generalization and efficiency gains.

Proposed method

• Design a pooling layer for ViT that reshapes 2D token maps into 3D tensors, applies depthwise convolution-based spatial reduction and channel expansion, then reshapes back to 2D tokens.
• Incorporate two pooling layers in PiT to create three spatial scales, with channel expansion implemented by increasing multi-head attention heads.
• Construct PiT variants at scales corresponding to ViT (tiny, small, base) and adjust FLOPs/parameters to be competitive with ViT while reducing latency.
• Compare PiT against ViT under vanilla, CutMix, DeiT, and Distillation training regimes on ImageNet-1k.
• Analyze attention patterns via entropy and spatial distance to understand how PiT reshapes token interactions across layers.
• Evaluate PiT as a backbone for Deformable DETR on COCO and assess robustness on ImageNet benchmarks.

Figure 1 : Schematic illustration of dimension configurations of networks. We visualize ResNet50 [ 13 ] , Vision Transformer (ViT) [ 9 ] , and our Pooling-based Vision Transformer (PiT); (a) ResNet50 gradually downsamples the features from the input to the output; (b) ViT does not change the spatial

Experimental results

Research questions

• RQ1 Does applying CNN-style spatial dimension reduction to Vision Transformers improve model capability and generalization?
• RQ2 Can a pooling-based ViT (PiT) outperform ViT across image classification, object detection, and robustness benchmarks?
• RQ3 What are the effects of PiT's pooling on attention patterns compared to ViT, and how does this relate to performance and robustness?
• RQ4 Is PiT competitive with convolutional backbones in terms of accuracy, throughput, and latency across scales?
• RQ5 How does PiT behave under extended training regimes (long training, large-resolution fine-tuning)?

Key findings

• PiT generally outperforms ViT across several ImageNet training regimes, achieving higher accuracy with lower FLOPs and often better throughput.
• On ImageNet, PiT-B achieves 84.0% vs ViT-B 83.4% under similar settings; PiT-S also shows gains over ViT-S in several configurations.
• PiT improves as model size increases, mitigating ViT’s generalization gap observed at larger scales and under standard training tricks.
• PiT as a backbone for object detection (Deformable DETR on COCO) yields higher AP than ViT-S (39.4 vs. 36.9) with competitive latency.
• In robustness benchmarks, PiT-S shows consistent improvements over ViT-S across standard, Occ, IN-A, BGC, and FGSM tests.
• Overall, PiT demonstrates that introducing pooling-based spatial reduction into ViT can enhance performance and robustness while achieving favorable efficiency trade-offs.

(a) Model capability