Survival In-Context: Amortized Bayesian Survival Analysis via Prior-Fitted Networks

TL;DR

This paper introduces Survival In-Context (SIC), a novel in-context learning model for survival analysis that leverages synthetic data pretraining to provide accurate, individualized survival predictions without task-specific training.

Contribution

It proposes a flexible synthetic survival data generation framework and a prior-fitted in-context learning approach that outperforms traditional models, especially with limited data.

Findings

SIC achieves competitive or superior performance on real datasets.

SIC requires no task-specific training or hyperparameter tuning.

The approach is particularly effective in small to medium data regimes.

Abstract

Survival analysis is crucial for many medical applications, but remains challenging for modern machine learning due to limited data, censoring, and the heterogeneity of tabular covariates. While the prior-fitted paradigm, which relies on pretraining models on large collections of synthetic datasets, has recently facilitated tabular foundation models for classification and regression, its suitability for time-to-event modeling remains unclear. We propose a flexible survival data generation framework that defines a rich survival prior with explicit control over covariates and time-event distributions. Building on this prior, we introduce Survival In-Context (SIC), a prior-fitted in-context learning model for survival analysis that is pretrained exclusively on synthetic data. SIC is trained to approximate Bayesian posterior predictive inference under the synthetic survival prior, enabling individualized survival prediction in a single forward pass, requiring no task-specific training or hyperparameter tuning. Across a broad evaluation on real-world survival datasets, SIC achieves competitive or superior performance compared to classical and deep survival models, particularly in small and medium-sized data regimes, highlighting the promise of a prior-fitted paradigm for survival analysis. The code and pretrained models will be made available upon publication.