[Paper Review] Time-to-Event Prediction with Neural Networks and Cox Regression

Håvard Kvamme, Ørnulf Borgan+1 more|arXiv (Cornell University)|Jul 1, 2019

Statistical Methods and InferenceMathematics32 references230 citations

TL;DR

This paper extends the Cox proportional hazards model with neural networks to enable time-to-event prediction, introduces a scalable case-control loss for both proportional and non-proportional hazards, and demonstrates competitive performance on real data sets with a Python package available.

ABSTRACT

New methods for time-to-event prediction are proposed by extending the Cox proportional hazards model with neural networks. Building on methodology from nested case-control studies, we propose a loss function that scales well to large data sets, and enables fitting of both proportional and non-proportional extensions of the Cox model. Through simulation studies, the proposed loss function is verified to be a good approximation for the Cox partial log-likelihood. The proposed methodology is compared to existing methodologies on real-world data sets, and is found to be highly competitive, typically yielding the best performance in terms of Brier score and binomial log-likelihood. A python package for the proposed methods is available at https://github.com/havakv/pycox.

Motivation & Objective

• Motivate time-to-event prediction and address censoring via survival analysis.
• Extend the Cox model by integrating neural networks to capture non-linear covariate effects.
• Provide a scalable training loss based on case-control sampling suitable for large data sets.
• Extend to non-proportional hazards through a time-dependent relative risk function.
• Release a Python package implementing the proposed methods for reproducibility.

Proposed method

• Replace the linear predictor g(x) with a neural network g(x) in a Cox framework.
• Introduce a case-control sampling based loss that approximates the Cox partial log-likelihood and scales to batches (loss forms in equations 8 and 9).
• Generalize to non-linear models (Cox-MLP) and introduce a non-proportional time component (Cox-Time) with time-dependent g(t, x).
• Incorporate a small penalty to regularize g(x) when using neural networks (equation 10).
• Adapt prediction by estimating the baseline hazard via Breslow’s method and computing survival functions; for non-proportional hazards, discretize predictions over time for computational efficiency.
• Provide a Python package (pycox) implementing these methods.

Experimental results

Research questions

• RQ1 Can neural networks improve time-to-event prediction when integrated with Cox modeling while handling censoring?
• RQ2 Does a case-control based loss provide a scalable and accurate approximation to the Cox partial likelihood in large data settings?
• RQ3 Can the Cox framework be extended to non-linear (Cox-MLP) and non-proportional (Cox-Time) hazards without prohibitive computational costs?
• RQ4 How do the proposed methods compare to existing approaches (DeepSurv, DeepHit, RSF) on real-world survival data?

Key findings

• The proposed case-control loss approximates the Cox partial log-likelihood and remains stable across batch sizes.
• Cox-SGD parameter estimates converge toward classical Cox estimates as data size increases; sampling more controls helps small datasets but adds computational cost.
• Non-proportional hazards via Cox-Time yield improved survival function estimates compared to proportional-hazards methods on simulated data.
• Across five real-world datasets, including SUPPORT, METABRIC, Rotterdam & GBSG, and FLCHAIN, the proposed methods are competitive or superior to existing methods in discrimination and calibration metrics such as C-index, Brier score, and binomial log-likelihood.
• The authors provide a PyTorch-based Python package (pycox) for implementing their methods and reproducing simulations and analyses.