Orthogonal Survival Learners for Estimating Heterogeneous Treatment Effects from Time-to-Event Data

Thank you for the in-depth evaluation of our manuscript. We will take all your comments and questions to heart and make improvements to our manuscript accordingly. Below, we provide answers to all the questions and concerns you raised.

1. Domain of T:

   Our framework applies to both discrete and continuous time domains. In the main paper, we focus on discrete time for ease of notation. In Appendix E, we extend all derivations for the continuous time domain.

2. Estimation in Theorem 5.1:

   Thank you for giving us the chance to explain the notation-heavy theorem in more detail. The function $\phi(Z, \eta_t)$ in Eq. 7 denotes the (unweighted) pseudo-outcome. It is estimated following Eq. 9 based on the nuisance functions $\eta_t(X) = \left( \pi(X), \left( \lambda^S_i(X, 1), \lambda^S_i(X, 0), \lambda^G_i(X, 1), \lambda^G_i(X, 0) \right)_{i=0}^t \right)$.

The derivatives $\frac{\partial f(\widetilde{\eta}_t(X))}{\partial m}$

for $m \in (\pi, S_{t-1}(\cdot, A), G_{t-1}(\cdot, A) )$ in Eq. 8 denote the (standard finite-dimensional) partial derivatives of the function $f(\pi, S_{t-1}, G_{t-1})$ w.r.t. one of its components. These formulas can be computed analytically and yield closed-form solutions, without the need to numerically estimate derivatives. As an example, we take a look at theT-Learner in Eq. 14:

To address the treatment overlap, we have $f(\widetilde{\eta}_t(X)) = \pi(X) (1 - \pi(X))$. In this case we have $\frac{\partial f(\widetilde{\eta}_t(X))}{\partial \pi} = 1-2\pi(X)$, resulting in the simplified $\rho(Z, \eta_t) = (A - \pi(X))^2$ according to Eq. 8. Note that the T-Learner (=treatment overlap learner) equals the survival R-Learner. As shown in Eq. 13, the survival DR-Learner is also an instantiation of our framework in Theorem 5.1. In conclusion, the derivative in Eq. 8 does not lead to higher variance compared to e.g., the DR-learner, and is easily computable in analytic form.

The (standard) T-Learner consists of two outcome regressions, each trading off bias for variance in its regression problem, which does not guarantee a joint optimal bias-variance trade-off (see, e.g., [1]). Our framework, however, provides statistically efficient learners and thus guarantees that our learners have the lowest variance among all unbiased estimators (see, e.g., [2]). Overall, the derivative in Eq. 8 does not lead to high variance.

Action: We will explain the definition of the derivative more clearly in the camera-ready version of our manuscript.

3. Consistency result:

   Our theorem 5.2. states that the orthogonal loss from Eq. 7 has the same minimizer as the oracle loss from Eq. 11, namely $\mathcal{L}_{f}(g, \eta_t) = \frac{1}{\mathbb{E}[f(\widetilde{\eta}_t(X))]} \mathbb{E}\left[f(\widetilde{\eta}_t(X))\left(\tau(X)- g(X) \right)^2\right]$.

   Hence, Theorem 5.2. immediately implies consistency of our learners: if the nuisance functions are consistent (i.e., well specified), the empirical loss in Eq. 12 converges in probability to the orthogonal loss in Eq. 7 by applying the standard law of large numbers. If $\mathcal{G}$ is well-specified, this implies $g^{\ast}_f = \tau_t$ for the minimizer of Eq. 12 asymptotically.

   Action: We will state consistency more clearly in the camera-ready version of our paper.

4. Benefits of orthogonality and rates:

   Thank you for raising this important point. Theoretical aspects and particularly rates for orthogonal learners are well-studied, we refer to the foundational work “Orthogonal statistical learning” [4] by Foster and Syrgkanis for details. In particular, their work provides slow and fast rates for general Neyman orthogonal losses under a convexity assumption. In Theorem 5.1. we proved Neyman orthogonality of our loss, which implies that their rates are readily applicable to our setting. The additional key assumption of convexity (Assumption 4 in their paper) is also satisfied as we consider a weighted squared loss. Hence, their Theorem 1 shows that nuisance estimation errors impact the estimation accuracy of our second-stage learners only in second order, as it is common for orthogonal learners.

   Action: We will add a discussion on rates to the camera-ready version of our paper.

5. Final effect estimation:

   Recall our definition of $\tau_t(x) = \mathbb{P}(T(1) > t \mid X = x) - \mathbb{P}(T(0) > t \mid X = x)$ for some fixed $t \in \mathcal{T}$ (see Section 3). Thus, $\tau_t(x)$ represents the difference in survival probability up to time $t$ for a unit with covariates $x$, which defines the survival treatment effect. Note that, in survival analysis, the HTE, which is our main effect of interest, does depend on both the covariates $x$ and the time $t$.

   In practice, one might also be interested in related but different treatment effects. In Appendix F, we thus also provide extensions of all our results to other causal quantities, such as the conditional means $\mathbb{E}[T(1) - T(0) \mid X = x]$ as well as the treatment-specific quantities $\mathbb{P}(T(a) > t \mid X = x)$ and $\mathbb{E}[T(a) \mid X = x]$.

[1] Morzywolek, P. et. al.. “On a general class of orthogonal learners for the estimation of heterogeneous treatment effects”. arXiv preprint, 2023

[2] Kennedy, E. H. Semiparametric doubly robust targeted double machine learning: A review. arXiv preprint, 2022.

[3] Victor Chernozhukov et al. “Double/debiased machine learning for treatment and structural 387 parameters”. The Econometrics Journal 21.1 (2018), pp. C1–C68

[4] Foster & Syrgkanis. "Orthogonal statistical learning". The Annals of Statistics (2023).

Thank you for your response and for acknowledging the points raised in our rebuttal. As we believe we have sufficiently addressed your main concerns, we would kindly ask you to reconsider your assessment and, if appropriate, revise your score accordingly.

Thank you once again for your thoughtful evaluation, we sincerely appreciate your time and effort in reviewing our work.

Thank you for your positive evaluation of our manuscript and your extensive comments and suggestions. We will take everything to heart and make improvements to our manuscript accordingly. Below, we provide answers to all questions and comments you raised.

Answers to weaknesses:

• Empirical validation:

  Thank you for this valuable suggestion. In general, our choice of datasets was motivated to show primarily the theoretical properties of our proposed learners and by a practical limitation: many survival datasets do not include treatment information, and, as a result, only a small number of publicly available datasets are commonly used in the literature => our selection reflects the availability of such established benchmarks. Nevertheless, to address your comment, we have strengthened our experimental evaluation as follows.

  First, we added new results evaluating our method across varying sample sizes using synthetic Dataset 1 with low censoring overlap (see below). Due to space constraints, we report the results for the C and the T+C learner, though we observe similar patterns for other learners as well. PEHE results are to the power $10^2$. As expected, the variance increases with smaller sample sizes, but overall, the proposed learners remain robust to changes in dataset size. We include a condensed version of the results here; we will provide a more comprehensive version (along with visualizations) in the camera-ready version.

  Learner	75% dataset	100% dataset	200% dataset
  C	0.0974 $\pm$ 0.0799	0.0677 $\pm$ 0.0506	0.0330 $\pm$ 0.0219
  T+C	0.1015 $\pm$ 0.0819	0.0732 $\pm$ 0.0535	0.0324 $\pm$ 0.0218

  Second, we chose established datasets from the causal inference literature, such as the Twins dataset (e.g., [1]), which has the rare property of allowing proxy access to both potential outcomes. This makes it ideal for validating our method. Importantly, we show that our method successfully recovers known effects from this dataset, which we regard as a strength to demonstrate the reliability of our survival learners.

  Third, to highlight the practical utility of our learners, we repeated our analysis using the ADJUVANT trial from oncology [2,3]. We provide details below.

• Minor weaknesses:

  Thank you for noting the typos in the figure and the notation. We will address these issues in the camera-ready version of our manuscript. We apologize for not including an in-depth discussion in our main paper and will include all relevant parts of the paper in the main section of the camera-ready version.

Answers to questions:

• Figures 5 & 6: The training dynamics in Figures 5 & 6 are as expected for our weighted learners. High learning rates are likely to result in oscillating behaviour on the loss, as the reweighting can be initially unstable to optimize. Therefore, we follow best practice and employ a low learning rate, as noted by the reviewer. To counteract the low learning rate, we employ the Adam optimizer, which adapts learning rates over time and may switch from small to larger gradients, which can thus result in larger drops in the validation loss. **

Action: We will include a discussion on the training dynamics in the updated version of our paper.

• Choice of learner: Thank you for this interesting question. Always opting for the (T+C+S) learner is not optimal, as ‘unnecessary’ reweighting can decrease the convergence of the learners (as seen in Fig. 5 & 6). In the extreme case, the training can become very unstable. Thus, it is important to consider the overlap present in the data when choosing/designing the targeted learner. In practice, we recommend using the estimated nuisance functions to inspect overlap (e.g., by visualizing the propensity score, censoring, and survival functions) to aid in deciding on the optimal learner for the data at hand.

  Action: We will discuss the benefits and disadvantages of the different learners and highlight a guideline on how to choose a good option for the data at hand in the camera-ready version of our paper.

• Expected performance on other real-world datasets:

  We initially chose the Twins dataset due to its widespread use in the literature (e.g., [1]), which allows for benchmarking against established findings. This choice was deliberate: by demonstrating that our survival learners recover known effects, we highlight the reliability of our framework, which we consider a key strength. A major benefit of this dataset is the ability to directly compare our results with prior work, and indeed, we observe that our method successfully reproduces these known treatment effects.

  Building on your suggestion, we realized that adding a second dataset from a medical domain could further show the practical value of our survival learners in clinical applications. However, many clinical datasets for survival analysis are not publicly available, which hampers reproducibility. To address this, we specifically sought out a public dataset and chose the ADJUVANT clinical trial [2,3]. This trial, conducted in patients with EGFR-mutant non-small cell lung cancer (NSCLC), compares the efficacy of gefitinib versus chemotherapy (vinorelbine plus cisplatin).

  While previous analyses of the ADJUVANT trial relied on linear survival models [3], we demonstrate how nonlinear survival learners provide new insights. Specifically, we observed more pronounced effects on disease-free survival (DFS) time in patient subgroups with co-alterations in TP53 and CDK4 or in NKX2-1 and CDK4 across all learners. Additionally, patient subgroups with co-amplification of NKX2-1 and SMAD4 showed a benefit in overall survival (OS) time when treated with gefitinib. Overall, we observe notable differences compared to the conventional linear models reported in [3], which can be attributed to the greater flexibility of our toolbox in handling nonlinear relationships. This improved accuracy may provide a more nuanced understanding of the relative benefit of cetuximab versus bevacizumab in metastatic lung cancer and thus may help in new insights for personalized cancer care.

  Action: In the revised version, we will explicitly explain our motivation for using established benchmark datasets to demonstrate reliability. While we are unfortunately not permitted to include new figures for the ADJUVANT dataset during the rebuttal phase (as per NeurIPS guidelines), we will incorporate Kaplan-Meier plots and related visualizations in the final version to illustrate survival differences captured by our learners.

[1] Christos Louizos et al. “Causal effect inference with deep latent-variable models”. NeurIPS. 2017.

[2] Zhong, W.-Z. et al. Gefitinib versus vinorelbine plus cisplatin as adjuvant treatment for stage II-IIIA (N1-N2) EGFR-mutant NSCLC (ADJUVANT/CTONG1104): a randomised, open-label, phase 3 study. Lancet Oncol. 19, 139–148 (2018)

[3] Data access details: https://www.nature.com/articles/s41467-023-41011-4

Thank you! We sincerely appreciate your time and effort in reviewing our work. We will make sure to incorporate your points into the camera-ready version if case of acceptance.

Thank you for your positive evaluation of our manuscript and the helpful comments. We will take all your comments and suggestions to heart and make improvements to our manuscript accordingly. Below, we provide answers to all the questions and concerns you raised. Will will incorporate all points marked with Action into the camera-ready version of our paper.

1. Thank you for pointing out the ambiguity related to the term censoring. Upon reading your comment, we realized that we should be more explicit that we deal with the problem of right censoring, which is very common in survival analysis settings. We thus assume that events have not happened before time t=0.

   Action: We will update the wording to emphasize that we are focusing on right censoring in the camera-ready version of our manuscript.

2. We are happy to follow your advice and also provide proof sketches in the main text, thereby fostering the reader's understanding.

   Action: We will include brief sketches of all proofs in the main section of the camera-ready version of our manuscript.

3. Thank you for this suggestion!

   Action: We will rephrase the respective parts in the text to emphasize that we are the first to provide a unified perspective on orthogonal HTE estimation in survival settings.

4. This is a very relevant question. For practitioners analyzing real-world data, we recommend to estimate the propensity score as well as the survival and censoring functions to inspect potential overlap violations. Then, depending on the type of overlap violation found (e.g., low treatment and censoring overlap), an appropriate weighting combination can be selected, and an appropriate orthogonal survival learner from our toolbox can be trained that accounts for said overlap violations. Finally, in the prediction step we recommend plotting the estimated overlap weights as a measure of “uncertainty” to inspect whether the predictions of the model can be trusted. This overall procedure is consistent with best-practices in standard causal inference using the propensity score.

Action: We will add a note stating the difficulties for HTE estimation on real-world data with complete overlap violations to the camera-ready version of our paper.

5. Thank you. We will fix the typo!

Thank you for the constructive feedback. We will take all your comments and suggestions to heart, and we will make improvements to our manuscript accordingly. Below, we provide answers to all the questions and concerns you raised.

Changes to our paper. In the camera-ready version of our paper, we will make the following key changes: (1) We will report plots of the survival and censoring curves in our real-world study to help contextualize the overlap for the reader. Unfortunately, the rebuttal guidelines prevent us from directly showing our revised plots during the rebuttal period. Where possible, we provide tables below, which will then convert to figures and be added to our paper. (2) We will fix the formatting issues in Fig. 4 (thank you for pointing this out!).

Novelty. Thank you for giving us the opportunity to elaborate on our contribution. While it is true that we leverage tools from existing literature to derive our learners (e.g., orthogonal learning), the following contributions are novel and non-trivial:

1. To the best of our knowledge, we are the first to extend the idea of overlap weighted learners beyond propensity scores. The survival setting is special as it allows for different sources of overlap beyond treatment overlap, but we suspect that such overlap-weighted learners could be beneficial in various other settings such as IV or causal mediation settings that also yield different sources of overlap. Hence, our paper potentially paves the way for substantial method innovation in causal inference.
2. The mathematical derivations of our learners are non-trivial as they required the construction of the efficient influence function for the overlap weighted averaged target quantity (see derivations in Appendix H).
3. We provide a complete toolbox of learners as well as various extensions (e.g., to continuous time, marginalized effects, ties, different estimands) in the appendix. Our results are readily applicable for practitioners in e.g., medicine, where accounting for censoring is crucial.

Additionally, we would like to emphasize that a large part of the literature in causal inference proposes influence-function based estimators in different settings and thus leverage existing tools to derive novel estimands [e.g., 2, 3, 4, 5]. Our novelty is thus in line with established works in that we leverage existing mathematical tools to derive novel results.

Reweighting justification:

Thank you. We would like to clarify that overlap-weighting does precisely align with your intuition: in low-overlap regions, estimation becomes indeed extrapolation and hence predictions are downweighted to not disturb the overall estimation. You may think about this as an adaptive version of your idea to “restrict estimation to regions with large overlap” via e.g., clipping. Advantages of our learners over clipping include: (1) there is no need for choosing a cutoff value, which is otherwise often done heuristically, and (2) the weighting term is considered in the orthogonal objective, which makes the learner more robust to estimation errors in the weighting function. Clipping on the other hand can be insensitive to errors in e.g., the propensity score. (3) In contrast to clipping, weighting is a smooth operation which allows for deriving an influence function and thus enables orthogonal learning in the first place (pathwise differentiability).

In standard causal inference, the R-learner can be seen as a propensity-overlap weighted learner, and is widely applied in practice. Our work enables applying such learners in survival analysis.

Action: We will include the discussion above in the camera-ready version of our manuscript.

Overlap tradeoffs:

Thank you for this important question. First, note that downweighting low-overlap samples does not introduce additional population bias for the weighted estimand if the model is correctly specified (see Theorem 5.2). In finite samples, overlap weighting might increase MSE in low-overlap regions as these regions are essentially "sacrificed" to reduce overall estimation variance. However, this is deliberate and similar in behaviour to established causal inference methods. For example, the same applies to fitting an R-learner in the standard (non-survival) setting, or fitting a CATE model using propensity clipping. In practice, it is always recommended to plot the weighting function (e.g., propensity or censoring overlap) as a measure of “trustworthiness” (or uncertainty) of the model predictions. Predictions in low-overlap regions may then be discarded or deferred to domain experts, while the model predictions in large-overlap regions benefit from weighting and may be leveraged.

Action: We will include the discussion and practical guidelines on how to interpret the results in low-overlap regions in the camera-ready version of our paper.

Estimand relevance:

Thank you for this thoughtful question. First, we would like to clarify that the benefit of our method is not restricted to settings with overall low survival or censoring overlap. Instead, overlap weighting can be beneficial even in settings in which a small subgroup of the population is affected by low overlap. For example, in a medical study, older participants may have a large probability of dropping out within a short time after the start of the study (low censoring overlap). Even if the majority of the population is younger, the small subpopulation of older people may yield a large variance when applying the standard DR-learner due to divisions by extremely small censoring functions (see e.g., Eq. 13). Overlap weighting focusses estimation accuracy on the younger population, as this is the population for which we have available outcome data and for which we can obtain causal insights.

Additionally, based on our collaborations in medicine, we frequently encounter settings where estimating τₜ(x) with low survival or censoring overlap is both common and clinically relevant.

(1) Clinical relevance:

Estimating τₜ(x) despite low survival overlap is important in many oncology settings. For example, the rapid adoption of new targeted therapies in advanced non-small cell lung cancer recently led to immediate uptake in some subgroups, resulting in limited data for others. Similarly, low censoring overlap may arise when treatments are discontinued due to toxicity, often linked to very heterogeneous reasons such as co-medications for some patients and genetic biomarkers for others. In both cases, the oncologists from our collaborations seek to generate clinically meaningful evidence about the treatment effectiveness while carefully addressing low survival or censoring overlap.

(2) Empirical relevance:

Low treatment overlap can reflect clear clinical decision rules and may be addressed through changes in the research design (e.g., excluding patients before estimation). In contrast, low survival or censoring overlap often stems from structural limitations, such as small high-risk subgroups or informative censoring. Here, simple fixes to the research design are typically not feasible for clinical researchers because ignoring these regions risks overlooking treatment effects for certain subpopulations. Therefore, developing methods that remain robust under limited survival or censoring overlap is crucial for valid inference in such real-world settings.

Importantly, our orthogonal learners aim to improve treatment effect estimation in survival settings (regardless of whether there is low overlap or not). It is well known in the literature that simple plug-in learners lead to so-called “plug-in bias” [1] and, thus, suboptimal estimation. A remedy is offered by our orthogonal two-stage learners, which are designed to prevent such plug-in bias. (ii) Existing learners lack the ability to address different types of low overlap and, as a result, exhibit a large variance under low censoring or survival overlap. As a remedy, our toolbox allows the specification of a custom weighting function for robust estimation under low treatment-, survival-, and/or censoring overlap.

Finally, we emphasize that we also provide several extensions of our toolbox to different settings and effects (continuous time, marginalized effects, different causal estimands, and unobserved ties) in the supplementary material.

Survival & censoring overlap:

We are happy to report the survival and censoring curves in the twins study to help contextualize the overlap. While we are unfortunately not permitted to include new figures showing the curves during the rebuttal phase (as per NeurIPS guidelines), we now report a summary of the curves. Overall, there is no sign of systematic low survival or censoring overlap on average, i.e., the mean-aggregated curves across all covariate realizations are bounded away from zero at all assessed times t<=30. Across the complete study time (of one year), we observe low overlap towards the end of the study time amplified through various covariates such as tobacco usage throughout the pregnancy for the survival curve. Rare risk factors such as cardiac disease are associated with low censoring-overlap.

[1] Edward H. Kennedy. “Semiparametric doubly robust targeted double machine learning: A review”. arXiv preprint (2022).

[2] Nathan Kallus. "What's the Harm? Sharp Bounds on the Fraction Negatively Affected by Treatment". NeurIPS 2022.

[3] Coston et al. "Counterfactual predictions under runtime confounding". NeurIPS 2020.

[4] Syrgkanis et al. "Machine learning estimation of heterogeneous treatment effects with instruments". NeurIPS 2019.

[5] Lan et al. "A Meta-learner for Heterogeneous Effects in Difference-in-Differences". ICML 2025.

Thank you for your response and for raising your score! We sincerely appreciate your time and effort in reviewing our work. We will make sure to incorporate your points into the camera-ready version if case of acceptance.