TRIAGE-4B-P12

SFT + RL checkpoints of TRIAGE for in-hospital mortality (P12), one per cross-validation split. Part of the TRIAGE collection.

• Paper: TRIAGE: Dialectical LLM Reasoning for Explainable Risk Prediction on Irregularly Sampled Medical Time Series
• Code, data and training pipeline: github.com/HyeongWon-Jang/TRIAGE
• Project page: hyeongwon-jang.github.io/triage-project-page

TRIAGE trains an LLM to reason dialectically over the competing clinical outcomes. Given a patient's irregularly sampled medical time series, the model writes a rationale for survival and a rationale for in-hospital death before it commits to a final decision, and the risk score is read from the model's probability over that decision. This mitigates risk polarization, the collapse of graded risk into overconfident predictions under conventional LLM reasoning, so that one model gives both a clinician-readable rationale and a risk score comparable across patients.

The SFT-only checkpoints these runs were post-trained from are released separately as TRIAGE-4B-P12-SFT.

Model description

Base model	Qwen/Qwen3-4B-Base
Task	In-hospital mortality prediction on P12 (PhysioNet Challenge 2012)
Training	supervised fine-tuning on dialectical-reasoning data (triage_sft/), then GRPO post-training with a batch separation reward (triage_rl/, built on verl)
Layout	split_1 … split_5, one checkpoint per split, each in its own subfolder
Selection	per split, the RL checkpoint with the best validation AUPRC
SFT warm start	the rl_init branch holds the SFT checkpoint each RL run started from

Input and output format

The model expects the chat-style prompt the training data uses: the patient's static features and per-variable (time, value) series, the question, and the reasoning order (which outcome to argue for first). The exact prompt and how it is built from the raw data are in triage_sft/data/P12/.

The answer has three headed sections, in the order the prompt asks for:

## Rationale for survival
...

## Rationale for in-hospital death
...

## Final Decision
1

0 means survival and 1 means in-hospital death. The risk score is the next-token probability of "1" at the decision position, renormalised over the two tokens "0" and "1", averaged over the two rationale orderings.

Quick start

from transformers import AutoModelForCausalLM, AutoTokenizer

split = "split_1"  # one of split_1 ... split_5
repo  = "Hyeongwon/TRIAGE-4B-P12"

tokenizer = AutoTokenizer.from_pretrained(repo, subfolder=split)
model     = AutoModelForCausalLM.from_pretrained(repo, subfolder=split, torch_dtype="auto", device_map="auto")

Scoring is two-step: sample the dialectical reasoning at temperature 0.7 up to the ## Final Decision header, then read the logits of "0" and "1" at the next position. A complete predict_risk function is in the GitHub README.

Data

• Raw: PhysioNet Challenge 2012, in-hospital mortality (Open Data Commons Attribution License v1.0).
• Processed: Raindrop-preprocessed release on figshare.
• Splits: the five cross-validation index files are in the GitHub repository (triage_sft/data/P12/splits/, phy12_split{1..5}.npy); each split's checkpoint here was trained on the training portion of the file with the same number.

No data is included in this repository.

License

Model checkpoints are released under CC BY-NC 4.0 (non-commercial). The base model and the datasets remain under their own licences.

Citation

@article{jang2026triage,
  title={TRIAGE: Dialectical LLM Reasoning for Explainable Risk Prediction on Irregularly Sampled Medical Time Series},
  author={Jang, Hyeongwon and Chu, Gyouk and Kim, Changhun and Yoon, Hangyul and Lee, Jeonguk and Yang, Eunho and Park, Joonhyung},
  journal={arXiv preprint arXiv:2606.09030},
  year={2026}
}