PaperlayerView exact revision on arXiv

Source-linked AI summary

Time-to-Event Prediction with Neural Networks and Cox Regression

Håvard Kvamme, Ørnulf Borgan, Ida Scheel

arXiv:1907.00825v2stat.MLcs.LG

TL;DR

Time-to-event prediction needs models that retain survival-analysis capabilities while improving predictive flexibility and scalability. The paper extends Cox regression with neural networks and a nested-case-control loss for proportional and non-proportional modeling. Across five real-world data sets, Cox-Time has the best overall integrated Brier and integrated binomial log-likelihood performance, while DeepHit has stronger discrimination but poorer calibration.

  • Problem

    Time-to-event prediction requires flexible survival models, but standard survival methods can emphasize interpretability over predictive accuracy and Cox training can be difficult to scale.

  • Method

    The paper parameterizes Cox relative-risk functions with neural networks and uses nested case-control sampling to construct a batch-computable loss for proportional and non-proportional extensions.

  • Results

    Cox-Time achieves the best overall performance in integrated Brier score and integrated binomial log-likelihood across five real-world data sets, while DeepHit generally has the best discrimination but poorly calibrated survival estimates.

  • Takeaways & Limitations

    The methods provide scalable neural survival modeling, and estimated survival curves can also support descriptive customer-churn segmentation.

  • Takeaways & Limitations

    Further work is needed for competing events, time-dependent covariates, dynamic and recurrent events, advanced network structures, and less expensive Cox-Time survival estimates.

Abstract

from arXiv · show

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.

1. Introduction

The paper addresses time-to-event prediction by combining survival models with machine learning, motivated by the predictive limitations of interpretability-focused survival analysis and the value of survival models under censoring.

  • Time-to-event prediction models when future events occur and applies across cancer survival, customer churn, credit scoring, and mechanical-system failures.
  • Survival analysis has emphasized interpretability, potentially at some cost of predictive accuracy, while binary classifiers are commonly used in applicable industrial settings.
  • Survival models tend to be advantageous when applications contain substantial censoring.
  • The proposed approach extends Cox proportional hazards with neural networks and removes its proportionality constraint.
  • Nested case-control methodology makes the extensions scalable, combining neural-network flexibility with continuous-time event modeling.

2. Related Work

Prior work extends Cox regression with neural networks, discrete-time neural survival models, and flexible tree ensembles, establishing the main comparison landscape for the paper.

  • Early neural-network extensions of Cox regression generally failed to outperform regular Cox models, whereas later deep-learning networks outperformed classical Cox models on the C-index.
  • SurvivalNet and image-based extensions use neural networks to fit proportional Cox models for tabular or pathological-image data.
  • Discrete-time alternatives compute hazards, survival probabilities, or probability mass functions on predetermined time grids using neural networks.
  • DeepHit combines a neural-network log-likelihood with a ranking loss and supports competing-risks applications.
  • Random survival forests use log-rank splitting and averaged leaf-node cumulative hazards, providing a flexible continuous-time method without the proportionality constraint.

3. Methodology

The methodology combines Cox survival modeling with neural networks and case-control sampling to support scalable proportional and non-proportional extensions trained with batched optimization.

  • Survival-analysis foundations: Survival analysis models event-time distributions continuously while accommodating right censoring, where only a possibly censored event time and event indicator are observed.
  • Cox regression: The Cox model specifies a hazard using a non-parametric baseline hazard and an exponentiated relative-risk function of covariates.
  • Cox regression: Cox regression estimates the parametric component by maximizing the partial likelihood, then estimates the baseline hazard from those results.
  • Cox with SGD: Mini-batch SGD is difficult because Cox losses sum over risk sets that can span the full data set, motivating an easily batched approximation.
  • Cox with SGD: The approximation samples weighted subsets of risk sets, always includes the event individual, omits non-events from loss computation, and averages over events.
  • Cox with SGD: Case-control sampling produces a fixed-size risk-set approximation that enables batched gradient descent and is often effective with one sampled control.
  • Non-linear Cox: The neural Cox model replaces the linear predictor with a neural network, yielding Cox-MLP and distinguishing case-control training from the DeepSurv loss.
  • Non-Proportional Cox-Time: The non-proportional extension lets the relative-risk function depend on time, allowing neural networks to model interactions between time and covariates without stratification.

4. Evaluation Criteria

The evaluation criteria address discrimination, calibration, and censoring in time-to-event prediction. They include concordance-based metrics, censoring-adjusted Brier scores, and censoring-adjusted binomial log-likelihoods.

  • Concordance metrics: The C-index estimates whether predicted survival times preserve the ordering of randomly paired individuals.It is commonly used because of its relationship to classification accuracy and ROC AUC.
  • Concordance metrics: For proportional hazards models, the regular C-index can use relative risk because prediction ordering does not change over time.For non-proportional hazards models, the evaluation uses a time-dependent C-index with adjustments for tied event times and survival estimates.
  • Brier score: The Brier score measures discrimination and calibration through squared probability error, extended to censored durations using inverse censoring weights.The censoring adjustment uses the Kaplan–Meier estimate of the censoring survival function and assumes censoring and survival times are independent.
  • Integrated metrics: Integrated Brier score and integrated binomial log-likelihood summarize performance across a time interval rather than at one fixed time.In practice, the integral is approximated numerically over the test-set duration using 100 grid points for stable scores.

5. Simulations

The simulations examine convergence, parameter estimates, and likelihood agreement for the proposed case-control loss. Results indicate that one sampled control is generally sufficient and that Cox-SGD approaches classical Cox behavior as data size increases.

  • Simulation design: The simulations were designed to verify that the proposed methods behave as expected rather than exhaustively evaluate them.Classical Cox denotes Cox regression with a linear relative-risk function fitted using Lifelines.
  • Simulation design: The proposed loss was tested under a proportional hazards model with constant baseline hazard, three specified coefficients, uniform covariates, and approximately 30% censoring.The simulation used independent censoring and censored individuals still under observation at time 30.
  • Loss behavior: With 10,000 training and 10,000 testing individuals, sampling 1, 2, 4, or 8 controls did not affect convergence rate, although computational complexity increased with controls.The experiment was repeated 100 times and evaluated using test-set mean partial log-likelihood.
  • Parameter estimates: For data sets of sizes 100, 1,000, and 10,000, Cox-SGD parameter estimates were compared with classical Cox estimates across 100 sampled data sets.For size 100, Cox-SGD estimates appeared slightly smaller, with a larger difference when fewer controls were sampled.
  • Likelihood comparison: As data sets grew, Cox-SGD mean log-likelihood appeared to converge to classical Cox regardless of control sample size.The classical Cox mean log-likelihood was approximately -2.2, while the difference for one sampled control was around 0.1% even for the smallest data sets.
  • Neural-network extensions: The simulations also verified that neural-network Cox models can represent nonlinear relative-risk functions in proportional hazards settings.Additional appendix simulations examined nonlinear and non-proportional models.

6. Experiments

The experiments compare neural-network Cox extensions with classical Cox regression, DeepHit, and random survival forests across common survival data sets and KKBox churn data. Cox-Time generally performs best on calibration-oriented metrics, while DeepHit leads concordance but often has poorly calibrated survival estimates.

  • 6.1 Four Common Survival Data Sets: Five survival data sets are evaluated against Classical Cox (Linear), DeepHit, Random Survival Forests, and proportional neural-network Cox methods.The experiments use four smaller literature data sets and one large data set for deeper analysis.
  • 6.1.1 Methods and Hyperparameter Tuning: Models are tuned with 5-fold crossvalidation, using partial log-likelihood for proportional Cox methods, Cox-Time loss for Cox-Time, and concordance for DeepHit and RSF.Neural networks use random searches over 300 configurations per fold and select models on held-out validation data.
  • 6.1.2 Results: DeepHit has the best overall discriminative performance, but this comes at the cost of poorly calibrated survival estimates.Its concordance advantage is contrasted with weaker integrated Brier score and binomial log-likelihood performance.
  • 6.1.2 Results: Cox-Time generally performs best on integrated Brier score and integrated binomial log-likelihood, although it struggles on FLCHAIN.For FLCHAIN, Cox-MLP (DeepSurv) performs best on both calibration metrics, while Cox-MLP (CC) has the best concordance.
  • 6.1.2 Results: Classical Cox (Linear) performs poorly under restrictive model assumptions, while the two Cox-MLP variants are generally close but can be unstable.The authors note that combining a flexible neural network with the proportionality constraint may be problematic for large data sets.
  • 6.2.3 Survival Curves: Cox-Time provides detailed survival curves, revealing heterogeneous churn patterns including immediate high risk, month-end hazard increases, and a day-400 drop.In the KKBox analysis, 18% of customers are in a high-immediate-risk cluster, 19% have nearly constant-slope curves, and the smallest cluster represents 3% of test data.

7. Discussion

The paper extends Cox regression with neural networks, including a scalable loss for proportional and non-proportional models. Across five real-world data sets, Cox-Time achieved the best overall IBS and IBLL, while survival-curve clustering offered a descriptive view of customer churn.

  • Extensions and scalability: Neural networks parameterize relative risk and interactions between covariates and time, removing the Cox proportionality constraint.A batch-computable loss based on nested case-control methods enables scaling to large data sets.
  • Extensions and scalability: The proposed loss produces Cox-model fits close to those obtained with the full Cox partial likelihood in simulation studies.
  • Real-world evaluation: Cox-Time had the best overall performance in integrated Brier score and integrated binomial log-likelihood across five real-world data sets.The comparison included classical Cox regression, RSF, DeepHit, and DeepSurv.
  • Real-world evaluation: DeepHit generally had the best discriminative performance but poorly calibrated survival estimates.
  • Survival-curve interpretation: Clustering estimated survival curves provided a descriptive view of customer churning behavior in a customer churn data set.
  • Future directions: Future work includes competing events, time-dependent covariates, dynamic predictions, recurrent events, advanced network structures, and cheaper Cox-Time survival estimates.

A.1 Four Common Survival Data Sets Tuning

The appendix specifies the hyperparameter search space and tuning procedure for four common survival data sets.

  • Hyperparameter search: Table A.1 defines the search space for Rot. & GBSG, SUPPORT, METABRIC, and FLCHAIN.The experiments sampled 300 random parameter configurations for each method, fold, and data set.

A.2 KKBox Tuning

The KKBox study used grid-search tuning, with some network parameters selected using a smaller network and fixed training settings for the proposed methods.

  • KKBox tuning: KKBox hyperparameters were selected by grid search over relevant configurations in Table A.2.
  • KKBox tuning: Parameters marked with an asterisk in Table A.2 were found using a two-layer network with 128 nodes.
  • Survival-curve output: Appendix Figure B.1 displays Cox-Time survival probabilities against days since first subscription for the churn case study.
  • KKBox tuning: The smaller-network settings were used to reduce hyperparameter search for the DeepHit loss and Cox-Time duration transformation.
  • Training settings: Cox-MLP (CC) and Cox-Time used a fixed penalty λ = 0.001 and batch size 1028.
  • Survival-curve output: Nine example KKBox test-set curves illustrate that Cox-Time learned a variety of survival-curve shapes.

B.1 KKBox Data Set

The KKBox data combine customer and transaction information to model churn, using a 30-day resubscription rule and time-fixed covariates. The study simplifies the data by excluding user logs and treats resubscribed churners as new customers.

  • Data construction: KKBox includes customer information, transaction logs, and aggregate user-usage logs.
  • Data construction: A customer is labeled churned after failing to resubscribe more than 30 days after the last subscription expired.
  • Data construction: Survival time is measured from first subscription or the earliest transaction record, and previously churned customers who resubscribed are treated as new customers.
  • Covariates and scope: All covariates are measured at subscription and held constant over time, while user logs are excluded.The authors identify this as a simplification because customer covariates are typically non-stationary.
  • Covariates and scope: The data include 15 covariates: 7 numeric and 8 categorical.Most numerical covariates were log-transformed.
  • Evaluation scope: The additional simulations are intended only to verify expected method behavior, not to provide a general evaluation.

C.1 Non-Linear Hazards

The simulations replace the Cox model’s linear predictor with a nonlinear function and compare classical Cox regression, Cox-SGD, and Cox-MLP. Cox-MLP closely recovers the true partial likelihoods, whereas Cox-SGD struggles with the nonlinear predictor.

  • Nonlinear predictor: The simulation replaces g(x) = βT x with a nonlinear predictor involving squared terms and pairwise interactions.The simulations otherwise retain the linear-case setup.
  • Experimental setup: The study uses 10,000 training samples, 10,000 test samples, and 1,000 validation samples, with validation data used for early stopping.
  • Models: Cox-MLP uses a one-hidden-layer MLP with 64 ReLU-activated hidden nodes to parameterize the Cox predictor.The model is evaluated alongside Cox-SGD and classical Cox regression.
  • Results: For 2,000 test individuals, Cox-MLP produces likelihood estimates very close to the true values, while Cox-SGD struggles to represent the nonlinear function.Estimates closer to the identity line indicate closer recovery of the true predictor and likelihoods.

C.2 Non-Proportional Hazards

The non-proportional simulations add time dependence to the hazard and compare Cox-MLP with Cox-Time. Cox-Time better captures both the true partial likelihoods and survival curves, while the simulation design imposes constraints to produce reasonable curves.

  • Non-proportional hazard: The hazard is extended to h(t | x) = h0 exp[g(t, x)] by adding a time-dependent term to the risk function.
  • Simulation assumptions: The simulation requires b(x) to be non-negative so hazards increase with time and adds sign(x3) to a(x), yielding two baseline hazards.The resulting baselines are 0.0074 and 0.054.
  • Experimental setup: Cox-MLP and Cox-Time use four-layer ReLU MLPs with 128 nodes per layer and dropout rate 0.1.The experiment samples 10,000 training, 10,000 test, and 2,000 validation observations.
  • Likelihood comparison: Cox-Time appears to capture the true partial log-likelihood values better than Cox-MLP, although the difference is not large.The comparison uses 2,000 test samples.
  • Survival-curve comparison: Cox-Time estimates the nine illustrated true survival curves better than Cox-MLP.Figure C.3 compares true curves with estimates from both methods.
  • Data generation: Survival times are generated by transforming exponential samples through the inverse cumulative hazard, with constant baselines used for analytical tractability.The proportional simulations use g(t, x) = g(x), while non-proportional simulations use g(t, x) = a(x) + b(x)t.
Loading 1907.00825v2…