Learning Affinity from Attention: End-to-End Weakly-Supervised Semantic Segmentation with Transformers

In plain words

Teaching a computer to colour every pixel of a photo with the object it belongs to usually needs a person to paint those masks by hand. This paper asks for less: a label per image that only says which objects are present, such as "dog" or "bus". The model has to work out the outlines itself.

Simple methods that do this first find a rough heat map of where each object appears, then use that heat map as a stand-in for the outline. Those heat maps can miss parts of an object that do not stand out. The authors use a Transformer as the image model, because its attention layers can relate pixels across the whole picture. They add a small extra part that learns from those attention links which pixels should share a label, and it uses this to fill in the rough heat map. A second, cheap step pulls the edges toward colour boundaries in the image. Everything trains in one pass, so there is no separate stage that first builds labels and then trains a second model on them. The reported scores are the authors' own. They are the best of the end-to-end methods the paper compares against, but several multi-stage pipelines score higher, some of them with saliency maps.

TL;DR

The paper trains one network end to end from image-level labels, with a MiT-B1 Transformer backbone §3.1 §4.1. Its affinity module (AFA) learns a semantic affinity from the last stage's self-attention, supervised by pseudo affinity labels §3.3 eq. 3 eq. 5. A pixel-adaptive refinement module (PAR) then cleans up the pseudo labels using image colour and position §3.4 eq. 8 eq. 9. The authors report 66.0 mIoU on VOC 2012 val (66.3 test) and 38.0 / 38.9 on MS COCO 2014 val, without and with CRF tab. 5 tab. 6. It is worth a look if you need an end-to-end weakly supervised segmenter. The code is public, but its repository has no LICENSE file.

Problem

Multi-stage WSSS pipelines train a classifier, extract class activation maps (CAM), refine them, and then train a separate segmenter §1. Their CNN backbones tend to activate only discriminative parts, so the masks are incomplete §1 §2.1. Earlier end-to-end methods avoid the extra stages but are built on convolutions and perform worse §1 §2.1. The paper argues that global attention in a Transformer can fix the incomplete regions. It also assumes that raw self-attention is not directly usable as an affinity, because nothing constrains it during training §3.3. The method's cost is the extra affinity and refinement modules, and the paper also keeps a CRF post-processing step in its headline numbers §4.3.

Method

Overview. A MiT-B1 Transformer encodes the image. Its last-stage features feed a linear classifier, and the classifier's weights turn them into a CAM per class §3.1 §3.2. The CAM is refined into two kinds of pseudo label, and the attention matrices of the last stage are turned into an affinity map. The affinity map is supervised by pseudo affinity labels built from the refined labels §3.3. The affinity then propagates the CAM, and the propagated, refined CAM supervises a segmentation head §3.5 fig. 3.

Steps

1. Backbone and attention. Patches are embedded as tokens. In each block, multi-head self-attention produces a matrix per head, and the block output passes through an FFN §3.1 eq. 1. The paper uses MiT-B1. Its final patch-merging step uses stride 1 rather than 2, so the last-stage feature map is h/16 by w/16 rather than h/32 §6.1.

2. CAM. The class map for class c is the ReLU of the classifier weights applied to the feature map, then min-max scaled to [0, 1] §3.2 eq. 2. A background score β separates foreground and background §3.2.

3. Pseudo label Y_p. A pixel takes the arg-max class when its CAM value is at least β_h, background when it is at most β_l, and "ignore" (255) otherwise §3.3 eq. 4. Default β_h = 0.55 and β_l = 0.35 §4.1.

4. Pseudo affinity label Y_aff. Two pixels are a positive pair when they share the class in Y_p, and a negative pair otherwise. Pairs with an ignored pixel are skipped. Only pairs within a local window are used §3.3.

5. Affinity A. For each block and head, the attention matrix S is added to its transpose, and a learned linear map combines all heads into one matrix A with a symmetric input §3.3 eq. 3. The paper calls this an MLP layer. The code applies a single 1×1 convolution with bias and a sigmoid model_attn_aff.py model_attn_aff.py.

6. Affinity loss. A binary loss with a positive term and a negative term. Each term is averaged over its own pairs §3.3 eq. 5. The code uses the same terms with a 0.5 weight on each losses.py.

7. Random-walk propagation. The transition matrix T is A raised to power α (here α = 2), with each row normalised by its sum: T = D⁻¹ Aᵅ §3.3 eq. 6. The CAM is propagated as M_aff = T · vec(M) §3.3 eq. 7.

8. PAR (pixel-adaptive refinement). For each pixel and each of its 8 neighbours, a Gaussian-style RGB term and a position term are computed. Each is normalised by a softmax over the neighbour set. Their sum gives the kernel §3.4 eq. 8 eq. 9. The CAM is updated by applying this kernel repeatedly, over neighbours at several dilation rates §3.4 eq. 10.

9. Losses. Classification uses a multi-label soft-margin loss §3.5 eq. 11. Segmentation uses cross-entropy against the propagated, PAR-refined label §3.5. The code computes it as the mean of a background-only and a foreground-only cross-entropy dist_train_voc.py. The total loss is L_cls + λ1·L_seg + λ2·L_aff + λ3·L_reg §3.5 eq. 12. The regularisation loss L_reg is described but its code is disabled (see Paper vs code).

Stage schedule. The paper trains the classifier alone for an initial phase, then turns on the affinity branch §4.1. The exact phase timings and the weight switch differ between the VOC and COCO code (see Pitfalls and Paper vs code).

# Training step, as described in the paper and in the code  [§3.1]-[§3.5]
# Input:  image batch (B, 3, H, W), image-level labels (B, C), iteration t
# Output: scalar loss

def training_step(images, cls_labels, t):
    # Backbone features, class logits, segmentation logits and attention  [§3.1]
    feats, attns = backbone(images)
    cls_logits, seg_logits = heads(feats)
    # Affinity from summed head pairs, then a linear map  [§3.3] [eq. 3]
    affinity = affinity_head(attns + transpose(attns))          # (B, hw, hw)
    # Initial CAM from classifier weights, scaled to [0, 1]  [§3.2] [eq. 2]
    cams = normalised_cam(feats, cls_weights)                   # (B, C, h, w)
    # Pseudo label with two background thresholds  [§3.3] [eq. 4]
    pseudo = threshold_cam(cams, beta_high, beta_low)           # (B, h, w)
    # Refine the pseudo label with the pixel-adaptive module  [§3.4]
    refined = par_refine(images, cams, pseudo)                  # (B, h, w)
    # Propagate the CAM with the learned affinity over a local window  [§3.3]
    propagated = random_walk(cams, affinity, radius=8, alpha=2)
    # Segmentation target: propagated CAM, then refined  [§3.5]
    seg_target = par_refine(images, propagated, pseudo)
    # Affinity target from the refined pseudo label, pairs within radius  [§3.3]
    aff_target = pairwise_labels(refined, radius=8, ignore=255)
    aff_loss = affinity_loss(affinity, aff_target)              # [eq. 5]
    seg_loss = cross_entropy(seg_logits, seg_target)            # [§3.5]
    cls_loss = multilabel_soft_margin(cls_logits, cls_labels)   # [eq. 11]
    # Weights from the paper's default setting  [§4.1] [tab. 10]
    return cls_loss + 0.1 * seg_loss + 0.1 * aff_loss

Hyperparameters

The paper's values are cited to the text, table or appendix. Values taken only from the released code are marked as such.

Results

Author-reported. Setup: PASCAL VOC 2012 trained on the augmented set (10,582 images), evaluated on val (1,449) and test (1,464); MS COCO 2014 trained on 82,081 images, evaluated on 40,137 val images; mIoU; image-level labels only; MiT-B1 backbone §4.1. Hardware is not stated.

Main VOC comparison tab. 5, mIoU in %.

The paper states that 66.0 on val is the value with CRF post-processing §4.3, so this is the CRF-included result. Whether the test value also includes CRF is not stated. The paper also notes its method reaches 83.8% of its fully supervised counterpart (SegFormer, MiT-B1, 78.7 val) §4.4 tab. 5.

Ablation on VOC val tab. 2:

Pseudo-label quality on VOC train / val tab. 3:

For reference, the paper lists 1Stage at 66.9 train and 65.3 val, and PSA at 59.7 train tab. 3.

PAR on CAM, VOC train tab. 4: CAM 48.2; PAMR 51.4; PAR with RGB term only 51.7; PAR with RGB and position terms 52.9.

Top-k pooling on CAM, VOC train / val tab. 1: gap 30.7 / 31.1; top-50% 34.5 / 34.8; top-25% 39.6 / 39.7; top-10% 43.5 / 43.6; gmp 48.2 / 48.3.

MS COCO 2014 val tab. 6: this paper 38.0 without CRF and 38.9 with CRF. Among multi-stage methods, only RIB (43.8) is above the non-CRF figure; EPS (35.7) and CGNet (36.4) are below it. The paper's own COCO table has no end-to-end competitor. The hyperparameter sweeps use VOC train or val only (tables 7 to 11).

Paper vs code

Repository: https://github.com/rulixiang/afa, first author's account, linked from the abstract. Commit adcfa29d2c84c09b5cf488fe8c44dbfdf9ff28fe (default branch master). The repository has no LICENSE file and the README names no license, so the code is not quoted and is cited by path and line only.

• Pseudo-label softmax. The paper defines the propagated CAM as T times the CAM eq. 7. The code applies a softmax over the class channels, including background, before propagation camutils.py. The paper does not mention this.
• Random-walk normaliser. The paper's row normalisation is exact eq. 6. The code adds 0.1 to the row sum, with a comment saying it avoids NaN camutils.py.
• Affinity source for propagation. The code propagates with the affinity from the largest scale, 1.5, only camutils.py. The paper does not state which scale is used.
• Affinity loss weights. The paper weights the positive and negative sums by 1/N⁺ and 1/N⁻ eq. 5. The code adds a 0.5 factor to each term losses.py, which halves the affinity loss relative to the printed equation.
• Affinity label resolution. The code resizes the refined label to h/16 by nearest-neighbour before building pair labels camutils.py. The paper does not state this step.
• Affinity head. The paper calls the affinity projection an MLP §3.3. The code uses one 1×1 convolution with bias on 16 channels (2 blocks × 8 heads, the last stage) and a sigmoid model_attn_aff.py model_attn_aff.py.
• Regularisation loss. The paper's default includes L_reg with λ3 = 0.01 §3.5 tab. 10. In the released training scripts the energy term is commented out dist_train_voc.py dist_train_coco.py, and its loss module needs a bilateral-filter wrapper that is not in the repository losses.py. The code therefore trains without L_reg.
• Loss weights. The code uses 1.0 for classification and 0.1 for segmentation and affinity dist_train_voc.py. This matches the paper's default.
• Warm-up. VOC: segmentation and affinity weights are 0 until iteration 2,000 dist_train_voc.py. The paper describes a further 4,000-iteration affinity warm-up §4.1. The VOC segmentation target switches to the propagated label only after 6,000 dist_train_voc.py. The VOC affinity label always uses the refined CAM label dist_train_voc.py. COCO: the affinity label uses the refined CAM label until 15,000 dist_train_coco.py.
• PAR iterations. The paper's default count is not stated in the text, and the table's best row uses 15 tab. 8. VOC code uses 10 dist_train_voc.py. COCO code uses 15 dist_train_coco.py.
• PAR scale terms. The paper names separate sigmas for RGB and position eq. 8. The code uses the same scale factor w1 = 0.3 for both, takes the standard deviation over the 8-neighbour set at each pixel, and averages squared per-channel RGB differences PAR.py PAR.py. The position weight is w2 = 0.01 in the code, which corresponds to w3 in the paper PAR.py PAR.py.
• PAR input resolution. Refinement runs on the image at half resolution camutils.py and uses replicate padding for neighbours PAR.py. Neither is stated in the paper.
• Normalisation layers. The code freezes all normalisation parameters (learning rate 0) dist_train_voc.py. The paper does not mention this.
• Batch size. The paper states 8. The VOC launch script gives 4 run_sbatch_attn_reg.sh. The COCO launch gives 8 run_sbatch_attn_reg_coco.sh.
• Test-time procedure. The paper's method does not describe evaluation. The released test script resizes the long side to 512, runs flip and three scales, takes the maximum over scales, and then applies dense CRF test_msc_flip.py test_msc_flip.py test_msc_flip.py. Its validation loop uses a background score of 0.35 for the affinity-propagated map dist_train_voc.py, which differs from training.

Pitfalls

• The headline number includes CRF. The VOC 66.0 and the COCO 38.9 both use dense-CRF post-processing §4.3 tab. 6. Without CRF the VOC ablation gives 63.8 tab. 2. CRF settings are not in the paper; use the code's values if you need to match it test_msc_flip.py.
• Affinity warm-up is not what the text says. Read the paper's 4,000-iteration affinity warm-up as the code's 6,000-iteration switch for the VOC segmentation target dist_train_voc.py. The affinity loss itself is active from iteration 2,000 in VOC dist_train_voc.py.
• Radius units. The radius r = 8 is in cells of the h/16 grid, so 128 input pixels at a 512 crop dist_train_voc.py. A radius counted in input pixels would be far too large.
• Affinity loss scale. The 0.5 factor on each term halves the affinity term relative to the printed loss losses.py. If you want the printed form, drop the 0.5s and retune λ2.
• Softmax before propagation. Propagating the softmax of the class channels (background included) is a design choice the paper does not describe camutils.py. Changing it changes what the random walk diffuses.
• Attention source. Only the last two blocks of the last stage are used. In MiT-B1 that stage has two blocks and 8 heads, so the affinity sees 16 maps mix_transformer.py model_attn_aff.py. The paper says "last stage" §6.1, which is consistent but easy to misread as a single block.
• Hyperparameter sweeps are on VOC. Tables 7 to 11 use VOC train or val, and no COCO setting is swept. Do not expect the same optimal values on COCO.
• Loss-weight typo. The text of §7.1 gives λ2 = 0.2 and then λ2 = 0.01 for λ3. The table default and §4.1 both give 0.1 and 0.01 tab. 10.
• Frozen normalisation and warm-up. A LayerNorm freeze and a 1,500-iteration LR warm-up are in the code and not the paper. Both change early training dist_train_voc.py voc_attn_reg.yaml.
• Compute. The affinity matrix covers the h/16 × w/16 grid, so 1,024 × 1,024 at a 512 crop §6.1. In the code, PAR runs its dilated neighbour passes (10 for VOC, 15 for COCO) on the image at half resolution camutils.py PAR.py. The paper does not report training time or memory.

Lineage

• Builds on: CAM §3.2 (Zhou et al., 2016, not linked here); random-walk propagation from PSA and IRN §3.3; pixel-adaptive convolution from 1Stage and PAMR §3.4; dense CRF §3.4; SegFormer's MiT backbone and MLP head §4.1.
• Compared against: 1Stage, RRM, AA&LR (end-to-end) and EPS, RIB, SEAM and others (multi-stage) tab. 5 tab. 6.
• Superseded by: not stated in this wiki.