Survival In-Context: Prior-fitted In-context Learning Tabular Foundation Model for Survival Analysis

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 produces 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 medium-sized data regimes, highlighting the promise of prior-fitted foundation models for survival analysis. The code will be made available upon publication.

Figures (7)

• Figure 1: Survival In-context (SIC) method contains two stages: (I) Synthetic Data Generation using Structural Causal Models (SCMs) for the generation of (A) covariates X and the related $(\eta_1, \eta_2)$ label nodes, and (B) the proposed Survival Prior with survival quantile $U$ and the survival function parametrization $H_0^{-1}(*, \alpha, \beta)$ to generate time-event labels; (II) In-context learning on synthetic data with a specialized time-event embedding and survival head and loss. Created in https://BioRender.com.
• Figure 2: Visualization of a random batch of 512 generated datasets: (upper) time-dependent C-index estimated from CoxPH models fit on the synthetic datasets and (lower) Kaplan–Meier curves on the standardized time scale. The proposed synthetic data generation method demonstrates the ability to generate diverse survival curve patterns and outcome complexity.
• Figure 3: Comparison of time-dependent C-index performance (mean and standard deviation) across 5 folds on 13 datasets. The proposed Survival In-context (SIC) method achieves comparable performance to the baseline models CoxPH, DeepHit, DeepSurv, XGBoost (XGB). SIC does not require hyperparameter tuning, whereas all baselines are tuned with 100 trials. * indicates p-value $<$ 0.05 for the two-sided t-test of SIC compared to the baselines.
• Figure 4: Dataset rank by method across 13 datasets. The proposed Survival In-context (SIC) method shows the highest mean and median rank compared to the baseline models CoxPH, DeepHit, DeepSurv, XGBoost (XGB).
• Figure 5: Time-dependent C-index of SIC model across training stages: stage 1 trains on a prior with fixed sample size of 1024, stage 2 expands to 1K-40K, stage 3 to 40K-60K. Expanding beyond stage 1 improves performance on larger datasets, while stages 2 and 3 are nearly indistinguishable.