Step by step reference on machine learning for survival analysis?

Question

I want to apply machine learning methods to survival analysis. This is, I have a sample of survival times $(t_1,...,t_m)$, censoring indicators $(\delta_1,...,\delta_m)$, and covariates ${\bf{x}_i}$, $i=1,...,m$.

I am interested in estimating the survival function, hazard function, making predictions, and etcetera, like I would do with the Cox model or a parametric model, using machine learning methods (in general, as I am mainly interested in learning machine learning in this context rather than a specific method). However, most references I have come across are a bit cryptic, only handwaving how the method is used, but not really explaining the calculations behind their proposals.

I would appreciate any references on detailed applications of machine learning to survival data, with explanations of methods that I can reproduce in a programming language (R, Python), even if it takes some effort (again, my goal is mostly self learning the magic behind the black box). I am not interested in packages that allow me to run black boxes.

Answer 1

Firstly, it seems like for some reason survival analysis using ML techniques does not seem to get done much. That seems weird, when you think about all the possible applications like, say, time to customer churn, time to equipment failure and so on, which would seem like interesting topics.

One of the more "obvious" approaches that I've seen used a few times is to use any neural network that you might use for such a tasks, if the goal were e.g. classification or regression, but to instead of outputting probabilities for each category or a single regression estimate to output the parameters of a survival distribution. I'm not sure I've seen this published, although it feels like someone must have done so and I probably just have not seen it. The datasets I've used it on are proprietary, so I don't have any nice walk-throughs.

This, obviously, does not really (or at least it's not obvious to me how) work for Cox regression/Kaplan-Meier type estimates due to their semi-parametric nature. However, this would e.g. work quite well for the Weibull distribution, you would have two outputs from the neural network that you would treat as the log-scale $\log \lambda$ and log-shape $\log k$ parameters. You would then use something like $$- \log k + \log \lambda - (k-1) ( \log x - \log \lambda) + (x/\lambda)^k$$ as your loss function for observations with an event (minus the log of the probability density function) and $$ (x/\lambda)^{k} $$ for observations that are right censored (minus the log of the complementary cumulative distribution function). It may very well be that the various re-parameterizations used in the statistical literature / in statistical software would turn out useful here, too.

Defining a neural network with two outputs is easy enough with e.g. the keras, PyTorch or fastai (on top of PyTorch) libraries in python, or the keras or torch packages in R. For what architecture to use before that, you can take inspiration from what people are generally doing for tabular data (assuming your inputs are tabular) including e.g. the (prediction) model architecture used in fastai (see e.g. the code accessible via the documentation or the excellent book using the library) or Dragonnet if you want to look into causal inference problems. Thus, the main thing you need to do to get this to work is to code up the custom loss function (I'm not aware that any libraries offer this as a default).

Answer 2

The vignettes provided with the R survival package are a good place to start. The main survival package vignette succinctly explains the principles of many different flavors of survival analysis. Examples are illustrated with the package's functions (whose open-source code is available for inspection). It's succinct, but most of the important issues are covered. The time dependence vignette addresses time-varying covariate values and proportional hazards assumptions in Cox models. Other vignettes go into more detail with issues like multi-state survival and competing risks. The presentation necessarily emphasizes the regression modeling approach used in that package, but the principles apply to other machine-learning methods.

The python lifelines documentation also gives a useful overview of the principles of survival modeling. Its organization might be more in line with the step-by-step approach you seek than the R survival vignettes.

References to book-length treatments are described on this page.

In response to comment:

These references are admittedly regression-centric. If you learn the fundamental principles of survival analysis from these references, however, it's a very straightforward extension to what you consider to be "machine-learning" methods. As much of survival analysis was developed in the regression context, explanations from that context will typically have a strong grounding in basic principles.

The main principle you have to learn is how to set up the proper objective function in situations where you need to deal with censoring or truncation of event times. For survival analysis, this means maximizing the likelihood (for parametric models) or partial likelihood (for semi-parametric models), or minimizing the (partial) deviance as a loss function. The formulas are fairly straightforward, but knowing when to apply each of them isn't always obvious. As survival modeling is based on covariate values in place among cases at risk at event times, there are also issues in the proper handling and coding of covariates whose values change over time, including how to avoid the dangers of survivorship bias.

Once that's mastered in the regression context, it's just a matter of using the machine-learning approach of choice to minimize the loss function. Approach this like An Introduction to Statistical Learning, The Elements of Statistical Learning, or Statistical Learning with Sparsity and similar texts approach things: start with regression models to develop the principles, then use random forests, boosted trees, neural nets, etc., to do the same optimization in more flexible machine-learning ways. The CRAN survival task view has links to many such implementations of survival modeling, often with vignettes that illustrate the particulars.

Answer 3

Check out scikit-survival (python).

It has excellent tutorial notebooks and support for linear and tree-based algorithms for survival data.

User guide: https://scikit-survival.readthedocs.io/en/stable/user_guide/index.html

Answer 4

A straightforward way of using any ML in a survival analysis framework is simply to do Discrete time survival analysis.

Take any model that outputs ML probabilities. Choose a time period, eg week. turn your dataset into weeks eg if you have customer churn, for every customer create a row for every week the customer had not yet churned ( including the last week when the customer churned) - label the churn event as 1

predict probability of churning for each customer and week using your choice of ml model and inputs (eg customer properties, time of year etc)

you can easily extend this to a multistate model using multinomial models, eg for defaults people can be 1 period in arrears, 2 period in arrears and then defaulted. encode the corresponding 3 categories as your target, and encode the corresponding states of the customers at each week.

(https://stats.idre.ucla.edu/r/faq/how-can-i-convert-from-person-level-to-person-period/)

https://stats.idre.ucla.edu/other/examples/alda/ gives code in r for performing these analyses (based on book Applied Longitudinal Data Analysis, written by Judith D. Singer and John B. Willett,)