PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Learning of Survival Analysis Models Under Outcome-Dependent Truncation and Censoring

Itai Zilberstein, Ioannis Anagnostides, Arman Kilic, Tuomas Sandholm

arXiv:2610.12538v2cs.LG

TL;DR

Outcome-dependent truncation can bias models trained on selected observations, and existing truncated-survival methods leave the optimization landscape largely uncharacterized. This paper proves efficient projected-SGD learning for Cox regression with a known baseline hazard despite nonconvexity, while experiments report reduced prediction error and substantial speedups over EM.

  • Problem

    Truncated-survival estimation has largely focused on identifiability and asymptotics, with optimization landscapes under outcome-dependent selection remaining less characterized and practical methods often relying on costly EM.

    Source: §1 · p. 3

  • Method

    The paper analyzes truncated and right-censored Cox likelihoods using weighted truncated-exponential properties to establish projected-SGD guarantees when the baseline hazard is known.

    Source: §1 · p. 3, §3 · p. 6

  • Results

    Across four datasets, truncation-aware likelihoods reduce prediction error, extend empirically to neural networks, and yield orders-of-magnitude speedups over EM.

    Source: §6 · p. 12

  • Takeaways & Limitations

    Efficient truncation-aware survival learning is a capability the authors identify as important for scaling models to higher dimensions in high-stakes applications.

    Source: §6 · p. 12

    With an unknown baseline, distinct hazards can induce the same sampled distribution, making the baseline non-identifiable in the constructed truncated model.

    Source: §4 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

Outcome-dependent truncation selects training observations according to their labels, so models fit without accounting for selection can be biased for the underlying population. This paper studies efficient Cox-model learning under truncation and right censoring, distinguishing known from unknown baseline hazards.

1 Introduction

  • In a graft-survival example, accounting for truncation recovers true Cox coefficients, while the standard likelihood produces biased coefficients and substantial prediction errors for some patients.The example uses simple one-sided left truncation and has no censoring.Source: §1 · p. 2, Figure 1 · p. 3
  • Prior truncated-survival methods focus largely on identifiability and asymptotics, while computationally challenging procedures such as EM motivate connecting survival analysis with efficient learning under truncation.The paper considers probabilistic retention across multiple intervals, covering standard left, right, and interval truncation as well as disconnected intervals.Source: §1 · pp. 2–3
  • With a known baseline hazard, the population loss has a unique stationary point at the true parameter and projected SGD converges despite nonconvexity.The guarantee applies to outcome-dependent truncation and right-censored data.Source: §1 · p. 3
  • With an unknown baseline, guarantees need not extend: models can be non-identifiable, or identifiable and correctly specified yet have suboptimal local minima.Profiling over baseline hazards can also yield a flat, suboptimal population loss across an interval of incorrect regression parameters.Source: §1 · p. 3
  • Across four clinical datasets, truncated Cox likelihood consistently lowers prediction error as truncation rises; neural variants benefit, and first-order training is roughly 100× faster than EM.On an enlarged dataset with hundreds of thousands of observations and thousands of covariates, training time falls by roughly 100× while predictive accuracy is maintained or improved.Source: §1 · p. 4

2 Preliminaries

The preliminaries specify the observed survival data, truncation and censoring setup, and assumptions used for population results. Likelihoods are conditional on observed covariates, while hazard, survival, and inclusion probabilities refer to the pre-inclusion population.

2 Preliminaries

  • The model observes covariates, the minimum of event and censoring times, and an event indicator; likelihoods are conditional on the observed covariates.The baseline hazard is assumed known and positive.Source: §2 · p. 4
  • Interval truncation uses disjoint time intervals and known retention probabilities, and the framework also includes right censoring.The inclusion indicator is conditionally independent of censoring time given the true event time and covariates.Source: §2 · p. 4
  • Population hazard, survival, and inclusion probabilities describe the population before selection, while sampled-person expectations condition on inclusion.The notation suppresses this conditioning; likelihoods remain conditional on observed covariates.Source: §2 · p. 4
  • Population results assume a positive probability of observing an event for almost every sampled covariate value, non-degeneracy, and regularity conditions on the loss and its gradient.The event-observation condition rules out covariate values for which every sample is censored.Source: §2 · p. 5

3 Main results

The main results establish an optimization route for Cox regression with truncation and right censoring when the baseline hazard is known, despite a nonconcave likelihood. They derive the result through weighted truncated exponentials and contrast it with failures of concavity and with guarantees that hold under specific conditions.

3 Main results

  • Although interval truncation can make the population log-likelihood nonconcave, its gradient points toward the true parameter at every incorrect value.This direction property rules out spurious stationary points and supports convergence of projected SGD.Source: §3 · pp. 5–6

3.1 Weighted truncated exponentials

  • For weighted truncated exponentials, expectations of fixed nondecreasing functions decrease with rate; right censoring can be represented by a nondecreasing transformation that preserves the mean.The transformation replaces censored values with conditional tail means and is nonconstant when exact observations have positive probability.Source: §3.1 · pp. 7–8

3.2 Connecting to Cox regression

  • Transforming event times by the cumulative baseline hazard maps truncation intervals to weighted truncated-exponential intervals, whose retention probabilities do not depend on the regression parameter.With covariates fixed, changing the parameter changes the exponential rate, enabling comparison of candidate and true event-time distributions.Source: §3.2 · p. 8

3.3 Population likelihood landscape

  • The transformed-time likelihood gradient yields a unique global minimizer and stationary point at the true parameter, with a directional gradient property away from that point.The proof uses the one-dimensional result for fixed covariates and lifts it to the multivariate population loss.Source: §3.3 · pp. 8–9

3.4 Convergence of SGD

  • Under the stated regularity and boundedness conditions, minibatch projected SGD has a finite-time convergence guarantee to the true optimum.The guarantee assumes a compact convex projection set containing the true parameter and prescribed step sizes.Source: §3.4 · p. 9

4 Impossibility results for unknown baselines

With an unknown baseline hazard, truncation can undermine both identifiability and the favorable optimization geometry established for known baselines. The section gives examples of non-identifiability, a suboptimal local minimum, and a flat suboptimal region after profiling over baselines.

4 Impossibility results for unknown baselines

  • Distinct baseline hazards can produce the same uniform sampled distribution under truncation, making the baseline and full model non-identifiable.In the example, θ = 0, there is no censoring, and truncation is to (0, 1).Source: §4 · p. 10
  • Even with correct specification and identifiable parameters, jointly estimating the regression coefficient and a smooth baseline can create a suboptimal local minimum.The example uses a single truncation interval and no censoring.Source: §4 · p. 10
  • Optimizing over a broad class of baselines can leave the profiled loss constant and suboptimal across an interval of incorrect regression coefficients.For the constructed model, the loss is constant on (−∞, 0] and every θ < 0 is a suboptimal stationary point; the infimum is not attained by any baseline in H.Source: §4 · p. 10

5 Experiments

Experiments on four survival datasets test truncated and standard Cox models alongside neural variants under increasing outcome-dependent truncation. Accounting for truncation improves accuracy, and first-order optimizers scale substantially better than EM on the enlarged dataset.

5 Experiments

  • The semi-synthetic tests generate event times from dataset-fitted models, impose censoring at original rates, then apply known outcome-dependent truncation.The four datasets are HEARTS, SUPPORT, METABRIC, and FLCHAIN; risk-score MSE is evaluated on untruncated test sets.Source: §5 · p. 11
  • Across four datasets, truncation correction improves accuracy as truncation grows; truncated neural networks perform best when the data-generating model is neural.When truncation is ignored, MSE increases rapidly with its severity.Source: §5 · pp. 11–12
  • With a linear Cox data-generating model, truncated Cox regression closely recovers true risk scores, while the neural model achieves similar performance.Figure 3 reports results across datasets and shows standard-deviation shading.Source: §5 · p. 11, Figure 3 · p. 11
  • Theoretical guarantees apply to linear Cox regression, but the same truncated-likelihood correction is effective empirically for neural survival models.The neural variants replace the linear risk score with a feed-forward network trained using Adam.Source: §5 · pp. 11–12
  • On HEARTS-Large, Adam and projected SGD achieve comparable MSE around 100× faster than EM, converging in 20–30 seconds versus roughly 45 minutes.The enlarged dataset has 750,000 samples and 2,000 covariates; on HEARTS, EM averages over 30 seconds versus 5.15 seconds for projected SGD and 2.05 seconds for Adam.Source: §5 · p. 12, Table 2 · p. 12

6 Conclusions

The paper characterizes efficient learning of Cox regression under outcome-dependent truncation and right censoring. It proves efficient projected-SGD learning with a known baseline hazard, while showing that unknown-baseline models can have unfavorable optimization landscapes.

6 Conclusions

  • With a known baseline hazard, projected SGD converges efficiently to true Cox parameter values despite outcome-dependent truncation and right censoring.The experiments also show first-order methods reduce prediction error relative to models that ignore selection, improve neural networks empirically, and yield orders-of-magnitude speedups over EM.Source: §6 · p. 12
  • When the baseline hazard is unknown, identifiable and correctly specified models can still have strict suboptimal local minima.Thus, a favorable optimization landscape is not guaranteed in that setting.Source: §6 · p. 12
  • Across four semi-synthetic datasets, first-order methods reduce prediction error relative to models that ignore selection and also improve neural networks empirically.The paper reports orders-of-magnitude speedups relative to EM.Source: §6 · p. 12

A Additional related work

Prior work establishes statistical properties for truncated survival estimators but leaves global optimization largely uncharacterized; this paper relates those challenges to efficient learning under truncation and survival analysis. The Cox-regression background derives likelihoods for known baselines and interval truncation.

A Additional related work

  • Unlike truncated regression and exponential-family work [Ilyas et al., 2020, Daskalakis et al., 2019, Lee et al., 2023], truncated Cox likelihoods are generally nonconvex and right-censored observations provide only lower bounds on event times.The paper shows that with a known baseline the population likelihood nevertheless has a unique stationary point at the true parameter, with gradients directed toward it.Source: A Additional related work · p. 17
  • For doubly truncated Cox regression, inverse-probability weighting, pseudo-likelihood, and conditional-likelihood approaches establish consistency and asymptotic properties but leave the global optimization landscape uncharacterized; some likelihood-based approaches use EM with only local convergence guarantees.Source: A Additional related work · p. 17
  • Truncated survival analysis includes left truncation, length-biased sampling, and doubly truncated Cox models, while survival prediction also spans right-censored models and neural architectures [Tsai et al., 1987, Wang, 1996, Qin and Shen, 2010, Kim et al., 2013, Yu et al., 2016, Mandel et al., 2018, Rennert and Xie, 2018, Shen and Liu, 2019, Rennert and Xie, 2022, Shen and Hsu, 2020, Chiou et al., 2019, Li et al., 2020, Ishwaran et al., 2008, Kaplan and Meier, 1958, Wei, 1992, Lee et al., 2018, Nagpal et al., 2021, Katzman et al., 2018].Cox regression remains widely used, including in high-stakes organ allocation [Cox, 1972, Papalexopoulos et al., 2024, OPTN, 2025, Dickerson and Sandholm, 2015, Zilberstein et al., 2026c,b,d, Anagnostides et al., 2026, Berrevoets et al., 2020].Source: A Additional related work · p. 17
  • Even with an identifiable, correctly specified model, learning an unknown baseline can produce suboptimal local minima, contrasting with classical survival nonidentifiability results.The paper also places its known-baseline unique-stationary-point result within work on favorable nonconvex optimization landscapes [Ge et al., 2015, Lee et al., 2016, 2019, Ge et al., 2016].Source: A Additional related work · p. 17

B Further background on Cox regression

  • For independent observations, the Cox-model likelihood is derived from separate event and censoring contributions, with the full likelihood combining those cases.When the baseline hazard is known, its event-time term is independent of the regression coefficient, and the resulting log-likelihood is concave because its Hessian is negative semidefinite.Source: B Further background on Cox regression · p. 18
  • For interval truncation, the likelihood distinguishes event and censored observations using each sample’s inclusion probability and the joint probability of inclusion and event time.The resulting known-baseline log-likelihood is concave.Source: B Further background on Cox regression · p. 18

C Omitted proofs

The omitted proofs establish a unique population optimum and finite-time optimization guarantee when the baseline hazard is known, and characterize strict suboptimal minima in the jointly learned and baseline-profiled Cox models. They also give the conditions and construction for these results, including flatness of the profiled loss below the true coefficient despite identifiability.

C Omitted proofs

  • With a known baseline hazard, the one-dimensional population loss has the true parameter as its unique global minimizer and stationary point, and the full population loss has the true parameter as its unique stationary point.The one-dimensional loss decreases to the true parameter and increases afterward.Source: C Omitted proofs · pp. 22–23
  • A finite-time convergence guarantee for minibatch projected SGD holds on a compact convex parameter set under the stated differentiability, finite-Hessian, and bounded-gradient-moment assumptions.The guarantee uses a fixed initialization and step sizes proportional to 1/(i + t0).Source: C Omitted proofs · p. 23
  • A correctly specified Cox model with an identifiable, smooth, strictly positive unknown-baseline family can have a suboptimal local minimum despite the true parameter being its unique global minimizer.The theorem’s construction uses a single truncation interval and no censoring; the incorrect interior point has a different regression coefficient.Source: C Omitted proofs · p. 24
  • The construction makes the incorrect parameter stationary by matching a group’s mean and tuning baseline shape, then uses a curvature parameter to make the Hessian positive definite.The incorrect point is stationary for every value of the curvature parameter.Source: C Omitted proofs · pp. 25–26
  • The identifiable joint model has a strict suboptimal interior local minimum β′, while β∗ is its unique global minimizer.The proof obtains positive-definite Hessian curvature at β′ and shows F(β′) > F(β∗).Source: C Omitted proofs · p. 27
  • The profiled loss is constant and strictly above its value at θ∗ for θ≤0; every θ<0 is therefore a suboptimal stationary point.For θ<0, the infimum over baseline hazards is not attained by any hazard in H.Source: C Omitted proofs · pp. 28, 30
  • For θ<0, the candidate likelihood ratio is strictly increasing while the true ratio is strictly decreasing, forcing loss above the profiled infimum for every h0 ∈ H.The strict comparison follows from the oppositely monotone likelihood ratio and η under the positive candidate density.Source: C Omitted proofs · pp. 28–29
  • The true coefficient θ∗ = log 2 is identifiable in the binary-covariate model truncated to 0 < T∗ < 1 without censoring.The model uses a true baseline hazard h∗0(t) = 1 and covariate values 0 and 1 with equal sampled probability.Source: C Omitted proofs · pp. 28, 30

D Further details on the experiments

The experiments use four clinical datasets for distinct survival-prediction applications. They describe cohort sizes and covariates, then specify consistent data splitting and preprocessing.

D Further details on the experiments

  • The four datasets are HEARTS for transplant graft survival, SUPPORT for time to death among hospitalized adults, METABRIC for breast cancer survival, and FLCHAIN for mortality prediction.HEARTS uses the UNOS patient registry; SUPPORT contains 8,873 patients and 14 clinical covariates, while METABRIC contains 1,903 patients and 9 clinical and molecular covariates.Source: D Further details on the experiments · p. 30
  • Experiments run on an M4 Pro processor with 24GB unified memory; all datasets except the UNOS registry are publicly available through standard Python libraries.The same preprocessing is applied without modification to validation and test sets.Source: D Further details on the experiments · p. 30
  • Patients are split into training, validation, and test sets of 70%, 15%, and 15%, with preprocessing determined using the training split only.Continuous covariates are median-imputed when needed and standardized; missing categorical values receive a separate category, and categorical variables are one-hot encoded.Source: D Further details on the experiments · p. 30

D.1 Semi-synthetic data generation

Semi-synthetic outcomes are generated from survival models fitted to the original datasets, with a fixed risk function and estimated baseline hazard defining the ground truth. Truncation and censoring are imposed afterward, enabling evaluation against known outcomes on untruncated test data.

D.1 Semi-synthetic data generation

  • The original survival times and censoring indicators fit the model used to define the semi-synthetic data-generating process, and the fitted model is not evaluated as an experimental method.The model is selected using validation data and then frozen.Source: §D.1 · pp. 30–31
  • The risk function is formed from the regression or neural-network component after mean-centering its log-risk scores, which fixes the arbitrary additive offset without changing the conditional survival distribution.The corresponding shift is absorbed into the baseline hazard.Source: §D.1 · p. 31
  • The cumulative baseline hazard is estimated from the original training data with the Breslow estimator, then interpolated and extrapolated to create a continuous, nondecreasing function with a computable inverse.The fixed risk function and cumulative baseline hazard are treated as the exact ground-truth model.Source: §D.1 · p. 31
  • Independent unit-exponential draws generate latent event times under the specified proportional-hazards model, after which truncation and right censoring are imposed synthetically.Evaluation uses the full, untruncated synthetic test set to measure recovery of the underlying risk function and survival distribution.Source: §D.1 · p. 31

D.2 Synthetic truncation and censoring

The experiments impose synthetic right censoring and tunable, event-time-dependent truncation, then use first-order optimization with fixed projected-SGD settings. Truncation strength is varied across six levels, from no removals to an expected 86% removed.

D.2 Synthetic truncation and censoring

  • Censoring times are independently drawn from an exponential distribution calibrated to match each dataset’s original censoring rate, then fixed across experiments.Source: §D.2 · p. 32
  • Retention is independently sampled within four event-time quartiles, with interval-specific probabilities that decrease at different rates as ζ increases.The second quartile is downsampled more aggressively than the third.Source: §D.2 · p. 32

D.3 Projected SGD settings

  • Projected SGD uses learning rate 10^-3, weight decay 10^-4, batch size 512, at most 500 epochs, and early-stopping patience 30.The projection set is a Euclidean ball of radius 5 for both data-generating model settings.Source: §D.3 · p. 32

D.4 Neural network model

The neural models replace the linear Cox risk score with a multilayer perceptron. Their standard and truncated versions differ in whether the likelihood conditions on inclusion.

  • DeepSurv uses a scalar-output multilayer perceptron with two width-64 hidden layers, ReLU activations, and dropout probability 0.1.Source: §D.4 · p. 32
  • Standard and truncated neural models use identical initial weights and differ only in their loss; training uses Adam with learning rate 10^-3 and batch size 512.Source: §D.4 · p. 32
  • The truncated neural model conditions the likelihood on inclusion, while the standard neural model minimizes the ordinary negative log-likelihood.The truncated loss is described as the likelihood of the original truncated Cox model with a linear predictor.Source: §D.4 · pp. 32–33

D.5 Enlarging the HEARTS dataset

The appendix describes a synthetic enlargement of HEARTS, a known-baseline EM procedure, and omitted figures and tables.

D.5 Enlarging the HEARTS dataset

  • HEARTS is enlarged synthetically by resampling training examples and adding Gaussian covariates, with coefficients scaled so the added features contribute half the original features’ risk-score variance.Censoring and truncation are applied as before.Source: §D.5 · p. 33

E Known-Baseline Expectation Maximization

  • Known-baseline EM estimates represented individuals and transformed event times in its E-step, then maximizes the expected complete-data Cox log-likelihood with a concave Newton-optimized objective.It uses the first-order methods’ stopping condition and projected-SGD regularization.Source: E Known-Baseline Expectation… · p. 33

F Omitted tables and figures

  • Figure 4 reports standard and truncated Cox-regression and neural-network results across datasets under a neural-network data-generating model; Table 4 compares optimizers on HEARTS-Large.The supplied captions identify the scope of these visuals but do not provide their results.Source: Figure 4 · p. 34, Table 4 · p. 34

Related papers

Loading 2610.12538v2…