TRIAGE: Dialectical Reasoning for Explainable Risk Prediction on Irregularly Sampled Medical Time Series with LLMs

Solving the Overconfidence Problem in Medical AI

When AI explains medical risks, it often becomes overconfident and ignores evidence that contradicts its first guess. This tendency can turn nuanced clinical assessments into rigid, binary predictions. These predictions fail to capture the true spectrum of patient health. A new study introduces TRIAGE, a framework that forces Large Language Models (LLMs) to write separate explanations for both possible outcomes—such as survival versus death—before making a final prediction. By doing so, the authors aim to create risk scores that are more accurate and more balanced. They hope these scores will be easier for clinicians to trust.

The Risk Polarization Problem

In intensive care units, doctors rely on early-warning systems to monitor patients via Electronic Health Records (EHRs). These records consist of irregularly sampled medical time series (ISMTS). These are sequences of vital signs and lab tests recorded at unpredictable intervals. Predicting adverse events from this data is difficult. The signals are often sparse and noisy.

Current approaches to this problem generally fall into two camps. Specialized deep learning models can predict risks with high accuracy. However, they are "black boxes" that offer no linguistic explanation for their decisions. Conversely, when researchers use LLMs to provide natural language rationales, they encounter "risk polarization."

The authors report that when an LLM is asked to "reason then predict," it tends to commit to a single outcome prematurely. The reasoning process acts as a self-imposed "hard prior." This is much like a juror who decides a verdict halfway through a trial. That juror then only listens for evidence that supports it. This creates two problems. First, the rationales become one-sided. They ignore contradictory clinical signals. Second, the model's predicted probability collapses toward 0 or 1. This means the model loses the ability to distinguish between a patient with a 60% risk and one with a 95% risk.

The Mechanics of Dialectical Reasoning

To break this cycle of overconfidence, the authors propose TRIAGE. This framework is centered on "dialectical reasoning." Instead of a linear path from observation to verdict, TRIAGE forces the model to deliberate over competing possibilities. As shown in, the process decomposes reasoning into outcome-conditioned rationales.

For any given patient, the model must first generate a rationale assuming the patient will survive. It must then generate a separate rationale assuming the patient will experience mortality.

This structure ensures the model weighs evidence for both sides of the clinical trajectory. In a medical context, a patient might show signs of both physiological stabilization and deterioration. Standard reasoning often picks one side and discards the other. TRIAGE surfaces both.

The authors implement this through a two-stage training pipeline. In Stage 1, they use supervised fine-tuning (SFT) to teach a small LLM how to produce these dual rationales. In Stage 2, they apply reinforcement learning (RL) using Group Relative Policy Optimization (GRPO). Crucially, the authors utilize a "batch-level reward" during this phase. Rather than rewarding the model simply for being right about one patient, the reward compares a patient's risk score against the average scores of other patients in the same batch. This encourages the model to produce risk scores that are comparable across different individuals. This prevents the "collapse" seen in earlier models.

Finally, the risk score itself is not taken from the model's spoken text. Instead, the authors extract the "implicit probability" from the model's internal logits. Logits are the raw numerical values the model assigns to tokens (the basic units of text). This avoids the trap of "verbalized" probabilities. In those cases, a model might say "70%" but actually be mathematically certain of its choice.

Improved Accuracy and Calibration

The results of the study suggest that this deliberative approach significantly improves the utility of LLM-based predictions. Evaluating the framework on three benchmarks (P12, P19, and MIMIC-III), the authors report an average AUPRC (Area Under the Precision-Recall Curve) improvement of 3.3% over competitive baselines. AUPRC is a critical metric here. Medical datasets are often heavily imbalanced, meaning there are far more survivors than deaths. In such cases, standard accuracy metrics can be misleading.

Perhaps more importantly, TRIAGE dramatically improves "calibration." A well-calibrated model is one where a predicted 20% risk actually corresponds to a 20% occurrence rate in reality. The authors report that TRIAGE reduces calibration error by 81% compared to previous methods. This massive reduction means the model's confidence levels are much closer to the actual frequency of events.

The qualitative impact is also notable. Using an "LLM-as-a-judge" assessment, the authors found that TRIAGE's rationales surpassed the quality of post-hoc explanations from baseline models by 20%. In case studies, the authors demonstrate that while traditional methods might incorrectly attribute a high heart rate to a reason for survival, TRIAGE correctly identifies such signals as markers of deterioration.

Limits of the Framework

Despite these gains, the authors highlight several boundaries to their work. TRIAGE is currently optimized for binary classification tasks. This means it predicts only one of two outcomes. The authors note that extending this to multi-class scenarios remains an open challenge. Examples include predicting several different types of organ failure.

There is also a practical trade-off regarding computational efficiency. TRIAGE requires the model to generate multiple lengthy reasoning chains. It does this before arriving at a single number. Therefore, it is more computationally expensive than lightweight models like GRU-D. It also has higher latency (the delay before a response is produced). This makes it potentially difficult to deploy in environments that require instantaneous inference.

Finally, the researchers acknowledge that the reasoning quality was assessed using other LLMs acting as judges. While they used multiple models to mitigate bias, they state that formal evaluation by clinical experts is a necessary next step. This is essential before such a system could move toward real-world clinical use.

Figures from the paper