Online Survival Analysis: A Bandit Approach under Cox PH Model

1. Problem Statement

The paper addresses the challenge of performing survival analysis in an online, sequential decision-making setting (specifically within the Contextual Bandit framework).

• Context: Traditional survival analysis (e.g., Cox Proportional Hazards model) is typically conducted in an offline manner, where researchers wait for all data to be collected before analyzing time-to-event outcomes. This is inefficient for applications like clinical trials, precision medicine, or customer retention, where decisions need to be updated dynamically as new data arrives.
• Core Challenges: Integrating survival analysis into bandits introduces three unique complexities not present in standard bandit problems:
  1. Staggered Entry: Subjects enter the study at different calendar times, meaning the "risk set" (the group of subjects currently at risk of an event) changes dynamically based on entry times and survival durations.
  2. Right Censoring: The outcome (event time) is often unobserved because the study ends or the subject is lost to follow-up before the event occurs. This creates incomplete feedback.
  3. Delayed Feedback: The "reward" (observed survival time) is only available after a delay determined by the event or censoring time, complicating the immediate update of the policy.
• Goal: To develop an algorithm that sequentially selects treatments (actions) to maximize a survival metric (e.g., survival probability at a fixed time $\tau_0$) while learning the underlying Cox Proportional Hazards (Cox PH) model parameters, balancing exploration (trying new treatments to learn) and exploitation (using current knowledge to treat patients optimally).

2. Methodology

The authors propose a unified framework that adapts the Cox PH model to an online bandit setting.

A. Theoretical Framework: Online Cox PH with Staggered Entry

• Risk Set Definition: Unlike offline settings where the risk set is static or defined by a single time scale, the authors define the risk set $R(\tau, s)$ using two time scales: calendar time ($\tau$) and survival time ($s$). A subject is at risk if they have entered the study and survived up to time $s$.
• Incremental Likelihood Update: Instead of refitting the Cox model from scratch at every round (which is computationally expensive, $O(N^2)$), the authors derive an incremental update for the partial likelihood function.
  • They decompose the change in the log-likelihood between round $t$ and $t+1$ into two parts:
    1. $P_1$: Contribution from newly observed events.
    2. $P_2$: Correction for changes in the risk set due to staggered entry (subjects entering the risk set as time passes).
  • This reduces computational complexity from $O(N^2)$ to roughly $O(N \cdot m)$, where $m$ is the number of observed events.

B. Algorithm Adaptation

The paper adapts three canonical bandit algorithms to this survival context:

1. Epsilon-Greedy (EG): Selects the optimal action based on the current Cox estimate with probability $1-\epsilon$, and a random action with probability $\epsilon$.
2. Upper Confidence Bound (UCB): Constructs a confidence ellipsoid around the estimated coefficient vector $\hat{\beta}$. The action is chosen to maximize the upper bound of the survival probability, balancing the estimated mean and the uncertainty (variance).
3. Thompson Sampling (TS): Approximates the posterior distribution of $\beta$ (using Laplace approximation or MCMC) and samples a parameter vector to determine the action.

C. Theoretical Guarantees

• Martingale Challenge: Standard martingale concentration tools fail in this setting because the score function is not a martingale under staggered entry.
• Solution: The authors define an adjusted score process $Q(\tau)$ by replacing the time-dependent average covariate $\bar{X}(\tau, s)$ with its asymptotic limit $\mu(s)$. This adjusted process is a martingale.
• Confidence Ellipsoid: They establish the first confidence ellipsoid for online survival settings, bounding the estimation error $\|\hat{\beta}_t - \beta\|_{I_t}$.
• Regret Bound: They prove a sublinear regret bound of $\tilde{O}(d\sqrt{T})$, where $d$ is the feature dimension and $T$ is the time horizon. This matches the rate of classical linear contextual bandits despite the added complexity of censoring and staggered entry.

3. Key Contributions

1. First Online Bandit Framework for Survival Analysis: The paper is the first to integrate the Cox PH model into a purely online learning setting that simultaneously handles right censoring, staggered entry, and delayed feedback.
2. Efficient Online Updates: They provide a computationally efficient method to update the Cox partial likelihood incrementally, avoiding the need to refit the model on the entire dataset at each step.
3. Theoretical Innovation: They overcome the breakdown of standard martingale properties in staggered entry settings by constructing an adjusted score process and deriving a novel confidence ellipsoid, leading to rigorous sublinear regret guarantees.
4. Algorithm Adaptation: They successfully extend EG, UCB, and TS strategies to the survival context, providing theoretical bounds for all three.

4. Results

A. Simulations

• Setup: Synthetic data generated with known Cox PH parameters, varying censoring rates, and staggered entry.
• Findings:
  • All three algorithms (EG, UCB, TS) converge rapidly to the true parameter $\beta$ and the optimal policy.
  • The proposed methods significantly outperform a "naive" baseline that ignores censoring (which discards unobserved data), demonstrating the importance of handling right-censored data correctly.
  • Sensitivity analysis shows robustness to varying censoring rates and model misspecification (e.g., when the true data follows an AFT model or a piecewise hazard model).

B. Semi-Real Experiments (SEER Cancer Data)

• Dataset: Used the SEER (Surveillance, Epidemiology, and End Results) database for malignant breast cancer patients (2004–2015).
• Task: Decide between two treatments: Breast-Conserving Surgery (BCS) vs. Mastectomy to maximize survival probability at specific time points ($\tau_0 = 25, 50, 75, 100$ months).
• Performance:
  • The bandit algorithms quickly learned the underlying survival patterns and converged toward the optimal survival probability.
  • Even when the "ground truth" was generated using a Random Survival Forest (RSF) to test robustness against Cox PH model misspecification, the algorithms continued to improve average survival probability over time, demonstrating practical stability.

5. Significance and Impact

• Clinical and Industrial Applications: This work enables adaptive clinical trials and precision medicine where treatment recommendations can be updated in real-time as patient outcomes become available, rather than waiting for the end of a study. This can significantly shorten feedback loops and improve patient outcomes.
• Efficiency: By handling staggered entry and censoring online, the approach reduces the time and cost associated with traditional offline survival analysis.
• Theoretical Foundation: It bridges the gap between survival analysis and reinforcement learning, providing a rigorous mathematical foundation for decision-making under uncertainty with censored time-to-event data.
• Scalability: The incremental update mechanism makes the approach scalable to large datasets, which is crucial for modern electronic health records and large-scale observational studies.

In summary, this paper presents a robust, theoretically grounded, and computationally efficient framework for making sequential treatment decisions in survival analysis, addressing critical real-world challenges like delayed and incomplete data.