Compositional Reward Models for Conditional Medical Image Generation

Aayush Kumar Tyagi; Prathosh A. P.; Mausam

cs.CVSep 4, 2026

Compositional Reward Models for Conditional Medical Image Generation

Authors: Aayush Kumar Tyagi, Prathosh A. P., Mausam

Organizations: Yardi School of Artificial Intelligence India Institute Of Technology, Delhi · Department of Electrical Communication Engineering Indian Institute of Science, Bengaluru · Department of Computer Science and Engineering Indian Institute of Technology, Delhi

Abstract

Acquiring high quality annotated medical image data is critical for training deep learning models; however, annotation is expensive, time consuming, and requires domain expertise. Conditional diffusion models, such as ControlNet, offer an alternative by generating images conditioned on semantic masks and text. However, existing approaches fail to capture fine grained properties (e.g., intensity and texture), as well as semantic consistency expected by domain experts, limiting their effectiveness for downstream tasks. Recent attempts to address these issues using reinforcement learning fine-tuning remain limited due to the reliance on a single scalar reward, which conflates diverse failure modes and provides weak corrective signals. We propose PRISM, a Compositional Reward Model (CRM) framework for conditional medical image generation. Instead of assigning a single reward, we decompose image quality into verifier grounded stages, each evaluating a distinct aspect of correctness from fine to coarse properties, including low level attributes (intensity and texture), structural alignment with conditioning inputs, and high level semantic fidelity. These stage wise rewards are composed through a Hierarchical Constrained Propagation (HCP) mechanism that enforces a fine to coarse notion of correctness, ensuring that lower level deficiencies are resolved before higher level rewards are accrued, preventing easier objectives from masking critical failures. We evaluate PRISM across three datasets spanning diverse medical imaging tasks: PanNuke (multi-class cell segmentation), CeDeM (villi/crypt detection and measurement), and ISIC (skin lesion classification). Training downstream models with data generated by PRISM yields improvements over closest baselines, including a 2.3% increase in mDice on PanNuke, a 8.5% reduction in Mean Relative Error (MRE) on CeDeM, and increases ISIC F1 by 5.9%.

Figures & tables

Figure 1: A: Conditional generator makes fine grain errors and a single score is not descriptive to provide good feedback. B: We decompose the reward into a composition of rewards (CRM) where, each reward constraint is checked before moving to next reward which we call, hierarchical constraint propagation (HCP). C: Using HCP to train generator with RL improves the image quality and D: eventually improves the performance of downstream task.

Figure 2: Once the generator $\pi_{\theta}$ is trained in a supervised manner, it is used to generate images conditioned on a segmentation mask $M$ and text prompt $T$ . Each generated image ( $\hat{I}$ ) is evaluated through set of verifiers. For the PanNuke dataset, these verifiers enforce a fine to coarse hierarchy: moving from low level attributes (intensity and texture matching by cell type) upward to structural and semantic agreement via pre-trained segmentation models. These stage wise sub rewards are composed using Hierarchical Constrained Propagation (HCP), ensuring that lower level constraints must be satisfied before higher level credit can accrue to yield the final gated reward $r_{CRM}$ .

Panel of Figure 3 — Figure 3: Qualitative comparisons for CeDeM, PanNuke, and ISIC. Each row uses the same conditioning mask across the supervised ControlNet (CN), ORM, and CRM generations. The original image is shown only as the paired real target for reference.

		CeDeM (MeasureNet)		PanNuke (SegFormer-B3)		ISIC (ResNet-50)
Method	Reward	MAE ${}_{\text{R}}\downarrow$	MRE ${}_{\text{R}}\downarrow$	mIoU $\uparrow$	mDice $\uparrow$	Acc. $\uparrow$	F1 $\uparrow$
SFT (Original only)	–	52.06	24.77	56.36	40.94	53.43	44.84
Generated (SFT)	–	47.69	22.15	58.26	42.24	55.82	47.30
Best-of-N	Inference	45.71	21.83	58.61	43.09	56.00	46.85
GDPO Liu et al. (2026)	ORM	41.95	18.85	58.88	43.27	63.00	45.50
GDPO Liu et al. (2026)	SUM	45.33	21.41	60.50	44.75	59.80	43.80

Table 1: Main results across all three datasets and downstream tasks. CeDeM uses MeasureNet ( Tyagi et al., 2024 ) (lower is better); PanNuke uses SegFormer-B3 ( Xie et al., 2021 ) (higher is better); ISIC uses ResNet-50 ( He et al., 2016 ) (higher is better). Best in bold , second best underlined . Full per dataset metrics are in Appendix B.1 .

$R_{1}$	$R_{2}$	$R_{3}$	Villi		Crypt		Ratio
			MAE	MRE	MAE	MRE	MAE	MRE
		✓	11.39	10.98	7.02	15.67	41.95	18.85
	✓		11.51	10.70	7.64	18.03	48.46	22.22
✓			10.02	10.54	7.76	19.84	46.00	20.72
	✓	✓	9.31	9.70	7.39	17.67	41.87	18.78
✓	✓	✓	10.01	9.68	6.64	14.43	39.58	17.24

Table 2: Subreward ablation on CeDeM (MeasureNet, GDPO).

R_{1}

: anchor,

R_{2}

: detection,

R_{3}

: ratio. Lower is better.

	CeDeM		PanNuke		ISIC
Order	MAE ${}_{\text{R}}\!\downarrow$	MRE ${}_{\text{R}}\!\downarrow$	mDice $\uparrow$	mIoU $\uparrow$	Acc. $\uparrow$	F1 $\uparrow$
$\boldsymbol{R_{1}R_{2}R_{3}}$	39.58	17.24	61.54	45.77	64.00	50.00
$R_{1}R_{3}R_{2}$	40.49	18.54	59.87	44.06	62.76	46.57
$R_{2}R_{1}R_{3}$	44.37	20.95	59.52	43.60	59.06	45.99
$R_{2}R_{3}R_{1}$	41.13	18.71	59.75	43.71	64.84	47.38
$R_{3}R_{1}R_{2}$	45.83	21.04	57.97	42.73	65.07	47.67

Table 3: CRM hierarchical order vs. primary downstream metrics. Best bold , second underlined .

Figure 4: CRM order ablation on CeDeM, PanNuke, and ISIC using the available permutation logs. The top row shows the mean CRM for last 15 epochs; the bottom row shows moving average training trajectories.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5: Dataset examples used for mask-conditional generation. Each row shows two representative conditioning masks and their paired real target images. CeDeM masks encode villus (blue), crypt (green), villi shoudler (teal) and crypt border (red) for measurement, PanNuke masks encode nuclei classes for segmentation, and ISIC masks encode lesion extent and dermoscopic attributes for diagnostic generation.

Figure 6: Supervised ControlNet training loss for the three datasets. Curves are EMA-smoothed over the raw per-step diffusion loss

	CeDeM	PanNuke	ISIC
Optimizer	AdamW	AdamW	AdamW
Learning rate	$5{\times}10^{-6}$	$5{\times}10^{-6}$	$5{\times}10^{-6}$
Batch size (per GPU)	4	4	4
Group size $N$	32	32	32
Clip ratio $\epsilon_{\mathrm{clip}}$	0.1	0.1	0.1
Inner iters	2	2	2

Appendix

Table 4: GDPO hyperparameters per dataset. “Group

N

” is the number of candidate samples per conditioning prompt; advantages are computed per component within each group and renormalised before clipping (Eq. 11 ). “Inner iters” is the number of policy-update passes per rollout batch.

Method	Reward	MAE ${}_{\text{V}}\downarrow$	MRE ${}_{\text{V}}\downarrow$	MAE ${}_{\text{C}}\downarrow$	MRE ${}_{\text{C}}\downarrow$	MAE ${}_{\text{R}}\downarrow$	MRE ${}_{\text{R}}\downarrow$
SFT (Original only)	–	13.88	14.56	8.41	21.52	52.06	24.77
Generated (SFT)	–	11.88	11.74	7.87	17.84	47.69	22.15
Best-of-N	Inference	10.26	10.49	7.55	17.84	45.71	21.83
GDPO Liu et al. [2026]	ORM (Ratio)	11.39	10.98	7.02	15.67	41.95	18.85
GDPO Liu et al. [2026]	SUM Liu et al. [2026]	9.89	10.55	7.18	16.58	45.33	21.41
GDPO Liu et al. [2026]	MaxMin	10.76	10.26	7.35	16.44	47.07	21.55

Appendix

Table 5: CeDeM full measurement metrics. Downstream model: MeasureNet Tyagi et al. [2024] . Subscripts denote villi (V), crypt (C), and villus-to-crypt ratio (R). Lower is better for all columns.

Method	Reward	Acc. $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$	F1 $\uparrow$
SFT (Original only)	–	53.43	49.91	50.55	44.84
Generated (SFT)	–	55.82	50.93	53.16	47.30
Best-of-N	Inference	56.00	49.11	52.74	46.85
GDPO Liu et al. [2026]	ORM	63.00	54.40	49.90	45.50
GDPO Liu et al. [2026]	SUM Liu et al. [2026]	59.80	47.20	46.60	43.80
GDPO Liu et al. [2026]	MaxMin	57.70	47.40	50.10	44.90

Appendix

Table 6: ISIC full skin-lesion classification metrics. Downstream model: ResNet-50 trained on real

+

synthetic dermoscopy images and evaluated on the held-out real ISIC-2019 test split (3 classes: nevus, melanoma, seborrheic keratosis). Precision, recall, and F1 are macro-averaged over the three classes. Higher is better.

Order	MAE ${}_{\text{V}}\downarrow$	MRE ${}_{\text{V}}\downarrow$	MAE ${}_{\text{C}}\downarrow$	MRE ${}_{\text{C}}\downarrow$	MAE ${}_{\text{R}}\downarrow$	MRE ${}_{\text{R}}\downarrow$
$\boldsymbol{R_{1}R_{2}R_{3}}$ (main)	10.01	9.68	6.64	14.32	39.58	17.24
$R_{1}R_{3}R_{2}$	10.22	9.83	6.50	15.59	40.49	18.54
$R_{2}R_{1}R_{3}$	11.08	10.68	6.99	15.41	44.37	20.95
$R_{2}R_{3}R_{1}$	10.27	10.40	6.88	15.12	41.13	18.71
$R_{3}R_{1}R_{2}$	11.75	11.16	6.99	14.99	45.83	21.04
$R_{3}R_{2}R_{1}$	10.09	10.13	6.96	15.93	44.83	20.14

Appendix

Table 7: CeDeM cascade-order ablation with fixed components and training budget.

R_{1}

: anchor,

R_{2}

: morphology,

R_{3}

: ratio; subscripts denote villus (V), crypt (C), and ratio (R). Lower is better. Best in bold , second underlined .

Order	mIoU $\uparrow$	mDice $\uparrow$
$\boldsymbol{R_{1}R_{2}R_{3}}$ (main)	45.77	61.54
$R_{1}R_{3}R_{2}$	44.06	59.87
$R_{2}R_{1}R_{3}$	43.60	59.52
$R_{2}R_{3}R_{1}$	43.71	59.75
$R_{3}R_{1}R_{2}$	42.73	57.97
$R_{3}R_{2}R_{1}$	43.38	58.91

Appendix

Table 8: PanNuke cascade-order ablation with fixed components and training budget.

R_{1}

: intensity,

R_{2}

: texture,

R_{3}

: SegFormer agreement. Higher is better. Best in bold , second underlined .

Order	Acc. $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$	F1 $\uparrow$
$\boldsymbol{R_{1}R_{2}R_{3}}$ (main)	64.00	51.60	53.60	50.00
$R_{1}R_{3}R_{2}$	62.76	46.86	49.40	46.57
$R_{2}R_{1}R_{3}$	59.06	45.72	47.07	45.99
$R_{2}R_{3}R_{1}$	64.84	49.27	48.76	47.38
$R_{3}R_{1}R_{2}$	65.07	47.50	49.75	47.67
$R_{3}R_{2}R_{1}$	64.07	46.42	45.54	45.16

Appendix

Table 9: ISIC cascade-order ablation with fixed components and training budget.

R_{1}

: intensity,

R_{2}

: attribute-mask agreement,

R_{3}

: diagnosis classifier. Precision, recall, and F1 are macro-averaged; higher is better. Best in bold , second underlined .

Panel of Figure 7 — Figure 7: Additional CeDeM qualitative comparisons. Mask overlays use blue (VilliM), red (Crypt Border), green (CryptM), and teal (Villi Shoulder). All generated columns share the conditioning mask of the corresponding original target image.

Panel of Figure 8 — Figure 8: Additional PanNuke qualitative comparisons. All generated columns use the same conditioning mask as the corresponding original target image.

Panel of Figure 9 — Figure 9: Additional qualitative ISIC comparisons. Each row shows the input mask, generations from the SFT ControlNet, ORM-aligned model, and CRM-aligned model, and the original image under the same conditioning.

Medical Image Synthesis

Medical image synthesis creates, translates, reconstructs, and edits clinical images to fill missing data and support analysis.

See all 312 papers

CompDiff enables fair and zero shot medical image generation across demographic intersections through compositional diffusion
Mar 17, 2026Mahmoud K. Ibrahim, Bart Elen, Chang Sun +5Healthcare Compositional Image Generation
Institute of Data Science, Faculty of Science and Engineering, Maastricht University, Maastricht, The Netherlands · Department of Advanced Computing Sciences, Faculty of Science and Engineering, Maastricht University, Maastricht, The Netherlands · VITO, Belgium
Generative Drifting for Conditional Medical Image Generation
Apr 21, 2026Zirong Li, Siyuan Mei, Weiwen Wu +33D Medical Imaging One-Step Generative Modeling
Department of Orthodontics and Orofacial Orthopedics, Friedrich-Alexander-University Erlangen-Nuremberg, Erlangen, Germany · Department Artificial Intelligence in Biomedical Engineering, Friedrich-Alexander-University Erlangen-Nuremberg, Erlangen, Germany · Pattern Recognition Lab, Friedrich-Alexander-University Erlangen-Nuremberg, Erlangen, Germany +1
PRISM: Prompt Refinement via Image-grounded Self-rewarding Mechanism for Text-to-Image Generation
Jul 27, 2026Guo Tang, HongJie Luo, Tianxu Wang +2T2I Generation RL for Image Generation
Harbin Institute of Technology, Shenzhen · Sun Yat-sen University · South China University of Technology +1

Mar 17, 2026J/K move · Enter open · S save

CompDiff enables fair and zero shot medical image generation across demographic intersections through compositional diffusion

Mahmoud K. Ibrahim, Bart Elen, Chang Sun, Ahmad Jiblawi+4

Institute of Data Science, Faculty of Science and Engineering, Maastricht University, Maastricht, The Netherlands · Department of Advanced Computing Sciences, Faculty of Science and Engineering, Maastricht University, Maastricht, The Netherlands · VITO, Belgium

Healthcare Compositional Image Generation Diffusion-Based Image Generation Medical Image Synthesis Intersectional Fairness

Open paper AskPDF arXiv

Medical image generators trained on imbalanced data can fail at demographic intersections absent from training. We introduce CompDiff, which encodes age, sex and race separately and composes supervised demographic tokens alongside clinical text. Across chest radiographs and fundus images, CompDiff improves overall and subgroup fidelity relative to prompt conditioning (RoentGen-v2) and loss reweighting (FairDiffusion). It generalises in zero-shot generation to 16 chest X-ray intersections excluded from training, achieving the lowest mean FID-RadImageNet in every intersection. In a blinded reader study of these unseen intersections, two radiologists gave CompDiff the highest mean scores among generators for anatomical realism and agreement with the clinical impression, and selected its images most often as the most realistic. Pretraining with CompDiff images improved downstream classification, while CompDiff audit cohorts reduced estimation error on rare intersections. These findings support compositional demographic conditioning for extending medical image synthesis to underserved populations. Code: https://github.com/mahmoudibrahim98/CompDiff

Abstract

Medical Image Synthesis

Explore similar work

CompDiff enables fair and zero shot medical image generation across demographic intersections through compositional diffusion

Generative Drifting for Conditional Medical Image Generation

PRISM: Prompt Refinement via Image-grounded Self-rewarding Mechanism for Text-to-Image Generation