SurvivalPFN

SurvivalPFN

SurvivalPFN is a prior-data fitted network for right-censored survival analysis. It predicts posterior predictive event-time distributions for query rows using an in-context training set of covariates, observed times, and censoring indicators.

• Paper: Amortizing Survival Prediction via In-Context Bayesian Inference
• Code: rgklab/SurvivalPFN
• License: Apache 2.0

Files

• survivalpfn_v0.1.pt: PyTorch checkpoint containing model weights and a sanitized model architecture config.
• checkpoint_info.json: Machine-readable summary of the uploaded checkpoint artifact.

Install

Clone the code repository and install one of the provided environments:

git clone https://github.com/rgklab/SurvivalPFN.git
cd SurvivalPFN

conda env create -f env-cpu.yaml
conda activate survivalpfn_cpu

For GPU/CUDA usage, use env.yaml instead. Alternatively, install from requirements.txt.

Quick Start

Run the release-checkpoint quickstart from the code repository:

python examples/quickstart.py

If CUDA is available:

python examples/quickstart.py --device cuda:0

The public inference API loads this Hugging Face checkpoint by repo id:

from survivalpfn import SurvivalEstimator

estimator = SurvivalEstimator(device="cpu", model_path="shi-ang/SurvivalPFN")
estimator.fit(X=X_train, T=T_train, delta=delta_train)
dist = estimator.predict_event_distribution(X_test)

To use a local checkpoint instead, pass the local .pt path as model_path.

Evaluation

The code repository includes dataset evaluation entrypoints. For example:

python -m survivalpfn.inference.evaluate --data PBC --n-exp 10 --seed 0 --device cpu --out output/pbc_survivalpfn.csv

See rgklab/SurvivalPFN for the full evaluation, training, testing, and reproduction instructions.

Checkpoint Details

• Format: PyTorch zip checkpoint (.pt)
• Size: 75,435,990 bytes (~72 MiB)
• SHA-256: a7b844809ad2dc5a675384704f7b6d85dfb88e96b82031c4542ce90a4fd439bb
• Top-level checkpoint keys: model_state_dict, model_config
• Model state: 151 tensors, 18,844,032 tensor values
• Model config: sanitized architecture fields only

Intended Use

This checkpoint is intended for research use with the SurvivalPFN implementation. It should not be used for clinical or other high-stakes decision making without independent validation in the target setting.

Citation

@article{qi2026survivalpfn,
  title         = {SurvivalPFN: Amortizing Survival Prediction via In-Context Bayesian Inference},
  author        = {Qi, Shi-ang and Balazadeh, Vahid and Cooper, Michael and Greiner, Russell and Krishnan, Rahul G.},
  journal       = {arXiv preprint arXiv:2605.15488},
  year          = {2026},
  eprint        = {2605.15488},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG}
}