Unsupervised Semantic Segmentation by Distilling Feature Correspondences

In plain words

Labelling every pixel of an image with its object type is slow and expensive. This paper asks whether a computer can find the object types with no labels at all.

The authors start from a model that was trained without labels to produce a feature vector for every small patch of an image. Two patches that show the same kind of thing tend to have similar vectors, across the image and across other images. STEGO keeps that model frozen and trains a small add-on network. The add-on is pushed so that pixel pairs which look alike in the frozen features become alike in its own output, and pairs that do not look alike are pushed apart. Afterwards, the outputs are grouped into clusters, and each cluster becomes one class.

The payoff is that the add-on trains quickly, needs no labels, and gives clearly better clusters than the earlier label-free methods the authors compare against.

TL;DR

• STEGO takes a frozen self-supervised DINO ViT, computes dense feature correspondences between images, and distills them into a 70-dimensional code per pixel §3.1 §3.3 app. Implementation Details.
• The loss is a contrastive-style correlation loss with a spatial centring step and a zero clamp, applied to a self pair, a k-nearest-neighbour pair and random pairs eq. 4 eq. 5.
• Codes are clustered and refined with a fully connected CRF §3.3.
• Author-reported on CocoStuff-27: 28.2 unsupervised mIoU and 41.0 linear-probe mIoU, against 14.4 and 14.8 for PiCIE+H §4.2.
• Relevant if you need label-free pixel classes and can afford precomputed nearest neighbours and a frozen ViT-B.

Problem

Dense labels are costly to produce, and in domains such as medicine or astrophysics the ground truth may be unknown or need expert knowledge §1. Earlier unsupervised segmentation methods learn features and clusters together, for example through transformation equivariance or mutual-information clustering (IIC, PiCIE) §1 §2. The authors argue that a frozen self-supervised backbone already has dense features whose correlations are semantically consistent, so feature learning and cluster compactness can be separated §1 §3.1.

Method

STEGO has three parts: a frozen backbone, a light segmentation head, and a correlation loss. Clustering and a CRF are applied afterwards §3.3.

1. Backbone. A DINO ViT-Base (teacher weights, patch size 8) produces dense final-layer features. The backbone is frozen and not fine-tuned §3.3 app. Implementation Details train_config.yml.
2. Correspondence tensor. For two images, the cosine similarity between every feature location of one and every location of the other gives a 4-index tensor F. Its entries are cosine similarities §3.1 eq. 1. The authors show that F predicts ground-truth label co-occurrence well on CocoStuff (about 50% recall at 90% precision for DINO, author-reported) without using labels to tune anything §3.1.
3. Segmentation head. A map S projects backbone features to K-dimensional codes with K < C. The head is a linear 1x1 path plus a two-layer ReLU MLP path, summed. K = 70 §3.2 app. Implementation Details modules.py.
4. Simple correlation loss. Multiply F and the code correlation S element-wise, subtract a shift b, and minimise the negative sum. Here b sets a uniform negative pressure §3.2 eq. 2.
5. Clamp. Clamp S at zero before the product. This keeps weakly correlated codes orthogonal instead of anti-aligned, which the authors found stabilises training §3.2.
6. Spatial centring. Subtract the mean over the second image's locations from each row of F. This rebalances small objects whose correlation pattern is concentrated §3.2 eq. 3.
7. Final correlation loss. Combine steps 4 to 6 into L_corr eq. 4.
8. Three pair types. Per training step, compute L_corr for (a) an image with itself, (b) an image with a random one of its 7 nearest neighbours, and (c) an image with a randomly shuffled batch partner. Combine them with weights lambda and shifts b §3.2 §3.3 eq. 5.
9. Nearest neighbours. Each image is represented by the spatial average of its backbone features, normalised §3.3. For each image the top 7 neighbours are stored, using cosine similarity §3.3 precompute_knns.py.
10. Five-crop. Before neighbours are mined, each training image is split into five crops of half its height and width, and each crop counts as a separate image §3.3 app. Implementation Details app. A note on 5-Crop Nearest Neighbors.
11. Clustering. The codes are clustered with minibatch k-means on cosine distance §3.3.
12. CRF. A fully connected CRF refines the cluster labels with pairwise pixel terms §3.3 app. Implementation Details.

The pair loss in pseudocode (steps 4 to 8), with the training step that combines it:

# Correspondence loss for one pair of images  [eq. 4]
# feats_*: frozen features at P = 121 sampled points, shape (B, P, C)
# codes_*: head outputs at the same points, shape (B, P, K)
def corr_loss(feats_a, feats_b, codes_a, codes_b, shift):
    # Cosine correlation of frozen features, no gradient  [eq. 1]
    corr = cosine_matrix(feats_a, feats_b)           # (B, P, P)
    old_mean = corr.mean()
    # Spatial centring: remove each row's mean  [eq. 3]
    corr = corr - corr.mean(dim=-1, keepdim=True)
    # Restore the global mean  [code: modules.py#L331-L333]
    corr = corr - corr.mean() + old_mean
    # Cosine correlation of the learned codes  [eq. 2]
    code_corr = cosine_matrix(codes_a, codes_b)      # (B, P, P)
    # Clamp code correlations at zero; push codes together or apart  [eq. 4]
    return -(corr - shift) * clamp(code_corr, min=0.0)

# Training loss for one batch  [§3.2] [eq. 5]
# images: (B, 3, 224, 224); nn_images: one nearest neighbour per image
def training_step(images, nn_images):
    # Frozen features with channel dropout  [app. Implementation Details]
    feats = dropout(backbone(images), p=0.1)         # (B, C, h, w)
    nn_feats = dropout(backbone(nn_images), p=0.1)   # (B, C, h, w)
    # Segmentation head: linear path plus two-layer ReLU MLP  [§3.2]
    codes = head(feats)                              # (B, 70, h, w)
    nn_codes = head(nn_feats)                        # (B, 70, h, w)
    # 121 random locations per image (11 x 11)  [app. Implementation Details]
    coords = random_points(B, 11, 11)
    a_feats, a_codes = sample(feats, codes, coords)
    k_feats, k_codes = sample(nn_feats, nn_codes, coords)
    # Self term: weight lambda_self, shift b_self  [eq. 5]
    loss = lambda_self * corr_loss(a_feats, a_feats,
                                   a_codes, a_codes, b_self)
    # Nearest-neighbour term: weight lambda_knn, shift b_knn  [eq. 5]
    loss += lambda_knn * corr_loss(a_feats, k_feats,
                                   a_codes, k_codes, b_knn)
    # Random term: five shuffled partners, averaged  [code: modules.py#L382]
    rand_losses = []
    for _ in range(5):
        partner = shuffled_batch_index(B)            # see Pitfalls
        rand_losses.append(corr_loss(a_feats, feats[partner],
                                     a_codes, codes[partner], b_rand))
    loss += lambda_rand * mean(rand_losses)
    return loss

The head and loss are small, so training is cheap. The authors report under two hours on a single V100 §3.3. The paper also shows that the full loss is equivalent to maximum-likelihood inference in a Potts (Ising-type) energy over all pixels of the dataset, which links it to CRF-style refinement §3.4 app. Relationship with Graph Energy Minimization eq. 6 eq. 7.

Hyperparameters

Core training settings. The paper's appendix states the values. The code column gives the default in src/configs/train_config.yml where it differs.

Loss weights and shifts. The paper's table is per dataset. Subscripts: self, knn (nearest neighbour), rand (random).

Sources: paper values from app. Implementation Details (table "Hyperparameters used in STEGO"). The paper's CocoStuff column matches the commented block "Cocostuff27 10/3 vit_base" train_config.yml and the Cityscapes column matches the commented block "Cityscapes 10/3" train_config.yml, both to rounding. The active code defaults are the "IAROA vit small" block train_config.yml. The paper says the lambda and b values tended to be dataset- and network-specific §3.3.

Results

Author-reported. Setup: 27-class CocoStuff validation set and 27-class Cityscapes validation set, backbone DINO ViT-Base, unsupervised metrics with Hungarian matching, linear probe trained with cross-entropy on the frozen codes §4.1 §4.2. The paper states that the validation set used by IIC is a subset of the full CocoStuff validation set, and that the full set does not change results much app. Implementation Details. CocoStuff-27, validation §4.2 (table "cocostuff_results" in the paper):

Cityscapes-27, validation, unsupervised only §4.2 (no linear probe reported for prior work):

Ablation on CocoStuff-27 §4.3 (table "ablation" in the paper; columns: 0-Clamp, 5-Crop, SC = spatial centring, CRF; check marks as in the paper):

Additional ablation (ViT-S, one component removed at a time) app. Additional Ablation Study:

Potsdam-3 (unsupervised accuracy) app. Additional Results on the Potsdam-3 Dataset: STEGO 77.0; IIC 65.1; K-means 45.7; Deep Cluster 41.7; SIFT 38.2; random CNN 38.2. The paper reports about +12 points over IIC.

Qualitative §4.2: the paper shows example segmentations on CocoStuff and Cityscapes, and a failure section (ground vs. trees, barn floor vs. outdoor, food as thing vs. stuff). The authors note the label ontology itself is ambiguous, so the linear probe is the more reliable quality signal §4.2.

Paper vs code

The official code is mhamilton723/STEGO (MIT, copyright Mark Hamilton, the first author). The paper's link aka.ms/stego-code redirects to it. The card was checked against commit eb4d6b5 (shallow clone). Discrepancies found:

1. Default config is not the paper's setting. The active defaults are a ViT-Small run (model_type: "vit_small" train_config.yml) with weights from the "IAROA vit small" block train_config.yml. The paper's CocoStuff and Cityscapes values sit in commented blocks train_config.yml train_config.yml. Running the defaults does not reproduce the paper's numbers. The paper's appendix says ViT-Base app. Implementation Details.
2. Batch size. Paper: 32 app. Implementation Details. Code default: 16 train_config.yml.
3. Loss normalisation. Paper's eqs. 2 and 4 use sums over all location pairs eq. 2 eq. 4. Code takes .mean() over each term: the self and KNN terms are averaged over the sampled points, and the random terms are concatenated over 5 partners and then averaged modules.py train_segmentation.py. The effective scale of the loss, and therefore the interaction with lr and lambda, differs from the sum form. 4. Global mean restored after centring. Paper eq. 3 only centres each row eq. 3. The code also adds back the global mean of the uncentred matrix modules.py. A commented-out passage in the paper's source describes a renormalisation step; eq. 3 omits it §3.2.
4. Number of random partners. Paper does not state how many random images are used per step. Code uses 5 train_config.yml modules.py.
5. Dropout placement. Paper: one channel-wise dropout on the backbone features app. Implementation Details. Code applies dropout to each head input path modules.py and also to the features returned to the loss modules.py.
6. Head hidden width. Paper: "two-layer ReLU MLP" with output 70 app. Implementation Details; hidden width not stated. Code uses the backbone width as hidden width: 384 for ViT-Small, 768 otherwise modules.py.
7. Nearest-neighbour list length. Paper: top 7 §3.3. Code precomputes the top 30 per image precompute_knns.py but samples only from entries 1 to 7 data.py.
8. Cluster probe. Paper: minibatch k-means on cosine distance §3.3 app. Implementation Details. Code: a learned ClusterLookup trained with a probability-weighted cosine loss modules.py modules.py. These may be equivalent, but this card has not verified that.
9. Stabilised variant. Code has an option stabalize that clamps code correlations at 0.8 modules.py. It is off by default train_config.yml. The paper does not describe it.
10. CRF parameters. The paper gives a, b, theta_alpha, theta_beta, theta_gamma app. Implementation Details. The config uses different names and values (alpha 0.5, beta 0.15, gamma 0.05, w1 10, w2 3) train_config.yml. The mapping is not stated, so the CRF settings used in the paper cannot be confirmed from the code here.

Points checked and matching: lr 5e-4, K = 70, 7 neighbours, five-crop ratio 0.5, 224 training resolution, patch size 8, 11 x 11 sampled points (121), channel dropout p = 0.1, probe lr 5e-3 train_config.yml train_segmentation.py.

Pitfalls

• Start from the paper's table, not the defaults. Running the shipped config gives the ViT-S setup. Use the ViT-B CocoStuff row, or the Cityscapes row, for the paper's numbers, and keep the batch size at 32 app. Implementation Details.
• Lambda and b are dataset- and network-specific. The authors say they tuned b to keep mean nearest-neighbour similarity near 0.3 and mean random similarity near 0.0 §3.3. They call hyperparameter setting without ground-truth cross-validation an open problem app. A Heuristic for Setting Hyper-parameters. Their diagnostic is the histogram of correspondence similarities over training: a balanced run is bimodal, peaking at 1 and at 0 app. A Heuristic for Setting Hyper-parameters The paper describes this diagnostic in words and in a figure in that subsection.
• Random partner selection. The paper says random images are made by shuffling and ensuring no image matches itself §3.3. Code's super_perm takes a random permutation and adds 1 to any fixed point, then takes the result modulo the batch size modules.py. Read as code, this is not a strict permutation: a fixed point i can be moved onto i+1 while another index already holds i+1, so one image can appear twice in a batch. This is inferred by reading the code and was not run.
• Nearest neighbours are a hidden cost. You need an all-pairs cosine top-K over the training set before training. Training raises an error if the neighbour file is missing data.py. Five-crop multiplies the number of "images" by five before neighbours are mined §3.3.
• Five-crop neighbours can be from the same image. The authors note that most patches have no neighbour from their own image, but some do app. A note on 5-Crop Nearest Neighbors.
• Feature resolution depends on input size. The paper states a feature resolution of (40, 40) §3.3, which matches a 320-pixel input with patch 8 (320 / 8 = 40). Training uses 224 pixels, so the map is smaller app. Implementation Details. The 28 x 28 size is derived from those numbers, not stated.
• Ablation trade-offs. Removing SC (pointwise) raises unsupervised accuracy to 50.2 but lowers unsupervised mIoU to 22.3 compared with the full ViT-S run (48.3 and 24.5) app. Additional Ablation Study. Pick the metric you care about.
• Metric depends on matching. Unsupervised scores use Hungarian matching of clusters to labels, which is invariant to permutation but depends on the label set §4.1. The authors say the 27-class ontology is partly arbitrary §4.2.
• Training length and variance. Training steps are not stated in the paper. Code default is 5000 steps train_config.yml. Seeds and run-to-run variance are not stated.
• Equation and table numbers are counted from the LaTeX source (equations in order: correspondence = 1, simple loss = 2, centring = 3, final loss = 4, full loss = 5, energy = 6, Boltzmann = 7). They were not checked against the compiled PDF. Table numbers are cited by section instead.
• Frozen backbone only. Results are for a frozen backbone. The paper does not test fine-tuning §3.3.

Lineage

• Builds on: DINO self-supervised ViT features §3.1 §2; correlation volume from optical flow work §3.1; fully connected Gaussian CRF §3.3; graph and Potts-model energy view §3.4.
• Compared against: PiCIE and PiCIE + H, MDC, IIC, and clustering on raw ResNet50, MoCoV2 and DINO features §4.2. Related work also discusses SegSort and MaskContrast §2.
• Superseded by: not stated in this wiki.