[Paper Review] Survival analysis as a classification problem

Chenyang Zhong, Robert Tibshirani|arXiv (Cornell University)|Sep 24, 2019

Statistical Methods and Inference6 references4 citations

TL;DR

This paper proposes a stacking-based classification framework for survival analysis, transforming survival data into a binary classification problem by treating each uncensored event as a classification task within its risk set. The method enables the use of diverse machine learning algorithms—such as random forests, gradient boosting, and neural networks—yielding improved survival curve estimation over the Cox proportional hazards model, especially in settings with complex interactions and non-linear effects.

ABSTRACT

In this paper, we explore a method for treating survival analysis as a classification problem. The method uses a "stacking" idea that collects the features and outcomes of the survival data in a large data frame, and then treats it as a classification problem. In this framework, various statistical learning algorithms (including logistic regression, random forests, gradient boosting machines and neural networks) can be applied to estimate the parameters and make predictions. For stacking with logistic regression, we show that this approach is approximately equivalent to the Cox proportional hazards model with both theoretical analysis and simulation studies. For stacking with other machine learning algorithms, we show through simulation studies that our method can outperform Cox proportional hazards model in terms of estimated survival curves. This idea is not new, but we believe that it should be better known by statistiicians and other data scientists.

Motivation & Objective

• To bridge survival analysis and machine learning by recasting survival problems as classification tasks.
• To develop a stacking-based framework that enables the use of diverse classification algorithms for survival data.
• To demonstrate that stacking with logistic regression approximates the Cox proportional hazards model.
• To show that stacking with advanced ML algorithms improves survival curve estimation in complex, non-linear settings.
• To extend the framework to time-varying covariates and validate its performance.

Proposed method

• Construct a large data frame by stacking risk sets for each uncensored event, where each row represents a subject at risk at a failure time.
• Define a binary response vector with 1 for the subject who failed and 0 for others in the same risk set.
• Form a predictor matrix from covariates of all subjects in each risk set, with additional binary columns for subject-specific intercepts.
• Apply classification algorithms (logistic regression, random forests, GBM, neural networks) to the stacked data to estimate survival probabilities.
• Use predicted survival curves from the classification models to estimate individual survival functions.
• Extend the method to time-varying covariates by using the most recent covariate values before each event time.

Experimental results

Research questions

• RQ1 Can survival analysis be effectively reformulated as a classification problem using a stacking framework?
• RQ2 How does the stacking method with logistic regression compare to the Cox proportional hazards model in terms of estimation and inference?
• RQ3 Can machine learning algorithms applied via stacking improve survival curve estimation compared to the Cox model, especially in non-linear or interaction-heavy scenarios?
• RQ4 How does the stacking method perform with time-varying covariates compared to standard Cox regression?
• RQ5 What is the predictive performance of different stacking-based models using Harrell’s C-index and AUC metrics?

Key findings

• Stacking with logistic regression produces estimation and inference results that are approximately equivalent to the Cox proportional hazards model, both theoretically and empirically.
• For model 1 with linear effects, stacking with logistic regression and Cox regression achieved identical AUC values of 0.743.
• For model 2 with non-linear and interaction effects, stacking with GBM achieved an AUC of 0.781, outperforming Cox regression (0.727) and logistic regression (0.727).
• Random forests and neural networks showed mixed performance: random forests had AUC 0.686–0.691, while neural networks had lower AUC (0.620–0.615) in model 2, indicating sensitivity to model complexity.
• The stacking method with GBM and neural networks provided more accurate survival curve estimates than Cox models in simulations with complex effects, as shown in visual comparisons.
• The method extended to time-varying covariates produced coefficient and p-value estimates very close to those of Cox regression, validating its consistency in dynamic covariate settings.