WEViT: Weight-Entangled Vision Transformers with Class-Specific Attention for Weakly Supervised Semantic Segmentation

• We propose WEViT, the first framework that integrates Neural Architecture Search (NAS) with vision transformers for Weakly Supervised Semantic Segmentation (WSSS). • A one-shot weight-sharing strategy is employed to efficiently train an overparameterized supernet, enabling fast and effective transformer architecture search. • The framework utilizes multi-class tokens to extract class-specific attention from transformers, enhancing the precision of localization maps. • We introduce a novel Refinement Patch Affinity strategy that suppresses background noise and improves class focus in multi-class images. • A regularization loss function is designed to promote class-discriminative attention, with experiments highlighting the importance of transformer layer selection for maximizing segmentation accuracy. Weakly Supervised Semantic Segmentation (WSSS) is a challenging task in computer vision, as it relies on limited supervision to generate precise object localization maps, often using Class Activation Maps (CAMs). Traditional methods struggle with balancing localization accuracy and scalability due to their reliance on fixed network architectures and handcrafted strategies. Neural Architecture Search (NAS), despite its proven success in optimizing network designs across tasks, has not yet been explored in WSSS due to the need for efficient weight sharing. To address these limitations, we propose WEViT, a novel framework that integrates NAS with transformers to optimize network architectures and generate accurate and class-specific object localization maps for WSSS. Our approach leverages the weight entanglement strategy, enabling the supernet to train multiple subnets simultaneously while ensuring high-quality weight inheritance. This eliminates the need for retraining subnets from scratch, significantly reducing computational cost. The best-performing architecture, obtained through the evolutionary algorithm, is then utilized to extract attention weights from transformer heads. These weights are further refined using a Refinement Patch Affinity strategy, effectively removing background noise and enhancing focus on relevant classes in multi-class images. We also incorporate a regularization loss function during training to enhance the generation of class-discriminative localization maps, with experiments highlighting the critical role of transformer layer selection in this process. WEViT achieves state-of-the-art performance on PASCAL VOC 2012 and MS COCO, demonstrating the efficacy of applying NAS to WSSS for the first time and paving the way for scalable, efficient, and accurate segmentation solutions.