SegViTv2: Exploring Efficient and Continual Semantic Segmentation with Plain Vision Transformers

TL;DR

SegViTv2 introduces a lightweight, efficient Vision Transformer-based framework for semantic segmentation, featuring a novel Attention-to-Mask decoder and a cost-effective encoder structure, with strong performance and continual learning capabilities.

Contribution

The paper presents SegViTv2, a novel semantic segmentation model that significantly reduces computational cost while improving accuracy, and extends it for continual learning with minimal forgetting.

Findings

Outperforms UPerNet with various ViT backbones

Reduces encoder computation by up to 50%

Achieves state-of-the-art results on ADE20k, COCO-Stuff-10k, PASCAL-Context

Abstract

This paper investigates the capability of plain Vision Transformers (ViTs) for semantic segmentation using the encoder-decoder framework and introduces \textbf{SegViTv2}. In this study, we introduce a novel Attention-to-Mask (\atm) module to design a lightweight decoder effective for plain ViT. The proposed ATM converts the global attention map into semantic masks for high-quality segmentation results. Our decoder outperforms the popular decoder UPerNet using various ViT backbones while consuming only about $5\%$ of the computational cost. For the encoder, we address the concern of the relatively high computational cost in the ViT-based encoders and propose a \emph{Shrunk++} structure that incorporates edge-aware query-based down-sampling (EQD) and query-based upsampling (QU) modules. The Shrunk++ structure reduces the computational cost of the encoder by up to $50\%$ while maintaining competitive performance. Furthermore, we propose to adapt SegViT for continual semantic segmentation, demonstrating nearly zero forgetting of previously learned knowledge. Experiments show that our proposed SegViTv2 surpasses recent segmentation methods on three popular benchmarks including ADE20k, COCO-Stuff-10k and PASCAL-Context datasets. The code is available through the following link: \url{https://github.com/zbwxp/SegVit}.