Nonparametric Estimation of Heterogeneous Treatment Effects: From Theory to Learning Algorithms

Alicia Curth; Mihaela van der Schaar

Source-linked AI summary

Nonparametric Estimation of Heterogeneous Treatment Effects: From Theory to Learning Algorithms

Alicia Curth, Mihaela van der Schaar

arXiv:2101.10943v2Jan 26, 2021stat.MLcs.LG

TL;DR

Flexible CATE estimation requires choosing among model-agnostic meta-learners despite unobserved counterfactuals. The paper analyzes four strategy classes theoretically and implements them with neural-network architectures, finding that learner and architecture performance depends on the DGP and sample size rather than one method dominating universally.

Problem
Choosing among many CATE meta-learners is difficult because counterfactual information needed for direct model selection is unavailable.
Method
The paper theoretically analyzes plug-in and pseudo-outcome regression strategies and implements them with neural-network architectures.
Results
Learner and architecture performance depends on the underlying DGP and available data: the DR-learner is asymptotically optimal, while RA- and information-sharing plug-in learners can outperform it in finite samples.
Takeaways & Limitations
Expert knowledge about the DGP and sample size should inform the practical choice of learner and architecture.
Takeaways & Limitations
The analysis excludes the R-learner because it requires specialized model-fitting procedures rather than arbitrary machine-learning methods.

Abstract

from arXiv · show

The need to evaluate treatment effectiveness is ubiquitous in most of empirical science, and interest in flexibly investigating effect heterogeneity is growing rapidly. To do so, a multitude of model-agnostic, nonparametric meta-learners have been proposed in recent years. Such learners decompose the treatment effect estimation problem into separate sub-problems, each solvable using standard supervised learning methods. Choosing between different meta-learners in a data-driven manner is difficult, as it requires access to counterfactual information. Therefore, with the ultimate goal of building better understanding of the conditions under which some learners can be expected to perform better than others a priori, we theoretically analyze four broad meta-learning strategies which rely on plug-in estimation and pseudo-outcome regression. We highlight how this theoretical reasoning can be used to guide principled algorithm design and translate our analyses into practice by considering a variety of neural network architectures as base-learners for the discussed meta-learning strategies. In a simulation study, we showcase the relative strengths of the learners under different data-generating processes.

1 INTRODUCTION

The paper addresses the growing need for flexible CATE estimation and the difficulty of choosing among many meta-learners when counterfactual outcomes are unobserved. It combines theoretical analysis with neural-network implementations to clarify when different learners may perform better and how DGP knowledge can guide model choice.

CATE estimation supports flexible investigation of treatment-effect heterogeneity for personalized policy and treatment design.
Selecting among CATE algorithms is difficult because the ground-truth treatment effects are unobserved.
The paper theoretically compares four broad, model-agnostic meta-learning strategies and translates the analysis into neural-network architectures.
The practical goal is not to identify one universally superior learner, but to use expert knowledge about the DGP to narrow algorithm choices and guide model design.
The proposed taxonomy distinguishes one-step plug-in learners from three two-step learners based on regression adjustment, propensity weighting, or doubly robust pseudo-outcomes.

2 PROBLEM DEFINITION

The paper defines CATE within the Neyman–Rubin potential-outcomes framework and formulates its estimation from observational data as a nonparametric problem. Identification relies on consistency, unconfoundedness, and overlap.

CATE is the expected treatment effect for an individual with covariate values X = x.
Only one potential outcome is observed for each individual, creating the fundamental counterfactual problem of causal inference.
Under consistency, unconfoundedness, and overlap, CATE equals the difference between conditional outcome regressions, τ(x) = µ1(x) − µ0(x).
The estimation problem assumes no known parametric form for CATE, nuisance parameters, or error distributions.
The meta-learners can use generic nonparametric regressors, with neural networks used in the experiments and other methods such as random forests also applicable.

3 CATEGORIZING CATE META-LEARNERS

The paper classifies CATE meta-learners by whether they difference estimated outcome regressions or regress unbiased pseudo-outcomes, yielding plug-in, RA, PW, and DR strategies. This taxonomy organizes their statistical analysis and practical comparison.

One-step plug-in learners estimate separate potential-outcome regressions and obtain CATE by differencing them.
S- and T-learners are plug-in special cases distinguished by whether nuisance-estimation tasks share information.
Two-step learners first estimate nuisance parameters and then regress a pseudo-outcome on covariates to estimate CATE directly.
The RA-learner uses a regression-adjusted pseudo-outcome, while the X-learner performs separate regressions before combining its CATE estimates.
The PW-learner uses a pseudo-outcome based on the Horvitz–Thompson transformation, whereas the DR-learner uses a doubly robust pseudo-outcome.
The R-learner is excluded because it requires specialized model-fitting procedures rather than arbitrary off-the-shelf machine-learning methods.

4 THEORETICAL ANALYSES OF CATE META-LEARNERS

The paper compares four CATE meta-learner strategies theoretically, deriving conditions for their relative performance under smoothness and overlap assumptions. It also explains how finite-sample variance, selection bias, and second-stage regularization affect practical behavior.

Theoretical scope: The analysis compares one-step plug-in learners with RA-, PW-, and DR-learners using asymptotic and finite-sample arguments.The RA, PW, and DR strategies use regression adjustment, propensity weighting, or both through pseudo-outcomes.
Assumptions: The theory assumes bounded nuisance functions and propensity scores, smooth outcome regressions, propensity scores, and treatment effects estimated at nonparametric minimax rates.Propensity estimates are bounded away from zero and one by δ, while the true propensity score satisfies an overlap condition involving ω.
Asymptotic behavior: Asymptotically, the DR-learner is preferred among two-step learners, while PW can outperform RA when propensity scores are easier to estimate than outcome regressions.RA and plug-in learners are expected to perform similarly when treatment effects are simpler than potential-outcome regressions.
Asymptotic behavior: Known propensity scores in experiments make PW- and DR-learner remainder terms exactly zero, yielding the oracle rate.This conclusion applies when propensity scores are known rather than estimated.
Finite-sample behavior: Finite-sample performance can favor RA because inverse-propensity weighting has high variance, especially with extreme or small propensity scores.The PW pseudo-outcome remains highly variable even when propensity scores are constant and known; small δ can also enlarge PW and DR error-bound constants.
Finite-sample behavior: Selection bias may be severe in finite samples when overlap is weak, whereas second-stage regression can correct sub-optimal first-stage weighting and regularize finite-sample bias.The paper characterizes second-stage regression as removing some bias from regularization and overfitting, while noting that estimating more parameters can increase variance.

5 CATE ESTIMATION USING NEURAL NETWORKS

The practical analysis instantiates CATE meta-learners with neural-network architectures that either separate or share information across nuisance-estimation tasks. The proposed SNet architecture flexibly interpolates between task sharing and task-specific modeling.

Neural-network implementations: A separate network for each regression task, called TNet, allows arbitrarily different regression surfaces and is a natural asymptotic implementation.TNet corresponds to the simplest neural-network implementation of the learners.
Neural-network implementations: Sharing information between nuisance-estimation tasks can be more efficient in finite samples when the potential-outcome regressions are similar.This motivates architectures based on representation learning rather than fully separate networks.
Existing shared architectures: TARNet, DragonNet, and DR-CFR share representations to different extents across potential-outcome and propensity-score tasks.They are referred to as SNet-1, SNet-2, and SNet-3 in the paper.
Underlying assumptions: The architecture choices encode assumptions about common feature spaces, propensity-score representation, and features affecting only one potential-outcome function.The last case is motivated by distinguishing prognostic markers from predictive markers in medical applications.
Proposed architecture: The proposed SNet learns five representations, allowing potential-outcome regressions to use subsets of shared features while encompassing existing SNet architectures and TNet as special cases.Changing the widths of shared and outcome-specific representations moves the architecture toward TNet or more strongly shared designs.
Practical considerations: In practice, SNet architectures can estimate nuisance parameters for two-step learners, and using all data for both stages may outperform sample splitting in small samples.The latter observation differs from the sample-splitting setup used in the theoretical analysis.

6 EXPERIMENTS

Experiments compare plug-in architectures and meta-learners across synthetic DGPs and the adapted IHDP benchmark, showing that performance depends on outcome structure, confounding, imbalance, and architecture choice. The results support using DGP knowledge to guide learner and model selection.

Experimental design: Synthetic experiments evaluate learner performance across stylized DGPs, sample sizes, predictive-feature counts, treatment proportions, and learner–architecture combinations.Performance is measured by RMSE of τ(x), with synthetic evaluations averaging results over independently generated test observations and repeated runs.
Plug-in architectures: S-architectures improve on TNet in small samples when the two potential-outcome functions share structure, while the flexible SNet performs best on average.SNet is also the only architecture outperforming TNet when the two potential-outcome functions are very different, and its relative strength is clearest at larger sample sizes.
Meta-learners: DR-learner performs substantially best under confounding without treatment effect, RA-learner performs best with confounding and a non-trivial effect, and plug-in learners perform best when CATE is more complex than either potential outcome.PW-learner performs poorly because its pseudo-outcome has low signal-to-noise ratio and high variance; improving it would require stronger regularization and smaller networks.
Feature complexity: As predictive features for τ(x) increase, S-architectures degrade relative to TNet and DR-learner loses its advantage as CATE becomes less sparse.The performance gap between TNet and RA-learner remains virtually constant in this interpolation experiment at n = 2000.
Treatment imbalance: Under treatment imbalance, information sharing adds more value for plug-in architectures, while RA-learner outperforms T-learner only under moderate to no imbalance.The performance gap between T-learner and DR-learner is not impacted by the treated proportion in this setting.
IHDP benchmark: On adapted IHDP, information sharing significantly improves plug-in performance over TNet, but simpler SNet-1 and SNet-2 outperform more complex SNet-3 and SNet models.Except for RA-learner applied to the best-performing plug-in model, two-step learners underperform plug-in learners, attributed to the complexity of simulated τ(x).

7 CONCLUSION

The paper finds that meta-learner and neural-network performance depends on theory, sample size, and the underlying data-generating process, rather than one universally superior learner.

The DR-learner is asymptotically optimal, but RA- and plug-in learners sharing information can outperform it in finite samples.Finite-sample performance differs from the asymptotic ranking.
Sophisticated nuisance-estimator architectures can improve small-sample performance for two-step learners.
The relative performance of learners and architectures depends on the data-generating process and the available sample size.
The authors recommend incorporating expert assessment of the likely data-generating process when choosing a learner in practice.
The study leaves the optimal choice of underlying machine-learning method, such as neural networks versus random forests, for future work.

SUPPLEMENTARY MATERIALS

The supplementary material revisits the theoretical assumptions and extends the analysis from smoothness to sparsity, while examining assumptions behind pseudo-outcome regression.

The supplementary section revisits assumptions made in the main theoretical analysis.
It analyzes estimation under assumed sparsity instead of smoothness.
It discusses assumptions associated with Kennedy’s theorem on pseudo-outcome regression.

A.1 Additional Analyses on Minimax Performance Using Assumptions on Sparsity

The supplementary minimax analysis replaces smoothness with sparsity assumptions to compare learner performance and derive conditions favoring different estimators.

Motivation: The analysis uses minimax error rates to quantify the difficulty of nonparametric estimation problems under sparsity.
Sparsity assumptions: Additive sparsity represents functions as sums of components, each depending on a small number of predictors.
Sparsity assumptions: Under equal component smoothness, dimension, and magnitude, the minimax rate combines a nonparametric estimation term with variable-selection uncertainty.The stated rate is kn^{-2p/(2p+d∗)} + k d∗log(d/d∗)/n.
Assumptions: The analysis uses linear sparse-function assumptions and regularity conditions for second-stage regression estimators to obtain pseudo-outcome error bounds.
Learner performance: Two-step learners can outperform plug-in learners when the treatment effect depends on fewer covariates than each potential outcome function.
Learner performance: Under sparsity, RA and plug-in learners can have equal performance, PW can outperform RA under a propensity-score condition, and DR can attain the oracle rate under a stated dimensionality condition.
Scope: A more general formulation using nonlinear sparse functions could derive broader conditions for when each learner outperforms the others.

B PROOFS

The proof and implementation materials describe two-step learner construction, loss functions, representation architectures, and regularization choices used for the neural-network experiments.

Two-step learner algorithm: Cross-fitting or a single sample split can implement two-step learners while preserving the analyzed theoretical guarantees.
Two-step learner algorithm: Two-step learners estimate nuisance models first, generate pseudo-outcomes, and regress those pseudo-outcomes in a second stage.
SNet architectures: SNet-1 uses a shared representation with separate potential-outcome hypotheses and regularization.
SNet architectures: SNet-2 adds a propensity-score hypothesis function learned from the shared representation.
SNet architectures: SNet-3 separates outcome-only, treatment-only, and confounding representations, while SNet adds treatment-group-specific representations.
Representation regularization: The orthogonalization term penalizes variables entering multiple representations to encourage disentanglement without enforcing a hard decomposition.

C.3 IMPLEMENTATION DETAILS

The implementations use ELU-based neural networks trained with Adam, minibatches, and early stopping. The architectures and regularization settings are specified, with code released in two repositories.

All networks use dense layers with ELU activations, Adam optimization, minibatches of 100, and early stopping after a 30% validation split.
SNet-1 and SNet-2 use three 200-unit representation layers, while SNet-3 uses 150 units in ΦC and 50 in ΦO and ΦW.The SNet configuration uses 100 units in ΦC and ΦW.
The experiments set λ = 0.0001 throughout, use γ = 0 in IHDP and γ = 0.01 in simulations, and implement all models in jax.Sklearn-style implementations are available through the two cited CATENets repositories.

D.1 Simulation set-up

The simulation varies covariate roles, confounding, treatment effects, and sample size across synthetic settings, alongside an adapted IHDP benchmark. A rescaled IHDP variant addresses cross-run differences in CATE scale so RMSE comparisons are more comparable.

Simulation design: The simulation uses d = 25 multivariate-normal covariates divided into disjoint subsets, with training sizes n ∈ {1000, 2000, 5000, 10000} and 10 replicates per size.Each replicate uses a test sample of size 500.
Simulation settings: Settings (i) and (ii) distinguish five confounders affecting outcomes and treatment from five outcome-only covariates.
Treatment assignment: Treatments are generated as Bernoulli variables from propensity scores, with ξ = 3 controlling selection bias and ω adaptively centered by the median in each run.The centering prevents the squared specification from producing a much larger treatment group than control group.
Simulation settings: Setting (i) has no treatment effect, whereas setting (ii) models treatment effects using five additional covariates.
Simulation settings: Setting (iii) removes confounding with π(x) = 0.5 and determines the two potential-outcome functions from non-overlapping covariate sets of dimension 10.
IHDP benchmark and rescaling: The adapted IHDP benchmark derives from a randomized premature-infant intervention study with 25 child- and mother-related covariates, while CATE-scale variation motivates rescaling.For runs with σ_CATE > 1, expected potential outcomes are rescaled while retaining N(0, 1) error terms; Figure D.2 compares CATE variation and SNet-1 RMSE distributions before and after adaptation.