TRIAGE: Teaching LLMs Dialectical Reasoning for Calibrated Clinical Risk Prediction

Forum topic · 小凯 · 2026-06-20

Summary

TRIAGE is a framework from researchers at KAIST, AITRICS, and the University of Wisconsin-Madison that addresses a key flaw in LLM-based medical risk prediction: when a model reasons before answering, it pre-commits to an outcome and its predicted probabilities polarize toward 0 or 1. Experiments with gpt-oss-120b on MIMIC-III mortality prediction show predicted class probabilities exceeding 99.98% with near-zero variance, destroying clinically necessary risk gradients. TRIAGE solves this by asking the model to generate separate, independent reasoning chains supporting each candidate outcome—a dialectical, debate-style structure—then extracts continuous, calibrated risk scores from the model's implicit token probabilities at the answer position. A two-stage training pipeline (supervised fine-tuning on GPT-5.1/Kimi K2 Thinking-generated rationales, followed by GRPO reinforcement learning with a batch-level ranking reward) trains a Qwen3-4B-Base model. On P12, P19, and MIMIC-III benchmarks, TRIAGE achieves an average AUPRC gain of 3.3% over the strongest baselines, reduces expected calibration error by roughly 81%, and produces explanations rated as more clinically sound than post-hoc baselines.

TRIAGE: Dialectical Reasoning for Explainable Risk Prediction on Irregularly Sampled Medical Time Series with LLMs — by Hyeongwon Jang, Gyouk Chu, Changhun Kim, Joonhyung Park, Hangyul Yoon, and Eunho Yang (KAIST, AITRICS, University of Wisconsin-Madison). arXiv: 2606.09030. Code: https://github.com/HyeongWon-Jang/TRIAGE

Key points

Problem — Risk Polarization. When an LLM generates reasoning before its diagnosis, 71.7% of chains explicitly pre-commit to a conclusion (e.g., "Therefore, this patient is likely to die"), and the reasoning shows one-sided confirmation bias. On MIMIC-III mortality prediction, gpt-oss-120b's predicted class probabilities exceed 99.98% with near-zero variance (vs. 86.4%±18.8% without reasoning), eliminating the risk gradient needed for triage.
Method — Dialectical reasoning. Irregularly sampled medical time series (ISMTS) are serialized into text via set-based encoding. For each candidate outcome, the model produces an independent, self-contained supporting rationale (empty if no supporting evidence exists), with no intermediate conclusions before the final decision token. This removes pre-commitment so the model's implicit distribution reflects genuine belief.
Risk estimation. TRIAGE ignores verbalized probabilities and instead normalizes implicit token probabilities at the answer position: P(death) = softmax(logit("1")) / (softmax(logit("0")) + softmax(logit("1"))), yielding continuous, calibrated, cross-patient-comparable scores.
Two-stage training. Stage 1: supervised fine-tuning of Qwen3-4B-Base on dialectical rationales generated by GPT-5.1 or Kimi K2 Thinking. Stage 2: GRPO reinforcement learning with loss L = L_GRPO + λ·L_CE and a batch-level contrastive ranking reward (hinge-style surrogate) that pushes positive samples' log-odds above the batch's negative average and vice versa, encouraging cross-patient risk ordering.

Results

Datasets: P12 (in-hospital mortality, ~12k samples), P19 (sepsis within 6h, ~40k), MIMIC-III (in-hospital mortality, ~36k).
Discrimination: TRIAGE_SFT+RL achieves the best average rank (1.58 across all six AUROC/AUPRC metrics), including P12 AUPRC 59.0, P19 AUROC 89.3, MIMIC-III AUPRC 54.1 — a 3.3% relative AUPRC improvement over the strongest baseline (GRU-D).
Zero-shot frontier LLMs fail: GPT-5.1 and gpt-oss-120b rank last (average ranks 10.50 and 11.67), showing general reasoning does not transfer directly to clinical ISMTS prediction.
Calibration: ECE drops to 0.03–0.04 (vs. 0.16–0.21 for GRU-D/STraTS), an ~80% reduction; Brier Score drops ~49%.
Explanation quality: Using the IDEA rubric (LLM-as-judge, 200 samples), TRIAGE scores 7.744 vs. 6.474 for a STraTS+IG+GPT post-hoc pipeline, with the largest gains in interpretive summary (+0.902) and alternative-diagnosis explanation (+0.288).
Ablations: One-sided rationales are worse than answer-only prediction (AUPRC 43.1 vs. 53.4 on P12), directly demonstrating risk polarization; dialectical reasoning reaches 86.9 AUROC / 56.4 AUPRC. The batch-level reward outperforms sample-level rewards on AUPRC, ECE, and Brier Score.
Robustness: With 10–50% of variables randomly hidden, TRIAGE remains competitive, leading on MIMIC-III AUPRC at all missing rates.

Limitations

Validated only on binary outcomes; multi-class/multi-label extension is future work.
Rationale generation adds latency, which may not suit strict low-latency settings.
Explanation quality was assessed via LLM-as-judge rather than clinical experts.
Generated reasoning may contain errors and must not replace qualified clinical judgment.

Ethics

MIMIC-III data was used under PhysioNet protocols and never sent to third-party LLM APIs; reasoning-generation models were run locally with open weights (Kimi K2 Thinking).

Takeaway

TRIAGE's core insight is that an LLM's apparent overconfidence in diagnosis is often an artifact of reasoning structure, not genuine belief. By requiring the model to argue both sides before deciding, it restores calibrated, comparable risk scores — a shift from "decide first, find evidence after" to evidence-first deliberation, potentially applicable to any risk-assessment domain such as credit approval, insurance pricing, or threat detection.