D3S2: Diffusion-Guided Dataset Distillation for Semantic Segmentation

2026-05-24

AI Summary

This work addresses the challenges of long-tailed class imbalance, difficulty in pixel-level label alignment, and high computational cost of high-resolution optimization in semantic segmentation dataset distillation. The authors propose D3S2, a novel framework that introduces diffusion models to segmentation distillation for the first time. D3S2 employs a two-stage strategy: it first constructs a class-balanced mask set via a greedy algorithm, then leverages a pretrained layout-to-image diffusion model to generate spatially aligned synthetic images, enhanced by dual-objective guided sampling to improve data utility. Pixel-wise alignment and semantic consistency are effectively preserved through segmentation consistency and class feature matching losses. Under a 1% compression rate, D3S2 achieves mIoU scores of 24.99% on ADE20K and 35.49% on COCO-Stuff, substantially outperforming existing baselines.

Abstract

Dataset distillation (DD) aims to compress large-scale datasets into compact synthetic sets while preserving training efficacy. However, existing studies mainly focus on image classification, leaving dense prediction tasks such as semantic segmentation largely underexplored. In this work, we identify three key challenges for segmentation DD: (i) long-tailed class imbalance, (ii) the need for strict pixel-wise alignment between images and dense labels, and (iii) the high computational cost of optimizing high-resolution data with complex models. To address these challenges, we propose D3S2, a Diffusion-guided Dataset Distillation framework for Semantic Segmentation. Our method adopts a two-stage design. In Class-Balanced Mask Selection, we construct a representative mask set via a greedy strategy that prioritizes underrepresented classes. In Diffusion-Guided Image Synthesis, we employ a pretrained layout-to-image diffusion model to generate images conditioned on the selected masks, naturally ensuring spatial alignment. To further enhance the training utility of synthesized data, we introduce guided diffusion sampling with two complementary objectives: a segmentation-consistency loss for pixel-level alignment, and a class-wise feature matching loss for aligning per-class feature statistics across layers. Extensive experiments demonstrate the superiority of D3S2. Notably, at an extremely compression rate of 1%, our method achieves 24.99% and 35.49% mIoU on ADE20K and COCO-Stuff with Mask2Former (Swin-S), outperforming random selection by 9.34% and 5.70%, respectively.

Problem

• dataset distillation
• semantic segmentation
• class imbalance
• pixel-wise alignment
• computational cost

Innovation

• dataset distillation
• semantic segmentation
• diffusion model
• class imbalance
• pixel-wise alignment