Combining T-learning and DR-learning: a framework for oracle-efficient estimation of causal contrasts

1. Problem Statement

The paper addresses the challenge of estimating heterogeneous causal contrasts, such as the Conditional Average Treatment Effect (CATE) and Conditional Relative Risk (CRR), using machine learning. The goal is to find a function $\theta_0$ that minimizes a population risk function $R_0(\theta)$ over a function class $\mathcal{F}$.

Existing methods face a trade-off between stability/convexity and statistical efficiency:

• T-learners (Plug-in): Estimate nuisance parameters (outcome regression $\mu$) and plug them directly into the contrast formula. They are stable and use convex loss functions but suffer from poor efficiency because they do not account for the bias introduced by estimating nuisance parameters.
• Neyman-Orthogonal Learners (DR-learner, R-learner): Use debiased loss functions (often involving inverse probability weighting) to achieve oracle efficiency (fast convergence rates even with flexible nuisance estimators). However, they suffer from two major drawbacks:
  1. Non-convexity: The resulting loss functions can be non-convex, making optimization difficult for standard algorithms (e.g., random forests, gradient boosting).
  2. Instability: They rely on pseudo-outcomes that involve inverse propensity weighting. When propensity scores are close to 0 or 1, these pseudo-outcomes can take extreme values, leading to unstable estimates and violations of known bounds (e.g., CATE estimates outside $[-1, 1]$ for binary outcomes).

2. Methodology: Efficient Plug-in (EP) Learning

The authors propose EP-learning, a framework that combines the stability of plug-in estimators with the efficiency of orthogonal learning.

Core Concept

Instead of minimizing a debiased loss function directly (as in DR/R-learners), EP-learning constructs a debiased plug-in estimator of the risk function.

1. Efficient Influence Function: The population risk $R_0(\theta)$ is pathwise differentiable. Its efficient influence function allows for the construction of a one-step debiased estimator.
2. Implicit Debiasing: The key innovation is constructing a refined outcome regression estimator, $\mu^*_n$, such that the debiasing term in the one-step estimator becomes negligible (asymptotically zero) across the function class $\mathcal{F}$.
   • Standard plug-in: $R_n(\theta) = \frac{1}{n}\sum L_{\mu_n}(\theta)$.
   • One-step debiased: $R_{n, \pi, \mu}(\theta) = \frac{1}{n}\sum [L_{\mu_n}(\theta) + \Delta_{\pi, \mu}(\theta)]$.
   • EP-learning: Constructs $\mu^*_n$ such that $\frac{1}{n}\sum \Delta_{\pi, \mu^*_n}(\theta) \approx 0$. Thus, the simple plug-in risk $\frac{1}{n}\sum L_{\mu^*_n}(\theta)$ becomes asymptotically equivalent to the efficient one-step estimator.

Algorithmic Steps

The EP-learner follows a two-stage process (often with cross-fitting):

1. Initial Estimation: Obtain initial estimates of the propensity score ($\pi_n$) and outcome regression ($\mu_n$) using flexible machine learning (e.g., Random Forests, XGBoost).
2. Sieve-Based Adjustment (Debiasing): Construct $\mu^*_n$ by adjusting $\mu_n$ within a sieve space (a finite-dimensional approximation of the function space).
   • The adjustment solves a weighted regression problem (often logistic or linear) where the target is to satisfy the score equations that make the debiasing term zero for the basis functions of the sieve.
   • Specifically, $\mu^*_n$ is constructed to satisfy:
     $$\frac{1}{n} \sum_{i=1}^n \frac{1}{\pi_n(A_i|W_i)} \sum_{m} H_{m, \mu^*_n}(A_i, W_i) \psi(W_i) \{Y_i - \mu^*_n(A_i, W_i)\} = 0$$
     for all basis functions $\psi$ in the sieve.
3. Final Estimation: Minimize the plug-in risk using the refined estimator $\mu^*_n$:
   $$\hat{\theta}_n = \arg\min_{\theta \in \mathcal{F}} \frac{1}{n} \sum_{i=1}^n L_{\mu^*_n}(\theta, W_i)$$

3. Key Contributions

1. Novel Framework: Introduction of EP-learning, which achieves oracle efficiency without explicitly minimizing non-convex, unstable Neyman-orthogonal loss functions.
2. Theoretical Guarantees:
   • Oracle Efficiency: Under regularity conditions (smoothness of nuisance estimators and the action space), the EP-learner is asymptotically equivalent to the minimizer of an oracle-efficient one-step debiased estimator.
   • Double Robustness: The risk estimator is doubly robust; it remains consistent if either the propensity score or the outcome regression is consistently estimated.
   • Convergence Rates: The paper establishes that ERM-based EP-learners achieve the minimax optimal rate (oracle rate) for the causal contrast, provided the nuisance estimators converge at sufficient rates.
3. Resolution of Practical Issues:
   • Convexity: By using the plug-in principle with $\mu^*_n$, the resulting loss function retains the convexity of the original loss (e.g., squared error for CATE, logistic loss for CRR), allowing the use of standard, stable supervised learning algorithms.
   • Stability: The pseudo-outcomes used in the final stage are bounded (or at least well-behaved) because they are derived from $\mu^*_n$ rather than raw inverse-weighted residuals, preventing the extreme values seen in DR-learners.
4. Specific Applications: The framework is instantiated for CATE (using squared error loss) and CRR (using logistic loss), providing the first doubly robust, efficient estimator for log-CRR that corresponds to a convex loss function.

4. Results

The authors validate the method through extensive simulations and theoretical analysis:

• Simulation Experiments:
  • CATE Estimation: EP-learners significantly outperformed DR-learners and R-learners when using tree-based methods (Random Forests, XGBoost), particularly in terms of Mean Squared Error (MSE).
  • Stability: In scenarios with limited overlap (extreme propensity scores), DR-learners produced unstable estimates and failed to capture the true CATE shape due to extreme pseudo-outcomes. EP-learners remained stable and accurate.
  • CRR Estimation: For Conditional Relative Risk, the orthogonal loss is non-convex and often yields negative weights or outcomes outside $[0,1]$, rendering standard logistic regression software unusable. EP-learners successfully used standard logistic regression, achieving better performance than T-learners and IPW-learners.
• Comparison with State-of-the-Art:
  • EP-learners matched or exceeded the performance of Causal Forests and R-learners across various settings (different sample sizes, dimensions, and smoothness levels).
  • They demonstrated superior robustness to tuning parameter choices (e.g., tree depth) compared to DR-learners.

5. Significance

This paper bridges a critical gap in causal inference methodology:

• Practical Usability: It enables the use of powerful, flexible machine learning algorithms (like Random Forests and Gradient Boosting) for causal effect estimation without the computational and statistical pitfalls of non-convex, unstable orthogonal losses.
• Theoretical Rigor: It proves that "plug-in" methods can be made oracle-efficient through careful construction of the nuisance estimator, challenging the notion that one must sacrifice efficiency for stability.
• Broad Applicability: The framework is general, applicable to a wide class of causal contrasts (including CATE, CRR, and log-odds ratios) and can be extended to continuous treatments and longitudinal settings.
• Implementation: The authors provide an open-source R package (hte3) implementing these methods, facilitating adoption by practitioners.

In summary, EP-learning offers a "best of both worlds" solution: it inherits the stability and convexity of T-learners while achieving the statistical efficiency and double robustness of DR/R-learners.