ResT: An Efficient Transformer for Visual Recognition

In plain words

Image models have long been built from convolutions. A convolution looks at a small neighbourhood, is cheap, and works well on pictures. A Transformer can relate any two parts of an image, which helps with wider context, but its attention gets expensive on large images, and the standard version treats the image as a flat list of fixed-size patches.

ResT is a Transformer laid out like a convolutional network. A few small convolutions look at the raw image first. Then four stages follow, each at a lower resolution and with more channels, so the network produces a pyramid of feature maps that detectors and segmenters can use.

Two changes make the attention cheaper and more useful. Keys and values are computed from a shrunken version of the image, made by a strided convolution. Position is handled by a learned gate that looks at each point's neighbours, so the network does not need a fixed input size. The authors report better accuracy than other backbones of similar size on image classification and on COCO detection.

TL;DR

ResT is a multi-scale Transformer backbone for image classification and dense prediction. Its efficient attention compresses keys and values with a strided depth-wise convolution, and a 1x1 convolution with instance normalisation lets the heads interact while keeping them different. Its positional encoding is a depth-wise-conv gate, so it accepts any input size. Author-reported: 79.6% ImageNet-1k top-1 for ResT-Small at 224 px, and 40.3 box AP for RetinaNet on COCO val2017. Implement from the code and this card, not from the paper alone (see Paper vs code).

Problem

The paper names four weaknesses of plain Transformer backbones §1:

1. Low-level structure such as edges and corners is hard to extract, because the model cuts raw images straight into patches §1, item 1 §2.3.
2. Memory and compute of multi-head self-attention grow quadratically with the number of tokens and with the channel count §1, item 2 §2.2.
3. Each head sees only a slice of the channels, which is weak when the per-head dimension is small §1, item 3 §2.2.
4. Tokens and positional encodings have one fixed scale, which does not suit dense prediction (detection, segmentation) §1, item 4 §2.4.

Method

Overall structure

ResT follows the ResNet layout §2 fig. 2:

• Stem: three 3x3 convolutions with strides 2, 1 and 2, reducing height and width by 4 §2.3 tab. 1. Batch norm and ReLU follow the first two; a positional gate (PA, below) follows the last tab. 1.
• Four stages. Each stage has a patch embedding (stages 2 to 4 only; the stem serves stage 1), a positional encoding inside that embedding, and a stack of L efficient Transformer blocks §2.
• Head: global average pooling over the last stage, then a linear classifier §2.5 tab. 1.

Stage widths are C, 2C, 4C and 8C, with C = 64 for Lite and Small and 96 for Base and Large tab. 1. Spatial sizes for a 224 px input are 56, 28, 14 and 7 tab. 1.

Efficient Transformer block

Each block is pre-norm with residual connections §2.2 eq. 5. A standard Transformer block has the same shape §2.1 eq. 1. The attention sub-layer is replaced by EMSA, described below §2.2 fig. 1.

y = x' + \mathrm{FFN}(\mathrm{LN}(x')), \quad x' = x + \mathrm{EMSA}(\mathrm{LN}(x)) eq. 5

The feed-forward network has two linear layers with GELU between them and a hidden width of 4C §2.1 eq. 3.

EMSA (efficient multi-head self-attention)

EMSA changes three things relative to multi-head self-attention §2.2 fig. 3:

1. Queries are computed from all N tokens.
2. The token grid is reduced by a depth-wise convolution with kernel s+1, stride s and padding s/2. The reduction s is 8, 4, 2 and 1 for stages 1 to 4 §2.2 tab. 1. The reduced map has N' = (h/s)(w/s) tokens.
3. Keys and values are computed from the reduced tokens. Their attention scores are mixed across heads by a 1x1 convolution and normalised with instance normalisation §2.2 eq. 4. The 1x1 convolution lets every head depend on every key; instance normalisation restores head diversity §2.2 app. A.2.

\mathrm{EMSA}(Q,K,V) = \mathrm{IN}\big(\mathrm{Softmax}\big(\mathrm{Conv}(QK^\top/\sqrt{d_k})\big)\big)\,V eq. 4

The heads are concatenated and projected as in standard attention §2.2. The paper's cost formula still has an n² term, divided by s² §2.2.

# Efficient Transformer block  [§2.2] [eq. 5]
# tokens: (B, N, C) with N = height * width
def efficient_block(tokens, height, width, reduction):
    # Attention sub-layer with a residual connection  [eq. 5]
    tokens = tokens + emsa(
        layer_norm(tokens), height, width, reduction)
    # Feed-forward sub-layer with a residual connection  [eq. 5] [eq. 3]
    tokens = tokens + ffn(layer_norm(tokens))
    return tokens

# Feed-forward network: GELU between two linear layers  [eq. 3] [§2.1]
def ffn(tokens):
    # Expand to 4C channels  [§2.1]
    hidden = gelu(linear_expand(tokens))          # (B, N, 4C)
    # Project back to C channels  [eq. 3]
    return linear_reduce(hidden)                  # (B, N, C)

# EMSA: efficient multi-head self-attention  [§2.2] [eq. 4] [fig. 3]
# tokens: (B, N, C); reduction is s for this stage  [tab. 1]
def emsa(tokens, height, width, reduction):
    # Queries use every token  [§2.2 step 1]
    queries = split_heads(linear_q(tokens))       # (B, k, N, d_k)
    heads = queries.shape[1]                      # k
    # Shrink the token grid with a strided depth-wise conv  [§2.2 step 2]
    grid = reshape_to_map(tokens, height, width)  # (B, C, H, W)
    reduced = depthwise_conv(
        grid, kernel=reduction + 1, stride=reduction,
        padding=reduction // 2)                   # (B, C, H/s, W/s)
    # LayerNorm after the reduction: code only, not in paper
    # [code: models/rest.py#L78] [code: models/rest.py#L88-L91]
    reduced = layer_norm(flatten_to_tokens(reduced))  # (B, N', C)
    # Keys and values come from the reduced tokens  [§2.2 step 3]
    keys = split_heads(linear_k(reduced))         # (B, k, N', d_k)
    values = split_heads(linear_v(reduced))       # (B, k, N', d_k)
    # Scaled dot-product scores  [§2.2] [eq. 4]
    scores = queries @ transpose(keys) / sqrt(d_k)  # (B, k, N, N')
    if heads > 1:
        # 1x1 conv mixes the heads; skipped for a single head  [eq. 4]
        # [code: models/rest.py#L98-L101]
        scores = conv1x1_across_heads(scores)     # (B, k, N, N')
    # Softmax over the N' keys  [eq. 4]
    weights = softmax(scores, axis=-1)            # (B, k, N, N')
    if heads > 1:
        # Instance norm per head restores diversity  [eq. 4] [§2.2]
        # [code: models/rest.py#L101]
        weights = instance_norm(weights)          # (B, k, N, N')
    # Weighted sum of values, heads merged  [§2.2 step 5]
    output = merge_heads(weights @ values)        # (B, N, C)
    # Output projection  [§2.2 step 5]
    return linear_out(output)                     # (B, N, C)

Patch embedding and stem

The stage-entry patch embedding is an overlapping 3x3 convolution with stride 2 and padding 1. It doubles the channels and halves height and width, with the positional gate applied afterwards §2.3 §2.4 fig. 4.

# Patch embedding for stages 2-4  [§2.3] [tab. 1] [§2.4]
def patch_embed(grid, out_channels):
    # Overlapping 3x3 conv: stride 2, padding 1  [§2.3]
    grid = conv3x3(grid, out_channels, stride=2, padding=1)
    # Batch norm after the conv: code only
    # [code: models/rest.py#L158] [code: models/rest.py#L166-L167]
    grid = batch_norm(grid)                       # (B, 2C, H/2, W/2)
    # Positional gate applied to the embedded grid  [§2.4] [eq. 8]
    return pixel_attention(grid)

# Stem: three 3x3 convs with strides 2, 1, 2  [§2.3] [tab. 1]
def stem(image):
    # Hidden width C/2; batch norm and ReLU after the first two convs  [tab. 1]
    hidden = widths[0] // 2
    grid = relu(batch_norm(conv3x3(image, hidden, stride=2)))
    grid = relu(batch_norm(conv3x3(grid, hidden, stride=1)))
    # Final conv to C channels, stride 2, then positional gate  [tab. 1] [§2.4]
    grid = conv3x3(grid, widths[0], stride=2)
    return pixel_attention(grid)

Positional encoding as pixel attention (PA)

Positional information is a per-channel gate: each value is multiplied by a sigmoid of a 3x3 depth-wise convolution of the same grid §2.4 eq. 8. The paper compares this with two other choices: a learned absolute table (LE, eq. 6) and a per-channel linear map (GL, eq. 7; appendix definition in eqs. 10 and 11) §2.4 app. A.4 tab. 8.

\hat{x} = \mathrm{PA}(x) = x \cdot \sigma\big(\mathrm{DWConv}(x)\big) eq. 8

# Positional encoding as pixel attention (PA)  [§2.4] [eq. 8]
def pixel_attention(grid):
    # Per-channel gate from a 3x3 depth-wise conv with padding 1  [eq. 8]
    gate = sigmoid(depthwise_conv3x3(grid))       # (B, C, H, W)
    # Scale each position by its gate  [eq. 8]
    return grid * gate                            # (B, C, H, W)

Whole network

# Whole network  [§2] [tab. 1]
# widths = [C, 2C, 4C, 8C]; depths per stage [tab. 1]
# reductions = [8, 4, 2, 1] per stage  [§2.2] [tab. 1]
def rest_forward(image):
    # Stem: low-level features at 1/4 resolution  [§2.3]
    grid = stem(image)                            # (B, C, H/4, W/4)
    for stage_index in range(4):
        # Stages 2-4 begin with a patch embedding  [§2.3]
        if stage_index > 0:
            grid = patch_embed(grid, widths[stage_index])
        height, width = grid.shape[2:]
        tokens = flatten_to_tokens(grid)          # (B, H*W, C_i)
        for block_index in range(depths[stage_index]):
            # Efficient blocks with this stage's reduction  [§2.2] [tab. 1]
            tokens = efficient_block(
                tokens, height, width, reductions[stage_index])
        grid = reshape_to_map(tokens, height, width)
    # Final LayerNorm before pooling: code only  [code: models/rest.py#L318]
    tokens = layer_norm(tokens)
    # Global average pooling, then the linear classifier  [§2.5]
    pooled = mean_over_tokens(tokens)             # (B, 8C)
    return linear_classifier(pooled)              # (B, num_classes)

Figures

The paper's figures are described here and cited, not reproduced. fig. 1 compares a ResNet bottleneck, a standard Transformer block and the efficient block; the efficient block only replaces MSA with EMSA. fig. 2 shows the pipeline: stem, four stages and a head. fig. 3 shows the EMSA module. fig. 4 shows patch embedding and positional encoding in one stage. fig. 5 (appendix) visualises attention maps before and after the 1x1 conv and instance norm. fig. 6 (appendix) shows Group-CAM heatmaps for ResT-Lite and ResNet-50.

Hyperparameters

Architecture (ImageNet-1k) tab. 1

Notation tab. 1: EMSA_n_r = n heads, reduction r. MLP_c = feed-forward network with hidden width 4c and output width c. Stage 1 has one head in every model tab. 1 rest.py. The paper benchmarks Lite and Small against ResNet-18, Base against ResNet-50 and Large against ResNet-101 §2.5.

Image classification, ImageNet-1k §3.1

Object detection and instance segmentation, COCO §3.2

Ablation recipe, ResT-Lite §3.3

Results

Author-reported. Setup: ImageNet-1k, 224x224 input, throughput on one V100 GPU tab. 2. The paper's text compares ResT-Small with PVT-T and ResT-Base and -Large with Swin-T and Swin-S §3.1.

ImageNet-1k tab. 2

The paper's improvement markers (blue, relative to ResNets) are omitted.

COCO object detection, RetinaNet tab. 3

COCO detection and instance segmentation, Mask R-CNN tab. 4

Ablations, ResT-Lite, ImageNet-1k (Top-1 / Top-5 %)

Simple recipe §3.3; these numbers are not comparable with the table above.

Other comparisons, from the appendix app. A.3:

Paper vs code

Code paths are relative to the ResT repository at code_commit. The repo has the v1 classification model in models/rest.py, the v1 detection backbone in object_detection/restv2/rest.py (the file also holds v2 code) and v1 detection configs in object_detection/configs/ResTv1/. The repository's models/rest_v2.py and object_detection/restv2/restv2.py are ResT V2 and are not covered here.

• Single-head stage 1 skips the 1x1 conv and instance norm. The paper's Eq. 4 applies both in every stage §2.2 eq. 4, and Table 1 gives stage 1 one head tab. 1. The code sets apply_transform = apply_transform and num_heads > 1 rest.py, so with one head neither operation runs rest.py.
• LayerNorm after the spatial reduction. Not in the paper. Code applies LayerNorm to the reduced tokens before computing keys and values rest.py rest.py.
• Batch norm in stage patch embeddings. The paper describes the stage down-sampling as a plain 3x3 stride-2 convolution §2.3. The code adds batch norm after it rest.py rest.py.
• Final LayerNorm before pooling. Not stated in the paper. Code applies self.norm to the stage-4 tokens before average pooling rest.py rest.py rest.py.
• Instance norm and 1x1 conv defaults. The paper does not say whether instance norm has learnable affine parameters or whether the 1x1 conv has a bias. Code uses PyTorch defaults: InstanceNorm2d(num_heads) with no affine parameters and no running statistics, and Conv2d with bias rest.py.
• Biases on Q, K and V. The paper does not state them. All four classification factories pass qkv_bias=True rest.py rest.py rest.py rest.py.
• Stem convolutions have no bias rest.py rest.py rest.py. The paper does not state this.
• Stochastic depth is not set in the classification factories. ResT defaults drop_path_rate to 0 rest.py, and the four factories do not pass a value rest.py. The paper's 0.1, 0.1, 0.2 and 0.3 values §3.1 must be passed in training.
• No ResT v1 training command. TRAINING.md documents ResT V2 settings TRAINING.md. The defaults in main.py differ from the paper's recipe: learning rate 4e-3 main.py, 20 warm-up epochs main.py, no gradient clipping main.py, drop path 0 main.py, label smoothing 0.1 main.py, and a crop_pct default of 1.0 main.py. Label smoothing is not in the paper.
• Evaluation crop. The paper resizes the shorter side to 256 and crops 224 §3.1. datasets.py uses 224/256 only when crop_pct is None datasets.py; the main.py default is 1.0, which resizes to 224 with no 256 margin unless --crop_pct is passed main.py. The model's own _cfg sets crop_pct to 0.9 rest.py.
• Positional gate (PA) and gate-only encoding. Matches the paper's Eq. 8: x times sigmoid of a 3x3 depth-wise conv rest.py.
• GL is defined but unused. The paper defines GL as a per-channel linear map app. A.4 eq. 10 eq. 11. The code's GL is a residual 3x3 depth-wise conv, and no model in the repo calls it rest.py. The LE and GL variants of the ablation tab. 8 are not in the repo.
• Detection: training scale. The paper states 480 to 800 §3.2. The configs set 640 to 800 Base-RCNN-FPN.yaml Base-RetinaNet.yaml. The max side is not set in these configs.
• Detection: learning-rate steps. Both base configs decay the rate at 60k and 80k iterations of 90k total Base-RCNN-FPN.yaml Base-RetinaNet.yaml. Batch 16 for 90k iterations is about 12 COCO epochs (derived). The paper states only "1x (12 epochs)" §3.2.
• Detection: gradient clipping. The Mask R-CNN config clips the full model's gradient norm at 0.01 Base-RCNN-FPN.yaml. The RetinaNet config does not clip. The paper says nothing about clipping for detection.
• Detection: crop. The Mask R-CNN base config crops to 0.9 relative size Base-RCNN-FPN.yaml. The paper says nothing about cropping.
• Detection: stochastic depth. The detection backbone uses 0.1 (Small), 0.2 (Base) and 0.2 (Large) rest.py. The paper gives stochastic depth only for classification §3.1.
• Detection: feature norms. Code adds a LayerNorm on each output stage before the FPN, as in the paper §3.2. Stage 4 always passes through the final norm rest.py rest.py.
• Detection: stage outputs. Mask R-CNN takes stages 1 to 4, RetinaNet takes stages 2 to 4 mask_rcnn_rest_small_FPN_1x.yaml retinanet_rest_small_FPN_1x.yaml. The paper does not state the stage choice.
• Positional tables. The code computes height and width from the input, with no learned table tied to input size rest.py, which supports the paper's arbitrary-size claim.

Pitfalls

• Stage 1 has one head, so Eq. 4's conv and instance norm do not run in the released code. Implement the one-head case as the code does rest.py, and note that the paper does not say so.
• The paper says patch embedding down-samples by "4x" §2.3, but the operation is a stride-2 convolution, and the example there changes h/4 to h/8, which is 2x per side. Implement stride 2 per stage.
• Attention memory per head is N x N'. With N = H x W queries and N' keys after reduction, stage 1 at 224 px has 56 x 56 = 3136 queries against 49 keys per head (derived from tab. 1 and the reduction formula §2.2). The paper's cost formula still has an n² term, divided by s² §2.2.
• Key count is 49 at 224 px in every stage. Reduction gives 56/8, 28/4, 14/2 and 7/1 = 7 per side (derived).
• Instance norm has no learnable affine and no running statistics in the code rest.py. Training and evaluation use the same statistics, per image and per head, over the N x N' map.
• Evaluation crop. Passing no --crop_pct gives a 224 resize and a 224 crop main.py datasets.py. Use 0.875 (224/256) to match the paper §3.1.
• Two ResNet-18 numbers. Table 2 gives 69.7% tab. 2; the appendix text gives 69.76% app. A.4.
• Throughput comparisons use one V100 GPU tab. 2. The MSA comparison used batch 512 because each V100 GPU handles only 64 images at a time for MSA, so it is not directly comparable with the EMSA batch app. A.3.
• Ablations use a different, simpler recipe (SGD, 100 epochs, flip and crop only) §3.3 app. A.4. Do not compare them with Table 2.
• Training hyperparameters are incomplete in the paper. Warm-up shape, final learning rate, Mixup and CutMix strengths, and label smoothing are not stated. The repo's main.py defaults are not the paper's recipe (see Paper vs code).
• Position encoding claims more than the ablation shows. The paper argues that PA encodes position and allows any input size §2.4. The ablation compares only ResT-Lite at 224 px tab. 8.
• Attribution to the repo. The code is Apache-2.0 LICENSE. This card quotes no code verbatim.

Lineage

• Builds on: ResNet, for the stem and stage layout §1 §2; ViT, for patch tokenisation, which the paper departs from §2.3; and the DeiT training recipe, which it follows for classification §3.1.
• Compared against: PVT, Swin, DeiT, RegNet and ResNet on ImageNet-1k tab. 2; RetinaNet and Mask R-CNN backbones on COCO tab. 3 tab. 4.
• Follow-up: ResT V2 (arXiv:2204.07366) is linked from the repository README, which says it removes the multi-head interaction part of EMSA README.md. It has no card yet.