# Know Your Attention Maps: Class-specific Token Masking for Weakly Supervised Semantic Segmentation

arXiv: https://arxiv.org/abs/2507.06848 · HTML: https://mlsanity.com/paper/2507.06848/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 34 references that are in the map.
- First posted on arXiv 9 Jul 2025.

## Builds on (most cited first, 26 of 26)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [Segment Anything](https://mlsanity.com/paper/2304.02643/) (2023)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [Learning Deep Features for Discriminative Localization](https://mlsanity.com/paper/1512.04150/) (2015)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Vision Transformers Need Registers](https://mlsanity.com/paper/2309.16588/) (2023)
- [Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned](https://mlsanity.com/paper/1905.09418/) (2019)
- [Learning Sparse Neural Networks through L0 Regularization](https://mlsanity.com/paper/1712.01312/) (2017)
- [Semi-Supervised Semantic Segmentation with Cross Pseudo Supervision](https://mlsanity.com/paper/2106.01226/) (2021)
- [ScribbleSup: Scribble-Supervised Convolutional Networks for Semantic Segmentation](https://mlsanity.com/paper/1604.05144/) (2016)
- [Learning Pixel-level Semantic Affinity with Image-level Supervision for Weakly Supervised Semantic Segmentation](https://mlsanity.com/paper/1803.10464/) (2018)
- [Self-supervised Equivariant Attention Mechanism for Weakly Supervised Semantic Segmentation](https://mlsanity.com/paper/2004.04581/) (2020)
- [Simple Does It: Weakly Supervised Instance and Semantic Segmentation](https://mlsanity.com/paper/1603.07485/) (2016)
- [Object Region Mining with Adversarial Erasing: A Simple Classification to Semantic Segmentation Approach](https://mlsanity.com/paper/1703.08448/) (2017)
- [Multi-class Token Transformer for Weakly Supervised Semantic Segmentation](https://mlsanity.com/paper/2203.02891/) (2022)
- [Learning Affinity from Attention: End-to-End Weakly-Supervised Semantic Segmentation with Transformers](https://mlsanity.com/paper/2203.02664/) (2022)
- [Box-driven Class-wise Region Masking and Filling Rate Guided Loss for Weakly Supervised Semantic Segmentation](https://mlsanity.com/paper/1904.11693/) (2019)
- [Class Re-Activation Maps for Weakly-Supervised Semantic Segmentation](https://mlsanity.com/paper/2203.00962/) (2022)
- [Scribble-Supervised LiDAR Semantic Segmentation](https://mlsanity.com/paper/2203.08537/) (2022)
- [Segment Anything Model (SAM) Enhanced Pseudo Labels for Weakly Supervised Semantic Segmentation](https://mlsanity.com/paper/2305.05803/) (2023)
- [Max Pooling with Vision Transformers reconciles class and shape in weakly supervised semantic segmentation](https://mlsanity.com/paper/2210.17400/) (2022)
- [DDANet: Dual Decoder Attention Network for Automatic Polyp Segmentation](https://mlsanity.com/paper/2012.15245/) (2020)
- [DuPL: Dual Student with Trustworthy Progressive Learning for Robust Weakly Supervised Semantic Segmentation](https://mlsanity.com/paper/2403.11184/) (2024)

## Built on by (most cited first, 3 of 3)

- [GeoSANE: Learning Geospatial Representations from Models, Not Data](https://mlsanity.com/paper/2603.23408/) (2026)
- [ModuSeg: Decoupling Object Discovery and Semantic Retrieval for Training-Free Weakly Supervised Segmentation](https://mlsanity.com/paper/2604.07021/) (2026)
- [DiCLIP: Diffusion Model Enhances CLIP's Dense Knowledge for Weakly Supervised Semantic Segmentation](https://mlsanity.com/paper/2605.04593/) (2026)

## Abstract (arXiv)

Weakly Supervised Semantic Segmentation (WSSS) is a challenging problem that has been extensively studied in recent years. Traditional approaches often rely on external modules like Class Activation Maps to highlight regions of interest and generate pseudo segmentation masks. In this work, we propose an end-to-end method that directly utilizes the attention maps learned by a Vision Transformer (ViT) for WSSS. We propose training a sparse ViT with multiple [CLS] tokens (one for each class), using a random masking strategy to promote [CLS] token - class assignment. At inference time, we aggregate the different self-attention maps of each [CLS] token corresponding to the predicted labels to generate pseudo segmentation masks. Our proposed approach enhances the interpretability of self-attention maps and ensures accurate class assignments. Extensive experiments on two standard benchmarks and three specialized datasets demonstrate that our method generates accurate pseudo-masks, outperforming related works. Those pseudo-masks can be used to train a segmentation model which achieves results comparable to fully-supervised models, significantly reducing the need for fine-grained labeled data.