Paper of the day, 2026-10-28

Unsupervised Semantic Segmentation by Distilling Feature Correspondences

Listen to the episode, or read on. The paper: arXiv:2203.08414.

STEGO: the frozen network already knew which pixels belong together

Labelling every pixel of a street scene by hand costs a lot. STEGO skips the labelling. It takes a vision transformer that was trained without labels, keeps it frozen, and trains a small add-on until the network's existing sense of which patches look alike becomes a clean pixel-by-pixel map of classes. On the 27-class CocoStuff benchmark, the authors report doubling the best earlier label-free score while training for under two hours on one GPU.

The problem in one picture

Drawing made for wikillm on 2026-10-28 by the AI editor. It is not taken from the paper.

How it works

Most earlier label-free segmentation methods, such as IIC and PiCIE, had to do two jobs together: learn good features and make those features form tight clusters. STEGO separates the jobs. A self-supervised DINO ViT-Base with patch size 8 provides the features and is never fine-tuned. Only a small segmentation head is trained, and its job is to make the clusters compact.

What the head learns from is a correspondence table. Pick two images, take the feature vector at every location in each, and compute the cosine similarity of every location in one with every location in the other. Before anything is trained, the authors check that this table carries meaning. On CocoStuff, high DINO similarities predict that two pixels share a ground-truth label at roughly 50% recall and 90% precision, and no labels were used to tune anything.

The head maps each feature vector to a 70-dimensional code. It sums two paths: a linear 1x1 projection and a two-layer ReLU MLP. The loss compares the frozen similarity table with the similarity of the codes and adds three refinements:

A shift b. It is subtracted from the frozen similarities, so any pair below b gets pushed apart. This works as a steady outward pressure.
A zero clamp on the code similarities. Weakly related codes end up orthogonal rather than pointing in opposite directions, which the authors found keeps training stable.
Spatial centring. Each row of the table has its mean removed. This keeps small objects, whose high similarities sit in a few places, from being drowned out by large regions.

That loss is applied to three kinds of image pair at every step: an image with itself, an image with one of its 7 nearest neighbours, and an image with a random partner from the shuffled batch. Each pair type has its own weight and shift. Neighbours are found before training, using the spatial average of each image's features. Each training image is first cut into five crops of half its height and width, and every crop counts as a separate image. After training, minibatch k-means on cosine distance groups the codes into classes, and a fully connected CRF cleans up the edges.

The paper also shows that the full loss amounts to maximum-likelihood inference in a Potts-style energy over all pixels of the dataset. That puts STEGO in the same family as the CRF it uses for clean-up.

Try it yourself

You can run the core of the loss on a toy example. Below, 90 fake "patches" from three hidden classes stand in for frozen features. The code builds the correspondence table, centres it, and then trains 3-dimensional codes directly with the shifted, clamped loss on the self pair.

Python

import numpy as np
import matplotlib.pyplot as plt
rng = np.random.default_rng(0)
# Toy "frozen features": 3 hidden classes, 30 patches each, 16 channels
centres = rng.normal(size=(3, 16))
labels = np.repeat([0, 1, 2], 30)
feats = centres[labels] + 0.9 * rng.normal(size=(90, 16))
unit = lambda x: x / np.linalg.norm(x, axis=1, keepdims=True)
F = unit(feats) @ unit(feats).T                 # correspondence (self pair)
F = F - F.mean(1, keepdims=True) + F.mean()     # spatial centring, global mean kept
shift = 0.12                                    # b: uniform push-apart pressure
Z = rng.normal(size=(90, 3))                    # learned codes, K = 3 < C = 16
for step in range(300):
    Zn = unit(Z)
    S = Zn @ Zn.T
    grad = -2 * ((F - shift) * (S > 0)) @ Zn    # zero clamp: no push below 0
    Z = Zn - 0.002 * grad                       # step, renormalised next loop
S = unit(Z) @ unit(Z).T
same = labels[:, None] == labels[None, :]
print("feature sim, same class :", round(F[same].mean(), 3))
print("feature sim, other class:", round(F[~same].mean(), 3))
print("code sim, same class    :", round(S[same].mean(), 3))
print("code sim, other class   :", round(S[~same].mean(), 3))
fig, ax = plt.subplots(1, 2, figsize=(9, 4))
ax[0].imshow(F - shift, cmap="coolwarm", vmin=-1, vmax=1)
ax[0].set_title("F - b: pull (red) or push (blue)")
ax[1].imshow(S, cmap="coolwarm", vmin=-1, vmax=1)
ax[1].set_title("learned code similarity")
plt.show()

What to expect. The left panel should show three warm blocks on the diagonal against a cool background: those are the pull and push signals. The right panel should look like a sharper copy of it. Same-class code similarity should approach 1, and other-class similarity should settle near zero rather than going strongly negative, because the clamp stops pushing once two codes are orthogonal. The start value of 0.12 is the self-pair shift the paper uses on CocoStuff.

Explore. Raise shift towards 0.5 and watch real neighbours get pushed apart, then set it to 0 and see the blocks blur. The authors tuned b so that mean nearest-neighbour similarity stayed near 0.3 and mean random similarity near 0.0. You are now doing that tuning by hand.

What the authors report, and what they leave open

All numbers below come from the authors and use a DINO ViT-Base backbone. Unsupervised scores use Hungarian matching of clusters to labels, and the linear probe is trained with cross-entropy on the frozen codes.

CocoStuff-27: 56.9 unsupervised accuracy and 28.2 unsupervised mIoU, against 50.0 and 14.4 for PiCIE+H. With a linear probe: 76.1 accuracy and 41.0 mIoU.
Cityscapes-27: 73.2 unsupervised accuracy and 21.0 mIoU, against 65.5 and 12.3 for PiCIE.
Potsdam-3: 77.0 unsupervised accuracy, against 65.1 for IIC.

The ablations show what matters most. On ViT-S, turning on the zero clamp alone raises unsupervised mIoU from 7.3 to 21.3. Removing the random-pair loss drops it from 24.5 to 12.8. There is one trade-off: without spatial centring, unsupervised accuracy rises to 50.2 but mIoU falls to 22.3, so which run counts as better depends on the metric you care about.

The open questions are real ones:

Setting λ and b without ground truth is, in the authors' own words, an open problem, and the best values differ by dataset and network.
The number of training steps, seeds and run-to-run variance are not stated.
Fine-tuning the backbone is not tested.
The authors call the 27-class label set partly arbitrary, which is why they treat the linear probe as the more reliable quality signal.

If you plan to reproduce the results: the shipped config is a ViT-Small run, not the paper's setup.

Ask yourself

The shift b sets a single "similar enough" threshold for every pair. How would you choose it on a dataset with no labels at all, such as satellite tiles or microscope slides?
The random-pair term matters more than the nearest-neighbour term in the ablation. Why might pushing apart unrelated images teach more than pulling neighbours together?
If the frozen backbone carries a bias, for example treating trees and ground as one thing, can any head on top fix it, or does the bias just get distilled into the codes?