PiCIE: Unsupervised Semantic Segmentation using Invariance and Equivariance in Clustering

In plain words

Labelling every pixel of an image is costly. PiCIE asks whether a network can learn to split the pixels of an image into groups that mean something, with no labels at all. It starts from a simple idea: pixels that look alike should land in the same group. On its own that is not enough, because the network can latch onto small details such as colour and texture rather than objects.

The extra ingredient is consistency. The same image is shown twice with different colour changes and different crops or flips. The network is pushed to give the same group to the same part of the scene in both versions, and the groups themselves should come out the same. That pushes it toward groups that follow real things and surfaces, not just their colour. The method works both for countable objects (people, cars) and for background material (sky, grass, road). It is evaluated on the COCO and Cityscapes street and scene datasets.

TL;DR

• Unsupervised semantic segmentation: each pixel gets a cluster label, with no ground-truth labels in training §1 §3.
• A parametric pixel classifier is replaced by a non-parametric one that scores pixels against cluster centroids §3.2.
• Two photometric views and one geometric transform per image give within-view and cross-view losses that enforce invariance and equivariance §3.3 eq. 10-14.
• Author-reported: best COCO-All result 49.99 Acc / 14.36 mIoU with auxiliary over-clustering, against 32.21 / 9.79 for modified DeepCluster; the paper says it handles both things and stuff tab. 1 §4.4.

Problem

Pixel labels are expensive to collect §1. Earlier unsupervised methods that cluster images into classes do not extend to scene images, where one image holds many classes, and where background stuff dominates the pixel count §1 §2. Methods that maximise mutual information between two pixel views assume roughly balanced clusters, which scene data violates. Such methods are limited to stuff categories in the paper's account, with the reasons given in §2. Learning features and clusters together is unstable when a learned classifier chases noisy pseudo-labels §3.2. The paper states no hyperparameter tuning in this setting; it still tunes the IIC baseline's λ and learning rate [supp. S1].

Method

Baseline (DeepCluster-style, pixel level) §3.1. Alternate two steps. (i) Run k-means over all pixel features in the dataset to get centroids and pixel labels eq. 1. (ii) Train the network with cross-entropy against those labels through a learned pixel classifier eq. 2-3.

Non-parametric prototype classifier §3.2. The learned classifier is removed. Each pixel is scored against the current centroids, and training minimises the negative log of a softmax over those scores eq. 4-5. The paper defines the distance as cosine distance eq. 5; the released code differs (see Paper vs code).

Invariance to photometric changes §3.3.1. For each image, sample two photometric transforms P¹ and P² and compute pixel features for each view eq. 6-7. Cluster each view separately, giving two sets of labels and centroids eq. 8-9. Two losses follow:

• Within-view: each view's features are trained against its own labels and centroids eq. 10.
• Cross-view: each view's features are trained against the other view's labels and centroids eq. 11. This is the term that makes the clustering itself invariant, not only the features.

Equivariance to geometric changes §3.3.2. Sample one geometric transform G per image, a random crop plus a horizontal flip [supp. S1]. View 1 uses the features of the transformed image, G(P¹(x)). View 2 uses the transformed features of the photometric view, G applied to f(P²(x)) eq. 12-13. The two views are clustered and trained exactly as in the photometric case. The total objective adds the within and cross terms eq. 14.

Per-epoch schedule and balancing §3.3 §4.1 [supp. S1]. Clustering and training alternate once per epoch. The transforms are sampled once and reused for both steps of that epoch, so the labels match the inputs seen in training. Each cluster's cross-entropy term is weighted by its share of pixels, which the paper describes as balancing the clusters §4.1. Auxiliary over-clustering (the "+H" variant) adds a second, finer clustering with a loss scale chosen so that its magnitude is comparable to the coarse one §4.1 eq. 15.

Evaluation [supp. S1]. Predicted clusters are matched one-to-one to ground-truth classes with Hungarian matching, and pixel accuracy and mIoU are reported. For COCO, the things and stuff subsets are scored by slicing the confusion matrix of the 27-class model [supp. S1].

# PiCIE, one training epoch  [§3.1-3.3] [eq. 1-14]
# Input:  images (N, 3, H, W); feature network f; K clusters
# Output: updated f; pixel labels per view, reused in training

def picie_epoch(images, f, K):
    # Draw a photometric pair P1, P2 and a geometric G per image; keep
    # them fixed for the whole epoch  [§3.3] [supp. S1]
    transforms = sample_transforms(images)

    # View 1 is G(f(P1(x))); view 2 is G(f(P2(x))) at feature level; the
    # features are L2-normalised  [§3.3.2] [eq. 12-13]
    feats_1 = view_features(images, transforms, f, view=1)  # (N, h, w, D)
    feats_2 = view_features(images, transforms, f, view=2)  # (N, h, w, D)

    # Cluster each view on its own with mini-batch k-means  [§3.1] [eq. 8-9]
    centroids_1 = minibatch_kmeans(feats_1, K)
    centroids_2 = minibatch_kmeans(feats_2, K)

    # Label each pixel with its nearest centroid, separately per view  [§3.3]
    labels_1 = nearest_centroid(feats_1, centroids_1)
    labels_2 = nearest_centroid(feats_2, centroids_2)

    # Class weight per cluster: its share of pixels  [§4.1]
    weight_1 = cluster_share(labels_1, K)
    weight_2 = cluster_share(labels_2, K)

    # Train f; the centroids and labels stay fixed for the epoch  [§3.2]
    for batch in epoch_batches(images, transforms, labels_1, labels_2):
        z_1 = view_features(batch.images, batch.transforms, f, view=1)
        z_2 = view_features(batch.images, batch.transforms, f, view=2)

        # Within-view: each view against its own labels and centroids  [eq. 10]
        loss_within = (
            cross_entropy(z_1, centroids_1, batch.labels_1, weight_1)
            + cross_entropy(z_2, centroids_2, batch.labels_2, weight_2))

        # Cross-view: each view against the other view's labels  [eq. 11]
        loss_cross = (
            cross_entropy(z_1, centroids_2, batch.labels_2, weight_2)
            + cross_entropy(z_2, centroids_1, batch.labels_1, weight_1))

        # Mean of the four terms; the paper's eq. 14 sums the two  [code]
        loss = (loss_within + loss_cross) / 4
        f = adam_step(f, loss)

Key equations, as printed in the paper:

\mathcal{L}_{clust} = -\log \frac{\exp(-d(f_\theta(x_i)[p], \mu_{y_{ip}}))}{\sum_{l} \exp(-d(f_\theta(x_i)[p], \mu_l))}, \quad d = \text{cosine distance} eq. 5

\mathcal{L}_{within} = \sum_{i,p} \mathcal{L}_{clust}(z^{(1)}_{ip}, y^{(1)}_{ip}, \boldsymbol{\mu}^{(1)}) + \mathcal{L}_{clust}(z^{(2)}_{ip}, y^{(2)}_{ip}, \boldsymbol{\mu}^{(2)}) eq. 10

\mathcal{L}_{cross} = \sum_{i,p} \mathcal{L}_{clust}(z^{(1)}_{ip}, y^{(2)}_{ip}, \boldsymbol{\mu}^{(2)}) + \mathcal{L}_{clust}(z^{(2)}_{ip}, y^{(1)}_{ip}, \boldsymbol{\mu}^{(1)}) eq. 11

\mathcal{L}_{total} = \mathcal{L}_{within} + \mathcal{L}_{cross} eq. 14

\mathcal{L} = \lambda_{K_1} \mathcal{L}_{K_1} + \lambda_{K_2} \mathcal{L}_{K_2}, \quad \lambda_{K_1} = \frac{\log K_2}{\log K_1 + \log K_2}, \quad \lambda_{K_2} = \frac{\log K_1}{\log K_1 + \log K_2} eq. 15

The paper's pseudocode is an unnumbered figure in §3.3. The archive also holds an algorithm version (picie_pseudocode.tex), but main.tex does not include it, so the pseudocode is cited by section here.

Hyperparameters

Results

Author-reported. Setup: COCO-Stuff curated to 27 classes (15 stuff, 12 things), FPN with ResNet-18 pretrained on ImageNet, 320×320 inputs, evaluated on pixel accuracy and mIoU after Hungarian matching §4.1 §4.3 [supp. S1]. Hardware: not stated in the paper.

COCO-All, 27 classes tab. 1.

Partitions of COCO tab. 2. Same trained model, scored on each subset.

Table labels are numbered in the order the tables are included in the source (see Pitfalls for a numbering conflict).

Cityscapes, K=27 tab. 3.

COCO-Stuff without ImageNet pretraining, 128×128, stuff only tab. 4. The table does not label its metric; the caption places it on the benchmark of the IIC and AC papers.

Ablation: components tab. 5. Columns mark which components are on.

Ablation: alternatives tab. 6. Columns mark which alternative is used.

Linear probe on frozen features tab. 7. A new linear classifier trained on each embedding, with ground-truth labels.

Re-training as initialisation tab. 8. Each trained network is then fine-tuned with supervised labels. C-Acc and C-mIoU are clustering results after that supervised training.

Robustness at test time [supp. S2]. PiCIE (48.09 Acc, 13.84 mIoU, no augmentation) is evaluated with the same transforms applied at test time. Applying all photometric transforms gives 47.42 / 13.39; all photometric and geometric transforms together give 46.28 / 13.16.

Text claims checked against the tables.

• "+17.5 Acc" on COCO §4.4 [abstract]: the table gives 49.99 − 32.21 = 17.78 against modified DC, the best Acc baseline in tab. 1. The paper's figure cannot be reproduced from the table.
• "+4.5 mIoU" on COCO: 14.36 − 9.79 = 4.57 tab. 1.
• Cityscapes "+18 Acc, +5.3 mIoU" §4.4: 65.50 − 47.88 = 17.62 against IIC; 12.31 − 7.06 = 5.25 against modified DC tab. 3.
• "+10 mIoU" on things §4.4: 23.83 − 13.64 = 10.19 against IIC; against modified DC the gap is 12.28 tab. 2.
• The paper says PiCIE keeps "better or compatible" stuff performance §4.4. In tab. 2, PiCIE + H stuff mIoU is 17.32, below modified DC's 22.24, while its stuff accuracy is higher.

Paper vs code

The official repository is janghyuncho/PiCIE, linked from the abstract; its owner is the first author. code_commit is the HEAD of a shallow clone made at fetch time. The paper does not name a commit, so these comparisons assume the HEAD matches the released results. The code is MIT-licensed LICENSE.

1. Learning rate. Paper: Adam at 1e-3 for PiCIE and modified DC [supp. S1]. Code: launch script sets LR=1e-4 train_picie.sh, and the argparse default is 1e-4 train_picie.py. The three released logs record lr=0.0001 picie_coco.out. Treat the paper's 1e-3 as unverified for PiCIE.

2. Distance and training score. Paper: softmax over negative cosine distance eq. 5. Code, training: the classifier is a 1×1 convolution with its weight set to the raw centroids and zero bias utils.py train_picie.py. Logits are therefore dot products between unit-norm features and unnormalised centroids, not cosine distances. Code, pseudo-labels: assignment minimises squared Euclidean distance, computed as −(1 − 2 f·μ + ‖μ‖²) utils.py. Argmax of that score includes a −‖μ‖² term that the training logit does not, so the labels and the training scores can disagree when centroid norms differ.

3. Total loss. Paper: L_total = L_within + L_cross eq. 14. Code averages them: loss = (loss_across + loss_within) / 2, where each term is itself an average of two views train_picie.py train_picie.py train_picie.py. Supplement: "we simply average the cross and within losses" [supp. S1]. The algorithm version in the archive (picie_pseudocode.tex, not included by main.tex) also divides by 4. Implement the average.

4. Equivariance composition. Paper text: view 2 is G applied to the features of P²(x) eq. 13. The algorithm version in the archive writes o_2 ← f(G(P_2(x))), transforming the image before the network. Code: view 2 takes the features of the photometric view and applies the feature-level transform train_picie.py. Follow the text, which matches the code.

5. Grayscale never applied. Paper: grayscale with p=0.2 [supp. S1]. The argument parser appends the key 'grey' utils.py. The dataset checks for 'gray' coco_train_dataset.py. The two strings differ, so grayscale is silently skipped in the released training. The released ablation numbers therefore do not include grayscale augmentation.

6. Blur σ range. Paper: σ in [0.1, 2] [supp. S1]. Code: the blur is applied when a uniform draw is below p; the radius is an affine function of that draw with slope (max−min)/(1−p) and intercept min, so the radius is at most 0.1 for every blurred image and can go negative custom_transforms.py (with the code's own values: 3.8·x + 0.1 for x in [−0.5, 0)). Not run to confirm how PIL treats a negative radius.

7. Over-clustering "+H". Paper: two clustering levels with K2 = 100 and the λ balancing of eq. 15 §4.1. Code: the training script has a single K_train train_picie.py, and no second clustering or λ weighting appears in the cloned repository. The "+H" rows in tab. 1, tab. 2, tab. 5, tab. 7 and tab. 8 cannot be reproduced from this release.

8. Class weights. Paper: each cluster is weighted "balanced with the size of each cluster" [supp. S1] and "balance term for each cluster" §4.1. Code: weight is the cluster's share of pixels, count / sum commons.py and train_picie.py. That makes large clusters weigh more, which is the opposite of what the balancing rationale suggests; the paper's words do not settle the direction.

9. k-means batch size, init and update. Paper: 50 initial batches of 128; centroids updated every 20 iterations [supp. S1]. Code launch: 20 initial batches of 256, one batch per update train_picie.sh train_picie.sh. Code default: 30 initial batches train_picie.py. The released COCO reproduce log shows batch_size_cluster=128 picie_coco.out.

10. Input size. Paper: all images resized and center-cropped to 320×320 during training §4.1. Code: the base transform resizes and crops to res2, 640 in the launch script coco_train_dataset.py train_picie.sh. View 2 is then resized to 320 coco_train_dataset.py coco_train_dataset.py. View 1 reaches 320 only through the random crop's resize custom_transforms.py.

11. FPN fusion. Paper supplement: each level is projected, bilinearly upsampled to 1/4 scale and summed [supp. S1]. Code: a top-down cascade in which each level is upsampled to the next level's size and added to that level's projection fpn.py fpn.py. The two give different sums; the code is what the released results used.

12. Learning-rate decay. Not stated in the paper. Code: the decay function exists but its call is commented out, so the rate is constant train_picie.py.

13. Final numbers. The released reproduce log's last COCO evaluation reports 51.61 Acc and 15.23 mIoU picie_coco.out, which differs from tab. 1. Its arguments include batch size 128. The reproduce log does not match the paper's 49.99 / 14.36.

No other discrepancies were found in the model definition, the feature normalisation or the cross/within structure of the loss.

Pitfalls

• Two learning rates in print. The supplement gives 1e-3 for PiCIE; every released run uses 1e-4. Expect instability if you follow the supplement and the code's logs disagree. Start from 1e-4 and say which you used.
• Grayscale is dead code in the release. The key mismatch ('grey' vs 'gray') means the grayscale augmentation at p=0.2 never runs. If you want the paper's augmentation set, fix the key.
• Blur radius out of range. The σ mapping gives at most 0.1 for applied blurs. Use an explicit uniform draw in [0.1, 2] if you want the paper's setting. Test how your image library handles a negative radius before you rely on it.
• Labels and scores come from different rules. The pseudo-label uses squared Euclidean distance; the training logit is a plain dot product with unnormalised centroids. The paper's cosine formula matches neither exactly. Pick one rule and use it for both labels and training, normalising the centroids if you follow the cosine formulation.
• Eq. 14 sums; the code averages. A factor of two or four on the total loss changes the effective learning rate, and the paper's eq. 14 and its algorithm version disagree on it.
• Class weights run the wrong way for balancing. Weight is proportional to cluster size. If the intent is to counter imbalance, the code is not doing that. PyTorch's weighted cross-entropy with reduction='mean' also normalises by the sum of the weights of the batch's targets; replicate that normalisation or the loss scale will differ.
• Batch norm changes during clustering. The k-means pass sets the model to training mode under no_grad commons.py, so BatchNorm running statistics update during clustering. Label computation then switches to eval mode commons.py. The paper does not mention this.
• Hidden costs. Clustering is offline, once per epoch. The supplement says about 20 minutes per epoch on COCO for the pseudo-labels [supp. S1]. Per-image transform parameters and per-image label files are stored for the whole dataset utils.py coco_train_dataset.py. The code needs a CUDA GPU and a GPU FAISS index utils.py.
• Labels are at 1/4 resolution. Pseudo-labels and the trained head live at about 80×80 for 320×320 inputs [supp. S1]. Boundaries are imprecise; the paper says so and suggests an edge detector or CRF as a fix [supp. S3]. Evaluation upsamples predictions bilinearly to label size commons.py.
• Partition metrics assume a label order. Things and stuff are scored by slicing the confusion matrix at index 12 after matching commons.py. This is correct only if the 27 class ids are ordered things-first, which the code's comment asserts.
• Numbering conflict in the supplement. The supplement calls the 128×128 benchmark "Table 3" [supp. S1]. The source inputs the Cityscapes table before that one, so by order it would be Table 4. Check the compiled PDF before citing a table number.
• "No hyperparameter tuning". The paper states it for the unsupervised setting §4.1, but tunes the IIC baseline's λ and learning rate [supp. S1].
• Seeds are not enough. The code sets deterministic cuDNN, but the authors note that the interpolation backward pass is not deterministic on GPU utils.py. Expect small run-to-run variance.

Lineage

• Builds on: DeepCluster (alternating clustering and training) §3.1 and IIC (mutual-information pixel clustering, used as a baseline) §2 §4.2. Neither has a card in this wiki; arXiv ids not stated in the paper's text.
• Compared against: IIC, IIC - res12, modified DeepCluster, and No Train baselines tab. 1; SIFT, Doersch 2015, Isola 2016, AC on the stuff-only benchmark tab. 4.
• Compared against in this wiki: 2203.08414 reports PiCIE and PiCIE + H numbers in its own comparison tables.
• Superseded by: not stated in this card.