Token Contrast for Weakly-Supervised Semantic Segmentation

In plain words

Teaching a segmentation model to colour every pixel usually needs masks drawn by hand. This paper uses a cheaper label: one tag per image that says which objects appear in it. The model has to find the outlines itself. A common first step is a heat map that shows where each object seems to be. Those maps tend to light up only the most obvious part of an object.

The authors use a Transformer as the image model, because its attention can relate pixels across the whole picture. They found a problem with that choice. As the layers stack up, the representation of neighbouring parts of the image becomes nearly identical, so the heat map cannot tell the parts apart. The middle layers still keep more variety, so the paper takes its signal from one of them. It asks the last layer to agree with the middle layer on which parts of the image belong together. A second idea pulls the summary vector of a small crop of an unclear area toward the summary vector of the whole image, so that the unclear parts join the object. Both ideas are added as extra training losses, and the segmentation model is trained in one stage.

TL;DR

The paper adds two contrastive losses to a ViT trained with image-level labels only. The patch-token loss (PTC) supervises the final patch tokens with pseudo token relations from block 10 §4.2. The class-token loss (CTC) contrasts local crops with the global class token §4.3. The single-stage model reaches 69.8 mIoU on VOC 2012 val with ViT-B and 71.1 with ImageNet-21k weights, both author-reported and with multi-scale test and CRF tab. 2. The public code has no LICENSE file, so use it as a reference only.

Problem

Class activation maps (CAM) derived from image labels usually cover only the most discriminative part of each object §1. ViT was proposed as a fix because self-attention sees the whole image §2. The paper observes over-smoothing in ViT: the final patch tokens become nearly uniform, with pairwise cosine similarities close to 1 §1 fig. 2. The authors measure this per block and find that late blocks are highly similar while early blocks keep semantic diversity §3.2 fig. 2. Because CAM is computed from the final patch tokens §3.1 eq. 1, uniform tokens give a uniform, poor CAM.

Method

Overview §4.1 fig. 3. ToCo keeps the ViT and adds one auxiliary classifier after an intermediate block. The auxiliary CAM supplies pseudo token labels for PTC and proposal regions for CTC. The final CAM gives the final pseudo labels. In the single-stage setting, these labels are refined with PAR and used to train a segmentation decoder §4.4.

CAM. For class c, the CAM is the ReLU of the class score map divided by its maximum §3.1 eq. 1. A background threshold separates background from foreground §3.1.

Auxiliary CAM. Patch tokens from block 10 are max-pooled over space and passed through a fully connected classifier §4.2 eq. 2. Two thresholds, beta_l < beta_h, split the auxiliary CAM into reliable background, uncertain and reliable foreground §4.2.

Patch Token Contrast (PTC) §4.2 eq. 3. Two patch tokens that share a reliable label form a positive pair. Two tokens with different reliable labels form a negative pair. Pairs that touch an uncertain token are ignored. The loss pushes positive pairs together and negative pairs apart on cosine similarity. The paper uses the absolute cosine so that strongly anti-correlated tokens are not treated as diverse §4.2.

# Patch Token Contrast loss  [§4.2] [eq. 3]
# Input:  final_tokens (B, N, d), aux_cam (B, C, h, w), cls_label (B, C)
# Output: scalar loss

def ptc_loss(final_tokens, aux_cam, cls_label):
    # 1. Reliable background and foreground from the auxiliary CAM  [§4.2]
    #    cam <= 0.25 background, cam > 0.7 class, else ignored
    token_label = pseudo_token_label(aux_cam, cls_label, low=0.25, high=0.7)
    # 2. Pair label: 1 same reliable class, 0 different reliable, else ignore
    pair_label = pairwise_label(token_label)              # (B, N, N)
    # 3. Absolute cosine similarity between final patch tokens  [§4.2]
    sim = abs(cosine_matrix(final_tokens))                # (B, N, N)
    # 4. Mean (1 - sim) over positive pairs, mean sim over negative pairs
    pos_term = mean(1 - sim[pair_label == 1])
    neg_term = mean(sim[pair_label == 0])
    return pos_term + neg_term

Class Token Contrast (CTC) §4.3 eq. 4 fig. 4. Each image gives one global view and several local crops of size 96 §5.1. Crops are sampled from uncertain regions (positive) and from background regions (negative), using the auxiliary CAM §4.3. The global class token is projected by a head whose weights are an exponential moving average (EMA) of the local head, and the gradient is stopped on the global branch §4.3. InfoNCE then pulls each positive local class token toward the global one and pushes negatives away, with temperature tau §4.3 eq. 4.

Key equations as printed in the paper:

\mathcal{L}_{ptc} = \frac{1}{N^+}\sum_{\mathbf{Y}_i = \mathbf{Y}_j}\bigl(1-\mathtt{CosSim}(\mathbf{F}_i,\mathbf{F}_j)\bigr) + \frac{1}{N^-}\sum_{\mathbf{Y}_i \neq \mathbf{Y}_j}\mathtt{CosSim}(\mathbf{F}_i,\mathbf{F}_j) eq. 3

\mathcal{L}_{ctc} = \frac{1}{N^+}\sum_{\mathbf{q}^+}\log\frac{e^{\mathbf{p}^\top\mathbf{q}^+/\tau}}{e^{\mathbf{p}^\top\mathbf{q}^+/\tau}+\sum_{\mathbf{q}^-}\bigl(e^{\mathbf{p}^\top\mathbf{q}^-/\tau}+\epsilon\bigr)} eq. 4

The printed form of eq. 4 has no leading minus sign and places epsilon inside the negative sum, as typeset. The objective is minimised, so implement the negative log (see Paper vs code).

Training objective §4.4 eq. 5. The classification loss on the final head and the auxiliary head are both multi-label soft margin losses §4.4.

\mathcal{L}_{toco} = \mathcal{L}_{cls} + \mathcal{L}^{m}_{cls} + \lambda_1\mathcal{L}_{ptc} + \lambda_2\mathcal{L}_{ctc} eq. 5

Single-stage training adds the segmentation loss and a regularisation loss, weighted by lambda_3, so that L = L_toco + lambda_3 L_seg §4.4. The paper gives no weight for the regularisation term and uses a standard cross-entropy for L_seg §4.4. The pseudo labels used for L_seg come from the final CAM, refined by PAR §4.4.

# Single-stage training objective  [§4.4] [eq. 5]
# Input:  scalar losses from the current batch
# Output: scalar loss for backward()

def total_loss(cls, cls_aux, ptc, ctc, seg, reg):
    # 1. Classification on final and auxiliary heads  [§4.4]
    loss = cls + cls_aux
    # 2. Token contrast with lambda_1 = 0.2, lambda_2 = 0.5  [§5.1]
    loss = loss + 0.2 * ptc + 0.5 * ctc
    # 3. Segmentation and regularisation, lambda_3 = 0.1  [§5.1]
    #    The regularisation weight 0.05 is code only  [code: train#L66]
    loss = loss + 0.1 * seg + 0.05 * reg
    return loss

Hyperparameters

Values in the paper are cited to the paper. Values marked "code only" appear in the official code and are not stated in the paper. Each one can change results.

Results

Author-reported. Setup: ViT-B backbone, single-stage, image-level labels only, VOC 2012 (augmented train, 10582 images) and MS COCO 2014 §5.1. Segmentation rows use multi-scale test and CRF §5.1. The table lists the paper's own numbers only; baselines are as the paper reports them.

Multi-stage methods with image labels only, on the same table: MCTformer 71.9 val, 71.6 test, 42.0 COCO; some multi-stage methods that also use saliency maps score higher, e.g. L2G 72.1 val, 71.7 test, 44.2 COCO tab. 2. The paper says its single-stage ToCo beats multi-stage methods that use image-level labels only §5.2. Table 2 does not fully support that: on VOC val, ESOL (69.9) is above ToCo (69.8), and MCTformer (71.9) is above ToCo-dagger (71.1). On COCO val, ReCAM (45.0), RIB (43.8) and ESOL (42.6) are above ToCo-dagger (42.3) tab. 2. The VOC test column is the only one where ToCo-dagger (72.2) leads all image-label-only rows.

Ablation on VOC val, final CAM (M), auxiliary CAM (M^m) and segmentation (Seg.) tab. 4:

Other analyses: auxiliary block 10 is best among blocks 8 to 11 tab. 5; local crop 96 is best among 64 to 112 tab. 5; EMA momentum 0.9 is best among 0 to 0.99 tab. 5; the chosen threshold pair is (0.25, 0.7) tab. 7; the absolute cosine beats plain cosine (36.5 vs 70.5 on M) tab. 6. ViT-S reaches 65.2 val segmentation and ViT-L-dagger 71.2 fig. 8 fig. 9. Per-class VOC results are in tab. 8.

Paper vs code

Official repository: rulixiang/ToCo at commit e13514d9a89a6c0b4131392a465fcf368a562b4e, as given in the paper's abstract README.md. The repository has no LICENSE file, so the code is cited only and nothing is copied from it. Each difference below was checked by reading the code at that commit.

• PTC scaling. Eq. 3 has no 0.5 factors and no +1 in the denominators eq. 3. The code multiplies each term by 0.5 and adds 1 to both counts losses.py. With lambda_1 = 0.2 §5.1, the effective PTC weight is half that of the printed form.
• PTC form. The code uses the absolute cosine, as the paper's text says §4.2 losses.py.
• CTC positives. The paper's positives are the uncertain local crops, and its negatives are background crops §4.3 eq. 4. The code also includes the two global views as positives, so each teacher global view is contrasted with student global views, and with itself losses.py.
• CTC extra term. The code's for ... else clause always runs, because there is no break. It adds log(1 + sum of negatives) for every sample, a term that eq. 4 does not contain losses.py.
• CTC loss sign. The code uses the negative log losses.py and normalises by two teacher views and the positive count losses.py. Eq. 4 as printed has no minus sign eq. 4.
• CTC crop rule. Crop centres come from background and uncertain pixels together, not only uncertain pixels camutils.py. A crop is negative when the mean of its region values (0 background, 1 uncertain, 2 foreground) is at most 0.2, and positive otherwise camutils.py. The paper describes the split as "uncertain" versus "background" §4.3.
• CTC local views. All eight local crops are cut from one local view of the image, and the loop that would make per-crop views is commented out voc.py.
• Projection heads. The paper describes separate global and local heads, P^g and P^l §4.3. The code uses one student head for both global and local views, and an EMA teacher copy of it model_seg_neg.py model_seg_neg.py. The code's head also has a 256-d bottleneck, a 2048-wide hidden stack and a 1024-way weight-normalised output model_seg_neg.py.
• CAM normalisation. Eq. 1 uses ReLU then division by the maximum eq. 1. The code also subtracts the minimum before dividing camutils.py. Training and evaluation both use it.
• CAM evaluation threshold. Training pseudo labels use the band (0.25, 0.7) with an ignored middle, and the background step at 0.5 is overwritten camutils.py. The CAM scores printed by validation and by the CAM inference script use a single background threshold of 0.5 dist_train_voc_seg_neg.py infer_cam.py. The paper does not state which threshold its CAM mIoU uses §5.2.
• Auxiliary tokens. The auxiliary CAM uses block 10 output before the final LayerNorm, and the final tokens use the normalised output vit.py. The paper does not say which is used.
• Positional embedding. The paper says bilinear interpolation §5.1. The code uses bicubic vit.py.
• Pretraining. Section 5.1 cites an ImageNet-21k source for the ViT-B backbone §5.1, but §5.2 and the Table 1 caption describe the plain ViT-B row as ImageNet-1k DeiT §5.2 tab. 1. The code's default backbone is deit_base_patch16_224 dist_train_voc_seg_neg.py, which loads the DeiT checkpoint deit.py. The ViT-B-dagger configuration uses vit_base_patch16_224, which is the one in the released VOC log toco_vit-b_voc_20k.log and in the evaluation command README.md.
• Learning rate groups. The paper has one peak rate §5.1. The code uses 10 times that rate for the classifiers, the projection head and the decoder dist_train_voc_seg_neg.py.
• Schedule. The poly decay is computed from step 0 over the full run optimizer.py, so when warm-up ends the rate drops from about its peak to about 0.93 of it on VOC (computed from the code's defaults). The paper describes warm-up and a polynomial decay only §5.1.
• Segmentation loss. The paper says "the common cross-entropy" §4.4. The code averages two cross-entropies, one over background-labelled pixels and one over foreground pixels losses.py.
• Regularisation weight. The paper gives no weight §4.4. The code uses 0.05, multiplied by the energy loss weight 1e-7 dist_train_voc_seg_neg.py dist_train_voc_seg_neg.py.
• Loss schedules. The paper gives no warm-up for the losses. VOC switches on only PTC and CTC from the start, and segmentation and regularisation after 2000 iterations dist_train_voc_seg_neg.py. COCO uses three stages (8000 and 12000) and pseudo labels from the auxiliary CAM before 12000 dist_train_coco_seg_neg.py.
• Weight decay. Weight decay 0.01 is code only dist_train_voc_seg_neg.py.
• Results. The README's ViT-B COCO value is 42.2 with multi-scale and CRF, and the paper's ToCo-dagger COCO value is 42.3 [README.md#L121] tab. 2. The README's DeiT-B and ViT-B values match the paper's ToCo and ToCo-dagger VOC columns (69.8 and 71.1 val with multi-scale and CRF) [README.md#L118-L119].

No discrepancies found: the classification loss, the auxiliary block index (10th block, which is aux_layer = -3), the thresholds (0.25 and 0.7; 0.65 on COCO), the temperature 0.5, the EMA momentum 0.9 and the loss weights 0.2, 0.5 and 0.1 match the code.

Pitfalls

• Which ViT-B. The paper calls both the plain and the dagger row ViT-B, and gives two different pretraining sources §5.1 §5.2 tab. 2. The default training command in the README does not pass --backbone, so it trains the DeiT-B backbone [README.md#L99]. Pass --backbone vit_base_patch16_224 for the ViT-B-dagger setting [README.md#L107].
• Eq. 4 sign and form. The printed loss has no minus sign eq. 4 and places epsilon inside the negative sum. Implement the negative log of the ratio, with epsilon once in the denominator, as the code does losses.py.
• Positives in CTC. Including the global views as positives and adding the log(1 + sum) term change the loss landscape losses.py. Decide whether you reproduce the paper's form or the code's, and say which in your report.
• Crop sampling is random per iteration. Crop centres come from a random permutation of background and uncertain pixels, so runs differ unless the seed is fixed camutils.py. Training uses seed 0 dist_train_voc_seg_neg.py.
• Background threshold 0.5 is dead in training. For training pseudo labels, the 0.5 step is overwritten by the ignore band camutils.py. Only 0.25 and 0.7 matter there.
• CAM scores differ by threshold. The validation CAM score uses one threshold of 0.5 dist_train_voc_seg_neg.py. Comparing your CAM mIoU with paper numbers needs the same threshold.
• Dependency outside the repo. The regularisation loss imports a compiled bilateral filter from a path under ./wrapper that the repository does not contain, and the path is tied to Python 3.8 losses.py. The README says to build the rloss extension [README.md#L92-L94].
• Compute. Each iteration runs the multi-scale CAM with three scales and flips, under no gradient camutils.py. The global views go through the backbone twice, once for the teacher and once for the student model_seg_neg.py. The PTC affinity matrix covers the token grid: a 448 crop with patch size 16 gives 28 x 28 = 784 tokens, so a 784 x 784 matrix per image deit.py §5.1.
• Positional embedding. The positional embedding is frozen vit.py and resized by bicubic interpolation at every forward pass vit.py.
• Single-stage claim. The text says ToCo beats image-label-only multi-stage methods §5.2, but Table 2 shows ESOL, MCTformer, ReCAM and RIB above ToCo on some columns tab. 2. Quote the table, not the claim.

Lineage

• Builds on: Learning Affinity from Attention (AFA), whose PAR refinement and single-stage framework ToCo reuses §4.4 §4.2. The classifier pooling follows AFA §4.2.
• Builds on: Emerging Properties in Self-Supervised Vision Transformers (DINO) for the class-token multi-crop idea and the EMA teacher §4.3 §5.1.
• Builds on: the Vision Transformer §2 (cited in the paper, no card here).
• Compared against: AFA, SLRNet, ViT-PCM, MCTformer, L2G, ESOL and others, as listed in tab. 1 and tab. 2.
• Superseded by: not known from content in this wiki.