[Paper Review] Deep Neural Networks for Survival Analysis Based on a Multi-Task Framework

Stephane Fotso|arXiv (Cornell University)|Jan 17, 2018

Machine Learning in HealthcareComputer Science15 references110 citations

TL;DR

The paper introduces Neural Multi-Task Logistic Regression (N-MTLR), a deep learning extension of MTLR for survival analysis, which improves survival predictions over MTLR and CoxPH, especially when nonlinear dependencies are present, and is implemented in TensorFlow/Keras.

ABSTRACT

Survival analysis/time-to-event models are extremely useful as they can help companies predict when a customer will buy a product, churn or default on a loan, and therefore help them improve their ROI. In this paper, we introduce a new method to calculate survival functions using the Multi-Task Logistic Regression (MTLR) model as its base and a deep learning architecture as its core. Based on the Concordance index (C-index) and Brier score, this method outperforms the MTLR in all the experiments disclosed in this paper as well as the Cox Proportional Hazard (CoxPH) model when nonlinear dependencies are found.

Motivation & Objective

• Motivate survival analysis for time-to-event forecasting in industries beyond medicine (e.g., churn, purchases, defaults).
• Overcome CoxPH limitations (proportional hazards, linear risk, unspecified baseline) by leveraging MTLR and deep learning.
• Propose the Neural MTLR (N-MTLR) to capture nonlinear dependencies while maintaining MTLR's interval-based survival modeling.

Proposed method

• Base on Multi-Task Logistic Regression to model survival across J time intervals with density and survival functions.
• Replace the linear core of MTLR with a deep neural network (MLP) to capture nonlinear relationships via a nonlinear mapping psi(x) to a J-dimensional output.
• Compute density f(a_s|x) and survival S(t|x) using the nonlinear psi(x) with the same Delta and Y constructs from MTLR and a normalization Z(psi(x)).
• Implement the model in Python using TensorFlow 2 and Keras with standard deep learning techniques (initialization, optimizer, activations, batch norm, dropout).
• Train with approaches compatible with survival data and evaluate with C-index and Brier score (and IBS) under right-censoring.
• Provide examples with 2-hidden-layer networks and various activations to illustrate nonlinear modeling benefits.

Experimental results

Research questions

• RQ1 Does N-MTLR outperform linear MTLR and CoxPH on synthetic data with linear, nonlinear square, and nonlinear Gaussian risk functions?
• RQ2 Do nonlinear dependencies in real datasets (WHAS, veteran) yield clear advantages for N-MTLR over traditional survival models?
• RQ3 How do C-index and Brier score (IBS) compare across CoxPH, MTLR, and N-MTLR under censoring?
• RQ4 Is the Python/TensorFlow/Keras implementation practical and advantageous for applying MTLR in real-world tasks?

Key findings

• On linear risk data, CoxPH, MTLR and N-MTLR show similar performance (C-index around 0.74; IBS around 0.10–0.12).
• On nonlinear square risk data, CoxPH and MTLR underperform while N-MTLR achieves higher C-index (0.70) and lower IBS (0.12) compared to CoxPH (0.53) and MTLR (0.53).
• On nonlinear Gaussian risk data, N-MTLR attains the best C-index (0.60) and IBS (0.13) among the three models.
• Real-data results show comparable performance across models on WHAS and veteran datasets, with N-MTLR often at least as competitive (e.g., WHAS: C-index 0.80 for N-MTLR vs 0.79 for CoxPH) and similar IBS (~0.17).
• Across experiments, N-MTLR consistently outperforms MTLR and matches or exceeds CoxPH when nonlinear dependencies are present, demonstrating the advantage of the deep, nonlinear core.