SurvivalPFN: Amortizing Survival Prediction via In-Context Bayesian Inference

2026-05-14

AI Summary

This work addresses the limitations of existing survival analysis approaches—namely, their reliance on expert knowledge for model selection, cumbersome hyperparameter tuning, and restrictive parametric assumptions—by introducing, for the first time, contextual Bayesian inference within the foundation model paradigm. The proposed method leverages a pre-trained Prior-Data Fitting Network (PFN) that, at inference time, performs Bayesian inference on right-censored data through a single forward pass, eliminating the need for task-specific fine-tuning. It automatically adapts to data complexity and yields well-calibrated survival distributions. Evaluated on a large-scale benchmark encompassing 61 datasets, 21 competing methods, and 5 performance metrics, the approach significantly outperforms current state-of-the-art survival models.

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).

Problem

• survival analysis
• censoring
• estimator selection
• time-to-event prediction

Innovation

• SurvivalPFN
• amortized inference
• in-context learning
• Bayesian survival analysis
• right-censored data