Project · Generative Representation Learning

DiffuseSeg

Synthetic data and semantic segmentation from a single diffusion model, turning one DDPM into a labelled-data factory

Harishwar Rao J.

Code 🤗 Dataset 🤗 DDPM Weights Colab Demo

DiffuseSeg summary — A summary of the DiffuseSeg pipeline.

Overview

A diffusion model learns a lot about image structure while it is being trained to generate. DiffuseSeg reuses that instead of discarding it after sampling: the same features that let a DDPM draw a face also let it parse one.

An unconditionally trained DDPM[1] serves as a backbone for both high-fidelity image synthesis and label-efficient semantic segmentation. By extracting the rich multi-scale features from its U-Net decoder, a lightweight pixel-level segmentation head can be trained with very few labelled examples, effectively converting the generative model into a source of paired images and annotations.

> 0.35 mIoU on
CelebA-HQ

100 Labelled images
needed to train the head

0 Backbone parameters
trained for segmentation

The approach was inspired by Baranchuk et al.[2], who showed that diffusion features support label-efficient segmentation.

The Two-Stage Pipeline

Stage 1: Train an unconditional DDPM

A DDPM with a U-Net core is trained from scratch on unlabelled face images (CelebA-HQ[3] at 64×64), learning to reverse a diffusion process that turns Gaussian noise into a realistic face. The training scripts generalise to any dataset with minor dataset-specific configuration changes, or you can skip straight to stage 2 with a pretrained model.

Unconditional DDPM sampling on CelebA-HQ 64×64.

Stage 2: Train a segmentation head on frozen features

With the DDPM U-Net frozen, it becomes a pure feature extractor:

Feature extraction: for a given image (real or synthetic), activations are pulled from specific decoder blocks of the U-Net at selected timesteps (t = 50, 150, 250). The choice of all up-blocks and these timesteps was motivated by the source paper[2] and the architecture used here.
Pixel descriptors: the multi-scale feature maps are upsampled to a common resolution and concatenated, giving every pixel one rich feature vector.
MLP ensemble: an ensemble of small pixel-wise MLPs is trained on a handful of labelled images to classify each pixel descriptor into a semantic class (hair, skin, nose, …).

On the choice of timestep. Features at different noise levels encode different things: high-noise timesteps carry coarse layout, low-noise timesteps carry fine texture and boundaries. Concatenating across timesteps is what gives a single linear-ish head enough signal to segment without any convolutional decoder of its own.

Results

The segmentation head performs well on the CelebA-HQ validation set. That is itself a sign that the frozen diffusion features are semantically organised, since a nearly linear probe would not work otherwise.

End-to-end generation

A synthetic image produced by the DDPM, alongside the segmentation map produced by the MLP head:

End-to-end sample 1 — End-to-end: generated image + generated mask, with no human annotation anywhere in the loop.

End-to-end sample 2 — End-to-end: generated image + generated mask, with no human annotation anywhere in the loop.

Validation on real images

Validation sample 1 — Real image, CelebA-HQ ground-truth mask, and the DiffuseSeg prediction.

Validation sample 2 — Real image, CelebA-HQ ground-truth mask, and the DiffuseSeg prediction.

Released Artefacts

Colab notebook: run the pipeline end to end.
Synthetic dataset: image/mask pairs on Hugging Face.
DDPM weights: trained on resized CelebA-HQ 256 → here.
Segmentation head weights: here.

References

J. Ho, A. Jain, P. Abbeel. Denoising Diffusion Probabilistic Models. NeurIPS, 2020.
D. Baranchuk, I. Rubachev, A. Voynov, V. Khrulkov, A. Babenko. Label-Efficient Semantic Segmentation with Diffusion Models. ICLR, 2022.
Z. Liu, P. Luo, X. Wang, X. Tang. Deep Learning Face Attributes in the Wild (CelebA-HQ). ICCV, 2015.

DDPM Semantic Segmentation Feature Reuse Label Efficiency