Conditional Counterfactual Mean Embeddings: Doubly Robust Estimation and Learning Rates

TL;DR

This paper introduces the Conditional Counterfactual Mean Embeddings framework for estimating heterogeneous treatment effects by embedding conditional distributions into RKHS, providing robust estimators with proven convergence rates.

Contribution

It proposes a novel RKHS-based framework for conditional distribution estimation, along with three practical estimators and finite-sample convergence guarantees.

Findings

Estimators accurately recover multimodal conditional distributions.

The proposed methods exhibit double robustness.

Finite-sample convergence rates are established.

Abstract

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.