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

Thank you very much for the detailed review, constractive comments, and your support for the acceptance of our paper. The following is our point-by-point response.

• The max time for RIST and the RMST horizon are very close, so it seems that the imputation procedure should have a very minimal effect.

Right censoring can occur at any time between 0 to $t_{max}$. As a result, the censoring rate may be high even when $h$ is very close to $t_{max}$, and therefore imputation has a substantial impact.

• The main weakness of the paper is the lack of baseline methods in the experiments.

Golmakani & Polley, (2020) propose two algorithms for constructing super learners in survival data prediction where the individual algorithms are based on proportional hazards. Bo et al., (2024) studied two meta-learning algorithms, T-learner and X-learner, each combined with three types of machine learning methods: random survival forest, Bayesian accelerated failure time model and survival neural network. These additional baselines will be incorporated in the revised version.

The existing literature for this problem is not sparse and there are baseline methods included in Cui et al., (2023). In particular, Cui et al. showed that their method is superior to the random survival forest (Iswaran & Kogalur, 2019), S-learner (Kunzel, 2019), enriched random survival forest (Lu et al., 2018), and the IPCW causal forest. Therefore, we focus on comparing our method with the CSF method of Cui et al. (2023). This point will be added to the revised version of the paper.

We believe that comparing to methods that ignore right censoring or causal inference principles adds little value, as their limitations are well established in the literature, even if not in our exact setting.

• It would be helpful to include a derivation or explanation for some of these in an appendix, especially the estimator for the target quantity in equation (2).

Thank you for bringing this up. The revised version will include detailed explanations, as suggested.

• It also seems that equation (4) is only valid when $t > C_i$.

Thank you for noticing this mistake, it will be corrected.

• The analysis on the Illinois unemployment insurance dataset seemed incomplete. Metrics or other analysis should be provided to indicate that the MISTR-IV results are actually better.

Thank you for this important comment. Indeed, as the reviewer points out, lacking ground truth makes any solid model evaluation and comparison challenging. Indeed, censoring and hidden confounding make the problem especially challenging in our setting. Nonetheless, following the reviewer's suggestion we have added qualitative comparisons between the top 10% and bottom 10% of the population expected to benefit the most and the least from the treatment, as rated by CSF, MISTR, and MISTR-IV; see Table 3 in the following link: https://drive.google.com/file/d/1tT_ROACNebxVOC09ty2ng_q4luaBD8Ay/view?usp=sharing . Such comparisons may reveal differences between the model results that may be explained by domain knowledge and will help to guide model selection. We see that MISTR-IV arrives quite distinct conclusions regarding the populations most and least benefitting from the treatment. We leave a full interpretation of these results to domain experts in the field.

Thank you very much for your thorough review and insightful feedback.

• Why does replacing censoring rate estimation in methods such as CSF with the conditional survival distribution proposed by RIST mitigate the impact of extreme cases?

Thank you for raising this important issue. It is well known that inverse probability weighting estimators, while often consistent, can suffer from high variance due to instability in the weights. The CSF method combines IPCW with a doubly robust approach to avoid discarding observations with $\delta_i^h = 0$ and to mitigate bias of the censoring distribution. However, despite this robustness, the use of unstable weights can still result in substantial variance. In contrast, our proposed multiple imputation approach avoids IPW altogether. Our extensive numerical studies clearly demonstrate the superior performance of MISTR compared to CSF. We agree that this point should be better explained in the revision.

• How IV and MISTR are integrated

We agree and apologize for not clearly explaining it. Consider first the unconfounded setting and as a start assume a constant treatment effect $\tau$. Consider the partially linear model of Robinson (1988): $g(\widetilde{T}_i) = \tau W_i + f(X_i) + \zeta_i$ where $E(\zeta_i | W_i , X_i)=0$, $W_i \perp \zeta_i$, and $E(W_i|X_i ) = \Pr(W_i=1|X_i)$. Then, it is easy to verify that $g(\widetilde{T}_i) - E\{g(\widetilde{T}_i)|X_i \} = \tau \{ W_i - E(W_i|X_i) \} + \zeta_i$. In the absence of censoring, $\tau$ can be estimated by the score function provided in line 199 of the paper. Eq. (2) (of the paper) then goes beyond constant treatment effects and accommodates a heterogeneous effects function $\tau(x)$. In the confounded setting, we relax the independence assumption between $\zeta_i$ and $W_i$, while the instrument $Z_i$ is independent of $\zeta_i$ given $X_i$, along with Assumptions B.1--B.5. Hence in the absence of censoring, Eq. (2) is replaced by

$$S^{IV}_n(\tau(x)) =$$ \sum_{i=1}^n \alpha_i(x) \cdot \left(Z_i - \widehat{h}(X_i) \right\) \cdot \left[ g(\widetilde{T}_i) - \widehat{m}(X_i) - \tau(x) \{ \left(W_i - \widehat{e}_i(X_i) \right) \} \right] = 0

where $\widehat{h}(X_i)=E(Z_i|X_i)$. We accommodate right-censoring by multiple imputation, and in Step 16 of Algorithm 1 causal forests are applied using either $S_n(\tau(x))$ or $S^{\sf IV}_n(\tau(x))$, corresponding to the MISTR and MISTR-IV estimators, respectively. We'll include a detailed explanation in the revised version.

• What is the author’s original theoretical contribution?

We view our main contribution to be the introduction of new methods. MISTR and MISTR-IV are new nonparametric estimators for HTE, designed for settings without and with unobserved confounding, respectively. Our methods build upon the foundations of RIST and Causal Forests (Athey et al., 2019), effectively merging their strengths to yield estimators with lower variance. This combination results in a notably flexible framework. As evidence of this flexibility, we show how the core approach can be easily modified to accommodate instrumental variables (resulting in MISTR-IV), extending its applicability in a way that is often challenging for other nonparametric estimators. While rigorous theoretical guarantees for our methods are currently lacking, extensive numerical studies, including using real-world datasets, demonstrate that MISTR and MISTR-IV consistently outperform state-of-the-art non-parametric alternatives, such as CSF without unobserved confounding and IPCW-IV with unobserved confounding, in terms of estimation efficiency.

• could papers [3]--[8] also be applied to survival analysis problems?

Paper [3] uses linear semiparametric or parametric regression models to relate the outcome and covariates, whereas our approach is fully nonparametric. Moreover, applying linear models to time-to-event data typically requires outcome transformation or additional constraints. Papers [4] and [5] address non-monotone missingness mechanisms, while right censoring is a special case of monotone coarsening (see Section 9.3 of ``Semiparametric Theory and Missing Data'', Tsiatis A.A., 2006)

Paper [6] focuses on HTE under sample selection without exclusion restrictions, and paper [7] deals with estimating mean functionals under non-ignorable non-response, where missingness depends on unobserved values. Paper [8] tackles both latent confounding bias—stemming from unmeasured variables influencing both treatment and outcome—and collider bias, which arises from non-random sample selection affected by both treatment and outcome. In contrast to these works, right-censored data provide partial information—the event has not occurred up to the censoring time—which must be incorporated into estimation. As such, these works are not directly applicable to our setting, though adapting them to the specific settings and constraints of survival analysis could be a valuable direction for future research.

Thank you very much for the careful review, thoughtful comments, and your support for the acceptance of our paper. The following is our point-by-point response.

• Explain the equations presented in the main body of the paper.

We apologize for this omission. The revised version of the paper will include additional explanations for the equations presented in the main text.

• How does MISTR perform when the assumption of ignorable censoring is only approximately met?

In survival analysis, Assumption A.4, stating that ${\widetilde{T}}_i \perp C_i \mid X_i, W_i$, is very widely used and is considered standard. In our context, see for example, Cui et al. (2023), Bo et al., (2024), Golmakani & Polley, (2020) Iswaran & Kogalur, (2019) and Lu et al., (2018). The reason for its ubiquity is that this assumption is essential, as for each individual only one of the two times (event or right-censoring) is observed. Consequently, their joint distribution is unidentifiable. This assumption is empirically untestable and must be accepted or rejected based on subject-matter knowledge. When there is a reason to believe this assumption may be substantially violated, for example due to the occurrence of other well-defined events, several strategies can be considered. One approach is to treat such events as competing risks and perform a competing-risks analysis. Alternatively, if one has a good knowledge of the specifics of the violation, one may employ parametric or semiparametric models that explicitly specify a joint distribution of the event and censoring times, such as copula-based models.

• What methods do you suggest for assessing the strength and validity of the instrumental variables used, and how does MISTR-IV behave with weak instruments?

Following this comment we extended our simulation study to assess the sensitivity of MISTR-IV to weak instruments by running Setting 200 (47% censoring) and Setting 204 (88% censoring) with IV of varying strength. We repeat the analysis of Section 6.2, modifying only the coefficient of $Z$ in the model of $W^{*}$, as reported in Table 1. The results are reported in Table 2 and Figure 1 in the following link: https://drive.google.com/file/d/1tT_ROACNebxVOC09ty2ng_q4luaBD8Ay/view?usp=sharing . As expected, the mean absolute error of all methods increases as the instrument weakens. Nonetheless, MISTR-IV outperforms the alternative approaches.

Validating instrumental variable strength in partially linear IV regression is a challenging topic and an active field of research (Windmeijer, (2025), Burauel (2023), Florence (2012), Hahn and Hausman (2002), Stock et al. (2002)). In linear models, IV strength is commonly evaluated using the effective first-stage F-statistic of Montiel, Olea, and Pflueger (2013). Windmeijer (2025) extends this approach to the Generalized Method of Moments (GMM) framework by proposing the Robust F-statistic. In the future we plan to investigate the best way to incorporate it in our approach.

The validity of the instrumental variable requires the standard IV assumptions: 1. Exclusion restriction: the instrument affects the outcome only through its influence on treatment assignment; 2. Independence: the instrument is independent of any unobserved confounders; 3. Relevance: the instrument is correlated with the treatment assignment. The relevance assumption can be empirically assessed by examining the correlation between the instrumental variable $Z$ and the treatment assignment $W$. In contrast, the exclusion restriction and independence assumptions cannot be tested directly and must be justified using domain knowledge. For example, in the Illinois Unemployment Insurance Experiment, the proposal to join the experiment was randomized, supporting the independence assumption. However, validating the exclusion restriction requires establishing that the proposal itself did not influence the outcome directly regardless if the individual chooses to participate or not - a condition that is inherently unidentifiable from the observed data. Another approach gaining prominence recently is falsification tests, or negative controls (Eggers et al. 2024). Constructing such tests for censored time-to-event data is an interesting avenue for future research

• The paper does not make strong theoretical improvements to the literature.

Indeed, we view our main contribution as introducing new methods. While our design choices -- particularly addressing heavy censoring -- are strongly motivated, we do not claim a theoretical result. Instead, we demonstrate the advantages of our method through extensive experiments, including several with real-world data, and through empirical comparisons with existing approaches in realistic survival analysis settings.