Hierarchical vision transformers

Overview

A plain Vision Transformer (ViT) keeps the same number of spatial tokens in every layer 2103.16302 §1. A hierarchical variant instead reduces the spatial grid in steps and increases the channel count, so later layers see a smaller map with more features per position, as in a CNN [arxiv:2103.16302 §2] 2103.16302 §3.2. The shared idea is a sequence of stages, each with its own spatial size, joined by a downsampling step. Whether the downsampling is a pooling layer, a strided convolution or a patch merge is a design choice that varies by variant. HiFormer, for example, uses patch merging in its Swin branch 2207.08518 §3.1.2.

Most later variants also change the attention itself, so that the cost falls with the grid size. They do this by pooling the keys and values [arxiv:2106.12011 §III-B.1], by attending from a sparse set of delegate tokens [arxiv:2205.03436 §3.2], or by splitting the heads between local windows and pooled global attention 2205.13213 §4.1. The sections below describe each one.

Variants

Pooling-based ViT (PiT)

PiT splits the transformer blocks into three stages with three spatial scales 2103.16302 §3.2. Between stages, a pooling layer reshapes the spatial tokens into a map, applies a depth-wise convolution that halves the spatial size and widens the channels, and flattens the map back into tokens. The class token has no position, so it passes through a linear layer that changes its width instead 2103.16302 §3.2. Channel width per stage is set by the head count of that stage 2103.16302 §3.4 2103.16302 tab. 1. Card: Rethinking Spatial Dimensions of Vision Transformers.

Reported ImageNet-1k result at 224 x 224, same DeiT-style recipe: PiT-S reaches 80.9% top-1 with 2.9 B FLOPs, against 79.8% and 4.6 B FLOPs for ViT-S 2103.16302 tab. 2 2103.16302 tab. 3. Author-reported, single runs, no variance given.

Pyramid pooling attention (P2T)

P2T keeps global attention but shortens the keys and values. It average-pools the token map at four ratios, adds a 3 x 3 depth-wise convolution to each pooled map, flattens and concatenates the four maps, and uses the result as the keys and values 2106.12011 §III-B.1 2106.12011 eq. 2-4. Its four stages have feature strides 4, 8, 16 and 32 2106.12011 §III-A. The feed-forward block is an inverted bottleneck with a depth-wise convolution 2106.12011 §III-B.2. Card: P2T: Pyramid Pooling Transformer for Scene Understanding.

Reported results, all author-reported: at matched size P2T-Tiny reaches 79.8% top-1 with 11.6 M parameters and 1.8 GFLOPs, against 78.7% and 13.1 M / 2.1 GFLOPs for PVTv2-B1. P2T-Small reaches 82.4%, against 82.0% for PVTv2-B2 2106.12011 tab. 2. With Semantic FPN on ADE20K, P2T-Small reports 46.7 mIoU against 46.1 for PVTv2-B2 2106.12011 tab. 3. With RetinaNet on COCO, P2T-Small reports 44.4 AP against 43.8, and with Mask R-CNN box AP it is 45.5 against 45.3 2106.12011 tab. 4. The gains over PVTv2 are therefore between about 0.2 and 1.1 points in these tables. In the paper's 100-epoch ablation, adding the depth-wise convolution to the pooled maps and the inverted-bottleneck feed-forward raises top-1 from 74.7 to 79.5 2106.12011 tab. 9. Those ablation numbers are not comparable to the 300-epoch main results.

Local-global block per layer (EdgeViT)

Each EdgeViT layer is a local-global-local block 2205.03436 §3.2. First, depth-wise and point-wise convolutions mix each token with its neighbours. Then self-attention runs over one sampled delegate token per r x r window, so the attention cost depends on the number of windows. Last, a transposed convolution with kernel and stride r spreads each delegate's context back over its window. The sample rates are 4, 2, 2 and 1 for the four stages 2205.03436 §3.3. The paper contrasts this with PVT and Twins, where every token issues a query 2205.03436 §3.2. Card: EdgeViTs: Competing Light-weight CNNs on Mobile Devices with Vision Transformers.

The paper's target is on-device latency, measured on a phone CPU at FP32 with batch 1 2205.03436 §4.1. On that setup EdgeViT-XXS reports 74.4% top-1 at 32.8 ms, XS 77.5% at 54.1 ms and S 81.0% at 85.3 ms [arxiv:2205.03436 tab. 2]. MobileNet-v2 reports 72.0% at 33.3 ms, and EfficientNet-B0 77.1% at 52.1 ms. MobileNet-v3 0.75 is faster, at 23.0 ms, with 73.3% [arxiv:2205.03436 tab. 2]. So EdgeViT-XXS is more accurate than that CNN but slower on this phone. In the ablation at about 0.5 GFLOPs, removing local aggregation drops top-1 from 74.4 to 72.7, and removing local propagation drops it to 73.9 2205.03436 tab. 4.

Head-split local and global attention (LITv2, HiLo)

HiLo splits the heads of one attention layer into two groups [arxiv:2205.13213 §4.1] 2205.13213 eq. 6. The Lo-Fi heads, a fraction alpha of the total, average-pool keys and values over s x s windows, while their queries still come from every token. The Hi-Fi heads run standard attention inside non-overlapping s x s windows, without shifting. The authors set alpha to 0.9 by grid search on ImageNet-1K and use s = 2 for classification 2205.13213 §5.1. LITv2 places HiLo in stages 3 and 4 and removes the attention from stages 1 and 2 2205.13213 §4.3 2205.13213 app. A. It replaces relative position bias with a 3 x 3 depth-wise convolution inside each feed-forward block 2205.13213 §4.2. Card: Fast Vision Transformers with HiLo Attention.

Reported on an RTX 3090 with batch 64: LITv2-S reaches 82.0% top-1 at 1,471 images/s, against 81.3% at 961 images/s for Swin-Ti 2205.13213 tab. 1. On LITv2-S, plain multi-head self-attention reaches 82.3% at 1,293 images/s, so HiLo buys speed for a small accuracy loss 2205.13213 tab. 4. On ADE20K with Semantic FPN, HiLo reports 44.3 mIoU against 43.7 for plain attention 2205.13213 tab. 4. The paper's own GPU table reports LITv2-S at between 928 and 1,874 images/s across four GPUs, so absolute speed depends on the card 2205.13213 tab. 6.

Swin-style pyramid with a CNN front end (HiFormer)

HiFormer is a 2D segmentation model with two pyramids side by side. A CNN with three levels extracts local features, and a Swin Transformer with three levels builds a multi-scale view on top of them 2207.08518 §3.1.1 2207.08518 §3.1.2. Patch merging concatenates 2 x 2 neighbouring tokens, doubling the width and halving the resolution 2207.08518 §3.1.2. The largest and the smallest levels are then fused by a Double-Level Fusion (DLF) module. Each outer level gets a class token, and that token attends to the other level's patch tokens through cross-attention 2207.08518 §3.1.3. A decoder turns the fused levels into a full-resolution label map 2207.08518 §3.2.

Three sizes, S, B and L, differ in the CNN depth and the DLF settings 2207.08518 tab. 1. The reference B model reports 80.39 Dice and 14.70 HD on Synapse, against 79.13 Dice and 21.55 HD for Swin-Unet 2207.08518 tab. 2. Removing DLF lowers Dice to 77.15 [arxiv:2207.08518 tab. 6]. Card: HiFormer. All numbers are author-reported.

Two-stream Swin-B pyramid for salient object detection (SwinNet)

SwinNet is a salient-object detector that needs a second, spatially paired input: depth for RGB-D, or thermal for RGB-T 2204.05585 §3.1 2204.05585 §4.2. Each input goes through its own Swin-B encoder, which supplies the hierarchy 2204.05585 §3.2. At each level the two streams are aligned with a shared spatial map and then re-weighted by channel attention 2204.05585 §3.3. A separate edge branch built from the shallow depth features guides a top-down decoder 2204.05585 §3.4 [arxiv:2204.05585 §3.5]. Card: SwinNet: Swin Transformer drives edge-aware RGB-D and RGB-T salient object detection.

Reported RGB-D S-measure on NLPR is 0.941, against 0.931 for VST, and on NJU2K it is 0.935, against 0.922 2204.05585 tab. 1. In the backbone ablation on NLPR, Swin-B gives 0.941, ResNet-50 gives 0.925 and ResNet-50 with a ViT-16 branch gives 0.932 2204.05585 tab. 3. The model has 198.7 M parameters and about 124.3 GFLOPs, and runs at about 10 FPS with I/O included 2204.05585 §4.4.5 2204.05585 tab. 7. Its fusion and edge modules are isolated in ablations 2204.05585 tab. 4 [arxiv:2204.05585 tab. 5], so the hierarchy is an input to the method rather than the method itself.

When to use which

Implementation sketch

The common core is a loop over stages: run the transformer blocks at one resolution, then downsample the spatial tokens and widen the channels. Where variants differ is the downsampling step and how the class token is treated. P2T and EdgeViT leave this loop unchanged and replace the attention inside each block with a pooled or sampled version. HiLo keeps the loop and splits the heads of each attention layer. HiFormer does not follow the loop exactly. It runs a CNN pyramid and a Swin pyramid, and its fusion step replaces the final downsample 2207.08518 §3.1.3. Its own pseudocode is on its card.

# Hierarchical ViT forward pass, PiT-style [arxiv:2103.16302 §3.2]
# Input:  images (B, 3, 224, 224)
# Output: class logits (B, num_classes)

def hierarchical_vit(images):
    # Patch embedding gives a 27 x 27 grid in the paper's text [arxiv:2103.16302 §3.4]
    feature_map = patch_embed(images)
    # One learned class token, carried alongside the spatial tokens
    class_tokens = cls_token.expand(batch_size)
    for stage in range(num_stages):
        # Transformer blocks at this resolution [arxiv:2103.16302 tab. 1]
        feature_map, class_tokens = blocks[stage](feature_map, class_tokens)
        if stage < num_stages - 1:
            # Halve the grid and widen the channels between stages [arxiv:2103.16302 §3.2]
            feature_map, class_tokens = downsample[stage](feature_map, class_tokens)
    # Classify from the first class token [arxiv:2103.16302 §3.2]
    return head(layer_norm(class_tokens)[:, 0])

The downsample step is the variant-specific part. For PiT it is a depth-wise convolution on the map and a linear layer on the class token. The paper does not state the convolution's kernel size or padding 2103.16302 §3.2. For P2T the variant-specific part is the pooling pyramid inside attention, and the paper does not state how pooled sizes are rounded when they do not divide evenly 2106.12011 §III-B.1.

Contradictions

• PiT-B patch size, within one paper. The text gives patch 16 and stride 8 for all PiT sizes 2103.16302 §3.4. Table 1 lists a 31 x 31 map for PiT-B, which those values do not produce at 224 input [arxiv:2103.16302 tab. 1]. The reference code uses different values for B. Follow the table or the code, and check the choice before building B.
• Robustness table label, within one paper. The caption of the robustness table names PiT-B, but the PiT row is labelled PiT-S, and the text matches PiT-S 2103.16302 tab. 6.
• Where DLF output goes, within one paper. Eq. 3 keeps the patch tokens of both levels unchanged, so the cross-level information is carried only in the class tokens 2207.08518 eq. 3. The decoder is said to consume the two patch grids 2207.08518 §3.2. The paper does not say how the class-token output reaches the decoder.
• PVTv2-B1 detection baseline, across two papers. P2T reports PVTv2-B1 at 40.2 AP with RetinaNet on COCO 2106.12011 tab. 4. EdgeViT reports 41.2 for the same backbone and the same parameter count of 23.8 M 2205.03436 tab. 5. The two cards do not explain the gap, and the PVTv2 paper was not checked. Do not cite either number as the PVTv2-B1 baseline without that check.
• PVTv2-B1 segmentation baseline, across two papers. P2T reports 41.5 mIoU on ADE20K at 33.1 GFLOPs 2106.12011 tab. 3. EdgeViT reports 42.5 mIoU at 34.2 GFLOPs for the same backbone 2205.03436 tab. 6. Both use a 512 x 512 input and a Semantic FPN head. The same caution applies.
• SwinNet, claim against table. The text says the RGB-T results are best on every metric 2204.05585 §4.3.2. Table 2 shows CGFNet with a higher F-measure on VT1000, 0.906 against 0.896 for SwinNet [arxiv:2204.05585 tab. 2]. Use the table when ranking.

No other cross-paper disagreement is known. The variants report on different tasks, ImageNet classification, COCO detection, ADE20K segmentation, medical segmentation and salient object detection, so their numbers are not comparable across tasks.

Pitfalls

• Gains depend on the recipe. At the vanilla setting PiT-S leads ViT-S by 4.6 points. With full DeiT plus distillation the lead is 0.7 points 2103.16302 tab. 2.
• Hidden cost at higher resolution. PiT-B throughput falls to 82 img/s against 91 for ViT-B at 384 x 384, so the saving does not carry over to larger inputs 2103.16302 tab. 4.
• Under-specified downsampling. Kernel size, padding and bias of the pooling convolution are not in the text 2103.16302 §3.2. Write the choice down explicitly when you implement it.
• Detection comparisons. The paper halves the input resolution for all backbones and does not state the exact size 2103.16302 §4.2. Its numbers are not directly comparable to standard COCO results.
• Architecture search. PiT was not found by neural architecture search, and the authors say it could improve with one 2103.16302 §3.4.
• Head count matters in the fusion branch. HiFormer's ablation with 12 heads on the small level scores lower Dice than its default of 6 heads on both levels 2207.08518 tab. 8. The reported DLF gain is +3.24 Dice 2207.08518 §5 2207.08518 tab. 6.
• Two HiFormer rows match exactly. The ResNet18 row of the backbone ablation and the no-DLF row of the DLF ablation both report 77.15 Dice and 16.88 HD 2207.08518 tab. 6 2207.08518 tab. 7. The paper does not explain this, so check it before citing either row.
• Encoders dominate the cost. HiFormer's cross-attention is linear in the token count, because only class tokens query. The level encoders before it are global self-attention over every token of a level [arxiv:2207.08518 §3.1.3] 2207.08518 tab. 1.
• P2T text and table disagree on a detection gain. The text claims P2T-Base leads Swin-S by 1.0 AP with RetinaNet 2106.12011 §IV-C. The table gives 46.1 against 44.5, a gap of 1.6 2106.12011 tab. 4.
• P2T ablations are not the main recipe. Its ablations train for 100 epochs and the baseline row has no pooling-convolution, no feed-forward convolution and no overlapping patch embedding 2106.12011 tab. 9. Use the 300-epoch table rows for comparison with other models.
• Pooling type matters. In P2T's ablation, average pooling gives 74.7 top-1 against 73.0 for max pooling and 73.8 for depth-wise convolution as the pooling operation 2106.12011 tab. 7.
• Energy metric favours small models. EdgeViT's efficiency value is top-1 divided by energy, which the paper notes favours less accurate models. That is why it compares only Pareto-optimal points 2205.03436 §4.1 2205.03436 tab. 3.
• Latency and energy come from different phones. Latency is measured on a Galaxy S21 and energy on a Snapdragon HDK 2205.03436 §4.1. Do not derive energy from latency and power.
• Figure text and table disagree. The EdgeViT Figure 1 caption gives 33.3 ms for the lightest model, while its table gives 32.8 ms for EdgeViT-XXS and 33.3 ms for MobileNet-v2 2205.03436 fig. 1 2205.03436 tab. 2.
• EdgeViT detection gain in the text. The text gives +1.2 AP^m for EdgeViT-S over PVTv2-B1 on Mask R-CNN 2205.03436 §4.2.1. The table gives 41.0 against 38.8, a gap of 2.2 2205.03436 tab. 5.
• EdgeViT ablation rescales the model. Ablation variants are rescaled to about 0.5 GFLOPs by changing layer counts, so they are not the released models 2205.03436 §4.1.1 2205.03436 app. A.2.
• Plain attention is the accuracy reference for HiLo. It scores 0.3 top-1 more on ImageNet, and the paper's gain is speed and FLOPs 2205.13213 tab. 4.
• Window size and resolution in HiLo. The paper says a window of 4 speeds up high-resolution inputs with a small accuracy loss 2205.13213 §5.2 2205.13213 app. C. In its detection table, LITv2-S with window 4 drops RetinaNet AP from 44.0 to 43.7 2205.13213 tab. 2.
• Lo-Fi head count is rounded down. Alpha times the head count is not an integer in general. The paper does not say how to round it [arxiv:2205.13213 §4.1].
• HiLo recipe is partly missing. The paper gives warm-up, clipping, label smoothing and mixup only as "the same as in LITv1" 2205.13213 §5.1.
• SwinNet feature levels. The paper names four levels but does not say which Swin outputs they are 2204.05585 §3.2. An implementation that takes the four stage outputs may not match the released model, so check this choice against the card.
• SwinNet checkpoint choice. The paper does not say how the reported checkpoint was picked 2204.05585 §4.2. Ask which validation split was used before you trust a reported number.
• SwinNet speed. About 10 FPS is reported with I/O and preprocessing included, on hardware not stated for that figure 2204.05585 §4.4.5. Training took about 26 hours on one RTX 2080 Ti 2204.05585 §4.2.
• Hardware-specific speed. Throughput depends on the GPU or CPU. The throughput rankings in each paper hold only for the setup it names 2205.13213 §5.1 2205.03436 §4.1 2106.12011 tab. 2.
• Input size. P2T and EdgeViT classification use 224 x 224 inputs 2106.12011 §IV-A 2205.03436 §4.1. Check the input size of each detection and segmentation row before comparing it with another.

Papers

• Rethinking Spatial Dimensions of Vision Transformers: pooling between transformer stages, shrinking the token grid while widening channels; reports gains over ViT at equal FLOPs on ImageNet, detection and robustness.
• P2T: Pyramid Pooling Transformer for Scene Understanding: multi-scale average pooling of keys and values, with depth-wise convolutions on each pooled map; reports ImageNet, ADE20K and COCO results against PVTv2 and Swin.
• EdgeViTs: Competing Light-weight CNNs on Mobile Devices with Vision Transformers: a local-global-local block with delegate-token attention; reports phone-CPU latency and energy against light CNNs and transformers.
• Fast Vision Transformers with HiLo Attention: splits heads into high-frequency window attention and low-frequency pooled attention inside a four-stage backbone; reports ImageNet, COCO and ADE20K results against Swin and other efficient attention.
• HiFormer: Hierarchical Multi-scale Representations Using Transformers for Medical Image Segmentation: a CNN pyramid and a Swin pyramid fused by class-token cross-attention; reports Synapse, skin-lesion and SegPC results, with an official repository that differs from the paper in several places.
• SwinNet: Swin Transformer drives edge-aware RGB-D and RGB-T salient object detection: two Swin-B encoders, per-level alignment and channel re-calibration, and an edge-guided decoder; reports RGB-D and RGB-T results against other salient object detectors.