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

Abstract: We investigate whether information about time-to-event risk estimated by a Cox proportional hazards model can be transferred into a generative LLM. We propose a text-based survival modelling pipeline in which structured clinical covariates are converted into text prompts and a Qwen-based LLM is fine-tuned to generate patient-specific survival risk using Cox model predictions as a training target. Across GBSG2, ACTG320, and WHAS500, the model achieves competitive held-out discrimination and calibration despite being trained as a text-generation task rather than with a conventional survival-analysis loss. We further analyse the geometry of the model's hidden states, where t-SNE visualisations reveal smooth risk gradients in latent space, suggesting that the model represents survival risk as a continuous structure rather than isolated risk categories. Together, these findings suggest that LLMs can internalise survival-risk structure while supporting calibrated prediction, providing a route towards time-to-event reasoning in LLMs.

• The paper introduces a novel framework that distills survival risk from Cox models into a large language model via prompt-based supervised fine-tuning.
• It demonstrates that LLM predictions achieve comparable discrimination (ΔC-index < 0.01) and calibration across multiple clinical datasets.
• The approach leverages LoRA fine-tuning to embed continuous risk representations, enabling flexible clinical risk assessment without bespoke survival heads.

Cox-Supervised Distillation of Survival Risk into LLMs

Introduction

The paper "From Hazard Functions to Language Space: Cox-Supervised Distillation of Survival Risk into a LLM" (2606.08945) presents a framework for transferring clinically calibrated, time-to-event risk information from a Cox proportional hazards model to a generative LLM via prompt-based supervised fine-tuning. The pipeline reformulates survival analysis as a text-generation task by encoding structured clinical covariates as natural-language prompts, using the Cox model to generate patient-level survival risk targets, and fine-tuning a Qwen-based LLM to autoregressively output Cox-consistent survival risk estimates. This research rigorously evaluates model discrimination, calibration, and latent-state risk geometry across multiple clinical datasets.

Methodology

The pipeline entails several critical stages:

1. Cox Model Fitting and Risk Target Generation: For each dataset (WHAS500, GBSG2, ACTG320), a Cox proportional hazards model is fitted on the training partition, producing patient-specific survival risk estimates (e.g., 1-year risk $r_i = 1 - S_i(\tau)$) serving as supervision targets.
2. Prompt Engineering: Tabular clinical covariates are rendered as canonicalised text prompts including a fixed task instruction (standardising prediction as 1-year risk between 0 and 1), followed by the clinical fields and a "Risk:" cue.
3. LoRA-based Fine-tuning: A Qwen2.5-1.5B-Instruct LLM, quantised to 4-bit, is fine-tuned using Low-Rank Adaptation (LoRA), with all supervision presented as prompt--completion pairs where the "completion" is a tokenised string representation (e.g., "0.2374") of the Cox-derived risk estimate. The fine-tuning occurs with completion-only loss, keeping base model weights fixed except for injected low-rank adapters.
4. Inference and Risk Recovery: Autoregressive decoding on held-out samples; predicted risk values are parsed from generated text using regular-expression extraction on the output.
5. Evaluation: Predictive discrimination (Harrell’s C-index), numerical calibration error, and percentile-based NRI are evaluated against ground-truth time-to-event data and the Cox teacher model.

The overall workflow is illustrated in (Figure 1).

Figure 1: Text-based survival risk prediction via Cox-supervised LLM training.

Experimental Results

Evaluation spans three datasets representing cardiovascular (WHAS500), oncology (GBSG2), and HIV (ACTG320) populations, each with differing covariate complexity and event distributions.

• Discrimination: Across all datasets, the LLM achieves C-index performance comparable to the Cox teacher ($\Delta$C-index $< 0.01$ on average), confirming that autoregressive SFT is sufficient for the LLM to learn time-to-event risk ranking to classical survival model standards.
• Calibration: On WHAS500, cohort-level calibration error (absolute deviation from perfect calibration slope) for the Qwen model is $0.0711\pm0.0238$, marginally higher than the Cox ($0.0583\pm0.0368$). Calibration bias is moderate but does not suggest widespread model miscalibration.
• Subgroup Calibration: Subgroup D21 calibration error is generally consistent between Cox and LLM, with the LLM occasionally achieving lower error (e.g., AF, MI type) and at other times higher (CVD, MI order).
• Reclassification: Percentile-based NRI is negative, indicating some loss of case discrimination in the LLM compared to the Cox, with greater degradation for event cases than controls. The LLM fails to improve upon the teacher in risk stratification.
• Latent Representation Geometry: t-SNE projection of final-layer hidden states reveals a smooth, monotonic gradient corresponding to model-predicted 1-year survival risk rather than clustering into discrete bins, implying the LLM encodes survival risk as a dense, continuous latent structure.

  Figure 2: t-SNE projection of Qwen hidden states on the WHAS500 test set, illustrating continuous risk-gradient structure.

Notable Claims and Strong Results

• The LLM learns competitive survival risk discrimination and calibration via Cox-based language supervision despite lacking a task-specific survival objective.
• The model’s latent space supports smooth, continuous risk encoding, as opposed to discrete class partitioning, demonstrating effective internalisation of survival-risk geometry from the Cox output.
• No evidence is found for meaningful NRI improvement over the Cox model; performance is consistently competitive but does not surpass the teacher.
• Substantial parameter efficiency is achieved; only $\sim$1.2% of model parameters are updated during LoRA fine-tuning.

Theoretical and Practical Implications

This work establishes that LLMs, when supervised with explicitly calibrated Cox model outputs via prompt-based SFT, can internalise survival risk structure at clinically competitive levels in the absence of domain-specific risk objectives. The approach circumvents requirements for bespoke survival heads or regression-specific tuning, leveraging natural-language input to unify disparate tabular schemas. This could enable greater flexibility in clinical deployment—especially where feature types, documentation conventions, or data pipeline idiosyncrasies impede the standardisation of tabular ML frameworks.

However, the Cox teacher necessarily bounds the predictive information available for distillation; further, the computational overhead for LLM fine-tuning far exceeds Cox model training, and the absence of explicit censored-data loss could limit extrapolation in non-standard survival settings. There is also no evidence that this pipeline can reliably outperform established classical models in terms of calibration or discrimination. Generalisation to broader clinical event types, heterogeneous prompt structures, or missingness patterns without re-tuning remains unverified.

Future Directions

Open questions include:

• Robustness of the LLM to prompt schema drift, missing covariates, or out-of-distribution feature values.
• Extension to endogenous self-supervised risk pretraining, moving beyond teacher-student distillation.
• Mechanistic analysis of LLM interpretability in survival-risk stratification along clinical axes.
• Ultra-scale experiments with larger LLMs or in-context learning for time-to-event tasks without explicit supervision.
• Deployment studies assessing language-based survival models under real-world data harmonisation and regulatory considerations.

Conclusion

The Cox-supervised LLM risk-distillation pipeline demonstrates that generative LLMs can act as clinically meaningful risk estimators when supervised by calibrated survival-model output, achieving strong discrimination and calibration without specialised survival losses or model heads. The capacity for smooth continuous risk encoding in language-model latent space is empirically established. While the approach does not outperform classical Cox models, it provides a flexible substrate for future developments in LLM-based clinical prediction in the context of heterogeneous, partially structured datasets.