PaperlayerView exact revision on arXiv

Source-linked AI summary

Label-Efficient Semantic Segmentation with Diffusion Models

Dmitry Baranchuk, Ivan Rubachev, Andrey Voynov, Valentin Khrulkov, Artem Babenko

arXiv:2112.03126v3cs.CVcs.LG

TL;DR

The paper addresses whether DDPMs can provide effective representations for semantic segmentation when labeled data is scarce. It extracts intermediate U-Net activations, builds pixel-level features, and trains a simple classifier; the resulting method outperforms alternatives in few-shot settings and remains robust under severe distortions.

  • Problem

    It was unclear whether DDPMs could serve as representation learners for discriminative vision problems such as semantic segmentation.

  • Method

    The method extracts intermediate DDPM U-Net activations, upsamples and concatenates them into pixel-level representations, then trains an ensemble of MLPs for segmentation.

  • Results

    The DDPM-based method significantly outperforms alternatives on most datasets, including with significantly less supervision, and preserves an advantage under severe image distortions.

  • Takeaways & Limitations

    DDPMs can serve as representation learners for discriminative computer vision and provide state-of-the-art performance in few-shot semantic segmentation.

Abstract

from arXiv · show

Denoising diffusion probabilistic models have recently received much research attention since they outperform alternative approaches, such as GANs, and currently provide state-of-the-art generative performance. The superior performance of diffusion models has made them an appealing tool in several applications, including inpainting, super-resolution, and semantic editing. In this paper, we demonstrate that diffusion models can also serve as an instrument for semantic segmentation, especially in the setup when labeled data is scarce. In particular, for several pretrained diffusion models, we investigate the intermediate activations from the networks that perform the Markov step of the reverse diffusion process. We show that these activations effectively capture the semantic information from an input image and appear to be excellent pixel-level representations for the segmentation problem. Based on these observations, we describe a simple segmentation method, which can work even if only a few training images are provided. Our approach significantly outperforms the existing alternatives on several datasets for the same amount of human supervision.

1 INTRODUCTION

The paper asks whether DDPMs can learn useful representations for discriminative vision, focusing on semantic segmentation with scarce labels. It finds that DDPM representations capture high-level semantics and support few-shot segmentation that outperforms alternatives.

  • DDPMs had not previously been established as effective representation learners for discriminative computer vision.
  • The paper investigates intermediate U-Net activations from DDPM reverse-diffusion steps as semantic image representations.
  • The proposed semantic segmentation method exploits these representations and works with only a few labeled images.
  • These activations capture high-level semantic information valuable for downstream vision tasks.
  • The DDPM-based method outperforms existing alternatives in the few-shot operating point and shows advantages over GAN-based counterparts.

2 RELATED WORK

Prior generative-model segmentation work is mainly GAN-based, with limitations in general semantic segmentation. This paper positions DDPM representations as an alternative and compares them against GAN-based methods.

  • Generative-model image segmentation has mainly relied on GANs.
  • Some GAN latent-space methods support binary segmentation but their applicability to general semantic segmentation remains unclear.
  • Other GAN methods use intermediate representations to train pixel-class predictors and demonstrate label efficiency.
  • The paper compares its DDPM-based method with GAN-based representations on the same datasets.
  • Generative models have also been studied as representation learners for global and dense prediction tasks, though less extensively than alternative unsupervised methods.

3 REPRESENTATIONS FROM DIFFUSION MODELS

The paper extracts pixel-level representations from intermediate DDPM U-Net activations at selected diffusion steps and decoder blocks. Their semantic usefulness varies across the reverse process, object sizes, and network depth, motivating a few-shot segmentation pipeline.

  • Extracting representations: A noisy image x_t is passed through the DDPM noise predictor, whose intermediate U-Net activations are upsampled to image resolution as pixel-level representations.
  • Representation analysis: Later reverse-diffusion steps typically provide more informative features than early steps, while middle U-Net decoder blocks are generally most discriminative.
  • Representation analysis: Large objects become predictable earlier and favor deeper blocks, whereas small objects favor shallower blocks; middle blocks remain most discriminative.
  • Representation analysis: K-means clusters from middle blocks span coherent semantic objects and parts, while block depth trades coarse masks against fine-grained face-part discrimination.
  • Few-shot semantic segmentation: The segmentation setup uses many unlabeled domain images and only a few labeled images, with the diffusion model trained unsupervised on the full image collection.
  • Few-shot semantic segmentation: The method concatenates upsampled activations from selected decoder blocks and later diffusion steps, then trains an ensemble of MLPs for pixel classification.

4 EXPERIMENTS

The experiments compare DDPM-based segmentation with generative and self-supervised alternatives across several few-shot datasets. DDPM representations generally outperform competing methods, remain effective with less supervision, and show robustness under distortions.

  • Methods and evaluation: The proposed DDPM method is compared with synthetic-data approaches, self-supervised models, GAN encoders, and VDVAE using mean IoU.Table 2 reports results averaged over five independent runs across different data splits.
  • Main results: The segmentation method is more straightforward than DatasetGAN and DatasetDDPM because it avoids synthetic-dataset generation and subsequent segmentation-model training.It directly exploits DDPM representations rather than requiring those auxiliary pipeline stages.
  • Sample-efficiency: DDPM remains able to outperform most baselines with significantly less supervision in the sample-efficiency experiment.The study evaluates reduced annotated-data settings using mean IoU on four datasets.
  • Robustness and feature extraction: Stochastic feature extraction changes performance only marginally, while DDPM preserves its advantage over SwAV and MAE under severe image distortions.The stochastic-feature comparison contrasts deterministic and multiple-noise extraction strategies.

5 CONCLUSION

The paper establishes DDPMs as representation learners for discriminative vision and reports strong few-shot segmentation performance, while requiring high-quality diffusion models for each target dataset.

  • DDPMs can serve as representation learners for discriminative computer vision problems.
  • DDPM representations can be computed directly for real images without learning an additional image-to-latent encoder.
  • The method achieves state-of-the-art performance in few-shot semantic segmentation.
  • High-quality diffusion models trained on the target dataset are required, limiting application to complex domains such as ImageNet or MSCOCO.

A EVOLUTION OF PREDICTIVE PERFORMANCE

Figures 7 and 8 track predictive performance across DDPM UNet blocks, diffusion steps, datasets, and semantic classes grouped by average area.

  • Figure 7 shows predictive performance evolving across different UNet blocks and diffusion steps on LSUN-Cat and LSUN-Bedroom.
  • Figure 8 shows predictive performance on FFHQ-256, LSUN-Cat, and LSUN-Bedroom for classes with the smallest and largest average areas.

B DATASETDDPM & DATASETGAN SATURATION

DatasetDDPM and DatasetGAN performance saturates when the synthetic training set reaches approximately 30K–50K images.

  • Mean IoU for DatasetDDPM and DatasetGAN saturates at 30K–50K synthetic training images.The table compares synthetic datasets ranging from 10K to 50K images.

C TRAINING SETUP

The training setup uses a fixed ensemble of MLP classifiers, while increasing MLP width or depth does not noticeably improve performance on the evaluated datasets.

  • The ensemble contains 10 independent MLPs trained with Adam for approximately 4 epochs at a 0.001 learning rate and batch size 64.This setting is used for all methods and datasets.
  • MLPs use two ReLU and batch-normalized hidden layers, sized 128 and 32 for fewer than 30 classes or 256 and 128 otherwise.
  • Twice wider or deeper MLPs show no noticeable performance difference on Bedroom-28 and FFHQ-34.

D PER CLASS IOUS

The section presents per-class intersection-over-union results for three methods and class-instance statistics for annotated real and synthetic training sets.

  • Figure 9 reports per-class IoUs for DatasetGAN, DatasetDDPM, and DDPM.
  • The per-class evaluation and class-instance statistics provide complementary views of segmentation performance and training-set composition.
  • Figure 10 reports the number of instances for each semantic class in the annotated real and synthetic training sets.

E DATASET DETAILS

The datasets span indoor scenes, faces, animals, and celebrity images, with class statistics computed for annotated real and GAN- or DDPM-generated synthetic images.

  • The study defines label sets for Bedroom-28, FFHQ-34, Cat-15, Horse-21, CelebA-19, and ADE-Bedroom-30.These sets contain 28, 34, 15, 21, 19, and 30 listed semantic classes, respectively.
  • Bedroom-28 and ADE-Bedroom-30 represent indoor scenes, while FFHQ-34 and CelebA-19 represent faces and Cat-15 and Horse-21 represent animals.
  • Figure 10 computes class statistics over annotated real images and annotated synthetic images produced by GAN and DDPM.

F EXTRACTING REPRESENTATIONS FROM MAE

The representation pipeline extracts and concatenates deep encoder features from a fully observed 256-resolution image, while alternative extraction strategies offer no improvement.

  • The method extracts feature maps from the deepest 12 ViT-L blocks after applying the model to a fully observed 256-resolution image.Each block produces 1024×32×32 features.
  • Upsampling the feature maps to 256×256 and concatenating them yields a 12288-dimensional pixel representation.
  • Including decoder activations produced no noticeable gains, while post-self-attention and every-second-encoder-block extraction performed slightly worse.
Loading 2112.03126v3…