Sample set dual reinforcement learning improves semantic segmentation dataset distillation
S$^2$-RL: Sample-Set Dual Reinforcement Learning for Generative Semantic Segmentation Dataset Distillation
Computer Vision and Pattern Recognition
Summary
Training machines to recognize and label every pixel in images (semantic segmentation) usually needs large datasets, which are slow and costly to handle. The authors propose a new technique called S2-RL that uses AI image generation combined with reinforcement learning to create small but rich sets of labeled images. This approach avoids expensive computations and generates diverse, meaningful scenes that help teach machines better. Tests on standard datasets show this method improves performance compared to earlier approaches.
What this means in practice
- •For computer vision engineers: Create compact training datasets with diverse, informative images for semantic segmentation tasks, reducing computational resources needed.
- •For ai model trainers: Generate custom image-label pairs using textual prompts for training segmentation models without requiring large original datasets.
Authors
Haoyu Wang, Fei Zhou, Qingqing Qiu, Lei Zhang, Wei Wei, Chen Ding
Abstract
While dataset distillation has witnessed significant advances in image classification, its extension to dense prediction tasks such as semantic segmentation is plagued by two core bottlenecks: bi-level optimization-based approaches suffer from prohibitive computational costs stemming from pixel-wise gradient unrolling, whereas proxy-based generative methods merely infuse visual textures into fixed semantic masks derived from real data, thus failing to condense rich semantic knowledge into a compact set of informative, novel scene generations. To address both bottlenecks simultaneously, we propose a novel Sample-Set Reinforcement Learning (S$^2$-RL) framework for generative semantic segmentation dataset distillation (SSDD). Specifically, S$^2$-RL formulates SSDD as a diffusion model based text-to-image generation paradigm, enabling flexible generation of images with informative, novel semantic distributions via a concise text prompt. Furthermore, we fine-tune a diffusion policy via Group Relative Policy Optimization (GRPO), leveraging a sample-set dual reward paradigm: for individual samples, it enforces semantic alignment with input text prompts and intra-class feature diversity; for sample sets, it maximizes inter-sample semantic distribution diversity by solving a contextual multi-armed bandit problem. These design choices enable S$^2$-RL to distill large-scale semantic segmentation datasets into a compact set of generated samples that encapsulate diverse, informative semantic knowledge without the need of pixel-wise gradient unrolling. Extensive experiments on the ADE20K and COCO datasets demonstrate that S$^2$-RL achieves substantial and consistent improvements over state-of-the-art baselines, establishing a new benchmark for SSDD.