Kernel Treatment Effects (KTE)

• Kernel Treatment Effects are nonparametric methods that embed full counterfactual outcome distributions in an RKHS to capture effects beyond mean shifts.
• The approach addresses adaptive data collection challenges using doubly robust scores and variance stabilization, ensuring valid asymptotic inference.
• Leveraging a martingale CLT, KTE enables rigorous hypothesis testing for complex distributional changes, providing higher power over mean-based tests.

Kernel Treatment Effects (KTE) are a class of methods for nonparametric causal inference that utilize kernel embeddings of outcome distributions to enable inference on distributional treatment effects, especially when data are collected adaptively or the research focus extends beyond average effects. The KTE framework represents counterfactual distributions in a reproducing kernel Hilbert space (RKHS) and quantifies treatment effects via distances (such as maximum mean discrepancy, MMD) between these mean embeddings. This approach provides rigorous inference with valid asymptotic distributions under adaptive data collection settings, addressing key challenges ignored by classical mean-focused tests.

1. Kernel Mean Embeddings and Distributional Causal Inference

KTE operates by embedding the full counterfactual outcome distributions for each treatment level ($a$) in an RKHS associated with a characteristic, bounded kernel $k(\cdot, \cdot)$. Specifically, the outcome distribution under treatment $a$ is represented by its mean embedding:

$$n(a) = \mathbb{E}_X[M_{Y|A,X}(a, X)] \in \mathcal{H}_y,$$

where $M_{Y|A,X}(a, X)$ is the conditional mean embedding of the outcome given $A=a$ and $X$. The kernel treatment effect between $a$ and $a'$ is then defined as the RKHS distance

$$\mathrm{KTE}(a, a') = \| n(a) - n(a') \|_{\mathcal{H}_y},$$

which is a generalization of the two-sample MMD to the context of counterfactual inference. If the kernel is characteristic, this distance is zero if and only if the counterfactual distributions are equal, thus allowing for hypothesis tests of distributional differences beyond mean shifts.

By working in the RKHS, KTE methods encode information about all moments and structural properties (such as variance, skewness, or even modes) of the outcome distribution, which makes possible the detection of treatment effects manifesting as changes in dispersion or higher moments.

2. Challenges with Adaptive Data Collection

Adaptive experiments—such as multi-armed bandits or adaptive clinical trials—allocate treatments dynamically in response to accruing data, making treatment assignment at time $t$ dependent on past outcomes. This dependence violates the i.i.d. assumption underlying classical asymptotic theory (including the standard central limit theorem), leading to random and history-dependent sample sizes and allocation probabilities.

In such settings, empirical means and classical estimators can converge to non-Gaussian mixtures or exhibit inflated variance, thus invalidating naive inference procedures. This is especially problematic for kernel-based functionals, which rely on summing over non-independent influence scores or weights.

3. Doubly Robust Scores and Variance Stabilization

The KTE methodology adapts to the non-i.i.d. setting by using influence functions that are both doubly robust and variance-stabilized. The canonical gradient for the mean embedding at treatment $a$ is

$$D'(\pi, \mu; a)(X, A, Y) = \frac{1\{A=a\}}{T(a|X)} [\phi_y(Y) - \mu(a,X)] + \mu(a, X),$$

where $T(a|X)$ is the allocation probability (potentially history-dependent), $\mu(a,X)$ is the outcome regression (nuisance model), and $\phi_y(Y)$ is the feature map for $Y$ in the RKHS.

The estimator for the difference of mean embeddings $\Psi(a,a') = n(a) - n(a')$ is constructed as a sum of increments of the form

$$Z_t = w_{t-1} \left[ O_t - \mathbb{E}(O_t \mid \mathcal{F}_{t-1}) \right],$$

where the stabilizing weights $w_{t-1}$ are chosen so that the predictable quadratic variation converges to a deterministic trace-class operator, thereby ensuring the total variance is stabilized across rounds. This variance normalization enables the use of a Hilbert-space martingale CLT as opposed to those based on independence.

By leveraging double robustness, the estimator remains consistent if either the propensity model $T(a|X)$ or the outcome regression $\mu(a,X)$ is correctly specified.

4. Asymptotic Theory: Hilbert-Space Martingale CLT

A central technical development is the demonstration that—under regularity conditions (such as strong positivity, boundedness of the kernel, suitable rates for nuisance estimation, and certain tail conditions)—the stabilized estimator has a limiting Gaussian distribution in the Hilbert space:

$$T^{-1/2} \sum_{t=1}^T Z_t \ \overset{d}{\longrightarrow}\ \mathcal{N}_{\mathcal{H}_y}(0, \mathcal{T}),$$

where $\mathcal{T}$ is the limiting covariance operator. In particular, three technical conditions are imposed:

• (B1) Negligibility (no big jumps),
• (B2) Predictable quadratic variation converges,
• (B3) Suitable tail or tightness condition.

Under these, the normalized sum of stabilized, mean-zero RKHS-valued increments asymptotically approaches a Gaussian element. This allows for the construction of valid confidence regions and hypothesis tests even in adaptive data settings, overcoming clustering issues due to dependence.

5. Sample-Split Stabilized Testing for Distributional Effects

A key practical advance is the sample-split stabilized test for the null hypothesis $H_0: n(a) = n(a')$, i.e., no distributional treatment effect. The test proceeds by partitioning the sample into two sequential folds. For each fold, a stabilized, doubly robust estimator is constructed using nuisance models fit independently on the opposite fold (cross-fitting ensures independence). Setting $T_1$ and $T_2$ as the two fold-wise stabilized estimators, the test statistic is the normalized cross inner product

$$T_{\mathrm{cross}} = \frac{\langle T_1, T_2 \rangle_{\mathcal{H}_y}}{V_{\mathrm{cross}}},$$

where $V_{\mathrm{cross}}$ is an estimator of the cross-fold variance. As $T_1$ and $T_2$ are asymptotically independent, under the null hypothesis, $T_{\mathrm{cross}}$ is standard normal, yielding a test with provably correct Type I error.

This design overcomes degeneracies of the classic MMD statistic under the null (typical in kernel two-sample testing) and neatly translates the martingale CLT from the estimator to the test.

6. Empirical Performance—Calibration, Power, and Robustness

The variance-stabilized DR KTE estimator demonstrates nominal calibration under the null for both mean-shift and higher-moment-shift scenarios. Under mean shifts, it performs similarly to adaptive mean-based approaches (CADR, AW-AIPW). However, only VS-DR-KTE retains high power for distributional shifts that leave the mean unchanged, as mean-based approaches are by construction insensitive to higher-order or structural changes in the outcome distribution.

The sample-split stabilized test achieves controlled false positive rates and is computationally tractable since cross-fitting removes the need for expensive permutations. Compared to scalar outcome testing, the kernel-based test provides substantially higher empirical power for detecting complex treatment effect patterns.

7. Broader Implications and Future Perspectives

KTE extends the reach of causal inference to adaptive experimental designs and to questions where effect heterogeneity manifests at the level of distributional differences. The RKHS embedding ensures that all moments and functional features of the outcome are considered, which is vital in modern applications (e.g., reinforcement learning, personalized medicine) where mean-based decisions are inadequate.

The stabilized, doubly robust framework may be generalized to more complex conditional KTEs, to high-dimensional settings via improved operator theory, or to settings with even weaker identifiability. Its use cases include not only clinical trials and economics, but also online learning and policy evaluation in environments where adaptation and non-i.i.d. sampling are inherent.

In conclusion, the KTE approach—via RKHS representation, doubly robust stabilized estimation, and martingale-based inference—provides a statistically rigorous and practically effective methodology for non-i.i.d. distributional effects analysis under adaptive data collection (Zenati et al., 11 Oct 2025).