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

Motivation

• Censoring makes naive time-to-event regression waste information and bias estimates toward earlier events, while survival analysis uses both complete and censored outcomes.This matters in clinical settings where risk changes over time.Source: §1 · p. 1
• Tabular values derive meaning from column names and units, making heterogeneous covariates a challenge for foundation models and favoring task-specific survival methods.Classical approaches include Cox Proportional Hazards and Accelerated Failure Time models.Source: §1 · p. 2

Research gap

• Prior-fitted learning had advanced tabular classification and regression, but survival requires modeling event-time distributions under censoring and time-dependent risk.The paper brings prior-fitted in-context learning to survival analysis and evaluates it on heterogeneous public cohorts.Source: §1 · p. 2

Contributions

• The proposed synthetic survival-data framework explicitly controls covariate and survival-time distributions.It is designed to support synthetic pretraining without relying on large real-world cohorts.Source: §1 · p. 2
• SIC is pretrained solely on synthetic data and predicts in one forward pass without task-specific training or hyperparameter tuning.The authors identify small- and medium-sized datasets as a particularly attractive setting.Source: §1 · p. 2

Evaluation

• The evaluation compiles public real-world survival datasets and compares SIC with established baselines using extensive hyperparameter tuning.The dataset collection draws on prior work and major libraries to cover heterogeneous settings.Source: §1 · p. 2

Tabular foundation models

• Tabular foundation models represent tables through joint input representations, language models used as tabular learners, or column semantics.The cited approaches aim to improve performance and reduce implementation time through model reuse.Source: §2 · p. 2

Prior-fitted networks

• Prior-fitted networks train on synthetic datasets drawn from a prior and use in-context learning to approximate its Bayesian posterior predictive distribution.After training, the model amortizes inference for new datasets without task-specific parameter updates or hyperparameter tuning.Source: §2 · pp. 2–3

SIC overview

• The SIC framework has two components: synthetic data generation with a survival prior and an in-context learning architecture with specialized time-event modeling.The figure describes covariate generation followed by continuous and discrete survival-prior options for time-event labels.Source: Figure 1 · p. 3, §2 · p. 3

Survival analysis

• Survival data pair covariates with observed time and event status, including right-censoring; the target survival function gives the probability of surviving beyond time t.Methods may model time continuously or discretely, with corresponding assumptions about hazard or survival functions.Source: §2 · p. 3

Survival methods

• Continuous-time approaches include CoxPH, AFT, and neural variants; event-time distribution assumptions can help with limited data but depend on application knowledge.DeepEH unifies PH and AFT formulations under the extended hazard assumption [23].Source: §2 · p. 4
• Discrete-time methods range from population-level Kaplan-Meier to covariate-aware models and DeepHit, while grid sensitivity motivates continuous hazard alternatives.Neural ODEs and monotone positive neural networks [66] [35] are cited as such alternatives.Source: §2 · p. 4

Synthetic data generation

• Earlier synthetic survival-data work modeled event times or generated covariates conditionally, but some methods require real-world datasets that are scarce in survival analysis.The paper instead aims to generate covariates without real data while modeling time-to-event outcomes using statistical properties and prior knowledge.Source: §2 · p. 4

3.1 Synthetic Data Generation

• Synthetic datasets use sampled DAGs with MLP- or tree-based relations and noisy features, with multiple SCM-generated targets for survival data.The target and train features are selected randomly, and survival-prior targets receive z-score normalization.Source: §3.1 · p. 4

3.1.1 Survival Prior

• The survival prior combines a flexible continuous Extended Hazard model with a more function-agnostic discrete prior based on positive monotone neural networks.Extended Hazard settings recover PH, AFT, and AH assumptions; the discrete prior normalizes survival curves to start at 1 and end at 0.Source: §3.1.1 · p. 5
• Figure 2 illustrates diverse C-index, Kaplan–Meier, and normalized event-survival patterns across 512 generated datasets; both prior types are used to train SIC.The figure visualizes continuous and discrete priors.Source: Figure 2 · p. 6, §3.1.1 · p. 6
• α controls curve scale, k controls weight heterogeneity, µt shifts the time reference, and λ controls covariate-effect strength; censoring is generated independently.The methods also apply administrative censoring at a random quantile of the latent event-time distribution.Source: §3.1.1 · pp. 5–6

3 Survival In-context

• Across 12 real-world datasets, SIC achieves comparable time-dependent C-index to CoxPH[17], DeepHit[46], DeepSurv[37], XGB[15], RSF[34], ConSurv[47], and TabICL[63] + Stacking[18].Baselines were tuned with 100 trials; SIC requires no hyperparameter tuning.Source: Figure 3 · p. 7

3.2 In-context Learning

• SIC uses TabICL’s encoder, a time-event embedding, and a DeepHit survival head with 10 fixed quantile bins; training gradually increases sample counts in two stages.The DeepHit loss combines negative log-likelihood and ranking losses.Source: §3.2 · p. 7

4 Experiments

• The evaluation covers public datasets with 137–9,105 samples and 6–40 raw features, using nested 5-fold cross-validation and tuned baselines against a single-pass, untuned SIC.The comparison includes seven strategies: CoxPH, DeepSurv, DeepHit, RSF, XGBoost, Consurv, and TabICL + Stacking.Source: §4 · pp. 7–8

4.1 Results

• Across 12 datasets, SIC is better than or on par with most dataset-specific baselines; FRAMINGHAM and SUPPORT are exceptions with small overall variability, except for TabICL + Stacking.TabICL + Stacking has substantially lower Ctd, particularly on larger datasets.Source: §4.1 · p. 8
• SIC achieves the best mean and median rank across datasets ordered by size, suggesting generalizability across dataset sizes and survival tasks with one model.The ranks summarize performance across the 12 datasets.Source: §4.1 · pp. 8–9
• On SUPPORT, SIC completes 5-fold cross-validation in 0.003h (10.25s), versus 0.028h for XGBoost, 0.281h for TabICL+Stacking, 1.82h for CoxPH, 3.63h for DeepSurv, 3.37h for DeepHit, 4.65h for RSF, and >25h for Consurv.Consurv exceeded the 5h constraint per fold.Source: §4.1 · p. 9

4.2 Discussion and Ablations

• Across prior-family ablations, the mixed prior performs better on average and is never the worst configuration, although no prior uniformly dominates across datasets.It combines complementary inductive biases to balance prior expressivity against spending Monte Carlo training effort on unrealistic regions.Source: §4.2 · p. 9
• On restricted-access UNOS and SEER, with 62,645 and 280,845 samples, SIC is not the best-performing model; extending its prior may help but increases computational burden.The stated burdens include GPU memory limitations and longer training times.Source: §4.2 · p. 9
• After stage 1, no discretization-bin count consistently wins, while DeepEH and DeepSurv heads are comparable or inferior to DeepHit, selected for flexibility and non-parametric formulation.Ten bins were chosen for subsequent stages because they train faster; alternative heads can predict in continuous time.Source: §4.2 · p. 9
• SIC is miscalibrated on larger datasets, similarly to DeepHit, but CiPOT post-training calibration yields performance comparable to CoxPH and DeepSurv.The authors identify calibration as a limitation and suggest calibration-aware losses, alternative heads, or post-training calibration for future work.Source: §4.2 · p. 9

5 Conclusion

• SIC is presented as the first prior-fitted in-context learning model for survival analysis.Its approach is enabled by a method for synthesizing survival data.Source: §5 · p. 9
• The synthetic generator uses structural causal models, the extended hazard assumption [23], and monotonic networks [66] to create covariates and time-to-event outcomes without real-world datasets.The authors state that the generated data retain statistical properties.Source: §5 · p. 10
• Across multiple real-world datasets, one SIC model outperforms or matches dataset-specific classical and deep baselines without tuning and runs orders of magnitude faster.The authors connect this to reduced training overhead and prediction in a single inference pass.Source: §5 · p. 10

A.1 Time Derivation under Extended Hazard Assumption

• The time derivation starts from the extended hazard assumption to derive event time T [23].The supplied passage does not state the derivation steps.Source: §A.1 · p. 16

A.2 Inverse Baseline Cumulative Hazard Distributions

• Table 1 presents inverse cumulative-hazard functions for baseline survival distributions.Source: Table 1 · p. 16

A.3 Control Parameters.

• Scale anchoring makes each baseline survival curve pass through a shared randomly selected point, while the shape parameter is sampled independently.The procedure is applied to Weibull and other distributions.Source: §A.3 · p. 16

B.1 Anchoring Endpoints of Survival Curve.

• The network output is normalized so the survival curve equals 1 at time 0 and 0 at time 1.Source: §B.1 · p. 17

B.2 Weight Initialization.

• Positive weights are initialized from a Gamma distribution, with a Xavier fan-in correction based on the second moment to control pre-activation variance.The passages derive the Gamma scale from the target second moment.Source: §B.2 · p. 17

B.3 Censoring and observed outcomes.

• Censoring times come from a separate time-only monotone network; event and censoring times are scaled by dataset-specific tmax, then administrative censoring determines observed times and event indicators.Observed time is the minimum of event time, censoring time, and cutoff; the event indicator records whether event time occurs first.Source: §B.3 · p. 17

B.4 Control Parameters.

• The prior controls α, k, µt, and λ set global steepness, weight concentration and tail-heaviness, transition location, and covariate-driven separation, respectively.Figure 5 visualizes survival and Kaplan–Meier curves for a batch of 100 datasets.Source: §B.4 · pp. 17–18, Figure 5 · p. 19

C DeepHit Loss

• The DeepHit loss uses predicted time-point probabilities and cumulative distributions, admissible event-survivor pairs, and a hyperparameter trade-off between likelihood and ranking.A second hyperparameter controls the steepness of the ranking penalty.Source: C DeepHit Loss · p. 20

D Training Details

• Training uses 10K steps with 1,024 samples per table in stage 1, followed by 1K steps with variable table sizes from 1K to 10K in stage 2.Source: D Training Details · p. 20
• Each step uses 512 synthetic datasets with at most 100 covariates, and the encoder starts from TabICL classification weights.Source: D Training Details · p. 20
• Stage 1 took 91h and stage 2 took 56h on one A100 80GB GPU, using AdamW with different stage-specific learning-rate schedules.Automatic mixed precision was applied globally.Source: D Training Details · p. 20

E Datasets

• Evaluation cohorts span cancer, cardiovascular disease, HIV, transplantation, and critical care outcomes.Source: E Datasets · pp. 20–21
• Competing events are excluded for the Framingham, SEER, and PBC targets, while UNOS and SEER have restricted access.Source: E Datasets · pp. 20–21

F Details on Hyperparameter Tuning for Baseline Comparison

• Baseline tuning allows up to 100 trials or 5 hours per fold, using model-specific software implementations.Source: F Details on Hyperparameter Tuning for… · p. 22

• Time-dependent C-index measures pairwise ordering of relative risks.Source: G Evaluation Metrics · p. 23
• Integrated Brier Score integrates Brier score across the evaluation interval, with lower scores indicating better performance.Source: G Evaluation Metrics · p. 23
• D-CAL tests whether predicted survival probabilities at true event times are uniformly distributed across 20 intervals.The report counts folds with p > 0.05 as calibrated.Source: G Evaluation Metrics · p. 23

H Numerical Results

• Table 4 reports time-dependent C-index means and standard deviations across five folds.Source: H Numerical Results · p. 23, Table 4 · p. 24

I.1 Priors

• Figure 6 compares time-dependent C-index for discrete-only, continuous-only, and mixed survival priors, with the mixed prior used for evaluation.Source: Figure 6 · p. 24

I.2 Generalization to larger datasets.

• Table 5 reports time-dependent C-index across five folds for the larger-dataset evaluation.Source: Table 5 · p. 25

I.3 Discretization

• Figure 7 compares 10, 50, and 100 discretization bins; 10 bins are the main evaluation setting.Source: Figure 7 · p. 25

I.4 Survival Heads

• Figure 8 compares DeepHit, DeepSurv, and DeepEH survival heads; DeepEH’s FLCHAIN output is NaN and excluded from comparison.Source: Figure 8 · p. 26

I.5 Calibration

• SIC is miscalibrated, while SIC-PostCal has stronger calibration similar to continuous-time CoxPH and DeepSurv; its discrimination change is not significant.Figure 10 reports Integrated Brier Score and follows a trend similar to D-calibration.Source: §I.5 · p. 26

J Limitations and Outlook.

• Comparisons across prior-fitted architectures remain untested because the analysis focuses on TabICL; future work could assess TabDPT or TabPFNv2.TabICL was chosen for scalability and public availability, while broader comparisons require substantial computational resources.Source: J Limitations and Outlook · p. 28
• Supporting recurrent events, competing risks, or time-varying effects would require changing both the in-context architecture and the prior formulation.The current approach does not cover these settings.Source: J Limitations and Outlook · p. 28

K Broader Impact

• Clinical predictions should be validated by domain experts; improper use or overreliance could contribute to inequitable care or discriminatory outcomes.The paper calls for transparency, accountability, and sustained dialogue in healthcare use.Source: K Broader Impact · p. 28