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

2026-04-11

AI Summary

This study addresses the challenge of stably estimating the time trajectory of the conditional average treatment effect (CATE) in survival analysis with right censoring and time-varying treatment effects. The authors propose Deep Survival Learner (DSL), a method that constructs doubly robust pseudo-outcomes to train a multi-output deep neural network, jointly modeling heterogeneous treatment effects across time points. By integrating cross-fitting with shared representation learning, DSL reduces both bias and overfitting. This approach provides the first joint and stable estimation of CATE trajectories over time, with theoretical guarantees showing that leveraging temporal structure incurs no substantial increase in approximation error and maintains robustness under model misspecification. Simulations demonstrate superior finite-sample performance even when models are misspecified, and an application to a Boston lung cancer study reveals dynamic heterogeneity in perioperative chemotherapy effects across patient subgroups and time.

Technology Category

• Machine Learning: Causal Learning
• Search and Optimization: Learning to Search
• Reasoning under Uncertainty: Causality

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• User Modeling, Personalization and Recommendation: Studies of user behavior, including longitudinal effects of personalized systems

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.

Problem

• heterogeneous treatment effects
• survival analysis
• right censoring
• conditional average treatment effect
• time-varying effects

Innovation

• heterogeneous treatment effects
• survival analysis
• deep survival learner
• doubly robust estimation
• multi-output neural network