Survival In-Context: Amortized Bayesian Survival Analysis via Prior-Fitted Networks
Organizations: Institute of Diagnostic and Interventional Radiology, Technical University of Munich, Germany · Chair for AI in Healthcare and Medicine, Technical University of Munich, Germany · Hasso Plattner Institute for Digital Engineering, University of Potsdam, Germany · University Hospital Hamburg-Eppendorf, Germany ·
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.
Figures & tables
Appendix figures & tables12 assets
Supplementary material from the paper’s appendix.
Appendix
| Distribution | |
| Weibull | |
| Gompertz | |
| Lognormal | |
| Log-logistic | |
| Birnbaum–Saunders |
| Dataset | #samples | #features | Event rate | |
| VETERAN | 137 | 6 | 93.4% | |
| LUNG | 288 | 10 | 72.4% | |
| PBC | 312 | 20 | 44.9% | |
| WHAS500 | 500 | 14 | 43.0% | |
| GBSG2 | 686 | 9 | 43.6% | |
| AIDS | 1,151 | 11 | 8.3% |
| Model | Hyperparameter | Domain |
| Cox [ 16 ] | l1_ratio | [0.0, 1.0] |
| penalizer | [1e-6, 1] | |
| DeepSurv [ 36 ] | optimizer | {adam, sgd} |
| weight_decay | [1e-6, 1e-2] | |
| momentum | [0.0, 0.9] | |
| lr | [1e-5, 1e-2] |
| Dataset | SIC | CoxPH | DeepSurv | DeepHit | XGB | RSF | TabICL + Stacking | Consurv |
| VETERAN | 0.720 (0.020) | 0.724 (0.039) | 0.687 (0.044) | 0.632 (0.105) | 0.669 ∗ (0.031) | 0.686 (0.044) | 0.660 (0.076) | 0.586 ∗ (0.056) |
| LUNG | 0.615 (0.060) | 0.615 (0.053) | 0.577 (0.106) | 0.514 (0.075) | 0.562 ∗ (0.074) | 0.587 ∗ (0.072) | 0.565 (0.032) | 0.606 (0.027) |
| PBC | 0.819 (0.067) | 0.790 (0.049) | 0.793 (0.060) | 0.812 (0.064) | 0.814 (0.056) | 0.808 (0.062) | 0.747 ∗ (0.072) | 0.790 (0.062) |
| WHAS500 | 0.757 (0.026) | 0.766 (0.037) | 0.754 (0.039) | 0.765 (0.026) | 0.758 (0.043) | 0.773 (0.009) | 0.665 ∗ (0.036) | 0.747 (0.041) |
| GBSG2 | 0.691 (0.040) | 0.680 (0.024) | 0.656 (0.076) | 0.667 (0.050) | 0.671 (0.039) | 0.681 (0.037) | 0.558 ∗ (0.034) | 0.658 ∗ (0.035) |
| AIDS | 0.743 (0.046) | 0.743 (0.030) | 0.662 (0.072) | 0.685 (0.073) | 0.735 (0.042) | 0.715 (0.048) | 0.670 (0.057) | 0.755 (0.026) |
| Dataset | SIC | CoxPH | DeepSurv | DeepHit | XGB | RSF | TabICL + Stacking | Consurv |
| UNOS | 0.571 (0.004) | 0.589 ∗ (0.005) | 0.595 ∗ (0.005) | 0.611 ∗ (0.006) | 0.599 ∗ (0.006) | – | 0.529 ∗ (0.013) | 0.595 ∗ (0.004) |
| SEER | 0.840 (0.004) | 0.820 ∗ (0.003) | 0.848 ∗ (0.005) | 0.853 ∗ (0.004) | 0.867 ∗ (0.003) | – | 0.804 ∗ (0.004) | 0.850 ∗ (0.004) |
| Dataset | SIC | SIC-PostCal | Delta mean (SIC-PostCal - SIC) |
| VETERAN | 0.720 0.020 | 0.721 0.020 | +0.001 |
| LUNG | 0.615 0.060 | 0.615 0.057 | 0.000 |
| PBC | 0.819 0.067 | 0.821 0.065 | +0.002 |
| WHAS500 | 0.757 0.026 | 0.757 0.030 | 0.000 |
| GBSG2 | 0.691 0.040 | 0.690 0.040 | -0.001 |
| AIDS | 0.743 0.046 | 0.707 0.092 | -0.036 |
Survival Analysis
Survival analysis models when events such as disease progression, failure, or death occur, accounting for incomplete follow-up and censoring.
See all 106 papersSurvICL: a Foundation Model for in-context Survival Analysis
SurvPFN: Towards Foundation Models for Survival Predictions
SurvPFN, a prior-data fitted network (PFN), for survival prediction tasks. We pretrain SurvPFN on millions of synthetic survival prediction tasks to learn survival via distributional regression that accounts for censored data. SurvPFN works by (1) generating data with Weibull event times and a non-informative censoring mechanism; (2) integrating a censored event indicator; and (3) minimizing a censored negative log-likelihood. On SurvSet, a collection of real-world survival tasks, SurvPFN is highly competitive with classical and deep survival baselines without per-dataset fitting, a survival-specific architecture, or feature engineering. We show that survival can be treated as a continuous-time distributional regression problem with censored loss, unlocking the power of PFNs for time-to-event predictions.