SegDiff: Image Segmentation with Diffusion Probabilistic Models

In plain words

Image segmentation means drawing a mask over the pixels that belong to an object, such as a building in an aerial photo or a cell in a microscope slide. This paper uses the kind of diffusion model that generates pictures and points it at that job instead. The model starts from random noise shaped like the mask and removes the noise in many small steps. At every step it also looks at the input image, so the noise it removes is guided toward the correct mask.

The image goes through its own network, and the resulting features are added to the features of the current noisy mask. The whole system is trained on the segmentation data from scratch, without a model that was first trained on other images. Because each run starts from different noise, the output varies from run to run. The authors run the process several times and average the results, which improves accuracy.

The authors report better scores than earlier methods on street scenes, aerial building images and microscope images. Those numbers are their own and come from specific settings, so they are not a neutral ranking.

TL;DR

The paper turns a denoising diffusion model into an image segmenter. The noise predictor is conditioned on the input image by adding image features to the features of the noisy mask, and the whole model is trained end-to-end without a pretrained backbone §4 eq. 17. At inference, several independent reverse chains are averaged and thresholded §4.1. Author-reported results are state of the art on Cityscapes per-object masks, Vaihingen building masks and MoNuSeg nuclei tab. 1 tab. 2 tab. 3. The official repository has no LICENSE file and uses 9 chains where the paper says 30, so check the Paper vs code section before reproducing.

Problem

Diffusion models are strong generators, but before this work they were rarely used on tasks with one unique correct answer such as segmentation §1. Segmentation is usually solved by encoder-decoder networks, many of them built on pretrained backbones §1 §2. The authors want to test whether a diffusion model can compete when trained end-to-end without that backbone §1.

Cityscapes is evaluated per object: each sample is a crop around one object's bounding box, and the model returns that object's mask §5, Datasets. The paper frames this as interactive segmentation, motivated by the need to speed up object annotation §5, Datasets.

A diffusion sampler needs many sequential steps per output. The paper notes that its inference time grows linearly with the number of chains and that it is slower than Segformer and STDC §5, Results.

Method

The method keeps the standard diffusion process over a binary segmentation map. The forward process adds Gaussian noise to the ground-truth mask over T steps eq. 4 eq. 5. The reverse process removes noise step by step using a learned noise predictor eq. 7 eq. 16. The change is in the predictor: it also receives the input image §4.

Conditioning on the image (§4)

The noise predictor is written as D(E(F(x_t) + G(I), t), t) eq. 17. The three networks are:

• G, the image encoder. It is a convolution, then a stack of Residual-in-Residual Dense Blocks (RRDB), then a convolution, a leaky ReLU and an output convolution §4.3 RRDB.py.
• F, the mask encoder. A single convolution that takes the noisy mask (one channel) to C channels §4.3 script_util.py.
• E and D, the encoder and decoder of a conventional U-Net. D has skip connections from every E layer §4.3.

F(x_t) and G(I) are summed, not concatenated, and the sum enters E §4 §5, Ablation. The timestep t enters both D and E §4.

Training (§4.2)

Each step draws an image and its mask, a timestep t uniformly from 1 to T, and Gaussian noise. The noisy mask x_t is built with the closed-form forward step eq. 5. The loss is the mean squared error between the noise and the prediction, conditioned on the image eq. 18 alg. 2.

Inference and multiple chains (§4.1)

Inference starts from Gaussian noise and applies the reverse update for T steps eq. 16 alg. 1. The noise added during each step makes separate runs disagree fig. 2. The method runs the chain several times and averages the outputs. The paper uses 30 chains in all experiments except the ablation §4.1. Averaging improves mIoU and the calibration score §5 fig. 5 fig. 6.

Architecture (§4.3)

Each U-Net level has residual blocks. Attention layers follow each block at the 16x16 and 8x8 resolutions, and the bottleneck has two residual blocks with attention between them. Each residual block takes the time embedding through a linear layer, SiLU and another linear layer §4.3. The encoder downsamples with a stride-2 convolution. The decoder upsamples with nearest interpolation followed by a convolution §4.3.

Pseudocode

The first block is the conditioned noise predictor, which is the paper's only architectural change eq. 17 §4.

# Noise predictor eps_theta(x_t, I, t)  [eq. 17] [§4]
# x_t: (B, 1, H, W) noisy mask; image: (B, 3, H, W); t: timestep

def predict_noise(x_t, image, t):
    # Encode the image with the RRDB network G  [§4.3]
    image_features = image_encoder(image)        # (B, C, H, W)
    # Encode the noisy mask with the single convolution F  [§4.3]
    mask_features = mask_encoder(x_t)            # (B, C, H, W)
    # Sum the two streams, then run the U-Net encoder E  [§4] [eq. 17]
    h, skips = encoder(mask_features + image_features, t)
    # Run the decoder D with skip connections, timestep t  [§4.3] [eq. 17]
    return decoder(h, skips, t)

The second block is one training step alg. 2 §4.2 eq. 18. The mask is encoded as -1 and +1 in the released code, which the paper does not state city.py.

# One training step  [alg. 2] [§4.2] [eq. 18]
# image: (3, H, W); mask: (1, H, W), values in {-1, +1}

def training_loss(model, image, mask, num_steps):
    # Draw a timestep uniformly from 1 to T  [§4.2] [alg. 2]
    t = sample_uniform_int(1, num_steps)
    # Draw Gaussian noise with the shape of the mask  [alg. 2]
    noise = sample_gaussian(mask.shape)
    # Noise the ground-truth mask to step t  [eq. 5] [alg. 2]
    noisy_mask = (
        sqrt(alpha_bar[t]) * mask
        + sqrt(1 - alpha_bar[t]) * noise
    )
    # Predict the noise, conditioned on the image  [eq. 17] [§4.2]
    predicted = model(noisy_mask, image, t)
    # Mean squared error between the noise and the prediction  [eq. 18]
    return mean((noise - predicted) ** 2)

The third block samples several chains and averages them alg. 1 §4.1. The released code maps each chain to [0, 1], clamps it, averages and rounds sampling_util.py.

# Sample one mask chain  [alg. 1] [eq. 16]
# Output: (1, H, W) values in [-1, 1]

def sample_chain(model, image, num_steps):
    # Start from standard Gaussian noise  [alg. 1]
    x_t = sample_gaussian(mask_shape)
    for t in range(num_steps, 0, -1):
        # Add noise only for t > 1  [alg. 1]
        z = sample_gaussian(mask_shape) if t > 1 else zeros
        # Predict the noise for this step  [eq. 17]
        eps = model(x_t, image, t)
        # Reverse update  [eq. 16] [eq. 13]
        x_t = (x_t - beta[t] / sqrt(1 - alpha_bar[t]) * eps)
        x_t = x_t / sqrt(alpha[t]) + sqrt(sigma_sq[t]) * z
    return x_t

# Average N chains and threshold  [§4.1] [§5]
# The released code averages on [0, 1] and rounds
# [code: improved_diffusion/sampling_util.py#L155-L163]

def segment(model, image, num_steps, num_chains):
    # Map each chain to [0, 1] and clamp
    # [code: improved_diffusion/sampling_util.py#L155-L160]
    chains = [
        clamp((sample_chain(model, image, num_steps) + 1) / 2, 0, 1)
        for _ in range(num_chains)
    ]
    # Average the chains, then round to a binary mask  [§4.1]
    # [code: improved_diffusion/sampling_util.py#L163]
    return round(mean(chains))

Hyperparameters

Values are from the paper unless the source column says code. Where the paper is silent, the code value is shown and labelled.

Results

Author-reported. Setup: Cityscapes per-object crops of eight classes, with one model per class and Cityscapes validation as the evaluation set; Vaihingen 68 test images; MoNuSeg 14 test images at 512x512 §5. Baselines are mostly taken from earlier papers, and several use ImageNet, PASCAL or COCO pretraining. The paper trains its own model from random initialization §5.

The paper states that its method beats every baseline in every Cityscapes class in both settings §5, Results. Baseline lists: Cityscapes, Polygon-RNN++, PSP-DeepLab, Polygon-GCN, Spline-GCN, Deep contour, Segformer-B5, STDC1 and STDC2 tab. 1. Vaihingen, FCN-UNet, FCN-ResNet34, FCN-HarDNet, DSAC, DarNet, Deep contour, TDAC, Segformer-B5, STDC1 and STDC2 tab. 2. MoNuSeg, FCN, U-Net, U-Net++, Res-UNet, A.A U-Net and MedT tab. 3.

Conditioning ablation, on Vaihingen with 9 chains rather than the main setting §5, Ablation: the full method scores F1 94.95 and IoU 90.64, against 91.60 and 85.45 for channel concatenation (variant one). Replacing RRDB with FC-HarDNet (variant two) and removing the image encoder (variant three) both lower the score. Adding the image features later in the U-Net (variants four to six) changes little for variant four, and variants five and six lose accuracy as the insertion point moves away from the first layer tab. 4 §5, Ablation.

Averaging chains: on Vaihingen and Cityscapes the mIoU rises with more chains, and the best count differs per class (10 chains for Cityscapes Bus, 3 for Train, 25 for MoNuSeg) §5 fig. 5. Reducing diffusion steps to 25 costs at most 2 mIoU on Vaihingen and 1 on Cityscapes Bus, which the paper reports as a roughly fourfold speedup §5 fig. 9.

Paper vs code

Code is the repository at commit e6592c983b5bae6eeaf714d91e4aaebc14ca7c21. Its owner handle matches the first author's name, and its README calls it the official repository of the paper. The README also says it is based on openai/improved-diffusion README.md.

• Chains per sample. Paper: 30 chains in all experiments except the ablation §4.1. Code: the major-vote count is hard-coded to 9 for every dataset, with no flag sampling_util.py. The paper's own ablation also uses nine §5, Ablation. The code cannot show which count produced the main tables.
• Beta schedule. Paper: beta_t runs linearly from 1e-4 to 2e-2 for any T eq. 14. Code: endpoints are multiplied by 1000/T, so they equal the paper's only at T = 1000 gaussian_diffusion.py. At T = 100, betas run from 1e-3 to 0.2.
• Sampling variance. Paper: sigma_t^2 = tilde beta_t eq. 13. Code: with learned sigma off and small-sigma off, the model uses the FIXED_LARGE variance, which is beta_t for steps above 1 script_util.py gaussian_diffusion.py.
• Time embedding. Paper: the timestep uses a shared learned look-up table in D and E §4. Code: a sinusoidal embedding followed by a two-layer MLP, shared by all blocks nn.py unet.py unet.py.
• Training batch size and RRDB blocks. Paper: totals of 30 (Cityscapes), 8 (Vaihingen) and 8 (MoNuSeg), with 15, 6 and 12 RRDB blocks §5. README examples give 2 x 15, 2 x 4 and 4 x 2 across processes, with RRDB counts 15, 6 and 12 README.md. These match if the paper's batch sizes are totals across processes. The paper does not say so.
• Learning rate. Paper: not stated. Code: default 2e-5, README examples 1e-4 image_train_diff_city.py README.md.
• Mask encoding. Paper: the mask is x_0 = M_i with no encoding stated §4.2. Code: masks are mapped to {-1, +1} for training city.py.
• Augmentation. Vaihingen values match the paper exactly vaih.py. Cityscapes and MoNuSeg match the listed transforms city.py monu.py.
• Metric. Paper: mIoU is a sum of per-class IoU terms over N classes eq. 19. Code: foreground IoU is computed per sample on the binary mask and averaged over samples sampling_util.py sampling_util.py sampling_util.py.
• Ablation variants. Paper: six variants with reported numbers tab. 4. Code: no code for variants one to six was found. Searching the Python files for concatenation or variant code returned nothing.
• Cityscapes class names. Paper: eight classes, including Truck tab. 1. README: the option list has seven names, with no truck and with "bike" README.md. Dataset class tuple: includes truck and bicycle city.py. The annotation JSON is not in the repository, so the key names cannot be checked.
• Not in the paper. Checkpoint selection on the evaluation split train_util.py. See Pitfalls.

Pitfalls

• Model selection on the reported split. During training the loop evaluates the validation loader every 5000 steps from 75000 (and at 60000) and keeps the checkpoint with the best mean IoU train_util.py. For Vaihingen and MoNuSeg, the loader's "val" mode loads the test files vaih.py monu.py. So the numbers in the tables may be selected on the test data. The paper does not say how its reported checkpoint was chosen §5. Use a separate split before comparing.
• Cost scales with chains and steps. Each mask needs N full reverse chains of T network evaluations. The code uses N = 9 and the paper's main setting uses 30 §4.1 sampling_util.py. Plan for the 30-chain cost if you want the paper's numbers.
• Sampling variance choice. Use the paper's tilde beta or the code's beta, and say which. The choice changes the reverse-step noise and the paper gives no ablation of it.
• Beta schedule depends on T in code. Fixed 1e-4 to 2e-2 endpoints (as in the paper) and T-scaled endpoints (as in the code) differ for any T other than 1000 eq. 14 gaussian_diffusion.py.
• Rounding at 0.5. The code rounds the average with torch.round, which sends exact halves to the nearest even value. With the 9 chains in the code, an exact half cannot occur. With an even chain count it can.
• Clamping. Chains are clamped to [0, 1] before averaging, and the result is bilinearly resized to the ground-truth size if the shapes differ sampling_util.py.
• Sampling clip setting is inherited. The sample scripts default to clip_denoised = True, but they overwrite it from the training args.json, where it is False image_sample_diff_city.py image_train_diff_city.py. Set it explicitly.
• RRDB sensitivity does not match the main configuration. The paper's main Cityscapes model uses 15 RRDB blocks, but its sensitivity baseline for Bus is 10 and the reported optimum is 10 §5, Parameter sensitivity. For Vaihingen the main model uses 6 blocks, but the reported optimum is 3 §5. Treat these as separate studies.
• Figure 4 axis label. The caption says the x-axis is the number of diffusion steps, while the text calls it iterations fig. 4 §5.
• Loss equation typo. Eq. 18 takes the expectation over an undefined variable x_e. The implemented loss is the plain epsilon mean squared error eq. 18 gaussian_diffusion.py.
• Mixed precision and an unlisted optimizer setting. The released training uses fp16 with a gradient-scale growth of 1e-3, and none of this is stated in the paper image_train_diff_city.py image_train_diff_city.py.
• Pretraining caveat. Most baselines use ImageNet, PASCAL or COCO pretraining. This model is trained from random initialization on the target data only §5. Cross-method comparisons confound architecture with pretraining.
• Cityscapes is per instance. Each sample is one object crop, with the crop taken from ground-truth boxes. The results do not measure whole-image segmentation §5, Datasets.
• MoNuSeg tiling. The 512x512 test image is split into four non-overlapping 256x256 patches, and the code hard-codes stride 256 §5 sampling_util.py.
• No license. The repository has no LICENSE file, so its code cannot be reused. Cite by path and line only.

Lineage

• Builds on: the DDPM formulation §3, the improved-diffusion code base README.md and the U-Net design of Improved DDPM §4. The image encoder uses the RRDB block from the ESRGAN work §4.3.
• Compared against: the Cityscapes, Vaihingen and MoNuSeg baselines listed under Results tab. 1 tab. 2 tab. 3.
• Superseded by: not known from content already in this wiki.