[Paper Review] Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions

Wenhai Wang, Enze Xie+7 more|arXiv (Cornell University)|Feb 24, 2021

Advanced Neural Network Applications58 citations

TL;DR

The paper introduces Pyramid Vision Transformer (PVT), a convolution-free Transformer backbone with a multi-scale feature pyramid and spatial-reduction attention, enabling high-resolution dense prediction and competitive performance across detection, segmentation, and classification tasks.

ABSTRACT

Although using convolutional neural networks (CNNs) as backbones achieves great successes in computer vision, this work investigates a simple backbone network useful for many dense prediction tasks without convolutions. Unlike the recently-proposed Transformer model (e.g., ViT) that is specially designed for image classification, we propose Pyramid Vision Transformer~(PVT), which overcomes the difficulties of porting Transformer to various dense prediction tasks. PVT has several merits compared to prior arts. (1) Different from ViT that typically has low-resolution outputs and high computational and memory cost, PVT can be not only trained on dense partitions of the image to achieve high output resolution, which is important for dense predictions but also using a progressive shrinking pyramid to reduce computations of large feature maps. (2) PVT inherits the advantages from both CNN and Transformer, making it a unified backbone in various vision tasks without convolutions by simply replacing CNN backbones. (3) We validate PVT by conducting extensive experiments, showing that it boosts the performance of many downstream tasks, e.g., object detection, semantic, and instance segmentation. For example, with a comparable number of parameters, RetinaNet+PVT achieves 40.4 AP on the COCO dataset, surpassing RetinNet+ResNet50 (36.3 AP) by 4.1 absolute AP. We hope PVT could serve as an alternative and useful backbone for pixel-level predictions and facilitate future researches. Code is available at https://github.com/whai362/PVT.

Motivation & Objective

• Develop a pure Transformer backbone suitable for dense prediction tasks (detection, segmentation) without convolutions.
• Introduce a pyramid-based feature hierarchy to deliver multi-scale, high-resolution representations.
• Reduce computational and memory costs for high-resolution attention via spatial-reduction attention (SRA).
• Demonstrate PVT's effectiveness as a drop-in backbone across object detection, instance/semantic segmentation, and image classification.

Proposed method

• Represent the input as fine-grained patches (4x4) and create a four-stage pyramid to produce multi-scale feature maps (F1..F4).
• Use patch embedding at each stage to progressively shrink feature map resolution (4x, 8x, 16x, 32x strides).
• Replace standard multi-head attention with spatial-reduction attention (SRA) to lower compute and memory by reducing K and V before attention.
• Share a Transformer encoder per stage with specific L_i, E_i, N_i, and R_i configurations to balance accuracy and efficiency.
• Enable convolution-free end-to-end pipelines by integrating PVT with DETR for object detection and with standard heads for segmentation/detection tasks.

(a) CNNs: VGG [ 54 ] , ResNet [ 22 ] , etc .

Experimental results

Research questions

• RQ1 Can a pure Transformer backbone with a pyramid multi-scale structure replace CNN backbones for dense prediction tasks?
• RQ2 How can attention mechanisms be redesigned to handle high-resolution feature maps efficiently?
• RQ3 What are the trade-offs in accuracy and efficiency when using a pyramid Transformer versus CNNs or ViT on dense prediction benchmarks?

Key findings

• PVT variants outperform equivocal CNN backbones at similar parameter counts on COCO object detection with RetinaNet (e.g., PVT-Small 40.4 AP vs ResNet50 36.3 AP).
• PVT-Large achieves 42.6 AP on COCO with 30% fewer parameters than ResNeXt101-64x4d.
• For instance segmentation, PVT-Tiny/Small/Medium surpass ResNet-18/50/101 baselines in APm on COCO with comparable FLOPs.
• Semantic segmentation on ADE20K shows PVT backbones yielding higher mIoU than CNN backbones, with PVT-Large reaching 42.1 mIoU and 44.8 with multi-scale testing.
• A pure Transformer DETR pipeline with PVT (PVT+DETR) achieves 34.7 AP on COCO val2017, outperforming ResNet50-based DETR.
• Across ImageNet classification, PVT models are competitive with ViT/DeiT and traditional CNNs, though gains are more pronounced for dense prediction tasks.

(b) Vision Transformer [ 13 ]