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 proposes a novel transformer-based unsupervised feature distillation method for semantic segmetation- STEGO (Self-supervised Transformer with Energy-based Graph Optimization). The pipeline has two separate stages, feature learning and cluster compactification (feature distillation). The feature network is DINO (Caron et al., 2021) - pretrained, frozen.   The main contribution is the cluster compactification(distillation) network. It features a novel 3-term loss function: Knn, self correltation and random image corelation. The authors also propose a number of strategies for performance improvement: clamping the segmentation feature correspondence tensor at 0, five-cropping the dataset (corners+center), spatial centering on the feature correspondences and conditional random fields for better segmentation at the object edges. Training the segmentation head is fast (a couple of hours on a recent GPU with the ViT-small transformer backbone) and the results are state-of-the art on semantic segmentation on CocoStuff and Cityscapes (+14mIou@CocoStuff, +7mIoU@Citiscapes).