PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Lixiang Ru, Yibing Zhan, Baosheng Yu, Bo Du

arXiv:2203.02664v2cs.CV

TL;DR

Image-level WSSS is challenging because end-to-end CNN-based methods inadequately capture global feature relations, producing incomplete object regions. This paper uses a Transformer-based framework with AFA and PAR to refine pseudo labels, achieving 66.0% mIoU on PASCAL VOC 2012 and 38.9% mIoU on MS COCO 2014 while surpassing recent end-to-end methods and several multi-stage competitors.

  • Problem

    End-to-end WSSS with image-level labels remains challenging because CNN-based methods fail to explore global feature relations properly, affecting pseudo-label quality.

  • Method

    A Transformer-based end-to-end framework generates initial CAM pseudo labels, learns semantic affinity from MHSA with AFA, and refines labels using affinity propagation and PAR.

  • Results

    The method achieves 66.0% mIoU on PASCAL VOC 2012 and 38.9% mIoU on MS COCO 2014, surpassing recent end-to-end methods and several multi-stage competitors.

  • Takeaways & Limitations

    The framework demonstrates new state-of-the-art performance for end-to-end WSSS on PASCAL VOC and MS COCO.

Abstract

from arXiv · show

Weakly-supervised semantic segmentation (WSSS) with image-level labels is an important and challenging task. Due to the high training efficiency, end-to-end solutions for WSSS have received increasing attention from the community. However, current methods are mainly based on convolutional neural networks and fail to explore the global information properly, thus usually resulting in incomplete object regions. In this paper, to address the aforementioned problem, we introduce Transformers, which naturally integrate global information, to generate more integral initial pseudo labels for end-to-end WSSS. Motivated by the inherent consistency between the self-attention in Transformers and the semantic affinity, we propose an Affinity from Attention (AFA) module to learn semantic affinity from the multi-head self-attention (MHSA) in Transformers. The learned affinity is then leveraged to refine the initial pseudo labels for segmentation. In addition, to efficiently derive reliable affinity labels for supervising AFA and ensure the local consistency of pseudo labels, we devise a Pixel-Adaptive Refinement module that incorporates low-level image appearance information to refine the pseudo labels. We perform extensive experiments and our method achieves 66.0% and 38.9% mIoU on the PASCAL VOC 2012 and MS COCO 2014 datasets, respectively, significantly outperforming recent end-to-end methods and several multi-stage competitors. Code is available at https://github.com/rulixiang/afa.

1. Introduction

The paper targets image-level WSSS, where efficient end-to-end methods struggle to model global relations and produce complete object regions. It introduces a Transformer-based framework with AFA and PAR to improve pseudo labels and segmentation.

  • 1. Introduction: Image-level WSSS avoids labor-intensive pixel annotations but remains the most challenging weak-supervision scenario.Weak labels include image-level labels, points, scribbles, and bounding boxes.
  • 1. Introduction: Multi-stage WSSS trains separate classification and segmentation models, complicating training and reducing efficiency.End-to-end alternatives avoid this pipeline but commonly rely on CNNs.
  • 1. Introduction: CNN-based end-to-end methods fail to explore global feature relations properly, affecting the completeness and quality of generated pseudo labels.Transformers are introduced because self-attention can model global relations and discover more integral object regions.
  • 1. Introduction: The framework uses Transformers for initial CAM pseudo labels, AFA to learn affinity from MHSA, and random-walk propagation to revise labels.AFA addresses the coarseness of semantic affinity in MHSA.
  • 1. Introduction: PAR integrates local RGB and position information to refine pseudo labels, enforce low-level appearance alignment, and support end-to-end training.Experiments on PASCAL VOC 2012 and MS COCO 2014 report gains over recent end-to-end and several multi-stage competitors.

2. Related Work

Prior WSSS work commonly uses multi-stage CNN pipelines, while Transformer research has shown promise for global visual modeling and localization. This paper extends that direction by learning reliable semantic affinity from MHSA to propagate CAM labels.

  • 2. Related Work: Multi-stage WSSS methods generate CAM-based pseudo labels with classification networks before training a separate segmentation model.Erasing strategies and other refinements seek to overcome incomplete CAM activation.
  • 2. Related Work: CNN-based WSSS approaches use refinements such as pixel-adaptive convolutions, CRF, regularization, and affinity modeling to improve pseudo labels.These methods commonly retain convolutional backbones.
  • 2. Related Work: Vision Transformers have expanded from image classification into segmentation, depth estimation, and video understanding.TS-CAM applied Transformers to weakly-supervised object localization using image-level supervision.
  • 2. Related Work: This work learns reliable semantic affinity from MHSA and propagates CAM with the learned affinity for WSSS.The approach connects Transformer self-attention with semantic affinity rather than using attention directly.

3. Methodology

The framework uses a Transformer backbone to generate initial CAM pseudo labels, learns semantic affinity from MHSA with AFA, and refines labels using pixel-adaptive local appearance information.

  • Transformer Backbone: Transformer patches are processed with MHSA to capture global feature dependencies, producing feature maps for subsequent WSSS modules.Input images are split into h × w tokens, and stacked Transformer blocks generate the backbone features.
  • Initial Pseudo Labels: CAMs generated from Transformer feature maps serve as the initial pseudo labels for efficient WSSS.ReLU removes negative activations, Min-Max normalization scales each class activation map to [0, 1], and a background score separates foreground from background.
  • Affinity from Attention: AFA combines multi-head attention and its transpose through an MLP to predict symmetric semantic affinity.The predicted affinity is supervised with pseudo affinity labels derived from refined labels, using positive and negative local pixel pairs while ignoring uncertain regions.
  • Affinity from Attention: Affinity supervision encourages reliable semantic relations, and random-walk propagation uses the learned affinity to revise the initial CAM.Propagation diffuses high-affinity semantic regions and dampens wrongly activated regions, improving alignment with semantic boundaries.
  • Pixel-Adaptive Refinement: PAR refines initial and propagated CAMs with RGB and spatial affinity kernels computed from local neighbors over multiple iterations.The module addresses local inconsistency while avoiding the training slowdown associated with dense CRF refinement.
  • Network Training: Training jointly optimizes classification, affinity, and segmentation losses, together with a regularization loss for local prediction consistency.The segmentation branch uses affinity-propagated and PAR-refined labels as supervision.

4. Experiments

Experiments evaluate pooling, ablations, pseudo-label refinement, and segmentation performance on PASCAL VOC and MS COCO. The proposed Transformer-based framework benefits from AFA and PAR and achieves strong benchmark results.

  • 4.1. Setup: Experiments use PASCAL VOC 2012 and MS COCO 2014, with image-level labels in the training sets.PASCAL VOC has 21 classes; MS COCO has 81 classes.
  • 4.1. Setup: Global max-pooling produces favorable CAM performance for the Transformer-based classification network.The study compares generalized top-k pooling choices on PASCAL VOC train and validation sets.
  • 4.3. Ablation Study and Analysis: 66.0% mIoU is achieved on the PASCAL VOC validation set after combining PAR, AFA, auxiliary regularization, and CRF postprocessing.The Transformer baseline reaches 46.7%, while PAR and AFA raise performance to 56.2% and 62.6%, respectively.
  • 4.3. Ablation Study and Analysis: 66.3% mIoU is obtained for generated pseudo labels on the PASCAL VOC train set after AFA affinity learning, compared with 54.4% without it.Propagation further improves pseudo-label reliability, whereas propagation using naive MHSA significantly reduces accuracy.
  • 4.3. Ablation Study and Analysis: 52.9% is obtained after PAR refines Transformer-baseline CAM, improving over 48.2% and outperforming PAMR.The position kernel in PAR is also reported as beneficial for CAM refinement.
  • 4.4. Comparison to State-of-the-art: 38.9% mIoU is achieved on the MS COCO 2014 validation set, while PASCAL VOC results surpass previous state-of-the-art end-to-end methods.On PASCAL VOC, the method reaches 83.8% of its fully supervised SegFormer counterpart.

5. Conclusion

The conclusion presents an end-to-end Transformer framework for WSSS that uses semantic relations from self-attention to refine pseudo labels. It reports new state-of-the-art performance for end-to-end WSSS on PASCAL VOC and MS COCO.

  • 5. Conclusion: The framework generates initial CAM pseudo labels with a Transformer backbone and refines them using AFA-derived semantic affinity.AFA supervises MHSA with reliable affinity labels and applies random-walk propagation.
  • 5. Conclusion: The method achieves new state-of-the-art performance for end-to-end WSSS on PASCAL VOC and MS COCO.The conclusion also frames semantic-relation guidance as a perspective for improving feature aggregation in vision Transformers.

6. More Technical Details

The implementation uses SegFormer’s MiT-B1 backbone and extracts last-stage MHSA matrices for semantic-affinity prediction. Affinity computation is restricted to local pixel windows.

  • 6. More Technical Details: The MiT-B1 backbone is modified by changing the final patch-merging stride from 2 to 1 to increase CAM feature resolution.SegFormer produces multi-scale feature maps using overlapped patch merging.
  • 6. More Technical Details: Last-stage MHSA matrices are concatenated to predict semantic affinity across spatial locations, blocks, and heads.The matrices form S with dimensions involving hw, 256, and the number of blocks and heads.
  • 6. More Technical Details: Affinity loss considers pixel pairs only when they lie within the same local window of radius r.Pairs outside the window are ignored through a mask-based implementation.

7. More Experimental Results

Additional experiments examine affinity-loss radius, PAR configurations and weights, background-score filtering, per-category performance, and qualitative pseudo-label refinements. The results identify preferred hyperparameters and show that PAR and AFA improve pseudo-label quality.

  • Affinity from Attention: r = 8 is a proper local-window radius when computing the affinity loss.The study notes that smaller radii provide too few affinity pairs, whereas larger radii may reduce the reliability of distant pairs.
  • Pixel-Adaptive Refinement: PAR remarkably outperforms PAMR [4] under the same dilation rates, demonstrating the necessity of its position kernel.The ablation varies dilation rates, position kernel, and iteration number on PASCAL VOC 2012 train-set mIoU.
  • Weight Factors: w1 = 0.3, w2 = 0.3, w3 = 0.01 is a favorable choice for PAR weight factors.The experiment sets w1 = w2 for simplicity.
  • Background Scores: βh and βl trade pseudo-label reliability against the number of valid labels and introduced noise, with an average background score of 0.45 preferred for generated CAM.Larger βh and smaller βl yield more reliable labels but fewer valid labels; the reverse introduces noise.
  • More Quantitative Results: The method achieves the best results for most PASCAL VOC categories and qualitatively outperforms CNN and Transformer CAM baselines.Visualizations report more integral Transformer CAM regions, improved CAM from learned affinity, and further refinements from PAR and AFA.
  • Qualitative Results: Directly using MHSA as semantic affinity is not beneficial because self-attention heads contribute unequally, with some contributing negatively.The learned weights for heads #2, #3, and #5 are identified as negative contributors.
Loading 2203.02664v2…