🤖 machine learning

From Hazard Functions to Language Space: Cox-Supervised Distillation of Survival Risk into a Large Language Model

This paper demonstrates that a Qwen-based large language model, when fine-tuned on text-converted clinical data using Cox model predictions as targets, can effectively internalize survival risk structures and achieve competitive discrimination and calibration for time-to-event reasoning.

Original authors: Nicholas I-Hsien Kuo, Blanca Gallego, Louisa Jorm

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Nicholas I-Hsien Kuo, Blanca Gallego, Louisa Jorm

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Idea: Teaching a Storyteller to Predict the Future

Imagine you have two experts trying to predict how long a patient might live after a serious illness.

The Old Expert (The Cox Model): This is a traditional, highly respected statistician. They look at a patient's data (age, blood pressure, history) and use strict mathematical formulas to calculate a risk score. They are very good at this, but they only speak "math." They can't explain why they think that, and they can't chat with you.
The New Student (The Large Language Model): This is a powerful AI trained to read and write stories, like a super-smart librarian who has read almost everything on the internet. It understands natural language but doesn't inherently know how to do complex medical math.

The Experiment: The researchers asked a simple question: Can we teach the "Storyteller" to act like the "Math Expert" just by having the Math Expert write down the answers, and the Storyteller trying to copy them?

How They Did It: The "Teacher-Student" Game

Instead of forcing the AI to learn complex math equations directly (which is hard for a language model), they used a clever trick called Knowledge Distillation. Think of it like a master chef teaching a sous-chef.

The Teacher (Cox Model): The researchers first let the traditional Math Expert look at thousands of patient records and calculate the exact "1-year survival risk" for each person.
The Translation: They took the patient's raw data (numbers like "Age: 69," "Heart Rate: 80") and turned them into a simple sentence, like a doctor's note: "Patient is 69 years old, has a heart rate of 80, and a history of heart failure."
The Lesson: They showed the AI this sentence and said, "Here is the patient. The answer is 0.2374."
The Practice: The AI (the student) practiced reading these notes and writing down the number. It didn't learn the math formula; it learned the pattern of how the story of a patient's life connects to a specific number.

They tested this on three different groups of patients: those with breast cancer, those with HIV, and those who had heart attacks.

What They Found: The Student Got It Right

The results were surprisingly good. Even though the AI was trained just to "write the next word" (like completing a sentence) rather than "solve a math problem," it learned to predict survival risks almost as well as the traditional Math Expert.

Accuracy: When they checked how well the AI ranked patients (e.g., "Is this patient at higher risk than that one?"), the AI was nearly as accurate as the expert.
Calibration: The AI's numbers were also realistic. If it said a patient had a 20% risk, about 20% of similar patients actually had an event.

The Hidden Magic: The "Risk Landscape"

The researchers also looked inside the AI's "brain" (its hidden states) to see how it understood the data.

They used a visualization tool to map out the AI's thoughts. Instead of seeing the AI treat patients as separate, isolated boxes (e.g., "Low Risk" vs. "High Risk"), they saw a smooth, continuous landscape.

The Analogy: Imagine a hill. At the bottom are healthy patients with low risk. As you walk up the hill, the risk gets higher and higher. The AI didn't build a staircase with distinct steps; it built a smooth ramp. It understands that risk is a sliding scale, not a switch that flips on or off. This suggests the AI truly "got" the concept of survival risk, not just memorized the numbers.

The Catch: It's Not a Perfect Replacement Yet

The paper is honest about the limitations:

The Ceiling: The AI can't be smarter than the teacher. Since the AI was just copying the Math Expert's answers, it couldn't discover anything the Math Expert missed. It is limited by the quality of the teacher.
Cost: Training this AI is much more expensive and computationally heavy than running the simple Math Expert.
Data Issues: The AI works well when the data is clean, but the researchers noted that in the real world, hospital records are messy and inconsistent. They haven't proven yet if this AI can handle the "noise" of real-world data as well as the traditional methods.

The Bottom Line

This paper proves that a Large Language Model can learn to understand time-to-event risks (like survival) just by reading text descriptions of patients and copying the answers from a traditional model. It shows that these AI models can internalize complex medical patterns and represent them as a smooth, continuous understanding of risk, opening the door for future tools that might reason about patient outcomes using natural language.

Technical Summary: From Hazard Functions to Language Space

Problem Statement

Survival analysis is a fundamental clinical prediction task involving time-to-event outcomes under censoring. While classical methods like the Cox proportional hazards model remain interpretable and robust, modern neural survival models typically rely on explicit predictive heads and specialized survival loss functions (e.g., Cox partial likelihood) rather than language-based supervision. Conversely, recent work in tabular large language models (LLMs) has demonstrated the ability to treat structured data as natural language prompts for classification tasks. However, it remains unclear whether generative LLMs can internalize the continuous, time-to-event risk structures estimated by statistical survival models when trained solely via standard text-generation objectives, without explicit survival-specific loss functions.

Methodology

The authors propose a Cox-supervised distillation framework that transfers survival risk information from a statistical teacher model to a generative LLM student. The pipeline consists of four primary stages:

Cox Teacher Training: A standard Cox proportional hazards model is fitted on structured clinical covariates ( $x_i$ ), follow-up times ( $t_i$ ), and event indicators ( $\delta_i$ ). This model generates patient-specific 1-year survival risk estimates ( $r_i = 1 - S_i(\tau)$ ), which serve as the ground-truth training targets.
Prompt Construction: Structured patient records are converted into natural-language prompts. Each prompt includes a task instruction, a structured list of clinical covariates (e.g., "Age: 69", "BMI: 26.4"), and a cue ("Risk:").
LLM Fine-Tuning (Student): A Qwen2.5-1.5B-Instruct model is fine-tuned using Low-Rank Adaptation (LoRA). The training objective is standard autoregressive supervised fine-tuning (SFT) with a completion-only loss. The model learns to map the textual prompt to a scalar risk value formatted as a four-decimal string (e.g., "0.2374"). Crucially, the model is not optimized with a survival-specific loss (like Eq. 1 in the paper); it is trained purely as a text-generation task where the target is the Cox-derived risk string.
Inference and Evaluation: During inference, the fine-tuned model generates risk strings autoregressively. These strings are parsed via regular expressions to extract numeric predictions. Performance is evaluated using standard survival metrics: Harrell's C-index (discrimination), calibration error ( $D_{21}$ ), and percentile-based Net Reclassification Improvement (NRI).

The framework was evaluated on three public datasets: GBSG2 (breast cancer), ACTG320 (HIV/AIDS), and WHAS500 (acute myocardial infarction).

Key Results

Across all three datasets, the Qwen-based LLM achieved competitive performance compared to the Cox teacher model, despite being trained as a text-generation task:

Discrimination: On the WHAS500 dataset, the LLM achieved a C-index of 0.7563, closely matching the Cox model's 0.7658. Similar results were observed on GBSG2 (0.6801 vs. 0.6780) and ACTG320 (0.7066 vs. 0.7038).
Calibration: The LLM maintained modest calibration, though with slightly higher error than the Cox model on WHAS500 (0.0711 vs. 0.0583). In some subgroups (e.g., Age > 60y in GBSG2), the LLM achieved lower calibration error than the Cox model.
Latent Space Geometry: t-SNE visualizations of the model's hidden states revealed smooth risk gradients in the latent space. Patients with similar predicted risks clustered together, suggesting the model represents survival risk as a continuous structure rather than discrete categories.
Reclassification: Percentile-based NRI analysis indicated that while the LLM generally preserved risk rankings, it did not fully preserve the Cox model's ranking, showing degradation particularly for event cases (NRI cases = -0.2067 on WHAS500).

Significance and Claims

The paper claims that large language models can internalize survival-risk structure through autoregressive supervised fine-tuning alone, without explicit survival objectives. The key contributions are:

Feasibility of Text-Based Survival Modeling: The study demonstrates that converting structured covariates to text and training an LLM to generate risk scores is a viable alternative to traditional neural survival architectures.
Continuous Representation: The observation of smooth risk gradients in the latent space suggests that LLMs can learn clinically meaningful, continuous representations of time-to-event risk, rather than treating risk as isolated classification bins.
Flexibility for Heterogeneous Data: The authors suggest that a text-based framework may offer practical advantages in settings with heterogeneous clinical data or incompatible feature schemas, as it conditions on natural language descriptions rather than rigid numerical feature vectors.

The authors maintain a modest stance, acknowledging that the Cox model defines the "predictive ceiling" for this distillation approach and that the LLM incurs higher computational costs. They note that the framework does not currently outperform the Cox model but offers a novel route toward integrating time-to-event reasoning into language models. Future work is proposed to address robustness to informative missingness and distribution shifts across institutions.