Unsupervised Semantic Segmentation by Distilling Feature Correspondences

Mark Hamilton; Zhoutong Zhang; Bharath Hariharan; Noah Snavely; William T. Freeman

Source-linked AI summary

Unsupervised Semantic Segmentation by Distilling Feature Correspondences

Mark Hamilton, Zhoutong Zhang, Bharath Hariharan, Noah Snavely, William T. Freeman

arXiv:2203.08414v1Mar 16, 2022cs.CVcs.AIcs.LGstat.ML

TL;DR

Unsupervised semantic segmentation seeks meaningful pixel-level categories without annotations, requiring features that are both semantically meaningful and compact. STEGO separates feature learning from cluster compactification by distilling semantically consistent unsupervised features with a contrastive loss. It reports state-of-the-art improvements on CocoStuff and Cityscapes, including +14 mIoU and +9 mIoU respectively.

Problem
Unsupervised semantic segmentation must produce semantically meaningful, compact pixel features and discrete labels without any annotation.
Method
STEGO distills pretrained unsupervised visual features into semantic clusters using a transformer-based framework and a contrastive loss that preserves feature correspondences.
Results
+14 unsupervised mIoU on CocoStuff and +8.7 unsupervised mIoU on Cityscapes are reported over the next-best baseline.
Takeaways & Limitations
Unsupervised feature correspondences can provide a learning signal for accurate semantic segmentation without fine-tuning the backbone on the target datasets.
Takeaways & Limitations
Evaluation is difficult when label ontologies contain arbitrary or poorly defined class boundaries, such as ceilings versus walls or indoor-related classes.

Abstract

from arXiv · show

Unsupervised semantic segmentation aims to discover and localize semantically meaningful categories within image corpora without any form of annotation. To solve this task, algorithms must produce features for every pixel that are both semantically meaningful and compact enough to form distinct clusters. Unlike previous works which achieve this with a single end-to-end framework, we propose to separate feature learning from cluster compactification. Empirically, we show that current unsupervised feature learning frameworks already generate dense features whose correlations are semantically consistent. This observation motivates us to design STEGO ($\textbf{S}$elf-supervised $\textbf{T}$ransformer with $\textbf{E}$nergy-based $\textbf{G}$raph $\textbf{O}$ptimization), a novel framework that distills unsupervised features into high-quality discrete semantic labels. At the core of STEGO is a novel contrastive loss function that encourages features to form compact clusters while preserving their relationships across the corpora. STEGO yields a significant improvement over the prior state of the art, on both the CocoStuff ($\textbf{+14 mIoU}$) and Cityscapes ($\textbf{+9 mIoU}$) semantic segmentation challenges.

1 INTRODUCTION

Unsupervised semantic segmentation must assign meaningful categories to pixels without human labels, despite the high cost and expertise required for segmentation annotation. STEGO separates feature learning from cluster compactification by distilling semantically consistent unsupervised features into discrete semantic clusters.

Motivation: Semantic segmentation classifies every image pixel into a known ontology, but creating segmentation labels can require over 100× more effort than classification or bounding-box annotation.Labels may also be unknown, ill-defined, or require substantial domain expertise in complex fields.
Motivation: Few prior methods attempt semantic segmentation without human supervision or motion cues, with IIC and PiCIE learning semantic features through transformation equivariance and clustering.
Approach: Unsupervised feature correlations, including those learned by DINO, are already semantically consistent within images and across image collections.This observation motivates using pretrained unsupervised features as the basis for segmentation.
Approach: STEGO distills pretrained unsupervised visual features into semantic clusters with a novel contrastive loss while jointly discovering and segmenting objects without human supervision.
Results: STEGO predictions are described as more consistent and detailed than PiCIE, including better resolution of fine-grained objects and fewer omitted key objects.Examples include horse legs and individual birds in CocoStuff images.

2 RELATED WORK

Related work includes self-supervised feature learning, unsupervised semantic segmentation, and visual transformers. These approaches use surrogate tasks, invariance or correspondence objectives, and self-attention to learn useful visual representations and model long-range interactions.

Self-supervised Visual Feature Learning: Self-supervised visual feature learning seeks meaningful representations without human annotations, often through surrogate tasks such as denoising, inpainting, colorization, rotation prediction, or contrastive learning.
Self-supervised Visual Feature Learning: Contrastive approaches maximize similarity between images and augmentations while minimizing similarity to negative samples, often relying on many negatives.Positive pairs can include adjacent video frames, image augmentations, and local crops.
Unsupervised Semantic Segmentation: IIC maximizes mutual information between patch-level cluster assignments across image augmentations, while PiCIE improves semantic segmentation using photometric and geometric invariance.
Unsupervised Semantic Segmentation: Contrastive Clustering and SCAN improve IIC’s image clustering with negative-sample and nearest-neighbor supervision but do not attempt semantic segmentation.
Visual Transformers: CNN-based systems use self-attention to model long-range interactions, while visual transformers rely on purely self-attentive networks for computer vision tasks.

3 METHODS

STEGO treats dense unsupervised feature correspondences as a semantic learning signal, then distills them into compact pixel-wise embeddings through a frozen-backbone segmentation head. Its correspondence loss combines image-pair relationships with stabilization mechanisms to produce clustered semantic representations.

Feature correspondences: Dense feature maps support a correlation volume whose entries measure cosine similarity between spatial positions in two feature tensors.When the tensors are identical, these correspondences measure similarity between regions of the same image.
Feature correspondences: DINO feature correspondences strongly predict true label co-occurrences, recalling 50% at 90% precision and outperforming MoCoV2 correspondences and a CRF kernel.The paper evaluates correspondences as probability logits and computes average precision for predicting ground-truth label co-occurrences.
Distilling correspondences: STEGO freezes a backbone and trains a lightweight segmentation head that projects its features into a lower-dimensional code space intended to form compact clusters.The head maps C-dimensional backbone features into K dimensions, with K < C.
Distilling correspondences: The correspondence loss pushes segmentation features together when source feature pairs are strongly coupled, using element-wise multiplication of feature and segmentation correlations.The loss applies attractive or repulsive forces whose strength is proportional to feature correspondences.
Distilling correspondences: Zero-clamping and spatial centering address instability from anti-alignment and imbalanced correlation patterns, with their effects evaluated in a CocoStuff ablation.Weakly correlated features are optimized toward orthogonality rather than total anti-alignment, while spatial centering balances the signal for small objects.
STEGO architecture: STEGO applies correspondence losses to self-images, nearest neighbors, and random image pairs, yielding mainly attractive signals from the first two and repulsive signals from random pairs.The architecture uses these loss instantiations to distill feature relationships into the segmentation head.

4 EXPERIMENTS

STEGO is evaluated on CocoStuff, Cityscapes, and related settings using linear-probe and unsupervised clustering metrics. It substantially outperforms prior methods, while qualitative and confusion-matrix analyses reveal both fine-detail strengths and persistent category ambiguities.

Evaluation setup: STEGO is evaluated on 27-class CocoStuff and Cityscapes settings using mIoU and Accuracy, with an additional Potsdam-3 comparison.Validation images are resized to 320 pixels along the minor axis and center-cropped to 320 × 320; the Cityscapes setting has fewer available baselines.
CocoStuff results: +14 unsupervised mIoU and +6.9 unsupervised accuracy over the next-best baseline are reported on CocoStuff.STEGO also improves by +26 linear-probe mIoU and +21 linear-probe accuracy on the same benchmark.
Cityscapes results: +8.7 unsupervised mIoU and +7.7 unsupervised accuracy are reported on the Cityscapes validation set.The reported results use DINO self-supervised ImageNet weights without fine-tuning the backbone for these datasets.
Qualitative analysis: STEGO captures small objects and fine details more effectively than PiCIE in qualitative CocoStuff examples.The paper attributes this in part to higher-resolution DINO features, 5-crop training, and CRF post-processing that aligns predictions to image edges.
Failure analysis: The confusion analysis identifies errors between food categories, ceilings and walls, and several inconsistently segmented classes.The authors also note that arbitrary label ontologies make some unsupervised segmentation boundaries difficult to define and evaluate.
Ablation study: Ablations show that ViT-Base is the strongest tested feature extractor and that the evaluated architectural modifications improve cluster and linear-probe metrics.The modifications include 0-Clamp, 5-Crop, spatial centering, and CRF post-processing.

5 CONCLUSION

The paper concludes that modern self-supervised visual backbones can support state-of-the-art unsupervised semantic segmentation when their feature correspondences are distilled into compact representations. STEGO uses a contrastive loss, achieves large gains on CocoStuff and Cityscapes, and is supported by an ablation study.

Conclusion: Modern self-supervised visual backbones can be refined into state-of-the-art unsupervised semantic segmentation methods.The conclusion links this capability to correspondences between deep features and ground-truth label co-occurrence.
Conclusion: STEGO distills feature correspondences with a novel contrastive loss into low-rank representations that cluster into semantic segmentation predictions.The loss is connected to maximum-likelihood estimation in Potts models over the dataset's pixels.
Conclusion: +14 mIoU on CocoStuff and +9 mIoU on Cityscapes summarize STEGO's improvement over prior state of the art.The paper also justifies its architectural decisions with an ablation study on CocoStuff.

A.1 VIDEO AND CODE

The appendix provides a short video description of the work and releases training and evaluation code.

Video and code: A short video description of the work is provided online.The video is linked at https://aka.ms/stego-video.
Video and code: Training and evaluation code are provided online.The code is linked at https://aka.ms/stego-code.

A.2 ADDITIONAL RESULTS ON THE POTSDAM-3 DATASET

On the Potsdam-3 aerial image segmentation task, STEGO is compared with prior art and achieves a reported accuracy improvement over IIC. Qualitative examples are included for this setting.

Additional results: +12% accuracy over IIC is reported on the Potsdam 3-class aerial image segmentation task.The comparison follows the Potsdam-3 setting presented in Ji et al. (2019).
Additional results: Qualitative STEGO segmentation results for Potsdam-3 are shown in Figure 7.The figure is described as a qualitative comparison on the Potsdam-3 segmentation challenge.

A.3 ADDITIONAL ABLATION STUDY

The additional ablation study examines STEGO’s architectural components on CocoStuff, finding that removing most components hurts performance across most metrics.

Additional architecture ablation study: Removing each major architectural component hurts performance in most metrics.The study isolates architectural decisions to assess their individual effects.

A.4 ADDITIONAL QUALITATIVE RESULTS

Additional qualitative predictions compare STEGO with PiCIE on the CocoStuff 27-class segmentation challenge using uncurated images.

Additional qualitative results: STEGO and PiCIE are compared on additional CocoStuff 27-class segmentation predictions.The images are not curated.

A.5 FAILURE CASES

STEGO’s unsupervised segmentations fail in several recurring visual and ontology-boundary cases, including backgrounds, object boundaries, food categories, and small people.

Failure cases: STEGO improperly separates ground from trees and backgrounds.This failure appears in the first column of Figure 9.
Failure cases: STEGO assigns a barn floor to “outdoor” and a barn wall to “building.”PiCIE makes the same type of barn error.
Failure cases: STEGO misses the boundary between a wall and ceiling.This is shown in the third column of Figure 9.
Failure cases: STEGO struggles to distinguish food (thing) from food (stuff).The fourth column demonstrates this characterization challenge, which PiCIE also shares.
Failure cases: STEGO misses a human in the lower-left of a challenging grayscale image.The passage attributes the difficulty probably to the person being hard to spot.

A.6 FEATURE CORRESPONDENCES PREDICT STEGO’S ERRORS

The analysis finds that DINO feature correspondences predict many STEGO errors, linking correspondence structure to confusion patterns while framing the analysis through graph-based optimization.

Feature correspondences and errors: DINO feature correspondences serve as an excellent proxy for true label co-occurrence information.The section uses this relationship to analyze where correspondence structure differs from ground-truth labels.
Feature correspondences and errors: DINO image-to-KNN correspondences predict many areas where STEGO fails on CocoStuff27.Thresholded strongest couplings are evaluated against whether paired objects share the same class.
Feature correspondences and errors: Food (things) versus food (stuff), appliance versus furniture, and window versus wall are confused in both DINO analysis and STEGO predictions.The shared confusion appears in the CocoStuff correspondence and STEGO confusion visualizations.
Feature correspondences and errors: Many STEGO errors originate from DINO feature structure rather than other architectural aspects.Whether the source is DINO itself or ambiguity in the CocoStuff ontology remains unresolved.
Feature correspondences and errors: The correspondence analysis can predict a fully trained STEGO confusion matrix and may help select better backbones without training STEGO.The paper presents this as a potential use of the analysis.
Graph energy interpretation: Maximum-likelihood estimation connects STEGO’s correlation-distillation loss to optimization on an undirected graphical or Ising model.The appendix develops this connection using a Potts-model energy formulation.
Graph energy interpretation: The graph’s nodes represent spatial locations across all images, with segmentation-head outputs serving as node functions and cosine distance as compatibility.The formulation rewrites the graph objective using image pairs and spatial-location pairs.
Graph energy interpretation: Gradient descent makes optimization tractable for massive datasets, unlike spectral methods requiring O(|V|^3) time for graph-Laplacian eigenvalues.In practice, the minimization is approximated with minibatch SGD, while KNN and self-correspondence distillation alter the weights but not the functional form.