S2-UniSeg: Scalable Self-Supervised Segmentation

• The paper introduces S2-UniSeg, a framework that replaces slow offline pseudo-label generation with fast online pseudo-mask extraction via Fast Universal Agglomerative Pooling (UniAP).
• It employs Query-wise Self-Distillation to align teacher and student segmentation queries for achieving class-agnostic, semantic, and panoptic outputs.
• The approach demonstrates significant performance gains on benchmarks like COCO and UVO by enabling continuous, large-scale self-supervised pretraining.

Searching arXiv for S2-UniSeg and closely related universal/self-supervised segmentation papers. First, locating the S2-UniSeg paper by arXiv ID. S2-UniSeg, short for Scalable Self-Supervised Universal Segmentation, is a self-supervised segmentation framework designed to learn a single segmentation model from unlabeled images, without manual masks or class labels, and to do so in a way that scales to large datasets. In this formulation, “self-supervised universal segmentation” means pretraining on unlabeled images so that one model can support class-agnostic instance segmentation, class-aware instance segmentation, semantic segmentation, and panoptic or universal segmentation. The method is built around two coupled contributions: Fast Universal Agglomerative Pooling (UniAP), which generates semantic-level, instance-level, and multi-granular pseudo-masks online, and Query-wise Self-Distillation (QuerySD), which trains a student model from a momentum teacher by matching segmentation queries rather than relying on a single global representation (Xu et al., 9 Aug 2025).

1. Concept and problem setting

S2-UniSeg is positioned against prior self-supervised segmentation methods that follow a cut-and-learn pipeline. In that earlier pattern, a model first extracts frozen self-supervised features, then runs an expensive graph-partitioning or mask-cutting algorithm offline on every image to generate pseudo-masks, then trains a segmentation model on those masks, then regenerates pseudo-masks from an updated checkpoint, and repeats this multi-round process. The paper names UnSAM, CutLER, U2Seg, MaskCut, and TokenCut as representative examples of this family (Xu et al., 9 Aug 2025).

The method identifies two core problems in that paradigm. First, pseudo-mask generation is too slow. On COCO val2017 with DINO ViT-B/8 features, the reported average per-image pseudo-mask generation times are 5.27 s for Divide-and-Conquer/UnSAM, 4.72 s for MaskCut, and 2.25 s for TokenCut, versus 0.045 s for Universal Agglomerative Pooling. Second, optimization is discontinuous: pseudo-masks are generated offline between training rounds, so the model alternates between “freeze model, label the whole dataset” and “train on those fixed labels.” This suggests that stale pseudo-labels and multi-stage refresh cycles can impede scalable continuous improvement (Xu et al., 9 Aug 2025).

Within this setting, “universal” follows the usage in U2Seg and Mask2Former-style formulations: the model is intended to represent both “thing” regions and “stuff” regions, and ultimately support panoptic-style outputs. This places S2-UniSeg in direct dialogue with unsupervised universal image segmentation systems such as U2Seg, which also target semantic, instance, and panoptic outputs from unlabeled data but rely on an offline pseudo-labeling and self-training pipeline built from DINO, MaskCut/CutLER, STEGO, and a Panoptic Cascade Mask R-CNN architecture (Niu et al., 2023). A plausible implication is that S2-UniSeg’s main novelty is not merely universal output coverage, but the replacement of offline pseudo-label generation by fast online pseudo-mask extraction coupled to continuous teacher-student optimization.

2. Overall framework and training loop

S2-UniSeg is organized as a teacher-student self-distillation system with online pseudo-mask generation. For each training image, the framework forms one global view, fed to the momentum teacher, and $\delta$ local views, produced by multi-crop augmentation and fed to the student. Both teacher and student use the same multi-scale encoder architecture. The teacher generates pseudo-masks online from its global-view features using Fast Universal Agglomerative Pooling (UniAP). These pseudo-masks include instance-level masks from spatially connected agglomeration, semantic-level masks from a fully connected semantic pooling stage that can merge disjoint regions, and multi-granular masks across multiple UniAP layers via progressively relaxed thresholds. The pooled mask features serve as teacher targets and are described as teacher queries (Xu et al., 9 Aug 2025).

The student receives local crops and predicts segmentation masks using a Mask2Former-style mask decoder. The decoder uses object queries partitioned into semantic queries, marked by a learnable token $[\mathrm{SEM}]$, and instance queries, marked by a learnable token $[\mathrm{INS}]$. The two groups share one decoder and are distinguished by query tokens rather than by separate decoding branches (Xu et al., 9 Aug 2025).

Because the teacher sees the full image and the student sees local crops, the teacher masks are cropped to the student view. Teacher queries or masks that do not overlap the student crop are discarded. Teacher masks are then matched to student predictions using bipartite matching, independently for semantic and instance levels, with only mask Dice similarity used as the matching criterion. After matching, the student query embeddings are trained to predict the corresponding teacher query embeddings through Query-wise Self-Distillation. The training loop is therefore: encode the global view with the teacher, run UniAP on teacher stride-4 features, encode local crops with the student, decode student masks and queries, crop teacher masks into each local view, match teacher and student masks separately for semantic and instance levels, apply query-wise distillation on matched queries, update the student by gradient descent, and update the teacher by momentum from the student (Xu et al., 9 Aug 2025).

This design differs substantially from supervised universal segmentation systems such as UniMedSeg, which unify visual examples, interactive prompts, language instructions, and 2D/3D medical images through a shared sequence-conditioned Transformer but remain explicitly supervised rather than self-supervised (Li et al., 14 Jul 2026). It also differs from prompt-driven supervised medical universal models such as MedUniSeg and UniSeg, which require task identities or modality-aware prompts and dense labels during training (Ye et al., 2024, Ye et al., 2023). S2-UniSeg instead frames universality as a consequence of unlabeled pretraining plus region-level self-distillation.

3. Fast Universal Agglomerative Pooling (UniAP)

UniAP is the key systems component that makes the framework scalable. It is a nonparametric hierarchical graph pooling and agglomerative region merging algorithm designed to generate universal pseudo-masks quickly enough to run inside every training step. UniAP takes as input the teacher’s largest-scale feature map, specifically the stride-4 feature map from the teacher encoder. Let

$\mathbf{F}\in \mathbb{R}^{HW\times d}$

be the L2-normalized feature map, where $H\times W$ is the spatial grid and $d$ is the feature dimension. At initialization, each token in the $H\times W$ feature map is treated as one graph node, and edges connect only horizontally or vertically adjacent nodes. Each node $i$ at layer $t$ has a binary mask

$\mathbf{M}_i^t \in \{0,1\}^{HW}$

indicating which original tokens belong to that node’s subtree. Initially,

$[\mathrm{SEM}]$0

and the initialized graph is

$[\mathrm{SEM}]$1

with $[\mathrm{SEM}]$2 and $[\mathrm{SEM}]$3. The spatial affinity matrix is

$[\mathrm{SEM}]$4

These definitions formalize UniAP as graph coarsening over teacher features rather than as a separate learned segmentation head (Xu et al., 9 Aug 2025).

At layer $[\mathrm{SEM}]$5, UniAP computes feature similarity for each connected pair of adjacent nodes:

$[\mathrm{SEM}]$6

It also defines a spatial similarity based on how similarly nodes relate to the original token map:

$[\mathrm{SEM}]$7

The final edge score is

$[\mathrm{SEM}]$8

and edges with $[\mathrm{SEM}]$9 are selected for coarsening. Once high-scoring edges are selected, UniAP uses SCC or connected-components grouping to identify node groups in parallel, encoded by

$[\mathrm{INS}]$0

The graph and masks are updated by

$[\mathrm{INS}]$1

A naive feature update is avoided because it can merge semantically different adjacent regions too early. Instead, supernode features are re-estimated directly from original token features:

$[\mathrm{INS}]$2

The paper describes this as a weighted voting mechanism that suppresses noisy boundary tokens (Xu et al., 9 Aug 2025).

UniAP separates instance pooling from semantic pooling. Because the initial graph uses only local adjacency and the update rule preserves connectivity, instance pooling naturally yields connected region masks appropriate for instance segmentation. To support semantic segmentation, where one class may appear in disconnected regions, UniAP constructs a fully connected graph

$[\mathrm{INS}]$3

after instance pooling and repeats the same identify-and-merge process to generate semantic masks. The graph derived from instance pooling is also used as input to the next UniAP layer. This is the mechanism by which the same pseudo-mask generator yields both connected instance masks and disconnected semantic-level masks (Xu et al., 9 Aug 2025).

UniAP is also hierarchical and multi-granular. With thresholds

$[\mathrm{INS}]$4

and default schedule

$[\mathrm{INS}]$5

earlier layers capture fine, highly similar regions, while later layers yield coarser groupings. The appendix algorithm uses inputs $[\mathrm{INS}]$6, temperature $[\mathrm{INS}]$7, similarity weights $[\mathrm{INS}]$8, and mask area threshold $[\mathrm{INS}]$9, with default $\mathbf{F}\in \mathbb{R}^{HW\times d}$0. This produces instance-level masks, semantic-level masks, and multi-granular masks online from the teacher feature map (Xu et al., 9 Aug 2025).

The practical speed claim is central: measured average pseudo-mask generation time per $\mathbf{F}\in \mathbb{R}^{HW\times d}$1 image on COCO val2017 is 5.27 s for Divide-and-Conquer/UnSAM, 2.25 s for TokenCut, 4.72 s for MaskCut, 4.86 s for MaskCut+STEGO, and 0.045 s for UniAP. This is the paper’s core argument that online pseudo-mask extraction is computationally feasible at scale (Xu et al., 9 Aug 2025).

4. Query-wise Self-Distillation (QuerySD)

QuerySD is the segmentation-oriented pretext task used to train S2-UniSeg. The motivation is that standard self-supervised learning methods such as DINO use a single global representation per image, whereas segmentation requires structured object- and region-level understanding. S2-UniSeg therefore condenses each image into a set of universal object queries and distills at the query level (Xu et al., 9 Aug 2025).

In this formulation, a query is a decoder-side object token used by the Mask2Former-style mask decoder. Teacher queries come from UniAP-pooled region features. Student queries come from the final decoder-layer query embeddings associated with student mask predictions. Since the teacher sees the full image and the student sees local crops, teacher pseudo-masks are cropped into the student view, non-overlapping teacher masks are removed, and teacher and student masks are matched independently for semantic and instance levels. The paper states explicitly that only mask Dice similarity is used as the matching criterion and that the two matching processes do not interfere with each other (Xu et al., 9 Aug 2025).

Let the matched teacher and student query embeddings be

$\mathbf{F}\in \mathbb{R}^{HW\times d}$2

where $\mathbf{F}\in \mathbb{R}^{HW\times d}$3 is the number of matched queries and $\mathbf{F}\in \mathbb{R}^{HW\times d}$4 is the query feature dimension. A 3-layer MLP projection head transforms queries to distribution logits. The paper writes QuerySD as

$\mathbf{F}\in \mathbb{R}^{HW\times d}$5

The extracted typesetting is incomplete, and the surrounding explanation states that this should be interpreted as the sum of DINO-style self-distillation loss over each matched teacher-student query pair. The paper includes the DINO reference loss

$\mathbf{F}\in \mathbb{R}^{HW\times d}$6

and presents QuerySD as its query-level analogue (Xu et al., 9 Aug 2025).

This suggests that S2-UniSeg reinterprets local-to-global invariance as a region correspondence problem. Rather than aligning one image embedding across crops, it aligns matched semantic and instance queries between a global-view teacher and local-view students. In comparison, ACSeg also treats segmentation as adaptive grouping on top of self-supervised ViT features, but its supervision is a modularity objective over image-specific concepts rather than a teacher-student distillation process over matched mask queries (Li et al., 2022). Likewise, UnSAMv2 learns granularity-controllable segmentation from unlabeled data by converting pseudo-hierarchies into mask-granularity pairs for SAM-2 adaptation, but it addresses promptable binary segmentation at variable scale rather than universal thing/stuff or panoptic-style segmentation (Yu et al., 17 Nov 2025).

5. Architecture, implementation, and scaling behavior

S2-UniSeg uses DINO-pretrained ViT-Base/8 as backbone, augmented with ViT-Adapter to produce multi-scale features. Appendix details include ViT-B/8 hidden dimension 768, 12 self-attention layers, interaction blocks inserted every 4 DINO blocks, and Deformable Attention in interaction blocks for linear complexity. To reduce cost, the image is resized by half before entering the DINO ViT branch, “as if using a patch size of 16,” while the adapter branch still receives the original image and yields a feature pyramid with strides 4, 8, 16, and 32. UniAP specifically uses the stride-4 teacher feature map (Xu et al., 9 Aug 2025).

The mask decoder is Mask2Former with 6 decoder layers, each containing masked cross-attention, self-attention, and feed-forward layers. Queries are partitioned into 50 semantic queries and 150 instance queries. The original Mask2Former classification head is replaced during pretraining by the QuerySD projection head. Teacher and student queries are taken as the query features before LayerNorm. Following DINO, the projection head is a 3-layer MLP with hidden dimension 2048, followed by L2 normalization and a linear layer with output dimension $\mathbf{F}\in \mathbb{R}^{HW\times d}$7, default $\mathbf{F}\in \mathbb{R}^{HW\times d}$8 (Xu et al., 9 Aug 2025).

Initialization and optimization are chosen to stabilize online pseudo-mask generation. The adapters are zero-initialized so that the model initially behaves like the pretrained DINO representation, allowing UniAP to generate sensible pseudo-masks from the start. The teacher is updated by momentum from the student, with a cosine schedule changing from 0.996 to 1. Data augmentation follows multi-crop training with one global view and local multi-crops; the default local crop scale is 0.05 to 0.4. Ablations test $\mathbf{F}\in \mathbb{R}^{HW\times d}$9, with 2 local views chosen in practice to save compute. Default hyperparameters are

$H\times W$0

Optimization uses AdamW, batch size 16, 8 A800 GPUs, learning-rate warmup to 0.000625 over 10k iterations, and total training of 160k iterations (Xu et al., 9 Aug 2025).

The main pretraining datasets are ImageNet-1k with 1.3M images and SA-1B subsets of 0.4M, 1.2M, and 2.0M images. After training, the teacher encoder plus mask decoder are used for inference, following official Mask2Former inference rules for semantic, instance, and panoptic tasks (Xu et al., 9 Aug 2025).

These choices distinguish S2-UniSeg from pseudo-label pipelines such as S-Seg, which also decouple segmentation structure from semantics but train open-vocabulary semantic segmentation from image-text pairs plus DINO+K-means pseudo-masks on a MaskFormer backbone (Lai, 2024). S2-UniSeg does not rely on captions or open-vocabulary text supervision; its universality is derived from thing/stuff pseudo-mask generation and query-level self-distillation from unlabeled images.

6. Empirical results, related comparisons, and limitations

S2-UniSeg is evaluated on COCO val2017 for class-agnostic instance segmentation and instance segmentation, UVO val for instance segmentation, COCOStuff-27 for semantic segmentation, and COCO val2017 and Cityscapes val for panoptic segmentation. Under ImageNet pretraining, it reports clear gains over U2Seg across all benchmark families. On COCO class-agnostic instance segmentation, CutLER (IN) attains AP 9.7 and AR 27.1, U2Seg (IN) AP 7.3 and AR 19.2, and S2-UniSeg (IN) AP 14.2 and AR 34.3. On class-aware instance segmentation on COCO, U2Seg (IN) reaches AP 6.4 and AR 18.5, while S2-UniSeg (IN) reaches AP 15.3 and AR 35.5. On UVO instance segmentation, U2Seg (IN) yields AP 6.2 and AR 21.0, whereas S2-UniSeg (IN) yields AP 16.2 and AR 32.1. On COCOStuff-27 semantic segmentation, STEGO gives PixelAcc 56.9 and mIoU 28.2, U2Seg (IN) gives PixelAcc 63.9 and mIoU 30.2, and S2-UniSeg (IN) gives PixelAcc 68.4 and mIoU 36.1. On COCO panoptic segmentation, U2Seg (IN) reports PQ 16.1, SQ 71.1, RQ 19.9, and S2-UniSeg (IN) reports PQ 20.2, SQ 80.6, RQ 26.7. On Cityscapes panoptic segmentation, U2Seg (IN) gives PQ 17.6, SQ 52.7, RQ 21.7, while S2-UniSeg (IN) gives PQ 20.5, SQ 60.1, RQ 29.7 (Xu et al., 9 Aug 2025).

The paper emphasizes the gains over U2Seg as AP +6.9 on COCO, AR +11.1 on UVO, PixelAcc +4.5 on COCOStuff-27, and RQ +8.0 on Cityscapes under the same setting. Relative to U2Seg’s multi-stage pseudo-labeling pipeline, this supports the claim that online universal pseudo-mask generation plus continuous teacher-student training can improve both efficiency and final segmentation quality (Xu et al., 9 Aug 2025).

Scaling on larger unlabeled data is another major result. Under the same 0.4M SA-1B setting, UnSAM reports AP 31.4 and AR 42.0 on class-agnostic COCO instance segmentation, while S2-UniSeg reports AP 33.6 and AR 43.9. As pretraining data grows, S2-UniSeg continues to improve: at SA-1B 1.2M it reaches AP 35.1 and AR 44.6, and at SA-1B 2.0M it reaches AP 36.4 and AR 46.3 for class-agnostic COCO instance segmentation. The same positive trend appears on COCO instance segmentation, UVO instance segmentation, COCOStuff-27 semantic segmentation, COCO panoptic segmentation, and Cityscapes panoptic segmentation (Xu et al., 9 Aug 2025).

The ablations clarify which design choices matter most. For similarity weights, $H\times W$1 is best, while removing spatial similarity by setting $H\times W$2 causes a large drop. For threshold schedules, very low thresholds such as 0.5–0.1 merge too aggressively too early, whereas the default 0.8–0.4 schedule balances accuracy and efficiency. For local crop number $H\times W$3, the method is not highly sensitive, and 2 local views are selected because more crops increase matching overhead. For projection dimension $H\times W$4, increasing from 128 to 512 to 1024 improves performance, suggesting that finer online clustering helps, though the very large output dimensions typical of DINO are unnecessary here (Xu et al., 9 Aug 2025).

Several constraints remain. The paper explicitly notes that UniAP has several hyperparameters that may require tuning, including threshold schedule, feature and spatial similarity weights, temperature, and area threshold. It does not provide a detailed memory benchmark, even though UniAP uses the dense affinity matrix

$H\times W$5

so memory growth with teacher token count is not analyzed in the reported experiments. The paper also does not present a full explicit total objective beyond QuerySD and the matching procedure, nor detailed ablations that isolate teacher-student training or multi-granularity as separate numeric components (Xu et al., 9 Aug 2025).

In the broader landscape, S2-UniSeg can be understood as a self-supervised counterpart to supervised universal segmentation blueprints such as UniMedSeg, which collapse visual examples, geometric prompts, text, and 2D/3D medical images into one shared sequence-conditioned Transformer but do not address self-supervised learning directly (Li et al., 14 Jul 2026). It is also more end-to-end than U2Seg’s staged clustering and self-training formulation (Niu et al., 2023). Its main transferable insight is architectural and systems-oriented: universal pseudo-mask extraction must be fast enough to move from offline dataset relabeling into the inner loop of training. Within that framing, S2-UniSeg demonstrates that online universal pseudo-mask generation through UniAP, coupled with query-level local-to-global self-distillation through QuerySD, is a viable route to scalable self-supervised universal segmentation (Xu et al., 9 Aug 2025).