Semi-Supervised Semantic Segmentation with Pixel-Level Contrastive Learning from a Class-wise Memory Bank

In plain words

Semantic segmentation means labelling every pixel of an image, for example road, car or sky. Drawing those labels is slow, so a method that can learn from a few labelled images plus many unlabelled ones is useful. This paper's teacher-student setup does that by letting a slowly moving copy of the model label the unlabelled images, and by training the student on both kinds of images.

The new part is a memory. The method keeps a small store of pixel features for each class, taken from the labelled images where the model was right and confident. The student is then encouraged to make each of its pixel features resemble stored features of the same class. This pushes classes apart in feature space across the whole dataset, not only within one image, and that helps most when there are very few labels.

TL;DR

Semi-supervised segmenter with a pixel-level, positive-only contrastive loss against a per-class FIFO memory bank of teacher features from labelled data. The student is trained on labelled and unlabelled images, the teacher is an exponential moving average (EMA) of the student, and inference uses only the student. The authors report state-of-the-art mIoU in most Cityscapes and Pascal VOC settings, with the largest margins at the lowest label fractions tab. 2 tab. 3.

Problem

Supervised segmentation needs dense per-pixel labels, which are expensive §1. Earlier semi-supervised segmentation methods use per-sample entropy minimisation and per-sample consistency regularisation, and the authors note these do not shape the feature space across the dataset to separate classes §1. The paper targets the regime where labelled data is scarce §1 tab. 2 tab. 3.

Method

The teacher f_\xi produces pseudo-labels for unlabelled images and fills the memory bank from labelled images. The student f_\theta is trained with four loss terms: supervised, pseudo-label, entropy and contrastive §3 eq. 1. The teacher is updated by EMA of the student weights, not by gradient descent §3 eq. 2.

\mathcal{L} = \lambda_{sup}\mathcal{L}_{sup} + \lambda_{pseudo}\mathcal{L}_{pseudo} + \lambda_{ent}\mathcal{L}_{ent} + \lambda_{contr}\mathcal{L}_{contr} eq. 1

\xi = \tau\xi + (1-\tau)\theta eq. 2

Supervised loss §3.1. Weighted cross-entropy on weakly augmented labelled images. The per-pixel weight is 1 and the per-class weight is set in §4.2 eq. 3 eq. 4.

\mathcal{H}(y_1, y_2) = -\frac{1}{N}\sum_{n=1}^{N}\sum_{c=1}^{C} y_2^{(n,c)}\log(y_1^{(n,c)})\,\alpha^c\,\beta^n eq. 3

Pseudo-label loss §3.2. Teacher pseudo-labels, recomputed every iteration, supervise two strong augmentations of each unlabelled image. Each pixel is weighted by the teacher's confidence raised to the power s=6 eq. 5 eq. 6 §4.2.

Entropy loss §3.3. Direct entropy minimisation of the student on unlabelled predictions eq. 7.

\mathcal{L}_{ent} = -\frac{1}{A}\frac{1}{N}\sum_{a=1}^{A}\sum_{n=1}^{N}\sum_{c=1}^{C} f_\theta(x_u^{a,n,c})\log f_\theta(x_u^{a,n,c}) eq. 7

Contrastive loss §3.4. Student features go through a projection head g_\theta and a prediction head q_\theta, as in BYOL and SimSiam §3.4. The memory bank stores teacher projections z' of labelled pixels per class §3.4. Each feature in a batch and each memory entry gets a learned weight from a class-specific attention module, rescaled so the weights sum to the count eq. 8. The loss is a weighted cosine distance between student predictions and memory entries, averaged over pairs and classes eq. 9 eq. 10 eq. 11.

w_{p_c} = \frac{N_{P_c}}{\sum_{p_i \in P_c} S_{c,\theta}(p_i)}\, S_{c,\theta}(p_c) eq. 8

\mathcal{C}(p_c, z'_c) = \frac{\langle p_c, z'_c \rangle}{\lVert p_c \rVert_2\, \lVert z'_c \rVert_2} eq. 9

\mathcal{D}(p_c, z'_c) = w_{p_c}\, w_{z'_c}\, (1 - \mathcal{C}(p_c, z'_c)) eq. 10

\mathcal{L}_{contr} = \frac{1}{C}\frac{1}{N_{p_c}}\frac{1}{N_{z'_c}} \sum_{c=1}^{C}\sum_{p_c \in P_c}\sum_{z'_c \in Z'_c} \mathcal{D}(p_c, z'_c) eq. 11

The stated rationale is that the memory holds high-quality features from all labelled samples, so the contrastive loss separates classes across the whole dataset and also aligns the unlabelled feature distribution with the labelled one §3.4.

Memory bank §3.4. Each iteration, the teacher runs on labelled images. A feature quality filter keeps pixels whose predicted class matches the label and whose confidence is above \phi = 0.95 §3.4 §4.2. The teacher's projection of the kept features is ranked per class by the teacher's class-specific attention, and the top-K vectors are inserted. Each class is a FIFO queue of 256 entries, so the newest entries stay and the oldest are dropped §3.4 §4.2. Per image and per class, the number of inserted vectors is \max(1, \psi/|X_l|), with \psi = 256 §4.2.

The concrete loop, from the equations above and the schedule in §4.2:

# One student update  [§3] [eq. 1]
# Input:  labeled images and labels; unlabeled images; iteration it
# Output: scalar loss

def training_step(images, labels, unlabeled, it):
    # Pseudo-labels from the teacher, recomputed every iteration  [§3.2] [eq. 5]
    teacher_probs = softmax(teacher(unlabeled))        # (B_u, C, H, W)
    pseudo, confidence = max_and_argmax(teacher_probs)  # (B_u, H, W)
    # Supervised weighted cross-entropy on weak augmentations  [§3.1] [eq. 4]
    loss = weighted_ce(student(weak(images)), labels, class_weights)
    # Pseudo-label loss over A=2 strong views, pixel weight = confidence**6
    # [§3.2] [eq. 6] [§4.2]
    loss = loss + lambda_pseudo * pseudo_loss(unlabeled, pseudo, confidence)
    # Entropy minimisation of the student on unlabeled views  [§3.3] [eq. 7]
    loss = loss + lambda_ent * entropy(student(strong(unlabeled)))
    # Contrastive terms start after 2K iterations  [§4.2]
    if it >= 2000:
        # Labeled features against the memory bank, true labels  [§3.4] [eq. 11]
        loss = loss + lambda_contr * contrastive(labeled_feats, labels, memory)
        # Unlabeled features against the memory bank, pseudo-labels  [§3.4]
        loss = loss + lambda_contr * contrastive(
            unlabeled_feats, pseudo, memory)
    return loss

# Memory bank update, once per iteration from it >= 1000  [§3.4] [§4.2]
# Holds 256 projected teacher vectors per class, first in first out

def update_memory(teacher, images, labels, memory, batch_size):
    # Teacher predictions and features on labeled images  [§3.4]
    preds, probs = max_and_argmax(softmax(teacher(images)))
    feats = teacher.features(images)                     # (B, H', W', 256)
    # Keep pixels predicted correctly with confidence above 0.95  [§3.4] [§4.2]
    keep = (preds == labels) and (probs > 0.95)
    feats = teacher.projection(feats[keep])
    for c in range(num_classes):
        feats_c = feats[labels[keep] == c]
        # Rank candidates by the teacher's class-c attention  [§3.4]
        scores = teacher.attention[c](feats_c)
        # Insert the top-K per image: K = max(1, 256 / |X_l|) per image  [§4.2]
        new = top_k(feats_c, scores, k=K * batch_size)
        # FIFO: newest first, then keep at most 256  [§3.4] [§4.2]
        memory[c] = concat(new, memory[c])[:256]

# Positive-only class-wise contrastive loss  [§3.4] [eq. 8-11]
# feats: student features of one batch; classes: label per feature

def contrastive(feats, classes, memory):
    total = 0
    for c in range(num_classes):
        # Student prediction vectors of class c  [§3.4]
        preds = prediction_head(projection_head(feats[classes == c]))
        # Teacher projections of class c from the memory bank  [§3.4]
        targets = memory[c]
        # Attention weights, rescaled to sum to the count  [eq. 8]
        w_p = count(preds) * S_student[c](preds) / sum(S_student[c](preds))
        w_z = count(targets) * S_mem[c](targets) / sum(S_mem[c](targets))
        # Cosine similarity for every pair  [eq. 9]
        cos = normalize(preds) @ normalize(targets).T    # (N_p, N_z)
        # Weighted distance  [eq. 10]
        dist = w_p[:, None] * w_z[None, :] * (1 - cos)
        total = total + dist.mean()
    # Average over the C classes  [eq. 11]
    return total / num_classes

Hyperparameters

Values in the paper come from §4.2 unless marked. Where the paper gives no value, the code's value is shown and marked, and the code is the only source for it.

Results

Author-reported, on the Cityscapes and Pascal VOC validation sets, mean IoU (mIoU). "FS" is fully supervised with all labels. * means ImageNet pre-training, + means COCO pre-training tab. 2 tab. 3. Cityscapes, labelled fraction 1/30, 1/8, 1/4:

Pascal VOC, labelled fraction 1/50, 1/20, 1/8:

Semi-supervised domain adaptation, GTA5 to Cityscapes, DeepLabv2 R-101 tab. 4. "Adapt" uses GTA5 as extra labelled data; "no adapt" uses only Cityscapes:

Ablation on Cityscapes 1/30, DeepLabv2 R-101, COCO pre-trained tab. 5:

Sensitivity tab. 6 tab. 7 tab. 8, same setting:

• \lambda_{contr}: 10^4 gives 50.3; 10^2 gives 51.4; 10^1 gives 54.8; 10^0 gives 59.1; 10^-1 gives 59.4; 10^-2 gives 58.7; 10^-4 gives 57.6 tab. 6.
• Memory size \psi per class: 32 gives 58.7; 64 gives 58.9; 128 gives 59.2; 256 gives 59.4; 512 gives 59.3 tab. 7.
• Components, where the columns are teacher features (f_\xi), class-specific attention (S_{c,\theta}) and the feature quality filter (FQF): base 58.3; base + f_\xi 58.7; base + S_{c,\theta} 58.6; base + FQF 59.0; all three 59.4 tab. 8.

Cost, author-reported §4.4: at inference only the student runs, so there is no extra cost. One training step with one labelled and one unlabelled image at 512 x 512 costs 1151.19 GFLOPs, compared with 1488.16 for the DMT method and 1116.12 for ClassMix. The contrastive module adds 35.07 GFLOPs, of which 8.59 are the projection and prediction heads, 15.96 the attention modules and 10.52 the distance computation §4.4.

Paper vs code

The official repository is linked from the arXiv abstract (https://github.com/Shathe/SemiSeg-Contrastive). The README contact is the first author's address, and the README says the code is based on the ClassMix code. Commit dc53c6da64f91c1ebce320f6e2ce3cc9f3df7414, Apache-2.0 [LICENSE]. The main loop is trainSSL.py and the memory bank is utils/feature_memory.py. The domain-adaptation script trainSSL_domain_adaptation_targetCity.py was not read; the domain-adaptation claims are not checked against code.

Where the code differs from the paper, each item cites both sides:

1. Memory selection keeps the lowest scores, not the highest. The paper says the top-K highest-scoring vectors are inserted §3.4. The code sorts the attention scores ascending and keeps the first K, which are the lowest. A comment in the code says "highest rankings" feature_memory.py.
2. The memory-side attention weights cancel to 1. For memory entries the code computes the sigmoid output as a (1, N) tensor, and then rescales with shape[0] / sum(dim=0). Here shape[0] is 1, so every entry becomes 1/x * x = 1. The student-side rescaling is correct contrastive_losses.py, but the memory-side factor contrastive_losses.py has no effect on the loss as written. This is read from the code and was not executed. If correct, the learned weighting only acts on the student side in the released code.
3. Two contrastive terms, each weighted 0.1. The paper has one \lambda_{contr}\mathcal{L}_{contr} eq. 1 §4.2. The code adds the labelled and unlabelled contrastive losses separately, each multiplied by 0.1 trainSSL.py trainSSL.py. This is code-only; the paper's wording "features from both labeled and unlabeled data" can be read either way.
4. The warm-up is a ramp, not a switch. The paper sets \lambda_{pseudo}=0 for 2K iterations §4.2. The code multiplies the pseudo-label weight by \exp(-5(1 - t/2000)^2), which is near zero at the start and reaches 1 at 2000 trainSSL.py trainSSL.py. The memory bank starts filling at iteration 1000 and the contrastive loss starts at 2000 trainSSL.py trainSSL.py.
5. Strong ClassMix probability is 0.75, not 0.80. Table 1 lists 0.80 tab. 1. The code draws random() < 0.75 for each strong view, and only after 2K iterations trainSSL.py trainSSL.py. Weak ClassMix at 0.2 matches trainSSL.py.
6. Colour jitter magnitudes differ. Table 1 gives per-property maxima of 0.15 (weak) and 0.30 (strong) for brightness and contrast, and so on tab. 1. The code passes a single jitter_vale of 0.125 for weak and 0.25 for strong trainSSL.py trainSSL.py. Mapping between the two is not stated.
7. Entropy normalisation differs. The code takes base-2 logarithms, divides by \log_2 C and by the fraction of valid pixels trainSSL.py. The paper's equation uses the natural logarithm and divides by AN eq. 7.
8. DeepLabv3+ Pascal learning rate. The paper says 4e-4 for DeepLabv3+ §4.2. The Pascal DeepLabv3+ config uses 2e-4 configSSL_pascal_1_50_split0_v3.json. The Cityscapes DeepLabv3+ config uses 4e-4 configSSL_city_1_30_split0_v3.json.
9. Memory insertion count is rounded and scaled differently. The paper gives \max(1, \psi/|X_l|) per image and class §4.2. The code rounds \psi/|X_l| to an integer feature_memory.py, and for Pascal VOC multiplies by n_classes / 3 (7 for 21 classes) feature_memory.py.
10. Class balancing applies only to Cityscapes and only after 2K iterations. The paper says Cityscapes and GTA5 §4.2. The code checks dataset == 'cityscapes' and i_iter > RAMP_UP_ITERS trainSSL.py. The frequency window is a fixed number of iterations, and the statistics sum counts over that window class_balancing.py. The paper does not give the window.
11. Teacher EMA schedule. The paper says "from 0.995 to 1" §4.2. The code uses m = 1 - (1-0.995)(\cos(\pi t/T)+1)/2 trainSSL.py, plus the early cap at 1 - 1/(10t+1) trainSSL.py.
12. Background is removed for Pascal ClassMix. Class 0 is dropped before ClassMix mixing on Pascal VOC trainSSL.py. The paper does not say so.
13. Batch split. The paper gives one batch size per setting §4.2. The code uses the full batch for labelled images and half of it for unlabelled images, each with two views trainSSL.py.
14. Labelled counts. The paper gives only fractions. The configs use 100 labelled Cityscapes images for 1/30 configSSL_city_1_30_split0_v3.json and 212 Pascal images for 1/50 configSSL_pascal_1_50_split0_v3.json.

Agreements worth noting: pseudo-labels come from the EMA teacher by default trainSSL.py trainSSL.py, the memory is filled from teacher projections with the teacher's attention trainSSL.py, the confidence threshold is 0.95 trainSSL.py, and the loss weights in §4.2 match the code trainSSL.py trainSSL.py.

Pitfalls

• Memory direction. Implement top-K by score as the paper says §3.4, not the code's lowest-K. The released code does the opposite, so its behaviour differs from the paper's description.
• Memory-side attention. If you port the released code literally, the memory-side weights are a no-op contrastive_losses.py. The paper's eq. 8 does not say which attention module (S_{c,\theta} or S_{c,\xi}) weights memory entries at loss time §3.4. The code uses the student's memory-side module.
• Which contrastive term to use. The paper's single \mathcal{L}_{contr} with \lambda_{contr}=0.1 eq. 1 and the code's two 0.1-weighted terms give different effective weights (0.1 against 0.2 in total). The paper does not say how the labelled and unlabelled features are combined into one term.
• Warm-up. Zeroing versus ramping the pseudo-label weight changes early training. The memory bank starts filling at 1K iterations, before the loss starts at 2K trainSSL.py trainSSL.py.
• Class-wise loss normalisation. Classes with fewer than two features in the batch or memory are skipped but still counted in the 1/C factor contrastive_losses.py contrastive_losses.py. The paper does not say how to handle empty classes.
• Resolution. DeepLab's output is 8x smaller than the input. Labels are downsampled with nearest-neighbour interpolation to the feature size before the feature quality filter and the memory update trainSSL.py.
• Hidden cost. The contrastive term computes a dense cosine matrix between every batch feature and every memory entry, per class, every iteration contrastive_losses.py. Memory and compute grow with \psi; the paper notes this §4.4 tab. 7.
• Strong ClassMix and jitter. The table values and the code values differ (see Paper vs code, items 5 and 6). Which set reproduces the reported numbers was not checked here.
• Entropy scale. The log base and the \log_2 C normalisation change the effective weight of the 0.01 entropy term. The paper's equation gives no normalisation factor eq. 7.
• Abstract claims. The text says the method "outperforms the current state-of-the-art" [abstract]. The tables do not show that everywhere: DMT*+ beats the method on Pascal 1/8 with DeepLabv2 (72.7 against 71.8) tab. 3, and Lai et al. beat it on Cityscapes 1/4 (72.7 against 71.7) and Pascal 1/8 with DeepLabv3+ (72.4 against 72.0) tab. 2 tab. 3. The paper says Lai et al. use a higher FS baseline and a larger batch and crop §4.3.1.
• Ablation numbers. The text gives 54.32 mIoU for the contrastive-only row, while the table gives 54.4 §4.4 tab. 5.
• Non-stated details. Weight decay, the poly power, the LeakyReLU slope, the EMA shape and the labelled counts are not in the paper. Their values come from the code, and the code is the only source for them.

Lineage

• Builds on: positive-only contrastive learning without negatives, as in SimSiam (Chen and He, 2020) and BYOL (Grill et al., 2020) §2 §3.4; the Mean Teacher EMA teacher (Tarvainen and Valpola, 2017) §2 §3; augmentation anchoring from FixMatch (Sohn et al., 2020) §3.2; ClassMix for augmentation (Olsson et al., 2021) §4.2; CBC for pixel confidence weighting and class balancing (Feng et al., 2020) §4.2; L1 normalisation of attention weights following Sun et al. (2020) §3.4.
• Related pixel-level contrastive work: Wang et al. (2021) apply standard contrastive learning at pixel and region level in supervised settings §2; Lai et al. (2021) use a self-supervised pixel contrast §2. Both are compared against in tab. 2 and tab. 3.
• Compared against: adversarial and GAN-based methods (Hung et al., 2018; Mittal et al., 2019), consistency methods (French et al., 2019), DMT (Feng et al., 2020), CBC, ClassMix, and error-correction (Mendel et al., 2020) tab. 2 tab. 3.
• Superseded by: not known from wiki content.