PaperlayerView exact revision on arXiv

Source-linked AI summary

MediX-R1: Open Ended Medical Reinforcement Learning

Sahal Shaji Mullappilly, Mohammed Irfan Kurpath, Omair Mohamed, Mohamed Zidan, Fahad Khan, Salman Khan, Rao Anwer, Hisham Cholakkal

arXiv:2602.23363v1cs.CV

TL;DR

Medical open-ended reasoning lacks reliable training and evaluation signals because exact-match and string-overlap methods mishandle clinical paraphrases, reasoning quality, and multimodal evidence. MediX-R1 addresses this with Group Based RL, composite semantic and structural rewards, and unified Reference-based LLM judging, achieving strong results across text-only and image+text benchmarks while remaining a research prototype rather than a deployment-ready clinical system.

  • Problem

    MCQ- and string-matching-based medical pipelines poorly support open-ended answers because they under-reward valid paraphrases and fail to assess reasoning quality or modality recognition.

  • Method

    MediX-R1 fine-tunes a multimodal backbone with Group Based RL using LLM accuracy, medical embedding, format, and modality rewards, plus a unified Reference-based LLM-as-judge evaluation framework.

  • Results

    MediX-R1 achieves strong results across diverse medical LLM and VLM benchmarks, with consistent gains across Group Based RL optimizers and backbone models using approximately 51K instructions.

  • Takeaways & Limitations

    The results support open-ended RL with comprehensive reward signals and LLM-based evaluation as a practical approach for multimodal medical reasoning.

  • Takeaways & Limitations

    MediX-R1 is a research prototype not intended for clinical or commercial deployment and may hallucinate findings, omit differentials, or overstate certainty.

Abstract

from arXiv · show

We introduce MediX-R1, an open-ended Reinforcement Learning (RL) framework for medical multimodal large language models (MLLMs) that enables clinically grounded, free-form answers beyond multiple-choice formats. MediX-R1 fine-tunes a baseline vision-language backbone with Group Based RL and a composite reward tailored for medical reasoning: an LLM-based accuracy reward that judges semantic correctness with a strict YES/NO decision, a medical embedding-based semantic reward to capture paraphrases and terminology variants, and lightweight format and modality rewards that enforce interpretable reasoning and modality recognition. This multi-signal design provides stable, informative feedback for open-ended outputs where traditional verifiable or MCQ-only rewards fall short. To measure progress, we propose a unified evaluation framework for both text-only and image+text tasks that uses a Reference-based LLM-as-judge in place of brittle string-overlap metrics, capturing semantic correctness, reasoning, and contextual alignment. Despite using only $\sim51$K instruction examples, MediX-R1 achieves excellent results across standard medical LLM (text-only) and VLM (image + text) benchmarks, outperforming strong open-source baselines and delivering particularly large gains on open-ended clinical tasks. Our results demonstrate that open-ended RL with comprehensive reward signals and LLM-based evaluation is a practical path toward reliable medical reasoning in multimodal models. Our trained models, curated datasets and source code are available at https://medix.cvmbzuai.com

1. Introduction

MediX-R1 addresses the mismatch between open-ended clinical reasoning and MCQ- or string-matching-based training and evaluation. It combines Group Based RL, multi-signal rewards, and Reference-based LLM judging to support free-form multimodal medical answers.

  • Open-ended clinical tasks require context-sensitive, multimodal answers, but MCQ and string-matching pipelines under-reward paraphrases and poorly assess reasoning or modality recognition.
  • MediX-R1 fine-tunes a baseline multimodal backbone with Group Based RL and a composite reward for clinical reasoning.
  • The reward combines strict LLM-based semantic correctness, medical embedding similarity, structured-output formatting, and explicit modality recognition.
  • Reference-based LLM-as-judge evaluation unifies text-only and image+text benchmarking while capturing semantic correctness, reasoning adequacy, and contextual alignment.
  • MediX-R1 uses approximately 51K instructions and reports strong results across medical LLM and VLM benchmarks, with consistent gains across Group Based RL optimizers and backbone models.

2. Open Ended Medical RL

MediX-R1 trains structured clinical responses with Group Based RL, using group-relative rewards and a composite objective that combines semantic, format, and modality signals. The design targets stable open-ended learning where exact-match or executable rewards are inadequate.

  • Given an image and question, the model generates a modality-tagged response containing <think> reasoning and a concise <answer>.
  • Group Based RL samples candidate completions, assigns each a scalar composite reward, and computes a standardized group-relative advantage without a learned value function.
  • GRPO uses PPO-style clipping and KL regularization, while DAPO and GSPO provide alternative Group Based RL optimizers for the same composite reward.
  • GSPO: GSPO replaces token-wise importance ratios with a length-normalized sequence-level ratio, reducing intra-sequence variance for noisy or long outputs.
  • Reward Design: The composite reward is r = wfmt Rformat + wllm Rllm + wemb Rembed + wmod Rmodality, with weights selected through staged ablation.
  • Reward Design: LLM judging and medical embeddings provide paraphrase-robust semantic feedback, while format and modality rewards regularize structure and imaging-modality grounding.

3. Evaluation Framework

The evaluation framework processes model outputs through generation, Reference-based LLM judging, and score aggregation. It supports binary assessment for short tasks and rubric-based assessment for long-form medical reports across text-only and multimodal benchmarks.

  • The pipeline has three stages: vLLM generation, Reference-based LLM-as-judge evaluation, and aggregation of judgment outputs.
  • Generation evaluates both text-only and image+text tasks spanning QA, MCQ, and long-form report generation.
  • The BASE template produces binary decisions for open-ended, one-word, and MCQ-style questions, while the MIMIC template scores long-form reports against clinical criteria.
  • A clinically equivalent paraphrase such as “Low blood flow or less perfusion” can receive a correct score against the reference “Low perfusion.”
  • String-overlap metrics often miss correct clinical paraphrases and cannot assess justification quality or contextual alignment, motivating Reference-based judging.

4. Experiments and Results

MediX-R1 is evaluated across unified text-only and image+text medical benchmarks, with strong performance across standard, open-ended, and clinically complex tasks. Ablations indicate that composite rewards and multi-model validation improve robustness and training stability.

  • State-of-the-art comparisons: The evaluation spans LLM benchmarks including MMLU variants, MedQA, MedMCQA, USMLE-SA, PubMedQA, and MIMIC-CXR summarization, plus VLM benchmarks including SLAKE-VQA, RadVQA, PathVQA, PMC-VQA, and report generation.
  • State-of-the-art comparisons: MediX-R1 achieves the highest average score across a unified suite of text-only and image+text medical benchmarks.The suite covers medical QA, multiple-choice, report generation, visual question answering, and clinical image interpretation.
  • Reward ablations: Composite-reward training transfers across RL algorithms, with DAPO achieving the best overall average of 0.610 versus 0.590 for GRPO, 0.600 for GSPO, and 0.570 for the baseline.
  • Reward ablations: Combining LLM and embedding rewards improves text-only scores to 0.686 and image+text scores to 0.410, raising the overall average to 0.589.The combination improves robustness to paraphrases and terminology variants, while the LLM-only signal does not improve VLM performance.
  • Reward ablations: The composite reward produces the strongest validation performance, reaching 0.431 on image+text tasks, 0.687 on text-only tasks, and 0.597 overall.The full composite combines LLM accuracy, embedding semantics, modality recognition, and shared format control.
  • Human evaluation: Human experts selected MediX-R1 as the best response in 72.7% of cases, compared with 13.6% for Llama3.2-Vision, 9.2% for MedGemma, and 4.5% for HuatuoGPT-Vision.

5. Conclusion

MediX-R1 combines structure-aware composite rewards with standardized Reference-based LLM-as-judge evaluation for open-ended medical multimodal reasoning. The framework shows strong benchmark performance, improved stability, and clinically faithful outputs with interpretable reasoning traces.

  • MediX-R1 trains a baseline VLM with Group Based RL and a composite reward for open-ended medical multimodal reasoning.
  • Its composite reward couples LLM judging, medical semantic alignment, format control, and modality recognition.
  • A unified vLLM-based evaluation pipeline provides consistent, paraphrase-robust scoring across text-only and image+text tasks.
  • MediX-R1 achieves strong results across diverse medical benchmarks with improved stability and resistance to reward hacking compared with single-signal RL variants.
  • Human expert preferences and reward ablations support its clinical answer quality, semantic alignment, and the benefits of multi-signal rewards.

Impact Statement

The proposed framework targets more realistic and transparent evaluation of free-form clinical responses, but MediX-R1 remains a research prototype with substantial safety and fairness risks. Responsible use requires auditing and clinician-oriented safeguards before redistribution or deployment.

  • Composite-reward RL and unified LLM-as-judge evaluation may reduce reliance on brittle string matching for free-form clinical responses.
  • MediX-R1 is a research prototype not intended for clinical or commercial deployment.
  • The model may hallucinate findings, omit key differentials, overstate certainty, and reinforce subtle biases or false positives through its judge-based reward and evaluation.
  • Residual privacy, demographic-bias, and health-disparity risks remain despite using publicly available de-identified datasets, requiring auditing before redistribution or deployment.

A.1. Training Data and Modality Distribution

MediX-R1 was trained on 51,335 multimodal medical instruction samples from four official dataset train splits, covering 16 modality tags. The supplied table identifies the source-dataset composition.

  • 51,335 multimodal medical instruction samples were used for training across 16 modality tags.
  • The training data came from the official train splits of PMC-VQA, SLAKE, RadVQA, and PathVQA.
  • The source-dataset counts were PMC VQA 25,000, SLAKE 4,919, Rad VQA 1,793, and Path VQA 19,623.
  • Table 8 is labeled “Modality Breakdown and Source Dataset composition.”

A.2. Training Configuration

MediX-R1 used a GRPO training configuration with KL-regularized advantage computation, dynamic batching, vLLM rollouts, and mixed-precision sharded optimization. Training ran for 200 steps with specified rollout, optimizer, and hardware settings.

  • Prompts and responses were capped at 4,352 and 4,096 tokens, with rollout and validation batch sizes of 512 and 1,024.
  • The configuration enabled a hybrid engine with actor optimization using global batch size 128 and dynamic batching.
  • AdamW optimization used learning rate 1e-6, weight decay 0.01, a constant scheduler, and 200 training steps.
  • Full parameter sharding, rank-zero initialization, bf16 parameters, fp32 reductions and buffers, and parameter and optimizer offloading were enabled.
  • vLLM rollouts sampled 5 responses at temperature 1.0 with tensor parallelism size 2 and maximum 8,448 batched tokens.
  • MediX-R1 used GRPO with gamma 1.0, lambda 1.0, and low-variance KL regularization at coefficient 0.01.

A.3. Reward Coefficient Selection Details

The reward-coefficient procedure prioritizes task-facing correctness while retaining parseable outputs and avoiding exhaustive hyperparameter search. Semantic-reward variants showed low sensitivity, leading to selection of an LLM-favored default and subsequent modality grounding.

  • Selection objectives: The coefficient-selection goal was to preserve parseable outputs, emphasize task-facing correctness, and avoid an expensive hyperparameter search.The procedure was designed to make coefficient selection transparent and reproducible.
  • Ablations: The embedding-only and LLM-only ablation settings were defined as separate single-signal configurations.
  • Semantic-reward combinations: Three combined-semantic variants allocated the non-format reward equally, toward LLM correctness, or toward embedding similarity.Each variant retained a 0.10 format term while splitting the remaining 0.90 between Rllm and Rembed.
  • Semantic-reward combinations: v1, v2, and v3 achieved validation averages of 0.582, 0.589, and 0.579, respectively, indicating low sensitivity across the coarse coefficient range.
  • Default configuration: The authors selected v2 as the default combined-semantic configuration because it slightly improved aggregate performance while favoring the stricter Rllm signal.
  • Modality grounding: Modality grounding reserved 5% of the non-format budget for Rmodality, equal to 0.045 of total reward mass, and renormalized the remaining weights.
  • Final coefficients: The final coefficients were wfmt = 0.10, wmod = 0.045, wllm = 0.5175, and wemb = 0.3375, summing to 1.0.
  • Limitation: The study did not exhaustively search reward coefficients because of computational constraints, leaving broader coefficient exploration for future work.

A.4. Reward Function Source Code

The source implementation defines four reward components that parse model outputs, check structural compliance and modality tags, and score answer correctness through LLM and embedding signals.

  • Reward components: The implementation includes format, modality-recognition, LLM-based accuracy, and embedding-based semantic reward components.
  • Format reward: The format reward returns 1.0 only when the output matches the required <think>...</think><answer>...</answer> structure.
  • LLM accuracy reward: The LLM accuracy reward extracts the answer and ground truth, then returns the external LLM matching score, with a fallback of 1.0 shown for the successful path.
  • Embedding reward: The embedding reward computes tensor embeddings and cosine similarity, returning a binary score when similarity reaches the threshold.The shown implementation uses a threshold of 0.8.
  • Modality reward: The modality reward compares the predicted modality tag before <think> with the ground-truth modality tag, returning 1.0 for a case-insensitive match.

A.5. Human Expert Comparative Evaluation Protocol

The human-expert evaluation compares anonymized model responses against a reference using clinical accuracy, completeness, relevance, reasoning quality, and practical usefulness. Experts select the strongest response, while report-generation evaluation uses a reference-based 0–5 clinical rubric.

  • Expert comparison protocol: Experts compare four anonymized responses with a reference description and select the single best response based on clinical correctness, relevance, and reasoning clarity.
  • Expert comparison protocol: The protocol instructs evaluators to read the ground truth, assess model responses for accuracy and completeness, and choose the response best aligned with the reference.
  • Expert criteria: The evaluation criteria explicitly include clinical accuracy and whether reasoning traces are correct and well explained.
  • Expert findings: MediX-R1 reasoning was preferred over MedGemma in 74.2% of cases, and 92.4% of cases rated its reasoning steps as acceptable.
  • Output structure: The training prompt requires a modality tag, structured <think> reasoning, and a concise <answer>, while only <answer> is scored by semantic rewards.The <think> content is ignored for scoring but retained for interpretability.
  • Short-form judging: The short-form judge assigns a binary score by comparing the predicted answer with the reference while allowing paraphrases and option-label matches.Inference uses a separate deterministic LLM judge with strict JSON parsing.
  • Report-generation judging: For long-form reports, the judge assigns a 0–5 score for clinical accuracy, completeness, and relevance after stripping hidden reasoning.Scores are averaged across items for dataset-level metrics.
  • Report-generation judging: The report rubric ranges from 0 for completely incorrect outputs to 5 for clinically accurate, complete, and instruction-faithful matches.
Loading 2602.23363v1…