Exploring Token-Level Augmentation in Vision Transformer for Semi-Supervised Semantic Segmentation

TL;DR

This work tackles semi-supervised semantic segmentation with Vision Transformers by introducing TokenMix, a token-level augmentation that aligns with ViT's patch-based processing. It combines TokenMix with a dual-branch teacher-student framework to enhance consistency regularization while expanding augmentation diversity. TokenMix exchanges tokens between labeled and unlabeled images after patch embedding, enabling cross-image contextual learning without disrupting patch semantics. Empirical results on Pascal VOC 2012, Cityscapes, and COCO show strong gains under low-label regimes, demonstrating ViT's potential for S4 with the proposed augmentation strategy.

Abstract

Semi-supervised semantic segmentation has witnessed remarkable advancements in recent years. However, existing algorithms are based on convolutional neural networks and directly applying them to Vision Transformers poses certain limitations due to conceptual disparities. To this end, we propose TokenMix, a data augmentation technique specifically designed for semi-supervised semantic segmentation with Vision Transformers. TokenMix aligns well with the global attention mechanism by mixing images at the token level, enhancing learning capability for contextual information among image patches. We further incorporate image augmentation and feature augmentation to promote the diversity of augmentation. Moreover, to enhance consistency regularization, we propose a dual-branch framework where each branch applies image and feature augmentation to the input image. We conduct extensive experiments across multiple benchmark datasets, including Pascal VOC 2012, Cityscapes, and COCO. Results suggest that the proposed method outperforms state-of-the-art algorithms with notably observed accuracy improvement, especially under limited fine annotations.

Figures (6)

• Figure 1: The proposed framework. Labeled images undergo weak augmentation and are fed into the student model to generate predictions, which are then used to compute the supervised loss against finely annotated labels. Blue and green lines denote two independent branches for unlabeled images. Unlabeled images undergo weak augmentation and are input into the teacher model to generate pseudo-labels. Subsequently, the same unlabeled images undergo strong augmentation, TokenMix, and Dropout within the student model to calculate the unsupervised loss using the pseudo-labels.
• Figure 2: Process of TokenMix. Tokens are mixed subsequent to the patch embedding and are reverted to their initial arrangement before the decoder.
• Figure 3: Comparison of consistency regularization for optimization stability on Pascal (92 Labels).
• Figure 4: Effectiveness of TokenMix in different architectures on Pascal.
• Figure 5: Effectiveness of TokenMix in larger models.