TorchSurv: A Lightweight Package for Deep Survival Analysis

TL;DR

TorchSurv addresses the rigidity of existing survival-analysis libraries by enabling deep survival models defined via custom PyTorch networks that map covariates $x$ to survival-model parameters $\theta$, with log-likelihoods computed entirely in PyTorch and trained via backpropagation. It introduces Cox and Weibull loss functions, a Momentum variant for training stability with small batches, and a suite of evaluation metrics including time-dependent AUC, C-index, and Brier score, along with confidence intervals and model-comparison tests. The framework emphasizes end-to-end flexibility, automatic differentiation, and numerical stability through log-scale computations, validated against open-source and synthetic datasets and complemented by thorough documentation and benchmarks. This yields a practical, high-signal toolkit for building and evaluating complex, high-dimensional survival models in Python, with cross-language context provided via an R-library appendix.

Abstract

TorchSurv is a Python package that serves as a companion tool to perform deep survival modeling within the PyTorch environment. Unlike existing libraries that impose specific parametric forms, TorchSurv enables the use of custom PyTorch-based deep survival models. With its lightweight design, minimal input requirements, full PyTorch backend, and freedom from restrictive survival model parameterizations, TorchSurv facilitates efficient deep survival model implementation and is particularly beneficial for high-dimensional and complex input data scenarios.