Models & Math · ML Systems Lab

Survival Analysis: Kaplan-Meier, Cox Proportional Hazards, and Why Censoring Matters

Churn modelling, LTV prediction, clinical trial analysis, and time-to-failure modelling all share a common structure: you are modelling the time until an event occurs, and some observations never experience the event during the observation window. Standard regression ignores this censoring and produces biased estimates. Survival analysis handles it correctly.

Survival analysis is the statistical framework for modelling time-to-event data where some observations are censored — the event has not occurred by the end of the observation period. Using standard regression in these settings produces systematically biased estimates. Every DS team working on subscription products, healthcare, or reliability engineering needs this toolkit.

The censoring problem

Consider modelling customer churn (time until a customer cancels). Some customers cancel during your observation window — you observe their exact churn time. Others are still active at the end of your observation window — you do not know when (or if) they will churn. These are right-censored observations. If you simply exclude censored customers, you introduce survivorship bias — the model only learns from churned customers, ignoring the information that uncensored customers survived at least until the observation end. If you treat the observation end date as the churn date, you underestimate survival times. Survival analysis uses all observations correctly: censored customers contribute information up to their censoring time.

Key functions

Survival function S(t) = P(T > t): the probability that the event occurs after time t. S(0) = 1, S(∞) = 0, S is monotone non-increasing. For a churn model, S(t) is the probability a customer is still active at time t. Hazard function h(t) = lim_{Δt→0} P(t ≤ T < t+Δt | T ≥ t) / Δt: the instantaneous rate of the event occurring at time t, given survival until t. Not a probability — it is a rate (can exceed 1). Cumulative hazard H(t) = ∫₀ᵗ h(s) ds. Relationship: S(t) = exp(-H(t)). Given the survival function, you can compute the hazard and vice versa.

Kaplan-Meier estimator

Kaplan-Meier is a non-parametric estimator of S(t). At each observed event time tⱼ: S(tⱼ) = S(tⱼ₋₁) × (1 - dⱼ/nⱼ), where dⱼ is the number of events at tⱼ and nⱼ is the number at risk (still alive and not censored) just before tⱼ. Between event times, S(t) is constant (step function). The KM estimator is the maximum likelihood estimator of S(t) for censored data with no parametric assumptions. Censored observations "leave" the risk set at their censoring time but are included in all earlier risk sets — this is how they contribute information. KM produces a survival curve, not a model with covariates.

Log-rank test: comparing survival curves

To test whether two groups (e.g., treatment vs control in a clinical trial) have the same survival distribution, the log-rank test computes: χ² = (Σⱼ (O₁ⱼ - E₁ⱼ))² / Σⱼ V_{1j}, where O₁ⱼ is the observed events in group 1 at time tⱼ and E₁ⱼ is the expected events under H₀ (no group difference). The log-rank test is most powerful when the hazard ratio is constant over time (proportional hazards assumption). For non-proportional hazards, use the weighted log-rank test or compare restricted mean survival time.

Cox proportional hazards model

Cox (1972): a semi-parametric regression model for survival data. The hazard for individual i with covariates xᵢ: h(t|xᵢ) = h₀(t) exp(βᵀxᵢ). h₀(t) is the baseline hazard — a non-parametric function of time shared by all individuals. exp(βᵀxᵢ) is the covariate effect, which is time-constant (the proportional hazards assumption). "Semi-parametric" because h₀(t) is left completely unspecified. Cox fits β without ever estimating h₀(t), using the partial likelihood: L(β) = Πⱼ exp(βᵀxⱼ) / [Σ_{i∈R(tⱼ)} exp(βᵀxᵢ)], where R(tⱼ) is the risk set at event time tⱼ. The denominator sums over everyone still at risk — censored observations contribute to the denominator of every event time before their censoring.

Interpreting Cox coefficients

exp(β_j) is the hazard ratio: the multiplicative change in hazard for a one-unit increase in feature xⱼ, holding all others constant. exp(β_j) > 1: feature increases the hazard (risk factor). exp(β_j) < 1: feature decreases the hazard (protective factor). Example in churn: if β_tenure = -0.3, then exp(-0.3) ≈ 0.74 — each additional month of tenure reduces the hazard of churn by 26%.

Proportional hazards assumption

The Cox model assumes that the hazard ratio between any two individuals is constant over time. To test: plot log(-log(S(t))) vs log(t) for different groups — parallel lines indicate proportional hazards. Or compute Schoenfeld residuals and test their correlation with time. When violated: use time-varying covariates, stratified Cox (separate h₀(t) per stratum), or flexible parametric models (piecewise exponential, flexible splines).

ML approaches to survival

Gradient boosted survival models: XGBoost and LightGBM support survival objectives (accelerated failure time or Cox partial likelihood). Random survival forests (Ishwaran et al., 2008): ensemble of survival trees that output cumulative hazard functions. DeepSurv (Katzman et al., 2018): neural network that outputs β parameters for a Cox model. All preserve the censoring mechanism while using ML's flexibility.

Interview questions on this topic

"Why can't you just use linear regression to model time-to-churn?" — Censored observations: customers still active at the observation end don't have a true churn time — you'd have to either exclude them (bias) or use the end date as churn time (underestimates survival). Survival analysis correctly uses all observations, including censored ones, through the likelihood formulation.

"Interpret a Cox coefficient of β = 0.5 on a binary feature 'is_high_usage'." — exp(0.5) ≈ 1.65. High-usage customers have a 65% higher hazard rate at any given time than low-usage customers. Put differently: at any point in time, high-usage customers are 1.65× more likely to experience the event (churn, death, failure) in the next instant, conditional on having survived until that point.

"What is the difference between KM and Cox? When would you use each?" — KM is non-parametric: it estimates the survival curve for the full population or pre-defined subgroups. No covariate modelling. Use KM for: descriptive analysis, comparing two groups, clinical reporting. Cox adds covariate regression: it models how individual features affect the hazard. Use Cox for: understanding which features predict the event, adjusting for confounders, predicting individual risk.

"What does the proportional hazards assumption mean and how do you test it?" — The ratio of hazards between any two individuals is constant over time: h(t|xᵢ)/h(t|xⱼ) = exp(β(xᵢ - xⱼ)) for all t. Test: Schoenfeld residual test (significant correlation of residuals with time indicates violation). Graphically: log-log survival curves should be parallel. Common violation: treatment effect diminishes over time (common in clinical trials).

Try on Colab: use the lifelines library with the Telco customer churn dataset (augment with synthetic survival times). Fit a Kaplan-Meier curve for churned vs non-churned. Run a log-rank test comparing high vs low tenure groups. Fit a Cox model with 5 covariates (tenure, monthly charges, contract type, payment method, internet service). Interpret the hazard ratios. Check the proportional hazards assumption using Schoenfeld residuals.

Continue interactively
Read this post inside ML Systems Lab — with Simplify toggle, interview Q&As, inline glossary, and the MLE Path forward pointer.
Open in MSL →