[Paper Review] Max Pooling with Vision Transformers reconciles class and shape in weakly supervised semantic segmentation

Simone Rossetti, Damiano Zappia+3 more|arXiv (Cornell University)|Oct 31, 2022

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes ViT-PCM, a novel end-to-end weakly supervised semantic segmentation method based on Vision Transformers that replaces class activation maps (CAM) with a Global Max Pooling (GMP)-based Patch-Class Mapping (PCM) mechanism. By leveraging ViT’s localized feature representations and optimizing patch-to-class relevance end-to-end with a multi-label BCE loss, ViT-PCM achieves 69.3% mIoU on Pascal VOC 2012 val set—surpassing state-of-the-art CAM-based methods with fewer parameters and no multi-stage refinement.

ABSTRACT

Weakly Supervised Semantic Segmentation (WSSS) research has explored many directions to improve the typical pipeline CNN plus class activation maps (CAM) plus refinements, given the image-class label as the only supervision. Though the gap with the fully supervised methods is reduced, further abating the spread seems unlikely within this framework. On the other hand, WSSS methods based on Vision Transformers (ViT) have not yet explored valid alternatives to CAM. ViT features have been shown to retain a scene layout, and object boundaries in self-supervised learning. To confirm these findings, we prove that the advantages of transformers in self-supervised methods are further strengthened by Global Max Pooling (GMP), which can leverage patch features to negotiate pixel-label probability with class probability. This work proposes a new WSSS method dubbed ViT-PCM (ViT Patch-Class Mapping), not based on CAM. The end-to-end presented network learns with a single optimization process, refined shape and proper localization for segmentation masks. Our model outperforms the state-of-the-art on baseline pseudo-masks (BPM), where we achieve $69.3\%$ mIoU on PascalVOC 2012 $val$ set. We show that our approach has the least set of parameters, though obtaining higher accuracy than all other approaches. In a sentence, quantitative and qualitative results of our method reveal that ViT-PCM is an excellent alternative to CNN-CAM based architectures.

Motivation & Objective

• To address the limitations of CAM-based methods in weakly supervised semantic segmentation (WSSS), particularly poor localization and shape estimation.
• To explore whether Vision Transformers (ViT) can bypass the CAM bottleneck by leveraging their inherent feature locality and global context.
• To develop a simple, end-to-end method that jointly optimizes pseudo-mask quality and class localization without multi-stage refinement.
• To demonstrate that Global Max Pooling (GMP) enhances ViT’s ability to map patch features to class-relevant regions, improving segmentation accuracy.
• To achieve state-of-the-art performance on Pascal VOC 2012 with minimal parameters and no post-processing boosting.

Proposed method

• Introduces ViT-PCM, a Vision Transformer-based architecture that replaces CAM with a Patch-Class Mapping (PCM) head using Global Max Pooling (GMP) to extract the most relevant patch features per class.
• Employs a multi-label binary cross-entropy (BCE) loss to train the PCM head end-to-end, directly optimizing for pixel-level class relevance from image-level labels.
• Uses a two-branch architecture with shared weights to enforce translation and scale equivariance, improving robustness in pseudo-mask prediction.
• Applies a CRF as post-processing to refine the baseline pseudo-mask (BPM), but does not use CRF during training, keeping the method end-to-end and lightweight.
• Maps ViT patch features to segmentation masks via a learnable upsampling layer, avoiding complex refinement heads or auxiliary supervision.
• Leverages self-supervised pretraining (DINO) to enhance feature quality, which is then fine-tuned end-to-end with image-level labels.

Experimental results

Research questions

• RQ1 Can Vision Transformers (ViT) achieve better localization and shape estimation in weakly supervised semantic segmentation than CNNs using CAM?
• RQ2 Does Global Max Pooling (GMP) effectively extract class-relevant patch features from ViT, enabling accurate pseudo-mask generation without CAM?
• RQ3 Can a simple, end-to-end ViT-based method outperform complex, multi-stage CAM-based approaches in WSSS with only image-level supervision?
• RQ4 How does the proposed PCM mechanism compare to CAM in terms of mIoU, parameter efficiency, and robustness across different backbones?
• RQ5 Does the absence of multi-stage refinement or auxiliary supervision (e.g., saliency maps) compromise performance, or can it be avoided with proper feature mapping?

Key findings

• ViT-PCM achieves 69.3% mIoU on the Pascal VOC 2012 validation set using only image-level labels and end-to-end training, outperforming all prior state-of-the-art methods.
• The method achieves 71.4% mIoU when combined with CRF post-processing, demonstrating strong baseline pseudo-mask quality.
• ViT-PCM sets a new SOTA on the end-to-end benchmark, with 70.3% mIoU on the validation set and 70.9% on the test set, outperforming even multi-stage methods in some cases.
• The model uses significantly fewer parameters than competing methods, with the smallest parameter count among all SOTA approaches, as shown in Figure 5.
• On the MS-COCO 2014 dataset, ViT-PCM achieves 45.03% mIoU on the validation set, demonstrating generalization beyond Pascal VOC.
• Ablation studies show that ViT-PCM with PCM outperforms CAM on all backbones tested, with the largest gain on ViT-S/16 (43.3% vs. 29.3% mIoU), proving the superiority of the PCM mechanism.