Estimating heterogeneous treatment effects with survival outcomes via a deep survival learner

Abstract: Estimating heterogeneous treatment effects in survival settings is complicated by right censoring as well as the time-varying nature of the estimand. While the conditional average treatment effect (CATE) provides a natural target, most existing approaches focus on a single prespecified time point and do not account for the temporal trajectory, leading to instability in estimation. We propose a deep survival learner (DSL) for estimating heterogeneous treatment effects with right-censored outcomes. The method is based on a doubly robust pseudo-outcome whose conditional expectation identifies time-specific CATEs under standard assumptions. This construction remains unbiased if either the outcome model or the treatment assignment model is correctly specified, when properly accounting for censoring. To estimate CATEs over a clinically relevant time spectrum, DSL employs a multi-output deep neural network with shared representations, enabling joint estimation of treatment effect trajectories. From a theoretical perspective, we derive error bounds for both pointwise and joint estimation over time. We show that joint estimation can leverage temporal structure to control estimation error without incurring much additional approximation cost under smoothness conditions, leading to improved stability relative to separate estimation. Cross-fitting is incorporated to reduce overfitting and mitigate bias arising from flexible nuisance estimation. Simulation studies demonstrate favorable finite-sample performance, particularly under nuisance model misspecification. Applied to the Boston Lung Cancer Study, DSL reveals heterogeneity in the effects of perioperative chemotherapy across patient characteristics and over time.

• The paper introduces a doubly robust pseudo-outcome approach to estimate conditional average treatment effects in survival settings with right censoring.
• It employs a multi-output deep neural network architecture to jointly model treatment trajectories over time, enhancing estimation stability.
• Empirical evaluations, including a lung cancer study, demonstrate superior performance with minimal bias in high-dimensional scenarios.

Deep Survival Learner for Heterogeneous Treatment Effects with Survival Data

Introduction

The analysis of heterogeneous treatment effects (HTEs) in survival settings is challenged by right censoring and the intrinsic time-dependence of the outcome. The conditional average treatment effect (CATE), which quantifies the impact of treatment conditional on covariates, is a key estimand for individualized decision making in contexts such as precision medicine. Classical methods often target CATE at a fixed time, neglecting the temporal evolution of effects and leading to suboptimal estimation with increased uncertainty. The presented work addresses these deficiencies through the introduction of the Deep Survival Learner (DSL), a framework leveraging deep neural networks (DNNs) for temporally joint estimation of heterogeneous treatment effect trajectories in the presence of censoring (2604.10398).

Methodological Framework

Doubly Robust Pseudo-Outcome for Survival Causal Inference

CATEs at time $t$ are unidentifiable due to latent counterfactuals and censoring. The DSL constructs a doubly robust pseudo-outcome, parameterized by time, whose conditional mean with respect to baseline covariates identifies the CATE under general conditions. The pseudo-outcome guarantees unbiasedness provided either the survival or treatment (propensity) model is correctly specified, even in the presence of nuisance model misspecification and non-trivial censoring. The pseudo-outcome formulation extends inverse probability weighting (IPW) and augmentation from the canonical missing data literature and ensures robustness in practical high-dimensional settings.

Multi-Output Deep Neural Architecture

The estimation problem is recast as a vector-valued regression, learning the CATE surface over both the covariate and time dimensions. The DSL employs a fully connected DNN with shared hidden representations and multiple output nodes corresponding to the grid of time points. This architecture facilitates information sharing across the temporal index, effectively regularizing HTE estimation by exploiting expected smoothness and functional dependence in the trajectory of treatment effects.

Cross-fitting is systematically incorporated to ensure decoupling of nuisance estimation and target regression, controlling overfitting and first-order biases induced by sample-adaptive nuisance regression.

Theoretical Results

The authors rigorously characterize statistical properties of the DSL, establishing identification, double robustness, and finite-sample error rates.

• Identification and Double Robustness: The pseudo-outcome admits unbiased conditional mean estimation of CATE if either the survival or propensity model is correct, provided accurate censoring modeling and standard causal assumptions hold.
• Non-asymptotic Error Decomposition: The error in DSL HTE estimates decomposes into approximation error (DNN expressivity), statistical error, and product/first-order terms reflecting nuisance estimation quality. The double robustness property is explicit in the error, wherein first-order terms vanish under correct specification.
• Temporal Joint Estimation: Leveraging functional smoothness in time, the multi-output DNN architecture is shown to be statistically advantageous over pointwise (separately fit) estimation, achieving strictly sublinear error scaling with respect to the number of time points and improved stability under plausible data-generating processes.

Empirical Evaluation

Simulations

Comprehensive simulations benchmark DSL against established meta-learner causal inference approaches (M-, R-, and X-learners), doubly robust Bayesian tree ensembles, and survival forests. Scenarios with correct and misspecified nuisance specifications are considered, as are high-dimensional covariate scenarios and increasing time evaluation grids.

DSL exhibited minimal bias in all configurations and strong robustness to model misspecification. In scenarios where either the survival or treatment model was incorrect, competing methods showed substantial bias. Numerical results indicate DSL achieves competitive or superior mean squared error (MSE) relative to the strongest competitors, especially when nuisance estimators are challenged by dimensionality or nonlinearity. The bias and MSE superiority under increasing numbers of estimation grid points and high-dimensional covariate sets is quantitatively demonstrated (see below).

Figure 1: Estimated conditional average treatment effects (CATE) across gender, tumor stage, age, BMI, and smoking intensity, illustrating both temporal and covariate-driven effect heterogeneity.

Application: Boston Lung Cancer Study

The DSL framework is applied to a real-world analysis of perioperative chemotherapy in stage II/III NSCLC patients from the Boston Lung Cancer Study. The causal estimand is survival probability difference at a given time point (~ CATE), stratified by treatment receipt and baseline variables.

The results reveal dynamic effect heterogeneity: perioperative chemotherapy improves survival probabilities, with greatest effect among women, lower-stage disease, higher BMI, and lower smoking intensity. Temporal smoothness is evident, with effect magnitude peaking within the first 10 years post-surgery and attenuating thereafter.

Figure 2: Kaplan–Meier survival curves delineating improved survival in the chemotherapy group versus surgery alone, motivating temporal causal effect estimation.

Implications and Future Directions

DSL advances the estimation of HTEs for survival data by directly addressing limitations in time-specific and censoring-robust causal estimation. The doubly robust construction and temporal joint estimation approach controls instability, enhances statistical efficiency, and accommodates modern high-dimensional covariate structures. Practically, DSL facilitates patient-specific causal conclusions that are critical for individualized treatment allocation and policy assessment in clinical contexts.

Theoretically, the analysis opens several avenues:

• Uncertainty Quantification: Extensions to formalize time-indexed inference (e.g., uniform confidence bands) for DSL estimates, accounting for deep network complexity and dependence across time.
• Relaxing Censoring Assumptions: DSL relies on correctly specified censoring models; methods for sensitivity analysis or robustification under informative censoring are needed.
• Competing Risks and Multistate Extensions: The extension of DSL to multi-outcome or multistate survival processes, potentially via multi-task architectures, would generalize its applicability.

Conclusion

The DSL framework represents a technically sophisticated integration of causal inference, robust learning, and deep architectures for time-to-event data. The combination of a doubly robust pseudo-outcome, temporally-joint neural estimation, and cross-fitted nuisance regression guarantees robust, data-adaptive HTE estimation with controlled statistical properties. Empirical evidence underscores both the practical and theoretical relevance, demonstrating clear advantages under challenging, realistic scenarios. The method paves the way for advanced, individualized causal survival analysis in biostatistics and allied domains.