survivalGPU: Scalable Survival Analysis for R and Python

Abstract

This paper introduces survivalGPU, a high-performance Python and R package designed to accelerate survival analysis computations using graphics processing units (GPUs). While the Cox proportional hazards model remains a cornerstone of survival analysis, it can impose substantial computational burdens when applied to large longitudinal data sets. To address this limitation, survivalGPU implements Cox proportional hazards models and weighted cumulative exposure (WCE) models in Python, using PyTorch and KeOps for efficient back-end computation, with interfaces in both Python and R. Numerical validation shows that survivalGPU reproduces the results of reference implementations within the tolerances expected for the selected floating-point precision, with relative errors below 10 -2 in float32 and below 10 -5 in float64. Benchmark experiments demonstrate substantial speedups. For Cox models, coxphGPU() was 5.6x faster than coxph() from the survival package on 1,000,000 patients, and 14.0x faster based on runtimes extrapolated to 1,000 bootstrap replicates on 100,000 patients. Gains were larger for WCE models: wceGPU() was 920.7x faster than WCE() from the WCE package on 100,000 patients, and 1,321.4x faster based on runtimes extrapolated to 1,000 bootstrap replicates. These results show that survivalGPU makes flexible time-dependent survival modeling and bootstrap-based estimation of confidence intervals feasible for large-scale pharmacoepidemiological applications.