Scalable Random Survival Forests via Risk Set Sampling

Technical Summary: Scalable Random Survival Forests via Risk Set Sampling

Problem Statement
Random Survival Forests (RSF) are a powerful machine learning extension of random forests designed for censored time-to-event data. However, their application to large or high-dimensional datasets is often computationally prohibitive. The primary bottleneck arises from the repeated evaluation of "risk sets" required to compute the log-rank splitting rule during tree construction. For every event time, the algorithm must evaluate all individuals still at risk, causing computational costs to scale rapidly with sample size and the number of distinct event times. While risk set sampling has been successfully integrated into Cox-based neural networks (e.g., Cox-Time) and is a standard design in classical epidemiology (nested case-control), it has not yet been adapted for RSF.

Methodology
The authors propose a modification to the RSF algorithm that integrates risk set sampling directly into the tree-building process.

• Core Mechanism: Instead of utilizing the full risk set ($R_k$) at each event time $t_k$ to calculate the log-rank splitting statistic, the method samples a fixed fraction of individuals from the risk set ($\tilde{R}_k$).
• Statistical Adjustment: The log-rank statistic relies on the expected number of events and their variance within a node. The authors modify these calculations to use the sampled risk set $\tilde{R}_k$ rather than the full set. Specifically, the expected number of events $E(D_{\alpha,k})$ and the variance $Var(D_{\alpha,k})$ for a daughter node $\alpha$ are recalculated using the ratio of the sampled at-risk individuals ($\tilde{R}_{\alpha,k}$) to the total sampled risk set ($\tilde{R}_k$).
• Implementation: The approach was implemented in the ranger R package, a high-performance implementation of random forests. The authors tested sampling fractions ranging from 10% to 100% (where 100% represents the standard full-risk set).
• Experimental Design:
  • Simulations: Heterogeneous survival data were generated using Exponential and Weibull mixture distributions with sample sizes of $n=500, 1000,$ and $2000$ and $p=15$ covariates.
  • Real-World Data: Two datasets were utilized: METABRIC (1904 breast cancer observations) and the German Breast Cancer Study Group (GBSG) dataset (686 observations).
  • Metrics: Performance was evaluated using the Concordance Index (C-index) via 5-fold cross-validation and out-of-bag (OOB) estimates, alongside computational runtime.

Key Results

• Predictive Performance: Across all simulated and real-world datasets, RSF models trained on sampled risk sets (specifically using 25% to 40% of the risk set) achieved C-index values virtually indistinguishable from models trained on full risk sets.
• Computational Efficiency: Reducing the risk set to the 25–40% range resulted in a reduction of training time by more than 50%. The runtime increased monotonically with the risk-set fraction, with the most significant gains observed in larger datasets.
• Stability and Bias: The study found that risk set sampling primarily affects the variance of the estimates rather than introducing systematic bias. Confidence intervals for the C-index showed substantial overlap across all sampling fractions, indicating no statistically meaningful difference in predictive discrimination.
• Regularization Effect: In several scenarios, models trained with moderate sampling fractions (25–40%) exhibited marginally higher C-index values than full-risk models, suggesting a potential beneficial regularization effect that reduces noise.

Significance and Claims
The paper claims that risk set sampling is a simple yet effective modification that substantially improves the scalability of Random Survival Forests without compromising predictive performance or model stability. The authors argue that:

1. Bridging Disciplines: The approach provides a methodological bridge between classical epidemiologic sampling designs (nested case-control) and modern machine learning for survival analysis.
2. Practical Utility: Full-risk computation often offers little practical advantage for achieving good predictive performance in RSF modeling while incurring significantly higher computational costs.
3. Accessibility: By implementing this methodology in the widely used ranger package, the authors facilitate the application of scalable RSF in practice, particularly for large-scale and high-dimensional survival data.

The authors conclude that while the method is robust, future work could explore its application to more complex frameworks (competing risks, multi-state models) and investigate interactions with other RSF hyperparameters.