Time-to-Event Prediction with Neural Networks and Cox Regression

In plain words

Survival models answer "when will this event happen?" for people or machines, and they can use records where the event has not happened yet because observation stopped early. The classic tool for this is the Cox model, which assumes that each covariate multiplies the risk by a fixed factor over time. That assumption is restrictive, and the standard way to fit the model does not scale well to very large data sets.

The paper makes two changes. First, the part of the model that turns covariates into a risk score becomes a neural network. Second, the training objective is replaced by a cheaper version that compares each event with one randomly chosen person who was still at risk at that moment. This version can be trained in mini-batches, and it borrows an old sampling design from epidemiology.

A further change lets the network see the time as an input. The risk then no longer has to be proportional over time, and the model can capture how covariates and time interact. The authors test these models on simulated data and on five real survival data sets, and they apply them to a large subscription-churn data set.

TL;DR

• Cox model with a neural risk term, trained by mini-batch SGD on a sampled-control loss §3.2 eq. 9 eq. 8.
• Cox-Time feeds time to the network, so proportional hazards no longer hold §3.4 eq. 11 eq. 12.
• Author-reported: Cox-Time has the best integrated Brier score (IBS) and integrated binomial log-likelihood (IBLL) on the KKBox churn data tab. 7, and the best IBS and IBLL on three of the four small data sets tab. 3 tab. 4. DeepHit has the best concordance on KKBox and most small sets tab. 2 tab. 7.
• The loss is an approximation to the Cox partial likelihood. Simulations show that it stays close to the exact likelihood §5 fig. 3.
• Prediction for Cox-Time is costly: each survival curve needs a network call for every distinct event time up to the horizon §3.5.

Problem

The Cox partial likelihood sums over risk sets, which can contain almost the whole data set. Newton-type fitting therefore does not batch, and mini-batch SGD needs a loss that can be computed on a batch §3.2. Earlier neural Cox models either failed to beat linear Cox regression §2 or kept the proportional-hazards restriction §2 §3.4. Discrete-time neural methods such as DeepHit discretise time and need a time grid §2 app. D.

The paper asks for a Cox-type model that (a) uses a neural risk term, (b) trains on mini-batches, (c) can drop proportional hazards, and (d) gives survival curves in continuous time §1 §3.

Method

Notation. Individual i has covariates x_i, observed time T_i = min(T*_i, C*_i) and event indicator D_i. The risk set R_i holds everyone with T_j >= T_i §3. The model is a relative-risk model h(t | x) = h0(t) exp[g(x)] eq. 3. The paper's two new forms are g(x) as a neural network (Cox-MLP) §3.3 and g(t, x) with time as an input (Cox-Time) eq. 11.

Exact losses, for reference. The negative Cox partial log-likelihood is the loss in eq. 5. The Breslow estimator gives the cumulative baseline hazard eq. 6. Both are in §3.1.

Sampled-control loss (Cox-MLP (CC)). For each event i, one control j is drawn from the risk set, and the loss is the log of 1 + exp(g(x_j) - g(x_i)) eq. 9, averaged over events eq. 8. Controls are sampled with replacement and a new control is drawn in each iteration §3.2. The authors note that a sampled-set weight does not affect the gradient and can be dropped §3.2.

Penalty. The loss has no unique minimiser for g, so a penalty on the sampled outputs is added. It is lambda times the sum of the absolute values of g over the sampled sets eq. 10 §3.3.

Cox-Time. The network receives (t, x). Each event i is paired with a control j drawn from R_i, and both are evaluated at the event time T_i. The loss takes the same form as above eq. 12. Its cost is O(n) per epoch, against O(n^2) for the full risk set, because the sampled set has fixed size §3.4.

Breslow estimator for Cox-Time. The increment at each distinct event time uses the full risk set, not a sample eq. 13. The cost is quadratic in the number of individuals §3.4 §3.5.

Prediction. S(t | x) = exp(-H(t | x)) eq. 1. For proportional models H factorises into a baseline and exp[g(x)] §3.5. For Cox-Time the cumulative hazard is a sum over event times up to t, each with a network call §3.5 eq. 13. The paper approximates this by using a grid of distinct times, and it estimates the baseline on a random subset of the training data to control the grid size §3.5.

Steps (Cox-Time training). Pseudocode below follows the paper's equations. The official code differs in some details; see "Paper vs code".

# Cox-Time training with one sampled control per event  [§3.2] [§3.4]
# Input:  covariates (N, d); durations (N,); events (N,) in {0, 1}
# Output: trained network g(t, x)

def train_cox_time(net, covariates, durations, events, lam, batch_size):
    # Cases are the individuals whose event was observed  [§3.2]
    case_idx = indices where events == 1
    for each mini-batch of case_idx:
        # Draw one control per case from its risk set, with replacement  [§3.2]
        # Risk set: everyone whose duration is >= the case's duration [§3.1]
        control_idx = sample_from_risk_set(durations, batch)
        # Both members of a pair are evaluated at the case's own time  [§3.4]
        time = durations[batch]
        g_case = net(time, covariates[batch])
        g_control = net(time, covariates[control_idx])
        # Pair loss log(1 + exp(g_control - g_case)), averaged  [eq. 9] [eq. 12]
        loss = mean(log(1 + exp(g_control - g_case)))
        # Penalty on the sampled outputs, pulls them toward zero  [§3.3] [eq. 10]
        penalty = lam * (sum(abs(g_case)) + sum(abs(g_control)))
        update(net, loss + penalty)

# Breslow increments for a time-dependent relative-risk model  [§3.5] [eq. 13]
# Input:  net; covariates (N, d); durations (N,); events (N,)
# Output: one increment per distinct event time t_k

def breslow_increments(net, covariates, durations, events):
    for each distinct event time t_k:
        # Number of events at t_k  [§3.1] [eq. 6]
        num_events = count of i with durations[i] == t_k and events[i] == 1
        # Denominator sums over the full risk set, not a sample  [§3.5] [eq. 13]
        at_risk = indices j with durations[j] >= t_k
        denominator = sum(exp(net(t_k, covariates[at_risk])))
        increment[k] = num_events / denominator
    return increment

# Survival curve for one covariate vector x  [§3.5] [eq. 13] [eq. 1]
# Input:  net; x; event_times (sorted ascending); increments (same length)
# Output: S_hat(t_k | x) for each t_k

def predict_survival(net, x, event_times, increments):
    cumulative = 0.0
    for each k in event_times:
        # Add the increment weighted by the relative risk at t_k  [eq. 13]
        cumulative += increments[k] * exp(net(event_times[k], x))
        # Survival from the cumulative hazard  [eq. 1]
        survival[k] = exp(-cumulative)
    return survival

Key equations, verbatim where the exact form matters:

\text{loss} = \frac{1}{n}\sum_{i: D_i=1} \log \left (\sum_{j \in \tilde{\mathcal{R}}_i} \exp[g(\mathbf{x}_j) - g(\mathbf{x}_i)]\right ) eq. 8

\text{loss} = \frac{1}{n}\sum_{i: D_i=1} \log \left ( 1 + \exp[g(\mathbf{x}_j) - g(\mathbf{x}_i)]\right ), \quad j \in \mathcal{R}_i \backslash \{i\} eq. 9

Here n is the number of events and the sampled set is written with a tilde in the paper's notation §3.2.

Hyperparameters

Small data sets (SUPPORT, METABRIC, Rot. & GBSG, FLCHAIN), random search over the space in app. A.1:

KKBox case study, grid search over app. A.2 (no weight decay §6.2.1):

Simulation settings:

Results

Author-reported. Setup: four small survival data sets from the DeepSurv package, plus the KKBox churn data (2017 Kaggle challenge data, subscription-level survival times) §6.1 §6.2. Small-set scores are means over five cross-validation folds §6.1.2. KKBox scores are medians over five fits §6.2.2.

Data set sizes, small sets tab. 1:

The paper replaces DeepSurv's WHAS set, which contains duplicated individuals, with FLCHAIN and removes rows with missing values and the cause-of-death covariate §6.1. Note: the FLCHAIN columns in tab. 1 are the paper's values; they do not state whether the 6,524 rows are before or after missing-value removal §6.1.

KKBox: 15 covariates (7 numerical, 8 categorical) app. B.1. Split by tab. 5:

Metrics (defined in §4): time-dependent concordance C^td eq. 14 (with tie handling from Ishwaran et al., Section 5.1 step 3 §4.1); IBS eq. 15 eq. 16; IBLL eq. 17. The paper reports IBS and IBLL weighted by a Kaplan-Meier estimate of the censoring distribution tab. 3 tab. 4. The paper's text says lower IBS is better and closer to zero is better for IBLL §6.1.2.

Small data sets, C^td, five-fold mean tab. 2:

Small data sets, integrated Brier score (IBS), censoring-weighted tab. 3:

Small data sets, integrated binomial log-likelihood (IBLL), censoring-weighted tab. 4:

Bold entries in the paper mark the best value per column; this card does not repeat them. The paper's own reading: Cox-Time is generally best on IBS and IBLL, DeepHit best on C^td but with the worst calibration on three of four sets, and the proportional-hazards methods are close to each other §6.1.2.

KKBox churn case study, medians of five fits tab. 7:

Brier score by duration fig. 4: DeepHit is poor at all durations. Cox-MLP (CC) is unstable at short durations. Cox-Time is clearly good at all durations. Cox-MLP (DeepSurv) is the best model at longer durations. The IBS table already reports these tab. 7.

Simulations §5 fig. 1 fig. 2 fig. 3:

• Loss-to-likelihood fit: more sampled controls give lower training error, but the convergence speed does not change; computation grows with the number of controls fig. 1 §5.
• Coefficient agreement with classical Cox: small samples (100) give slightly smaller Cox-SGD estimates; the gap shrinks with sample size fig. 2.
• Full-likelihood agreement: with one control, the mean log-likelihood gap is about 0.1% for the smallest sets (classical Cox MLL about -2.2) §5. The gap shrinks with n fig. 3.
• Non-linear and non-proportional simulations: Cox-MLP and Cox-Time estimate the true partial likelihoods and survival curves better than the linear, proportional models app. C.1 app. C.2.

Survival-curve clustering, KKBox §6.2.3 fig. 5: K-means with 10 clusters on Cox-Time survival curves. Reported: 19% of the test set in a flat-slope cluster, 18% at high immediate risk, and a 3% cluster with a sharp drop near day 400 §6.2.3. The authors describe these as observations, not a validated segmentation §6.2.3.

Paper vs code

Paper: arXiv v2 (JMLR version, 20(129), 2019) [raw metadata]. Official code: the abstract and footnote link github.com/havakv/pycox, which is by an author (H. Kvamme) and is the package the paper names §1. Pinned commit 3eccdd7fd9844a060f50fdcc315659f33a2d2dc1, license BSD-2-Clause (LICENSE).

Differences found, each with citations:

• Control can be the case itself. The paper samples controls from R_i \ {i} eq. 9, but also says i is always in the sampled set §3.2. The code builds each at-risk set from the first index of a duration onward, so it contains the case data.py and samples uniformly from it with replacement data.py data.py. Consequence: a self-sampled control gives softplus(0) = log 2 and zero gradient, so it dilutes the batch mean but adds no signal.
• Penalty normalisation. The paper's penalty sums |g| over events and sampled controls eq. 10. The code adds lambda times the batch mean of |g_case| plus the batch mean of |g_control| loss.py. The two scale differently with batch size, so the paper's lambda values do not transfer directly.
• Loss normalisation. The paper divides by n, the number of events in the whole data set eq. 8. The code takes a mean over the mini-batch of event pairs loss.py. The gradient scale is therefore set per batch, not per data set.
• Cox-Time duration transform. The paper says durations may be log-transformed app. A.1 and does not mention standardisation. The code offers log(1 + t) label_transforms.py and standardises durations by default label_transforms.py. The network sees the transformed time, and the code maps predictions back to the original scale cox_time.py.
• Cox-Time network. The code's Cox-Time MLP removes the output bias and output activation and concatenates time to the covariates, with batch normalisation on by default cox_time.py. The paper does not state these details beyond "time as a regular covariate" §3.4.
• Breslow baseline. The code computes the at-risk sums over all rows with duration at or after each event time, matching eq. 13 cox_time.py. It can use a random subsample of training rows, as the paper describes cox_time.py §3.5.
• Simulation covariate formula. The paper's non-linear risk is (2/3)(x1^2 + x3^2 + x1x2 + x1x3 + x2x3) with 1-based indices app. C.1, eq. sim_non_lin_ph. The code's version is (2/3)(x0^2 + x2^2 + x0x1 + x1x2 + x1x2) relative_risk.py. The paper's x1x3 term (x0x2 in 0-based indices) appears as a second x1x2 in the code. The code is therefore not the paper's function. Use the paper's formula if you reproduce the simulation, and say so.
• Experiment settings. The paper's experiments use AdamWR with cycle restarts and 300 random configurations per fold §6.1.1. This repository has no script for the paper's experiments. Its Cox-Time example uses tt.optim.Adam with a learning-rate finder, batch size 256 and 512 epochs cox-time.ipynb cox-time.ipynb cox-time.ipynb.
• IBLL sign. The code's integrated_nbll returns the negative of the integrated binomial log-likelihood eval_surv.py, so it is positive. The paper reports IBLL as negative tab. 4 tab. 7.
• Integration. The paper approximates the integral with 100 grid points §4.2. The code applies Simpson's rule on the time grid the caller supplies ipcw.py. The grid is the caller's choice.
• Concordance variant. The paper says it modified C^td following Ishwaran et al. §4.1. The code's default is method 'adj_antolini' eval_surv.py concordance.py. The paper does not name its variant, so this card cannot confirm they match.
• DeepHit. The loss is alpha times the negative log-likelihood plus (1 - alpha) times the ranking loss, as the paper states app. D loss.py (the single-risk class). The ranking term uses the same exp((S_i - S_j)/sigma) form app. D.

No discrepancies found in the Cox partial-likelihood sign convention and the risk-set construction for Cox-Time §3.4 data.py.

Pitfalls

• Self-sampled controls. The paper states both R̃_i includes i §3.2 and that controls come from R_i \ {i} eq. 9. Pick one. Excluding the case is the cleaner reading of eq. 9. The code allows self-sampling data.py data.py; with 1/|R_i| probability per draw, this matters most on small data sets.
• Penalty scale. Eq. 10 is a sum over the sampled sets, so its size grows with batch size. A mean (as in the code) gives a different effective lambda. Keep lambda as a tuned value for your normalisation, and do not copy the paper's values across normalisations.
• Loss range. With one control and no penalty, the per-pair loss lies in (0, log 2] §3.2. The lower bound is approached as g(x_i) grows and g(x_j) falls. With the penalty, the optimum is pulled toward zero §3.3.
• Non-unique minimiser. For non-linear g, the loss has no unique minimiser §3.3. Without the penalty, outputs can drift. The paper's experiments use lambda in {0.1, 0.01, 0.001, 0}, and the KKBox runs fix 0.001 app. A.1 app. A.2.
• Cox-Time prediction cost. Each individual needs one network call per distinct event time up to the horizon, and the Breslow denominators use the full risk set §3.5. Prediction and baseline estimation are quadratic in the number of individuals unless you subsample §3.4. Subsampling changes the time grid and the baseline estimate; the paper gives no error analysis.
• Batch size in the KKBox text. The appendix says "batch size of 1028" app. A.2. The grid in app. A.1 includes 1024 and the KKBox grid has no batch size. The value is probably a typo, but the paper does not confirm it. The code's examples use 256 cox-time.ipynb.
• Unstable proportional models on KKBox. The paper reports rather unstable results for both Cox-MLP methods and attributes it to the proportionality constraint in large data sets §6.2.2. It reports medians for this reason. Use medians across several seeds.
• Concordance is not a calibration check. DeepHit has the best C^td on KKBox but the worst IBS and IBLL tab. 7. The paper treats C^td and the censoring-weighted scores as separate criteria §6.2.2.
• Censoring weights. IBS and IBLL weight by a Kaplan-Meier censoring estimate and assume censoring is independent of survival times §4.2. Violations are not discussed.
• Feature handling. FLCHAIN rows with missing values are removed, and the cause-of-death covariate is dropped §6.1. Results on FLCHAIN are not comparable to published numbers that keep those rows.
• Cox-Time on FLCHAIN. The paper says Cox-Time struggles on FLCHAIN, with the largest IBS and IBLL gap of any method there §6.1.2 tab. 3 tab. 4. It attributes this to proportionality being reasonable on that set §6.1.2.
• KKBox survival times. Times start at the first subscription or earliest transaction record, and a customer who churned and came back is treated as a new customer app. B.1 §6.2. This differs from the Kaggle task, which uses calendar dates §6.2.
• Time-varying covariates. All KKBox covariates are fixed at subscription time, which the authors call a simplification app. B.1.

Lineage

• Builds on: the Cox proportional-hazards model and the Breslow estimator §3.1 eq. 3 eq. 6; the nested case-control design of Thomas (1977) and the consistency result of Goldstein and Langholz (1992) for sampled risk sets §3.2; the time-dependent C-index of Antolini et al. §4.1; AdamWR §6.1.1.
• Earlier neural Cox: Faraggi and Simon (1995), who used an MLP in place of the linear term; the paper says it was not better than regular Cox §2.
• Compared against: DeepSurv §2 §6.1.1 (proportional, batched negative partial likelihood; the paper's Cox-MLP (DeepSurv) is its own variant, not the original code); DeepHit §2 app. D; random survival forests §2 §6.1.1; classical linear Cox §6.1.1.
• Not compared: Luck et al. (2017) and neural MTLR were not compared because they share the proportional restriction §2 §6.1.1.
• Superseded: not stated.