Year: 2021

Abstract:
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 (Self-supervised Transformer with Energy-based Graph Optimization), a novel framework that distills unsupervised features into highquality 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 (+14 mIoU) and Cityscapes (+9 mIoU) semantic segmentation challenges. Figure 1: Unsupervised semantic segmentation predictions on the CocoStuff (Caesar et al., 2018) 27 class segmentation challenge. Our method, STEGO, does not use labels to discover and segment consistent objects. Unlike the prior state of the art, PiCIE (Cho et al., 2021), STEGO’s predictions are consistent, detailed, and do not omit key objects.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 METHODS

heading: 3.1 FEATURE CORRESPONDENCES PREDICT CLASS CO-OCCURRENCE
text: 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.

heading: 3.2 DISTILLING FEATURE CORRESPONDENCES
text: In Section 3.1 we have shown that feature correspondences have the potential to be a quality learning signal for unsupervised segmentation. In this section we explore how to harness this signal to create pixel-wise embeddings that, when clustered, yield a quality semantic segmentation. In particular, we seek to learn a low-dimensional embedding that “distills” the feature correspondences. To achieve this aim, we draw inspiration from the CRF which uses an undirected graphical model to refine noisy
or low-resolution class predictions by aligning them with edges and color-correlated regions in the original image.
More formally, let N : RC′H′W ′ → RCHW represent a deep network backbone, which maps an image x with C ′ channels and spatial dimensions (H ′,W ′) to a feature tensor f with C channels and spatial dimensions (H,W ). In this work, we keep this backbone network frozen and focus on training a light-weight segmentation head S : RCHW → RKHW , that maps our feature space to a code space of dimension K, where K < C. The goal of S is to learn a nonlinear projection, S(f) =: s ∈ RKHW , that forms compact clusters and amplifies the correlation patterns of f . To build our loss function let f and g be two feature tensors from a pair of images x, and y and let s := S(f) ∈ RCHW and t := S(g) ∈ RCIJ be their respective segmentation features. Next, using Equation 1 we compute a feature correlation tensor F ∈ RHWIJ from f and g and a segmentation correlation tensor S ∈ RHWIJ from s and t. Our loss function aims to push the entries of s and t together if there is a significant coupling between two corresponding entries of f and g. As shown in Figure 4, we can achieve this with a simple element-wise multiplication of the tensors F and S:
Lsimple−corr (x, y, b) := − ∑ hwij (Fhwij − b)Shwij (2)
Where b is a hyper-parameter which adds uniform “negative pressure” to the equation to prevent collapse. Minimizing L with respect to S encourages elements of S to be large when elements of F − b are positive and small when elements of F − b are negative. More explicitly, because the elements of F and S are cosine similarities, this exerts an attractive or repulsive force on pairs of segmentation features with strength proportional to their feature correspondences. We note that the elements of S are not just encouraged to equal the elements of F but rather to push to total anti-alignment (−1) or alignment (1) depending on the sign of F − b. In practice, we found that Lsimple−corr is sometimes unstable and does not provide enough learning signal to drive the optimization. Empirically, we found that optimizing the segmentation features towards total anti-alignment when the corresponding features do not correlate leads to instability, likely because this increases co-linearity. Therefore, we optimize weakly-correlated segmentation features to be orthogonal instead. This can be efficiently achieved by clamping the segmentation correspondence, S, at 0, which dramatically improved the optimization stability.
Additionally, we encountered challenges when balancing the learning signal for small objects which have concentrated correlation patterns. In these cases, Fhwij − b is negative in most locations, and the loss drives the features to diverge instead of aggregate. To make the optimization more balanced, we introduce a Spatial Centering operation on the feature correspondences:
FSChwij := Fhwij − 1
IJ ∑ i′j′ Fhwi′j′ . (3)
Together with the zero clamping, our final correlation loss is defined as: Lcorr(x, y, b) := − ∑ hwij (FSChwij − b)max(Shwij , 0). (4)
We demonstrate the positive effect of both the aforementioned “0-Clamp” and “SC” modifications in the ablation study of Table 2.

heading: 3.3 STEGO ARCHITECTURE
text: STEGO uses three instantiations of the correspondence loss of Equation 4 to train a segmentation head to distill feature relationships between an image and itself, its K-Nearest Neighbors (KNNs), and random other images. The self and KNN correspondence losses primarily provide positive, attractive, signal and random image pairs tend to provide negative, repulsive, signal. We illustrate this and other major architecture components of STEGO in Figure 4.
STEGO is made up of a frozen backbone that serves as a source of learning feedback, and as an input to the segmentation head for predicting distilled features. This segmentation head is a simple feed forward network with ReLU activations (Glorot et al., 2011). In contrast to other works, our
method does not re-train or fine-tune the backbone. This makes our method very efficient to train: it only takes less than 2 hours on a single NVIDIA V100 GPU card.
We first use our backbone to extract global image features by global average pooling (GAP) our spatial features: GAP (f). We then construct a lookup table of each image’s K-Nearest Neighbors according to cosine similarity in the backbone’s feature space. Each training minibatch consists of a collection of random images x and random nearest neighbors xknn. In our experiments we sample xknn randomly from each image’s top 7 KNNs. We also sample random images, xrand, by shuffling x and ensuring that no image matched with itself. STEGO’s full loss is:
L = λselfLcorr(x, x, bself ) + λknnLcorr(x, xknn, bknn) + λrandLcorr(x, xrand, brand) (5)
Where the λ’s and the b’s control the balance of the learning signals and the ratio of positive to negative pressure respectively. In practice, we found that a ratio of λself ≈ λrand ≈ 2λknn worked well. The b parameters tended to be dataset and network specific, but we aimed to keep the system in a rough balance between positive and negative forces. More specifically we tuned the bs to keep mean KNN feature similarity at ≈ 0.3 and mean random similarity at ≈ 0.0. Many images within the CocoStuff and Cityscapes datasets are cluttered with small objects that are hard to resolve at a feature resolution of (40, 40). To better handle small objects and maintain fast
training times we five-crop training images prior to learning KNNs. This not only allows the network to look at closer details of the images, but also improves the quality of the KNNs. More specifically, global image embeddings are computed for each crop. This allows the network to resolve finer details and yields five times as many images to find close matching KNNs from. Five-cropping improved both our Cityscapes results and CocoStuff segmentations, and we detail this in Table 2.
The final components of our architecture are the clustering and CRF refinement step. Due to the feature distillation process, STEGO’s segmentation features tend to form clear clusters. We apply a cosine distance based minibatch K-Means algorithm (MacQueen et al., 1967) to extract these clusters and compute concrete class assignments from STEGO’s continuous features. After clustering, we refine these labels with a CRF to improve their spatial resolution further.

heading: 3.4 RELATION TO POTTS MODELS AND ENERGY-BASED GRAPH OPTIMIZATION
text: Equation 4 can be viewed in the context of Potts models or continuous Ising models from statistical physics (Potts, 1952; Baker Jr & Kincaid, 1979). We briefly overview this connection, and point interested readers to Section A.7 for a more detailed discussion. To build the general Ising model, let G = (V, w) be a fully connected, weighted, and undirected graph on |V| vertices. In our applications we take V to be the set of pixels in the training dataset. Let w : V × V → R represent an edge weighting function. Let φ : V → C be a vertex valued function mapping into a generic code space C such as the probability simplex over cluster labels P(L), or the K-dimensional continuous feature space RK . The function φ can be a parameterized neural network, or a simple lookup table that assigns a code to each graph node. Finally, we define a compatibility function µ : C × C → R that measures the cost of comparing two codes. We can now define the following graph energy functional:
E(φ) := ∑
vi,vj∈V w(vi, vj)µ(φ(vi), φ(vj)) (6)
Constructing the Boltzmann Distribution (Hinton, 2002) yields a normalized distribution over the function space Φ:
p(φ|w, µ) = exp(−E(φ))∫ Φ exp(−E(φ′))dφ′ (7)
In general, sampling from this probability distribution is difficult because of the often-intractable normalization factor. However, it is easier to compute the maximum likelihood estimate (MLE), arg maxφ∈Φ p(φ|w, µ). In particular, if Φ is a smoothly parameterized space of functions and φ and µ are differentiable functions, one can compute the MLE using stochastic gradient descent (SGD) with highly-optimized automatic differentiation frameworks (Paszke et al., 2019; Abadi et al., 2015). In Section A.7 of the supplement we prove that the finding the MLE of Equation 7 is equivalent to minimizing the loss of Equation 4 when |V | is the set of pixels in our image training set, φ = S ◦N , w is the cosine distance between features, and µ is cosine distance. Like STEGO, the CRF is also a Potts model, and we use this connection to re-purpose the STEGO loss function to create continuous, minibatch, and unsupervised variants of the CRF. We detail this exploration in Section A.8 of the Supplement.

heading: 4 EXPERIMENTS
text: We evaluate STEGO on standard semantic segmentation datasets and compare with current state-ofthe-art. We then justify different design choices of STEGO through ablation studies. Additional details on datasets, model hyperparameters, hardware, and other implementation details can be found in Section A.9 of the Supplement.

heading: 4.1 EVALUATION DETAILS
text: Datasets Following Cho et al. (2021), we evaluate STEGO on the 27 mid-level classes of the CocoStuff class hierarchy and on the 27 classes of Cityscapes. Like prior art, we first resize images to 320 pixels along the minor axis followed by a (320× 320) center crops of each validation image. We use mean intersection over union (mIoU) and Accuracy for evaluation metrics. Our CocoStuff evaluation setting originated in Ji et al. (2019) and is common in the literature. Our Cityscapes evaluation setting is adopted from Cho et al. (2021). The latter is newer and more challenging, and thus fewer baselines are available.
Linear Probe One way we evaluate the quality of the distilled segmentation features is to measure transfer learning effectiveness. As in Van Gansbeke et al. (2021); Cho et al. (2021); Chen et al. (2020b), we train a linear projection from segmentation features to class labels using the cross entropy loss. Note that this is solely for evaluating feature quality and is not part of the STEGO training process.
Clustering Unlike the linear probe, the clustering step does not have access to ground truth supervised labels. As in prior art, we use a Hungarian matching algorithm to align our unlabeled clusters and the ground truth labels for evaluation and visualization purposes. This measures how consistent the predicted semantic segments are with the ground truth labels and is invariant to permutations of the predicted class labels.

heading: 4.2 RESULTS
text: We summarize our main results on the 27 classes of CocoStuff in Table 1. STEGO significantly outperforms the prior state of the art, PiCIE, on both linear probe and clustering (Unsupervised) metrics. In particular, STEGO improves by +14 unsupervised mIoU, +6.9 unsupervised accuracy, +26 linear probe mIoU, and +21 linear probe accuracy compared to the next best baseline. In Table 3, we find a similarly large improvement of +8.7 unsupervised mIoU and +7.7 unsupervised accuracy on the Cityscapes validation set. These two experiments demonstrate that even though we do not fine-tune the backbone for these datasets, DINO’s self-supervised weights on ImageNet (Deng et al., 2009) are enough to simultaneously solve both settings. STEGO also outperforms simply clustering the features from unmodified DINO, MoCoV2, and ImageNet supervised ResNet50 backbones. This demonstrates the benefits of training a segmentation head to distill feature correspondences.
We show some example segmentations from STEGO and our baseline PiCIE on the CocoStuff dataset in Figure 1. We include additional examples and failure cases in Sections A.3 and A.4. We note that STEGO is significantly better at resolving fine-grained details within the images such as the legs of horses in the third image from the left column of Figure 1, and the individual birds in the right-most column. Though the PiCIE baseline uses a feature pyramid network to output high resolution predictions, the network does not attune to fine grained details, potentially demonstrating the limitations of the sparse training signal induced by data augmentations alone. In contrast, STEGO’s predictions capture small objects and fine details. In part, this can be attributed to DINO backbone’s higher resolution features, the 5-crop training described in 3.3, and the CRF post-processing which helps to align the predictions to image edges. We show qualitative results on the Cityscapes dataset in Figure 5. STEGO successfully identifies people, street, sidewalk, cars, and street signs with high detail and fidelity. We note that prior works did not publish pretrained models or linear probe results on Cityscapes so we exclude this information from Table 3 and Figure 5.
Table 2: Architecture ablation study on the CocoStuff Dataset (27 Classes).
Arch. 0- C
la m p 5- C ro p SC C R F
Unsup. Linear Probe Acc. mIoU Acc. mIoU
MoCoV2 X 48.4 20.8 70.7 26.5 ViT-S 34.2 7.3 54.9 15.6 ViT-S X 44.3 21.3 70.9 36.8 ViT-S X X 47.6 23.4 72.2 36.8 ViT-S X X X 47.7 24.0 72.9 38.4 ViT-S X X X X 48.3 24.5 74.4 38.3 ViT-B X X X 54.8 26.8 74.3 39.5 ViT-B X X X X 56.9 28.2 76.1 41.0
Table 3: Results on the Cityscapes Dataset (27 Classes). STEGO improves significantly over all baselines in both accuracy and mIoU.
Unsup. Model Acc. mIoU
IIC (Ji et al., 2019) 47.9 6.4 MDC (Cho et al., 2021) 40.7 7.1 PiCIE (Cho et al., 2021) 65.5 12.3
STEGO (Ours) 73.2 21.0
To better understand the predictions and failures of STEGO, we include confusion matrices for CocoStuff (Figure 6) and Cityscapes (Figure 10 of the Supplement). Some salient STEGO errors include confusing the “food” category from the CocoStuff “things”, and the “food” category from CocoStuff “stuff”. STEGO also does not properly separate “ceilings” from “walls”, and lacks consistent segmentations for classes such as “indoor”, “accessory”, “rawmaterial” and “textile”. These errors also draw our attention to the challenges of evaluating unsupervised segmentation methods: label ontologies can be arbitrary. In these circumstances the divisions between classes are not well defined and it is hard to imagine a system that can segment the results consistently without additional information. In these regimes, the linear probe provides a more important barometer for quality because the limited supervision can help disambiguate these cases. Nevertheless, we feel that there is still considerable progress to be made on the purely unsupervised benchmark, and that even with the improvements of STEGO there is still a measurable performance gap with supervised systems.

heading: 4.3 ABLATION STUDY
text: To understand the impact of STEGO’s architectural components we perform an ablation analysis on the CocoStuff dataset, and report the results in Table 2. We examine the effect of using several different backbones in STEGO including MoCoV2, the ViT-Small, and ViT-Base architectures of DINO. We find that ViT-Base is the best feature extractor of the group and leads by a significant margin both in terms of accuracy and mIoU. We also evaluate the several loss function and architecture decisions described in Section 3.3. In particular, we explore clamping the segmentation feature correspondence tensor at 0 to prevent the negative pressure from introducing co-linearity (0-Clamp), five-cropping the dataset prior to mining KNNs to improve the resolution of the learning signal (5- Crop), spatially centering the feature correspondence tensor to improve resolution of small objects (SC), and Conditional Random Field post-processing to refine predictions (CRF). We find that these modifications improve both the cluster and linear probe evaluation metrics.

heading: 5 CONCLUSION
text: We have found that modern self-supervised visual backbones can be refined to yield state of the art unsupervised semantic segmentation methods. We have motivated this architecture by showing that correspondences between deep features are directly correlated with ground truth label cooccurrence. We take advantage of this strong, yet entirely unsupervised, learning signal by introducing a novel contrastive loss that “distills” the correspondences between features. Our system, STEGO, produces low rank representations that cluster into accurate semantic segmentation predictions. We connect STEGO’s loss to CRF inference by showing it is equivalent to MLE in Potts models over the entire collection of pixels in our dataset. We show STEGO yields a significant improvement over the prior state of the art, on both the CocoStuff (+14 mIoU) and Cityscapes (+9 mIoU) semantic segmentation challenges. Finally, we justify the architectural decisions of STEGO with an ablation study on the CocoStuff dataset.

heading: A APPENDIX
text: A.1 VIDEO AND CODE
We include a short video description of our work at https://tinyurl.com/wrkx4rz9.
We also provide training and evaluation code at https://tinyurl.com/w435zpzc

heading: A.2 ADDITIONAL ABLATION STUDY
text: 0- C
la m p 5- C ro p Po in tw
is e
C R F Se
lfL
os s
K N
N -L
os s
R an
dL
os s
Unsupervised Linear Probe Backbone Acc. mIoU Acc. mIoU ViT-Small X X X X X X X 48.3 24.5 74.4 38.3 MoCoV2 X X X X X X X 43.1 19.6 65.9 26.0 ViT-Small X X X X X X 42.8 10.3 59.3 19.3 ViT-Small X X X X X X 48.0 23.1 73.9 38.9 ViT-Small X X X X X X 50.2 22.3 73.7 37.7 ViT-Small X X X X X X 47.7 24.0 72.9 38.4 ViT-Small X X X X X X 43.0 20.2 73.0 36.2 ViT-Small X X X X X X 47.0 22.2 74.0 37.7 ViT-Small X X X X X X 39.8 12.8 65.5 29.9

heading: A.3 ADDITIONAL QUALITATIVE RESULTS

heading: A.4 FAILURE CASES
text: Unsupervised Segmentation is prone to a variety of issues. We include some of the following to segmentations to demonstrate cases where STEGO breaks down. In the first column of Figure 8 we can see that STEGO improperly segments ground from trees and backgrounds. In the second column we see that STEGO makes an understandable error and assigns the barn floor to the “outdoor” class and the barn wall to the “building” class. In the third column STEGO misses the boundary between wall and ceiling. The fourth column demonstrates the challenge between food (thing) and food (stuff) characterization. Interestingly PiCIE makes the same type of error both here, and in the barn case. The last column shows an example of STEGO missing a human in the lower left. In this image it is challenging to spot the person, probably because it is grayscale.

heading: A.5 FEATURE CORRESPONDENCES PREDICT STEGO’S ERRORS

heading: A.6 HIGHER RESOLUTION CONFUSION MATRICES

heading: A.7 RELATIONSHIP WITH GRAPH ENERGY MINIMIZATION
text: In section 3.4 we briefly mention that STEGO’s feature correlation distillation loss defined in Equation 4 can be seen as a particular case of Maximum Likelihood (ML) estimation on a undirected graphical model or Ising model. In this section we demonstrate this connection in greater detail using the formalism defined in 3.4. In particular, we recall the energy for a Potts model:
E(φ) := ∑
vi,vj∈V w(vi, vj)µ(φ(vi), φ(vj)) (8)
We then construct the Boltzmann Distribution (Hinton, 2002) yields a normalized distribution over the function space Φ:
p(φ|w, µ) = exp(−E(φ))∫ Φ exp(−E(φ′))dφ′ (9)
In general, sampling from this probability distribution is difficult because of the often-intractable normalization factor. However, it is easier to compute the maximum likelihood estimate (MLE):
arg max φ∈Φ p(φ|w, µ) = arg max φ∈Φ
1 Z exp(−E(φ)) (10)
Where Z is the unknown constant normalization factor. Simplifying the right-hand side yields:
arg max φ∈Φ p(φ|w, µ) = arg min φ∈Φ E(φ) = arg min φ∈Φ ∑ vi,vj∈V w(vi, vj)µ(φ(vi), φ(vj)) (11)
We are now in the position to connect this to the STEGO loss function. First, we take our nodes V to be the set of all spatial locations across our entire dataset of images. For concreteness we can represent v ∈ V by the tuple (n, h,w) where h,w represent height and width n represents the image number. We now let φ(vi) be the output of the segmentation head, svi , at the image and spatial location vi. Using cosine distance, dcos(x, y) = 1 − x|x| y |y| as the compatibility function, µ, yields the following:
= arg min S ∑ vi,vj∈V −w(vi, vj) svi |svi | svj |svj |
(12)
Wherte the argmin now ranges over the parameters of the segmentation head S . We can now observe that the sum over all pairs vi, vj ∈ V can be written as a sum over pairs of images x, y ∈ X and pairs of spatial locations (h,w), (i, j) where we note that (i, j) in this context refers to the spatial coordinates of image y as in 3.1 and not the indices of the vertices.
= arg min S ∑ x,y∈X ∑ hwij −W (x, y)hwijS(x, y)hwij (13)
Where we define S(x, y) to be the segmentation feature correlation tensor for images x, y as defined in Section 3.2. Finally letting W (x, y)hwij = Fhwij − b we recover our loss:
arg max φ∈Φ p(φ|w, µ) = arg min S ∑ x,y∈X Lsimple−corr (x, y, b) (14)
Finally we note that in practice we approximate the minimization using minibatch SGD, and our inclusion of KNN and Self-correspondence distillation changes the weight function w, but does not change its functional form.
Switching to the ML formulation of this problem allows us to solve this optimization for φ by gradient descent on the parameters of the segmentation head, S, and makes this computationally tractable. For large image datasets that can contain millions of high-resolution images, the induced graph can contain billions of image locations. Other graph embedding and clustering approaches such as Spectral methods require solving for eigenvalues of the graph Laplacian, which can take O(|V|3) time (Yan et al., 2009). More recent attempts to accelerate Spectral clustering such as (Yan et al., 2009) and (Han & Filippone, 2017) further assume a “Nonparametric” structure on the function φ, where a separate cluster assignment is learned for each vertex. This assumption of a “nonparametric” function φ can be undesirable as one cannot cluster or embed new data without recomputing the entire clustering. In contrast, STEGO’s backbone and segmentation head act as a parametric form for the function φ allowing the approach to output predictions for novel images.

heading: A.8 CONTINUOUS, UNSUPERVISED, AND MINI-BATCH CRF
text: Fully connected Gaussian Conditional Random Fields (CRFs) (Lafferty et al., 2001) are an extremely popular addition to semantic segmentation architectures. The CRF has the ability to improve initial predictions of locations, and can “sharpen” predictions to make them consistent with edges and areas with consistent color in the original image. CRF post-processing for refining supervised and weakly supervised semantic segmentation predictions is ubiquitous in the literature (Lafferty et al., 2001; Chen et al., 2014; Long et al., 2015; Liu et al., 2020; Ahn et al., 2019). Recently, new connections between CRF message passing and convolutional networks have allowed CRFs to be embedded into existing models (Chen et al., 2017; Teichmann & Cipolla, 2018) and trained jointly for better performance. By connecting the STEGO correspondence distillation loss to the energy of an undirected model on image pixels we can use the same minibatch MLE strategy to estimate other similar graphical models. For example, in the fully connected Gaussian edge potential CRF, one forms a pairwise potential function potential function for the pixels of a single image:
wcrf (vi, vj) = a exp
( −|pi − pj | 2
2θ2α − |Ii − Ij | 2 2θ2β
) + b exp ( −|pi − pj | 2
2θ2γ
) (15)
Where pi represent the pixel coordinates associated with node vi and Ii represents pixel colors associated with node vi. The parameters a, b, θα, θβ , θγ are hyperparameters and control the behavior of the model. These parameters balance the effect of long- and short-range color similarities against
smoothness. The CRF directly learns a pixel-wise array of probabilistic class assignments over k labels corresponding to the probability simplex code space C = P(l) and a non-parametric clustering function f . For a compatibility function µ the CRF chooses the Potts Model (Potts, 1952): µpotts(φ(vi), φ(vj)) := P(φ(vi) 6= φ(vj)). With this setting of the weights and compatibility function, we directly recover the binary potentials of the fully connected Gaussian edge potential CRF (Krähenbühl & Koltun, 2011). We can also add the unary potentials which are often the outputs of another model. However, for our analysis we explore the case without unary potentials which yields an “unsupervised” variant of the CRF. However, without external unary potential terms, the strictly positive similarity kernel encourages the maximum likelihood estimator (MLE) of the graph to be the constant function. To rectify this, we can add small negative constant, −b, to the weight tensor to push unrelated pixels apart. This negative force is the direct analogue of the negative pressure hyper-parameter in STEGO and can be interpreted through the lens of negative sampling (Mikolov et al., 2013). This negative shift also appears in the word2vec and graph2vec embedding techniques (Narayanan et al., 2017; Levy & Goldberg, 2014). Our shifted CRF potential encourages natural clusters to form that respect the structure of the potentials that capture similarities in pixel colors and locations. In the discrete case, solutions to this equation resemble superpixel algorithms such as SLIC (Zhang et al., 2015). Additionally lifting this to the continuous code space and provide a natural continuous generalization of superpixels and seems to avoid challenging local minima. We illustrate these solutions to just the unsupervised CRF potential in Figure 12. Finally, we note that the second term of Equation 15, referred to as the smoothness kernel, matches IIC’s notion of local class consistency. However, we found that adding these CRF terms to the self-correspondence loss of STEGO did not improve performance.
A.9 IMPLEMENTATION DETAILS
Model STEGO uses the “ViT-Base” architecture of DINO pre-trained on ImageNet. This backbone was trained using self-supervision without access to ground-truth labels. We use the “teacher” weights when creating our backbone. We take the final layer of spatially varying features and apply a small amount (p = 0.1) of channel-wise dropout (Srivastava et al., 2014) before using them throughout the architecture during training. Our segmentation head consists of a linear network and a two-layer ReLU MLP added together and outputs a 70 dimensional vector. We use the Adam optimizer (Kingma & Ba, 2014) with a learning rate of 0.0005 and a batch size of 32. To make our losses resolution independent we sample 121 random spatial locations in the source and target implementations and use grid sampling (Jaderberg et al., 2015) to sample features from the backbone and segmentation heads. Our cluster probe is trained alongside the STEGO architecture using a minibatch k-means loss where closeness is measured by cosine distance. Cluster and linear probes are trained with separate Adam optimizers using a learning rate of .005
Datasets We use the training and validation sets of Cocostuff described first in Ji et al. (2019) and used throughout the literature including in Cho et al. (2021). We note that the validation set used in Ji et al. (2019) is a subset of the full CocoStuff validation set and we use this validation subset to be consistent with prior benchmarks. We note that using the full validation set does not change results significantly. When five-cropping images we use a target size of (.5h, .5w) for each crop where h,w are the original image height and width. Training images are then scaled to have minor axis equal to 224 and are then center cropped to (224, 224), validation images are first scaled to 320 then are center cropped to (320, 320). All image resizing uses bilinear interpolation and resizing of target tensors for evaluation uses nearest neighbor interpolation.
CRF We use PyDenseCRF (Krähenbühl & Koltun, 2011) with 10 iterations with parameters a = 4, b = 3, θα = 67, θβ = 3, θγ = 1 as written in Section A.8.
Compute All experiments use PyTorch (Paszke et al., 2019) v1.7 pre-trained models, on an Ubuntu 16.04 Azure NV24 Virtual Machine with Python 3.6. Experiments use PyTorch Lightning for distributed and multi-gpu training when necessary (Falcon et al., 2019).
Hyperparameters We use the following hyperparameters for our results in Tables 1 and 3:

heading: A.10 A HEURISTIC FOR SETTING HYPER-PARAMETERS

heading: A.11 A NOTE ON 5-CROP NEAREST NEIGHBORS

Review summaries:
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).

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). 

This paper proposed STEGO, an unsupervised approach for semantic clustering/segmentation using feature refinement on top of self-supervised neural networks. A distilled version of self-supervised features is learned segmentation specifically through a feed-forward network via SGD.  The main contribution comes from an idea of distilling powerful deep features from strong self-supervised backbones to further improve their semantic discriminativeness. Though the idea is intuitive, several loss functions and regularisations are proposed to avoid trivial solutions and make the idea really work in practice. Extensive qualitative and quantitatve results are done to demonstrate the performance of STEGO and its design choices.