S-RL: Sample-Set Dual Reinforcement Learning for Generative Semantic Segmentation Dataset Distillation
Organizations: Northwestern Polytechnical University · The Chinese University of Hong Kong · Zhejiang University · Xi’an University of Posts & Telecommunications
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-RL) framework for generative semantic segmentation dataset distillation (SSDD). Specifically, S-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-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-RL achieves substantial and consistent improvements over state-of-the-art baselines, establishing a new benchmark for SSDD.
Figures & tables
| Method | 0.625‰ | 1.25‰ | 2.5‰ | |||
| mAcc | mIoU | mAcc | mIoU | mAcc | mIoU | |
| Random | 28.82 1.32 | 19.71 0.67 | 38.00 0.40 | 27.75 0.46 | 44.89 0.40 | 34.31 0.37 |
| Uniform | 33.45 0.83 | 23.37 0.97 | 43.37 0.97 | 32.18 0.75 | 50.07 0.43 | 37.39 0.46 |
| K-Center | 34.77 0.67 | 24.99 0.57 | 43.35 0.72 | 32.30 0.46 | 48.92 0.47 | 37.19 0.38 |
| Herding | 32.07 0.35 | 23.04 0.24 | 40.20 0.25 | 30.49 0.22 | 47.56 0.23 | 37.10 0.20 |
| SD3.5 Baseline | 38.18 0.73 | 25.98 0.67 | 46.65 0.48 | 32.88 0.87 | 51.81 0.35 | 37.74 0.34 |
| Model Configuration | Ratio = 0.25% | Ratio = 0.5% | Ratio = 1.0% | Full | |||
| Uniform | S 2 -RL (Ours) | Uniform | S 2 -RL (Ours) | Uniform | S 2 -RL (Ours) | ||
| (a) Impact of Model Capacity | |||||||
| ResNet-18 + UPerNet | 5.32 0.20 | 5.90 0.29 0.58 | 8.33 0.42 | 9.12 0.36 0.79 | 12.10 0.61 | 12.45 0.25 0.35 | 35.58 |
| ResNet-50 + UPerNet | 6.72 0.35 | 7.69 0.36 0.97 | 10.25 0.39 | 11.84 0.51 1.59 | 14.83 0.57 | 15.43 0.20 0.60 | 42.05 |
| ResNet-101 + UPerNet | 7.40 0.49 | 8.64 0.49 1.24 | 10.99 0.32 | 13.14 0.46 2.15 | 15.89 0.63 | 16.82 0.28 0.93 | 43.82 |
| (b) Impact of Backbone Architecture | |||||||
| Method | Occupancy C/P | Pair R/P | Geometry C/P |
| Uniform | .506/.929 | .592/.932 | .544/.916 |
| SD3.5 | .434/.888 | .539/.941 | .549/.880 |
| UniDD ∗ | .427/.901 | .500/.955 | .464/.862 |
| S 2 -RL | .538/.929 | .665/.898 | .693/.939 |
| Method Configuration | 0.25% | 0.5% | 1.0% |
| SD3.5 Baseline | 13.54 0.56 | 18.09 0.90 | 21.66 0.40 |
| + Alignment ( ) | 14.82 0.43 | 19.15 0.57 | 22.91 0.23 |
| + Diversity ( ) | 15.17 0.61 | 20.83 0.42 | 24.73 0.28 |
| + Set ( ) | 15.60 0.59 | 21.52 0.56 | 25.47 0.24 |
Appendix figures & tables22 assets
Supplementary material from the paper’s appendix.
Appendix
| Hyperparameter | ADE20K | COCO |
| Optimization (GRPO) | ||
| Optimizer | AdamW [ 24 ] | |
| Learning Rate | ||
| Batch Size (Prompts) | 16 | |
| Group Size ( ) | 8 | 16 |
| Ratio Clipping Width ( ) | ||
| Method | Construction Cost | Peak GPU Memory |
| Random | selection 1 min | N/A |
| Uniform | selection 1 min | N/A |
| K-Center | selection 10 min | 4 GB |
| Herding | selection 10 min | 4 GB |
| SD3.5 Baseline | inference 1 GPU-hr | 40 GB |
| UniDD ∗ | FreestyleNet training 48 GPU-hrs inference 10 GPU-hrs | 30 GB |
| Method | GPU-hours | mIoU (%) | mAcc (%) |
| (a) Shared generator backbone | |||
| UniDD ∗ (SD3.5 control) | – | 23.95 0.41 | 33.80 0.92 |
| S 2 -RL | – | 25.47 0.24 | 35.58 0.57 |
| (b) Approximately matched generative computation | |||
| SD3.5 best-of- | 50 | 23.29 0.36 | 33.04 0.83 |
| S 2 -RL | 49 | 25.47 0.24 | 35.58 0.57 |
| Method | Selection Feature | mAcc (%) | mIoU (%) |
| K-Center | ResNet-50 | 19.73 0.51 | 14.35 0.46 |
| K-Center | DINOv3 | 21.53 0.63 | 14.98 0.37 |
| Herding | ResNet-50 | 16.21 0.23 | 11.88 0.14 |
| Herding | DINOv3 | 17.55 0.26 | 12.66 0.15 |
| Random | N/A | 16.61 0.67 | 12.00 0.46 |
| Uniform | N/A | 21.03 0.92 | 14.83 0.57 |
| mAcc | mIoU | |||
| 1.0 | 3.0 | 3.0 | 21.65 0.29 | 14.66 0.16 |
| 1.0 | 3.0 | 5.0 | 22.40 0.36 | 15.13 0.16 |
| 1.0 | 5.0 | 3.0 | 23.07 0.34 | 15.43 0.20 |
| 1.0 | 5.0 | 5.0 | 22.47 0.35 | 15.29 0.19 |
| 1.0 | 10.0 | 3.0 | 22.18 0.37 | 15.04 0.25 |
| 1.0 | 10.0 | 5.0 | 22.91 0.27 | 15.31 0.23 |
| 4 | 8 | 16 | |
| mIoU (%) | 14.96 0.25 | 15.43 0.20 | 15.52 0.19 |
| mAcc (%) | mIoU (%) | |
| 25 | 22.30 0.38 | 15.00 0.16 |
| 50 | 23.07 0.34 | 15.43 0.20 |
| 100 | 22.58 0.87 | 15.24 0.55 |
| Prototype bank | mIoU (%) | mAcc (%) |
| Random same-class regions | 14.82 0.37 | 22.21 0.51 |
| K-means nearest real regions | 15.43 0.20 | 23.07 0.34 |
| Reward Evaluator(s) | mIoU (%) |
| UPerNet only | 14.52 0.22 |
| SegFormer-B5 only | 15.01 0.36 |
| DPT-Large only | 14.85 0.25 |
| Ensemble | 15.43 0.20 |
| Dataset/method | Feature C | Feature P | Occupancy C | Occupancy P | Vendi |
| ADE Uniform | .525 .015 | .878 .010 | .506 .016 | .929 .010 | 2.664 .023 |
| ADE SD3.5 | .370 .015 | .777 .018 | .434 .016 | .888 .019 | 2.397 .005 |
| ADE UniDD ∗ | .443 .018 | .834 .019 | .427 .011 | .901 .024 | 2.587 .025 |
| ADE S 2 -RL | .447 .010 | .796 .011 | .538 .016 | .929 .005 | 2.389 .025 |
| COCO Uniform | .314 .011 | .875 .008 | .304 .014 | .967 .005 | 2.659 .005 |
| COCO SD3.5 | .240 .005 | .884 .015 | .262 .003 | .967 .008 | 2.183 .029 |
| Dataset/method | Pair R | Pair P | Geometry C | Geometry P |
| ADE Uniform | .592 .008 | .932 .011 | .544 .012 | .916 .005 |
| ADE SD3.5 | .539 .019 | .941 .008 | .549 .020 | .880 .012 |
| ADE UniDD ∗ | .500 .017 | .955 .002 | .464 .015 | .862 .016 |
| ADE S 2 -RL | .665 .015 | .898 .010 | .693 .006 | .939 .004 |
| COCO Uniform | .534 .024 | .975 .005 | .334 .022 | .836 .023 |
| COCO SD3.5 | .475 .007 | .981 .002 | .288 .013 | .778 .010 |
| Dataset/method | Pair R/P | Geometry C/P | Feature micro P | Geometry micro P | Real-weighted geometry P |
| ADE Uniform | .571/.926 | .537/.907 | .851/.852 | .822/.804 | .959/.955 |
| ADE SD3.5 | .529/.935 | .545/.879 | .792/.793 | .856/.833 | .948/.948 |
| ADE UniDD ∗ | .492/.951 | .475/.867 | .839/.837 | .857/.838 | .940/.941 |
| ADE S 2 -RL | .645/.896 | .680/.937 | .795/.800 | .822/.799 | .971/.971 |
| COCO Uniform | .517/.972 | .336/.836 | .868/.866 | .918/.909 | .946/.944 |
| COCO SD3.5 | .464/.977 | .296/.778 | .865/.864 | .927/.914 | .924/.919 |
| Dataset/method | FID | KID | PRDC recall | PRDC coverage | Matched LPIPS |
| ADE Uniform | 101.31 0.74 | 1.83 0.66 | 0.791 0.014 | 0.381 0.019 | 0.6996 0.0043 |
| ADE SD3.5 | 118.08 1.72 | 10.10 1.24 | 0.592 0.012 | 0.265 0.009 | 0.6566 0.0048 |
| ADE UniDD ∗ | 106.75 1.68 | 4.49 1.25 | 0.749 0.016 | 0.352 0.014 | 0.6993 0.0052 |
| ADE S 2 -RL | 122.21 1.05 | 15.71 1.27 | 0.550 0.042 | 0.264 0.027 | 0.5754 0.0023 |
| COCO Uniform | 108.20 1.58 | 0.99 0.36 | 0.791 0.012 | 0.253 0.020 | 0.7043 0.0113 |
| COCO SD3.5 | 120.98 1.33 | 9.00 0.89 | 0.587 0.023 | 0.218 0.012 | 0.6867 0.0087 |
Dataset Distillation
Dataset distillation compresses training data into a small synthetic or selected set that retains the information needed to train useful models.
See all 67 papers