S²-RL: Sample-Set Dual Reinforcement Learning for Generative Semantic Segmentation Dataset Distillation

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

Generative dataset distillation

S²-RL: Sample-Set DualReinforcement Learningfor Generative Semantic Segmentation
Dataset Distillation

Haoyu Wang¹ Fei Zhou²Qingqing Qiu³Lei Zhang^1,†Wei Wei¹Chen Ding⁴

¹Northwestern Polytechnical University²The Chinese University of Hong Kong³Zhejiang University⁴Xi'an University of Posts & Telecommunications

† Corresponding author

arXiv Code Hugging Face Explore the data

Learning to generate compact, diverse datasets
for semantic segmentation.

Lakeside

Street

Bathroom

Living room

Useful individual samples. Complementary sample sets. One dual-reward objective.

+2.06mIoU points on ADE20K0.25% data budget · vs. SD3.5
+2.48mIoU points on COCO0.0625% data budget · vs. SD3.5

The idea

Learn to build a better training set.

A compact segmentation dataset needs both informative images and diverse semantic coverage. S²-RL optimizes these two goals together.

We fine-tune a text-to-image diffusion policy with group relative policy optimization (GRPO). A sample-level reward encourages semantic alignment and intra-class feature diversity. A set-level reward encourages exploration of underrepresented semantic prototypes. The resulting images receive teacher-generated pseudo-labels and form a small training set for downstream segmentation models.

Read the research summary

While dataset distillation has advanced substantially in image classification, extending it to semantic segmentation remains challenging. Pixel-wise gradient unrolling is computationally expensive, and mask-conditioned generative approaches depend on existing spatial layouts. S²-RL formulates generative semantic segmentation dataset distillation as reward-guided diffusion policy optimization. At the sample level, it encourages prompt alignment and intra-class feature diversity. At the set level, a contextual bandit formulation rewards coverage of underrepresented semantic prototypes. Experiments on ADE20K and COCO show consistent gains over the evaluated selection and generative baselines across the reported data budgets.

01 / The samples

The same prompt. Richer semantic content.

Explore the paper's paired examples and their predicted segmentation maps.

PROMPT

sky, tree, stone, bridge, lake

Same text conditioning, different generated scenes. Examples from Figure 4.

02 / The method

Two levels of feedback. One generation policy.

Reward each sample for its content, and each set for what it covers.

S squared RL framework: a diffusion policy generates an image batch; sample and set rewards drive a GRPO update. Sample reward checks semantic alignment and intra-class diversity. Set reward uses class-wise prototype memory, scarcity and exploration.

SAMPLE-LEVEL REWARD

Make each sample informative.

Align the image with the requested categories and encourage varied local features within each class.

Semantic alignmentIntra-class diversity

SET-LEVEL REWARD

Give each new sample a purpose.

Track class-specific semantic prototypes in a global memory and reward exploration of underrepresented modes.

Prototype memoryScarcity + exploration

Generated imagesTeacher pseudo-labelsTrain a segmentation studentEvaluate on real data

All synthetic baselines use the same dataset-specific pseudo-label teacher and downstream evaluation protocol.

03 / The evidence

Better segmentation at the same data budget.

Consistent gains across both datasets and all six evaluated budgets.

ADE20K mIoU versus data budget. S squared RL reaches 15.60, 21.52 and 25.47 at 0.25%, 0.5% and 1% budgets, above all evaluated baselines.

COCO mIoU versus data budget. S squared RL reaches 28.46, 35.03 and 38.83 at 0.0625%, 0.125% and 0.25% budgets, above all evaluated baselines.

Downstream model: SegFormer-B5 · mIoU (%) ↑ · 5 dataset groups × 5 training runs. Error bars reproduce the manuscript's reported ± values. Data budgets are displayed as percentages.

View all reported results

UniDD* denotes the reproduced UniDD-inspired baseline. All values are from the paper's main tables.

Download table data

Ablation study

Each reward contributes.

Rewards are added cumulatively to the SD3.5 baseline.

ADE20K budget

mIoU (%) ↑ · Same downstream model and budget for all four configurations.

Beyond a single sample

Explore a broader semantic space.

Set-level feedback encourages the generator to cover diverse semantic modes. The paper's class-wise t-SNE visualizations compare real data, uniform selection and S²-RL at the 1% ADE20K budget.

Feature-space views from Figure 3; quantitative performance is reported above.

Seat class t-SNE: real data as gray dots, uniform selection as blue triangles, S squared RL as red stars — Real dataUniform★S²-RL

04 / Independent evaluation

More ways for a scene to come together.

Measure where classes appear, which classes coexist, and how they relate in space.

Spatial occupancy

Scroll the chart to compare coverage/recall and precision →

Mean ± sample SD over five stored sets · Scores range from 0 to 1 · Higher is better

Numbers and evaluation protocol

The same frozen OneFormer annotates constructed sets and held-out real images: 2,000 ADE20K and 5,000 COCO references. Regions smaller than 0.1% of an image are excluded. Occupancy and geometry use fixed real-reference local support, macro-averaged over eligible classes or pairs. These diagnostics measure set coverage; they are separate from downstream mIoU.

The results do not imply improvement in every form of diversity. Uniform has stronger generic feature coverage; Vendi and micro-averaged precision show mixed results. Variation is across stored sets, not independent policy runs.

Download diagnostic data ↓

See what the descriptors measure

Read the layout of an actual sample.

Stored example

The first three S²-RL images in the evaluated set, in stored order. Choose two classes to inspect their occupied area and relative position.

Independent OneFormer annotations · Dots mark class-union centroids

Class AClass B

Each cell is the class area in that cell, as a percentage of the whole image.

Mask overlay35%

05 / Explore the training data

A compact set. Every image and label.

Explore 868 image–label pairs across both datasets and all three budgets.

First stored set · Original 1024 × 1024 images · Click to inspect

Build on this work

Citation

@misc{wang2026s2rl,
  title={S$^2$-RL: Sample-Set Dual Reinforcement Learning for Generative Semantic Segmentation Dataset Distillation},
  author={Haoyu Wang and Fei Zhou and Qingqing Qiu and Lei Zhang and Wei Wei and Chen Ding},
  year={2026},
  eprint={2610.35491},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2610.35491}
}

⚗️S2-RL: Sample-Set DualReinforcement Learningfor Generative Semantic Segmentation Dataset Distillation