Source-linked AI summary
Multi-class Token Transformer for Weakly Supervised Semantic Segmentation
Lian Xu, Wanli Ouyang, Mohammed Bennamoun, Farid Boussaid, Dan Xu
TL;DR
WSSS requires class-specific object localization from weak image-level labels, but standard transformer attention is class-agnostic. MCTformer learns class-linked tokens and refines their attention with patch affinities, achieving state-of-the-art results on PASCAL VOC and MS COCO.
Problem
Image-level labels lack localization information, while one-class-token transformer attention does not reliably associate attended regions with semantic classes.
Method
MCTformer learns one class-linked token per category and uses class-to-patch attention for localization, with patch-to-patch attention refining the resulting maps.
Results
MCTformer produces high-quality class-specific localization maps and achieves 71.6% mIoU on PASCAL VOC test and 42.0% mIoU on MS COCO.
Takeaways & Limitations
Class-to-patch attention provides class-specific localization, while patch-to-patch affinity and CAM complementarity improve pseudo-label quality for WSSS.
Abstract
from arXiv · showhide
This paper proposes a new transformer-based framework to learn class-specific object localization maps as pseudo labels for weakly supervised semantic segmentation (WSSS). Inspired by the fact that the attended regions of the one-class token in the standard vision transformer can be leveraged to form a class-agnostic localization map, we investigate if the transformer model can also effectively capture class-specific attention for more discriminative object localization by learning multiple class tokens within the transformer. To this end, we propose a Multi-class Token Transformer, termed as MCTformer, which uses multiple class tokens to learn interactions between the class tokens and the patch tokens. The proposed MCTformer can successfully produce class-discriminative object localization maps from class-to-patch attentions corresponding to different class tokens. We also propose to use a patch-level pairwise affinity, which is extracted from the patch-to-patch transformer attention, to further refine the localization maps. Moreover, the proposed framework is shown to fully complement the Class Activation Mapping (CAM) method, leading to remarkably superior WSSS results on the PASCAL VOC and MS COCO datasets. These results underline the importance of the class token for WSSS.
1. Introduction
WSSS needs better pseudo-label localization because image-level labels lack spatial information and existing transformer attention remains class-agnostic. MCTformer addresses this with class-specific tokens and attention refinement, achieving strong benchmark results.
- Image-level labels indicate class presence but provide no ground-truth localization, making high-quality pseudo-label generation a critical WSSS challenge.
- Single-class-token transformers produce noisy, non-discriminative localization because one token captures multiple object categories, context, and background.
- MCTformer learns multiple class tokens with class-aware supervision, linking each token to a class label and yielding class-specific class-to-patch localization maps.
- Patch-to-patch transformer attention supplies a pairwise affinity that refines class-specific attention maps without additional computation or supervision.
- 71.6% mIoU on PASCAL VOC test and 42.0% mIoU on MS COCO establish new state-of-the-art WSSS results.
2. Related works
WSSS research has addressed incomplete CAM supervision through specialized losses and affinity refinement, while transformer-based approaches introduce attention-derived localization. MCTformer extends transformer attention with class-specific localization using multiple class tokens.
- Raw CNN-based CAM maps are incomplete and have coarse boundaries, motivating specialized segmentation losses such as SEC, CRF, and contrastive loss.
- The framework splits images into patch tokens, adds C class tokens and position embeddings, processes them through transformer encoding layers, and uses the output class tokens for class scores.
- Affinity-based methods learn pairwise semantic relationships to refine CAM maps, including propagation through an affinity matrix and random walk.
- MCTformer uses transformer self-attention to extract class-specific object localization maps, unlike previous WSSS methods based on CNNs.
- Earlier ViT work found class-token attention reflected scene layout, but did not establish one-to-one class mappings or extend the finding to weakly supervised learning.
- TS-CAM uses class-agnostic ViT attention, whereas MCTformer exploits class-specific maps and better complements CAM, producing better localization maps than TS-CAM.
3. Multi-class Token Transformer
MCTformer learns class-specific transformer attention with multiple class tokens, refines it using patch affinity, and extends it with PatchCAM-based complementary supervision.
- Multi-class token structure design: MCTformer uses C class tokens alongside M patch tokens, with each class token intended to learn class-discriminative information.Class tokens and patch tokens are concatenated as transformer inputs, while class-aware supervision connects each class token to its corresponding class label.
- Class-specific transformer attention: Class-to-patch attention from each transformer layer yields class-relevant localization maps, which are fused across the last K layers and min-max normalized.Higher layers provide more discriminative representations, whereas earlier layers provide more general visual information, motivating the layer fusion.
- Class-specific attention refinement: Patch-to-patch attention provides a pairwise affinity without additional computation or supervision, refining localization maps with improved appearance continuity.The affinity is extracted from the patch-to-patch attention and applied to the class-specific transformer attention.
- Complementarity to Patch-Token CAM: MCTformer-V2 adds a CAM module that predicts classes from reshaped patch tokens while retaining class-token predictions, and trains with both classification losses.The CAM module uses a convolutional layer and global average pooling, creating PatchCAM maps that are combined with transformer attention at inference.
- Complementarity to Patch-Token CAM: PatchCAM maps and class-specific transformer attention are fused by element-wise multiplication before further patch-affinity refinement.The resulting maps are then used to produce the final object localization maps.
- Evaluation: The framework produces qualitative segmentation results on PASCAL VOC and MS COCO validation sets and evaluates pseudo-label seeds and masks on PASCAL VOC training data.The cited figures show input, ground truth, and model outputs; Table 1 evaluates initial seeds and corresponding pseudo segmentation masks using mIoU.
4. Experiments
Experiments evaluate MCTformer on PASCAL VOC and MS COCO, comparing segmentation, localization-map quality, model complexity, and design choices. The results show strong WSSS performance and benefits from class-specific attention, patch affinity, CAM complementarity, and attention fusion.
- Experimental settings: The approach is evaluated on PASCAL VOC 2012 and MS COCO 2014, using image-level supervision and standard WSSS segmentation settings.PASCAL VOC includes 20 object classes and an augmented training set of 10,582 images.
- State-of-the-art comparison: 71.9% and 71.6% mIoU are achieved on the PASCAL VOC validation and test sets, respectively, outperforming existing image-level-label methods.The method also performs comparably or better than methods using additional saliency maps.
- State-of-the-art comparison: 42.0% mIoU is achieved on MS COCO, while category-level comparisons show MCTformer outperforming other state-of-the-art methods on most object categories.The MS COCO result is reported on the validation set.
- Ablation studies: 47.2% mIoU for MCTformer-V1 class-specific attention exceeds TS-CAM* by 5.9 percentage points on the PASCAL VOC training set.TS-CAM obtains 29.9% mIoU, while adding ReLU produces an 11.4-point improvement for TS-CAM*.
- Ablation studies: 58.2% mIoU from MCTformer-V2 with CAM increases to 61.7% after patch-affinity refinement, while pseudo-label quality improves by 8% and segmentation by 3.2% for MCTformer-V1 refinement.Transformer attention localizes objects with lower responses and less noise, whereas PatchCAM activates more background pixels; fusion improves localization maps.
- Ablation studies: Average pooling yields 47.2% mIoU for class-specific localization, compared with 41.5% for fully connected projection and 26.8% for max pooling.The paper attributes average pooling’s advantage to encouraging attention to more relevant patches and better spatial context.
5. Conclusions
MCTformer uses multiple class tokens to generate class-specific localization and patch-to-patch attention to refine it. Combined with CAM, this framework produces high-quality pseudo labels and state-of-the-art WSSS results.
- 5. Conclusions: MCTformer produces class-specific object localization maps from class-to-patch attention associated with different class tokens.This framework is presented as a transformer-based approach for WSSS.
- 5. Conclusions: Patch-to-patch attention provides pairwise affinities that refine localization maps, while the framework complements CAM to generate high-quality pseudo ground-truth labels.The refinement and CAM combination are stated as complementary components of the framework.
A.1. Training and testing of MCTformer
MCTformer training integrates a CAM module and aggregates transformer attentions to produce class-specific localization maps and patch-level affinities at test time.
- Training and testing: MCTformer is trained with a 3 × 3 convolutional CAM module, AdamW, batch size 64, learning rate 5 × 10^-4, and 60 epochs.The CAM convolution uses C kernels, where C is the number of classes.
- Training and testing: Class-specific localization maps aggregate class-to-patch attention from the last three transformer layers, while patch affinity aggregates patch-to-patch attention from all twelve layers.Head attentions are first averaged within each layer, then combined across layers.
A.2. Training and testing for semantic segmentation
Semantic segmentation uses a ResNet38-based DeepLab-V1 model with specified augmentation, optimization, training, and multi-scale testing settings.
- Training and testing: The segmentation network is ResNet38-based DeepLab-V1, trained for 30 epochs with SGD, batch size 4, and initial learning rate 7 × 10^-4.Polynomial learning-rate decay uses power 0.9.
- Training and testing: Training applies random scaling, horizontal flipping, and 321 × 321 cropping, while testing uses multi-scale inference.The cited passage also specifies a scaling factor of ±0.3.
B. Additional quantitative results
Per-class IoU results are reported for PASCAL VOC validation and test sets and MS COCO validation. MCTformer outperforms other state-of-the-art methods on most object categories.
- Per-class IoU results are reported on PASCAL VOC validation and test sets and the MS COCO validation set.
- MCTformer outperforms other state-of-the-art WSSS methods on most object categories.The comparison is presented in Tables 8 and 9.
C. Additional qualitative results
The qualitative results show how MCTformer generates and refines class-specific localization maps across PASCAL VOC and MS COCO. PatchCAM and transformer attention exhibit complementary behavior for large, small, and irregular objects.
- Segmentation results: Qualitative segmentation results show clear boundaries on large objects and fine-grained details on small objects across indoor and outdoor scenes.The results are shown on the PASCAL VOC and MS COCO validation sets.
- MCTformer-V1: MCTformer-V1 uses patch-level pairwise affinity to refine class-specific transformer attention maps.The affinity propagates activations from similar regions to increase scores for otherwise inactive object patches.
- MCTformer-V1: The learned affinity map represents each selected patch’s similarity to all patches in the image.In the dog example, affinity highlights nearly the entire object region and propagates activation to an initially inactive patch.
- MCTformer-V2: PatchCAM benefits from the transformer’s global receptive field for full-context localization of large objects but can over-activate small or irregular objects.Examples include planes and trains for full-context localization, versus birds and plants for over-activation.
- MCTformer-V2: Transformer attention can assign small, evenly distributed values to large objects while preserving localization for small or slim objects.The paper attributes this behavior to class-token attention values summing to one.
- MCTformer-V2: MCTformer-V2 qualitative maps include PatchCAM, transformer attention, their fusion, affinity-refined fusion, and ground truth.These visualizations are provided for both PASCAL VOC and MS COCO training sets.