Label-Efficient Semantic Segmentation with Diffusion Models

In plain words

A diffusion model learns to remove noise from images, and to do that well its network has to understand what is in the picture. The paper asks whether the network's inner layers can act as a description of every pixel. If you feed the network a lightly noised copy of a photo and read out its internal values at a few points in the denoising process, the values at each pixel carry information about the object that pixel belongs to.

The method takes those per-pixel descriptions and trains a small classifier on them, using only a handful of hand-labelled images. The classifier then labels every pixel of new images. Because the diffusion model was already trained on unlabelled images of the same kind, few labels are needed. The authors report that this beats the other ways they tried, including GAN-based approaches and self-supervised encoders.

TL;DR

• Take a pretrained DDPM (ADM architecture), noise each labelled image to three diffusion steps, and read the decoder activations of five UNet blocks §3.2.
• Upsample and concatenate them: 8448 features per pixel §3.2.
• Train an ensemble of ten MLPs on pixels of the few labelled images, and predict by majority vote §3.2 app. C.
• Author-reported: best mean IoU among the compared methods on all six datasets tab. 2. The official code differs from the paper in MLP training length and noise handling (see Paper vs code).

Problem

Semantic segmentation needs a label for every pixel, which is costly when only a few annotated images exist. The compared few-shot methods either train a segmentation model on large sets of synthetic images produced by a GAN §4, or pull features out of real images through a GAN encoder or inversion §4. The paper's claim is that DDPM activations need neither a separate encoder nor a synthetic-data stage, and that they are discriminative enough to use directly §1 §5. The stated limit is that the method needs a high-quality diffusion model trained on the target domain, which the authors expect to be hard for complex domains such as ImageNet or MSCOCO §5.

Method

Background §3. The forward process adds Gaussian noise in small steps eq. 1; a noisy sample at any step can be written in closed form from the clean image eq. 2. A pretrained network predicts the reverse step, and its noise predictor is a UNet eq. 3 §3.

Steps of the representation extraction and segmentation pipeline:

1. Pick a diffusion step t. Corrupt the labelled image x0 to x_t with the closed form in eq. 2 §3.
2. Run the noise predictor on x_t once. Keep the outputs of the chosen UNet decoder blocks §3 fig. 1.
3. Upsample each kept activation to image size with bilinear interpolation, so each pixel gets a vector §3.
4. Repeat for the three steps t = 50, 150, 250 and concatenate all upsampled activations of the five blocks B = {5, 6, 7, 8, 12} into one 8448-dimensional vector per pixel §3.2.
5. Fix the noise: one noise draw is used for all steps §3.2. The paper does not say whether the same draw is reused across images; the code does (see Paper vs code).
6. Train an ensemble of independent MLPs on the pixel vectors and labels of the annotated images. The ensemble and training setup follow DatasetGAN §3.2 app. C.
7. For a test image, compute its pixel vectors the same way, let every MLP predict a label per pixel, and take the majority vote §3.2.

Analysis that motivated the choice of steps and blocks §3.1. The authors train an MLP on pixel features from each decoder block at each diffusion step and measure mean IoU on held-out images fig. 2. Later diffusion steps carry more semantic information, and early steps are weak §3.1. Middle decoder blocks are the most informative overall §3.1. Shallow blocks help more for small classes and deep blocks for large ones fig. 3. K-means clusters of the features on FFHQ show coarse masks at block 6 and fine face parts at block 12 fig. 4 §3.1. The chosen blocks and steps are not tuned per dataset §3.2.

Pseudocode for feature extraction (the noise draw and the step loop follow the code, see Paper vs code):

# Per-pixel features from a pretrained DDPM  [§3.2]
# [code: src/feature_extractors.py#L108-L123]
# Input:  image x0 (H, W, 3), scaled to [-1, 1]  [code: src/datasets.py#L12-L17]
# Output: features (H, W, 8448)

def extract_features(image, noise, steps, blocks):
    # One noise draw is reused for every step  [§3.2] [code: train_interpreter.py#L36-L41]
    # Collect one upsampled activation per (step, block) pair  [§3.2]
    features = []
    for step in steps:  # steps = [50, 150, 250]  [§3.2]
        # Corrupt x0 to step t with the closed form of q(x_t | x_0)  [eq. 2]
        noisy_image = (
            sqrt(alpha_bar[step]) * image
            + sqrt(1 - alpha_bar[step]) * noise
        )
        # One forward pass of the noise predictor; keep decoder outputs  [§3]
        decoder_outputs = unet_decoder(noisy_image, step)
        for block in blocks:  # blocks = [5, 6, 7, 8, 12]  [§3.2]
            # Bilinear upsampling to H x W gives per-pixel values  [§3]
            features.append(upsample_bilinear(decoder_outputs[block], H, W))
    # Stack along channels: 3 steps x 5 blocks  [§3.2]
    return concatenate(features, axis=channels)

Pseudocode for the ensemble and the prediction:

# Train an ensemble of independent pixel classifiers  [§3.2] [app. C]
# Input:  pixel features (N, 8448), pixel labels (N,) from annotated images
# Output: list of 10 trained MLPs

def train_ensemble(pixel_features, pixel_labels, num_classes):
    models = []
    for index in range(10):  # ensemble size 10  [app. C]
        # Two hidden layers of 128 and 32 units if fewer than 30 classes,
        # otherwise 256 and 128; ReLU and batch norm after each  [app. C]
        model = make_mlp(input_dim=8448, num_classes=num_classes)
        # Adam, learning rate 0.001, batch size 64  [app. C]
        model = fit_adam(model, pixel_features, pixel_labels, batch_size=64)
        models.append(model)
    return models

# Predict a label for every pixel of a test image  [§3.2]
# Input:  trained ensemble, pixel features (H * W, 8448)
# Output: label map (H, W)

def predict(models, pixel_features):
    # Each MLP gives its most likely label per pixel  [code: src/pixel_classifier.py#L81-L96]
    votes = [argmax(model(pixel_features), axis=1) for model in models]
    # Majority vote across the ensemble for each pixel  [§3.2]
    # [code: src/pixel_classifier.py#L105-L106]
    return majority_vote(votes)

Key equation, the closed-form noisy sample used in step 1 eq. 2:

x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1-\bar{\alpha}_t}\, \epsilon, \qquad \epsilon \sim \mathcal{N}(0, I) eq. 2

Hyperparameters

Results

Author-reported. Metric: mean IoU, averaged over 5 independent runs for different data splits §4 tab. 2. Datasets: Bedroom-28, FFHQ-34, Cat-15 and Horse-21 are annotated by the authors; CelebA-19 and ADE-Bedroom-30 use public annotations §4. The starred columns use models pretrained on other data: CelebA-19 uses FFHQ-256 models and ADE-Bedroom-30 uses LSUN-Bedroom models §4 tab. 2. Hardware: not stated.

Main result tab. 2:

Observations §4:

• The DDPM method is the highest mean IoU in every column of tab. 2.
• MAE is the closest competitor, and is comparable on FFHQ-34 and Cat-15 §4.
• SwAV is weaker, which the authors attribute to discriminative training suppressing fine detail §4.
• DatasetDDPM beats DatasetGAN on most benchmarks, which the authors attribute to higher synthetic quality §4.

Source of training data tab. 3. Training on real annotations is best where generator fidelity is lower, for example Cat-15 §4.1:

Sample efficiency tab. 4: DDPM with fewer annotated images §4.1:

The numbers in parentheses are training images. The authors say DDPM still beats most baselines of tab. 2 with the reduced supervision §4.1.

Feature-extraction randomness tab. 5. Sharing noise between train and test, or across steps, changes mIoU only marginally §4.1:

Robustness fig. 6 §4.1: mean IoU over 18 corruption types at severity levels 1, 3 and 5 on Bedroom-28 and Horse-21. Classifiers trained on clean images keep their advantage over SwAV and MAE at all levels. Figure content is not reproduced here.

Appendix results:

• DatasetDDPM and DatasetGAN, from 10K to 50K synthetic images tab. 6:

The authors say both saturate at 30K to 50K synthetic images app. B.

• Wider or deeper MLPs in the ensemble tab. 7: Bedroom-28 and FFHQ-34 move by at most 0.2 points app. C. Original 49.4 / 59.1; wider 49.5 / 59.1; deeper 49.3 / 58.9.

• Per-class IoUs for DatasetGAN, DatasetDDPM and DDPM on Horse-21, Cat-15, FFHQ-34 and Bedroom-28 are in fig. 9 app. D.

Paper vs code

Official code: https://github.com/yandex-research/ddpm-segmentation, linked from the paper's abstract. It is maintained by the authors' lab (Yandex Research). MIT license. Checked at commit 9a13776bb78d4753df6a561749f3274cd158297a (committed 2023-04-08).

Discrepancies:

• Training length. The paper says each MLP trains for about 4 epochs app. C. The code trains for up to 100 epochs train_interpreter.py and stops early once the training loss has not improved for 50 consecutive iterations, counted after epoch 3 train_interpreter.py. There is no fixed-epoch path in the code.
• Noise sharing. The paper says the noise is fixed across steps §3.2. The code draws one noise tensor of shape (1, 3, 256, 256) and reuses it for every training and test image, at every step train_interpreter.py train_interpreter.py. The paper does not state that the draw is shared across images.
• Number of runs. The paper reports 5 independent runs over data splits §4. The code exposes one seed argument, default 0 train_interpreter.py, and I found no loop over seeds or splits.
• Bedroom-28 label indexing. The config has 29 output classes and ignores label 0 ddpm.json. The paper lists 28 class names app. E. How the names map to labels 1 to 28 is not stated.
• Small-annotation filter. Classes with 1 to 19 pixels in a training image are set to the ignore label. The paper does not mention this train_interpreter.py.
• The SwAV row. The paper's SwAV uses a twice-wider ResNet-50 §4. The paper's SwAV row matches the README's SwAVw2 row, not its SwAV row README.md. The only differing value is ADE-Bedroom-30's standard deviation: ±1.6 in the paper, ±1.0 in the README.
• DatasetGAN standard deviations. The README gives ±2.7 for Bedroom-28 and ±1.0 for FFHQ-34 README.md; the paper gives ±2.3 and ±1.1 tab. 2.

Checked and consistent: diffusion steps {50, 150, 250}, decoder blocks {5, 6, 7, 8, 12}, feature dimension 8448, dataset sizes tab. 1, MLP widths, ensemble size 10, Adam with learning rate 0.001, batch size 64, majority vote, and bilinear upsampling. The README's dated note on the earlier configuration (four steps, blocks 6 to 9) predates the arXiv v3 text, which uses the current values README.md.

Pitfalls

• Training stop. "About 4 epochs" is not implementable as stated. Either follow the code's early stopping, which watches a noisy single-batch loss with a 50-iteration patience, or fix the epoch count and expect different results.
• Memory. Each image gives 256 x 256 x 8448 float32 values, about 2.2 GB by calculation. The code stores all training pixels, then permutes and reshapes the tensor, which copies it train_interpreter.py train_interpreter.py. The README reports about 210 GB for 50 training images README.md. Feature storage, not training, is the bottleneck.
• Hidden compute. Each image needs three full UNet forward passes, one per step, and each pass yields all five blocks at once feature_extractors.py.
• Per-block channels. The paper gives only the total of 8448. Across three steps the five blocks together contribute 2816 channels, by calculation. The split per block is not stated, and I did not derive it from the model.
• Timestep convention. The paper does not state the total number of diffusion steps or the index base. The code uses 1000 steps with a linear schedule train_interpreter.sh and passes the integer step directly feature_extractors.py.
• Block indexing. The code indexes the decoder's output blocks from 0 feature_extractors.py. The paper's figures number blocks "from deep to shallow" §3.1, which does not fix the base.
• Half precision. The code runs the diffusion model in fp16 when the flag is set feature_extractors.py. The paper does not mention it.
• Pooling. The code pools the pixels of all training images into one classifier training set and drops ignored pixels train_interpreter.py. The paper does not state pooling.
• Ties. The paper says "majority voting". The code takes the most frequent label with torch.mode pixel_classifier.py. Tie-breaking is not stated in the paper.
• Metric details. The code's mean IoU averages over all non-ignored class indices, including classes absent from both prediction and label, which score 0 pixel_classifier.py. Pixels labelled as the ignore value are not removed from the evaluation, so predicting a class there counts as a false positive pixel_classifier.py.
• Table bolding. In tab. 2 the MAE values 58.8 (FFHQ-34) and 52.4 (Cat-15) are bold, but DDPM's values in the same columns are higher, 59.1 and 53.7. Read the numbers, not the bold.
• Cross-column comparisons. CelebA-19 and ADE-Bedroom-30 use models pretrained on other data §4 tab. 2, so columns are not strictly like-for-like.
• Unused uncertainty. The code computes an ensemble uncertainty score from entropy, but it only prints it and does not use it in prediction pixel_classifier.py train_interpreter.py.

Lineage

• Builds on: DDPM [sohl2015deep, ho2020denoising] for the diffusion framework §2; the ADM pretrained models §4; DatasetGAN for the pixel classifier ensemble and its training setup §3.2 app. C.
• Compared against: MAE (strongest baseline in tab. 2); SwAV; GAN Inversion and GAN Encoder (e4e); ALAE; VDVAE; DatasetGAN; and DatasetDDPM, the authors' variant that trains DeepLabV3 on DDPM synthetic data §4.
• Superseded by: not known from content in this wiki.