Source-linked AI summary
SegDiff: Image Segmentation with Diffusion Probabilistic Models
Tomer Amit, Tal Shaharbany, Eliya Nachmani, Lior Wolf
TL;DR
Image segmentation presents a test of whether diffusion models can compete on tasks with a unique ground-truth result. The paper introduces an end-to-end diffusion segmentation method that sums separate encodings of the input image and current estimate, averages multiple generations, and achieves state-of-the-art results across diverse benchmarks.
Problem
Diffusion models had mainly been used for generation tasks, leaving their competitiveness on image segmentation with a unique ground-truth result uncertain.
Method
SegDiff learns segmentation end-to-end by summing encoded input-image and current-estimate representations in a diffusion denoising network, then averaging multiple stochastic generations.
Results
State-of-the-art segmentation results are reported on Cityscapes, Vaihingen building segmentation, and MoNuSeg.
Takeaways & Limitations
The approach achieves state-of-the-art results across street-view, aerial, and microscopy segmentation benchmarks.
Takeaways & Limitations
Direct comparison with PolyTrans-form [28] was not performed because it uses the entire image and does not operate on standard segmentation patches.
Abstract
from arXiv · showhide
Diffusion Probabilistic Methods are employed for state-of-the-art image generation. In this work, we present a method for extending such models for performing image segmentation. The method learns end-to-end, without relying on a pre-trained backbone. The information in the input image and in the current estimation of the segmentation map is merged by summing the output of two encoders. Additional encoding layers and a decoder are then used to iteratively refine the segmentation map, using a diffusion model. Since the diffusion model is probabilistic, it is applied multiple times, and the results are merged into a final segmentation map. The new method produces state-of-the-art results on the Cityscapes validation set, the Vaihingen building segmentation benchmark, and the MoNuSeg dataset.
1. Introduction
The paper extends diffusion models to image segmentation, addressing uncertainty about their competitiveness when segmentation has a unique ground-truth result. It proposes end-to-end conditioning and multiple stochastic generations, achieving state-of-the-art results across several benchmarks.
- Motivation: Diffusion models had rarely been applied to tasks with a unique ground-truth result, motivating their evaluation for image segmentation.The paper notes that diffusion models were used mainly for generation tasks evaluated by user studies or quality and diversity scores.
- Approach: The method applies a diffusion model to learn an image segmentation map end-to-end without a pre-trained backbone.This distinguishes the approach from recent segmentation methods that rely on pre-trained backbone networks.
- Approach: The input image and current binary segmentation estimate are encoded separately, summed, and passed through a U-Net to produce the next estimate.This provides the model’s proposed conditioning mechanism for iterative refinement.
- Approach: Multiple stochastic segmentation generations are averaged to improve overall accuracy and stabilize the final segmentation results.The experiments use thirty generated instances except in the ablation study measuring the gain from averaging.
- Results: State-of-the-art results are reported on Cityscapes, building segmentation, and nuclei segmentation benchmarks, with especially large margins on small datasets.The cited benchmarks include Cityscapes, Vaihingen building segmentation, and MoNuSeg.
2. Related work
The paper situates segmentation within established encoder-decoder and related conditional diffusion approaches, then introduces residual conditioning by adding encoded image information to the current segmentation estimate.
- Image segmentation: Image segmentation assigns each pixel a label indicating whether it belongs to a specific class.Related architectures include fully convolutional, encoder-decoder, U-Net, transformer-based, and hypernetwork-based models.
- Diffusion models: Diffusion models are generative models based on a Markov chain that transforms a simple distribution into samples from a complex data distribution.The paper also describes their ability to generate images competitive with or better than recent GAN methods.
- Conditional diffusion: Conditional diffusion models use an input such as an image, audio representation, text, or class embedding to guide generation.Prior image-conditioning methods include channelwise concatenation or convolutional processing before concatenation.
- SegDiff: SegDiff conditions diffusion differently by adding the encoded input image to the current segmentation estimate, forming a residual model rather than concatenating inputs.The paper describes this as a distinct conditioning approach for image segmentation.
3. Background
Diffusion models use forward and reverse Markov processes to transform simple noise into data samples. The paper describes their noise schedule, reparameterization, training, and inference procedures for segmentation.
- Diffusion models are generative models parametrized by a Markov chain with forward and backward processes.
- The forward process adds Gaussian noise over T steps, with βt controlling the schedule and In×n denoting the identity matrix.The schedule increases linearly from β1 = 10−4 to βT = 2 ∗ 10−2.
- The forward process can be sampled at an arbitrary timestamp t through a reparameterized formulation.
- The reverse process starts from pθ(xT) = N(xT; 0, In×n) and transforms the latent distribution toward the data distribution pθ(x0).Reverse steps use small Gaussian transitions.
- Training samples an image-mask pair, a diffusion timestamp, and Gaussian noise, while inference iterates from xT over the total diffusion steps.The training procedure computes xt from the ground-truth mask and optimizes the predicted noise; inference is detailed in Algorithm 1.
4. Method
The method conditions a diffusion denoising network on both the input image and the current segmentation estimate, then iteratively refines the map. Multiple stochastic inference runs are averaged to stabilize predictions and improve performance.
- The denoising network combines representations of the current estimate xt and input image I before U-Net processing.The two encoders produce tensors that are summed as F(xt) + G(I).
- The U-Net encoder is split into E, F, and G, where F encodes xt, G encodes I, and E processes their summed features.The decoder D remains conventional, and both E and D receive the embedded step index t.
- The conditioned output ϵθ replaces the unconditioned network in the reverse diffusion update to produce the next segmentation estimate.
- Stochastic reverse-process noise creates substantial variability across inference runs on identical inputs.
- Thirty inference instances are averaged in experiments to stabilize segmentation results and improve performance.The ablation study quantifies the gain from this averaging procedure.
- Training samples a noisy version of the ground-truth segmentation map and minimizes a noise-prediction loss conditioned on the paired image.The ground truth is used by setting x0 = Mi.
- The image encoder G uses residual-in-residual dense blocks, while F maps the single-channel segmentation estimate into C channels.
- The E-D U-Net uses residual blocks, attention at selected resolutions, downsampling, upsampling, and skip connections.
5. Experiments
Experiments evaluate SegDiff on Cityscapes, Vaihingen, and MoNuSeg using multiple protocols, metrics, baselines, and ablations. The method outperforms baselines across the reported benchmarks, with especially large gains on smaller datasets, while diffusion-step and instance-count studies examine efficiency and stability.
- Cityscapes: SegDiff outperforms all baseline methods across Cityscapes categories and both bounding-box protocols.The protocols use either a 15% bounding-box expansion or a tight bounding box.
- Benchmark results: The performance gap is especially large for datasets or classes with fewer training images.Figure 3 orders Cityscapes classes by their number of training images while normalizing mIoU relative to SegDiff.
- Benchmark results: SegDiff outperforms previous work on all four reported Vaihingen scores and both MoNuSeg segmentation metrics.MoNuSeg comparisons include recent U-Net and transformer variants developed specifically for the task.
- Instance averaging: Increasing generated instances generally improves mIoU, but the number needed for optimal performance varies across classes and datasets.Examples include 10 instances for Cityscapes “Bus”, 3 for “Train”, and 25 for MoNuSeg.
- Ablation study: Ablations favor summation conditioning and RRDB blocks, while removing the encoder significantly affects metrics.Changing the signal-integration position has negligible or task-specific effects, whereas other variants decrease performance.
- Diffusion-step analysis: Reducing diffusion steps to 25 preserves stable segmentation results, with degradation of up to 2 mIoU for Vaihingen and 1 mIoU for Cityscapes “Bus”.The reduction can speed performance by a factor of four, creating an accuracy–runtime tradeoff.
6. Conclusions
The paper applies diffusion models to image segmentation, using an encoder-conditioned U-Net to iteratively refine segmentation maps. It reports state-of-the-art results across diverse benchmarks, including street-view, aerial, and microscopy images.
- Diffusion models are used to incrementally improve the obtained segmentation.
- The input image is conditionally encoded through a second pathway whose activations are summed early in the U-Net encoder.
- The approach achieves state-of-the-art segmentation results across street-view, aerial, and microscopy benchmarks.