Technical Summary: Set-Distance Rewards for Radiology Report Generation (SDR)

Problem Statement

The generation of chest X-ray reports using vision–language models (VLMs) presents a unique challenge for Reinforcement Learning (RL) post-training, specifically regarding the design of reward functions. Standard RL approaches for reasoning tasks often rely on process reward models that verify intermediate steps in a causal chain. However, radiology reports do not follow a sequential or causal reasoning structure; instead, they consist of unordered and orthogonal findings (e.g., "heart size may be enlarged" and "right lung is clear" are independent observations).

Consequently, standard discrete rewards based on exact-match accuracy are incompatible with this domain. They induce noise by penalizing reports that contain correct findings but in a different order or with slight phrasing variations, failing to capture the clinical reality that the set of findings matters more than their sequence. Furthermore, step-level annotations required for process rewards are often unavailable, and verifying each step via external sources is computationally expensive.

Methodology

The authors propose Set-Distance Rewards (SDR), a framework that treats radiology reports as unordered sets of sentence embeddings rather than sequential text. The methodology operates in two phases: post-training (GRPO) and inference-time scaling.

1. Sentence-Level Representation

A chest X-ray report $y$ is split into two sections: Findings ($y_F$) and Impression ($y_I$). Each section is segmented into individual sentences. A frozen, pre-trained sentence transformer (specifically all-mpnet-base-v2) embeds each sentence $s$ into a fixed-dimensional vector $e \in \mathbb{R}^d$. The report is thus represented as two unordered sets of embeddings:
$$E_F(y) = \{E_\phi(s^F_i) \mid 1 \le i \le n_F\}, \quad E_I(y) = \{E_\phi(s^I_j) \mid 1 \le j \le n_I\}$$
This representation is permutation-invariant, aligning with the clinical intuition that the order of findings carries no diagnostic meaning.

2. Set-to-Set Distance Rewards (Post-Training)

During Group Relative Policy Optimization (GRPO), the model generates a report $\hat{y}$. The semantic reward is calculated by computing the distance between the set of generated embeddings and the set of reference (ground truth) embeddings.

• Distance Metrics: The authors utilize symmetric set-to-set distances normalized to $[0, 1]$, specifically Chamfer distance and Hausdorff distance, built upon cosine distance between unit-norm embeddings.
  • Chamfer: Averages the nearest-neighbor costs, acting as a "soft coverage" reward.
  • Hausdorff: Takes the maximum nearest-neighbor cost, acting as a "worst-case" penalty for missing critical findings.
• Reward Formulation: The total reward $R(\hat{y}, y)$ is a weighted sum of a binary format reward (ensuring the output follows a structured template) and the continuous semantic reward derived from the set distance:
  $$R(\hat{y}, y) = \lambda_{fmt} R_{fmt}(\hat{y}) + \lambda_{sem} (1 - D(E_S(\hat{y}), E_S(y)))$$
  where $D$ is the set distance. This provides a continuous, differentiable signal that rewards partial correctness and handles unordered findings.

3. Inference-Time Scaling and Pruning

The same set-distance metrics are applied at inference time without gradient updates:

• Best-of-N Selection: For a given image, the model generates $K$ candidate reports. Each candidate's embedding set is compared against the pre-computed embedding sets of the training corpus. The candidate with the minimum distance to the training distribution is selected. This serves as a gradient-free prior for clinically plausible outputs.
• Streaming Pruning: To reduce computational cost, the system computes the running set distance of partial generations during decoding. Candidates whose partial distance exceeds a threshold are pruned early. This "prune-then-decode" loop reduces token generation by over 50% while maintaining quality comparable to full Best-of-N selection.

Key Contributions

1. Set-Distance Reward Functions for GRPO: The paper addresses the infeasibility of process reward modeling for radiology by treating reports as unordered sets. Using Chamfer and Hausdorff distances as continuous rewards, the method consistently outperforms Supervised Fine-Tuning (SFT) and discrete exact-match GRPO across three VLM backbones (Qwen3-VL-2B/4B, Gemma3-4B) and two datasets (MIMIC-CXR, ReXGradient).
2. Set-Distance-Guided Test-Time Selection: The authors demonstrate that the same distance metrics can serve as a selection criterion for both fine-tuned models and closed-source generalist LLMs (Mistral-Small, Gemini-2.5 Flash-Lite, GPT-4o-mini, GPT-5 mini). This approach improves performance over random selection without requiring access to model gradients or domain-specific fine-tuning.
3. Distance-Based On-the-Fly Pruning: The paper introduces a method to prune low-scoring candidates mid-generation. This reduces the total number of generated tokens by >50% while preserving the quality of the final selection, offering a more efficient form of test-time scaling.

Experimental Results

The evaluation covers headline metrics including BERTScore, RadGraph F1, and CheXbert F1.

• Post-Training Performance:
  • On the ReXGradient dataset, Chamfer-based rewarding achieved the highest performance across all models and metrics.
  • On MIMIC-CXR, both Chamfer and Hausdorff rewards reached the highest average performance.
  • Relative improvements over SFT and exact-match GRPO were significant: average improvements of 6.80% (BERTScore), 7.82% (RadGraph F1), and 4.45% (CheXbert F1).
• Inference-Time Selection:
  • Applying set-distance selection to closed-source LLMs yielded substantial gains over random selection. For example, on BERTScore, improvements ranged from +11.2% to +18.7% depending on the model and prompt template.
  • The method successfully improved performance on generalist models (e.g., GPT-4o mini, Gemini) that were not fine-tuned on the specific medical domain.
• Pruning Efficiency:
  • The pruning strategy saved 42.1% to 60.1% of generation tokens.
  • Pruned generations maintained comparable quality to full Best-of-N selection, with average relative improvements over random selection of +12.7% (BS-F1), +17.1% (RG-F1), and +6.2% (CXB-F1).

Significance and Claims

The paper claims to establish set-distance rewards as a unified signal for both post-training and test-time scaling in chest X-ray report generation.

• Unified Signal: The same mathematical formulation (set-to-set distance) effectively drives the RL post-training process and serves as a robust, gradient-free selection criterion at inference time.
• Domain Adaptation: By acknowledging the unordered nature of radiology findings, the method avoids the noise introduced by sequential constraints, making it specifically suited for medical report generation where findings are often independent.
• Efficiency: The work demonstrates that high-quality report generation does not necessarily require generating all tokens for all candidates; early pruning based on set distance offers a viable path to reduce inference costs.
• Generalizability: The ability to apply these selection rules to closed-source LLMs suggests a pathway to improve clinical AI outputs without the need for expensive domain-specific fine-tuning or access to model weights.

The authors conclude that their approach provides a practical, continuous reward signal that aligns with the clinical structure of radiology reports, outperforming traditional discrete or sequential reward mechanisms.