[Paper Review] Survival stacking: casting survival analysis as a classification problem

Erin Craig, Chenyang Zhong+1 more|arXiv (Cornell University)|Jul 28, 2021

Statistical Methods and Inference18 references4 citations

TL;DR

This paper proposes survival stacking, a method that reformulates survival analysis as a binary classification problem by constructing a stacked matrix of risk set indicators and event indicators. By applying general classification algorithms—such as logistic regression, glinternet, random forests, or neural networks—on this matrix, the approach enables flexible modeling of time-varying effects and interactions, outperforming traditional Cox regression in capturing complex hazard structures.

ABSTRACT

While there are many well-developed data science methods for classification and regression, there are relatively few methods for working with right-censored data. Here, we present "survival stacking": a method for casting survival analysis problems as classification problems, thereby allowing the use of general classification methods and software in a survival setting. Inspired by the Cox partial likelihood, survival stacking collects features and outcomes of survival data in a large data frame with a binary outcome. We show that survival stacking with logistic regression is approximately equivalent to the Cox proportional hazards model. We further recommend methods for evaluating model performance in the survival stacked setting, and we illustrate survival stacking on real and simulated data. By reframing survival problems as classification problems, we make it possible for data scientists to use well-known learning algorithms (including random forests, gradient boosting machines and neural networks) in a survival setting, and lower the barrier for flexible survival modeling.

Motivation & Objective

• To address the limitations of proportional hazards assumptions in traditional survival models like Cox regression.
• To enable the use of general classification algorithms—such as tree-based models and neural networks—for survival analysis.
• To allow modeling of non-linear relationships and time-varying effects through interaction terms between covariates and risk set indicators.
• To provide a scalable framework for large datasets by supporting mini-batching and alternative time representations.
• To unify survival analysis with modern machine learning by casting it as a classification problem on a transformed data matrix.

Proposed method

• Transform survival data into a stacked matrix where each row corresponds to a subject at each observed event time, with risk set indicators marking which subjects are at risk.
• Construct a binary response vector indicating whether the subject experienced an event at that time point.
• Apply general classification algorithms (e.g., logistic regression, glinternet, random forests, gradient boosting, neural networks) to the stacked matrix.
• Use time as a categorical, ordinal, or continuous variable in the predictor matrix, optionally with basis expansions for smoother time effects.
• Formulate the log-likelihood of the stacked model to match the full Cox partial likelihood under a discrete baseline hazard.
• Support memory efficiency via mini-batching for large datasets, avoiding full instantiation of the stacked matrix.

Experimental results

Research questions

• RQ1 Can survival analysis be effectively reformulated as a binary classification problem using a stacked data representation?
• RQ2 To what extent can general classification algorithms capture time-varying effects and interactions not modeled by proportional hazards assumptions?
• RQ3 How do different learning algorithms (e.g., glinternet, random forests, neural networks) perform in the survival stacked setting compared to Cox regression?
• RQ4 What are the implications of representing time as continuous or categorical in the stacked matrix for model flexibility and interpretability?
• RQ5 Can scalable inference be achieved in large survival datasets using mini-batching with the stacked matrix framework?

Key findings

• The stacked matrix formulation allows general classification algorithms to model survival data while preserving the full Cox partial likelihood under a discrete baseline hazard.
• Linear models applied to the stacked matrix recover the proportional hazards assumption, while non-linear models relax it by discovering interactions between covariates and risk set indicators.
• Methods like glinternet enable discovery of pairwise interactions and time-varying effects through regularization in logistic regression.
• Tree-based models (e.g., random forests, gradient boosted trees) naturally capture non-linear relationships and interactions between covariates and time.
• Neural networks can model rich, non-linear hazard functions by learning complex patterns in the stacked data.
• The framework supports scalable inference via mini-batching, making it feasible for large datasets where full matrix instantiation is impractical.