S2-UniSeg: Fast Universal Agglomerative Pooling for Scalable Segment Anything without Supervision

TL;DR

S2-UniSeg introduces a fast, scalable self-supervised segmentation method that generates pseudo-masks efficiently, outperforming existing models on multiple benchmarks through a novel pretraining approach.

Contribution

The paper proposes UniAP, a fast pseudo-mask algorithm, and S2-UniSeg, a scalable self-supervised segmentation framework with a new pretext task, improving performance and scalability.

Findings

Outperforms SOTA UnSAM with +6.9 AP on COCO

Achieves +11.1 AR on UVO

Gains +4.5 Pixel Accuracy on COCOStuff-27

Abstract

Recent self-supervised image segmentation models have achieved promising performance on semantic segmentation and class-agnostic instance segmentation. However, their pretraining schedule is multi-stage, requiring a time-consuming pseudo-masks generation process between each training epoch. This time-consuming offline process not only makes it difficult to scale with training dataset size, but also leads to sub-optimal solutions due to its discontinuous optimization routine. To solve these, we first present a novel pseudo-mask algorithm, Fast Universal Agglomerative Pooling (UniAP). Each layer of UniAP can identify groups of similar nodes in parallel, allowing to generate both semantic-level and instance-level and multi-granular pseudo-masks within ens of milliseconds for one image. Based on the fast UniAP, we propose the Scalable Self-Supervised Universal Segmentation (S2-UniSeg), which employs a student and a momentum teacher for continuous pretraining. A novel segmentation-oriented pretext task, Query-wise Self-Distillation (QuerySD), is proposed to pretrain S2-UniSeg to learn the local-to-global correspondences. Under the same setting, S2-UniSeg outperforms the SOTA UnSAM model, achieving notable improvements of AP+6.9 on COCO, AR+11.1 on UVO, PixelAcc+4.5 on COCOStuff-27, RQ+8.0 on Cityscapes. After scaling up to a larger 2M-image subset of SA-1B, S2-UniSeg further achieves performance gains on all four benchmarks. Our code and pretrained models are available at https://github.com/bio-mlhui/S2-UniSeg