PaperlayerView exact revision on arXiv

Source-linked AI summary

SDR: Set-Distance Rewards for Radiology Report Generation

Halil Ibrahim Gulluk, Max Van Puyvelde, Wim Van Criekinge, Olivier Gevaert

arXiv:2606.00440v1cs.AI

TL;DR

Chest X-ray reports contain unordered, independent findings, limiting the fit of exact-match and step-level rewards designed for causal reasoning. The paper represents reports as sentence-embedding sets and uses set-to-set distances for GRPO rewards and test-time selection. Across training and inference settings, this signal improves report quality and enables token-saving candidate pruning.

  • Problem

    Chest X-ray findings are unordered and independent rather than a causal reasoning chain, making exact-match and step-level rewards poorly aligned with report generation.

  • Method

    The paper represents reports as unordered sentence-embedding sets and uses set-to-set distances as continuous, permutation-invariant GRPO rewards and inference-time signals.

  • Results

    Set-distance methods consistently outperform SFT and exact-match GRPO across three vision–language backbones and two datasets, while also improving candidate selection and reducing generation tokens.

  • Takeaways & Limitations

    Set distances provide a unified signal for post-training, best-of-N selection, and more efficient test-time scaling in chest X-ray report generation.

Abstract

from arXiv · show

Reinforcement learning with verifiable rewards has rapidly advanced reasoning in vision--language models. However, for chest X-ray report generation, the standard rewards (i.e. exact-match accuracy and step-level processes) are incompatible because the reports consist of unordered and orthogonal findings, rather than a causal reasoning chain. We address this gap with a set-based view: each report is split into sentences and embedded by a frozen sentence transformer, yielding unordered embedding sets. We propose the use of set-to-set distances between generated and reference embeddings as continuous, permutation-invariant rewards. Across two datasets and three vision--language models (Qwen3-VL-2B/4B, Gemma3-4B), post-training with set-to-set distance based rewards via GRPO consistently outperforms supervised fine-tuning and exact-match GRPO on all headline metrics (BERTScore, RadGraph F1 and CheXbert F1 by average \%6.80, \%7.82 and \%4.45 relative improvements respectively). The same set distances also enable test-time best-of-$N$ selection: scoring candidates by their distance to training-report embeddings outperforms random selection on our trained models as well as three closed-source LLMs (Mistral-Small, Gemini-2.5 Flash-Lite, GPT-4o-mini) with on average \%16.4 relative improvement on BERTScore. Used as a streaming signal, they support a more efficient form of test-time scaling: pruning low-scoring candidates mid-generation reduces generated tokens by over 50\% while preserving the Findings quality of full best-of-$N$ selection. Together these results establish set-distance rewards as a unified signal for both post-training and test-time scaling in chest X-ray report generation. Our code is publicly \href{https://anonymous.4open.science/r/Set-Distance-Rewards-CXR-BFDA}{available}.

1 Introduction

Chest X-ray findings are often independent and unordered, making step-level verification and discrete rewards poorly suited to report generation. The paper proposes set-distance rewards and extends the same signal to candidate selection and generation pruning.

  • Motivation: Chest X-ray findings may be independent and arbitrarily ordered, so step-by-step verification is less meaningful than report-level evaluation.Step-level annotations may also be unavailable, while external verification can be computationally expensive.
  • Method: Set-distance rewards embed report sentences and compare generated and reference embedding sets to provide continuous, permutation-invariant supervision during GRPO.The formulation is designed to account for the unordered and independent nature of chest X-ray findings.
  • Results: Across three vision–language backbones and two report-generation benchmarks, set-distance rewards consistently outperform SFT and discrete exact-match GRPO on evaluation metrics.
  • Test-time selection: Set distances improve closed-source generalist LLMs over random selection when choosing among sampled test-time report candidates.The procedure selects the candidate closest to training-report embedding sets.
  • Test-time pruning: Distance-based pruning during generation achieves comparable quality with substantially fewer generated tokens by stopping candidates whose running distance crosses a threshold.The method uses partial generations to reduce test-time scaling cost.

2 Related Work

Related work covers medical vision–language models, reinforcement-learning approaches for medical reasoning, and reward-design methods addressing limitations of discrete supervision.

  • Medical vision–language models: Medical vision–language models have supported medical diagnosis, image–report retrieval, visual question answering, and rationale generation.
  • Medical reasoning: MedReason, Med-R1, and MedVLM-R1 extend reasoning-oriented training across medical image modalities, including GRPO-based training.
  • Reward design: Continuous rewards are being investigated in GRPO post-training to reduce noise from binary supervision and preserve value from partially correct reasoning steps.
  • Process reward modeling: Process reward modeling methods such as reasoning-driven PRM, ER-PRM, and EDU-PRM explore structured or entropy-based ways to improve reward modeling.

3 Method

The method represents chest X-ray reports as unordered sentence-embedding sets and uses set-based semantic rewards in GRPO, alongside format validation. The same distances support inference-time candidate selection and pruning against the training-report distribution.

  • 3.1 Sentence-level report representation: Reports are split into Findings and Impression sentences, each embedded with frozen all-mpnet-base-v2 into separate unordered sets.The set representation is permutation-invariant, reflecting that individual findings are generally orthogonal rather than a causal chain.
  • 3.2 Set-based semantic reward: The format reward requires exactly one ordered <think>/<answer> template with non-empty Findings and Impression content.Malformed outputs receive no semantic reward because their sections cannot be extracted reliably.
  • 3.2 Set-based semantic reward: Chamfer and Hausdorff distances compare generated and reference sentence sets using cosine distance, preserving symmetry and permutation invariance.Chamfer averages nearest-neighbour costs, whereas Hausdorff uses worst-case costs and therefore penalizes uncovered or contradictory outliers more harshly.
  • 3.2 Set-based semantic reward: The GRPO scalar reward combines format and semantic rewards as R(ŷ, y) = λfmt Rfmt(ŷ) + λsem Rsem(ŷ, y), with λfmt = λsem = 1 in experiments.When formatting fails, the semantic reward is set to zero, penalizing the generation through both components.
  • 3.3 Inference-time response selection: At inference time, K reports are sampled and the candidate closest to cached training-report embedding sets is retained without updating the generation policy.The cached reference embeddings make the procedure applicable to both fine-tuned VLMs and closed-source generalist LLMs.
  • 3.3 Inference-time response selection: Running set distances enable early pruning of candidates whose partial generations cross a threshold, reducing test-time computation while preserving selection quality.The paper also introduces Hungarian-matching variants, Hung-NN and Hung-Pen, as additional set-distance choices.

4 Experiments

Experiments evaluate set-distance rewards across two chest X-ray datasets and three vision–language backbones. Chamfer and Hausdorff rewarding outperform supervised fine-tuning and discrete exact-match GRPO across the reported evaluation metrics.

  • Experimental setup: Three vision–language models—Qwen3-VL-2B, Qwen3-VL-4B, and Gemma3-4B—are fine-tuned and post-trained with GRPO on MIMIC-CXR and ReXGradient.The experiments use SFT, format-only, exact-match, Chamfer, and Hausdorff reward configurations.
  • GRPO post-training results: Chamfer Distance rewarding outperforms all other methods on average across evaluation metrics for ReXGradient Findings.Table 1 compares the different GRPO reward functions.
  • GRPO post-training results: Hausdorff Distance rewarding has the highest average performance on MIMIC-CXR Findings.The MIMIC-CXR comparison is summarized in Table 2.
  • Overall findings: Set-distance rewarding improves report-generation performance over supervised fine-tuning and discrete exact-match GRPO across the evaluated models and metrics.The reported headline metrics include BERTScore F1, COMET, RadGraph averaged F1, and CheXbert macro F1.

5 Inference-time response selection

Inference-time selection scores multiple candidate reports against the training distribution using set distances, while pruning applies the same signal during partial decoding. The resulting procedure improves selection over random picking and substantially reduces token usage, though it does not always match full-generation selection.

  • Response selection: Table 3 reports the best distance-based selection score for each model and metric alongside a matched random-selection baseline and percentage improvement.The table aggregates distance-metric and aggregation-method combinations across five clinically meaningful NLP metrics.
  • Response selection: Appropriate distance and aggregation choices improve overall response-selection performance substantially over random selection.The comparison is averaged over multiple runs and uses the training distribution to select among generated candidates.
  • Distance-guided pruning: Full-generation selection requires completing all K candidates before scoring, making each test-image generation cost exactly K single-decoding costs.Partial-generation pruning is introduced to reduce this cost by eliminating candidates before completion.
  • Distance-guided pruning: Pruning encodes active candidates after each generated sentence, scores partial outputs against the training distribution, and drops the bottom-scoring fraction before continuing.The experiments use a pruning fraction ρ = 0.5 and continue sentence by sentence until one candidate remains.
  • Distance-guided pruning: The pruning policy is also evaluated as a proof of concept on five closed-source LLMs, including Mistral-Small, Gemini variants, and GPT-4o mini and GPT-5 mini.The reported token savings are framed for an equivalent open-weights deployment with the same inference budget.
  • Distance-guided pruning: 42.1–60.1% of generation tokens are saved while pruning improves random selection by +12.7% on BS-F1, +17.1% on RG-F1, and +6.2% on CXB-F1.Pruning does not always match standard full-generation selection, but the reported metric gap is small and token savings are substantial.

6 Conclusion

The paper presents set distances as a unified signal for chest X-ray report generation across training and inference. The same signal supports GRPO rewards, best-of-N selection, and token-saving pruning.

  • Conclusion: Reports are represented as unordered sentence-embedding sets, and set-to-set distances provide a continuous, permutation-invariant reward.The distances score generated reports against references during GRPO training.
  • Conclusion: Chamfer- and Hausdorff-based rewards outperform supervised fine-tuning and discrete exact-match GRPO across three vision–language backbones and two datasets.This conclusion summarizes the post-training comparison.
  • Conclusion: At inference time, the same distances improve best-of-N selection for GRPO-fine-tuned models and closed-source generalist LLMs over random selection.The selection criterion is gradient-free and compares candidate reports with the training distribution.
  • Conclusion: Partial-decoding pruning retains the quality benefits of full best-of-N selection while cutting roughly half of generated tokens.The pruning policy applies the distance signal during generation rather than after all candidates are complete.

A Set-to-set distance metrics

The paper compares set distances that capture local coverage, transport, assignment, cardinality mismatch, and internal geometry between sentence-embedding sets. These distances are used as candidate semantic rewards for GRPO.

  • Nearest-neighbour metrics: Chamfer averages nearest-neighbour costs in both directions, rewarding generated-sentence coverage while penalising uncovered reference sentences.It permits multiple sources to reuse the same target sentence.
  • Nearest-neighbour metrics: Hausdorff replaces Chamfer’s averages with maxima, so one uncovered sentence can dominate the score.This makes Hausdorff a harsher reward than Chamfer.
  • Transport metrics: Optimal transport assigns uniform mass to both sets and models how mass moves between their sentence embeddings.The resulting distance is converted into the similarity reward 1 − D.
  • Transport metrics: Sinkhorn distance provides a smoother, faster proxy for optimal transport through entropy regularisation, with ε = 0.01, 0.1, or 0.5.Smaller ε approaches exact optimal transport at higher iteration cost, while larger ε increases smoothing.
  • Transport metrics: Unbalanced and partial optimal transport relax mass constraints when set sizes differ or some sentences lack counterparts.Unbalanced OT uses KL penalties, while partial OT discards a fraction 1 − ρ of mass at zero cost.
  • Geometry-based metrics: Gromov–Wasserstein compares pairwise internal distances, making it invariant to isometries when absolute embedding positions shift but internal structure remains.It is undefined when min(n, m) < 2 and then falls back to zero similarity.
  • Assignment metrics: Assignment-based metrics enforce one-to-one matching on the smaller set and handle unmatched elements through nearest-neighbour costs or a count penalty.The count-penalty variant evaluates α = 0.1 and 0.5, with larger α penalising incorrect sentence counts.

C.1 Datasets

The experiments use MIMIC-CXR and ReXGradient with separate training and validation splits, alongside specified H100 training configurations and cached reference embeddings for selection.

  • Datasets: MIMIC-CXR contributes 179,778 training samples and 45,364 validation samples.
  • Datasets: ReXGradient contributes 238,968 training samples and 17,007 validation samples from the original dataset split.
  • Evaluation setup: Response selection draws K = 5 generations per test image and caches reference embeddings for N = 5000 training samples per dataset, while evaluating N = 1000 test samples.The cache is used to accelerate sample-to-training-distribution distance calculations and limit API costs.

D Full GRPO post-training results

The appendix reports full GRPO post-training results across both datasets, report sections, model averages, random seeds, and multiple metric families. Results are organised to compare reward functions across NLP and clinical evaluation metrics.

  • Reporting protocol: All reported post-training results average across 5 random seeds and include sample standard deviations.
  • Aggregation: Final Mean blocks average each reward–metric cell across the models in the corresponding dataset.
  • ReXGradient: The ReXGradient tables cover Findings and Impression sections across BLEU, ROUGE, embedding-based or lexical metrics, and clinical metrics.
  • MIMIC-CXR: The MIMIC-CXR tables provide the same metric-family breakdown for Findings and Impression sections.

E Full selection-results breakdown

The selection-results breakdown evaluates distance-based policies against random selection across models, metrics, strata, and report sections. A method-by-metric heatmap visualises where policies outperform or lose to random selection.

  • Headline selection results: Headline selection tables report the best score from any distance-metric and aggregation policy for five clinically meaningful NLP metrics.Matched random-selection baselines and percentage improvements are shown for each model.
  • Heatmap view: Figure 3 encodes mean percentage improvement over random across 13 models and five metrics, with teal cells indicating gains and coral cells indicating losses.
  • Per-stratum breakdown: The appendix reports absolute scores and differences versus random for every model, selection policy, ground-truth-defined stratum, and per-sample metric.

F.1 Stratified tables (four metrics)

The stratified tables report selection performance across abnormality strata and four metrics, while accompanying figures visualize per-model improvements over random. The appendix also specifies a distance-guided pruning procedure that progressively eliminates candidates during decoding.

  • Stratified tables: The stratified results divide patients into no-finding, single-finding, and multiple-finding groups with 504, 206, and 290 cases, respectively.The best policy is selected by mean improvement over random across stratum-by-metric cells.
  • Tables: Tables 25–29 report mean percentage improvement over random across distance metrics, aggregations, pathology strata, and pruning policies.Table 25 covers all patients; Tables 26–28 cover the three strata; Table 29 evaluates distance-guided pruning on Findings.
  • Figures: Figures 4 and 5 visualize per-model, per-distance percentage improvements over random for BERTScore F1 and RadGraph F1 using Avg aggregation.Teal cells indicate improvement over random, whereas coral cells indicate the opposite.
  • Evaluation scope: The appendix evaluates NLP and clinical metrics for Findings and Impression under Random, Standard, and Pruning policies across every model–distance combination.The reported metrics include BERTScore F1, RadGraph F1, and CheXbert F1 for pruning comparisons.
  • Distance-guided pruning: The pruning algorithm samples candidates’ opening sentences in lock-step, scores partial outputs against the training distribution, and repeatedly drops the highest-scoring fraction.The surviving candidate is then decoded to its end-of-sequence token and returned.

H Qualitative examples

The qualitative examples examine whether distance-to-training-distribution selection chooses candidates that are also closest to the ground truth by BERTScore-F1. Across the shown cases, selected candidates are contrasted with rejected alternatives under several models and distance configurations.

  • Example setup: The examples use ReXGradient chest-X-ray reports and compare the ground truth, selected candidate, and a rejected candidate.The selected candidate is chosen by distance to the training distribution, while the rejected candidate has lower BERTScore-F1 against the ground truth.
  • Configurations: One Qwen3-VL-4B example uses GRPO with Partial OT distance, min aggregation, and K=10 candidates.The displayed selected candidate has BERTScore-F1 0.203 and distance 0.599.
  • Selection behavior: The qualitative rule selects the candidate with the lowest training-distribution distance and the highest ground-truth BERTScore-F1.Figures 6–8 show rejected alternatives with larger selected-metric distances and correspondingly lower BERTScore-F1.
  • Configurations: Two Gemini 2.5 Flash-Lite examples use Hungarian + Pen with avg aggregation or Hausdorff with min aggregation, each sampling K=5 candidates.Each panel reports candidate BERTScore-F1 and distance to the training distribution.
Loading 2606.00440v1…