Efficient Learning of Survival Analysis Models Under Outcome-Dependent Truncation and Censoring
Authors: Itai Zilberstein, Ioannis Anagnostides, Arman Kilic, Tuomas Sandholm
Organizations: Department of Computer Science, Carnegie Mellon University, Pittsburgh, PA · Department of Surgery, Division of Cardiothoracic Surgery, Medical University of South Carolina, Charleston, SC · Additional affiliations: Strategy Robot, Inc., Strategic Machine, Inc., Optimized Markets, Inc.
Modern machine learning models are typically trained on data assumed to be representative of the predictive tasks the model will face in the future, but practical datasets violate this assumption. We focus on outcome-dependent truncation in which an observation appears in the training data depending on its label. Our work connects two largely distinct literatures: efficient learning under truncation and survival analysis. Existing methods for truncated survival analysis do not characterize the optimization landscape and instead rely on iterative procedures such as expectation maximization. We address this gap and prove that Cox regression can be efficiently learned under truncated and right-censored data using first-order optimization when the baseline hazard is known (despite nonconvexity). In contrast, when the baseline hazard is learned jointly, we construct identifiable and correctly specified models whose negative log-likelihood contains strict suboptimal local minima. On experiments with four survival datasets, we show that accounting for truncation reduces the prediction bias induced by naive Cox regression and extends to neural networks. We also show orders of magnitude speedups compared to standard expectation maximization.
Figures & tables
Figure 1: Illustrative example of learning a Cox model with truncated samples. When the likelihood accounts for truncation, SGD converges to the true regression coefficients, whereas the standard Cox likelihood converges to biased parameters (left). The resulting parameter bias can translate into substantial errors in predicted graft survival for individual patients (right).
Figure 2: Population log-likelihood from Example 2 .
Figure 3: Results of learning standard and truncated Cox regression and neural networks with first-order optimizers across datasets when the data-generating model is a Cox model. Shaded regions show standard deviation.
METABRIC
SUPPORT
FLCHAIN
HEARTS
Method
Time (s)
MSE
Time (s)
MSE
Time (s)
MSE
Time (s)
MSE
Adam (Ours)
0.79±0.38
0.0542±0.0605
1.11±0.42
0.0240±0.0263
1.13±0.77
0.0625±0.0596
2.05±1.32
0.0271±0.0197
Projected SGD (Ours)
0.56±0.29
0.0446±0.0474
0.86±0.33
0.0231±0.0242
1.59±1.56
0.0545±0.0433
5.15±3.58
0.0316±0.0327
EM
0.16±0.10
0.0462±0.0555
0.70±0.52
0.0240±0.0277
1.40±0.55
0.0877±0.0822
34.29±13.65
0.0503±0.0475
Table 1: Computational efficiency of three optimizers for the same truncated Cox regression MLE across datasets with 10 independent runs per truncation level. Time is the mean wall-clock time to convergence (s) and MSE is the mean test risk-score error.
Method
Time (s)
MSE
Adam (Ours)
25.4±20.3
0.0281±0.0283
Projected SGD (Ours)
22.0±13.4
0.0276±0.0276
EM
2607.1±1210.6
0.0272±0.0271
Table 2: Average optimizer efficiency on HEARTS-Large across truncation levels.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Dataset
N
Covariates
Censoring rate
Median event time
HEARTS
59,570
95
49.8%
5.7 years
SUPPORT
8,873
14
32.0%
231 days
METABRIC
1,903
9
42.0%
115 months
FLCHAIN
7,871
8
72.5%
4,303 days
Appendix
Table 3: Details of datasets used in the semi-synthetic experiments.
Figure 4: Results of learning standard and truncated Cox regression and neural networks with first-order optimizers across datasets when the data-generating model is a neural network.
Survival Analysis (SA) is a statistical framework that models the time span until some event of interest occurs. Widely used in several domains, including healthcare and churn prediction, a central challenge in its applicability stems from the time of the event being partially observed or right-censoring. Tabular Foundation Models (TFM) have attracted significant interest in recent years due to their ability to perform prediction tasks in a single forward pass, requiring no dataset-specific parameter fitting. Despite their success, their application to prediction tasks on time-to-event data remains difficult due to right censoring. In this work, we present a training-free method to survival regression by leveraging TFMs to both predict the time of the event and iteratively impute right-censored data. Our method uses a TFM to construct an Accelerated Failure Time (AFT) model requiring no training beyond fitting a single scalar parameter. Subsequently, by building on the Buckley-James estimator, we introduce a non-parametric in-context estimator for right-censored data. Our experiments on standard survival analysis benchmarks show that our method is competitive with several parametric and semi-parametric survival regression models that require training, including Cox regression and parametric AFT models.
Survival analysis is a widely used statistical framework for modeling time-to-event data under censoring. Classical methods, such as the Cox proportional hazards (Cox PH) model, offer a semiparametric approach to estimating the effects of covariates on the hazard function. Despite its importance, survival analysis has been largely unexplored in online settings, particularly within the bandit framework, where decisions must be made sequentially to optimize treatments as new data arrive over time. In this work, we take an initial step toward integrating survival analysis into a purely online learning setting under the Cox PH model, addressing key challenges including staggered entry, delayed feedback, and right censoring. We adapt three canonical bandit algorithms to balance exploration and exploitation, with theoretical guarantees of sublinear regret bounds. Extensive simulations and semi-real experiments using SEER cancer data demonstrate that our approach enables rapid and effective learning of near-optimal treatment policies.
Yang Xu, Wenbin Lu, Rui Song
Department of Statistics · North Carolina State University
We revisit Cox's proportional hazards model to improve variable selection in survival analysis. A square-root transformation of the partial likelihood renders the selection of the regularization parameter pivotal, free of the unknown baseline hazard and censoring mechanism. The resulting criterion borrows from information criteria such as BIC and from penalized regression methods such as the lasso, taking the best of both. On simulated and real data, our method substantially improves upon state-of-the-art approaches used daily in support recovery.