PVT v2: Improved Baselines with Pyramid Vision Transformer

In plain words

Vision transformers cut an image into patches and let every patch look at every other patch. That works well, but the cost grows quickly with the image size, and the original pyramid version of this idea cut the image into patches that never touch each other. This paper keeps the same pyramid design and fixes three things.

First, the expensive "every patch looks at every patch" step is made cheap by first squeezing the patch grid down to a small fixed-size summary, so the work grows in step with the image instead of with its square. Second, the patches are cut so that neighbours overlap, which keeps local detail that hard cuts lose. Third, a small local filter is added inside the block that processes each token, which gives the network a sense of position without a fixed position table.

The authors present the result as a stronger set of baselines rather than a new idea. It improves image classification and the detection and segmentation tasks built on top of the backbone, and it is meant to be a reference point for later work on vision transformers.

TL;DR

PVT v2 is a four-stage pyramid vision transformer backbone. It changes the attention to pool keys and values to a 7x7 grid, which makes its cost linear in the number of image tokens §3.2 eq. 2, uses overlapping patch embedding §3.3, and adds a 3x3 depthwise convolution in the feed-forward block §3.4. Author-reported results include 83.8% ImageNet-1K top-1 for the largest model tab. 2 and 46.1 RetinaNet AP on COCO tab. 3. Use it when you need a transformer backbone for dense prediction; check the Pitfalls section before reproducing the training recipe.

Problem

PVT v1 has three limits, according to the authors §3.1:

• Its cost is high on high-resolution input, such as images whose shorter side is 800 pixels §3.1.
• It splits the image into non-overlapping patches, which loses some local continuity §3.1.
• Its position encoding has a fixed size, which does not suit inputs of other sizes §3.1.

The authors keep PVT v1's four-stage pyramid and fix these three limits §3.1. They state that the three changes are orthogonal to the PVT v1 framework §1.

Method

The network has four stages. Each stage starts with an overlapping patch embedding that downsamples the feature map by a stride S, followed by L encoder blocks. Resolution shrinks and channel count grows with depth, as in a ResNet §3.5. Stage 3 holds most of the computation §3.5.

Attention: linear spatial reduction (linear SRA). In PVT v1 the keys and values are shrunk by a strided convolution with ratio R. Here they are shrunk by average pooling to a fixed P x P grid, with P = 7 in all stages §3.2. The complexity of the original SRA is given by eq. 1 and that of linear SRA by eq. 2, with h x w x c the input size, R the reduction ratio and P the pool size §3.2. The attention cost no longer depends on the square of the image size. The queries still use every token pvt_v2.py.

\Omega(\mathrm{SRA}) = \frac{2h^2w^2c}{R^2} + hwc^2R^2 eq. 1

\Omega(\mathrm{Linear\ SRA}) = 2hwP^2c eq. 2

Overlapping patch embedding. Adjacent patches overlap by half their area, and the map is padded with zeros so that the resolution is kept §3.3. The authors implement this as a convolution with stride S, kernel size 2S-1 and padding S-1, producing an output of h/S x w/S §3.3. Stage 1 uses S = 4 and stages 2 to 4 use S = 2. The released code uses kernel 7 with stride 4 and padding 3 in stage 1, and kernel 3 with stride 2 and padding 1 in later stages pvt_v2.py pvt_v2.py.

Convolutional feed-forward network (CFFN). A 3x3 depthwise convolution with padding 1 sits between the first fully connected layer and the GELU of the feed-forward block §3.4 fig. 2. The authors use it to replace the fixed position encoding of PVT v1, since zero padding gives the network positional information §3.4.

Scaling. The series B0 to B5 changes the stage widths, the feed-forward expansion ratios (B5 uses 4 everywhere) and the depths; head counts and the reduction ratios are the same across variants §3.5 tab. 1. The linear variant B2-Li swaps the reduction ratio for the pooling size P tab. 1. The code builds the same variants with the same numbers [code: classification/pvt_v2.py#L329-L401].

Pseudocode. The paper gives no algorithm box. The sketch below follows the text and the released code. Lines marked as code-only are not described in the text.

# Overlapping patch embedding, the downsampling step between stages  [§3.3]
# Input:  feature map (B, C_in, H, W); stride S
# Output: tokens (B, N, C) and the new size height_out, width_out
def overlap_patch_embed(feature_map, stride, out_dim):
    # Kernel 2S-1 and padding S-1, so neighbouring patches overlap  [§3.3]
    # [code: classification/pvt_v2.py#L185-L186]
    projected = conv2d(
        feature_map, out_dim, kernel=2 * stride - 1,
        stride=stride, padding=stride - 1,
    )
    height_out, width_out = projected.shape[-2:]  # ceil(H/S), ceil(W/S)
    tokens = flatten_to_tokens(projected)  # (B, N, C)
    return layer_norm(tokens), height_out, width_out

# Linear spatial reduction attention  [§3.2] [eq. 2]
# Keys and values come from a fixed P x P grid, so their count is P^2 = 49
def linear_sra(tokens, height, width):
    # Put the tokens back on their grid  [code: classification/pvt_v2.py#L111]
    grid = to_grid(tokens, height, width)  # (B, C, H, W)
    # Average-pool to P x P, with P = 7  [§3.2]
    # [code: classification/pvt_v2.py#L77]
    pooled = adaptive_avg_pool(grid, output_size=7)  # (B, C, 7, 7)
    # Code only, not in the text: 1x1 conv, LayerNorm, GELU
    # [code: classification/pvt_v2.py#L78-L80]
    reduced = gelu(layer_norm(conv1x1(pooled)))  # (B, C, 7, 7)
    summary = flatten_to_tokens(reduced)  # (B, 49, C)
    # Queries use all N tokens; keys and values use the 49 summary tokens
    # [§3.2]
    queries = linear_q(tokens)  # (B, N, C)
    keys_values = linear_kv(summary)  # (B, 49, 2C)
    return multi_head_attention(queries, keys_values)  # scale = head_dim^-0.5

# Convolutional feed-forward network  [§3.4]
# Input:  tokens (B, N, C); height and width of the grid
# Output: tokens (B, N, C)
def conv_ffn(tokens, height, width, linear_variant):
    # Expand the width with the first linear layer  [tab. 1]
    # [code: classification/pvt_v2.py#L43]
    hidden = fc1(tokens)  # (B, N, E*C)
    # Code only, linear variants (B2-Li): ReLU after fc1
    # [code: classification/pvt_v2.py#L44-L45]
    if linear_variant:
        hidden = relu(hidden)
    # 3x3 depthwise conv, padding 1, gives local continuity and position
    # [§3.4]
    hidden = depthwise_conv3x3(to_grid(hidden, height, width))
    hidden = flatten_to_tokens(hidden)
    # GELU, then the second linear layer  [§3.4]
    # [code: classification/pvt_v2.py#L47-L49]
    return fc2(gelu(hidden))

# One encoder block, residual  [§3.2] [§3.4]
# Shown for linear variants. Other variants use the strided-conv SRA with
# ratio R instead  [§3.2] [eq. 1] [code: classification/pvt_v2.py#L72-L75]
# Code only: pre-norm residual layout
# [code: classification/pvt_v2.py#L163-L164]
def encoder_block(tokens, height, width, linear_variant):
    # Attention branch with the residual connection  [§3.2]
    tokens = tokens + linear_sra(layer_norm(tokens), height, width)
    # Feed-forward branch with the residual connection  [§3.4]
    tokens = tokens + conv_ffn(
        layer_norm(tokens), height, width, linear_variant,
    )
    return tokens

Stochastic depth is applied to both residual branches in the code; the paper does not describe it. See the Hyperparameters section for the rates.

Hyperparameters

Architecture (the same numbers for all tasks). Each cell gives C (channels), R (SRA reduction ratio), P (pool size for linear SRA), N (heads), E (FFN expansion ratio) and L (blocks). Stride S is 4 in stage 1 and 2 in stages 2 to 4 for every variant tab. 1.

Source: tab. 1. The code uses the same stage widths, heads and ratios pvt_v2.py. Depths differ for stage 2 of B2, B2-Li and B3; see Paper vs code.

Other architecture values:

Classification training on ImageNet-1K §4.1:

Object detection on COCO §4.2:

Semantic segmentation on ADE20K §4.3:

Results

All numbers are author-reported. The setup for each line is the task and split named in its row.

Full tables

Classification on ImageNet-1K validation, input 224 x 224 tab. 2. GFLOPs at 224 x 224. "-Li" is PVT v2 with linear SRA.

COCO val2017, RetinaNet 1x and Mask R-CNN 1x tab. 3. "#P" is parameters in millions. "-Li" is PVT v2 with linear SRA.

COCO val2017 box AP with four detectors, compared with Swin-T and ResNet-50 tab. 4. "#P" is parameters in millions; GFLOPs at 1280 x 800. The table does not state the schedule.

ADE20K validation, Semantic FPN, input 512 x 512 for GFLOPs tab. 5:

Ablation on ImageNet top-1 and COCO RetinaNet 1x AP tab. 6. The GFLOPs column is RetinaNet 1x.

Figures: fig. 1 compares SRA with linear SRA. fig. 2 shows the two convolution-based changes. fig. 3 shows qualitative detection, instance and segmentation outputs. fig. 4 plots GFLOPs against input size. Per the authors, the compute of PVT v2-B2-Li climbs far less steeply as the input grows than PVT v1-Small's, and follows a curve similar to ResNet-50's §4.4.2. No figure is reproduced on this page.

Paper vs code

Checked against the repository linked from the abstract, at the commit above (Apache-2.0; the README calls it the official implementation). The classification, detection and segmentation folders were all read.

• Linear SRA has extra layers in code. The text describes only average pooling to P x P §3.2. The code also applies a 1x1 convolution, LayerNorm and GELU to the pooled grid before the key and value projections pvt_v2.py pvt_v2.py. The released checkpoints are loaded into this code, so an implementation that matches them needs these layers.
• Linear variants add a ReLU in the feed-forward block. The text does not mention it. The code applies ReLU after the first linear layer when the variant is linear pvt_v2.py pvt_v2.py.
• Stage 2 depth differs for B2, B2-Li and B3. Table 1 gives L = 3 in stage 2 for these variants tab. 1. The code uses depths [3, 4, 6, 3] for B2 and B2-Li and [3, 4, 18, 3] for B3 pvt_v2.py pvt_v2.py pvt_v2.py. The README's B2 count (25.4M) matches Table 2, but I did not count parameters from the code, so I cannot tell from the sources which depth is right.
• Classification learning rate. The text gives 1e-3 for batch 128 §4.1. The code's default is 5e-4 and it scales the rate linearly by batch x world size / 512 main.py main.py. That gives 1e-3 only when 128 is the batch per GPU on 8 GPUs (a global batch of 1024). The text does not say which.
• Classification augmentation. The text lists crop, flip, label smoothing, mixup and random erasing §4.1. The code also uses RandAugment, colour jitter, CutMix and a repeated-augmentation sampler main.py main.py.
• Evaluation resize. The text states only a 224 centre crop §4.1. The code resizes the shorter side to 256 first datasets.py.
• Detection learning rate for Mask R-CNN 1x. The text gives 1e-4 §4.2. The 1x Mask R-CNN config sets 2e-4 mask_rcnn_pvt_v2_b2_fpn_1x_coco.py. The Sparse R-CNN 3x config sets about 1.8e-5 sparse_rcnn_pvt_v2_b2_fpn_300_proposals_crop_mstrain_480-800_3x_coco.py.
• Detection multi-scale range for 3x. The text gives a shorter side in 640-800 §4.2. The 3x configs use 480-800 in their multi-scale policy, with an AutoAugment-style crop branch mask_rcnn_pvt_v2_b2_fpn_3x_mstrain.py mask_rcnn_pvt_v2_b2_fpn_3x_mstrain.py.
• Stochastic depth and gradient clipping are set in code only: 0.1 for B0, B1, B2 and B2-Li, and 0.3 for B3, B4 and B5; clipping at 1.0 for B3 to B5 pvt_v2_b3.py pvt_v2_b5.py.
• Segmentation code. The repository's segmentation configs cover PVT v1 only fpn_pvt_s_ade20k_40k.py. Their learning rate is 1e-4 doubled for 8 GPUs, with 40k iterations and 4 images per GPU fpn_pvt_s_ade20k_40k.py [code: segmentation/configs/base/datasets/ade20k.py#L35]. This is consistent with the text's 1e-4 at batch 16 on 4 GPUs under linear scaling, but no PVT v2 segmentation config is released.
• Model sizes in the README differ from Table 2. The README lists 3.7M for B0 (table: 3.4M), 14.0M for B1 (table: 13.1M), and 83.1% top-1 for B3 (table: 83.2%) README.md.
• Classification head. The code averages the tokens after the final norm and applies one linear layer pvt_v2.py. The paper does not describe the head.

Pitfalls

• Text and table disagree on several numbers. The introduction says B5 "yields 83.8% top-1 error"; the table reports accuracy §1 tab. 2. The text says B5 is 0.5% above Swin and Twins §4.1; Twins-SVT-L is 83.7%, which is 0.1% lower than B5, while Swin-B is 83.3%.
• Segmentation sentence names the wrong model. The text compares "PVT-Large" with ResNeXt101-64x4d and reports 48.7 mIoU §4.3. The 48.7 figure is for PVT v2-B5; PVT v1-Large is 42.1 tab. 5. The 12% GFLOPs saving quoted there matches B5 (91.1 against 103.9) tab. 5.
• Overlap patch embedding output size. The text says the output is h/S x w/S §3.3. With kernel 2S-1 and padding S-1, the output is ceil(H/S) for any H. This matters for detection and segmentation, where inputs are not multiples of the stride. The padding is set by the code's patch_size // 2 pvt_v2.py.
• Only keys and values shrink. Queries still come from all N tokens, and keys and values number 49 per head. The cost is therefore linear in N, but the query and output projections still scale with N §3.2 eq. 2.
• The linear variant trades accuracy for compute. Against the convolutional SRA model, linear SRA cuts RetinaNet 1x GFLOPs from 290.7 to 227.4 and costs 1.0 AP (44.6 to 43.6) tab. 6. On ImageNet it gains 0.1 (82.0 to 82.1) tab. 6.
• The convolutional FFN and overlap embedding add cost. Going from row 2 to row 3 of Table 6 adds about 0.2M parameters and about 2 GFLOPs at RetinaNet input size tab. 6.
• No position embedding. Positional information comes from zero padding in the overlap embedding and the depthwise conv. The claim that the model handles variable resolution rests on the ablation rows and on the authors' statement §3.4 §4.4.1; no separate experiment varies the input size.
• Schedules are missing from some comparisons. Section 4.2 says "1x or 3x" and that all settings match the Swin comparison §4.2, but Table 4 does not state which schedule was used tab. 4.
• Numerical details the paper omits. Stochastic depth, warmup, minimum learning rate, LayerNorm epsilon, and the classification head are not stated in the paper. The values above come from the code main.py pvt_v2.py.
• Mixed precision is used in several 3x detection configs (ATSS, GFL, Cascade Mask R-CNN, Mask R-CNN-Li) atss_pvt_v2_b2_fpn_3x_mstrain_fp16.py. The paper does not mention it.

Lineage

• Builds on: PVT v1 (arXiv:2102.12122), whose four-stage pyramid and spatial reduction attention this paper keeps §1 §3.1. The README gives the PVT v1 id README.md.
• Compared against: Swin Transformer, which has its own card [/papers/2103.14030]; ViT and DeiT for classification tab. 2; CNNs such as ResNet, ResNeXt and RegNet; and T2T-ViT, TNT, CvT, CoaT and Twins tab. 2.
• Superseded by: not known from content in this wiki.