Segmenter: Transformer for Semantic Segmentation

In plain words

Semantic segmentation means giving every pixel of a photo a label: road, person, sky, chair. Whether a small patch of grey is road or wall often depends on what is elsewhere in the picture, so a model needs to see the whole scene. Convolutional networks build that wide view slowly, layer by layer, and lean on extra machinery to widen it.

Segmenter skips convolutions entirely. It cuts the image into small square patches, treats them like words in a sentence, and runs a standard Vision Transformer over them, so every patch can look at every other patch from the very first layer. The transformer's output for each patch is then turned into class scores, and those coarse patch-level scores are smoothly stretched back to full image resolution.

There are two ways to do that last step. The simple one is a single linear layer per patch. The better one adds one learned "class token" per category, lets these tokens talk to the patch tokens through two more transformer layers, and then scores each patch by how similar it is to each class token, giving one mask per class.

The practical lessons are that pre-training on ImageNet is essential, that bigger models and smaller patches both help a lot, and that smaller patches sharpen boundaries at a steep cost in compute and memory.

TL;DR

Segmenter applies an ImageNet-pre-trained ViT (or DeiT) to semantic segmentation by keeping the output token of every patch and decoding it to class scores, either with a point-wise linear layer or with a "mask transformer" that jointly processes patch tokens and K learnable class embeddings and returns K masks via a scalar product §3. Training is plain per-pixel cross-entropy with SGD and a poly schedule; stochastic depth 0.1 helps and dropout hurts §4.2 tab. 2. Accuracy grows with model size and with smaller patches (Seg-B/8 49.54 vs Seg-B/16 48.06 mIoU single-scale on ADE20K val) tab. 3. The best model, Seg-L-Mask/16, reports 53.63% multi-scale mIoU on ADE20K val, 59.0% on Pascal Context and 81.3% on Cityscapes val (author-reported) tab. 6 tab. 7 tab. 8. Relevant if you want a simple, convolution-free segmentation baseline built from off-the-shelf ViT checkpoints.

Problem

Fully convolutional encoder-decoder networks dominate semantic segmentation, but convolutions only see a local neighbourhood, while labelling a patch often requires the global image context; existing fixes (dilated convolutions, spatial pyramid pooling, attention modules on top of a CNN backbone) add specialised layers and keep a convolutional, locally biased backbone §1 §2. Applying a transformer directly to pixels is too expensive because attention cost is quadratic in sequence length §1. Segmenter casts segmentation as sequence-to-sequence over patches so that global context is available in every layer, and asks whether ViT models pre-trained for classification can be fine-tuned on moderately sized segmentation datasets §1.

Method

The model is an encoder-decoder made only of transformer layers, trained end to end with per-pixel cross-entropy; at inference, argmax over classes after upsampling gives one label per pixel §3.

Encoder §3.1

1. Split the image \mathbf{x} \in \mathbb{R}^{H \times W \times C} into N = HW/P^2 patches of size P \times P, flatten each and project linearly with \mathbf{E} \in \mathbb{R}^{D \times (P^2 C)} to get \mathbf{x_0} \in \mathbb{R}^{N \times D} §3.1. (Code implements the projection as a convolution with kernel and stride P vit.py.)
2. Add learnable position embeddings: \mathbf{z_0} = \mathbf{x_0} + \mathbf{pos} §3.1.
3. Apply L pre-norm transformer layers eq. 1 eq. 2 with standard scaled dot-product attention eq. 3 to get \mathbf{z_L} \in \mathbb{R}^{N \times D} §3.1.

\mathbf{a_{i-1}} = \text{MSA}(\text{LN}(\mathbf{z_{i-1}})) + \mathbf{z_{i-1}} eq. 1

\mathbf{z_{i}} = \text{MLP}(\text{LN}(\mathbf{a_{i-1}})) + \mathbf{a_{i-1}} eq. 2

\text{MSA}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\left(\frac{\mathbf{QK}^T}{\sqrt{d}}\right)\mathbf{V} eq. 3

Linear decoder (baseline) §3.2

4. A point-wise linear layer maps \mathbf{z_L} to patch logits \mathbf{z_{lin}} \in \mathbb{R}^{N \times K}, reshaped to H/P \times W/P \times K, bilinearly upsampled to H \times W \times K, then softmax over classes §3.2.

Mask transformer decoder §3.2 fig. 2

5. Introduce K learnable class embeddings \mathbf{cls} \in \mathbb{R}^{K \times D}, randomly initialised, one per semantic class §3.2.
6. Concatenate them with the patch encodings \mathbf{z_L} and run a transformer encoder of M layers over the joint sequence §3.2.
7. Take the output patch embeddings \mathbf{z'_M} \in \mathbb{R}^{N \times D} (L2-normalised) and class embeddings \mathbf{c} \in \mathbb{R}^{K \times D} and compute masks by scalar product eq. 4:

\text{Masks}(\mathbf{z'_M}, \mathbf{c}) = \mathbf{z'_M}\mathbf{c}^T \in \mathbb{R}^{N \times K} eq. 4

8. Reshape to H/P \times W/P \times K, bilinearly upsample to H \times W \times K, then (per the paper) softmax over classes followed by a layer norm §3.2. The code orders these steps differently; see Paper vs code.

The authors relate the design to DETR, Max-DeepLab and SOLO-v2 object embeddings and note that swapping class embeddings for object embeddings would allow panoptic segmentation; this is not evaluated §3.2.

Algorithm: Segmenter forward pass (as implemented)  [§3] [code: segm/model/segmenter.py#L32-L48] [code: segm/model/decoder.py#L80-L102]
Input: image batch im (B, 3, H0, W0), K classes
1. Zero-pad im on bottom/right so H, W are multiples of P.
2. Patch-embed (conv P x P, stride P) -> (B, N, D); prepend CLS token
   (and DIST token for DeiT backbones).
3. Add position embeddings, bilinearly resized to the (H/P, W/P) grid if
   the grid differs from the stored one.
4. Run L pre-norm transformer blocks (GELU MLP, stochastic depth), final LayerNorm.
5. Drop the CLS (and DIST) tokens -> z_L (B, N, D).
6a. Linear decoder: logits = Linear(D -> K)(z_L), reshape to (B, K, H/P, W/P).
6b. Mask transformer:
    x = Linear(D -> D)(z_L); append K class embeddings -> (B, N+K, D)
    run M blocks; LayerNorm
    split: patches (B, N, D), classes (B, K, D)
    patches = patches @ W_patch; classes = classes @ W_class   (D x D each)
    L2-normalise both along D
    masks = patches @ classes^T                                (B, N, K)
    LayerNorm over the K dimension; reshape to (B, K, H/P, W/P)
7. Bilinearly upsample to (H, W); crop away the padding.
Output: per-pixel class logits (B, K, H0, W0)
Training: cross-entropy(logits, labels, ignore_index) [code: segm/engine.py#L20-L34]

Training and inference recipe §4.2

• Initialise the encoder from ImageNet-pre-trained ViT (ImageNet-21k, trained with the augmentation/regularisation recipe of Steiner et al.) or DeiT (ImageNet-1k); bilinearly interpolate the pre-trained position embeddings to the fine-tuning sequence length; initialise decoders from a truncated normal §4.2.
• Loss: per-pixel cross-entropy, no class reweighting. Optimiser: SGD, weight decay 0, poly decay \gamma = \gamma_0 (1 - N_{iter}/N_{total})^{0.9} §4.2.
• Augmentation: MMSegmentation pipeline: mean subtraction, random resize with ratio in [0.5, 2.0], random horizontal flip, random crop / pad to a fixed size §4.2.
• Inference: sliding window at the training resolution; multi-scale uses scales (0.5, 0.75, 1.0, 1.25, 1.5, 1.75) plus horizontal flip, averaging results §4.2.

Hyperparameters

Model variants tab. 1 §4.2. Head size is 64, heads = token size / 64, MLP hidden size = 4 x token size §4.2.

Naming: Seg-B/16 means Base with 16x16 patches; † marks DeiT backbones; "-Mask" marks the mask transformer decoder §4.2. Patch sizes used: 8, 16, 32 §4.2.

Results

Author-reported. Ablations are on ADE20K val, single-scale, linear decoder unless stated §4.3. Throughput (Im/sec) is measured on one V100 at 512 resolution with the largest batch that fits §4.3.

Regularisation, Seg-S/16, ADE20K val mIoU tab. 2:

Model size and patch size, ADE20K val tab. 3:

Decoder and object size, ADE20K val mIoU tab. 4:

Training-set size, Seg-S/16, ADE20K val tab. 5: 4k images 38.31, 8k 41.87, 12k 43.42, 16k 44.61, 20k 45.37 mIoU (SS). The authors read this as a marked drop below 8k images §4.3.

ImageNet pre-training, Seg-S/16, ADE20K val tab. 9 app. A:

State of the art, ADE20K val tab. 6:

Pascal Context val, mIoU (MS) tab. 7:

Cityscapes val, mIoU (MS) tab. 8:

Other reported findings:

• Starting from the improved ViT-B/16 weights instead of the original ones raises ADE20K mIoU from 45.69% to 48.06% §4.2.
• Speed/accuracy plot: Seg-B-Mask/16 is described as matching SETR-MLA at twice the speed and beating Swin-B on both axes; Seg-B-Mask/32 runs at 400 images/s, about four times faster than a ResNet-50 FCN baseline at similar accuracy §4.3 fig. 4.
• Qualitatively, DeepLabv3+ gives sharper boundaries, while Segmenter labels large instances more consistently and handles partial occlusion better §4.4 fig. 5. Both fail on small instances app. C fig. 12.
• Attention fields of view adapt to instance size, and some first-layer heads already attend far away app. B fig. 6 fig. 7. An SVD projection of the learned class embeddings groups related categories app. B fig. 8.

Paper vs code

Checked against rstrudel/segmenter at 20d1bfa (MIT).

• Mask normalisation order. The paper upsamples the masks, applies softmax over classes, then a layer norm §3.2. The code applies a LayerNorm over the class dimension to the N x K patch-level masks before reshaping and upsampling, and the model returns these as logits with no softmax; softmax happens only inside the cross-entropy loss during training and in window merging at inference decoder.py segmenter.py utils.py.
• Extra projections in the mask transformer. Eq. 4 is a plain product of decoder outputs. The code adds a linear input projection from encoder width to decoder width, a LayerNorm after the decoder blocks, separate learned D x D projections for patches and for class embeddings, and L2-normalises the class embeddings as well as the patches, so the masks are cosine similarities decoder.py decoder.py. The paper mentions L2 normalisation only for the patch embeddings in the text §3.2.
• Decoder dropout. The paper says all models are trained without dropout §4.3; the code's mask transformer config sets dropout 0.1 and drop_path 0.0 config.yml.
• CLS/DIST tokens. The paper's encoder sequence contains only patch tokens §3.1; the code keeps the pre-trained CLS token (and DeiT distillation token) through the encoder and drops them before decoding vit.py segmenter.py.
• Schedule details. The paper's poly formula decays to 0 §4.2; the code decays to min_lr 1e-5 and uses SGD momentum 0.9, neither stated in the paper scheduler.py train.py. Training length is configured in epochs (64 / 256 / 216), which works out to roughly the stated 160K / 80K / 80K iterations at the stated batch sizes config.yml.
• Augmentation. Beyond what §4.2 lists, the ADE20K pipeline includes MMSegmentation's PhotoMetricDistortion and a random crop with cat_max_ratio=0.75; normalisation divides by a std as well as subtracting a mean, with mean/std chosen per backbone (0.5/0.5 for ViT, ImageNet statistics for DeiT) ade20k.py utils.py.
• Cityscapes decoder depth. The paper uses a 1-layer mask transformer for Seg-L-Mask/16 on Cityscapes §4.4; the shared config defaults to 2 layers, so this must be set by hand config.yml.

Pitfalls

• The softmax-then-layer-norm step in §3.2 is internally inconsistent: a layer norm after a softmax destroys the stated property that class scores sum to 1. Follow the code (LayerNorm on the patch-level masks, softmax only in the loss / at inference).
• Pre-training is not optional. Training Seg-S/16 from scratch gives 4.42 to 12.51 mIoU versus 45.37 with ImageNet-21k weights tab. 9; and which ImageNet weights matters (45.69 vs 48.06 for ViT-B/16) §4.2.
• Dropout hurts sharply (45.37 → 36.63 with dropout 0.2) tab. 2. Make sure the encoder's dropout is really 0 when loading timm/ViT configs whose defaults may differ (the code's VisionTransformer defaults to dropout 0.1 and relies on the training script passing 0 vit.py train.py).
• Patch size is a compute trap. Seg-B/8 gains 1.48 mIoU over Seg-B/16 but runs at 7 vs 92 images/s tab. 3. Since N = HW/P² §3.1, halving P quadruples the token count, and global attention cost grows with N².
• Position embeddings must be resized to the fine-tuning grid by bilinear interpolation §4.2; the code also re-interpolates on the fly whenever the input grid differs from the stored one, e.g. during sliding-window inference with padding vit.py. Inputs are zero-padded to a multiple of P and predictions cropped back utils.py.
• Inconsistent numbers across tables. Seg-L-Mask/16 is 51.30 (SS) in tab. 4 but 51.82 in tab. 6; the paper does not say why (plausibly the 640 vs 512 training resolution, not stated). Seg-B†/16 is 47.10 in tab. 4 vs 47.08 in tab. 3 and tab. 6. Throughput for Seg-B†/16 is 92 images/s in tab. 3 but 77 in tab. 6.
• The headline ADE20K comparison is not like-for-like in resolution: the best model trains at 640x640 while most Segmenter rows use 512 §4.2; its multi-scale margin over Swin-L UperNet is 0.13 mIoU tab. 6.
• The mask transformer does not always help per bucket: Seg-B-Mask/16 is worse than the linear Seg-B/16 on large objects (52.66 vs 53.28) tab. 4, although overall mIoU improves.
• Learning-rate remark is unclear. §4.2 says the schedule is like DeepLabv3+ "with learning rates divided by a factor 10 except for Cityscapes"; use the explicit values (1e-3, 1e-3, 1e-2) §4.2.
• Small datasets degrade results: below about 8k training images ADE20K mIoU drops noticeably tab. 5; Segmenter is less dominant on the smaller Cityscapes benchmark tab. 8.
• Sliding-window stride is not stated in the paper; the ADE20K config uses non-overlapping windows (stride = window size) config.yml.

Lineage

• Builds on: ViT (arXiv 2010.11929) for the encoder and DeiT (arXiv 2012.12877) for ImageNet-1k weights; improved ViT training recipe from Steiner et al. (arXiv 2106.10270) for initialisation; the mask transformer is inspired by DETR object queries, Max-DeepLab and SOLO-v2 §2 §3.2 §4.2.
• Compared against: SETR (ViT backbone with CNN-style decoders), Swin Transformer + UperNet, DeepLabv3+ with ResNet/ResNeSt backbones, OCR, DANet, CCNet and other FCN variants tab. 6 tab. 7 tab. 8.
• Superseded by: not known from content in this wiki.