Source-linked AI summary
Rethinking Spatial Dimensions of Vision Transformers
Byeongho Heo, Sangdoo Yun, Dongyoon Han, Sanghyuk Chun, Junsuk Choe, Seong Joon Oh
TL;DR
The paper asks whether CNNs’ spatial dimension conversion, largely absent from ViT, benefits transformer-based vision models. It introduces PiT with pooling-based spatial reduction and finds improved performance over ViT across classification, detection, and robustness evaluations.
Problem
ViT does not sufficiently reflect CNN design principles, including increasing channels while reducing spatial dimensions, whose effectiveness for transformer architectures remains to be investigated.
Method
The paper proposes PiT, a Vision Transformer with a newly designed pooling layer that applies ResNet-style spatial reduction, and analyzes its attention patterns.
Results
PiT improves over ViT across image classification, object detection, and robustness benchmarks, while using fewer FLOPs and achieving faster speed than corresponding ViT models.
Takeaways & Limitations
Spatial dimension transformation through pooling can be incorporated into ViT while harmonizing computational efficiency, generalization ability, and performance improvements.
Takeaways & Limitations
PiT was not designed with large-scale neural architecture search and could be further improved through network architecture search.
Abstract
from arXiv · showhide
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
1. Introduction
ViT brought self-attention into vision but retained transformer design conventions without sufficiently reflecting effective CNN principles. This paper investigates spatial dimension conversion, proposes PiT, and evaluates its performance and attention patterns against ViT.
- 1. Introduction: ViT applies self-attention to image patches while keeping spatial tokens unchanged throughout its layers.Unlike CNNs, ViT permits all image positions to interact through self-attention but does not include spatial reduction after embedding.
- 1. Introduction: CNNs increase channel size while decreasing spatial size, and this dimension conversion is linked to receptive fields, expressiveness, and generalization.Pooling performs the spatial reduction used by modern CNN architectures such as AlexNet, ResNet, and EfficientNet.
- 1. Introduction: PiT adds a newly designed pooling layer to ViT, enabling ResNet-style spatial reduction within the transformer structure.The paper first verifies the advantages of ResNet-style dimensions on CNNs before extending them to ViT.
- 1. Introduction: Attention-matrix entropy and average-distance analyses are used to examine attention patterns inside ViT and PiT layers.The analysis is intended to clarify the inner mechanism of the two architectures.
- 1. Introduction: PiT improves over ViT across image classification, object detection, and robustness evaluations.The experiments compare PiT and ViT at multiple scales and training environments, including COCO 2017 deformable DETR backbones and robustness benchmarks.
2. Related works
Prior work established spatial dimension reduction as a common CNN design pattern, while transformer work explored token pooling and self-attention for language and vision. The paper positions its approach against these conventions and compares spatial configurations in ResNet.
- 2. Related works: AlexNet, VGGnet, GoogLeNet, and ResNet use pooling or stride-based operations to reduce spatial dimensions while increasing channels.ResNet replaces max pooling with stride-2 convolutions, a method also used in recent architectures such as EfficientNet and MobileNet.
- 2. Related works: Most convolutional networks use spatial reduction, whereas ViT uniquely maintains the same spatial size across all layers.The supplied related-work discussion identifies this as a structural distinction between convolutional networks and ViT.
- 2. Related works: Funnel Transformer reduces tokens with pooling and skip connections, but its pooling strategy differs from PiT because NLP and vision architectures have different basic structures.The paper also situates self-attention-based backbones within broader computer-vision applications.
- 2. Related works: Experiments with ResNet50 compare ResNet-style and ViT-style spatial configurations for model capability, generalization performance, and model performance.The figure caption states that ResNet-style is better across these evaluated dimensions.
3. Revisiting spatial dimensions
The paper investigates whether CNN-style spatial dimension conversion benefits ViT and introduces PiT, which adds pooling-based spatial reduction. Experiments report improved capability, generalization, and performance compared with ViT.
- 3.2. Pooling-based Vision Transformer (PiT): PiT applies ResNet-style spatial reduction to ViT through a newly designed pooling layer.The paper investigates spatial dimensions, then designs PiT to introduce dimension conversion into ViT.
- 3.1. Dimension setting of CNN: ResNet-style dimension configuration improves model capability and generalization performance in ResNet, yielding higher validation accuracy.Compared with ViT-style dimensions, ResNet-style dimensions achieve lower training loss at equal FLOPs and higher validation accuracy.
- 3.2. Pooling-based Vision Transformer (PiT): PiT reshapes spatial tokens into a 3D structure, reduces spatial size, increases channels with depth-wise convolution, and reshapes responses for transformer blocks.Class and distillation tokens use an additional fully connected layer to adjust channel size.
- 3.2. Pooling-based Vision Transformer (PiT): At equal computation cost, PiT achieves lower training loss than ViT and improves validation accuracy as training accuracy increases.ViT does not improve validation accuracy when training accuracy increases, whereas PiT shows a corresponding validation-accuracy increase.
- 3.4. Architecture design: PiT changes attention patterns by increasing entropy and interaction distance in shallow layers and decreasing both in deep layers.The analysis measures attention entropy and relative interaction distance over spatial tokens, excluding the final transformer block.
- 3.4. Architecture design: PiT models are designed at scales comparable to ViT-Base, ViT-Small, and ViT-Tiny, with lower FLOPs, parameter size, and GPU latency than ViT.The architecture configuration compares spatial sizes, blocks, heads, channels, and FLOPs; the authors note that PiT was not developed through large-scale NAS.
4. Experiments
Experiments evaluate PiT against ViT and convolutional architectures across ImageNet classification, object detection, and robustness benchmarks. PiT generally improves performance and computation relative to ViT, while remaining weaker than CNNs at small scales and ResNet50 in detection.
- ImageNet classification: PiT models use fewer FLOPs and achieve faster speed than corresponding ViT models while showing higher or comparable performance across training settings.This pattern holds for vanilla, CutMix, DeiT, and distillation settings.
- ImageNet classification: PiT performance increases with model size across all training settings, whereas ViT sometimes shows no performance increase as model size grows.The authors associate this difference with alleviated generalization problems from pooling layers.
- ImageNet classification: At ResNet50 scale or higher, transformer architectures outperform convolutional architectures, but CNNs outperform ViT-Ti and PiT-Ti at small scale.PiT-B outperforms convolutional architectures, while PiT-S exceeds ResNet50 or matches the comparison through throughput against EfficientNet-b3.
- Object detection: PiT-S outperforms ViT-S as a Deformable-DETR backbone on COCO, with lower detector latency than both ResNet50- and ViT-S-based detectors.The PiT detector does not beat ResNet50 in performance, although it improves over ViT-S and has better latency than ResNet50.
- Robustness benchmarks: PiT performs better than ViT on all four robustness benchmarks despite comparable standard ImageNet accuracy, including a smaller occlusion drop than ResNet50.PiT drops from 80.8 to 74.6, while ResNet50 drops from 79.0 to 67.1; ResNet50 nevertheless outperforms ViT variants on BGC.
5. Conclusion
The paper argues that CNN-style spatial dimension transformation is important for self-attention vision architectures. PiT incorporates pooling into ViT and improves performance while supporting computational efficiency and generalization.
- Conclusion: CNN-style spatial dimensional transformation increases computational efficiency and generalization ability in the authors’ ResNet study.The transformation is performed by pooling or convolution with strides.
- Conclusion: PiT incorporates a pooling layer into ViT to harmonize spatial dimension transformation with self-attention architecture.The proposed design transfers the spatial reduction principle used in CNNs to ViT.
- Conclusion: The paper concludes that pooling, by considering spatial interaction ratio, is essential to self-attention-based architecture.This conclusion follows the paper’s extensive experiments on PiT and ViT.