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

Mar 17, 2026

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

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

Mahmoud K. Ibrahim, Bart Elen, Chang Sun +5

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

Ask PDF

Apr 21, 2026

Generative Drifting for Conditional Medical Image Generation

Conditional medical image generation plays an important role in many clinically relevant imaging tasks. However, existing methods still face a fundamental challenge in balancing inference efficiency, patient-specific fidelity, and distribution-level plausibility, particularly in high-dimensional 3D medical imaging. In this work, we propose GDM, a generative drifting framework that reformulates deterministic medical image prediction as a multi-objective learning problem to jointly promote distribution-level plausibility and patient-specific fidelity while retaining one-step inference. GDM extends drifting to 3D medical imaging through an attractive-repulsive drift that minimizes the discrepancy between the generator pushforward and the target distribution. To enable stable drifting-based learning in 3D volumetric data, GDM constructs a multi-level feature bank from a medical foundation encoder to support reliable affinity estimation and drifting field computation across complementary global, local, and spatial representations. In addition, a gradient coordination strategy in the shared output space improves optimization balance under competing distribution-level and fidelity-oriented objectives. We evaluate the proposed framework on two representative tasks, MRI-to-CT synthesis and sparse-view CT reconstruction. Experimental results show that GDM consistently outperforms a wide range of baselines, including GAN-based, flow-matching-based, and SDE-based generative models, as well as supervised regression methods, while improving the balance among anatomical fidelity, quantitative reliability, perceptual realism, and inference efficiency. These findings suggest that GDM provides a practical and effective framework for conditional 3D medical image generation.

Zirong Li, Siyuan Mei, Weiwen Wu +3

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

3D Medical Imaging One-Step Generative Modeling Sparse-View CT Reconstruction Medical Image Synthesis Conditional Generation

Ask PDF

Jul 27, 2026

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

Text-to-image generation models can synthesize high-quality images from natural language descriptions, but their performance remains highly sensitive to prompt formulation. Existing prompt optimization methods mainly rely on text-side rewriting, prompt expansion, or external reward signals, offering limited image-grounded diagnosis and weak support for learning reusable optimisation policies. In this paper, we propose PRISM, a Prompt Refinement framework via Image-grounded Self-rewarding Mechanism. PRISM closes the prompt-image-feedback loop by interpreting generated images with structured visual diagnosis and scoring them along semantic consistency, aesthetic quality, and human preference alignment. It first initializes a unified VLM through multi-task supervised fine-tuning, and then improves the prompt policy via self-rewarding optimization with a hybrid ideal-point and Chebyshev reward. Extensive experiments show that PRISM improves holistic image quality and fine-grained semantic alignment, while providing interpretable feedback for targeted prompt refinement. The code is available at https://anonymous.4open.science/r/PRISM-FF81.

Guo Tang, HongJie Luo, Tianxu Wang +2

Harbin Institute of Technology, Shenzhen · Sun Yat-sen University · South China University of Technology +1

T2I Generation RL for Image Generation Prompt Optimization for T2I Generation

Ask PDF

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