Rethinking Spatial Dimensions of Vision Transformers

In plain words

A Vision Transformer (ViT) cuts a picture into small square patches and treats every patch the same way from the first layer to the last. Convolutional networks do something different: as they go deeper they make the picture smaller and give it more channels. That shrinking is a long-standing design habit, and ViT skips it.

This paper asks whether the habit helps transformers too. Its answer is a variant called PiT. It adds small pooling steps between groups of transformer layers. Each step makes the grid of patches coarser and widens the features, which lets the later layers work on a smaller map with more information per position.

The authors report that this makes the model better at the same compute cost, and that it generalizes better. It also does better on object detection and on a few robustness tests. Results are reported for the ImageNet classification benchmark and are the authors' own.

TL;DR

PiT inserts two pooling layers into ViT, so the token grid shrinks in two steps and the channel count grows. The pooling uses a depth-wise convolution for the spatial tokens and a linear layer for the class token. On ImageNet, PiT beats ViT at comparable or lower FLOPs and throughput is also higher, under the authors' settings. Use it when you need a ViT-style backbone at a fixed compute budget. Check the Pitfalls first: the paper's text and its code disagree on patch size for the base model.

Problem

ViT keeps the same number of spatial tokens in every layer §1. Its self-attention lets every token interact with every other, but the size of the area a token can attend to is still set by the token count, and the paper argues that a fixed grid lacks the reduction that makes CNNs generalize well §3.2. The authors report that ViT's validation accuracy stops improving as training accuracy rises, which they read as weak generalization at ImageNet scale §3.2 fig. 3.

Method

PiT is a ViT with two pooling layers. The layers split the sequence of transformer blocks into three stages with three spatial scales §3.2. The Transformer blocks work on the same token layout as ViT. Between stages, the pooling layer:

1. turns the spatial tokens back into a 3-D map §3.2;
2. applies a depth-wise convolution that halves the spatial size and increases the channels §3.2 fig. 4;
3. flattens the map back into tokens for the next stage §3.2;
4. passes the class token (which has no spatial position) through an extra fully connected layer that changes its channel count §3.2.

Channel growth is produced by adding attention heads: the channels of each stage are the base width times the head count of that stage §3.4 tab. 1. The paper also analyses ViT and PiT through two attention statistics. The first is the entropy of each attention map, the second is the mean spatial distance between a query token and the tokens it attends to, using relative positions §3.3 eq. 1 eq. 2. Both are averaged over validation images and heads, skip the class token and skip the final block §3.3. The paper reports this as a figure, not as numbers fig. 5.

# PiT forward pass  [§3.2] [code: pit.py#L168-L187]
# Input:  images (B, 3, 224, 224)
# Output: class logits (B, num_classes)

def pit_forward(images):
    # Patch embedding: strided conv, patch 16, stride 8 in the text [§3.4]
    feature_map = patch_embed(images)                 # (B, C0, 27, 27)
    # Learned position embedding, added once before the first stage
    # [code: pit.py#L102-L105] [code: pit.py#L171-L172]
    feature_map = feature_map + pos_embed
    # One learned class token per image  [code: pit.py#L173]
    class_tokens = expand(cls_token, batch_size)      # (B, 1, C0)
    for stage in range(num_pools):
        # Transformer blocks on spatial plus class tokens  [§3.2]
        # [code: pit.py#L38-L51]
        feature_map, class_tokens = transformer[stage](
            feature_map, class_tokens)
        # Pooling halves the spatial size and widens channels  [§3.2]
        # [code: pit.py#L64-L69]
        feature_map, class_tokens = pool[stage](
            feature_map, class_tokens)
    # The last stage has no pooling after it  [code: pit.py#L178]
    feature_map, class_tokens = transformer[num_pools](
        feature_map, class_tokens)
    # LayerNorm is applied to the class token only  [code: pit.py#L180]
    class_tokens = layer_norm(class_tokens)
    # Linear classifier on the first class token  [code: pit.py#L186]
    return head(class_tokens[:, 0])

# Pooling layer between two stages  [§3.2] [fig. 4]
# Input:  feature_map (B, C0, H0, W0), class_tokens (B, n, C0)
# Output: feature_map (B, C1, H0/2, W0/2), class_tokens (B, n, C1)

def pool(feature_map, class_tokens):
    # Depth-wise conv, kernel 3, stride 2, padding 1, groups C0
    # Pattern from the code, not stated in the paper  [code: pit.py#L59-L61]
    feature_map = depthwise_conv(feature_map, groups=C0)
    # Linear map on the class tokens; no spatial change  [code: pit.py#L62]
    # [code: pit.py#L66-L67]
    class_tokens = linear(class_tokens)
    return feature_map, class_tokens

Hyperparameters

Architecture: the paper gives the spatial size, block count, head count and channel width of each stage tab. 1. FLOPs are measured at 224 x 224 input §3.4. The code sets a base width per model (32 for Ti, 48 for XS and S, 64 for B) and multiplies it by the head count of each stage [code: pit.py#L214-L285].

Patch and embedding settings. The paper's text says patch size 16 and stride 8 for all PiT sizes §3.4. The code uses those values for Ti, XS and S, and patch 14 with stride 7 for B pit.py. Table 1 lists 31 x 31 for PiT-B, which matches the code's value and not the text's value (see Paper vs code).

Results

Author-reported. ImageNet-1k classification at 224 x 224 unless noted. The vanilla, CutMix, DeiT and distillation columns are four training settings §4.1. FLOPs and throughput are in the paper's units. Throughput was measured on one V100 with batch 128 §4.1.

ImageNet, same training recipe as ViT tab. 2

Comparison with CNNs and ViT, best reported settings tab. 3

Throughput was re-measured on the same V100 §4.1. Entries marked with a dagger use distillation.

Longer training and larger resolution tab. 4

COCO object detection, Deformable DETR, half resolution tab. 5

Latency is single-image, on a random noise image at 600 x 400 §4.2.

ImageNet robustness tab. 6

The caption names PiT-B, but the PiT row is labelled PiT-S and the text matches PiT-S (see Paper vs code).

Headline claims, as the paper states them: PiT beats ViT in every robustness benchmark despite similar standard accuracy §4.3; PiT-S beats the ViT-S detector by 2.5 AP tab. 5; and at 384 x 384 fine-tuning PiT does not beat ViT §4.1 tab. 4. The paper says PiT is superior to ViT in the vanilla and CutMix settings, and comparable or better under the DeiT and distillation settings §4.1.

Paper vs code

The official code is naver-ai/pit, linked from the abstract. It was cloned at commit 9d97a62e6a2a72a86685003998fcae700f952e18. The repository contains only the model definitions in pit.py, plus a README. It has no training script, no evaluation script and no attention-analysis script. Pretrained weights are linked from the README, not in the repo.

1. Patch size and stride for PiT-B. The paper's text says patch 16 and stride 8 for all PiT sizes §3.4. The code uses patch 14 and stride 7 for pit_b and pit_b_distilled pit.py pit.py. Table 1 lists 31 x 31 spatial size for PiT-B tab. 1. The code's values give a 31 x 31 map at 224 input, and the text's values give 27 x 27. The code therefore matches Table 1. Follow the code for B.
2. Position embedding. The paper does not describe a position embedding. The code adds a learned position embedding of shape (1, C0, H0, W0) to the feature map before the first stage pit.py [code: pit.py#L171-L172]. It is not pooled again.
3. Class token pooling. The paper says a fully connected layer adjusts the class token §3.2. The code uses a plain linear layer with bias, with no activation pit.py pit.py.
4. Pooling convolution details. The paper says depth-wise convolution §3.2 fig. 4. The code uses kernel 3, stride 2, padding 1, and groups equal to the input channels, with a bias pit.py. With groups equal to the input channels and more output channels, each input channel feeds several outputs. The paper does not give kernel size or padding.
5. Final normalization. The code applies LayerNorm to the class token only, after the last stage pit.py. The spatial tokens are not normalized before the classifier. The paper does not say this.
6. Distillation head. The paper describes a distillation token trained with a distillation loss §4.1. The code keeps two class tokens and two linear heads pit.py. In training it returns both logits. In evaluation it returns the average of the two logits pit.py. The paper does not mention averaging at test time. The distillation loss itself is not in the repo.
7. Initialization. The code initializes position and class tokens with a truncated normal (std 0.02) and LayerNorm to weight 1 and bias 0 [code: pit.py#L145-L152]. Linear and convolution layers keep the framework default. The paper does not state an initialization.
8. Transformer block. The code builds each block from timm's vision_transformer.Block with qkv_bias=True, MLP ratio 4 and LayerNorm epsilon 1e-6 pit.py. Activation and normalization placement therefore come from timm, which is an external dependency. The paper does not state these. I did not check timm's source.
9. FLOPs units. The paper's tables write FLOPs in B tab. 1 tab. 2. The README writes 0.71 G for PiT-Ti, and the paper writes 0.7 B. The values agree.
10. Training recipe. The paper defers the learning rate, weight decay and warmup to the DeiT paper §4.1. The repository has no training code, so these values cannot be checked here.

The repository was taken from the link in the abstract, so no GitHub search was needed and no third-party reimplementation was used. Apart from the items above, I found no discrepancies between the paper's text and the code.

Pitfalls

• Patch size for PiT-B. Text and code disagree (see Paper vs code). Table 1 matches the code. A model built from the text alone gets the wrong grid for B.
• Attention normalization axis. The paper defines attention values with "sum over i of alpha_ij = 1" §3.3. The usual softmax normalizes each query row over keys. The paper's axis is the transpose. Which axis the softmax runs over is not stated in the text, so check both conventions against the entropy values in fig. 5.
• Matrix sizes. The paper writes A as M x N but then sums over N tokens with 1/N eq. 1 eq. 2. The dimension names are not tied to the query and key axes. The log base in the entropy is not stated.
• Attention distance uses relative coordinates. Positions are (x/W, y/H), and the distance is an L1 norm between them §3.3 eq. 2. A pixel-based distance would give different numbers.
• Pooling convolution is not fully specified. Kernel size, padding, bias and groups are not stated. The code uses kernel 3, padding 1, depth-wise groups and a bias pit.py.
• Hidden cost at 384 x 384. PiT-B throughput falls to 82 img/s against 91 for ViT-B at 384 x 384 tab. 4. The paper says the design targets 224 x 224 §4.1. Do not assume the saving carries to larger inputs.
• Gains depend on the recipe. At the vanilla setting PiT-S leads ViT-S by 4.6 points (73.3 vs 68.7). With full DeiT plus distillation the lead is 0.7 points (81.9 vs 81.2) tab. 2. Compare on the same recipe.
• Small models trail CNNs. The paper says CNNs at the Ti scale beat ViT-Ti and PiT-Ti §4.1 tab. 3.
• Detection does not beat ResNet50. PiT-S detector is 2.1 AP below ResNet50 but has lower latency tab. 5 §4.2.
• Resolution for detection. The paper halves the image resolution for all backbones and does not state the exact size §4.2. Results are not directly comparable to standard COCO numbers.
• Mixed measurement setups. Table 3 mixes CNN numbers reported by other papers with throughput the authors re-measured §4.1. Accuracy for CNNs uses their best published setting, while PiT and ViT use the DeiT-style recipe.
• Single numbers. Results are single runs as far as the paper reports. No variance is given.
• Architecture search. PiT was not chosen by neural architecture search, and the authors say it could improve with one §3.4.
• Figure-number references. Figure and table numbers in this card follow the order of the LaTeX source, not a compiled PDF. Fig. 1 is the schematic in the introduction. Fig. 2 is the ResNet experiment in related work. Fig. 3 is the ViT experiment in section 3, fig. 4 the pooling layer, fig. 5 the attention plots.

Lineage

• Builds on: ViT, the transformer baseline it modifies §1; the DeiT training recipe and distillation method used for the main tables §4.1.
• Motivated by: CNN dimension conventions with stride-2 downsampling and pooling, which the paper reviews §2.
• Compared against: ResNet, MobileNet, EfficientNet, RegNet and ResNeSt on ImageNet tab. 3; Deformable DETR detection with ViT-S backbone tab. 5; ViT-S and ResNet50 on robustness tab. 6.
• Superseded by: not stated.