🤖 machine learning

SurvPFN: Towards Foundation Models for Survival Predictions

The paper introduces SurvPFN, a prior-data fitted network pretrained on millions of synthetic tasks to enable foundation models to handle censored data for survival prediction, achieving competitive performance on real-world datasets without requiring per-dataset fitting or specialized architectures.

Original authors: Samuel Böhm (Institute of Epidemiology and Prevention, Medical Center - University of Freiburg, Faculty of Medicine, University of Freiburg, Freiburg, Germany), Lennart Purucker (Department of Compute

Published 2026-06-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Samuel Böhm (Institute of Epidemiology and Prevention, Medical Center - University of Freiburg, Faculty of Medicine, University of Freiburg, Freiburg, Germany), Lennart Purucker (Department of Computer Science, University of Freiburg, Freiburg, Germany, PriorLabs, Freiburg, Germany), Frank Hutter (Department of Computer Science, University of Freiburg, Freiburg, Germany, PriorLabs, Freiburg, Germany), Pascal Schlosser (Institute of Epidemiology and Prevention, Medical Center - University of Freiburg, Faculty of Medicine, University of Freiburg, Freiburg, Germany, Department of Epidemiology, Johns Hopkins Bloomberg School of Public Health, Baltimore, Maryland, US, CIBSS - Centre for Integrative Biological Signalling Studies, University of Freiburg, Freiburg, Germany)

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

Imagine you are trying to predict exactly when a lightbulb will burn out. In a perfect world, you could wait until every single bulb fails and record the exact time. But in the real world, you can't wait forever. Some bulbs are still working when you stop the experiment, and some get thrown away before they burn out. In statistics, this is called censoring. You know the bulb lasted at least until you stopped watching, but you don't know the exact moment it would have died.

For a long time, powerful new AI tools called Foundation Models (specifically "TabPFN") have been great at predicting things like "Will this email be spam?" or "How much will this house sell for?" But they struggled with survival predictions (like "When will this patient relapse?") because they didn't know how to handle those "still working" lightbulbs. If you fed them censored data, they would get confused and make biased guesses.

Enter SurvPFN. The authors built a new version of this AI specifically designed to understand the "lightbulb problem."

Here is how they did it, using some everyday analogies:

1. The "Training Gym" (Synthetic Data)

You can't train a fighter by only letting them fight in real matches; they need a gym to practice. The authors built a massive synthetic gym using millions of fake survival scenarios.

The Setup: They used a mathematical recipe (called a Structural Causal Model) to generate fake "lightbulbs" with different lifespans.
The Twist: They intentionally "censored" some of these fake bulbs—stopping the timer early for some, just like in real life.
The Lesson: They taught the AI that a "stopped" bulb isn't a broken one; it's just a bulb that survived at least that long. They used a special scoring system (a "censored loss") that rewards the AI for being honest about these lower bounds rather than guessing a specific death time.

2. The "Two-Track" Brain (The Model Architecture)

The AI they built is based on a smaller, efficient version of existing models called NanoTabPFN. To make it work for survival, they added a special feature:

The Event Indicator: Think of this as a little flag the AI looks at for every single data point.
- Flag = 0: "This bulb burned out at time X." (We know the exact time).
- Flag = 1: "This bulb was still working at time X." (We only know it survived past X).
The AI has two separate "paths" in its brain to process these two types of flags. This allows it to learn that a "survived" bulb is different from a "died" bulb, even if they happened at the same time.

3. The "One-Size-Fits-All" Approach

Usually, to predict survival, a doctor or data scientist has to spend weeks tuning a specific model for a specific dataset (like tuning a radio for one specific station).

SurvPFN's Superpower: It was trained once on that massive synthetic gym. After that, it can walk into a new, real-world dataset (like a hospital record or a machine reliability log) and start making predictions immediately. It doesn't need to be retrained or "fine-tuned" for that specific group of people. It's like a universal translator that speaks every language without needing a dictionary for each new country.

4. The Results: How Good Is It?

The authors tested SurvPFN on 22 real-world datasets (ranging from cancer studies to unemployment data) and compared it against the "gold standard" methods:

Cox Proportional Hazards: The classic, old-school statistical method.
Random Survival Forests: A popular machine learning method.
DeepSurv: A modern deep learning approach.
BinSurv: A newer method that tries to turn survival into a series of "yes/no" questions.

The Verdict:
SurvPFN performed just as well as all these specialized, heavily-tuned experts.

It didn't need to be retrained for each new dataset.
It didn't need special "feature engineering" (manually tweaking the data).
It didn't even need a complex, custom-built architecture.

The Catch (Limitations)

The paper is honest about where this AI currently has limits:

Size: It works best on smaller datasets (up to 1,000 rows and 10 features). It's like a sports car: great for quick, agile races, but not yet built for a massive trucking fleet.
The "Weibull" Bias: The AI was trained on a specific mathematical shape of survival curves (called Weibull). While it handles many real-world scenarios well, it might struggle if the real-world data looks nothing like the shapes it saw in the gym.
Static View: It assumes the factors influencing survival (like age or job type) don't change over time. It doesn't handle "competing risks" (e.g., a patient dying from a car accident before their disease progresses).

The Big Takeaway

The paper proves that survival analysis is just a type of distributional regression with a special rule for "lower bounds." By teaching a general-purpose AI to respect the "censoring" rule (the event indicator flag), they unlocked the ability to use these powerful, pre-trained foundation models for time-to-event predictions without needing to build a new model from scratch every time.

Technical Summary: SurvPFN: Towards Foundation Models for Survival Predictions

Problem Statement

Tabular Foundation Models (TFMs), such as TabPFN, have achieved rapid success in standard classification and regression tasks by leveraging prior-data fitted networks (PFNs) trained on synthetic data. However, survival prediction tasks have remained largely untouched by this paradigm. The primary barrier is censoring: in survival analysis, the event of interest (e.g., death, failure) is often not observed for all subjects within the study period (right-censoring). Standard TFM architectures and regression objectives cannot natively handle censored data; treating censored times as regression targets or discarding them introduces systematic bias and inaccurate predictions. Existing survival methods (e.g., Cox Proportional Hazards, Random Survival Forests, DeepSurv) explicitly model censoring in their likelihood functions, but they typically require per-dataset fitting and lack the "foundation model" capability of zero-shot adaptation.

Methodology

The authors propose SurvPFN, a prior-data fitted network specifically designed for survival prediction. The approach treats survival as a continuous-time distributional regression problem with a censored loss function.

1. Synthetic Prior Generation

To train the model without per-dataset fitting, the authors generate millions of synthetic survival tasks using Structural Causal Models (SCMs):

Event Times: Generated from a Weibull distribution, a standard choice in survival analysis. The model supports both proportional hazards (PH) and non-proportional hazards (non-PH) regimes. In non-PH scenarios, a second outcome node modifies the shape parameter per subject, inducing crossing survival curves.
Censoring Mechanism: A non-informative censoring mechanism is applied. This includes:
- Administrative censoring: A dataset-level follow-up horizon ($FU$) set at a quantile of the event-time distribution.
- Staggered enrollment: Variable start times for observations.
- Mid-study dropout: Observations may end at a uniformly drawn time before their administrative censoring time.
Data Structure: Synthetic datasets contain 2–10 covariates (including categorical features) and up to 1,000 rows.

2. Model Architecture and Training

SurvPFN is built upon NanoTabPFN, a lightweight, user-friendly TFM backbone (4.4M parameters, 6 layers).

Input Representation: The event indicator (0 for censored, 1 for observed event) is appended as an explicit feature column. This allows the model to distinguish between observed event times and lower bounds (censored times).
Target Encoding: A two-routed target encoder processes observed events and censored rows through separate encoders of identical architecture, enabling the model to learn distinct representations for true event times and lower bounds.
Training Objective: The model is trained to minimize a censored negative log-likelihood (NLL):
- For observed events: Weighted by the inverse probability of censoring (IPCW) to correct for informative censoring.
- For censored rows: Penalizes the model for placing probability mass before the censoring time.
- Ranking Term: A differentiable pairwise ranking term is optionally added to align predicted survival times with the underlying risk order induced by the SCM.
Inference: The model outputs a discrete density $\hat{p}_j$ , from which the survival function $\hat{S}(t)$ and predicted median time are derived directly.

3. Evaluation Protocol

The model was evaluated on SurvSet, a curated collection of 22 real-world survival datasets spanning clinical, economic, and reliability domains.

Constraints: Datasets were filtered to $\le 1000$ rows and $\le 10$ features to match the training regime of the small backbone.
Baselines: Comparisons were made against:
- Cox Proportional Hazards (CoxPH): The dominant classical baseline.
- Random Survival Forests (RSF): Standard non-parametric ensemble.
- DeepSurv: The most widely used deep learning approach.
- BinSurv: A recent TFM-based approach that discretizes time into binary classification tasks.
Metrics: Performance was assessed via the Uno's C-index (discrimination), Integrated Brier Score (IBS) (accuracy), and Integrated Calibration Index (ICI) (calibration).

Key Results

Competitive Performance: SurvPFN is highly competitive with all baselines across the 22 datasets. On average, it matches or slightly exceeds RSF, DeepSurv, and BinSurv, and performs comparably to CoxPH.
Zero-Shot Capability: Crucially, SurvPFN achieves this performance without per-dataset fitting, without a survival-specific architecture, and without feature engineering. It relies solely on the pre-trained prior.
Ablation Studies:
- Event Indicator: Removing the event indicator as a feature (forcing it to 1) preserves discrimination (C-index) but significantly degrades calibration (ICI). This confirms that while relative risk ordering can be inferred from time alone, absolute calibration requires explicit knowledge of which data points are censored.
- Loss Variants: Variants using oracle targets (available only during synthetic pretraining) performed comparably to the native censored loss on real-world data, suggesting the native loss is sufficient for transfer.

Significance and Claims

The paper claims that survival analysis can be successfully framed as a continuous-time distributional regression problem solvable by PFNs. The key contributions are:

Bridging the Gap: Demonstrating that foundation models can handle censored data, a task previously considered unsuitable for standard regression PFNs.
Simplicity: Showing that a standard regression PFN architecture, augmented only with an event indicator feature and a censored loss function, is sufficient to recover the performance of specialized survival models.
Efficiency: Providing a model that requires no hyperparameter tuning or retraining on new datasets, offering a "plug-and-play" solution for small-to-medium tabular survival tasks.

The authors remain modest regarding limitations, noting that the current evaluation is restricted to small datasets ( $\le 1000$ rows) and static covariates. They acknowledge that the synthetic prior is limited to Weibull distributions and does not yet address competing risks or time-varying covariates, suggesting these as necessary directions for future scaling and expressiveness.