Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers

In plain words

Semantic segmentation means labelling every pixel of an image: this pixel is road, that one is car, that one is sky. For years the standard recipe was a convolutional network that shrinks the image step by step, so that deep layers see a wide area and understand what they are looking at, followed by a decoder that blows the result back up to full size. The weak point is that early layers only ever see a small neighbourhood, so relating far-apart parts of the scene has to wait until late in the network.

SETR drops the shrinking encoder. It cuts the image into 16 by 16 pixel tiles, treats the tiles like words in a sentence, and feeds them to a standard Transformer, the same kind of model used for language and for image classification (ViT). Every Transformer layer lets every tile look at every other tile, so the whole image is in view from the first layer onward, and the grid of tiles never gets coarser. A small decoder then turns the per-tile features back into a per-pixel label map; the paper tries three decoders, from a single big upsampling step to gradual upsampling to mixing features from several Transformer layers.

The catch is that this only works with a Transformer pre-trained on large image classification data; trained from scratch it does badly. With pre-training, the authors report top results on ADE20K and Pascal Context and results close to the best on Cityscapes, and they show that the gain is not simply due to the bigger pre-training dataset.

TL;DR

SETR (SEgmentation TRansformer) encodes an image as a sequence of 16x16 patch embeddings plus learned position embeddings and runs a ViT encoder (T-Base: 12 layers, 768 hidden, 12 heads; T-Large: 24 layers, 1024 hidden, 16 heads) at a fixed 1/16 resolution §3.2 tab. 1. Three decoders map the 1/16 features to pixels: Naive (1x1 convs + one 16x bilinear upsample), PUP (four alternating conv + 2x upsample stages) and MLA (aggregates M evenly spaced encoder layers top-down) §3.3. Training follows mmsegmentation defaults with SGD, poly LR and auxiliary losses on intermediate layers §4.1. Author-reported: SETR-MLA reaches 50.28% mIoU on ADE20K val and 55.83% on Pascal Context with multi-scale testing; SETR-PUP reaches 82.15% on Cityscapes val (MS) tab. 4 tab. 5 tab. 6. Pre-training is essential: random initialisation gives 42.27% on Cityscapes tab. 2.

Problem

FCN-style encoders reduce spatial resolution layer by layer, and the receptive field of a convolutional stack grows only linearly with depth, so long-range context is available only in high layers §3.1. Prior fixes (dilated convolutions, larger kernels, pyramids, attention modules on top of the encoder) keep the encoder-decoder FCN structure; attention is added only on low-resolution high layers because of its quadratic cost in the number of pixels, leaving lower layers without long-range dependency learning §1 §3.1. Even convolution-free models such as Axial-DeepLab still downsample progressively §2. SETR asks whether a sequence-to-sequence encoder with global context in every layer and no downsampling can replace the FCN encoder §1.

Method

Image to sequence §3.2. Flattening raw pixels is infeasible (a 480x480x3 image would be a 691,200-long input). Instead, SETR fixes the sequence length to match the 1/16 feature map of a typical segmentation encoder: L = (H/16)(W/16) = HW/256. The image is split uniformly into an (H/16) x (W/16) grid of 16x16 patches; each flattened patch p is mapped by a linear projection f to e ∈ R^C, and a learned position embedding p_i is added per location:

E = \{e_1+p_1,\ e_2+p_2,\ \dots,\ e_L+p_L\} §3.2

Transformer encoder §3.2. L_e layers of multi-head self-attention (MSA) and MLP blocks with layer norm applied before both (omitted in the equations). For layer l with input Z^{l-1} ∈ R^{L×C}:

\text{query} = Z^{l-1}W_Q,\quad \text{key} = Z^{l-1}W_K,\quad \text{value} = Z^{l-1}W_V,\qquad W_Q, W_K, W_V \in \mathbb{R}^{C\times d} eq. 1

SA(Z^{l-1}) = Z^{l-1} + \operatorname{softmax}\!\left(\frac{Z^{l-1}W_Q (Z W_K)^\top}{\sqrt{d}}\right)(Z^{l-1}W_V) eq. 2

MSA(Z^{l-1}) = [SA_1(Z^{l-1}); \dots; SA_m(Z^{l-1})]\,W_O,\quad W_O \in \mathbb{R}^{md\times C},\ d = C/m §3.2

Z^l = MSA(Z^{l-1}) + MLP(MSA(Z^{l-1})) \in \mathbb{R}^{L\times C} eq. 3

The layer outputs {Z^1, ..., Z^{L_e}} are the features used by the decoders. All transformer layers and the patch projection are initialised from ViT or DeiT checkpoints; pre-trained position embeddings are 2D-interpolated for the fine-tuning input size §4.1.

Decoders §3.3. Each decoder first reshapes the needed features from (HW/256) x C to an (H/16) x (W/16) x C map.

1. Naive (SETR-Naive): Z^{L_e} → 1x1 conv → sync BN + ReLU → 1x1 conv to the number of classes (e.g. 19 for Cityscapes) → bilinear upsample to full resolution → pixel-wise cross-entropy §3.3.
2. PUP (SETR-PUP, progressive upsampling): alternate conv layers and 2x upsampling; four 2x steps take Z^{L_e} from 1/16 to full resolution. Further details (channel widths, kernel sizes) are only shown in fig. 1(b), not in the text §3.3 fig. 1.
3. MLA (SETR-MLA, multi-level aggregation) §3.3 fig. 1:

Algorithm: MLA decoder  [§3.3]
Input: encoder features Z^1..Z^{L_e}, number of streams M
1. Select Z^m for m in {L_e/M, 2·L_e/M, ..., L_e}  (M evenly spaced layers).
2. For each stream: reshape Z^m to (H/16, W/16, C).
3. Each stream: layer 1 = 1x1 conv halving channels.
4. Top-down aggregation: element-wise add each stream's layer-1 output to that of
   the stream above it, then apply an extra 3x3 conv to the summed feature.
5. Each stream: layer 2 = 3x3 conv; layer 3 = 3x3 conv halving channels,
   then 4x bilinear upsampling.
6. Concatenate all streams along channels; bilinear upsample 4x to full resolution.
Output: per-pixel features for classification

The text does not specify the exact ordering of "top-down" (which stream is top) or the final classifier; fig. 1(c) depicts it.

Auxiliary losses §4.1. Following PSPNet, auxiliary segmentation heads (each a 2-layer network) are attached to intermediate layers and trained concurrently with the main head: SETR-Naive at Z^10, Z^15, Z^20; SETR-PUP at Z^10, Z^15, Z^20, Z^24; SETR-MLA at Z^6, Z^12, Z^18, Z^24 (layer indices for T-Large).

Hybrid baseline §4.1. A ResNet-50 FCN encoder with output stride 1/16 whose output is fed into SETR-Naive-Base; transformer weights from ViT, DeiT or random init; ResNet-50 from ImageNet-1k.

Hyperparameters

Results

Author-reported. Baselines FCN and Semantic FPN are taken from mmsegmentation and trained/tested in the same settings §4.1. Ablations use Cityscapes train-fine, batch 8, single-scale val mIoU (%) tab. 2.

Variants on Cityscapes val, single-scale tab. 2 ("Pre": pre-training of transformer part; R = random init):

Pre-training vs FCN, single-scale (ADE20K val / Cityscapes val) tab. 3:

ADE20K val tab. 4 (SS single-scale, MS multi-scale):

A single model trained on ADE20K train+val for 160k iterations ranked first on the ADE20K test server leaderboard on the day of submission; the test score itself is not given §4.3.

Pascal Context val (60 classes) tab. 5:

Cityscapes val tab. 6:

Cityscapes test tab. 7 (‡: trained on fine + coarse data):

Observations the authors draw §4.2: PUP is best on Cityscapes, while MLA is slightly ahead of PUP and well ahead of Naive on ADE20K; T-Large beats T-Base; with ImageNet-21k pre-training, FCN improves but SETR still leads by a large margin on ADE20K, which they read as evidence that the gain comes from the sequence-to-sequence design rather than more pre-training data tab. 3. The authors note Axial-DeepLab trains at full 1024x2048 resolution with batch 32 for 60k iterations, versus SETR's 768x768 crops with batch 8 for 80k §4.3.

Qualitative figures compare SETR and dilated FCN on ADE20K, Pascal Context and Cityscapes fig. 2 fig. 3 fig. 4; fig. 5 visualises outputs of layers Z^1, Z^9, Z^17, Z^24; fig. 6 shows attention maps on Pascal Context; the appendix shows learned position embeddings that encode image distance, per-layer PUP features and recursive attention maps app. A fig. 7 fig. 8 fig. 9 fig. 10.

Paper vs code

Not checked. The paper points to a project page, fudan-zvg.github.io/SETR p. 1, but the authors' repository was not cloned for this card, so code, code_commit and code_license are left empty and no hyperparameter was verified against code. The paper states that data augmentation, training schedule and multi-scale testing follow mmsegmentation defaults §4.1, so any unstated value (poly power, normalisation constants, sliding-window stride, auxiliary loss weight) lives in that codebase's configs, not in the paper.

Pitfalls

• Pre-training is not optional. SETR-PUP from random init gets 42.27 vs 78.39 with ViT-21k at 40k iterations on Cityscapes tab. 2. Budget for loading ViT/DeiT weights and interpolating their position embeddings in 2D to the new grid (e.g. 768/16 = 48x48 for Cityscapes crops) §4.1.
• The block equations are loose. Eq. 2 puts a residual inside each single-head SA and mixes Z and Z^{l-1} in the key term; eq. 3 adds the MLP residual on the MSA output, and layer norm placement is only described as "before MSA and MLP" §3.2. Since the weights are loaded from ViT/DeiT checkpoints, implement the block exactly as those checkpoints expect rather than transcribing eq. 2-3 literally (with d = C/m, a per-head residual Z^{l-1} + (L×d) term is not even shape-consistent).
• Decoder internals are underspecified in text. PUP channel widths and kernel sizes, MLA's M and stream ordering, and the auxiliary head architecture beyond "2-layer" are not stated in prose §3.3 §4.1; fig. 1 is the only source.
• Naive decoder description is redundant. It projects to class count with a 1x1 conv, then mentions "a classification layer" after upsampling §3.3; whether a further layer exists is unclear.
• Auxiliary layer indices assume T-Large (24 layers) §4.1; for T-Base (12 layers) the paper does not say where auxiliary heads go.
• Inconsistencies in tab. 3. DeiT rows are labelled T-Large, but DeiT models are T-Base everywhere else tab. 2 tab. 4; SETR-PUP-DeiT on ADE20K is 46.24 in tab. 3 vs 46.34 in tab. 4; and the Cityscapes DeiT numbers (78.98, 79.45) match the 80k column of tab. 2 while the ViT rows (76.65, 78.39) match the 40k column, so rows in tab. 3 are not on the same schedule.
• Hidden costs. Sequence length is HW/256 with full global attention in all 24 layers and no downsampling: a 768x768 crop is 2,304 tokens per layer. T-Large SETR models have 305-318M parameters vs 68.59M for the FCN baseline tab. 2. Number and type of GPUs, memory use and training time are not stated.
• Output stride mismatch with baseline. SETR runs at stride 16 while the dilated FCN baseline runs at stride 8 §4.1.
• Cityscapes comparisons mix protocols. Tab. 6 compares SETR single- and multi-scale numbers against baselines whose test protocol is often unlabelled, and the Axial-DeepLab setup differs in crop size, batch and iterations §4.3.
• Weight decay is 0 for all experiments §4.1, unusual for fine-tuning and easy to overlook if copying a generic config.

Lineage

• Builds on: ViT (arXiv 2010.11929) for the patch-embedding transformer encoder and its ImageNet-21k weights; DeiT (arXiv 2012.12877) for ImageNet-1k weights; FCN and PSPNet-style auxiliary losses; FPN / Semantic FPN inspire the MLA decoder §3.3 §4.1.
• Compared against: dilated FCN and Semantic FPN (mmsegmentation), attention-augmented FCNs (Non-local, CCNet, DANet, OCRNet), DeepLab V3/V3+, and the attention-only but downsampling Axial-DeepLab tab. 4 tab. 5 tab. 6 tab. 7.