Reinforcement Learning Improves LLM Accuracy and Reasoning in Disease Classification from Radiology Reports

1. Problem Statement

The paper addresses the challenge of accurately classifying diseases from free-text radiology reports using Large Language Models (LLMs). While Supervised Fine-Tuning (SFT) on lightweight LLMs improves classification accuracy, it often suffers from a critical limitation: catastrophic forgetting of reasoning capabilities.

• The Trade-off: Standard SFT optimizes for the final label (disease list) but typically omits detailed reasoning steps during training. Consequently, models may achieve high accuracy but fail to generate explicit, coherent, and clinically relevant justifications for their predictions.
• Clinical Necessity: In radiology, transparency and explainability are paramount for clinical decision-making and trust. A "black box" prediction is insufficient; the model must demonstrate why a disease was detected based on evidence in the report.
• Constraints: Deploying massive models via external APIs raises privacy concerns, while local deployment of large models requires substantial hardware. Lightweight LLMs (e.g., 3B–8B parameters) offer a solution but require optimization to match the performance of larger models without sacrificing reasoning.

2. Methodology

The authors propose a two-stage training and inference framework designed to enhance both classification accuracy and reasoning quality without requiring explicit reasoning supervision labels (which are expensive to create).

A. Two-Stage Training Pipeline

1. Stage 1: Supervised Fine-Tuning (SFT)
   • Input: Radiology reports paired with ground-truth disease labels.
   • Process: The lightweight LLM is fine-tuned to predict the disease list.
   • Constraint: The reasoning section in the output is intentionally left empty or omitted during this stage. The model learns to map reports to labels but loses its ability to articulate the "thought process."
2. Stage 2: Group Relative Policy Optimization (GRPO)
   • Algorithm: A Reinforcement Learning (RL) algorithm applied to the same dataset used in Stage 1.
   • No Reasoning Labels: Crucially, this stage does not require ground-truth reasoning text.
   • Reward Function: The model is optimized using a composite reward function:
     • Accuracy Reward: Based on precision and recall of the predicted disease list against ground truth.
     • Format Reward: Enforces a specific output structure (<reasoning>...</reasoning><answer>...</answer>). If the reasoning section is empty or merely repeats the answer, the reward is zero.
     • Goal: This forces the model to generate explicit reasoning to satisfy the format constraint while simultaneously maximizing classification accuracy.

B. Inference Strategy

To maximize robustness and reasoning quality, the inference process involves:

1. Ensemble Generation: The model generates five independent predictions (including reasoning) for each report.
2. Majority Voting: The final disease list is determined by majority voting across the five predictions to improve classification reliability.
3. Reasoning Summarization: A separate prompt asks the original (unfined-tuned) lightweight LLM to synthesize the five generated reasoning paths into a single, coherent clinical rationale based on the majority-voted disease list.

3. Experimental Setup

• Datasets:
  • MIMIC-CXR: Used for training (2,000 reports) and evaluation (100 gold-standard, 1,000 silver-standard).
  • NIH-CXR: 50 reports (Gold-standard).
  • MIDRC: 50 reports (Gold-standard).
  • Annotation: Gold-standard labels were curated by board-certified radiologists. Silver-standard labels were generated by GPT-4o.
• Models Evaluated:
  • Lightweight LLMs: LLaMA 3.1-8B, Qwen 2.5-3B, Phi-3 Mini-128K.
  • Baselines: Pre-trained models, SFT-only models, BioClinicalBERT, and MedAlpaca-7B.
• Evaluation Metrics:
  • Classification: Micro-averaged Precision, Recall, and F1-score.
  • Reasoning Quality:
    • Reasoning Recall: Proportion of ground-truth diseases mentioned in the generated reasoning.
    • Reasoning Comprehensiveness: Proportion of predicted diseases supported by explicit reasoning.
  • Evaluation of Reasoning: Performed by GPT-4o and Gemini 2.0 Flash, and validated by two human radiologists.

4. Key Results

• Classification Accuracy:
  • SFT significantly outperformed baseline models across all datasets.
  • GRPO further improved performance: In 8 out of 9 evaluation cohorts, adding GRPO to SFT yielded higher Micro-F1 scores.
  • Notable Gains: Qwen-3B on MIMIC-CXR saw a 13.2% F1 improvement; Phi-3 on NIH-CXR saw 12.8%.
  • Model Scale: Larger models (Qwen-7B) outperformed smaller ones, but the GRPO strategy remained effective across scales.
• Reasoning Capabilities:
  • SFT Failure: SFT alone caused a collapse in reasoning capabilities (e.g., LLaMA's reasoning recall dropped to 0% due to catastrophic forgetting).
  • GRPO Recovery: GRPO successfully restored and enhanced reasoning. For Qwen and Phi-3, reasoning recall and comprehensiveness improved significantly in 8/9 cohorts.
  • Human Validation: In a blind study with two radiologists, the SFT+GRPO outputs were preferred over SFT-only outputs in the majority of cases (30/50 and 18/50 respectively) based on consistency with the report and ground truth.
• Comparison with Alternatives:
  • vs. SFT with Reasoning Labels: SFT+GRPO achieved comparable or better classification accuracy than SFT+Reasoning (fine-tuning on a small set of reasoning data) but required no explicit reasoning labels, making it more scalable.
  • vs. BioClinicalBERT: Lightweight LLMs with GRPO generally outperformed BERT-based models in classification and provided the added benefit of natural language reasoning.

5. Significance and Contributions

1. Solving the Accuracy-Reasoning Trade-off: The paper demonstrates that Reinforcement Learning (specifically GRPO) can decouple classification accuracy from reasoning capability. It allows lightweight models to achieve high accuracy and generate coherent clinical justifications without the cost of creating massive reasoning datasets.
2. Scalability for Clinical Deployment: By utilizing lightweight models (3B–8B) that can run locally, the approach addresses privacy concerns while maintaining state-of-the-art performance.
3. Novel Reward Design: The use of a format-based reward to incentivize reasoning generation without ground-truth reasoning labels is a key methodological innovation. It leverages the model's ability to follow instructions to "think" rather than just "answer."
4. Robust Inference: The ensemble approach (5 inferences + majority voting + summarization) significantly boosts reliability, making the system suitable for real-world clinical assistance.

6. Limitations

• Reward Function: The current reward function focuses on accuracy and format but does not explicitly supervise the quality or length of reasoning, nor does it ensure the reasoning is grounded in specific report phrases (though this is partially addressed by the format constraint).
• Evaluation Bias: Reasoning quality is evaluated by other LLMs (GPT-4o/Gemini), which may introduce "LLM-as-a-judge" biases, though human validation supported the findings.
• Uncertainty Quantification: The model does not currently provide calibrated probability estimates for predictions, which is crucial for clinical risk assessment.

Conclusion

This study establishes that Reinforcement Learning (GRPO) is a powerful tool for refining lightweight LLMs in medical NLP. By optimizing for accuracy and output format simultaneously, the proposed framework overcomes the reasoning degradation typically seen in SFT, producing models that are not only accurate disease classifiers but also capable of providing transparent, evidence-based clinical reasoning.