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

TL;DR

This paper introduces a deep survival learner (DSL) that estimates time-varying heterogeneous treatment effects in survival analysis, accounting for right censoring and leveraging neural networks for joint trajectory estimation.

Contribution

The paper proposes a novel doubly robust deep neural network approach for estimating time-specific treatment effects with censored survival data, improving stability and accuracy.

Findings

DSL accurately estimates treatment effect trajectories in simulations.

Joint estimation reduces error and improves stability over separate methods.

Application reveals heterogeneity in treatment effects over time in lung cancer data.

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.