Heterogeneous Treatment Effect in Time-to-Event Outcomes: Harnessing Censored Data with Recursively Imputed Trees

2025-02-03

AI Summary

Existing survival analysis methods suffer from bias in estimating heterogeneous treatment effects (HTE) under high right-censoring rates (e.g., ALS data) and unmeasured confounding, and lack principled integration of instrumental variables (IVs) for causal identification. To address this, we propose MISTR—the first nonparametric framework for IV-based HTE estimation in censored survival data. MISTR avoids modeling the censoring mechanism by leveraging recursive imputation survival trees and multiple imputation for robustness. It innovatively unifies IV adjustment, causal forests, and augmented inverse probability censoring weighting (AIPW-CPW) to achieve unbiased HTE estimation under unmeasured confounding. Extensive evaluations on heavily censored synthetic data, the ACTG 175 clinical trial dataset, and the Illinois unemployment study demonstrate substantial improvements over state-of-the-art methods. MISTR is currently the only unified framework enabling nonparametric IV-HTE estimation for right-censored outcomes.

Technology Category

• Machine Learning: Causal Learning
• Reasoning under Uncertainty: Causality
• Intelligent Robots: State Estimation

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Responsible Web: Measurement, analysis, and circumvention of Web censorship
• User Modeling, Personalization and Recommendation: User privacy protection in personalized systems

Abstract

Tailoring treatments to individual needs is a central goal in fields such as medicine. A key step toward this goal is estimating Heterogeneous Treatment Effects (HTE) - the way treatments impact different subgroups. While crucial, HTE estimation is challenging with survival data, where time until an event (e.g., death) is key. Existing methods often assume complete observation, an assumption violated in survival data due to right-censoring, leading to bias and inefficiency. Cui et al. (2023) proposed a doubly-robust method for HTE estimation in survival data under no hidden confounders, combining a causal survival forest with an augmented inverse-censoring weighting estimator. However, we find it struggles under heavy censoring, which is common in rare-outcome problems such as Amyotrophic lateral sclerosis (ALS). Moreover, most current methods cannot handle instrumental variables, which are a crucial tool in the causal inference arsenal. We introduce Multiple Imputation for Survival Treatment Response (MISTR), a novel, general, and non-parametric method for estimating HTE in survival data. MISTR uses recursively imputed survival trees to handle censoring without directly modeling the censoring mechanism. Through extensive simulations and analysis of two real-world datasets-the AIDS Clinical Trials Group Protocol 175 and the Illinois unemployment dataset we show that MISTR outperforms prior methods under heavy censoring in the no-hidden-confounders setting, and extends to the instrumental variable setting. To our knowledge, MISTR is the first non-parametric approach for HTE estimation with unobserved confounders via instrumental variables.

Problem

• Heterogeneous Treatment Effects
• Censored Data
• Instrumental Variables

Innovation

• MISTR
• Heterogeneous Treatment Effects
• Instrumental Variables