Rethinking Semantic Segmentation: A Prototype View

In plain words

Segmentation models label every pixel of an image. Most of them, whether built on convolutions or on Transformers, end with a set of learned "class cards": one vector per class that each pixel is compared against. The authors argue that these learned cards are the weak point. Each class gets only one of them, their number grows with the size of the label set, and nothing in the training loss makes the pixels of one class sit tightly together.

The proposed alternative learns no class cards. During training, the embeddings of each class's pixels are grouped into a few clusters, and the centre of each cluster becomes a "prototype" for that class. These centres are kept as running averages rather than trained by gradient descent. A new pixel is labelled with the class of the prototype it lies closest to, much like nearest-neighbour matching. The training loss pulls each pixel toward its assigned prototype and pushes it away from the others, so the shape of the feature space is learned directly.

TL;DR

• Most segmentation heads, from FCN softmax layers to query-based Transformers, can be read as learnable class prototypes that are trained with no regard to how representative they are §2.
• The method replaces them with K = 10 non-learnable prototypes per class. These are centres of clusters of training-pixel embeddings, found by online Sinkhorn clustering and updated by momentum §3 §5.1.
• Prediction is nearest-prototype retrieval. The count of learnable prototype parameters does not depend on the number of classes §5.3 tab. 4.
• Author-reported: +0.6 to +1.2 mIoU on ADE20K over four reproduced baselines tab. 1. Official code is MIT, but only a Cityscapes HRNet-W48 configuration is released.

Problem

The paper groups softmax heads (one weight vector per class) and pixel-query heads (one query vector per class) as parametric prototype learners §2, eq. 1–4. It names three limitations of that setting §2:

• One prototype per class. Each class is described by a single learned vector, so intra-class variation is not represented §2.
• Parameters grow with classes. A head needs at least D × C learnable values. For 800 classes and D = 512 that is at least 0.4 M parameters, and ten prototypes per class would need 4 M §2.
• Only relative distances are optimised. Cross-entropy shapes the gap between a pixel's own class and the others, but does not directly tighten intra-class distances §2.

The method assumes dense labels for the training images, which is the setting of all three benchmarks §5.1. Nothing is stated about label noise or about classes that never appear in a batch.

Method

Overview. An encoder φ maps an image to one D-dimensional embedding per pixel §3. Each class c owns K prototypes p_{c,k}. Each prototype is the centre of one sub-cluster of training pixels of class c in embedding space §3. A pixel is scored against every prototype by cosine similarity, and the class score is the best score among that class's prototypes §3, eq. 5–6. Training has five parts: cross-entropy on these class scores eq. 7; within-class online clustering, solved with an entropy-regularised transport problem and Sinkhorn iterations eq. 8–10 alg. 1; a pixel-prototype contrastive loss eq. 11; a compactness loss eq. 12; and a momentum update of the prototypes eq. 14. The prototypes receive no gradient §3.

Class probability from the nearest prototype eq. 5–6:

p(c \mid \bm{i}) = \frac{\exp\left(\max_{k} \bm{i}^\top \bm{p}_{c,k}\right)}{\sum_{c'=1}^{C} \exp\left(\max_{k} \bm{i}^\top \bm{p}_{c',k}\right)} eq. 5–6

Training objective eq. 7, eq. 13:

\mathcal{L}_{\text{CE}} = -\log p(c_i \mid \bm{i}) eq. 7

\mathcal{L}_{\text{SEG}} = \mathcal{L}_{\text{CE}} + \lambda_1 \mathcal{L}_{\text{PPC}} + \lambda_2 \mathcal{L}_{\text{PPD}} eq. 13

Online clustering. For the labelled pixels X^c of class c in a batch, the paper seeks an assignment L^c between pixels and the K prototypes of that class. The constraints are one prototype per pixel, and each prototype receiving about N/K pixels (equipartition) §3, eq. 8. The relaxation adds an entropy term weighted by κ and gives the Sinkhorn form §3, eq. 9–10. Here u and v are normalisation vectors computed by a few Sinkhorn–Knopp iterations §3, eq. 10.

\mathbf{L}^c = \mathrm{diag}(\bm{u})\,\exp\!\left(\frac{\bm{P}^{c\top}\bm{X}^c}{\kappa}\right)\mathrm{diag}(\bm{v}) eq. 10

Pixel-prototype contrast. Each pixel's hard prototype is k_i = argmax_k of its assignment §3. The positive is p_{c_i, k_i}. The negatives are the other C·K − 1 prototypes §3, eq. 11:

\mathcal{L}_{\text{PPC}} = -\log \frac{\exp(\bm{i}^\top \bm{p}_{c_i,k_i}/\tau)}{\exp(\bm{i}^\top \bm{p}_{c_i,k_i}/\tau) + \sum_{\bm{p}^- \in \mathcal{P}^-} \exp(\bm{i}^\top \bm{p}^-/\tau)} eq. 11

Compactness. The paper pulls each pixel to its assigned prototype with a squared-distance term, using normalised vectors §3, eq. 12:

\mathcal{L}_{\text{PPD}} = \left(1 - \bm{i}^\top \bm{p}_{c_i,k_i}\right)^2 eq. 12

Prototype update. Each prototype moves toward the normalised mean of the pixels assigned to it, with momentum μ §3, eq. 14:

\bm{p}_{c,k} \leftarrow \mu\,\bm{p}_{c,k} + (1-\mu)\,\bar{\bm{i}}_{c,k} eq. 14

The pseudocode below follows alg. 1 for the clustering and the equations above for the losses. Variable names are ours, and the step numbers are removed.

# One training step with non-learnable prototypes  [§3] [alg. 1] [eq. 5-14]
# Input:  images (B, 3, H, W); labels (B, H, W), -1 = ignore
# State:  prototypes (C, K, D), L2-normalised, updated, never trained by SGD
# Output: scalar loss

def training_step(images, labels, prototypes):
    # Embed every pixel and L2-normalise it  [§3]
    pixels = normalise(encoder(images))             # (N, D), N = B*H*W
    targets = labels.flatten()                      # (N,)

    # Cosine similarity of each pixel to every prototype  [§3] [eq. 5]
    sims = pixels @ prototypes.reshape(C * K, D).T  # (N, C*K)
    sims = sims.reshape(N, C, K)

    # Class logit = similarity to the closest prototype of that class  [eq. 6]
    class_logits = sims.max(dim=2)                  # (N, C)

    # Cross-entropy over the class logits  [eq. 7]
    loss_ce = cross_entropy(class_logits, targets)

    # Per class, softly assign its pixels to its K prototypes  [eq. 8-10]
    # Sinkhorn with kappa = 0.05 and 3 iterations  [alg. 1]
    for each class c present in the batch:
        soft = sinkhorn(sims[targets == c, c, :], kappa=0.05, iters=3)
        # Hard index per pixel  [§3]
        assigned[targets == c] = argmax(soft, dim=1)

    # Pixel-prototype contrast with temperature tau = 0.1  [eq. 11]
    loss_ppc = contrast(pixels, prototypes, assigned, tau=0.1)

    # Pull each pixel to its assigned prototype  [eq. 12]
    loss_ppd = mean((1 - dot(pixels, prototype_of(assigned)))**2)

    # Total loss; lambda1 = lambda2 = 0.01  [§5.1] [eq. 13]
    loss = loss_ce + 0.01 * loss_ppc + 0.01 * loss_ppd

    # Momentum update of prototypes, mu = 0.999; no gradient  [eq. 14]
    for each prototype (c, k) that received pixels:
        mean_k = mean(pixels[assigned == (c, k)])
        prototypes[c, k] = normalise(mu * prototypes[c, k]
                                     + (1 - mu) * mean_k)

    return loss

Hyperparameters

Results

Author-reported. ADE20K and COCO-Stuff baselines are the authors' own reproductions, using the MMSegmentation recipe §5.1. The paper reports mIoU (%), with parameter counts in millions.

Ablations on ADE20K val, MiT-B2, 40k iterations tab. 5:

Paper vs code

Official repository: https://github.com/tfzhou/ProtoSeg, MIT, at commit 1c4a7784bbce96c06fe72d55255af15e6cf1ca96. It is linked from the paper's title block and belongs to the first author. The README says it is built on openseg.pytorch README.md. The released code has one prototype model, HRNet-W48 on Cityscapes hrnet.py. No SegFormer, Swin, ADE20K or COCO-Stuff prototype configuration was found.

• Temperature. The paper uses τ = 0.1 in eq. 11 §5.1. The code takes the raw cosine similarity as logits and applies no temperature hrnet.py loss_proto.py.
• λ2. The paper sets λ2 = 0.01 §5.1. The released config sets loss_ppd_weight to 0.001 H_48_D_4_proto.json. λ1 = 0.01 matches H_48_D_4_proto.json.
• Which pixels update prototypes. Eq. 8 covers all labelled pixels of a class §3. The code keeps only pixels whose prediction is correct, both for the assignment and the update hrnet.py hrnet.py.
• Assignment used for the update. The paper describes a soft Sinkhorn solution and then a hard argmax index §3, eq. 10. The code draws a hard Gumbel-softmax sample (τ = 0.5) and uses it for the prototype update sinkhorn.py hrnet.py. The PPC target uses the argmax index hrnet.py.
• Learnable or not. The paper says prototypes are not learned by gradient descent §3. The code creates them as a trainable parameter hrnet.py, then replaces the parameter with a non-trainable one on every training step hrnet.py. From reading the code, the optimiser's first-step parameter object is no longer the one used afterwards. This was not run.
• Logits. Eq. 6 is a plain softmax over max cosine similarity §3. The code adds a projection head, a LayerNorm on embeddings and a LayerNorm over classes with learned affine parameters before the score hrnet.py hrnet.py. Inference takes argmax over these logits tester.py, so it is not pure nearest-prototype retrieval.
• Resolution of the losses. The paper does not say at which resolution the losses run. The code computes PPC and PPD on the feature map with nearest-downsampled labels hrnet.py, while cross-entropy upsamples the prediction to label resolution loss_proto.py.
• Crop. The paper uses 768 × 768 crops on Cityscapes §5.1. The config uses random crops of 1024 × 512 H_48_D_4_proto.json. Batch size 8 matches the run script run_h_48_d_4_proto.sh.
• Iterations. The paper trains Cityscapes for 160k iterations §5.1. The run script sets 80,000 run_h_48_d_4_proto.sh, and the README's table says 80K README.md.
• Released numbers. The README reports 80.1 mIoU for the proto model against a 79.0 HRNet baseline, both at 80K iterations and batch 8 README.md. The paper reports 81.1 against 80.4 on Cityscapes val tab. 2.
• Framework. The paper says the method is implemented in MMSegmentation §5.1. The repository is built on openseg.pytorch README.md.
• Parallel training. The code averages prototypes across processes after each step hrnet.py. The paper does not mention this.
• Empty classes. Classes without pixels in a batch are skipped hrnet.py. Prototypes that receive no pixels are not updated hrnet.py.

Pitfalls

• Equipartition in a small batch. Eq. 8 forces each prototype to receive about N/K pixels of its class in the batch §3. A class with few pixels in a batch is still spread over K prototypes. The paper does not discuss this, and the code does not change the constraint sinkhorn.py.
• Exponent range. The Sinkhorn step computes exp(sim / κ). With κ = 0.05 and cosine values in [−1, 1], the exponent spans about e^−20 to e^20 §3, eq. 10. The step runs under the mixed-precision autocast context used in training trainer.py. I did not check how autocast treats this operation.
• Undefined symbols in Alg. 1. Alg. 1 uses K and N without defining them, and its signature is missing a colon alg. 1. The code uses B for the batch dimension sinkhorn.py.
• Prototype initialisation. The paper does not say how prototypes are initialised. The code draws them from a truncated normal with std 0.02 and normalises them at the first forward pass hrnet.py hrnet.py.
• Momentum naming. The text calls the momentum m = 0.999 §5.1, while eq. 14 and the ablation use μ. The value is the same.
• Test-time cost. The paper says there is no extra test-time overhead §5.2. Retrieval compares each pixel to C·K prototypes instead of C class vectors. The paper gives no FLOPs comparison, so the claim cannot be checked from the paper.
• Table 4 parameter counts. The nonparametric rows show 0 learnable prototype parameters tab. 4. In the code, prototypes are stored as non-trainable nn.Parameter objects hrnet.py. A count that includes all parameters will include them.
• Inconsistent backbone name. The scalability study names MiT-V2 §5.3. The ablation names MiT-B2 §5.4. The table caption is also labelled MiT-B2 tab. 5.
• Inconsistent K. The text says K = 3 gives 46.0 §5.4. The table reports K = 5 at 46.0 and has no K = 3 row tab. 5.
• Inconsistent grouping. The paper places SETR and SegFormer in the standard softmax group §1 and also in the pixel-query group §2.
• Clustering cost. The discussion says the clustering costs about 2.5 ms for 10 K pixels into 10 prototypes app. D. The hardware for this figure is not stated.
• Released scope. Only the Cityscapes HRNet-W48 configuration is released, and the README says more results will come README.md. The ADE20K, COCO-Stuff, SegFormer and Swin rows cannot be reproduced from this release.
• Parameter tally. The 6.5 M prototype parameters quoted for the parametric baseline §5.3 are a function of D and C. D is not stated in the paper.

Lineage

• Builds on: the softmax and pixel-query decoders the paper groups under parametric prototype learning §2. SETR has a card: SETR, which the paper cites as a pixel-query style model §2.
• Compared against: FCN, HRNet, Swin and SegFormer, reproduced by the authors with the MMSegmentation recipe §5.1 tab. 1–3. The paper also compares with published numbers from DeepLabV3+, OCR, MaskFormer, SETR and Segmenter on ADE20K tab. 1.
• Superseded by: not known from content in this wiki.