PaperlayerView exact revision on arXiv

Source-linked AI summary

Compositional Reward Models for Conditional Medical Image Generation

Aayush Kumar Tyagi, Prathosh A. P., Mausam

arXiv:2609.05028v1cs.CV

TL;DR

Medical image generation seeks to reduce dependence on costly expert annotation, but scalar reward models provide weak guidance for diverse generation failures. PRISM uses verifier-grounded compositional rewards with hierarchical fine-to-coarse propagation, and its generated data improves downstream performance across PanNuke, CeDeM, and ISIC. The framework requires domain expertise to design task-specific subrewards and their hierarchy.

  • Problem

    Scalar reward models conflate distinct medical-image properties such as intensity, texture, and structural alignment, limiting guidance for correcting specific failure modes.

  • Method

    PRISM decomposes image quality into verifier-grounded reward components and uses hierarchical constrained propagation so lower-level attributes must be satisfied before higher-level rewards accumulate.

  • Results

    PRISM consistently improves downstream performance over the closest baseline across PanNuke, CeDeM, and ISIC, including 2.3% higher mDice, 8.5% lower ratio MRE, and 5.9% higher F1.

  • Takeaways & Limitations

    Compositional rewards with fine-to-coarse constraints improve conditional medical image generation and downstream segmentation, measurement, and classification performance.

  • Takeaways & Limitations

    PRISM requires domain expertise to choose task-specific subrewards, verifiers, and their hierarchical order.

Abstract

from arXiv · show

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%.

1 Introduction

Medical image annotation is costly, while conditional diffusion models can generate plausible mask-conditioned images but often miss fine-grained and structural properties. PRISM addresses this with staged compositional rewards and hierarchical constraints, improving downstream performance across three medical imaging tasks.

  • Motivation: Expert pixel-level annotation is expensive and time consuming, motivating synthetic medical image generation.Conditional diffusion models have shown promise for improving downstream performance.
  • Motivation: Conditional diffusion models can respect spatial layouts but frequently fail on intensity distributions, texture consistency, and structural correctness.These deficiencies can degrade models trained on synthetic data.
  • Motivation: Single scalar rewards conflate intensity fidelity, texture realism, and structural alignment, providing weak guidance for correcting specific failure modes.This limits reinforcement-learning-based improvement of medical image generation.
  • PRISM: PRISM decomposes image quality into verifier-grounded subrewards and uses HCP to require lower-level attributes before accumulating higher-level rewards.The hierarchy covers fine-grained properties first and structural or semantic correctness later, preventing easier objectives from masking critical failures.
  • Results: 2.3% higher mDice on PanNuke, 8.5% lower ratio MRE on CeDeM, and 5.9% higher F1 on ISIC were obtained over the closest baseline.The evaluated tasks cover cell segmentation, villi/crypt detection and measurement, and skin lesion classification.
  • Contribution: PRISM contributes a compositional reward model with hierarchical fine-to-coarse correctness and consistent improvements across segmentation, measurement, and classification benchmarks.These contributions summarize the framework and its cross-task evaluation.

2 Related Work

Related work covers diffusion models for conditional medical image synthesis, reinforcement-learning alignment of diffusion models, and reward-based optimization strategies. These approaches establish the generation and alignment context for PRISM.

  • Synthetic Medical Image Generation: Diffusion models support conditioning through text, class labels, and spatial controls, while ControlNet adds structural inputs suited to mask-conditioned medical image generation.Prior medical applications include nuclei and pathology synthesis.
  • Diffusion Models with Reinforcement Learning: Reinforcement-learning methods formulate diffusion denoising as a multi-step decision process and use policy optimization with reward signals.DDPO, DPOK, and DRaFT are identified as representative approaches, with some using KL regularization for training stability.
  • Reward Modeling for Diffusion: Reward-based optimization can select the best of multiple generated samples, but inference-time selection does not correct underlying generation errors.This motivates approaches that directly optimize the diffusion model with reward signals.

3 Method

PRISM trains conditional diffusion models with verifier-grounded rewards that progress from appearance properties to structural and semantic correctness. Hierarchical Constrained Propagation gates later rewards on earlier success, while task-specific components cover medical imaging attributes across PanNuke, CeDeM, and ISIC.

  • Conditional diffusion generator: The initial ControlNet generator models medical images conditioned on a semantic mask M and text prompt T.It is trained to model pθ0(I | M, T) using the standard denoising objective.
  • Compositional reward model: PRISM replaces a single outcome reward with verifier-grounded components targeting distinct image-quality attributes.Examples include distributional appearance matching, segmentation agreement, and classification-based scores.
  • Hierarchical constrained propagation: HCP orders rewards from low-level attributes to higher-level properties, preventing structural or semantic credit from masking failures in intensity and texture.Later components are gated by cumulative success in the hierarchy, with a minimum gate γ = 0.1 preserving gradient signal.
  • Task-specific reward design: PanNuke evaluates intensity, texture, and segmentation in that order using class-specific appearance descriptors and a frozen SegFormer.CeDeM instead combines anchor, detection, and morphological consistency, while ISIC evaluates appearance, attribute-mask agreement, and diagnosis.
  • Policy optimization: RL fine-tuning treats diffusion denoising as a multi-step MDP and optimizes the CRM reward vector with GDPO.The current policy πθ is initialized from π0, candidates are sampled per conditioning input, and updates use a PPO-style trust region.

4 Experiments

PRISM is evaluated on three medical imaging benchmarks covering segmentation, measurement, and classification, with task-specific datasets, models, and metrics. The experiments use ControlNet-based conditional generation at 512×512 resolution.

  • Datasets and Tasks: PRISM is evaluated on PanNuke, CeDeM, and ISIC across segmentation, measurement, and classification tasks.The benchmarks cover nuclei segmentation, villus-to-crypt measurement, and multiclass skin lesion diagnosis.
  • Datasets and Tasks: PanNuke uses SegFormer-B3 for five-class nuclei segmentation, with 1143 training, 245 validation, and 246 test samples.
  • Datasets and Tasks: CeDeM uses MeasureNet to compute villus-to-crypt length ratios from annotated villi and crypts, with 1169 training, 146 validation, and 146 test samples.
  • Datasets and Tasks: ISIC-2019 uses a ResNet-50 backbone for multiclass skin lesion classification, with 2339 training, 252 validation, and 1297 test images.
  • Implementation Details: ControlNet is initialized from Stable Diffusion 2.1 and trained at 512×512 resolution using AdamW with learning rate 10^-5 for up to 250 epochs.
  • Evaluation Metrics: Evaluation reports mean IoU and mean Dice for PanNuke, MAE and MRE for CeDeM, and accuracy and macro F1 for ISIC.

5 Results

PRISM’s CRM improves downstream performance across CeDeM, PanNuke, and ISIC, while ablations show that reward composition and hierarchical order matter. Qualitative and quantitative comparisons identify CRM as the strongest structured-reward approach.

  • CRM achieves the best downstream performance across CeDeM, PanNuke, and ISIC, reaching 39.58 MAER, 61.54 mIoU, and 64.00 accuracy, respectively.Compared with original-data-only training, these results correspond to a 24.0% MAER reduction on CeDeM, +5.18 mIoU on PanNuke, and +10.57 accuracy points on ISIC.
  • CRM samples better align with conditioning masks while maintaining realistic appearance, and improvements are statistically significant on CeDeM, PanNuke, and ISIC (p < 0.001).The qualitative comparisons include villus shoulder and crypt boundaries, nuclei texture, and lesion borders.
  • Subreward ablation: Ablating CeDeM subrewards exposes distinct errors: ratio misses villus shoulders, detection misplaces villus starting points, and anchor alone preserves locations but misses detection.The complete composition improves MAER to 41.87, with adding R1 yielding a further −0.17 MAER improvement.
  • Reward-order ablation: CRM is order sensitive: the canonical R1→R2→R3 sequence extracts the maximum reward, whereas placing R3 first prevents effective extraction of R1 and R2 signals.The order comparison uses mean CRM over the last 15 epochs and moving-average training trajectories.
  • Limitations: PRISM requires domain expertise to choose task-specific verifiers and define the hierarchical reward order.The paper reports fine-grained appearance cues before a final task reward as beneficial for class-level distinct attributes on PanNuke and ISIC.

6 Conclusion

PRISM aligns mask-conditional diffusion models with compositional rewards and hierarchical constrained propagation. Across PanNuke, CeDeM, and ISIC, it improves downstream performance over multiple baselines.

  • PRISM uses verifier-grounded rewards for distinct failure modes, combined with hierarchical constrained propagation to enforce fine-to-coarse correctness.
  • PRISM consistently improves over SFT, ORM, and SUM/MaxMin aggregation baselines across PanNuke, CeDeM, and ISIC downstream tasks.
  • The same PRISM recipe is applicable to other conditional medical image generation applications.

A.1 Dataset Visualization

Figure 5 pairs conditioning masks with real target images across CeDeM, PanNuke, and ISIC. The masks encode task-specific anatomical, cellular, or lesion information.

  • Figure 5 shows representative conditioning mask–target pairs for the three medical imaging benchmarks.
  • CeDeM masks encode villi, crypts, villus shoulders, and crypt borders for measurement.
  • PanNuke masks encode nuclei classes for segmentation, while ISIC masks encode lesion extent and dermoscopic attributes for diagnostic generation.

A.2 ControlNet Architecture and Training

ControlNet is trained separately for each dataset using semantic masks and text prompts, then fine-tuned with dataset-specific reward cascades. The cascades prioritize low-level appearance and structural prerequisites before semantic or measurement objectives.

  • ControlNet training: ControlNet is fine-tuned on Stable Diffusion 2.1 separately for each dataset using a color-mapped semantic mask and text prompt.
  • ControlNet training: Only the ControlNet branch is updated, while the pretrained diffusion backbone and text encoder remain frozen.
  • Checkpointing: EMA-smoothed supervised losses show that all three initial policies reach stable plateaus before GDPO fine-tuning.
  • Checkpointing: Checkpoints are selected by the lowest validation FID rather than raw diffusion loss, including within GDPO runs.
  • Dataset-specific rewards: PanNuke rewards assess class-conditional intensity, texture, and segmentation agreement, with averaging that gives rare foreground classes equal contribution.
  • Dataset-specific rewards: CeDeM rewards progress from anchor consistency to global morphology and then villus-to-crypt ratio consistency.
  • Dataset-specific rewards: ISIC rewards assess lesion and skin intensity, attribute-mask consistency, and diagnostic classification margin.
  • Reward propagation: MIN-GATE uses earlier component scores to limit later reward dominance, while GDPO normalizes components within each group.

B.1 Full Per-Dataset Results

This section provides complete per-dataset metric tables corresponding to the main paper’s Table 1, with additional CeDeM and ISIC metrics.

  • The section reports complete per-dataset metric tables summarized in Table 1 of the main paper.
  • Additional tables provide measurement metrics for CeDeM and complete classification metrics for ISIC.

C Qualitative Performance

Figures 7–9 provide additional qualitative comparisons for CeDeM, PanNuke, and ISIC. The comparisons use shared conditioning masks and, for ISIC, include SFT ControlNet, ORM-aligned, and CRM-aligned generations alongside originals.

  • Additional qualitative comparisons cover PanNuke, CeDeM, and ISIC.
  • CeDeM and PanNuke generated columns share the conditioning mask of their corresponding original target images.
  • ISIC comparisons show the input mask, SFT ControlNet, ORM-aligned, CRM-aligned, and original images under the same conditioning.
Loading 2609.05028v1…