Outcome-adapted Automatic Debiased Machine Learning

Technical Summary: Outcome-adapted Automatic Debiased Machine Learning

1. Problem Statement

In causal inference and econometrics, parameters of interest such as treatment effects (ATE) or policy effects are often expressed as linear functionals of an outcome regression function, $\gamma(x) = E[Y|X=x]$. While "plug-in" estimators using machine learning (ML) for $\gamma$ are intuitive, they often suffer from significant bias and fail to achieve $\sqrt{n}$-consistency.

To address this, Automatic Debiased Machine Learning (AutoDML) (Chernozhukov et al., 2022b, 2024) was introduced. This framework constructs asymptotically normal estimators by simultaneously estimating the outcome regression $\gamma$ and a Riesz representer $\alpha$ (a nuisance function derived from the functional's structure). Existing AutoDML architectures, such as RieszNet and MADNet, utilize a shared intermediate covariate representation $Z$ to estimate both $\gamma$ and $\alpha$.

The central problem addressed in this paper is the optimality of the shared representation. It remains unclear whether the shared representation $Z$ should be predictive of the Riesz representer $\alpha$ (as traditionally assumed in RieszNet) or the outcome $Y$. The authors investigate whether a representation that prioritizes outcome prediction while discarding information about the Riesz representer yields superior estimation efficiency.

2. Methodology and Theoretical Framework

Theoretical Results on Efficiency:
The authors establish two primary theoretical results regarding covariate representations $Z = h(X)$:

• Theorem 6 (Outcome-Adapted Efficiency): If a representation $Z$ preserves the predictive power of the outcome (i.e., $\gamma_{h,P}(Z) = \gamma_P(X)$ and $Var_P[Y|Z] = Var_P[Y|X]$), the resulting AutoDML estimator has an asymptotic variance $V_{h,P}$ that is smaller than or equal to the variance $V_P$ of the estimator using the full covariates $X$. Specifically, the variance reduction is proportional to $E_P[Var_P[\alpha_P(X)|Z] Var_P[Y|Z]]$. To maximize efficiency, $Z$ should be maximally predictive of $Y$ but minimally predictive of $\alpha_P(X)$.
• Theorem 8 (Riesz-Adapted Inefficiency): Conversely, for a broad class of functionals (including ATE and average policy effects), if a representation preserves the Riesz representer (i.e., $\alpha_{h,P}(Z) = \alpha_P(X)$), the asymptotic variance $V_{h,P}$ is generally larger than or equal to $V_P$. This suggests that learning a representation specifically for the Riesz representer can lead to efficiency losses.

The Outcome-Adapted AutoDML Estimator:
Motivated by Theorem 6, the authors propose the outcome-adapted AutoDML estimator. The procedure involves:

1. Sample Splitting: Partitioning data into sets for representation learning and final estimation.
2. Representation Learning: Estimating a representation $Z = \hat{h}(X)$ exclusively through an outcome regression step on the first data split.
3. Nuisance Estimation: Using the learned representation $Z$ to estimate both the outcome regression $\hat{\gamma}$ and the Riesz representer $\hat{\alpha}$ (via Riesz regression) on the same split.
4. Final Estimation: Computing the parameter estimate on the held-out data split using the doubly robust estimating function.

Neural Network Implementation:
The authors implement this using a neural network architecture similar to RieszNet but with a distinct training protocol:

• Architecture: A shared "trunk" (hidden layers) maps covariates $X$ to a representation $Z$, followed by separate branches for the outcome ($\gamma$) and the Riesz representer ($\alpha$).
• Training Strategy:
  1. Step 1: Optimize the shared trunk and the outcome branch by minimizing the Mean Squared Error (MSE) loss.
  2. Step 2: Freeze the shared trunk parameters and optimize only the Riesz branch parameters by minimizing the Riesz loss.
• Information Bottleneck: To ensure $Z$ is minimally predictive of $\alpha$, the authors introduce dimensionality reduction techniques:
  • Simple Dimensionality Reduction: Enforcing a small fixed dimension for $Z$.
  • Adaptive Dimensionality Reduction: Applying a group LASSO penalty on the weights of the $Z$-layer to encourage sparsity and effectively zero out neurons, reducing the dimensionality of $Z$.

3. Key Contributions

1. Theoretical Characterization: The paper proves that for AutoDML, the optimal shared representation is outcome-adapted (predictive of $Y$, not $\alpha$), contrary to the intuition behind existing methods like RieszNet which aim for Riesz-adapted representations.
2. New Estimator: Introduction of the outcome-adapted AutoDML estimator, which learns representations via outcome regression first, followed by Riesz regression on the fixed representation.
3. Neural Network Implementation: A specific two-step training procedure for neural networks that enforces the outcome-adapted property and utilizes information bottlenecks (via dimensionality reduction) to maximize efficiency gains.
4. Re-evaluation of RieszNet: The authors provide a theoretical explanation for the empirical success of RieszNet, hypothesizing that its performance stems from the specific choice of loss weights (specifically a small weight on the Riesz loss and the inclusion of TMLE loss) which inadvertently creates an outcome-adapted representation.

4. Experimental Results

The authors validate their method on synthetic and semi-synthetic datasets:

• Synthetic Data (Proof of Concept): In a setting where the true optimal representation is known, the outcome-adapted estimator with adaptive dimensionality reduction achieved the theoretical minimum variance, behaving as if it had oracle knowledge of the optimal representation. In contrast, estimators using separate networks or Riesz-adapted representations showed higher variance.
• IHDP Benchmark: On the semi-synthetic Infant Health and Development Program (IHDP) dataset, the outcome-adapted neural network with adaptive dimensionality reduction achieved state-of-the-art performance (lowest Mean Absolute Error) compared to RieszNet, MADNet, and separate neural networks.
• Average Shift Effect: The outcome-adapted estimator with adaptive dimensionality reduction consistently demonstrated the lowest variance and Mean Squared Error (MSE) across various sample sizes.
• Mean Missing Outcome (Adversarial Setting): In a challenging setting with extreme propensity scores (Kang & Schafer, 2007), the outcome-adapted estimator with adaptive dimensionality reduction achieved the lowest MSE, driven primarily by reduced squared bias. While Gaussian confidence intervals had low coverage in this difficult setting, bootstrap confidence intervals (specifically percentile bootstrap) significantly improved coverage, reaching nominal levels for large sample sizes.

5. Significance and Claims

The paper claims that outcome-adaptation is the key property for achieving estimator efficiency in the AutoDML framework. The authors argue that:

• Learning a representation that is predictive of the outcome but discards information about the Riesz representer leads to asymptotically more efficient estimators.
• The proposed outcome-adapted AutoDML estimator is a generalization of the debiased outcome-adapted propensity estimator (DOPE) to the broader AutoDML framework.
• The proposed method achieves state-of-the-art accuracy on standard benchmarks without requiring complex tuning of trade-off parameters between multiple objectives (unlike RieszNet, which requires balancing MSE, Riesz, and TMLE losses).
• The results suggest that the success of existing methods like RieszNet may be due to their implicit learning of outcome-adapted representations rather than Riesz-adapted ones.

The authors conclude that while their method offers efficiency gains, the representation error must be controlled (e.g., via variable selection or sufficient sample sizes) to ensure valid inference, and they recommend the use of bootstrapping for confidence intervals in challenging settings where representation error may not be asymptotically negligible.