Conditional Counterfactual Mean Embeddings

• CCME is a framework that embeds full counterfactual outcome distributions into reproducing kernel Hilbert spaces to capture complex causal effects beyond mean differences.
• It employs empirical kernel ridge regression and doubly robust estimators to achieve nonparametric consistency and optimal convergence rates under smoothness assumptions.
• The approach unifies methods from causal inference, off-policy evaluation, and Bayesian modeling, supporting high-dimensional and structured outcome spaces.

Conditional Counterfactual Mean Embeddings (CCME) formalize the representation and estimation of counterfactual outcome distributions given observed covariates and interventions, by embedding these distributions into reproducing kernel Hilbert spaces (RKHS). This enables nonparametric modeling, estimation, and hypothesis testing of full conditional counterfactual distributions and distributional treatment effects, extending causal inference beyond means to entire distributions, higher-order moments, and structured or high-dimensional outcomes. The CCME framework unifies kernel mean embedding methods, distributional causal inference, double-robust learning, and off-policy evaluation within a consistent mathematical and algorithmic structure.

1. Formal Definition and Mathematical Framework

CCME generalizes the notion of counterfactual mean embeddings (CME) to the conditional setting. Let $Y^a$ denote the potential outcome under treatment $A = a \in \{0,1\}$, $X$ covariates, and $V = \eta(X)$ a (possibly lower-dimensional) feature of $X$ relevant for stratification. Define an RKHS $\mathcal H_Y$ of functions on $\mathcal Y$ with positive definite kernel $k_Y$. The canonical feature map is $\phi(y) = k_Y(\cdot, y)$.

The CCME for $Y^1$ given $V = v$ is the element

$$\mu_{Y^1|V}(v) := \mathbb E\big[ \phi(Y^1) \mid V = v \big] \in \mathcal H_Y.$$

For a characteristic kernel $k_Y$, the mapping $v \mapsto \mu_{Y^1|V}(v)$ encodes the full conditional distribution of $Y^1 \mid V = v$ (Anancharoenkij et al., 4 Feb 2026).

This construction generalizes to the kernel embedding of the counterfactual distribution $P(Y \mid do(X = x))$, using the mean element

$$\mu_{Y|do(X = x)} := \int k_Y(\cdot, y) dP(Y = y \mid do(X = x)) \in \mathcal H_Y,$$

satisfying $$\langle f, \mu_{Y|do(X = x)} \rangle_{\mathcal H_Y} = \mathbb E[f(Y) \mid do(X = x)]$$ for all $f \in \mathcal H_Y$ (Muandet et al., 2018).

A key distinction is that conditional counterfactual mean embeddings model the distribution of potential outcomes given covariates or projected features $V$, rather than only marginalizing over $X$.

2. Estimation Procedures and Algorithms

2.1 Empirical CCME: Kernel Ridge Regression

Given a sample $\{(x_i, y_i)\}_{i=1}^n$, empirical estimation of $\mu_{Y|do(X = x)}$ or $\mu_{Y|X = x}$ follows via regularized kernel ridge regression (Muandet et al., 2018):

• Construct feature matrices $\Phi = [k_X(x_i, \cdot)]$, $\Psi = [k_Y(y_i, \cdot)]$.
• Form the Gram matrix $K = [k_X(x_i, x_j)]_{i,j=1}^n$, and set $w(x) = (K + n\lambda I)^{-1}k_X(x)$.
• Obtain the estimator:

$$\hat \mu_{Y|do(X = x)} = \Psi w(x) = \sum_{i=1}^n w_i(x) k_Y(y_i, \cdot).$$

2.2 Two-Stage Doubly-Robust CCME Estimator

The two-stage meta-estimator (Anancharoenkij et al., 4 Feb 2026) uses sample-splitting and orthogonalization:

• Learn nuisance functions on a held-out fold:
  • Estimated propensity $\hat \pi(x) \approx P(A = 1 \mid X = x)$,
  • Outcome regression $$\hat \mu_0(x) \approx \mathbb E\big[ \phi(Y) \mid X = x, A = 1 \big]$$.
• Form the doubly robust pseudo-outcome:

$$\hat \xi(Z) = \frac{A}{\hat \pi(X)}\left\{ \phi(Y) - \hat \mu_0(X) \right\} + \hat \mu_0(X).$$

• Regress $\hat \xi(Z)$ in $\mathcal H_Y$ onto $V$ via vector-valued kernel ridge, deep features, or neural-kernel architectures.

Three Practical Estimators (Anancharoenkij et al., 4 Feb 2026):

Name	Feature/Function Class	Key Optimization Formulation
Ridge CCME	$\mathcal H_\Gamma$ (operator-valued RKHS)	$$\arg\min_\mu \frac{1}{n}\sum_i \|\mu(V_i{-}\hat\xi(Z_i))\|_{\mathcal{H}_Y}^2 + \lambda_1\|\mu\|_{\mathcal{H}_\Gamma}^2$$
Deep-Feature CCME	Neural net feature $\psi_\theta(V)$ with linear $C$	$$\arg\min_{\theta, C}\frac{1}{n}\sum_i\|C\psi_\theta(V_i)-\hat\xi(Z_i)\|_{\mathcal{H}_Y}^2 + \lambda_1\|C\|_{\text{HS}}^2$$
Neural-Kernel CCME	Neural net coefficients over $\{\tilde y_j\}$ grid in $\mathcal{Y}$	Minimize finite-dimensional RKHS loss over output coefficients, SGD-optimized

Each instantiates the doubly-robust meta-estimator with corresponding function classes and regularization.

2.3 Alternative and Bayesian Extensions

The Bayesian CCME framework places a GP prior on the function-valued embedding, producing a posterior process for $\mu_{Y|X=x}$ with explicit epistemic uncertainty quantification. In the counterfactual context, Bayesian updates provide closed-form posterior mean and variance expressions for embeddings and resulting quantities (Martinez-Taboada et al., 2022).

3. Distributional Treatment Effects and Hypothesis Testing

CCME enables nonparametric quantification of conditional distributional treatment effects (CoDiTE), extending beyond mean differences (CATE) to distances between entire distributions:

$U_D(x) = D(P_0(x), P_1(x)),$

where $D$ is a probability metric such as MMD. With RKHS $k_Y$, the MMD simplifies to

$$U_{\text{MMD}}(x) = \| \mu_{Y|X,Z=1}(x) - \mu_{Y|X,Z=0}(x) \|_{\mathcal H_Y},$$

providing a natural effect-size measure (Park et al., 2021).

Hypothesis testing for no conditional distributional effect is formulated as testing $$t = \mathbb E_X[ \| \mu_{Y|X,Z=1}(X) - \mu_{Y|X,Z=0}(X) \|_{\mathcal H_Y}^2 ] = 0$$, estimated via plug-in and permutation approaches (Park et al., 2021). The witness function $w_x(y)$ visualizes local discrepancies between conditional densities.

Higher-order effects, such as conditional variance or Gini difference, are accommodated through U-statistic regression:

$$\theta(P_{Y|X=x}) = \mathbb E[ h(Y_1,\ldots,Y_r) \mid X_1 = \dots = X_r = x ],$$

solved via regularized RKHS regression on the $r$-fold product space (Park et al., 2021).

4. Theoretical Properties: Consistency, Rates, and Double Robustness

Consistency and Rates

Both classical and doubly-robust CCME estimators admit nonparametric finite-sample convergence rates in the RKHS norm. Under Sobolev or Gaussian kernels and regularity (smoothness) of conditional densities, the minimax optimal rate is

$n^{-2r/(2r + d_v)}$

(up to log factors), where $r$ denotes the smoothness and $d_v$ the dimensionality of $V$. Additional nuisance terms are added, controlled by $L^2$ risk or RKHS norm risk of propensity and outcome models (Anancharoenkij et al., 4 Feb 2026, Muandet et al., 2018).

Double Robustness

CCME estimators based on doubly-robust pseudo-outcomes remain consistent provided at least one of the first-stage nuisance models (propensity or outcome regression) is consistent, with bias controlled by

$\min\{ R_\pi^2(\hat \pi), R_\mu^2(\hat \mu_0) \}$

where $R_\pi^2$ and $R_\mu^2$ are the $L^2$ errors of nuisance models (Anancharoenkij et al., 4 Feb 2026).

Theoretical properties are established under boundedness and universality of RKHS kernels, regularization conditions, and unconfoundedness/overlap assumptions (Anancharoenkij et al., 4 Feb 2026, Park et al., 2021, Muandet et al., 2018).

5. Practical Considerations and Computational Aspects

CCME enables estimation and inference with structured, high-dimensional, or non-Euclidean outcome spaces, conditional on covariates or projected features:

• The only requirement on the outcome space $\mathcal Y$ is the existence of a positive definite kernel. Kernels for images, sequences, graphs, or sets integrate naturally (Muandet et al., 2018).
• Regularization ($\lambda$) is typically chosen by cross-validation or prescribed decay rates; kernel parameters may be tuned similarly.
• Computational complexity is dominated by matrix inversion or linear system solving in kernel ridge regression, scaling as $O(n^3)$. For U-statistic regression, complexity can grow as $n^r$ for order-$r$ statistics (Park et al., 2021). Nyström or random feature approximations can reduce cost.
• Sample generation from estimated counterfactual distributions is supported via kernel herding algorithms (Muandet et al., 2018).
• Bayesian CCME methods propagate epistemic uncertainty through posterior variances, interpretable as coverage for estimated effects or densities (Martinez-Taboada et al., 2022).

6. Applications and Empirical Results

CCME has demonstrated utility in:

• Detecting distributional effects missed by mean-based metrics (e.g., variance shifts detected via higher-order MMD) (Muandet et al., 2018).
• Visualization of treatment effect heterogeneity using witness functions (Park et al., 2021).
• Consistent estimation of conditional counterfactual distributions for complex outcomes (e.g., images in semi-synthetic MNIST tasks), accurately recovering multimodal structure and canonical representatives even with misspecified models (Anancharoenkij et al., 4 Feb 2026).
• Off-policy evaluation, where CCME outperforms direct, inverse-propensity, doubly robust, and slate estimators, particularly under strong covariate shift (Muandet et al., 2018).
• Bayesian CCME offers calibrated credible intervals for counterfactual estimates, including complex estimands involving unpaired data fusion and OPE scenarios, achieving nominal coverage when both sources of epistemic uncertainty are modeled (Martinez-Taboada et al., 2022).

A comparative table of prominent CCME variants and their features:

Variant	Key Estimator Type	Guarantees
Classical CCME	Kernel ridge regression	Nonparametric consistency,
		minimax rates
Doubly-robust CCME	Meta-estimator (any RKHS-valued)	Double robustness, rates
Bayesian CCME	GP prior/posterior on embeddings	Epistemic uncertainty, calibration
U-Statistic Regression	Higher-order moments	Consistency, for structured effects

7. Extensions and Connections

CCME unifies methodologies in nonparametric causal inference, kernel embedding, treatment effect heterogeneity, and off-policy evaluation. Recent developments include:

• Meta-learners using learned feature maps, deep-kernel architectures, and neural network surrogates (Anancharoenkij et al., 4 Feb 2026).
• Generalizations to alternative distributional metrics (Wasserstein, energy distance), and scalable solvers for high-dimensional or higher-order moment settings (Park et al., 2021).
• Bayesian CCME enables the integration of multiple data sources and propagation of uncertainty in both treatment assignment and outcome mapping (Martinez-Taboada et al., 2022).
• Potential connections to double machine learning, orthogonalized learners, and robust hypothesis testing.

By embedding full conditional counterfactual distributions into RKHS, CCME provides a rigorous, extensible, and computationally tractable framework for causal, distributional, and off-policy inference across arbitrary, possibly structured, outcome spaces.