Native Segmentation Vision Transformers

TL;DR

SeNaTra introduces a spatial grouping backbone that replaces uniform downsampling with content-aware token grouping, enabling native segmentation to emerge within the backbone itself. The method uses differentiable clustering with locality to produce a hierarchical, region-aware token representation, which can be interpreted as a Markov-chain decomposition across stages and upsample to full image resolution without pixel-level supervision. Empirically, SeNaTra achieves strong zero-shot segmentation from vision-language pretraining and outperforms several baselines on semantic and panoptic tasks when trained with mask supervision, while offering substantial parameter- and compute-efficiency and seamless compatibility with existing segmentation heads. This backbone-centric paradigm unlocks a segmentation-focused design space that couples boundary-preserving grouping with end-to-end differentiability, enabling scalable and efficient dense prediction with or without dedicated decoders.

Abstract

Uniform downsampling remains the de facto standard for reducing spatial resolution in vision backbones. In this work, we propose an alternative design built around a content-aware spatial grouping layer, that dynamically assigns tokens to a reduced set based on image boundaries and their semantic content. Stacking our grouping layer across consecutive backbone stages results in hierarchical segmentation that arises natively in the feature extraction process, resulting in our coined Native Segmentation Vision Transformer. We show that a careful design of our architecture enables the emergence of strong segmentation masks solely from grouping layers, that is, without additional segmentation-specific heads. This sets the foundation for a new paradigm of native, backbone-level segmentation, which enables strong zero-shot results without mask supervision, as well as a minimal and efficient standalone model design for downstream segmentation tasks. Our project page is https://research.nvidia.com/labs/dvl/projects/native-segmentation.

Figures (6)

• Figure 1: Downsampling in vision backbones via uniform downsampling (top) v.s. learned downsampling (bottom, this work): Vision backbones downsample feature maps using uniform-grid operators (e.g., pooling, top) and rely on uniform upsampling (e.g., bilinear interpolation, top) for image segmentation tasks. Our new backbone with spatial grouping layers learns to map pixels to a reduced set of tokens, aligning with image boundaries during downsampling (bottom). This enables scalable backbone-level native segmentation, i.e., without the need for dedicated segmentation heads.
• Figure 2: Overall model design. Visualization of our hierarchical architecture and its key components. (a) Our backbone architecture consists of four processing stages interconnected by content-aware grouping layers for downsampling. (b) Core operations of our Spatial Grouping Layer, which computes soft token assignments and updates group features iteratively (detailed in \ref{['alg:group_layer']}). (c) The composition of learned assignment matrices across grouping layers in consecutive backbone stages enables principled feature upsampling.
• Figure 3: Segmentation emerges from ImageNet pre-training. We visualize group decompositions across each backbone stage, along with their upsampled activations over the predicted class. We observe that even in the absence of mask supervision, super-pixel-like structures emerge in earlier layers, and are eventually grouped into semantically coherent regions in dense grouping layers.
• Figure 4: Qualitative zero-shot segmentation learned from image-text contrastive pre-training. We visualize hierarchical final decompositions along with their predicted semantic masks, obtained in a zero-shot setting on Pascal VOC validation images Everingham10IJCV, and corresponding ground truth masks. Note that these models did not receive any form of mask supervision during training, and were trained with a standard contrastive objective on image-text pairs. Final masks are obtained without any form of heuristic postprocessing.
• Figure : Grouping layer over an input feature map $X$ for $L$ iterations with sparsity.