SurvivalPFN: Amortizing Survival Prediction via In-Context Bayesian Inference

TL;DR

SurvivalPFN is a neural network model that efficiently performs survival analysis by amortizing Bayesian inference through in-context learning, adapting to various datasets without extensive tuning.

Contribution

It introduces a pretrained network that generalizes survival analysis across diverse data-generating processes without task-specific training.

Findings

Achieves strong predictive performance across 61 datasets.

Often outperforms established survival models.

Provides calibrated survival distributions.

Abstract

Survival analysis provides a powerful statistical framework for modeling time-to-event outcomes in the presence of censoring. However, selecting an appropriate estimator from the many specialized survival approaches often requires substantial methodological and domain expertise. We introduce SurvivalPFN, a prior-data fitted network that amortizes Bayesian inference for censored observations through in-context learning. SurvivalPFN is pretrained on a diverse family of synthetic, identifiable, and right-censored data-generating processes, enabling it to amortize survival analysis in a single forward pass during inference. As a result, the model adapts to the effective complexity of each dataset without task-specific training or hyperparameter tuning, avoids restrictive parametric assumptions, and produces calibrated survival distributions. In a large-scale benchmark spanning 61 datasets, 21 methods, and 5 evaluation metrics, SurvivalPFN achieves strong predictive performance and often improves upon established survival models. These results suggest that SurvivalPFN offers a principled and practical foundation model for survival analysis, with potential applications in high-impact domains such as healthcare, finance, and engineering (https://github.com/rgklab/SurvivalPFN).