Unsupervised Semantic Segmentation by Distilling Feature Correspondences

1 INTRODUCTION . Semantic segmentation is the process of classifying each individual pixel of an image into a known ontology . Though semantic segmentation models can detect and delineate objects at a much finer granularity than classification or object detection systems , these systems are hindered by the difficulties of creating labelled training data . In particular , segmenting an image can take over 100× more effort for a human annotator than classifying or drawing bounding boxes ( Zlateski et al. , 2018 ) . Furthermore , in complex domains such as medicine , biology , or astrophysics , ground-truth segmentation labels may be unknown , ill-defined , or require considerable domain-expertise to provide ( Yu et al. , 2018 ) . Recently , several works introduced semantic segmentation systems that could learn from weaker forms of labels such as classes , tags , bounding boxes , scribbles , or point annotations ( Ren et al. , 2020 ; Pan et al. , 2021 ; Liu et al. , 2020 ; Bilen et al. ) . However , comparatively few works take up the challenge of semantic segmentation without any form of human supervision or motion cues . Attempts such as Independent Information Clustering ( IIC ) ( Ji et al. , 2019 ) and PiCIE ( Cho et al. , 2021 ) aim to learn semantically meaningful features through transformation equivariance , while imposing a clustering step to improve the compactness of the learnt features . In contrast to these previous methods , we utilize pre-trained features from unsupervised feature learning frameworks and focus on distilling them into a compact and discrete structure while preserving their relationships across the image corpora . This is motivated by the observation that correlations between unsupervised features , such as ones learnt by DINO ( Caron et al. , 2021 ) , are already semantically consistent , both within the same image and across image collections . As a result , we introduce STEGO ( Self-supervised Transformer with Energy-based Graph Optimization ) , which is capable of jointly discovering and segmenting objects without human supervision . STEGO distills pretrained unsupervised visual features into semantic clusters using a novel contrastive loss . STEGO dramatically improves over prior art and is a considerable step towards closing the gap with supervised segmentation systems . We include a short video detailing the work at https : //tinyurl.com/wrkx4rz9 . Specifically , we make the following contributions : • Show that unsupervised deep network features have correlation patterns that are largely consistent with true semantic labels . • Introduce STEGO , a novel transformer-based architecture for unsupervised semantic segmentation . • Demonstrate that STEGO achieves state of the art performance on both the CocoStuff ( +14 mIoU ) and Cityscapes ( +9 mIoU ) segmentation challenges . • Justify STEGO ’ s design with an ablation study on the CocoStuff dataset . 2 RELATED WORK . Self-supervised Visual Feature Learning Learning meaningful visual features without human annotations is a longstanding goal of computer vision . Approaches to this problem often optimize a surrogate task , such as denoising ( Vincent et al. , 2008 ) , inpainting ( Pathak et al. , 2016 ) , jigsaw puzzles , colorization ( Zhang et al. , 2017 ) , rotation prediction ( Gidaris et al. , 2018 ) , and most recently , contrastive learning over multiple augmentations ( Hjelm et al. , 2018 ; Chen et al. , 2020a ; a ; c ; Oord et al. , 2018 ) . Contrastive learning approaches , whose performance surpass all other surrogate tasks , assume visual features are invariant under a certain set of image augmentation operations . These approaches maximize feature similarities between an image and its augmentations , while minimizing similarity between negative samples , which are usually randomly sampled images . Some notable examples of positive pairs include temporally adjacent images in videos ( Oord et al. , 2018 ) , image augmentations ( Chen et al. , 2020a ; c ) , and local crops of a single image ( Hjelm et al. , 2018 ) . Many works highlight the importance of large numbers of negative samples during training . To this end Wu et al . ( 2018 ) propose keeping a memory bank of negative samples and Chen et al . ( 2020c ) propose momentum updates that can efficiently simulate large negative batch sizes . Recently some works have aimed to produce spatially dense feature maps as opposed to a single global vector per image . In this vein , VADeR ( Pinheiro et al. , 2020 ) contrasts local per-pixel features based on random compositions of image transformations that induce known correspondences among pixels which act as positive pairs for contrastive training . Instead of trying to learn visual features and clustering from scratch , STEGO treats pretrained self-supervised features as input and is agnostic to the underlying feature extractor . This makes it easy to integrate future advances in self-supervised feature learning into STEGO . Unsupervised Semantic Segmentation Many unsupervised semantic segmentation approaches use techniques from self-supervised feature learning . IIC ( Ji et al. , 2019 ) maximizes mutual information of patch-level cluster assignments between an image and its augmentations . Contrastive Clustering ( Li et al. , 2020 ) , and SCAN ( Van Gansbeke et al. , 2020 ) improve on IIC ’ s image clustering results with supervision from negative samples and nearest neighbors but do not attempt semantic segmentation . PiCIE ( Cho et al. , 2021 ) improves on IIC ’ s semantic segmentation results by using invariance to photometric effects and equivariance to geometric transformations as an inductive bias . In PiCIE , a network minimizes the distance between features under different transformations , where the distance is defined by an in-the-loop k-means clustering process . SegSort ( Hwang et al. , 2019 ) adopts a different approach . First , SegSort learns good features using superpixels as proxy segmentation maps , then uses Expectation-Maximization to iteratively refine segments over a spherical embedding space . In a similar vein , MaskContrast ( Van Gansbeke et al. , 2021 ) achieves promising results on PascalVOC by first using an off-the-shelf saliency model to generate a binary mask for each image . MaskContrast then contrasts learned features within and across the saliency masks . In contrast , our method focuses refining existing pretrained self-supervised visual features to distill their correspondence information and encourage cluster formation . This is similar to the work of Collins et al . ( 2018 ) who show that low rank factorization of deep network features can be useful for unsupervised co-segmentation . We are not aware of any previous work that achieves the goal of high-quality , pixel-level unsupervised semantic segmentation on large scale datasets with diverse images . Visual Transformers Convolutional neural networks ( CNNs ) have long been state of the art for many computer vision tasks , but the nature of the convolution operator makes it hard to model longrange interactions . To circumvent such shortcomings , Wang et al . ( 2018 ) ; Zhang et al . ( 2019 ) use self-attention operations within a CNN to model long range interactions . Transformers ( Vaswani et al. , 2017 ) , or purely self-attentive networks , have made significant progress in NLP and have recently been used for many computer vision tasks ( Dosovitskiy et al. , 2020 ; Touvron et al. , 2021 ; Ranftl et al. , 2021 ; Caron et al. , 2021 ) . Visual Transformers ( ViT ) ( Vaswani et al. , 2017 ) apply self-attention mechanisms to image patches and positional embeddings in order to generate features and predictions . Several modifications of ViT have been proposed to improve supervised learning , unsupervised learning , multi-scale processing , and dense predictions . In particular , DINO ( Caron et al. , 2021 ) uses a ViT within a self-supervised learning framework that performs self-distillation with exponential moving average updates . Caron et al . ( 2021 ) show that DINO ’ s class-attention can produce localized and semantically meaningful salient object segmentations . Our work shows that DINO ’ s features not only detect salient objects but can be used to extract dense and semantically meaningful correspondences between images . In STEGO , we refine the features of this pre-trained backbone to yield semantic segmentation predictions when clustered . We focus on DINO ’ s embeddings because of their quality but note that STEGO can work with any deep network features . 3 METHODS . 3.1 FEATURE CORRESPONDENCES PREDICT CLASS CO-OCCURRENCE . Recent progress in self-supervised visual feature learning has yielded methods with powerful and semantically relevant features that improve a variety of downstream tasks . Though most works aim to generate a single vector for an image , many works show that intermediate dense features are semantically relevant ( Hamilton et al. , 2021 ; Collins et al. , 2018 ; Zhou et al. , 2016 ) . To use this information , we focus on the “ correlation volume ” ( Teed & Deng , 2020 ) between the dense feature maps . For convolutional or transformer architectures , these dense feature maps can be the activation map of a specific layer . Additionally , the Q , K or V matrices in transformers can also serve as candidate features , though we find these attention tensors do not perform as well in practice . More formally , let f ∈ RCHW , g ∈ RCIJ be the feature tensors for two different images where C represents the channel dimension and ( H , W ) , ( I , J ) represent spatial dimensions . We form the feature correspondence tensor : Fhwij : = ∑ c fchw |fhw| gcij |gij | , ( 1 ) whose entries represent the cosine similarity between the feature at spatial position ( h , w ) of feature tensor f and position ( i , j ) of feature tensor g. In the special case where f = g these correspondences measure the similarity between two regions of the same image . We note that this quantity appears often as the “ cost-volume ” within the optical flow literature , and Hamilton et al . ( 2021 ) show this acts a higher-order generalization of Class Activation Maps ( Zhou et al. , 2016 ) for contrastive architectures and visual search engines . By examining slices of the correspondence tensor , F , at a given ( h , w ) we are able to visualize how two images relate according the featurizer . For example , Figure 2 shows how three different points from the source image ( shown in blue , red , and green ) are in correspondence with relevant semantic areas within the image and its K-nearest neighbors with respect to the DINO ( Caron et al. , 2021 ) as the feature extractor . This feature correspondence tensor not only allows us to visualize image correspondences but is strongly correlated with the true label co-occurrence tensor . In particular , we can form the ground truth label co-occurrence tensor given a pair of ground-truth semantic segmentation labels k ∈ CHW , l ∈ CIJ where C represents the set of possible classes : Lhwij : = { 1 , if lhw = kij 0 , if lhw 6= kij By examining how well the feature correspondences , F , predict the ground-truth label cooccurrences , L , we can measure how compatible the features are with the semantic segmentation labels . More specifically we treat the feature correspondences as a probability logit and compute the average precision when used as a classifier for L. This approach not only acts as a quick diagnostic tool to determine the efficacy of features , but also allows us to compare with other forms of supervision such as the fully connected Conditional Random Field ( CRF ) ( Krähenbühl & Koltun , 2011 ) , which uses correspondences between pixels to refine low-resolution label predictions . In Figure 3 we plot precision-recall curves for the DINO backbone , the MoCoV2 backbone , the CRF Kernel , and our trained STEGO architecture . Interestingly , we find that DINO is already a spectacular predictor of label co-occurrence within the Coco stuff dataset despite never seeing the labels . In particular , DINO recalls 50 % of true label co-occurrences with a precision of 90 % and significantly outperforms both MoCoV2 feature correspondences and the CRF kernel . One curious note is that our final trained model is a better label predictor than the supervisory signal it learns from . We attribute this to the distillation process discussed in Section 3.2 which amplifies this supervisory signal and drives consistency across the entire dataset . Finally , we stress that our comparison to ground truth labels within this section is solely to provide intuition about the quality of feature correspondences as a supervisory signal . We do not use the ground truth labels to tune any parameters of STEGO .

Summary: The paper introduces STEGO (Self-supervised Transformer with Energy-based Graph Optimization), a novel feature correlation refinement method that builds on top of modern self-supervised visual backbones (visual-transformers) that generate dense semantically-correlated features in an effort to improve scene semantic segmentation without any type of labels (unsupervised). Different from previous works, the authors decouple the feature learning from cluster compactification and introduce a novel contrastive loss function (a combination of three correlation factors: KNN, self and random images) in order to further constrain the features to form compact semantic clusters without damaging their consistency throughout the dataset. The authors demonstrate state-of-the-art results on two popular benchmarks CocoStuff and Cityscapes. The results are quite striking in both a quantitative (method improves segmentation by a large margin (+14mIoU and +9mIoU) compared to recently published work) and also a qualitative manner (clearly distinguish the structure of the objects in the scene).