S$^2$-RL: Sample-Set Dual Reinforcement Learning for Generative Semantic Segmentation Dataset Distillation

Haoyu Wang; Fei Zhou; Qingqing Qiu; Lei Zhang; Wei Wei; Chen Ding

cs.CVOct 26, 2026

S $^2$ -RL: Sample-Set Dual Reinforcement Learning for Generative Semantic Segmentation Dataset Distillation

Authors: Haoyu Wang, Fei Zhou, Qingqing Qiu, Lei Zhang, Wei Wei, Chen Ding

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 $^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.

Figures & tables

Panel of Figure 1 — Figure 1: Comparison of paradigms. (Left) Optimization-based methods incur prohibitive memory costs from unrolling pixel-wise gradients. (Middle) Proxy-based generative methods require masks from real data, restricting synthesized datasets to existing layouts. (Right) S 2 -RL treats distillation as active scene construction, using RL to generate diverse spatial topologies and valid scene layouts from text alone.

Method	0.625‰		1.25‰		2.5‰
Method	mAcc	mIoU	mAcc	mIoU	mAcc	mIoU
Random	28.82 $\pm$ 1.32	19.71 $\pm$ 0.67	38.00 $\pm$ 0.40	27.75 $\pm$ 0.46	44.89 $\pm$ 0.40	34.31 $\pm$ 0.37
Uniform	33.45 $\pm$ 0.83	23.37 $\pm$ 0.97	43.37 $\pm$ 0.97	32.18 $\pm$ 0.75	50.07 $\pm$ 0.43	37.39 $\pm$ 0.46
K-Center	34.77 $\pm$ 0.67	24.99 $\pm$ 0.57	43.35 $\pm$ 0.72	32.30 $\pm$ 0.46	48.92 $\pm$ 0.47	37.19 $\pm$ 0.38
Herding	32.07 $\pm$ 0.35	23.04 $\pm$ 0.24	40.20 $\pm$ 0.25	30.49 $\pm$ 0.22	47.56 $\pm$ 0.23	37.10 $\pm$ 0.20
SD3.5 Baseline	38.18 $\pm$ 0.73	25.98 $\pm$ 0.67	46.65 $\pm$ 0.48	32.88 $\pm$ 0.87	51.81 $\pm$ 0.35	37.74 $\pm$ 0.34

Table 1: Quantitative comparison on COCO [ 20 ] using SegFormer-B5 [ 43 ] as the downstream student. We report mAcc and mIoU (%) at various distillation ratios. ∗ denotes our reproduced implementation. The highlighted row denotes our method, and green indicates the improvement over the strongest baseline.

Model Configuration	Ratio = 0.25%		Ratio = 0.5%		Ratio = 1.0%		Full
Model Configuration	Uniform	S 2 -RL (Ours)	Uniform	S 2 -RL (Ours)	Uniform	S 2 -RL (Ours)	Full
(a) Impact of Model Capacity
ResNet-18 + UPerNet	5.32 $\pm$ 0.20	5.90 $\pm$ 0.29 $\Delta$ 0.58	8.33 $\pm$ 0.42	9.12 $\pm$ 0.36 $\Delta$ 0.79	12.10 $\pm$ 0.61	12.45 $\pm$ 0.25 $\Delta$ 0.35	35.58
ResNet-50 + UPerNet	6.72 $\pm$ 0.35	7.69 $\pm$ 0.36 $\Delta$ 0.97	10.25 $\pm$ 0.39	11.84 $\pm$ 0.51 $\Delta$ 1.59	14.83 $\pm$ 0.57	15.43 $\pm$ 0.20 $\Delta$ 0.60	42.05
ResNet-101 + UPerNet	7.40 $\pm$ 0.49	8.64 $\pm$ 0.49 $\Delta$ 1.24	10.99 $\pm$ 0.32	13.14 $\pm$ 0.46 $\Delta$ 2.15	15.89 $\pm$ 0.63	16.82 $\pm$ 0.28 $\Delta$ 0.93	43.82
(b) Impact of Backbone Architecture

Table 2: Comprehensive robustness analysis on ADE20K [ 51 ] . We report mIoU (%) across varying model capacities, backbone architectures, and semantic segmentation mechanisms at three distillation ratios.

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

Table 3: Semantic/layout diagnostics on ADE20K 1.0%. C/P: coverage/precision; R/P: recall/precision. Values are means over five stored sets; metric definitions and SDs are in Appendix E .

Method Configuration	0.25%	0.5%	1.0%
SD3.5 Baseline	13.54 $\pm$ 0.56	18.09 $\pm$ 0.90	21.66 $\pm$ 0.40
+ Alignment ( $R_{align}$ )	14.82 $\pm$ 0.43	19.15 $\pm$ 0.57	22.91 $\pm$ 0.23
+ Diversity ( $R_{div}$ )	15.17 $\pm$ 0.61	20.83 $\pm$ 0.42	24.73 $\pm$ 0.28
+ Set ( $R_{set}$ )	15.60 $\pm$ 0.59	21.52 $\pm$ 0.56	25.47 $\pm$ 0.24

Table 4: Component effectiveness. ADE20K performance across distillation ratios.

Figure 3: Qualitative comparison of complex scene generation. We visualize images and pseudo-labels generated from multi-category prompts. In these examples, the SD3.5 baseline produces simplified scenes that omit some requested object categories, whereas S 2 -RL synthesizes semantically richer images with more complex scene structure.

Appendix figures & tables22 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	ADE20K	COCO
Optimization (GRPO)
Optimizer	AdamW [ 24 ]
Learning Rate	$1\times 10^{-4}$
Batch Size (Prompts)	16
Group Size ( $G$ )	8	16
Ratio Clipping Width ( $\epsilon$ )	$2\times 10^{-6}$

Appendix

Table 5: Detailed Hyperparameters for S 2 -RL Training.

Method	Construction Cost	Peak GPU Memory
Random	selection $<$ 1 min	N/A
Uniform	selection $<$ 1 min	N/A
K-Center	selection $\sim$ 10 min	$\sim$ 4 GB
Herding	selection $\sim$ 10 min	$\sim$ 4 GB
SD3.5 Baseline	inference $\sim$ 1 GPU-hr	$\sim$ 40 GB
UniDD ∗	FreestyleNet training $\sim$ 48 GPU-hrs inference $\sim$ 10 GPU-hrs	$\sim$ 30 GB

Appendix

Table 6: Construction cost and peak GPU memory on ADE20K at the 1.0% budget.

Method	GPU-hours	mIoU (%)	mAcc (%)
(a) Shared generator backbone
UniDD ∗ (SD3.5 control)	–	23.95 $\pm$ 0.41	33.80 $\pm$ 0.92
S 2 -RL	–	25.47 $\pm$ 0.24	35.58 $\pm$ 0.57
(b) Approximately matched generative computation
SD3.5 best-of- $N$	$\sim$ 50	23.29 $\pm$ 0.36	33.04 $\pm$ 0.83
S 2 -RL	$\sim$ 49	25.47 $\pm$ 0.24	35.58 $\pm$ 0.57

Appendix

Table 7: Additional controls on ADE20K 1.0% with SegFormer-B5. Both blocks use SD3.5; costs in (b) are approximate generative GPU-hours. Scores are mean

\pm

SD over 5 sets

\times

5 student seeds.

Method	Selection Feature	mAcc (%)	mIoU (%)
K-Center	ResNet-50	19.73 $\pm$ 0.51	14.35 $\pm$ 0.46
K-Center	DINOv3	21.53 $\pm$ 0.63	14.98 $\pm$ 0.37
Herding	ResNet-50	16.21 $\pm$ 0.23	11.88 $\pm$ 0.14
Herding	DINOv3	17.55 $\pm$ 0.26	12.66 $\pm$ 0.15
Random	N/A	16.61 $\pm$ 0.67	12.00 $\pm$ 0.46
Uniform	N/A	21.03 $\pm$ 0.92	14.83 $\pm$ 0.57

Appendix

Table 8: DINOv3-controlled coreset reruns on ADE20K at the 1.0% budget. This controlled analysis uses the canonical UPerNet + ResNet-50 downstream student.

$\gamma_{a}$	$\gamma_{d}$	$\gamma_{s}$	mAcc	mIoU
1.0	3.0	3.0	21.65 $\pm$ 0.29	14.66 $\pm$ 0.16
1.0	3.0	5.0	22.40 $\pm$ 0.36	15.13 $\pm$ 0.16
1.0	5.0	3.0	23.07 $\pm$ 0.34	15.43 $\pm$ 0.20
1.0	5.0	5.0	22.47 $\pm$ 0.35	15.29 $\pm$ 0.19
1.0	10.0	3.0	22.18 $\pm$ 0.37	15.04 $\pm$ 0.25
1.0	10.0	5.0	22.91 $\pm$ 0.27	15.31 $\pm$ 0.23

Appendix

Table 9: Impact of varying reward weights on downstream metrics (ADE20K 1%).

$G$	4	8	16
mIoU (%)	14.96 $\pm$ 0.25	15.43 $\pm$ 0.20	15.52 $\pm$ 0.19

Appendix

Table 10: Sensitivity to GRPO group size

G

on downstream mIoU (ADE20K 1%).

$K$	mAcc (%)	mIoU (%)
25	22.30 $\pm$ 0.38	15.00 $\pm$ 0.16
50	23.07 $\pm$ 0.34	15.43 $\pm$ 0.20
100	22.58 $\pm$ 0.87	15.24 $\pm$ 0.55

Appendix

Table 11: Effect of the number of semantic prototypes

K

on ADE20K at the 1.0% budget.

Prototype bank	mIoU (%)	mAcc (%)
Random same-class regions	14.82 $\pm$ 0.37	22.21 $\pm$ 0.51
K-means nearest real regions	15.43 $\pm$ 0.20	23.07 $\pm$ 0.34

Appendix

Table 12: Prototype construction at

K=50

on ADE20K 1.0%, evaluated with UPerNet–ResNet-50 over five student seeds.

Reward Evaluator(s)	mIoU (%)
UPerNet only	14.52 $\pm$ 0.22
SegFormer-B5 only	15.01 $\pm$ 0.36
DPT-Large only	14.85 $\pm$ 0.25
Ensemble	15.43 $\pm$ 0.20

Appendix

Table 13: Reward-evaluator robustness under a fixed SD3.5 generator prior on ADE20K at the 1.0% budget.

Dataset/method	Feature C	Feature P	Occupancy C	Occupancy P	Vendi
ADE Uniform	.525 $\pm$ .015	.878 $\pm$ .010	.506 $\pm$ .016	.929 $\pm$ .010	2.664 $\pm$ .023
ADE SD3.5	.370 $\pm$ .015	.777 $\pm$ .018	.434 $\pm$ .016	.888 $\pm$ .019	2.397 $\pm$ .005
ADE UniDD ∗	.443 $\pm$ .018	.834 $\pm$ .019	.427 $\pm$ .011	.901 $\pm$ .024	2.587 $\pm$ .025
ADE S 2 -RL	.447 $\pm$ .010	.796 $\pm$ .011	.538 $\pm$ .016	.929 $\pm$ .005	2.389 $\pm$ .025
COCO Uniform	.314 $\pm$ .011	.875 $\pm$ .008	.304 $\pm$ .014	.967 $\pm$ .005	2.659 $\pm$ .005
COCO SD3.5	.240 $\pm$ .005	.884 $\pm$ .015	.262 $\pm$ .003	.967 $\pm$ .008	2.183 $\pm$ .029

Appendix

Table 14: Reward-independent feature and occupancy diagnostics at the 0.1% area cutoff. Mean

\pm

sample SD across five stored sets. C: coverage; P: precision.

Dataset/method	Pair R	Pair P	Geometry C	Geometry P
ADE Uniform	.592 $\pm$ .008	.932 $\pm$ .011	.544 $\pm$ .012	.916 $\pm$ .005
ADE SD3.5	.539 $\pm$ .019	.941 $\pm$ .008	.549 $\pm$ .020	.880 $\pm$ .012
ADE UniDD ∗	.500 $\pm$ .017	.955 $\pm$ .002	.464 $\pm$ .015	.862 $\pm$ .016
ADE S 2 -RL	.665 $\pm$ .015	.898 $\pm$ .010	.693 $\pm$ .006	.939 $\pm$ .004
COCO Uniform	.534 $\pm$ .024	.975 $\pm$ .005	.334 $\pm$ .022	.836 $\pm$ .023
COCO SD3.5	.475 $\pm$ .007	.981 $\pm$ .002	.288 $\pm$ .013	.778 $\pm$ .010

Appendix

Table 15: Reward-independent category-pair and relative-geometry diagnostics at the 0.1% area cutoff. Mean

\pm

sample SD across five stored sets.

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

Appendix

Table 16: Threshold and aggregation sensitivity. Pair R/P and geometry C/P use the 0.2% area cutoff; the three precision columns list 0.1% / 0.2%. Values are descriptive five-set means.

Dataset/method	FID	KID $\times 10^{3}$	PRDC recall	PRDC coverage	Matched LPIPS
ADE Uniform	101.31 $\pm$ 0.74	1.83 $\pm$ 0.66	0.791 $\pm$ 0.014	0.381 $\pm$ 0.019	0.6996 $\pm$ 0.0043
ADE SD3.5	118.08 $\pm$ 1.72	10.10 $\pm$ 1.24	0.592 $\pm$ 0.012	0.265 $\pm$ 0.009	0.6566 $\pm$ 0.0048
ADE UniDD ∗	106.75 $\pm$ 1.68	4.49 $\pm$ 1.25	0.749 $\pm$ 0.016	0.352 $\pm$ 0.014	0.6993 $\pm$ 0.0052
ADE S 2 -RL	122.21 $\pm$ 1.05	15.71 $\pm$ 1.27	0.550 $\pm$ 0.042	0.264 $\pm$ 0.027	0.5754 $\pm$ 0.0023
COCO Uniform	108.20 $\pm$ 1.58	0.99 $\pm$ 0.36	0.791 $\pm$ 0.012	0.253 $\pm$ 0.020	0.7043 $\pm$ 0.0113
COCO SD3.5	120.98 $\pm$ 1.33	9.00 $\pm$ 0.89	0.587 $\pm$ 0.023	0.218 $\pm$ 0.012	0.6867 $\pm$ 0.0087

Appendix

Table 17: Generic image-distribution and matched perceptual-diversity diagnostics. Lower FID/KID and higher recall, coverage, and LPIPS are preferred by these metrics.

Figure 4: Generated samples on ADE20K with 0.25% data ratio.

Figure 5: Generated samples on ADE20K with 0.5% data ratio.

Figure 6: Generated samples on ADE20K with 1.0% data ratio. (Continued on next page)

Figure 7: (Continued) Generated samples on ADE20K with 1.0% data ratio.

Figure 8: Generated samples on COCO with 0.625‰ data ratio.

Figure 9: Generated samples on COCO with 1.25‰ data ratio.

Figure 10: Generated samples on COCO with 2.5‰ data ratio. (Continued on next page)

Figure 11: (Continued) Generated samples on COCO with 2.5‰ data ratio. (Continued on next page)

Figure 12: (Continued) Generated samples on COCO with 2.5‰ data ratio.

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

D3S2: Diffusion-Guided Dataset Distillation for Semantic Segmentation
May 24, 2026Wenjie Zheng, Haoji Hu, Jiali Lu +2Diffusion Model Guidance Dataset Distillation
Zhejiang University, China
SAS: Semantic-aware Sampling for Generative Dataset Distillation
May 18, 2026Mingzhuo Li, Guang Li, Linfeng Ye +4Diversity-Aware Sampling Training Data Selection
Hokkaido University, N-14, W-9, Kita-Ku, Sapporo, 060-0814, Japan · University of Toronto, 27 King’s College Circle, Toronto, Ontario M5S 1A1, Canada · The University of Tokyo, 7-3-1 Hongo, Bunkyo-ku, Tokyo, 113-8654, Japan
Self-Supervised Representation-Guided Generative Dataset Distillation
Aug 4, 2026Mingzhuo Li, Guang Li, Linfeng Ye +4Diffusion Model Guidance Dataset Distillation
Hokkaido University, Japan · University of Toronto, Canada · University of Tokyo, Japan

May 24, 2026J/K move · Enter open · S save

D3S2: Diffusion-Guided Dataset Distillation for Semantic Segmentation

Wenjie Zheng, Haoji Hu, Jiali Lu, Xingze Zou+1

Zhejiang University, China

Diffusion Model Guidance Dataset Distillation Semantic Segmentation

Open paper AskPDF arXiv

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. Our code is available at https://github.com/zwj084/D3S2.

Abstract

Dataset Distillation

Explore similar work

D3S2: Diffusion-Guided Dataset Distillation for Semantic Segmentation

SAS: Semantic-aware Sampling for Generative Dataset Distillation

Self-Supervised Representation-Guided Generative Dataset Distillation