Diffusion-Driven Two-Stage Active Learning for Low-Budget Semantic Segmentation

TL;DR

Semantic segmentation under labeling budgets is addressed by a diffusion-driven two-stage active learning framework that decouples diversity from uncertainty. In Stage 1, MaxHerding selects a representative, diverse pixel pool using diffusion-based representations; in Stage 2, entropy-augmented DALD (eDALD) selects the most informative pixels by combining mutual information with an entropy term computed on noisy multi-scale diffusion features, improving robustness to noise. The method operates under an extreme budget where the total annotation budget is $B=N$, distributed across $R=10$ rounds so that $b=B/R=0.1N$, yielding about $N$ labeled pixels in total, and achieves near fully-supervised performance on four standard benchmarks (CamVid, ADE-Bed, Cityscapes, Pascal-Context). Experimental results show consistent gains over strong baselines across all datasets and budgets, with released code for reproducibility.

Abstract

Semantic segmentation demands dense pixel-level annotations, which can be prohibitively expensive - especially under extremely constrained labeling budgets. In this paper, we address the problem of low-budget active learning for semantic segmentation by proposing a novel two-stage selection pipeline. Our approach leverages a pre-trained diffusion model to extract rich multi-scale features that capture both global structure and fine details. In the first stage, we perform a hierarchical, representation-based candidate selection by first choosing a small subset of representative pixels per image using MaxHerding, and then refining these into a diverse global pool. In the second stage, we compute an entropy-augmented disagreement score (eDALD) over noisy multi-scale diffusion features to capture both epistemic uncertainty and prediction confidence, selecting the most informative pixels for annotation. This decoupling of diversity and uncertainty lets us achieve high segmentation accuracy with only a tiny fraction of labeled pixels. Extensive experiments on four benchmarks (CamVid, ADE-Bed, Cityscapes, and Pascal-Context) demonstrate that our method significantly outperforms existing baselines under extreme pixel-budget regimes. Our code is available at https://github.com/jn-kim/two-stage-edald.

Figures (9)

• Figure 1: Overview of the proposed pipeline (illustrative hyperparameters). (a) For each pixel $x$, the pre‐trained diffusion model yields a feature representation $\mathbf{z}_x$, an uncertainty estimate $\mathrm{eDALD}(x)$, and the class prediction $\hat{y}(x)$ (here shown with $T=2$ diffusion steps and $L=3$ layers purely for illustration). (b) These outputs drive our two‐stage selection: Stage 1 picks $K=5$ representative pixels per image to form a global pool of size $M=10$, then Stage 2 applies the uncertainty criterion to select the final $b=5$ pixels (all parameter values are example settings, not the exact experimental hyperparameters).
• Figure 2: Active learning progression: mIoU (%) vs. AL rounds on (a) CamVid, (b) ADE‐Bed, (c) Cityscapes, and (d) Pascal‐Context. Curves compare Random sampling, single‐stage uncertainty methods, and two‐stage variants.
• Figure 3: mIoU vs. AL rounds (budget $b=0.1N$). Horizontal gray lines denote fully supervised mIoUs, and green lines denote the $90\%$ thresholds. We use two-stage eDALD for rounds $1\!-\!10$ and switch to Margin thereafter.
• Figure 4: Qualitative pixel‐selection comparison. (a) Input image; (b) baseline selections; (c) our selections. Light blue dots indicate selected pixels. Our method covers boundaries and fine details more broadly.
• Figure 5: Hyperparameter sensitivity on CamVid: mIoU achieved after 10 AL rounds for varying per-image representative count $K\in\{20,50,100\}$ and global pool size $M\in\{0.25M_0,0.4M_0,0.5M_0\}$. Each subplot shows one $M$ setting.