Overlap-weighted orthogonal meta-learner for treatment effect estimation over time

TL;DR

This work tackles the challenge of estimating heterogeneous treatment effects over time when treatment overlap decays with horizon. It introduces the overlap-weighted orthogonal (WO) meta-learner, a model-agnostic approach that up-weights high-overlap regions and minimizes a Neyman-orthogonal weighted population risk to obtain robust CATE and CAPO estimates. The authors prove that the WO risk minimizes a weighted oracle risk and is Neyman-orthogonal with respect to all nuisance functions, mitigating first-order bias from nuisance estimation. Empirical results across synthetic, semi-synthetic, and real-world datasets show that WO consistently outperforms existing meta-learners, especially under low overlap and complex nuisance specifications, validating its potential for reliable decision-making in longitudinal personalized medicine.

Abstract

Estimating heterogeneous treatment effects (HTEs) in time-varying settings is particularly challenging, as the probability of observing certain treatment sequences decreases exponentially with longer prediction horizons. Thus, the observed data contain little support for many plausible treatment sequences, which creates severe overlap problems. Existing meta-learners for the time-varying setting typically assume adequate treatment overlap, and thus suffer from exploding estimation variance when the overlap is low. To address this problem, we introduce a novel overlap-weighted orthogonal (WO) meta-learner for estimating HTEs that targets regions in the observed data with high probability of receiving the interventional treatment sequences. This offers a fully data-driven approach through which our WO-learner can counteract instabilities as in existing meta-learners and thus obtain more reliable HTE estimates. Methodologically, we develop a novel Neyman-orthogonal population risk function that minimizes the overlap-weighted oracle risk. We show that our WO-learner has the favorable property of Neyman-orthogonality, meaning that it is robust against misspecification in the nuisance functions. Further, our WO-learner is fully model-agnostic and can be applied to any machine learning model. Through extensive experiments with both transformer and LSTM backbones, we demonstrate the benefits of our novel WO-learner.

Figures (7)

• Figure 1: Overlap problems in time-varying settings. With longer prediction horizons, treatment sequences become more complex, and the probability of observing each treatment sequence decreases exponentially (low overlap). Hence, standard meta-learners often have a poor performance due to low overlap.
• Figure 2: Causal graph. Shown are the observed history, the interventional treatment sequence, and the potential outcome, along with their causal connections. When intervening on future treatments, we do not observe the future covariates (in purple), which leads to time-varying confounding.
• Figure 3: Weight functions. Our weight functions are designed to up-weight samples with large treatment overlap (CATE), or treatment propensity (CAPO). Thereby, we effectively overcome estimation variance issues in low-overlap and low-propensity regimes in a fully data-driven way.
• Figure 4: Complex propensity score$\mathcal{D}^\pi$: Reported are the average RMSEs for CATE estimation. Our WO-learner significantly improves upon existing meta-learners, even for large prediction horizons.
• Figure 5: Complex response function$\mathcal{D}^\mu$: Reported are the average RMSEs for CATE estimation for increasing dimensionality of the covariate space. Our WO-learner clearly outperforms the non-orthogonal RA learner.

Theorems & Definitions (35)

• Definition 3.1: Response functions and propensity scores
• Definition 3.2: DR pseudo-outcomes Frauen.2025
• Definition 4.1: Weight functions
• Definition 4.2: WO pseudo-outcomes
• Theorem 4.3: Weighted population risk
• proof
• Corollary 4.4: Time-varying adjustment
• proof
• Theorem 4.5: Neyman-orthogonality
• proof