RTFormer: Efficient Design for Real-Time Semantic Segmentation with Transformer

In plain words

Real-time segmentation labels every pixel of a video frame fast enough for a car or a phone. The fastest systems are convolutional networks. Transformers see far across an image, which helps, but their usual attention is slow on the graphics chips that run these systems.

RTFormer keeps two copies of the picture in its middle layers: a detailed high-resolution copy and a coarse, low-resolution copy. The coarse copy gets a cheap attention step that is friendly to graphics chips. The detailed copy then looks up the big-picture summary from the coarse copy, so each detailed pixel knows about the whole scene without a costly full attention. The early layers are plain convolutions. The authors report that this gives a better speed-accuracy trade-off than earlier real-time models on the benchmarks they test.

TL;DR

• RTFormer is a real-time semantic segmentation network with two resolution branches. Its building block combines a linear-cost attention on the low-resolution branch with cross-resolution attention that feeds global context to the high-resolution branch §3.1.
• Early stages are convolutional. The last two stages use the RTFormer block §3.2.
• Author-reported: RTFormer-Base reaches 79.3% val mIoU on Cityscapes at 39.1 FPS, and RTFormer-Slim 76.3% at 110.0 FPS, on one RTX 2080Ti tab. 2.
• The released PaddleSeg code differs from the paper in several training details, and the cross-attention scale differs from the paper's equation rtformer.py. See Paper vs code.

Problem

Transformers capture long-range context, which semantic segmentation needs §1. But the quadratic cost of standard self-attention, and the multi-head split that breaks a large matrix product into many small ones, both slow inference on GPU-like hardware §1. Earlier segmentation transformers reduce keys and values, or work only at high resolution, and still keep these costs §1 §2. The paper argues that, for real-time use on such devices, CNNs still dominate the state of the art §1.

Method

The network has two parts: a block that runs two resolution branches in a stepped layout §3.1, and a backbone built from that block with a segmentation head §3.2.

Stages and branches. The stem is two 3×3 convolutions §3.2. Stages 1 and 2 are residual blocks. From stage 3 on, the network runs a high-resolution branch at stride 8 next to a low-resolution branch at stride 16, then 32, and the two exchange features §3.2. The paper says stage 3 is still built from residual blocks, while stages 4 and 5 use the RTFormer block §3.2 fig. 4.

RTFormer block. The low-resolution branch runs GPU-Friendly Attention, then an FFN §3.1. The result is added, after compression and upsampling, into the high-resolution branch. The high-resolution branch then runs cross-resolution attention, whose keys and values come from the low-resolution branch, followed by its own FFN §3.1 fig. 2. The paper calls this a stepped layout: the low branch is computed first so that its output can feed the high branch §3.1. When the injection flag is on, the released code also feeds high-resolution features back into the low branch after each block; the paper does not describe this step (see Paper vs code). The pseudocode follows the code's order. The paper gives the stepped order only in outline rtformer.py.

# RTFormer block: one stage with two branches  [§3.1] [fig. 2]
# Input:  x_h (B, C_h, H, W) at stride 8; x_l (B, C_l, H/16, W/16)
# Output: x_h at stride 8; x_l at stride 16 or 32
# Order follows the code  [code: paddleseg/models/rtformer.py#L544-L577]

def rtformer_block(x_h, x_l):
    # 1. Low branch: GPU-Friendly Attention, with residual  [§3.1] [eq. 3]
    x_l = x_l + gfa(x_l)                       # (B, C_l, h, w)
    # 2. Low branch: FFN of two 3x3 convs, no width expansion  [§3.1]
    x_l = x_l + ffn(x_l)
    # 3. Stepped layout: add upsampled low branch into high branch  [§3.1]
    x_h = x_h + upsample(compress(x_l))        # 1x1 conv, then bilinear
    # 4. High branch: cross-resolution attention  [§3.1] [eq. 4]
    cross_keys, cross_values = cross_kv(x_l)   # pool to s x s, then 1x1 conv
    x_h = x_h + cross_attention(x_h, cross_keys, cross_values)
    # 5. High branch: FFN  [§3.1]
    x_h = x_h + ffn(x_h)
    # 6. Feed high-resolution features back to the low branch  [code only]
    x_l = x_l + inject(x_h)                    # two stride-2 3x3 convs
    return x_h, x_l

GPU-Friendly Attention (GFA). It starts from external attention, which applies a double normalization to the product of tokens and learned keys, then multiplies by learned values §3.1 eq. 1. Multi-head external attention splits its matrix products into groups; GFA removes that split §3.1 fig. 3. It keeps one large product, enlarges the set of learned keys and values, and replaces the second normalization by a grouped one that splits the keys into groups §3.1 eq. 3.

# GPU-Friendly Attention (low branch)  [§3.1] [eq. 3]
# Input: x (B, N, d), N = h*w tokens
# Learned keys and values: (M_g, d); the groups count is set per branch

def gfa(x, keys, values, groups):
    # 1. Similarity of every token with every learned key  [eq. 3]
    scores = x @ keys.T                         # (B, N, M_g)
    # 2. First normalization: softmax over the tokens  [§3.1]
    attn = softmax(scores, axis_tokens)         # (B, N, M_g)
    # 3. Second normalization: per token, sum over the keys of each group
    attn = attn / (group_sum(attn, groups) + 1e-6)   # eps: code only
    # 4. Weighted sum of the learned values  [eq. 3]
    return attn @ values                        # (B, N, d)

Cross-resolution attention (CA). The high-resolution tokens are the queries. The keys and values are computed from the low-resolution features, which are pooled to a fixed spatial size and projected §3.1 eq. 4. CA uses a single softmax over the pooled tokens and no multi-head split, because the keys and values are not learned parameters §3.1. The pooled size is the main knob: 8×8 for Slim and 12×12 for Base tab. 1.

# Cross-resolution attention (high branch)  [§3.1] [eq. 4]
# x_h: (B, N_h, C_h) queries from the high branch
# cross_keys, cross_values: (B, s*s, C_h) from the pooled low branch

def cross_attention(x_h, cross_keys, cross_values):
    # 1. Similarity of every high-resolution token with every pooled token
    scores = x_h @ cross_keys.transpose(0, 2, 1)   # (B, N_h, s*s)
    # 2. Scale. The paper divides by sqrt(d_h); the code divides by
    #    sqrt(s*s) instead. See Paper vs code.
    scores = scores / sqrt(C_h)
    # 3. Softmax over the pooled tokens; no multi-head split  [§3.1]
    attn = softmax(scores, axis_pooled_tokens)
    # 4. Weighted sum of the pooled values
    return attn @ cross_values                     # (B, N_h, C_h)

FFN. The FFN is two 3×3 convolutions with no expansion of the hidden width. A typical transformer FFN uses two MLP layers with a 2× or 4× expansion plus a depth-wise 3×3 convolution. The paper reports that its FFN is both faster and more accurate in its ablation §3.1 tab. 5.

Segmentation head. The low-resolution output goes through a DAPPM module, a pyramid-pooling module borrowed from DDRNet §3.2. Its output is fused with the high-resolution feature to give a stride-8 map. A classification head with a 3×3 convolution and a 1×1 convolution then predicts the labels, with the hidden width equal to the input width §3.2. The release also adds an auxiliary head on the stride-8 stage-3 feature, which the paper does not mention rtformer.py rtformer.py.

Variants. RTFormer-Slim and RTFormer-Base differ in the base width (32 vs 64 channels) and the pooled size, as in tab. 1 §3.2. The architecture description is given in §3.2; the released configs set the head widths and the injection flags per dataset (see Hyperparameters).

Key equations. External attention, multi-head external attention, GFA and CA, as printed eq. 1 eq. 2 eq. 3 eq. 4:

EA(X, K, V) = DN(X \cdot K^{T}) \cdot V eq. 1

MHEA(X) = Concat(h_1, \dots, h_H), \quad h_i = EA(X_i, K', V') eq. 2

GFA(X, K_g, V_g) = GDN(X \cdot K_g^{T}) \cdot V_g eq. 3

CA(X_h, K_c, V_c) = Softmax\left(\frac{X_h \cdot K_c^{T}}{\sqrt{d_h}}\right) \cdot V_c eq. 4

K_c, V_c = \phi(X_c), \quad X_c = \theta(X_l) eq. 4

Here X is an N×d input with N elements, K and V are learned M×d matrices, DN is double normalization, and H is the head count §3.1 eq. 1 eq. 2. K_g and V_g have M_g rows, with GDN as the grouped version of DN §3.1 eq. 3. φ is a set of splits, permutations and reshapes, θ is pooling and convolution, and d_h is the high-branch width §3.1 eq. 4. The paper states M_g = M×H in the main text §3.1 and M_g = d in the appendix app. B.1. The released code follows the second (see Paper vs code).

Hyperparameters

Paper defaults are in the first table; the released configuration values are cited to code. Where the two differ, both are listed and the conflict is in Paper vs code.

Architecture

Cityscapes training (paper §4.1; configs in configs/rtformer/)

CamVid training (paper only; no RTFormer CamVid config in the repository) §4.1

ADE20K training (paper §4.1; configs rtformer_*_ade20k_512x512_160k.yml)

COCOStuff training (paper only; no RTFormer config in the repository) §4.1

Ablation training (paper only) §4.4 app. B.1: RTFormer-Slim trained from scratch with learning rate 0.001, other settings as on ADE20K. No ablation configs in the repository.

ImageNet pretraining (paper, appendix only; the recipe is not in this repository) app. A tab. 6

Results

Author-reported. Setup: semantic segmentation on Cityscapes (val, single crop 1024×2048), CamVid (test, 720×960), ADE20K (val, 512×2048 for FLOPs and FPS), and COCOStuff (test, 640×640) tab. 2 tab. 3 tab. 4. FPS is measured on one RTX 2080Ti without TensorRT by default §4.1. Baselines are taken from their papers unless marked, with * for retrained models and † for FPS measured by the authors tab. 2 tab. 3.

Headline numbers

Table 2: Cityscapes and CamVid tab. 2. Columns: Cityscapes FPS and val mIoU; CamVid FPS and test mIoU. Blank cells are not reported. † FPS measured by the authors on one RTX 2080Ti. * retrained by the authors.

Table 3: ADE20K tab. 3. Params, FLOPs and FPS at 512×2048. † FPS on one RTX 2080Ti; * retrained by the authors.

Table 4: COCOStuff tab. 4. Params, FLOPs and FPS at 640×640.

Table 5: ablations on ADE20K, RTFormer-Slim tab. 5 §4.4. Each row gives FPS and mIoU (%), as the paper reports them.

Table 6: ImageNet pretraining settings app. A. See Hyperparameters.

Table 7: ImageNet classification, 224 px single crop app. A. Params and FLOPs at 224 px.

Table 8: attention types on ADE20K app. B.1. Appendix table with more variants. All on RTX 2080Ti. SA multi-head self-attention, EA multi-head external attention, LA Linformer, NA Nyströmformer, and the number is the landmark count.

Paper vs code

Official code: the paper links PaddleSeg (github.com/PaddlePaddle/PaddleSeg) [abstract], which is the PaddlePaddle organisation and the authors are from Baidu. The clone is at commit 3c4db66de1d9d59d0628ed87590b6308a2f4aa2a, Apache-2.0 LICENSE. The model is paddleseg/models/rtformer.py, and the configs are in configs/rtformer/.

Discrepancies found:

• Cross-attention scale. Paper: softmax of the scores divided by sqrt(d_h), the high-branch width eq. 4. Code: the scores are multiplied by the inverse square root of the number of pooled tokens, cross size squared, before the softmax rtformer.py rtformer.py. With the configured sizes, the code uses 1/12 for Base (12×12, 144 tokens) and for Slim on Cityscapes (default 12), and 1/8 for Slim on ADE20K (8×8, 64 tokens). The paper's d_h gives 1/sqrt(128) ≈ 1/11.3 for Base and 1/sqrt(64) = 1/8 for Slim, since the high-branch width at stage 5 is 128 for Base and 64 for Slim. The paper does not say which form its reported numbers used.
• Poly power, Cityscapes. Paper: power 0.9 §4.1. Code: power 1.0 in both Cityscapes configs rtformer_base_cityscapes_1024x512_120k.yml rtformer_slim_cityscapes_1024x512_120k.yml.
• Poly power and crop, ADE20K. Paper: other settings "identical to Cityscapes", which implies power 0.9 and a 512×1024 crop §4.1. Code: power 1.0 rtformer_base_ade20k_512x512_160k.yml and a 512×512 crop rtformer_base_ade20k_512x512_160k.yml.
• Cross size, Slim. Table 1 gives 8×8 for Slim tab. 1. The Slim Cityscapes config leaves it at the code default of 12 rtformer.py. Only the Slim ADE20K config sets 8 rtformer_slim_ade20k_512x512_160k.yml.
• Low-branch width M_g. Main text: M_g = M×H §3.1. Appendix: M_g = d app. B.1. Code: the learned key count equals the low-branch output width rtformer.py, which matches the appendix.
• Feedback from high to low branch. The paper does not describe it. The code injects high-resolution features into the low branch after each block, and the flag is on in stage 4 for all configs and in stage 5 only for Slim Cityscapes rtformer.py rtformer_slim_cityscapes_1024x512_120k.yml. This changes what stage 5 receives.
• Auxiliary head. Not in the paper. The code adds a head on the stride-8 stage-3 feature, weighted 0.4 in the loss, during training rtformer.py rtformer.py rtformer_base_cityscapes_1024x512_120k.yml.
• Drop path. Not stated for segmentation. The code default is 0.2 rtformer.py, and the ADE20K configs set 0.1 rtformer_base_ade20k_512x512_160k.yml.
• Cross-feature pooling. Paper: "pooling", type not stated §3.1. Code: adaptive max pooling rtformer.py.
• ADE20K Base result. Paper: 42.1 tab. 3. The repository README reports 42.02 for the same model [README: configs/rtformer/README.md]. Cityscapes: paper 79.3 and 76.3 against README 79.24 and 76.31, rounding only.
• Missing from the repository. No CamVid or COCOStuff RTFormer config, no ablation configs, and no ImageNet training code. Pretrained backbones are downloaded from URLs in the configs rtformer_base_cityscapes_1024x512_120k.yml.
• Stage-4 attention stride. The first projection of the stage-4 low-branch attention has stride 2, so that attention runs on the 1/32 grid rtformer.py. The paper does not say this.

Pitfalls

• Scale of the cross-attention scores. The paper's d_h and the code's pooled-token count give different scales (see Paper vs code). The paper reports no ablation of this scale, so either choice is untested by the authors' own experiments.
• Crop order. The paper writes 512×1024. The code writes [1024, 512]. The code's convention (width first) is inferred from the 2048×1024 validation resize written the same way; the transform was not checked. Confirm before training.
• Stage 4 and 5 timing. The low-branch attention in stage 4 runs on a grid downsampled by a stride-2 projection, so the attention and its shortcut work at stride 32 rtformer.py. Stage 5 runs at stride 32 without the projection rtformer.py.
• ReLU placement. The inputs to the stage-4 and stage-5 blocks pass through ReLU first, on both branches rtformer.py. The paper says nothing about ReLU placement.
• FFN details the paper omits. The released FFN is BatchNorm (epsilon 1e-6), then a 3×3 conv, GELU, dropout, and a second 3×3 conv rtformer.py. The paper specifies only two 3×3 convs with no expansion §3.1.
• Normalization epsilon. The second normalization in GFA adds 1e-6 to the denominator rtformer.py. The paper gives no epsilon.
• Pooling and max vs average. The cross-attention pooling is max in the code rtformer.py. The paper does not name the pooling type.
• DAPPM sizes. The paper borrows DAPPM from DDRNet without its pool sizes. The code uses DDRNet-style pool kernels of 5, 9 and 17 with strides 2, 4 and 8, plus a global pool rtformer.py.
• Groups in high-resolution branch. The paper's ablations vary groups in both branches tab. 5(c), but the main model uses CA on the high branch, where the group count does not apply rtformer.py. The reported group effect is for the GFA+GFA variant only.
• Cost of the high-resolution branch. The high-resolution branch stays at stride 8 through the network. The paper reports FPS and parameters, but no memory numbers.
• Speed numbers. FPS is measured on one RTX 2080Ti without TensorRT by default §4.1. Table 2 mixes GPUs, and its FPS values come from different sources, so the speed ranking is not a controlled comparison tab. 2.
• Single-scale results. Table 2 reports single-crop numbers, while the README also lists flip and multi-scale test results [README: configs/rtformer/README.md]. Do not compare those with single-scale numbers in the paper.
• Pretraining dependence. Results rely on ImageNet-pretrained backbones. The ImageNet recipe is only in the appendix app. A, and the released code does not include it.
• CamVid and COCOStuff. Those results come only from the paper, with no released config to check them against.

Lineage

• Builds on: external attention, arXiv:2105.02358 [guo2021beyond] (GFA is derived from it §3.1). Dual-resolution design from DDRNet, arXiv:2101.06085 [hong2021deep] (DAPPM and the dual-branch fusion §3.2).
• Compared against: the Cityscapes, CamVid, ADE20K and COCOStuff baselines in tab. 2 to tab. 4, including DDRNet-23 and -Slim, STDC, BiSeNet and SegFormer.
• Superseded by: not known from content already in this wiki.