Conditional Counterfactual Mean Embeddings: Doubly Robust Estimation and Learning Rates
This work addresses the challenge of estimating conditional counterfactual distributions in causal inference, which is crucial for understanding heterogeneous treatment effects, but it is incremental as it builds on existing kernel and neural network methods.
The paper tackles the problem of understanding heterogeneous treatment effects by characterizing the full conditional distribution of potential outcomes, proposing the Conditional Counterfactual Mean Embeddings (CCME) framework and developing three estimators with finite-sample convergence rates and double robustness, as demonstrated in experiments that accurately recover distributional features like multimodal structure.
A complete understanding of heterogeneous treatment effects involves characterizing the full conditional distribution of potential outcomes. To this end, we propose the Conditional Counterfactual Mean Embeddings (CCME), a framework that embeds conditional distributions of counterfactual outcomes into a reproducing kernel Hilbert space (RKHS). Under this framework, we develop a two-stage meta-estimator for CCME that accommodates any RKHS-valued regression in each stage. Based on this meta-estimator, we develop three practical CCME estimators: (1) Ridge Regression estimator, (2) Deep Feature estimator that parameterizes the feature map by a neural network, and (3) Neural-Kernel estimator that performs RKHS-valued regression, with the coefficients parameterized by a neural network. We provide finite-sample convergence rates for all estimators, establishing that they possess the double robustness property. Our experiments demonstrate that our estimators accurately recover distributional features including multimodal structure of conditional counterfactual distributions.