PaperlayerView exact revision on arXiv

Source-linked AI summary

Estimating heterogeneous treatment effects with right-censored data via causal survival forests

Yifan Cui, Michael R. Kosorok, Erik Sverdrup, Stefan Wager, Ruoqing Zhu

arXiv:2001.09887v5stat.MEcs.LGstat.ML

TL;DR

Existing heterogeneous treatment effect methods often cannot handle potentially censored survival outcomes, where realized survival times are not observed. Causal survival forests adjust for censoring with doubly robust estimating equations and outperform available baselines in experiments.

  • Problem

    Heterogeneous treatment effect methods often cannot be used with potentially censored survival outcomes, where only censored times and non-censoring indicators are observed.

    Source: §1 · pp. 1–2

  • Method

    Causal survival forests adapt causal forests to survival outcomes, using doubly robust estimating equations to adjust for censoring.

    Source: §1 · p. 1, §2.3 · p. 7

  • Results

    In almost all scenarios, causal survival forests achieve the best performance among competing methods; in setting 4, IPCW is best for RMST and SRC1 for survival probability.

    Source: §4.1 · p. 15

  • Takeaways & Limitations

    When the stated assumptions hold, causal survival forest predictions are consistent and asymptotically normal, supporting pointwise confidence intervals for τ(x).

    Source: §1 · p. 1, §3.1 · p. 11

    Pointwise intervals for complex τ(·) may be wide, and tuning for undersmoothing can be difficult to assess in practice.

    Source: §3.2 · p. 12

Abstract

from arXiv · show

Forest-based methods have recently gained in popularity for non-parametric treatment effect estimation. Building on this line of work, we introduce causal survival forests, which can be used to estimate heterogeneous treatment effects in a survival and observational setting where outcomes may be right-censored. Our approach relies on orthogonal estimating equations to robustly adjust for both censoring and selection effects under unconfoundedness. In our experiments, we find our approach to perform well relative to a number of baselines.

1 Introduction

Causal survival forests address the difficulty of estimating heterogeneous treatment effects when survival outcomes are right-censored. The paper adapts causal forests with doubly robust censoring adjustment and reports stronger simulation accuracy than non-doubly robust baselines.

1 Introduction

  • Causal survival forests adapt causal forests with doubly robust estimating equations to adjust for censoring in heterogeneous treatment-effect estimation.The paper also establishes conditions under which predictions are consistent and asymptotically normal.Source: §1 · p. 1
  • Right-censored survival outcomes complicate treatment-effect estimation because researchers observe censored times and event indicators rather than the target survival times.The target is the conditional average effect on survival time, or a relevant transformation, given covariates.Source: §1 · pp. 1–2
  • CSF is considerably more accurate than virtual twins and the S-learner in two simulation designs, and directly targets CATE heterogeneity when splitting.The baselines do not use doubly robust survival-analysis methods to target treatment effects.Source: §1 · p. 2
  • The paper also studies asymptotics, runs simulations, illustrates the approach with ACTG175 data, and provides a software implementation.The implementation is part of the grf package for R and C++.Source: §1 · p. 2

1.1 Related work

  • Prior survival trees and forests were not directly targeted at heterogeneous treatment effects in observational studies, while censored policy-learning methods may retain confounding and selection bias.The paper adapts forest-based treatment-effect methods to survival data as a step toward flexible estimation.Source: §1.1 · pp. 2–4

2 Causal survival forests

Causal survival forests estimate covariate-dependent effects on transformed survival outcomes despite right-censoring. They adapt causal forests with censoring adjustments, especially a doubly robust approach, under assumptions supporting causal and censoring identification.

2 Causal survival forests

  • The framework targets treatment effects on transformed survival time under a finite horizon, with potential-outcome, treatment-ignorability, and conditionally independent censoring assumptions.The paper notes that transformations can represent restricted mean survival time or survival probability.Source: §2 · pp. 4–5

2.1 Causal forests without censoring

  • Causal forests combine local forest weights with Robinson’s estimating equation to estimate treatment-effect heterogeneity across covariate-defined neighborhoods.Adaptive tree partitions can narrow along covariate directions where effects change faster.Source: §2.1 · pp. 5–6

2.2 Adjusting for censoring via weighting

  • IPCW adapts causal forests by weighting observed failure cases by the inverse censoring-survival probability, but discards censored cases and can be sensitive to censoring-model errors.The paper notes that censored observations still reveal survival up to their censoring time, so discarding them may reduce efficiency.Source: §2.2 · pp. 6–7, §2.3 · p. 7

2.3 A doubly robust correction

  • The preferred method combines doubly robust censoring-adjusted scores with forest weights and pseudo-outcome splitting to estimate heterogeneous effects.The implementation targets restricted mean survival time and survival probability.Source: §2.3 · pp. 7–9
  • For a constant treatment effect, the orthogonal estimating equation yields an asymptotically normal, root-n estimator under cross-fitting, the stated assumptions, and sufficiently fast nuisance estimation.The paper gives a 1/sqrt(n) rate under fourth-root rates for nuisance components.Source: §2.3 · p. 8

3 Asymptotics and inference

Under stated assumptions, the paper establishes asymptotic Gaussianity for causal survival forests, supporting pointwise confidence intervals, and develops inference for linear projections that summarize treatment-effect heterogeneity.

3 Asymptotics and inference

  • The asymptotic analysis assumes bounded covariate density, Lipschitz treatment and nuisance functions, honest symmetric trees, randomized splits, and subsampling.These conditions define the setting for the stated inference results.Source: §3 · p. 10
  • Under Assumptions 2–8 and the stated rate conditions, causal survival forests are asymptotically Gaussian, supporting valid pointwise confidence intervals for τ(x).The result applies when the theorem’s conditions hold.Source: §3 · p. 11

3.1 Pointwise confidence intervals

  • The bootstrap of little bags estimates asymptotic variance by using cluster structure in tree-training samples and within- versus between-cluster tree correlations.The resulting variance estimate is used to construct pointwise confidence intervals.Source: §3.1 · p. 11

3.2 Inference on linear projections of the CATE

  • Best linear projections can summarize treatment heterogeneity with robust inference even when nonparametric estimates of τ(X) are uncertain.The researcher can prespecify how the conditional mean of τ(X) varies with features such as age and gender.Source: §3.2 · p. 12
  • The projection procedure constructs doubly robust scores from forest CATE estimates, regresses them on features of interest, and uses heteroskedasticity-robust inference.The implementation uses the sandwich variance estimator with HC3 covariance weights.Source: §3.2 · pp. 12–13
  • Because forest nuisance-component rates may be slower, the authors cannot formally verify a √n-rate central limit theorem for their projection approach, though simulations show nearly normal sampling distributions.The nearly normal distribution is reported for the estimator in the simulation experiments.Source: §3.2 · p. 13

4 Simulation study

The simulations compare causal survival forests with survival-forest and causal-forest baselines across restricted mean survival time and survival-probability estimands. CSF performs best in almost all scenarios, while confidence-interval coverage is poor at some feature-space corners and the doubly robust BLP outperforms naive regression for inference.

4 Simulation study

  • The simulations assess treatment effects on RMST at horizon h and survival probability at h; each simulation specifies h, with survival evaluated as surviving past the 90-th percentile of U.Source: §4 · pp. 13–14
  • The study compares CSF and IPCW with VT, SRC1, and SRC2 across RMST and survival-probability estimands.The forests use 2,000 training samples and an independent 2,000-sample test set; each simulation is repeated 250 times.Source: §4 · p. 13, §4.1 · p. 15

4.1 Results

  • CSF performs best in almost all scenarios; in Setting 4, IPCW is best for RMST and SRC1 for survival probability.IPCW generally outperforms the other baselines; overall, CSF is superior to ordinary random survival forests.Source: §4.1 · p. 15, Table 1 · p. 15, Table 2 · p. 16
  • Confidence-interval coverage is poor at some tested points, especially x1 and x4 near feature-space corners.The study evaluates proposed 95% intervals at four deterministic points; coverage is mostly promising elsewhere.Source: §4.1 · pp. 15–16
  • The doubly robust BLP performs well, while naive direct regression can be biased and underestimate sampling variability, producing poor coverage.The comparison considers DR and naive regressions of fitted CATEs against projection variables.Source: §4.1 · p. 16

5 HIV data analysis

The ACTG175 analysis compares ZDV+ddI with ddI in a selected subset treated as observational because of apparent selection effects. The authors truncate survival time before three years; the age-related pattern is not confirmed by the reported BLP coefficients.

5 HIV data analysis

  • The analysis treats the selected ACTG175 subset as observational because treatment groups show apparent selection effects.It compares 522 patients receiving ZDV+ddI with 561 receiving ddI.Source: §5 · p. 16
  • Survival time is truncated at h = 1000, just before three years, to keep estimated censoring probabilities in a reasonable range and address positivity.The raw data have all observations with the largest survival times censored, so moments of survival time are not identified there.Source: §5 · p. 17

BLP (DR)

  • Although age appears positively associated with benefit from ZDV+ddI, none of the coefficients in the considered BLPs differs significantly from zero.The projections use doubly robust scores and include all covariates or age alone.Source: BLP (DR) · p. 18

A Proof of Proposition 1

The proof uses cross-fitting and doubly robust estimating equations to control estimation error, concluding that cross-fitting can yield an efficient estimator of the average treatment effect under an n^1/4 nuisance-estimation condition.

A Proof of Proposition 1

  • Cross-fitting splits observations into two halves and estimates each half’s nuisances using the other half.Conditioning on one half makes the other half’s summands mean-zero and independent in the stated argument.Source: A Proof of Proposition 1 · pp. 28–29
  • The error analysis separates survival-nuisance estimation errors from errors in the treatment propensity and outcome-regression estimates.The proof invokes double robustness and bounds the terms involving estimated rather than true nuisance functions.Source: A Proof of Proposition 1 · pp. 29, 31
  • Cross-fitting transforms an n^1/4-consistent machine-learning method into an efficient estimator of the average treatment effect τ.The stated conclusion concerns τ = E[y(T(1)) − y(T(0))].Source: A Proof of Proposition 1 · p. 33

B Proof of Lemma 2

The proof adapts the Proposition 1 argument to generalized random forest weights and then completes the lemma’s proof.

B Proof of Lemma 2

C Proof of Theorem 3

The theorem proof uses a forest-weighted linear approximation and the influence function, then checks regularity conditions for the estimating equation.

C Proof of Theorem 3

  • The proof begins with a linear approximation to the generalized-random-forest estimate using its forest weights.Source: C Proof of Theorem 3 · p. 34
  • The approximation relates the influence function for τ(x) to a pseudo-forest output whose outcomes are τ(x) + ρ*_i(x).Source: C Proof of Theorem 3 · p. 34
  • The passage states that Assumptions 2–6 in Athey et al. [2019] hold immediately from the definition of estimating equation ψτ(x), and specifically notes Lipschitz continuity in terms of τ(x) for Assumption 4.Source: C Proof of Theorem 3 · p. 34

D Figures

The figures compare methods across four scenarios using classification error and excess MSE for restricted-mean-survival-time and survival-probability estimands; the supplied passages do not report which method performs best.

Figures

  • Classification-error plots compare VT, SRC1, SRC2, IPCW, and CSF across four scenarios for both estimands.The figures show RMST and survival-probability estimands, with 15 covariates, a test set of 2,000, and 250 repetitions.Source: Figure 5 · p. 35, Figure 7 · p. 37
  • The RMST excess-MSE plot compares each method’s MSE with the minimum MSE across methods in each scenario.The reported setup uses training sizes 500, 1,000, 2,000, and 5,000, 15 covariates, a test set of 2,000, and 250 repetitions.Source: Figure 6 · p. 36, D Figures · p. 36
  • The survival-probability excess-MSE plot uses the same relative-to-best-method comparison across four scenarios.The setup uses training sizes 500, 1,000, 2,000, and 5,000, 15 covariates, a test set of 2,000, and 250 repetitions.Source: Figure 8 · p. 38, D Figures · p. 38

Related papers

Loading 2001.09887v5…