Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

OncoReason: Structuring Clinical Reasoning in LLMs for Robust and Interpretable Survival Prediction

Raghu Vamshi Hemadri, Geetha Krishna Guruju, Kristi Topollai, Anna Ewa Choromanska

TL;DR

OncoReason addresses the need for accurate and interpretable cancer outcome predictions by aligning autoregressive LLMs with clinical reasoning through multi-task learning. It integrates three training strategies—autoregressive fine-tuning, SFT with Chain-of-Thought distillation, and cold-start SFT with GRPO—on the MSK-CHORD dataset to predict binary survival and survival months while generating clinically grounded rationales. Key contributions include a reasoning-aligned training pipeline, empirical evidence that Chain-of-Thought prompts boost predictive performance and interpretability, and a GRPO-based approach that yields robust, coherent reasoning with state-of-the-art generation metrics. The work advances trustworthy clinical AI in precision oncology by delivering auditable explanations alongside predictions and provides a foundation for future multimodal and human-in-the-loop extensions.

Abstract

Predicting cancer treatment outcomes requires models that are both accurate and interpretable, particularly in the presence of heterogeneous clinical data. While large language models (LLMs) have shown strong performance in biomedical NLP, they often lack structured reasoning capabilities critical for high-stakes decision support. We present a unified, multi-task learning framework that aligns autoregressive LLMs with clinical reasoning for outcome prediction on the MSK-CHORD dataset. Our models are trained to jointly perform binary survival classification, continuous survival time regression, and natural language rationale generation. We evaluate three alignment strategies: (1) standard supervised fine-tuning (SFT), (2) SFT with Chain-of-Thought (CoT) prompting to elicit step-by-step reasoning, and (3) Group Relative Policy Optimization (GRPO), a reinforcement learning method that aligns model outputs to expert-derived reasoning trajectories. Experiments with LLaMa3-8B and Med42-8B backbones demonstrate that CoT prompting improves F1 by +6.0 and reduces MAE by 12%, while GRPO achieves state-of-the-art interpretability and predictive performance across BLEU, ROUGE, and BERTScore. We further show that existing biomedical LLMs often fail to produce valid reasoning traces due to architectural constraints. Our findings underscore the importance of reasoning-aware alignment in multi-task clinical modeling and set a new benchmark for interpretable, trustworthy LLMs in precision oncology.

OncoReason: Structuring Clinical Reasoning in LLMs for Robust and Interpretable Survival Prediction

TL;DR

OncoReason addresses the need for accurate and interpretable cancer outcome predictions by aligning autoregressive LLMs with clinical reasoning through multi-task learning. It integrates three training strategies—autoregressive fine-tuning, SFT with Chain-of-Thought distillation, and cold-start SFT with GRPO—on the MSK-CHORD dataset to predict binary survival and survival months while generating clinically grounded rationales. Key contributions include a reasoning-aligned training pipeline, empirical evidence that Chain-of-Thought prompts boost predictive performance and interpretability, and a GRPO-based approach that yields robust, coherent reasoning with state-of-the-art generation metrics. The work advances trustworthy clinical AI in precision oncology by delivering auditable explanations alongside predictions and provides a foundation for future multimodal and human-in-the-loop extensions.

Abstract

Predicting cancer treatment outcomes requires models that are both accurate and interpretable, particularly in the presence of heterogeneous clinical data. While large language models (LLMs) have shown strong performance in biomedical NLP, they often lack structured reasoning capabilities critical for high-stakes decision support. We present a unified, multi-task learning framework that aligns autoregressive LLMs with clinical reasoning for outcome prediction on the MSK-CHORD dataset. Our models are trained to jointly perform binary survival classification, continuous survival time regression, and natural language rationale generation. We evaluate three alignment strategies: (1) standard supervised fine-tuning (SFT), (2) SFT with Chain-of-Thought (CoT) prompting to elicit step-by-step reasoning, and (3) Group Relative Policy Optimization (GRPO), a reinforcement learning method that aligns model outputs to expert-derived reasoning trajectories. Experiments with LLaMa3-8B and Med42-8B backbones demonstrate that CoT prompting improves F1 by +6.0 and reduces MAE by 12%, while GRPO achieves state-of-the-art interpretability and predictive performance across BLEU, ROUGE, and BERTScore. We further show that existing biomedical LLMs often fail to produce valid reasoning traces due to architectural constraints. Our findings underscore the importance of reasoning-aware alignment in multi-task clinical modeling and set a new benchmark for interpretable, trustworthy LLMs in precision oncology.
Paper Structure (44 sections, 7 equations, 7 figures, 9 tables)

This paper contains 44 sections, 7 equations, 7 figures, 9 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Dataset Description
  4. Data Overview and Preparation
  5. Attribute Selection and Summary Construction
  6. Chain-of-Thought Reasoning Trace Generation
  7. Chain-of-Thought Evaluation via Embedding-Based Metrics
  8. Method
  9. Autoregressive Fine-Tuning for Treatment Outcome Prediction
  10. Supervised Fine-Tuning with Chain-of-Thought Distillation
  11. Cold-Start SFT with GRPO for Joint Reasoning and Prediction
  12. Cold-Start Supervised Fine-Tuning with CoT
  13. Generalized Reward Policy Optimization (GRPO)
  14. Reinforcement Learning via Generalized Reward Policy Optimization (GRPO)
  15. Multi-Faceted Reward Function
  16. ...and 29 more sections

Figures (7)

  • Figure 1: From Black-Box Models to Reasoning-Augmented LLMs for Cancer Outcome Prediction. Existing LLMs excel at general biomedical tasks but lack structured clinical reasoning and explainability. Our proposed framework augments LLMs with Chain-of-Thought (CoT) prompting and GRPO-based alignment to perform interpretable, multi-task survival prediction on MSK-CHORD.
  • Figure 2: This diagram illustrates the methodology applied to train survival predictor on the MSK-CHORD dataset, using GRPO.
  • Figure 3: Absolute error distribution (in months) for all models with and without CoT prompting. CoT-enhanced models show reduced median and variance.
  • Figure 4: Predicted vs. actual survival months for Ours (LLaMa3), Ours (Med42), and Med42 baselines. The dashed line indicates perfect prediction.
  • Figure 5: Generation quality under CoT prompting, evaluated via BLEU, ROUGE, and BERTScore.
  • ...and 2 more figures