SurvPFN: Towards Foundation Models for Survival Predictions

2026-06-03Machine Learning

Machine Learning
AI summary

The authors focus on improving survival prediction models, which estimate the time until an event happens, like a patient relapse, while handling incomplete data called censoring. They introduce SurvPFN, a new model trained on many fake datasets designed to mimic real survival data with censoring. SurvPFN learns to predict survival times using a special loss function that accounts for censored cases, making it effective without extra tuning for each dataset. Their results show it competes well with existing methods on real-world tasks. This work shows survival prediction can be treated as a regression problem that considers censoring properly.

Tabular foundation modelsSurvival predictionCensored dataPrior-data fitted networksDistributional regressionWeibull distributionNegative log-likelihoodTime-to-event analysisSurvival analysisSynthetic data
Authors
Samuel Böhm, Lennart Purucker, Frank Hutter, Pascal Schlosser
Abstract
Tabular foundation models (TFMs) have made rapid progress in standard classification and regression, but time-to-event survival prediction tasks have remained largely untouched. Unlike in standard regression tasks, survival prediction models must account for censored data. Standard TFMs cannot handle natively censored data, leading to biased and inaccurate predictions, making them unsuitable for real-world applications. To overcome this fundamental limitation, we propose \texttt{SurvPFN}, a prior-data fitted network (PFN), for survival prediction tasks. We pretrain \texttt{SurvPFN} on millions of synthetic survival prediction tasks to learn survival via distributional regression that accounts for censored data. \texttt{SurvPFN} works by (1) generating data with Weibull event times and a non-informative censoring mechanism; (2) integrating a censored event indicator; and (3) minimizing a censored negative log-likelihood. On SurvSet, a collection of real-world survival tasks, \texttt{SurvPFN} is highly competitive with classical and deep survival baselines without per-dataset fitting, a survival-specific architecture, or feature engineering. We show that survival can be treated as a continuous-time distributional regression problem with censored loss, unlocking the power of PFNs for time-to-event predictions.