# Doubly Robust Kernel Statistics for Testing Distributional Treatment Effects

Jake Fawkes<EMAIL_ADDRESS>
Department of Statistics
University of Oxford Robert Hu<EMAIL_ADDRESS>
Amazon Robin J. Evans<EMAIL_ADDRESS>
Department of Statistics
University of Oxford Dino Sejdinovic<EMAIL_ADDRESS>
School of Mathematical Sciences
University of Adelaide†
Work mainly done while the authors were with the Department of Statistics,
University of Oxford.

###### Abstract

With the widespread application of causal inference, it is increasingly
important to have tools which can test for the presence of causal effects in a
diverse array of circumstances. In this vein we focus on the problem of
testing for _distributional_ causal effects, where the treatment affects not
just the mean, but also higher order moments of the distribution, as well as
multidimensional or structured outcomes. We build upon a previously introduced
framework, Counterfactual Mean Embeddings, for representing causal
distributions within Reproducing Kernel Hilbert Spaces (RKHS) by proposing
new, improved, estimators for the distributional embeddings. These improved
estimators are inspired by doubly robust estimators of the causal mean, using
a similar form within the kernel space. We analyse these estimators, proving
they retain the doubly robust property and have improved convergence rates
compared to the original estimators. This leads to new permutation based tests
for distributional causal effects, using the estimators we propose as tests
statistics. We experimentally and theoretically demonstrate the validity of
our tests.

##  1 Introduction

In this work we focus on the problem of testing for distributional treatment
effects (Bellot & van der Schaar, 2021; Park et al., 2021; Chikahara et al.,
2022), where the aim is to test for causal effect which manifest as something
other than a mean shift. This can be especially useful when the target
variable is high dimensional or structured as a network, since in these cases
there is no natural mean to compare. Our contributions are as follows:

  1. 1.

We introduce new estimators to be used within the Counterfactual Mean
Embeddings framework, based upon the doubly robust estimator of the causal
mean from semi-parametric statistics (Bang & Robins, 2005; Tsiatis, 2006).
These may be applied to estimate kernelised versions of the average treatment
effect and effect of treatment on the treated. We prove these estimators
inherit the double robustness properties of well established semi-parametric
estimators of causal effects and that they converge to the correct value if
either of the models underlying them converge to the true model. This shows
that they have theoretically improved convergence results when compared to the
previous counterfactual mean embedding estimators.

  2. 2.

We apply these new estimators to permutation testing for distributional causal
effects. We propose a new permutation approach which allows for doubly robust
trainable statistics to be used within permutation testing and prove that
these tests are valid.

  3. 3.

We experimentally validate the performance of our test on synthetic, semi-
synthetic and real data. An implementation of our approach can be found at:
https://github.com/Jakefawkes/DR_distributional_test.

###  1.1 Related work

Our work builds upon the counterfactual mean embeddings framework introduced
in Muandet et al. (2021), which we detail further within Section 2.3. This
falls under the general area of applying kernels to test for distributional
causal effects of which Bellot & van der Schaar (2021) is an early example,
whose test statistic arises as a special case of our own for the
distributional effect of treatment on the treated. In a concurrent work,
Martinez-Taboada et al. (2023) develop a similar approach using doubly robust
statistics (Robins & Rotnitzky, 1995) within kernel spaces. However, by
applying the work of Shekhar et al. (2022) they are able to take a permutation
free approach to testing distributional causal effects. Due to the concurrency
and current lack of publicly available code we do not compare against this
method. This work has also been built to estimate counterfactual densities in
Martinez-Taboada & Kennedy (2023). In more general applications of kernel
spaces to testing distributional causal effectsPark et al. (2021) develop a
statistic targeting the conditional average treatment effect (CATE) using
conditional mean embeddings to estimate the RKHS distance between the expected
potential outcomes. Chikahara et al. (2022) use tests for distributional
causal effects to find which features are relevant to difference in treatment
effects. Testing for causal distributional effects falls within long tradition
of using kernel spaces for hypothesis testing as they faithfully capture all
features of a distribution (Gretton et al., 2005; 2012). Kernel methods have
also been widely applied within the larger causality literature. For example,
to instrumental variables (Singh et al., 2019; Muandet et al., 2020), to
causal learning with proxies and unmeasured confounders (Singh et al., 2020;
Mastouri et al., 2021), and to the orientation of causal edges (Mitrovic et
al., 2018).

##  2 Background and Notation

###  2.1 Causal Set Up

$T$$Y$$X$ Figure 1: Assumed DAG

Throughout, we will let $Y$ denote the outcome, which is affected by a binary
treatment $T$ in the presence of additional covariates $X$. We will use the
potential outcome notation so that $Y(t)$ corresponds to the outcome observed
when $T$ is given the value $t$. We assume that the causal relationships
between these variables are given by the basic confounding DAG in Figure 1.

#### The Propensity Score and Overlap

Throughout we will make reference to the propensity score, which is the
probability of treatment given a set of covariates: $P(T=1\mid X=x)$. For a
flexible notation we will use $e(x,t)$ to denote $P(T=t\mid X=x)$, however as
we are using binary treatment we will have that $e(x,t)=1-e(x,t^{\prime})$ if
$t+t^{\prime}=1$. We will also make use of the inverse propensity odds which
we will denote $w(x,t)=\frac{1-e(x,t)}{e(x,t)}$. Finally $\hat{e}(x,t)$ and
$\hat{w}(x,t)$ will be used to denote finite sample estimates of these
quantities.

An important assumption for causal inference from observational data is that
of overlap. This ensures that there are treated and untreated individuals we
can compare and it written formally as:

######  Definition 1.

We say there is overlap if we have:

$\displaystyle 0<e(X,t)<1$

with probability 1. We say the overlap is one-sided if only one-side of the
inequality holds and that any overlap is strict if there is some $\delta>0$
such that:

$\displaystyle\delta\leq e(X,t)\leq 1-\delta,$

again with probability 1. Under the one-sided overlap assumption that
$e(x,t)>0$ and the causal DAG in Figure 1, the distribution of the
interventional quantity $Y(t)$ is uniquely identified from observational data
as follows:

$\displaystyle P(Y(t)=y)=\mathbb{E}_{P(X)}\left[P(Y=y\mid X,T=t)\right].$

This allows us to estimate causal quantities of interest from observational
distributions alone, such as the average treatment effect (ATE):

$\displaystyle{\mathbb{E}}[Y(1)-Y(0)].$

It is important to note that both-sided overlap is required for the
identification of the average treatment effect, as each interventional
distribution requires overlap at a different $T=t$ for identification. In the
case where we are only likely to have one-sided overlap, such as when an
experimental new medical treatment is only given to the most severe cases,
practitioners instead focus on the effect of treatment on the treated (ETT):

$\displaystyle{\mathbb{E}}[Y(1)-Y(0)\mid T=1].$

This quantity is identified from observational data whenever $e(x,0)>0$ for
almost all $x$, so all individuals have some probability of not being treated.
This allows us to still get some estimate of treatment efficacy in cases where
both-sided overlap is violated. It is also worth noting that overlap is a
testable assumption, with a test developed for example in Lei et al. (2021).

###  2.2 Kernel Background

We use kernel embeddings of distributions to construct our test statistics and
so we now briefly introduce the relevant background material and notation. For
a more thorough engagement with this material we refer the reader to Muandet
et al. (2017).

#### Reproducing Kernel Hilbert Spaces

Let $\mathcal{X}$ be some non-empty space. A real-valued RKHS,
$(\mathcal{H}_{\mathcal{X}},\langle\cdot,\cdot\rangle_{\mathcal{H}_{\mathcal{X}}})$,
is a complete inner product space of functions $f:\mathcal{X}\to\mathbb{R}$
such that the evaluation function is continuous for all $x\in\mathcal{X}$. Due
to the Riesz representer theorem we have that for all $x\in\mathcal{X}$ there
is a function $k_{x}\in\mathcal{H}_{\mathcal{X}}$ which satisfies the
reproducing property, that $f(x)=\left\langle f,k_{x}\right\rangle$ for all
$f\in\mathcal{H}_{k}$. The function $k(x,x^{\prime})=\left\langle
k_{x},k_{x^{\prime}}\right\rangle$ is known as the reproducing kernel for the
space $\mathcal{H}_{\mathcal{X}}$. Conversely, via the Moore-Aronszajn Theorem
(Aronszajn, 1950), any symmetric positive definite function $K$ on
$\mathcal{X}$ defines a unique RKHS. We will suppose access to an RKHS on
$\mathcal{X}$ and $\mathcal{Y}$ denoted by $\mathcal{H}_{\mathcal{X}}$ and
$\mathcal{H}_{\mathcal{Y}}$, with with kernels $k$ and $\ell$ respectively.

Further, for a random variable $X$ on $\mathcal{X}$ with distribution $P_{X}$
we can define the mean embedding of $X$ as:

$\displaystyle\mu_{X}={\mathbb{E}}_{P_{X}}[k(X,\cdot)].$

Under the intergrability condition that
$\int_{\mathcal{X}}\sqrt{k(x,x)}dP_{X}(x)<\infty$ we have that
$\mu_{X}\in\mathcal{H}_{\mathcal{X}}$. If we have two random variables,
$X,X^{\prime}$, over $\mathcal{X}$ with distributions $P_{X},P_{X^{\prime}}$
respectively, with a slight abuse of notation we denote the maximum mean
discrepancy (MMD) by:

$\displaystyle\mathrm{MMD}[X,X^{\prime},\mathcal{H}_{\mathcal{X}}]\vcentcolon=\left\lVert\mu_{P_{X}}-\mu_{P_{X^{\prime}}}\right\rVert_{\mathcal{H}_{\mathcal{X}}}.$

The kernel $k$ is known as characteristic if
$\mathrm{MMD}[X,X^{\prime},\mathcal{H}_{\mathcal{X}}]=0$ if and only if
$P_{X}\overset{d}{=}P_{X^{\prime}}$. Many popular kernels such as the Gaussian
and Matérn are charecteristic kernels. The $\mathrm{MMD}$ between two
distributions can be estimated from finite samples by computing the
$\mathrm{MMD}$ between the empirical distributions. This was shown in Gretton
et al. (2012) to converge to the true MMD at the paramteric convergence rate,
$\mathcal{O}_{P}(n^{-\frac{1}{2}})$. For this reason the MMD is often used for
efficiently testing equality of distributions.

#### Conditional Mean Embedding

RKHSs also allow us to represent conditional distributions through the
Conditional Mean Embedding (CME) (Song et al., 2009; Muandet et al., 2017),
given by:

$\displaystyle\mu_{Y\mid X=x}={\mathbb{E}}[\ell(Y,\cdot)\mid X=x].$

In order to estimate this Grünewälder et al. (2012) propose to take a
regression point of view, seeing the CME as the solution to the following
regression problem:

$\left\\{\begin{aligned}
\hfil\displaystyle\begin{split}&C^{*}={\arg\min}_{C\in
B_{2}(\mathcal{H}_{\mathcal{X}},\mathcal{H}_{\mathcal{Y}})}{\mathbb{E}}_{P_{(Y,X)}}\left\lVert\ell(Y,\cdot)-Ck(X,\cdot)\right\rVert_{\mathcal{H}_{\mathcal{Y}}}\\\
&\mu_{Y\mid X=x}=C^{*}K(x,\cdot)\end{split}\end{aligned}\right.,$

where $B_{2}(\mathcal{H}_{\mathcal{X}},\mathcal{H}_{\mathcal{Y}})$ is the
space of Hilbert-Schmidt operators
$\mathcal{H}_{\mathcal{X}}\to\mathcal{H}_{\mathcal{Y}}$. This interpretation
leads to an estimate of the CME from a finite dataset,
$\mathcal{D}=\\{x_{i},y_{i}\\}^{n}_{i=1}$ as:

$\left\\{\begin{aligned}
\hfil\displaystyle\begin{split}\hat{C}^{*}&=\underset{C\in\mathsf{B}_{2}(\mathcal{H}_{\mathcal{X}},\mathcal{H}_{\mathcal{Y}})}{\arg\min}\,\frac{1}{n}\sum_{i=1}^{n}\|\ell(y_{i},\cdot)-Ck(x_{i},\cdot)\|_{\mathcal{H}_{\mathcal{Y}}}^{2}+\gamma\|C\|^{2}_{\mathsf{B}_{2}}\\\
\quad&=\boldsymbol{\ell}^{\top}\mathbf{W}\mathbf{k}\\\
\hat{\mu}_{Y|X=x}&=\boldsymbol{\ell}^{\top}\mathbf{W}\mathbf{k}(x)\end{split}\end{aligned}\right..$
(1)

where
$\boldsymbol{\ell}=\left(\ell(y_{i},\cdot)\right)_{i=1}^{n},\mathbf{k}=\left(k(x_{i},\cdot)\right)_{i=1}^{n}$,
and $\mathbf{W}=(\mathbf{K}+\lambda\mathbf{I}_{n})^{-1}$ for
$\mathbf{K}=\left(k(x_{i},x_{j})\right)_{i,j=1}^{n}$. We will often make use
of the CME conditional on a particular value of $T$, so $\mu_{Y\mid X=x,T=t}$.
In that case we let $\boldsymbol{\ell}_{t}$, $\mathbf{W}_{t}$ and
$\mathbf{k}_{t}$ denote the respective matrices used in the estimation of
$\hat{\mu}_{Y\mid X=x,T=t}$.

###  2.3 Counterfactual Mean Embeddings

Our work builds upon the counterfactual mean embeddings framework of Muandet
et al. (2021). They demonstrate that under the causal structure in Section 2.1
and with the one-sided overlap $e(x,t)>0$, the mean embedding of the potential
outcome, $Y(t)$, can be written as:

$\displaystyle\mu_{Y(t)}=\mathbb{E}_{P_{(Y,X,T)}}\left[\frac{\ell(Y,\cdot)\,\mathbbm{1}\\!\left\\{T=t\right\\}}{e(X,t)}\right],$

which may be seen as the kernel equivalent of the inverse probability
weighting estimator for the causal mean. Given a finite sample,
$\\{(t_{i},x_{i},y_{i})\\}_{i=1}^{n}$, and access to the true propensity
score, this quantity can be estimated as:

$\displaystyle\hat{\mu}_{Y(t)}=\frac{1}{n}\sum_{i=1}^{n}\frac{\ell({y}_{i},\cdot)\,\mathbbm{1}\\!\left\\{{t}_{i}=t\right\\}}{e({x_{i}},t)}.$

Muandet et al. (2021) prove that this converges to the true embedding at rate
$\mathcal{O}_{P}(n^{-\frac{1}{2}})$. If we have both-sided overlap both
$\hat{\mu}_{Y(1)}$ and $\hat{\mu}_{Y(0)}$ are therefore identified from data
and we can use ${\mathrm{MMD}}[Y(1),Y(0),\mathcal{H}_{\mathcal{Y}}]$ to
measure any distributional effects of treatment. Furthermore, if the kernel is
characteristic we have that
${\mathrm{MMD}}[Y(1),Y(0),\mathcal{H}_{\mathcal{Y}}]=0$ if and only if
$Y(1)\overset{d}{=}Y(0)$ . Finally, this quantity can be estimated from finite
samples at rate $\mathcal{O}_{P}(n^{-\frac{1}{2}})$ by:

$\displaystyle\widehat{\mathrm{MMD}}[Y(1),Y(0),\mathcal{H}_{\mathcal{Y}}]=\left\lVert\hat{\mu}_{Y(1)}-\hat{\mu}_{Y(0)}\right\rVert_{\mathcal{H}_{\mathcal{Y}}}.$

This leads to a permutation test for distributional causal effects following
the conditional permutation scheme for testing for causal effects (Rosenbaum,
1984), and using the estimated MMD between the potential outcomes as a test
statistic. Throughout we will refer to
$\widehat{\mathrm{MMD}}[Y(1),Y(0),\mathcal{H}_{\mathcal{Y}}]$ as the
_distributional average treatment effect_ (DATE), due to the similarity
between this and the average treatment effect statistic in Section 2.1.

Analogously to the average treatment effect, the distributional average
treatment effect is only identified from data under both-sided overlap.
Therefore we now introduce a second target based of the effect of treatment on
the treated, the _distributional effect of treatment on the treated_ (DETT):

$\displaystyle{\mathrm{MMD}}\Big{[}Y(1)_{\left\\{T=1\right\\}},Y(0)_{\left\\{T=1\right\\}},\mathcal{H}_{\mathcal{Y}}\Big{]}\coloneqq\left\lVert\mu_{\left\\{Y(1)\mid
T=1\right\\}}-\mu_{\left\\{Y(0)\mid
T=1\right\\}}\right\rVert_{\mathcal{H_{Y}}}.$

Estimation this quantity requires estimation of the embeddings
$\mu_{\left\\{Y(0)\mid T=1\right\\}},\mu_{\left\\{Y(1)\mid T=1\right\\}}$.
Since $P(Y(1)\mid T=1)=P(Y\mid T=1)$ we can estimate the latter embedding
using observed samples of $Y$ for individuals with $T=1$. This means the only
challenge is the estimation of $\mu_{\left\\{Y(0)\mid T=1\right\\}}$ which
Muandet esimtate using the fact that:

$\displaystyle\mu_{Y(0)\mid T=1}=\mathbb{E}_{P(X\mid T=1)}\left[\mu_{Y\mid
X,T=0}\right].$

This means a finite sample estimate of the conditional mean embedding from
Section 2.2 can be used to compute a finite sample estimate of the mean
embedding.

Alternatively, Bellot & van der Schaar (2021) also target the distributional
effect of treatment on the treated using the fact that:

$\displaystyle\mu_{Y(0)\mid
T=1}=\mathbb{E}\left[\ell(Y,\cdot)\mathbbm{1}\\{T=0\\}w(X,0)\right].$

Both of these estimators lead to test statistics based on the distributional
effect of treatment on the treated.

##  3 Doubly Robust Counterfactual Mean Embeddings

In this section we build on previous methodology by considering estimators for
the causal mean embeddings that are built on doubly robust (DR) estimators of
causal effects (Robins & Rotnitzky, 1995). We propose new estimators for the
distributional average treatment effect and the distributional effect of
treatment on the treated. We prove that both of these estimators have the
doubly robust property and so they will converge to the true causal mean
embedding if either of the two models underlying them converges to the true
model. We call this set of approaches Doubly Robust Counterfactual Mean
Embeddings.

###  3.1 Doubly robust estimator of the distributional average treatment
effect (DATE)

Firstly, based on the doubly robust estimator of the average treatment effect
we note that the embedding of the potential outcome, $Y(t)$, can also be
written as:

$\displaystyle\mu_{Y(t)}$
$\displaystyle={\mathbb{E}}[\ell(Y(t),\cdot)-\mu_{Y\mid
X,T=t}]+{\mathbb{E}}[\mu_{Y\mid X,T=t}]$
$\displaystyle={\mathbb{E}}\left[\frac{\mathbbm{1}\\!\left\\{T=t\right\\}\left(\ell(Y,\cdot)-\mu_{Y\mid
X,T=t}\right)}{e(X,t)}\right]+{\mathbb{E}}[\mu_{Y\mid X,T=t}]$
$\displaystyle=\mathbb{E}\left[\frac{\mathbbm{1}\\!\left\\{T=t\right\\}\left(\ell(Y,\cdot)-\mu_{Y\mid
X,T=t}\right)}{e(X,t)}+\mu_{Y\mid X,T=t}\right].$

This is a kernelised version of the doubly robust estimator of the treatment
mean (Robins & Rotnitzky, 1995); it can be estimated from finite samples by
fitting models ${\hat{e}(x,t)},\hat{\mu}_{Y\mid X=x,T=t}$ on a training set
and then averaging them over the test set as:

$\displaystyle\hat{\mu}^{\mathrm{DR}}_{Y(t)}=\frac{1}{n}\sum_{i=1}^{n}\left\\{\frac{\mathbbm{1}\\!\left\\{t_{i}=t\right\\}\left(\ell(y_{i},\cdot)-\hat{\mu}_{Y\mid
X=x_{i},T=t}\right)}{\hat{e}(x_{i},t)}+\hat{\mu}_{Y\mid
X=x_{i},T=t}\right\\}.$ (2)

Take a constant train/test split and let $\hat{e}_{n}(X,t)$ and
$\hat{\mu}^{(n)}_{Y\mid X,T=t}$ be the models that arise from a total sample
of size $n$. Now under the assumption that the propensity and estimated
propensity are uniformly bounded, the following theorem shows that
$\hat{\mu}^{\mathrm{DR}}_{Y(t)}$ has the _doubly robust_ property that it
converges to the true embedding as long as either of the propensity model or
the conditional mean embedding converges. Furthermore, it is also doubly rate
robust, in the sense that it achieves a convergence rate that is the product
of the convergence rates of these working models.

######  Theorem 1.

Assuming that $e(x,t)$ is strongly bounded away from zero, as well as
additional overlap assumptions in Appendix A.1, we have that if the estimators
$\hat{e}_{n}(X,t)$ and $\hat{\mu}^{(n)}_{Y\mid X,T=t}$ satisfy the following:

$\displaystyle\left\lVert\hat{e}_{n}(X,t)-e(X,t)\right\rVert_{2}=\mathcal{O}_{P}(\gamma_{e,n}),\hskip
20.00003pt\left\lVert\left\lVert\hat{\mu}^{(n)}_{Y\mid X,T=t}-\mu_{Y\mid
X,T=t}\right\rVert_{\mathcal{H}_{Y}}\right\rVert_{2}=\mathcal{O}_{P}(\gamma_{r,n})$

with $\gamma_{e,n}=o(1)$ and $\gamma_{r,n}=o(1)$. Then:

$\left\lVert\mu_{Y(t)}-\hat{\mu}^{\mathrm{DR}}_{Y(t)}\right\rVert_{\mathcal{H}_{Y}}=\mathcal{O}_{P}(\mathrm{max}\\{n^{-\frac{1}{2}},\gamma_{r,n}\gamma_{e,n}\\}).$

Therefore, if $\gamma_{r,n}\gamma_{e,n}=O(n^{-\frac{1}{2}})$ we obtain the
parametric convergence rate.

These statistics can be used to form a doubly robust estimator of the
distributional average treatment effect as:

$\displaystyle\widehat{\mathrm{MMD}}_{\mathrm{DR}}[Y(1),Y(0),\mathcal{H}_{\mathcal{Y}}]=\left\lVert\hat{\mu}_{Y(1)}^{\mathrm{DR}}-\hat{\mu}^{\mathrm{DR}}_{Y(0)}\right\rVert_{\mathcal{H}_{\mathcal{Y}}}.$

We derive the closed form of this statistic squared in Appendix B.1. This will
converge to the correct MMD if both $\hat{\mu}_{Y(1)}^{\mathrm{DR}}$ and
$\hat{\mu}_{Y(0)}^{\mathrm{DR}}$ converge to the true embedding at a rate that
is a maximum of both their rates.

###  3.2 Doubly robust estimator of the Distributional Effect of Treatment on
the Treated (DETT)

We now turn to estimating the distributional effect of treatment on the
treated111We leave the $t$ arbitrary, so if $t=1$ this would form an effect of
treatment on the control. For simplicity we do not distinguish the two cases.,
again we form a kernel version of the standard doubly robust estimator of the
distributional effect of treatment on the treated (Moodie et al., 2018). This
leads from the fact that the mean embedding of the counterfactual distribution
$P(Y(t)\mid T=t^{\prime})$ can be written as:

$\displaystyle{\mu}_{Y(t)\mid T=t^{\prime}}$
$\displaystyle={\mathbb{E}}[\ell(Y(t),\cdot)]$
$\displaystyle={\mathbb{E}}[\ell(Y(t),\cdot)-\mu_{Y\mid X,T=t}\mid
T=t^{\prime}]+{\mathbb{E}}[\mu_{Y\mid X,T=t}\mid T=t^{\prime}]$
$\displaystyle={\mathbb{E}}\left[\frac{e(X,t^{\prime})P(T=t)}{e(X,t)P(T=t^{\prime})}\left(\ell(Y,\cdot)-\mu_{Y\mid
X,T=t}\right)\mid T=t\right]+{\mathbb{E}}[\mu_{Y\mid X,T=t}\mid T=t^{\prime}]$
$\displaystyle={\mathbb{E}}\left[\frac{P(T=t)}{P(T=t^{\prime})}w(X,t)\left(\ell(Y,\cdot)-\mu_{Y\mid
X,T=t}\right)\mid T=t\right]+{\mathbb{E}}[\mu_{Y\mid X,T=t}\mid
T=t^{\prime}],$

where $w(x,t)=\frac{1-e(x,t)}{e(x,t)}$.

By fitting models $\hat{w}(x,t)$ and $\hat{\mu}_{Y\mid X=x,T=t}$ on training
samples we get a finite sample estimate of the distributional effect of
treatment on the treated as:

$\hat{\mu}^{\mathrm{DR}}_{Y(t)\mid
T=t^{\prime}}=\frac{1}{n_{t^{\prime}}}\sum_{i=1}^{n}\bigg{(}{\mathbbm{1}\\!\left\\{t_{i}=t\right\\}\hat{w}(x_{i},t)\left(\ell(y_{i},\cdot)-\hat{\mu}_{Y\mid
X=x_{i},T=t}\right)}+\mathbbm{1}\\!\left\\{t_{i}=t^{\prime}\right\\}\hat{\mu}_{Y\mid
X=x_{i},T=t}\bigg{)}$

where $n_{t}=\sum_{i}\mathbbm{1}\\!\left\\{t_{i}=t\right\\}$. Again letting
$\hat{e}_{n}(x,t)$ and222Any model for $w(x,t)$ leads to a model for $e(x,t)$
directly. $\hat{\mu}^{(n)}_{Y\mid X,T=t}$ be the respective models trained on
a set of size $n$ we have that:

######  Theorem 2.

Under overlap assumptions given in the appendix, we have that if the
estimators $\hat{e}_{n}(X,t)$ and $\hat{\mu}^{(n)}_{Y\mid X,T=t}$ satisfy the
following:

$\displaystyle\left\lVert\hat{e}_{n}(X,t)-e(X,t)\right\rVert_{2}=\mathcal{O}_{P}(\gamma_{e,n}),\hskip
20.00003pt\left\lVert\left\lVert\hat{\mu}^{(n)}_{Y\mid X,T=t}-\mu_{Y\mid
X,T=t}\right\rVert_{\mathcal{H}_{Y}}\right\rVert_{2}=\mathcal{O}_{P}(\gamma_{r,n})$

where $\gamma_{e,n}=o(1)$ and $\gamma_{r,n}=o(1)$ then:

$\left\lVert{\mu}_{Y(t)\mid T=t^{\prime}}-\hat{\mu}^{\mathrm{DR}}_{Y(t)\mid
T=t^{\prime}}\right\rVert_{\mathcal{H}_{Y}}=\mathcal{O}_{P}(\mathrm{max}\\{n^{-\frac{1}{2}},\gamma_{r,n}\gamma_{e,n}\\}).$

Applying this we can now estimate the relevant MMD for between treated and
untreated for this population as:

$\displaystyle\widehat{\mathrm{MMD}}_{\mathrm{DR}}\Big{[}Y(1)_{\left\\{T=1\right\\}},Y(0)_{\left\\{T=1\right\\}},\mathcal{H}_{\mathcal{Y}}\Big{]}=\left\lVert\hat{\mu}^{\mathrm{DR}}_{Y(t)\mid
T=t^{\prime}}-\frac{1}{n_{t^{\prime}}}\sum_{i=1}^{n}\mathbbm{1}\\!\left\\{t_{i}=t^{\prime}\right\\}\ell(y_{i},\cdot)\right\rVert_{\mathcal{H}_{\mathcal{Y}}}$

Which has the same convergence rate as $\hat{\mu}^{\mathrm{DR}}_{Y(t)\mid
T=t^{\prime}}$. Therefore, we only need strict one-sided overlap for
convergence, as $\hat{\mu}^{\mathrm{DR}}_{Y(t)\mid T=t^{\prime}}$ just
requires $e(x,t)>\epsilon$. Again we derive the squared statistic and give it
in closed form in Appendix B.2.

##  4 Permutation testing for Distributional Treatment Effects

We now apply both test statistics to the problem of testing for distributional
treatment effects, taking a permutation based approach. As is standard in this
context we test against Fisher’s sharp null (Fisher, 1936; Rosenbaum,
2002)333We note that while the test must be formulated against Fisher’s sharp
null, it will generally have power against alternatives which show
distributional causal effects, so $Y(1)\overset{d}{\neq}Y(0)$.:

$\displaystyle H_{0}:Y_{i}(1)=Y_{i}(0).$

In order to apply a permutation algorithm, we must choose permutations such
that the treatment vector is exchangeable under them. That is, if we let $n$
be the size of our dataset
$\mathcal{D}=\left(\mathbf{X},\mathbf{Y},\mathbf{T}\right)$ and $S_{n}$ be the
permutation group of size $n$ we want to restrict to permutations, $\sigma\in
S_{n}$, such that:

$\displaystyle\left(\mathbf{X},\mathbf{Y},\mathbf{T}\right)\overset{d}{=}\left(\mathbf{X},\mathbf{Y},\mathbf{T}_{\sigma}\right),$

where $\mathbf{T}_{\sigma}=\left(T_{\sigma(i)}\right)^{n}_{i=1}$ is the
permuted treatment vector.

A standard way to ensure that the permutations have this property is to
through matching (Stuart, 2010), where we split the data into matched sets,
such that within each set we have one treated individual. These matched sets
are formed using some measure of ‘similarity’ of individuals, such as the
Mahalanobis distance on covariates, or through the propensity score. Rosenbaum
(2002) argues that if this matching is exact, so that there are no differences
in propensity between matched individuals, then the treatment is exchangeable
under permutations that preserve the matched sets. That is if we let $M\subset
S_{n}$ be the subset of permutations that only permute within matched sets and
sample some $\sigma\in M$, then $\mathbf{T}$ is exchangeable under $\sigma$.
Therefore if we fix a statistic $S$, in our case either the DATE or DETT
statistic, and randomly sample $\sigma_{1},\ldots,\sigma_{m}$ from $M$, we
have that:

$\displaystyle
S\left(\mathbf{X},\mathbf{Y},\mathbf{T}\right),S\left(\mathbf{X},\mathbf{Y},\mathbf{T}_{\sigma_{1}}\right),\ldots,S\left(\mathbf{X},\mathbf{Y},\mathbf{T}_{\sigma_{m}}\right)\text{
are exchangeable.}$

This allows means that we can define the p-value of the test as:

$\displaystyle
p=\frac{1+\sum^{m}_{i=1}\mathbbm{1}\\!\left\\{S\left(\mathbf{X},\mathbf{Y},\mathbf{T}_{\sigma_{i}}\right)\geq
S\left(\mathbf{X},\mathbf{Y},\mathbf{T}\right)\right\\}}{m+1},$

where this is now a valid p-value in the sense that $P(p\leq\alpha)\leq\alpha$
under the null. However, permuting in this way creates computational
challenges when $S$ is one of the proposed statistics. Namely, the calculation
of these statistics requires the fitting of a propensity model and conditional
mean embedding, and this would need to be repeated for every permutation.

To resolve this we first form the matched sets, and then randomly split into
train/test sets, $\mathcal{D}_{{\rm Tr}},\mathcal{D}_{{\rm Te}}$, such that
each matching set is fully contained in one of the datasets. Now let
$M_{\lvert\mathcal{D}_{{\rm Tr}}\rvert}$ be the set of permutations on
$\mathcal{D}$ which leave $\mathcal{D}_{{\rm Te}}$ fixed and preserve the
matching on $\mathcal{D}_{{\rm Tr}}$, and vice versa for $\mathcal{D}_{{\rm
Te}}$. We then sample $N$ random permutations
$\\{\sigma_{1},\ldots,\sigma_{N}\\}$ from $M_{\lvert\mathcal{D}_{{\rm
Tr}}\rvert}$, and form the set $\tilde{M}$ of permutations as:

$\displaystyle\tilde{M}_{N}=\left\\{\sigma\circ\pi:\sigma\in\\{Id\\}\cup\\{\sigma_{1},\ldots,\sigma_{N}\\},\pi\in
M_{\mathcal{D}_{{\rm Te}}}\right\\},$

Then the following demonstrates by applying results from Ramdas et al. (2023)
that we can form a valid permutation test by randomly sampling permutations
from $\tilde{M}$:

######  Proposition 1.

Under exact matching, if we sample $\tau_{0},\ldots,\tau_{m}$ from
$\tilde{M}_{N}$ we have that the following is a valid p-value:

$\displaystyle
p=\frac{1+\sum^{m}_{i=1}\mathbbm{1}\\!\left\\{S\left(\mathbf{X},\mathbf{Y},\mathbf{T}_{\tau_{i}}\right)\geq
S\left(\mathbf{X},\mathbf{Y},\mathbf{T}\right)\right\\}}{m+1}$

for any statistic $S$ and number of sampled permutations $N$.

This alleviates the computational challenges associated with using the DATE or
DETT statistics as there is now a controllable parameter $N$ which determines
how many models must be trained. Increasing $N$ will mean training more models
and so will decrease the variance in $p$, however it will incur greater
computational cost.

To our knowledge, this is the first time this permutation strategy has
appeared in the literature. Moreover, this procedure could be applied to
efficiently use other trainable test statistics within permutation testing,
such as standard doubly robust estimators of the causal mean or more general
permutation tests for assessing the performance of predictive models as in
Ojala & Garriga (2010).

The result in Proposition 1 relies on exact matching, which is a common
assumption in proving the validity of conditional permutation tests Rosenbaum
(2002). However recent results in the area have shown that p-values arising
from matched permutations are approximately valid under inexact matching
(Berrett et al., 2020; Pimentel, 2022). We expect and empirically observe that
similar results hold in our case.

##  5 Experiments

###  5.1 Fit tests for Doubly Robust Counterfactual Mean Embeddings

To demonstrate the convergence properties implied by our theoretical results
in Section 3 we fit our statistics on simulated data from the following data
generating processes:

$\displaystyle X$ $\displaystyle\sim\mathcal{N}(0,I_{9})$ $\displaystyle p$
$\displaystyle=\left(\frac{1}{1+\exp{(\boldsymbol{a}^{\top}X)}}\right)^{2}-{\mathbb{E}}_{X}\left[\left(\frac{1}{1+\exp{(\boldsymbol{a}^{\top}X)}}\right)^{2}\right]$
$\displaystyle T$ $\displaystyle\sim\text{Ber}(p)$ $\displaystyle Y$
$\displaystyle=\boldsymbol{b}^{\top}X+\beta T+\epsilon,$

where $\epsilon\sim\mathcal{N}(0,\sigma^{2})$ and the values of
$\boldsymbol{a},\boldsymbol{b}\in\mathbb{R}^{9}$ and
$\beta,\sigma\in\mathbb{R}$ are given in Appendix C. This is similar to the
setting in Muandet et al. (2021). We also simulate from:

$\displaystyle T$ $\displaystyle\sim\text{Ber}$ $\displaystyle X$
$\displaystyle\sim\mathcal{N}(0,(1+\alpha T)I_{10})$ $\displaystyle Y$
$\displaystyle=f_{T}(X)+\epsilon^{\prime}$ $\displaystyle\epsilon^{\prime}$
$\displaystyle\sim\mathcal{N}(0,\left(\sigma^{\prime}\right)^{2})$

where $f_{0},f_{1}:\mathbb{R}^{10}\rightarrow\mathbb{R}$ and
$\alpha,\sigma^{\prime}$ are given in the Appendix. The setting here is the
same as in Bellot & van der Schaar (2021).

Figure 2: These plots demonstrate the convergence of the doubly robust
estimators when the propensity model is misspecified. The left plots show the
embeddings required for DATE, whilst the right plots show those needed for
DETT. We vary between two data generating process given in Section 5.1, aiming
to show both sided and one sided overlap of the propensity score.

For both settings we fit a linear logistic regression for the propensity score
so that the model is incorrectly specified. In first data generating process
we have both-sided overlap in the true propensity and so plots (a) and (b) in
Figure 2 demonstrate that the doubly robust embeddings converge for both
interventional values. Due to an incorrectly specified propensity model the
counterfactual mean embeddings of Muandet et al. (2021) do not converge.

For the second data generating process, increasing the value of $b$ creates a
setting where the overlap is more one-sided. As such plots (c) and (d) in
Figure 2 demonstrate that only $\hat{\mu}^{\mathrm{DR}}_{Y(1)}$ and
$\hat{\mu}^{\mathrm{DR}}_{Y(1)\mid T=0}$ converge quickly to the true
embeddings.

Plots of the true propensity for both simulations can be found in Appendix
C.1.

###  5.2 Testing For Distributional Effects on Simulated Data

We now apply these statistics to the testing of distributional causal effects
where the data is simulated from:

$\displaystyle
X\sim\mathcal{N}(0,I_{9}),\;\;\;p=\left(\frac{1}{1+\exp{(\boldsymbol{a}^{\top}X)}}\right),$
(3) $\displaystyle
T\sim\text{Ber}(p),\;\;\;Y=\boldsymbol{b}^{\top}X+\beta(2Z-1)T+\epsilon,\;\;\;\epsilon\sim\mathcal{N}(0,\sigma^{2})$
(4)

Where $\alpha,\beta\in\mathbb{R}^{9}$ and $\sigma\in\mathbb{R}$ are as in
Section 5.1 and the variable $Z$ is $1$, a $\mathrm{Ber}(\frac{1}{2})$ or
$\mathrm{Unif}([0,1])$. The second two settings capture when there is a
distributional causal effect which does not shift the causal means.

In Figure 3 we plot the rejection rate of the tests based on four statistics,
the DATE and DETT statistics estimated without conditional mean embeddings,
and the DATE and DETT statistics estimated with conditional mean embeddings
denoted by DR-DATE and DR-DETT respectively. The matching for all statistics
is done via logistic regression and we apply the permutation from Section 4.
We compare our methods against Double Machine Learning (Chernozhukov et al.,
2018) and Targeted Maximum Likelihood Estimation (Van Der Laan & Rubin, 2006)
as baselines to effectively pick up a shift in causal means. We run these
experiments with $2000$ data points, rejecting at the $0.05$ significance
level. Power plots showing the rejection rates over $100$ re-runs of this
experiment are displayed in Figure 3.

In this experiment we observe similar performance for all distributional test
statistics. We believe this is due to the fact that both the true and
estimated propensities are linear functions of the covariates, and so all
statistics will correctly fit the causal mean embeddings. Plot (a)
demonstrates that when the shift is in the mean, as expected our methods will
perform worse than the baselines. This is due to the cost of targeting
distributional effects over simply mean shifts. However, plots (b) and (c)
demonstrate that when the causal effect is only distributional, our statistics
can correctly reject the null, unlike these traditional methods.

Figure 3: Simulated tests for distributional causal effects where the $\beta$
parameter controls the size of the effect. In plot $(a)$, $\beta$ introduces a
shift in the mean only, whereas in $(b)$ and $(c)$ the shift only affects
higher moments of the distribution. All simulations with $n=2000$ samples.

###  5.3 Semi-Synthetic and Real World Data

Finally, to evaluate the performance of our tests within more realistic
settings we evaluate on a selection of semi-synthetic data. We evaluate on two
standard semi-synthetic tasks, the infant health and development program
(IDHP) introduced in Hill (2011), the linked births and deaths data (LBIDD)
(Shimoni et al., 2018).

As both of these datasets are semi-synthetic, we have access to the
counterfactuals and so can simulate realistic data under the null hypothesis.
In order to prevent problems with extreme propensity scores we remove any data
for which $e(x,t)<0.03$ for any value of $t$. We again use logistic regression
matching and weights model.

The results can be found in Table 1, which shows all kernel based test
statistics produce valid p-values under the null hypothesis. Further the DATE
based test statistics demonstrate strong power against the alternative for all
three datasets. The DETT test statistic on the other hand only shows a good
level of power on the IDHP dataset, with more samples required to reject on
the LBIDD dataset.

Table 1: Rejection rates at the $0.05$ level on semi-synthetic datasets. Dataset | Hypothesis | DATE | DR-DATE | DETT | DR-DETT | DML
---|---|---|---|---|---|---
IDHP ($n=747$) | $H_{0}$ | 0.02 ±0.03 | 0.05 ±0.04 | 0.01 ±0.02 | 0.03 ±0.03 | 0.04 ±0.04
$H_{1}$ | 1.00 ±0 | 1.00 ±0 | 1.00 ±0 | 0.99 ±0.03 | 1.00 ±0
LBIDD ($n=1000$) | $H_{0}$ | 0.02 ±0.03 | 0.03 ±0.03 | 0.00 ±0 | 0.00 ±0 | 0.17 ±0.10
$H_{1}$ | 0.85 ±0.07 | 0.82 ±0.07 | 0.13 ±0.06 | 0.08 ±0.06 | 1.00 ±0
LBIDD ($n=2500$) | $H_{0}$ | 0.03 ±0.03 | 0.03 ±0.03 | 0.00 ±0 | 0.00 ±0 | 0.00 ±0
$H_{1}$ | 0.82 ±0.08 | 0.75 ±0.09 | 0.19 ±0.08 | 0.13 ±0.06 | 1.00 ±0

##  6 Conclusion

We have proposed new doubly robust estimators for the kernel mean embedding of
causal distributions. Our theoretical experimental results show that these
converge in a wider array of circumstances than the original estimators of
these quantities. Further, we applied these embeddings to the problem of
testing for distributional causal effects. We used a permutation based
approach for testing, and proposed a permutation that allows us to use our
doubly robust statistics within permutation testing. We prove the validity of
this permutation under exact matching and experimentally validate this
approach under inexact matching. The results show that our test statistics are
able to pick up causal effects that only manifest as distributional shifts,
where traditional mean shift methods fail.

## References

  * Aronszajn (1950) Nachman Aronszajn.  Theory of reproducing kernels.  _Transactions of the American mathematical society_ , 68(3):337–404, 1950.
  * Bang & Robins (2005) Heejung Bang and James M. Robins.  Doubly robust estimation in missing data and causal inference models.  _Biometrics_ , 61(4):962–973, 2005.
  * Bellot & van der Schaar (2021) Alexis Bellot and Mihaela van der Schaar.  A kernel two-sample test with selection bias.  In _Uncertainty in Artificial Intelligence_ , pp. 205–214. PMLR, 2021.
  * Berrett et al. (2020) Thomas B Berrett, Yi Wang, Rina Foygel Barber, and Richard J Samworth.  The conditional permutation test for independence while controlling for confounders.  _Journal of the Royal Statistical Society Series B: Statistical Methodology_ , 82(1):175–197, 2020.
  * Chernozhukov et al. (2018) Victor Chernozhukov, Denis Chetverikov, Mert Demirer, Esther Duflo, Christian Hansen, Whitney Newey, and James Robins.  Double/debiased machine learning for treatment and structural parameters.  _The Econometrics Journal_ , 21(1):C1–C68, 01 2018.  ISSN 1368-4221.  doi: 10.1111/ectj.12097.  URL https://doi.org/10.1111/ectj.12097.
  * Chikahara et al. (2022) Yoichi Chikahara, Makoto Yamada, and Hisashi Kashima.  Feature selection for discovering distributional treatment effect modifiers.  _arXiv preprint arXiv:2206.00516_ , 2022.
  * Fisher (1936) Ronald Aylmer Fisher.  Design of experiments.  _British Medical Journal_ , 1(3923):554, 1936\.
  * Gretton et al. (2005) Arthur Gretton, Olivier Bousquet, Alex Smola, and Bernhard Schölkopf.  Measuring statistical dependence with hilbert-schmidt norms.  In _International conference on algorithmic learning theory_ , pp. 63–77. Springer, 2005.
  * Gretton et al. (2012) Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola.  A kernel two-sample test.  _The Journal of Machine Learning Research_ , 13(1):723–773, 2012.
  * Grünewälder et al. (2012) Steffen Grünewälder, Guy Lever, Luca Baldassarre, Sam Patterson, Arthur Gretton, and Massimilano Pontil.  Conditional mean embeddings as regressors.  In _Proceedings of the 29th International Coference on International Conference on Machine Learning_ , pp. 1803–1810, 2012.
  * Hill (2011) Jennifer L Hill.  Bayesian nonparametric modeling for causal inference.  _Journal of Computational and Graphical Statistics_ , 20(1):217–240, 2011.
  * Kennedy et al. (2017) Edward H Kennedy, Zongming Ma, Matthew D McHugh, and Dylan S Small.  Non-parametric methods for doubly robust estimation of continuous treatment effects.  _Journal of the Royal Statistical Society: Series B (Statistical Methodology)_ , 79(4):1229–1245, 2017.
  * Lei et al. (2021) Lihua Lei, Alexander D’Amour, Peng Ding, Avi Feller, and Jasjeet Sekhon.  Distribution-free assessment of population overlap in observational studies.  Technical report, Working paper, Stanford University, 2021.
  * Martinez-Taboada & Kennedy (2023) Diego Martinez-Taboada and Edward H Kennedy.  Counterfactual density estimation using kernel stein discrepancies.  _arXiv preprint arXiv:2309.16129_ , 2023.
  * Martinez-Taboada et al. (2023) Diego Martinez-Taboada, Aaditya Ramdas, and Edward H Kennedy.  An efficient doubly-robust test for the kernel treatment effect.  _arXiv preprint arXiv:2304.13237_ , 2023.
  * Mastouri et al. (2021) Afsaneh Mastouri, Yuchen Zhu, Limor Gultchin, Anna Korba, Ricardo Silva, Matt Kusner, Arthur Gretton, and Krikamol Muandet.  Proximal causal learning with kernels: Two-stage estimation and moment restriction.  In _International conference on machine learning_ , pp. 7512–7523. PMLR, 2021.
  * Mitrovic et al. (2018) Jovana Mitrovic, Dino Sejdinovic, and Yee Whye Teh.  Causal inference via kernel deviance measures.  _Advances in neural information processing systems_ , 31, 2018.
  * Moodie et al. (2018) Erica EM Moodie, Olli Saarela, and David A Stephens.  A doubly robust weighting estimator of the average treatment effect on the treated.  _Stat_ , 7(1):e205, 2018.
  * Muandet et al. (2017) Krikamol Muandet, Kenji Fukumizu, Bharath Sriperumbudur, Bernhard Schölkopf, et al.  Kernel mean embedding of distributions: A review and beyond.  _Foundations and Trends® in Machine Learning_ , 10(1-2):1–141, 2017.
  * Muandet et al. (2020) Krikamol Muandet, Arash Mehrjou, Si Kai Lee, and Anant Raj.  Dual instrumental variable regression.  _Advances in Neural Information Processing Systems_ , 33:2710–2721, 2020.
  * Muandet et al. (2021) Krikamol Muandet, Motonobu Kanagawa, Sorawit Saengkyongam, and Sanparith Marukatat.  Counterfactual mean embeddings.  _Journal of Machine Learning Research_ , 22(162):1–71, 2021.
  * Ojala & Garriga (2010) Markus Ojala and Gemma C Garriga.  Permutation tests for studying classifier performance.  _Journal of machine learning research_ , 11(6), 2010.
  * Park et al. (2021) Junhyung Park, Uri Shalit, Bernhard Schölkopf, and Krikamol Muandet.  Conditional distributional treatment effect with kernel conditional mean embeddings and u-statistic regression.  In _International Conference on Machine Learning_ , pp. 8401–8412. PMLR, 2021.
  * Pimentel (2022) Samuel D Pimentel.  Covariate-adaptive randomization inference in matched designs.  _arXiv preprint arXiv:2207.05019_ , 2022.
  * Ramdas et al. (2023) Aaditya Ramdas, Rina Foygel Barber, Emmanuel J Candès, and Ryan J Tibshirani.  Permutation tests using arbitrary permutation distributions.  _Sankhya A_ , pp. 1–22, 2023.
  * Robins & Rotnitzky (1995) James M Robins and Andrea Rotnitzky.  Semiparametric efficiency in multivariate regression models with missing data.  _Journal of the American Statistical Association_ , 90(429):122–129, 1995.
  * Rosenbaum (1984) Paul R Rosenbaum.  Conditional permutation tests and the propensity score in observational studies.  _Journal of the American Statistical Association_ , 79(387):565–574, 1984.
  * Rosenbaum (2002) Paul R Rosenbaum.  Observational studies.  In _Observational Studies_ , pp. 1–17. Springer, 2002.
  * Shekhar et al. (2022) Shubhanshu Shekhar, Ilmun Kim, and Aaditya Ramdas.  A permutation-free kernel two-sample test.  _Advances in Neural Information Processing Systems_ , 35:18168–18180, 2022.
  * Shimoni et al. (2018) Y. Shimoni, C. Yanover, E. Karavani, and Y. Goldschmnidt.  Benchmarking Framework for Performance-Evaluation of Causal Inference Analysis.  _ArXiv preprint arXiv:1802.05046_ , 2018.
  * Singh et al. (2019) Rahul Singh, Maneesh Sahani, and Arthur Gretton.  Kernel instrumental variable regression.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Singh et al. (2020) Rahul Singh, Liyuan Xu, and Arthur Gretton.  Generalized kernel ridge regression for nonparametric structural functions and semiparametric treatment effects.  _arXiv e-prints_ , pp. arXiv–2010, 2020.
  * Song et al. (2009) Le Song, Jonathan Huang, Alex Smola, and Kenji Fukumizu.  Hilbert space embeddings of conditional distributions with applications to dynamical systems.  In _Proceedings of the 26th Annual International Conference on Machine Learning_ , pp. 961–968, 2009.
  * Stuart (2010) Elizabeth A Stuart.  Matching methods for causal inference: A review and a look forward.  _Statistical science: a review journal of the Institute of Mathematical Statistics_ , 25(1):1, 2010.
  * Tsiatis (2006) Anastasios A. Tsiatis.  _Semiparametric Theory and Missing Data_.  Springer, 2006.
  * Van Der Laan & Rubin (2006) Mark J Van Der Laan and Daniel Rubin.  Targeted maximum likelihood learning.  _The international journal of biostatistics_ , 2(1), 2006\.

##  Appendix A Proofs

For ease of writing, throughout this appendix we will let:

$\displaystyle r(x,t)$ $\displaystyle:=\mu_{Y\mid X=x,T=t}$ (5)
$\displaystyle\hat{r}_{n}(x,t)$ $\displaystyle:=\hat{\mu}^{(n)}_{Y\mid
X=x,T=t}$ (6)

###  A.1 Proof of Theorem 1

This proof follows under the following standard assumptions444These are the
same as in for example Kennedy et al. (2017):

  * •

The propensity is uniformly bounded away from zero. So there exist an
$\epsilon>0$ such that $e(x,t)>\epsilon$ for all $x$ in the support of the
covariates.

  * •

We assume that the estimated propensity score is uniformly bounded away from
$0$ and $1$ on the support of $P(X)$. So for all $x$ in the support of $P(X)$
and for all $n$ we have $\delta<\hat{e}_{n}(x,t)$ for some $\delta>0$.

Now let $m$ be the size of the test set we average over. So $m=\mathcal{O}(n)$
where $n$ is the sample size.

First note due to the consistency property we can write
$k(y_{i},\cdot)\mathbbm{1}\\!\left\\{t_{i}=t\right\\}=k(y_{i}(t),\cdot)\mathbbm{1}\\!\left\\{t_{i}=t\right\\}$.
This allows us to rewrite the DR estimator as:

$\displaystyle\hat{\mu}^{\mathrm{DR}}_{Y(t)}=\frac{1}{m}\sum_{i=1}^{m}\left\\{\ell(y_{i}(t),\cdot)+\frac{\left(\mathbbm{1}\\{t_{i}=1\\}-\hat{e}_{n}(x_{i},t)\right)\left(\ell(y_{i}(t),\cdot)-\hat{r}_{n}(x_{i},t)\right)}{\hat{e}_{n}(x_{i},t)}\right\\}.$

Now letting:

$\displaystyle\epsilon^{\mathrm{DR}}_{Y(t)}=\frac{1}{m}\sum_{i=1}^{m}\frac{\left(\mathbbm{1}\\!\left\\{t_{i}=t\right\\}-\hat{e}_{n}(x_{i},t)\right)\left(\ell(y_{i}(t),\cdot)-\hat{r}_{n}(x_{i},t)\right)}{\hat{e}_{n}(x_{i},t)},$

we can write:

$\displaystyle\left\lVert{\mu}_{Y(t)}-\hat{\mu}^{\mathrm{DR}}_{Y(t)}\right\rVert_{\mathcal{H}_{Y}}\leq\left\lVert{\mu}_{Y(t)}-\frac{1}{m}\sum^{m}_{i=1}\ell(y_{i}(t),\cdot)\right\rVert_{\mathcal{H}_{Y}}+\left\lVert\epsilon^{\mathrm{DR}}_{Y(t)}\right\rVert_{\mathcal{H}_{Y}}.$

As $\frac{1}{m}\sum_{i=1}^{m}\ell(y_{i}(t),\cdot)$ is an unbiased estimator
for ${\mu}_{Y(t)}$ with parametric convergence rate
$\mathcal{O}(m^{-\frac{1}{2}})=\mathcal{O}(n^{-\frac{1}{2}})$, the convergence
rate of our estimator is entirely determined by the convergence rate of
$\left\lVert\epsilon^{\mathrm{DR}}_{Y(t)}\right\rVert_{\mathcal{H}_{Y}}$ so
consider:

$\displaystyle\mathbb{E}\left\lVert\epsilon^{\mathrm{DR}}_{Y(t)}\right\rVert_{\mathcal{H}_{Y}}^{2}$
$\displaystyle=\mathbb{E}\left[\frac{1}{m^{2}}\sum_{i,j=1}^{m}\frac{\left(\mathbbm{1}\\!\left\\{t_{i}=t\right\\}-\hat{e}_{n}(x_{i},t)\right)\left(\mathbbm{1}\\{t_{j}=t\\}-\hat{e}_{n}(x_{j},t)\right)\left\langle
y_{j}(t)-\hat{r}_{n}(x_{j},t),y_{i}(t)-\hat{r}_{n}(x_{i},t)\right\rangle_{\mathcal{H}_{Y}}}{\hat{e}_{n}(x_{i},t)\hat{e}_{n}(x_{j},t)}\right]$
$\displaystyle=\mathbb{E}\Bigg{[}\frac{1}{m^{2}}\sum_{i\neq
j}\frac{\left(\mathbbm{1}\\!\left\\{t_{i}=t\right\\}-\hat{e}_{n}(x_{i},t)\right)\left(\mathbbm{1}\\{t_{j}=t\\}-\hat{e}_{n}(x_{j},t)\right)\left(\langle
y_{j}(t)-\hat{r}_{n}(x_{j},t),y_{i}(t)-\hat{r}_{n}(x_{i},t)\rangle_{\mathcal{H}_{Y}}\right)}{\hat{e}_{n}(x_{i},t)\hat{e}_{n}(x_{j},t)}$
(7)
$\displaystyle\qquad+\frac{1}{m^{2}}\sum_{i=j}\left\\{\frac{\left(\mathbbm{1}\\!\left\\{t_{i}=t\right\\}-\hat{e}_{n}(x_{i},t)\right)\left\lVert
y_{i}(t)-\hat{r}_{n}(x_{i},t)\right\rVert_{\mathcal{H}_{Y}}}{\hat{e}_{n}(x_{i},t)}\right\\}^{2}\
\Bigg{]}$ (8)

Now let $(X,T,Y(t))$ be a random sample. We may write the term (8) as:

$\displaystyle\frac{1}{m}\mathbb{E}\left\\{\frac{\left(\mathbbm{1}\\{T=t\\}-\hat{e}_{n}(X,t)\right)\left\lVert
Y(t)-\hat{r}_{n}(X,t)\right\rVert_{\mathcal{H}_{Y}}}{\hat{e}_{n}(X,t)}\right\\}^{2}.$

By applying bounded propensity and bounded kernels we have that this
expectation is bounded, which gives the term is bounded by $C_{1}/n$. Now for
the first term let $(\tilde{X},\tilde{T},\tilde{Y}(t))$ be a second
independent sample. We may write (7) as:

$\displaystyle\frac{m-1}{m}\mathbb{E}\left[\frac{\left(\mathbbm{1}\\{T=t\\}-\hat{e}_{n}(X,t)\right)\left(\mathbbm{1}\\{\tilde{T}=t\\}-\hat{e}_{n}(\tilde{X},t)\right)\left\langle
Y(t)-\hat{r}_{n}(X,t),\tilde{Y}(t)-\hat{r}_{n}(\tilde{X},t)\right\rangle_{\mathcal{H}_{Y}}}{\hat{e}_{n}(X,t)\hat{e}_{n}(\tilde{X},t)}\right]$
$\displaystyle=$
$\displaystyle\frac{m-1}{m}\mathbb{E}\left[\mathbb{E}\left[\frac{\left(\mathbbm{1}\\{T=t\\}-\hat{e}_{n}(X,t)\right)\left(\mathbbm{1}\\{\tilde{T}=t\\}-\hat{e}_{n}(\tilde{X},t)\right)\left\langle
Y(t)-\hat{r}_{n}(X,t),\tilde{Y}(t)-\hat{r}_{n}(\tilde{X},t)\right\rangle_{\mathcal{H}_{Y}}}{\hat{e}_{n}(X,t)\hat{e}_{n}(\tilde{X},t)}\middle|X,\tilde{X}\right]\right]$
$\displaystyle=$
$\displaystyle\frac{m-1}{m}\mathbb{E}\left[\mathbb{E}\left[\frac{\left(\mathbbm{1}\\{T=t\\}-\hat{e}_{n}(X,t)\right)\left(\mathbbm{1}\\{\tilde{T}=t\\}-\hat{e}_{n}(\tilde{X},t)\right)}{\hat{e}_{n}(X,t)\hat{e}_{n}(\tilde{X},t)}\middle|X,\tilde{X}\right]\mathbb{E}\left[\left\langle
Y(t)-\hat{r}_{n}(X,t),\tilde{Y}(t)-\hat{r}_{n}(\tilde{X},t)\right\rangle_{\mathcal{H}_{Y}}\middle|X,\tilde{X}\right]\right]$
$\displaystyle=$
$\displaystyle\frac{m-1}{m}\mathbb{E}\left[\frac{\left(e(X,t)-\hat{e}_{n}(X,t)\right)\left(e(\tilde{X},t)-\hat{e}_{n}(\tilde{X},t)\right)}{\hat{e}_{n}(X,t)\hat{e}_{n}(\tilde{X},t)}\left\langle
r(X,t)-\hat{r}_{n}(X,t),r(\tilde{X},t)-\hat{r}_{n}(\tilde{X},t)\right\rangle_{\mathcal{H}_{Y}}\right]$
$\displaystyle\leq$
$\displaystyle\frac{m-1}{m}\left(\mathbb{E}\left[\left(\frac{\left(e(X,t)-\hat{e}_{n}(X,t)\right)\left(e(\tilde{X},t)-\hat{e}_{n}(\tilde{X},t)\right)}{\hat{e}_{n}(X,t)\hat{e}_{n}(\tilde{X},t)}\right)^{2}\right]\mathbb{E}\left[\left\langle
r(X,t)-\hat{r}_{n}(X,t),r(\tilde{X},t)-\hat{r}_{n}(\tilde{X},t)\right\rangle_{\mathcal{H}_{Y}}^{2}\right]\right)^{\frac{1}{2}}$
$\displaystyle\leq$
$\displaystyle\frac{m-1}{m}\mathbb{E}\left[\left(\frac{e(X,t)-\hat{e}_{n}(X,t)}{\hat{e}_{n}(X,t)}\right)^{2}\right]\mathbb{E}\left[\left\lVert
r(X,t)-\hat{r}_{n}(X,t)\right\rVert_{\mathcal{H}_{Y}}^{2}\right]$
$\displaystyle\leq$ $\displaystyle
C_{2}\mathbb{E}\left[\left(e(X,t)-\hat{e}_{n}(X,t)\right)^{2}\right]\mathbb{E}\left[\left\lVert
r(X,t)-\hat{r}_{n}(X,t)\right\rVert_{\mathcal{H}_{Y}}^{2}\right].$

Therefore we have:

$\displaystyle\mathbb{E}\left\lVert\epsilon^{\mathrm{DR}}_{Y(t)}\right\rVert_{\mathcal{H}_{Y}}^{2}\leq\frac{C_{1}}{m}+C_{2}\mathbb{E}\left[\left(e(X,t)-\hat{e}_{n}(X,t)\right)^{2}\right]\mathbb{E}\left[\left\lVert
r(X,t)-\hat{r}_{n}(X,t)\right\rVert_{\mathcal{H}_{Y}}^{2}\right].$

This gives that the rate of convergence is controlled by
$\mathbb{E}\left(e(X,t)-\hat{e}_{n}(X,t)\right)^{2}$ and
$\mathbb{E}\left\lVert
Y(t)-\hat{r}_{n}(X,t)\right\rVert_{\mathcal{H}_{Y}}^{2}$. By applying Jensen’s
inequality we have
$\mathbb{E}\left\lVert\epsilon^{\mathrm{DR}}_{Y(t)}\right\rVert_{\mathcal{H}_{Y}}=\mathcal{O}(\gamma_{r,n}\gamma_{e,n})$
and so by Markov’s inequality
$\left\lVert\epsilon^{\mathrm{DR}}_{Y(t)}\right\rVert_{\mathcal{H}_{Y}}=\mathcal{O}_{P}(\gamma_{r,n}\gamma_{e,n})$.

###  A.2 Proof of Theorem 2

The proof follows a similar structure to that of Theorem 1 where we now have:

$\displaystyle\left\lVert{\mu}_{Y(t)\mid
t=t^{\prime}}-\hat{\mu}^{\mathrm{DR}}_{Y(t)\mid
t=t^{\prime}}\right\rVert_{\mathcal{H}_{Y}}$
$\displaystyle\leq\left\lVert{\mu}_{Y(t)}-\frac{1}{n_{t^{\prime}}}\sum^{n}_{i=1}\mathbbm{1}\\!\left\\{t_{i}=t^{\prime}\right\\}\ell(y_{i}(t),\cdot)\right\rVert_{\mathcal{H}_{Y}}$
$\displaystyle+\left\lVert\left(\frac{1}{n_{t^{\prime}}}\sum^{n}_{i=1}\mathbbm{1}\\!\left\\{t_{i}=t^{\prime}\right\\}\ell(y_{i}(t),\cdot)\right)-\hat{\mu}^{\mathrm{DR}}_{Y(t)\mid
t=t^{\prime}}\right\rVert.$

Now we have:

$\displaystyle\left(\frac{1}{n_{t^{\prime}}}\sum^{n}_{i=1}\mathbbm{1}\\{t_{i}=t^{\prime}\\}\ell(y_{i}(t),\cdot)\right)-\hat{\mu}^{\mathrm{DR}}_{Y(t)\mid
t=t^{\prime}}$
$\displaystyle=\frac{1}{n_{t^{\prime}}}\sum_{i=1}^{n}\left\\{\mathbbm{1}\\!\left\\{t_{i}=t^{\prime}\right\\}\ell(y_{i}(t),\cdot)-\left({\mathbbm{1}\\!\left\\{t_{i}=t\right\\}w(x_{i},t)\left(\ell(y_{i}(t),\cdot)-r(x_{i},t)\right)}+\mathbbm{1}\\{t_{i}=t^{\prime}\\}r(x_{i},t)\right)\right\\}$
$\displaystyle=-\frac{1}{n_{t^{\prime}}}\sum_{i=1}^{n}\left\\{{\mathbbm{1}\\!\left\\{t_{i}=t\right\\}w(x_{i},t)\left(\ell(y_{i}(t),\cdot)-r(x_{i},t)\right)}+\mathbbm{1}\\{t_{i}=t^{\prime}\\}\left(r(x_{i},t)-\ell(y_{i}(t),\cdot)\right)\right\\}$
$\displaystyle=-\frac{1}{n_{t^{\prime}}}\sum_{i=1}^{n}\left\\{\left(\ell(y_{i}(t),\cdot)-r(x_{i},t)\right)\left({\mathbbm{1}\\!\left\\{t_{i}=t\right\\}w(x_{i},t)}-\mathbbm{1}\\!\left\\{t_{i}=t^{\prime}\right\\}\right)\right\\}$
$\displaystyle=-\frac{1}{n_{t^{\prime}}}\sum_{i=1}^{n}\left\\{\left(\ell(y_{i}(t),\cdot)-r(x_{i},t)\right)\left({\mathbbm{1}\\!\left\\{t_{i}=t\right\\}w(x_{i},t)}-\mathbbm{1}\\{t_{i}=t^{\prime}\\}\right)\right\\}$
$\displaystyle=\frac{1}{n_{t^{\prime}}}\sum_{i=1}^{n}\left\\{\left(\ell(y_{i}(t),\cdot)-r(x_{i},t)\right)\left[\mathbbm{1}\\!\left\\{t_{i}=t\right\\}\left(\frac{1-e(x_{i},t)}{e(x_{i},t)}\right)-\mathbbm{1}\\{t_{i}=t^{\prime}\\}\frac{e(x_{i},t)}{e(x_{i},t)}\right]\right\\}$
$\displaystyle=\frac{1}{n_{t^{\prime}}}\sum_{i=1}^{n}\left\\{\left(\ell(y_{i}(t),\cdot)-r(x_{i},t)\right)\frac{\mathbbm{1}\\!\left\\{t_{i}=t\right\\}-\mathbbm{1}\\!\left\\{t_{i}=t\right\\}e(x_{i},t)-\mathbbm{1}\\!\left\\{t_{i}=t^{\prime}\right\\}e(x_{i},t)}{e(x_{i},t)}\right\\}$
$\displaystyle=\frac{1}{n_{t^{\prime}}}\sum_{i=1}^{n}\left\\{\left(\ell(y_{i}(t),\cdot)-r(x_{i},t)\right)\frac{\mathbbm{1}\\!\left\\{t_{i}=t\right\\}-e(x_{i},t)}{e(x_{i},t)}\right\\}$
$\displaystyle=\frac{n_{t^{\prime}}}{n}\epsilon^{\mathrm{DR}}_{Y(t)}$

Now as $n_{t^{\prime}}/n$ converges to a constant at rate $n^{-\frac{1}{2}}$
the convergence rate just depends on the rate that
$\epsilon^{\mathrm{DR}}_{Y(t)}$ tends to zero. Therefore the analysis in the
previous proof establishes the rate.

###  A.3 Proof of Proposition 1

Ramdas et al. (2023) demonstrate that we can test the null hypothesis

$\displaystyle H_{0}:X_{1},\ldots,X_{n}\text{ are exchangeable}$

from a fixed set of permutations $S\subset S_{n}$ (not necessarily a group)
using a statistic $T$ by randomly sampling permutations
$\sigma_{0},\ldots,\sigma_{M}$ from $S$ uniformly and computing:

$\displaystyle
p=\frac{1+\sum^{M}_{i=1}\mathbbm{1}\\!\left\\{T(X_{\sigma_{i}\circ\sigma^{-1}_{0}})\geq
T(X)\right\\}}{M+1},$ (9)

where $X_{\sigma}=\left(X_{\sigma(1)},\ldots,X_{\sigma(n)}\right)$ is the
permuted data. Ramdas et al. show that $p$ is a valid p-value in the sense
that $P(p\leq\alpha)\leq\alpha$ under the null hypothesis.

In our method the data is placed into bins, denoted by $B_{i}$, where $i$
ranges over the total number of bins $k$, and under the assumption of exact
matching the null hypothesis is:

$\displaystyle H_{0}:T_{B_{i_{1}}},\ldots,T_{B_{i_{n}}}\text{ are exchangeable
for each $i$.}$

Further we have split the bins into training a $\mathcal{B}_{{\rm
Tr}}=\\{B_{1},\ldots,B_{m}\\}$ and $\mathcal{B}_{{\rm
Te}}=\\{B_{m+1},\ldots,B_{k}\\}$. For $B_{i}\in\mathcal{B}_{{\rm Te}}$ we set
the possible permutations to be $S_{\lvert B_{i}\rvert}$ i.e. the full set of
permutations on $B_{i}$. For each set $B_{j}\in\mathcal{B}_{{\rm Te}}$ we
sample ${\sigma_{1}}_{j},\ldots,{\sigma_{N+1}}$ from $S_{\lvert B_{j}\rvert}$
and set $S_{j}=\\{{\sigma_{1}}_{j},\ldots,{\sigma_{N+1}}\\}$. We then set the
total set of permutations $S$ to be:

$\displaystyle
S=\left(\bigtimes^{m}_{j=1}S_{j}\right)\times\left(\bigtimes^{k}_{j=m+1}S_{\lvert
B_{i}\rvert}\right)$

As the samples in each bin are exchangeable, and samples in distinct bins are
independent, we have that our data is exchangeable under the permutations in
$S$. Therefore we can sample from $S$ as in (9) to produce a valid p-value
under $H_{0}$.

To go from this to the form in the proof note that sampling the random
permutations for $S_{j}$, choosing a random permutation $\sigma_{0}$ from
$S_{j}$ and then taking $\sigma_{i}\circ\sigma^{-1}_{0}$ is equivalent to
sampling $N$ random permutations, ${\sigma_{1}}_{j},\ldots,{\sigma_{N}}$ and
then sampling a random permutation from
$\tilde{S}=\\{\operatorname{Id},{\sigma_{1}}_{j},\ldots,{\sigma_{N}}\\}$.

##  Appendix B Derivation of Test Statistics

###  B.1 Derivation of
$\widehat{\mathrm{MMD}}^{2}_{\mathrm{DR}}[Y(1),Y(0),\mathcal{H}_{\mathcal{Y}}]$

To derive
$(\widehat{\mathrm{MMD}}_{\mathrm{DR}}[Y(1),Y(0),\mathcal{H}_{\mathcal{Y}}])^{2}$
we let $e(x)=e(x,1)=P(T=1\mid X=x)$. We now write:

$\displaystyle\hat{\mu}_{Y(1)}^{\mathrm{DR}}-\hat{\mu}^{\mathrm{DR}}_{Y(0)}$
$\displaystyle=\frac{1}{m}\sum_{i=1}^{m}\left\\{\frac{t_{i}\left(\ell(y_{i},\cdot)-\hat{r}(x_{i},1)\right)}{\hat{e}(x_{i})}+\hat{r}(x_{i},1)-\left(\frac{\left(1-t_{i}\right)\left(\ell(y_{i},\cdot)-\hat{r}(x_{i},0)\right)}{1-\hat{e}(x_{i})}+\hat{r}(x_{i},0)\right)\right\\}$
$\displaystyle=\frac{1}{m}\sum_{i=1}^{m}\frac{t_{i}-e(x_{i})}{e(x_{i})(1-e(x_{i}))}\left(\ell(y_{i},\cdot)-(1-e(x_{i}))\hat{r}(x_{i},1)+e(x_{i})\hat{r}(x_{i},0)\right).$

Now we have
$\hat{r}(x_{i},1)=\mathbf{k}^{\top}_{1}(x_{i})\mathbf{W}_{1}\boldsymbol{\ell}_{1}$
and
$\hat{r}(x_{i},0)=\mathbf{k}^{\top}_{0}(x_{i})\mathbf{W}_{0}\boldsymbol{\ell}_{0}$.
Further for ease of writing we let:

$\displaystyle\tilde{\mathbf{k}}^{\top}_{0}(x_{i})$
$\displaystyle=e(x_{i})\mathbf{k}^{\top}_{0}(x_{i}),$
$\displaystyle\tilde{\mathbf{k}}^{\top}_{1}(x_{i})$
$\displaystyle=(1-e(x_{i}))\mathbf{k}^{\top}_{1}(x_{i}),$
$\displaystyle\alpha(x_{i})$
$\displaystyle=\frac{(t_{i}-e(x_{i}))}{(e(x_{i})(1-e(x_{i}))}.$

This allows us to write the above as

$\displaystyle=\frac{1}{m}\sum_{i=1}^{m}\alpha(x_{i})\left(\ell(y_{i},\cdot)+\tilde{\mathbf{k}}^{\top}_{0}(x_{i})\mathbf{W}_{0}\boldsymbol{\ell}_{0}-\tilde{\mathbf{k}}^{\top}_{1}(x_{i})\mathbf{W}_{1}\boldsymbol{\ell}_{1}\right),$

and plugging this in we obtain

$\displaystyle(\widehat{\mathrm{MMD}}_{\mathrm{DR}}[Y(1),Y(0),\mathcal{H}_{\mathcal{Y}}])^{2}:=\left\lVert\hat{\mu}_{Y(1)}^{\mathrm{DR}}-\hat{\mu}^{\mathrm{DR}}_{Y(0)}\right\rVert_{\mathcal{H}_{\mathcal{Y}}}^{2}$
$\displaystyle=\sum_{i,j}\alpha(x_{i})\alpha(x_{j})\left\langle\ell(y_{i},\cdot)+\tilde{\mathbf{k}}^{\top}_{0}(x_{i})\mathbf{W}_{0}\boldsymbol{\ell}_{0}-\tilde{\mathbf{k}}^{\top}_{1}(x_{i})\mathbf{W}_{1}\boldsymbol{\ell}_{1},\ell(y_{j},\cdot)+\tilde{\mathbf{k}}^{\top}_{0}(x_{j})\mathbf{W}_{0}\boldsymbol{\ell}_{0}-\tilde{\mathbf{k}}^{\top}_{1}(x_{j})\mathbf{W}_{1}\boldsymbol{\ell}_{1}\right\rangle$
$\displaystyle=\sum_{i,j}\alpha(x_{i})\alpha(x_{j})\left(\ell(y_{i},y_{j})+\tilde{\mathbf{k}}^{\top}_{0}(x_{i})\mathbf{W}_{0}\boldsymbol{\ell}_{0}(y_{j})+\tilde{\mathbf{k}}^{\top}_{0}(x_{j})\mathbf{W}_{0}\boldsymbol{\ell}_{0}(y_{i})-\tilde{\mathbf{k}}^{\top}_{1}(x_{i})\mathbf{W}_{1}\boldsymbol{\ell}_{1}(y_{j})-\tilde{\mathbf{k}}^{\top}_{1}(x_{j})\mathbf{W}_{1}\boldsymbol{\ell}_{1}(y_{i})\right)$
$\displaystyle+\sum_{i,j}\alpha(x_{i})\alpha(x_{j})\left(\tilde{\mathbf{k}}^{\top}_{0}(x_{i})\mathbf{W}_{0}\mathbf{L}_{0}\mathbf{W}_{0}\tilde{\mathbf{k}}_{0}(x_{j})+\tilde{\mathbf{k}}^{\top}_{1}(x_{i})\mathbf{W}_{1}\mathbf{L}_{1}\mathbf{W}_{1}\tilde{\mathbf{k}}_{1}(x_{j})-\tilde{\mathbf{k}}^{\top}_{0}(x_{i})\mathbf{W}_{0}\mathbf{L}_{0,1}\mathbf{W}_{1}\tilde{\mathbf{k}}_{1}(x_{j})-\tilde{\mathbf{k}}^{\top}_{1}(x_{i})\mathbf{W}_{1}\mathbf{L}_{0,1}\mathbf{W}_{0}\tilde{\mathbf{k}}_{0}(x_{j})\right).$

Now we introduce some new notation: for a variable we will use a subscript to
denote the value of $T$ we condition and a superscript to denote if the
variable is in training or test. So $X^{{\rm Te}}_{0}$ is the test points with
$T=0$. Further let
$K(X,\tilde{X})=\left(k(x_{i},\tilde{x}_{j})\right)_{i,j=1}^{i=n_{X},j=n_{\tilde{X}}}$
, so the kernel matrix where rows. This allows us to write the test statistic
as:

$\displaystyle\mathbf{\alpha}^{\top}\bigg{(}L(Y^{{\rm Te}},Y^{{\rm Te}})$
$\displaystyle+2\mathrm{diag}(\mathbf{e})K(X_{0}^{{\rm Tr}},X^{{\rm
Te}})^{\top}\mathbf{W}_{0}L(Y_{0}^{{\rm Tr}},Y^{{\rm
Te}})-2\mathrm{diag}(\mathbf{1-e})K(X_{1}^{{\rm Tr}},X^{{\rm
Te}})^{\top}\mathbf{W}_{1}L(Y_{1}^{{\rm Tr}},Y^{{\rm Te}})$
$\displaystyle+\mathrm{diag}(\mathbf{e})K(X_{0}^{{\rm Tr}},X^{{\rm
Te}})^{\top}\mathbf{W}_{0}L(Y^{{\rm Tr}}_{0},Y^{{\rm
Tr}}_{0})\mathbf{W}_{0}K(X_{0}^{{\rm Tr}},X^{{\rm
Te}})\mathrm{diag}(\mathbf{e})$
$\displaystyle+\mathrm{diag}(\mathbf{1-e})K(X_{1}^{{\rm Tr}},X^{{\rm
Te}})^{\top}\mathbf{W}_{1}L(Y^{{\rm Tr}}_{1},Y^{{\rm Tr}}_{1})K(X_{1}^{{\rm
Tr}},X^{{\rm Te}})\mathbf{W}_{1}\mathrm{diag}(\mathbf{1-e})$
$\displaystyle-2\mathrm{diag}(\mathbf{e})K(X_{0}^{{\rm Tr}},X^{{\rm
Te}})^{\top}\mathbf{W}_{0}L(Y^{{\rm Tr}}_{0},Y^{{\rm Tr}}_{1})K(X_{1}^{{\rm
Tr}},X^{{\rm
Te}})\mathbf{W}_{1}\mathrm{diag}(\mathbf{1-e})\bigg{)}\mathbf{\alpha}.$

###  B.2 Derivation of
$\widehat{\mathrm{MMD}_{\mathrm{DR}}}\Big{[}\left\\{Y(1)\mid
T=0\right\\},\left\\{Y(0)\mid T=0\right\\},\mathcal{H}_{\mathcal{Y}}\Big{]}$

We now derive the closed form of the estimator based on the effect of
treatment on the treated for the case of a binary treatment. Again we begin by
deriving a simpler form of the difference between mean embeddings:

$\displaystyle\hat{\mu}_{\left\\{Y(1)\mid
T=0\right\\}}-\hat{\mu}_{\left\\{Y(0)\mid T=0\right\\}}$
$\displaystyle=\hat{\mu}_{\left\\{Y(1)\mid
T=0\right\\}}-\hat{\mu}_{\left\\{Y\mid T=0\right\\}}$
$\displaystyle=\frac{1}{n_{0}}\sum_{i=1}^{n}\bigg{(}{t_{i}\hat{w}(x_{i},1)\left(\ell(y_{i},\cdot)-\hat{r}(x_{i},1)\right)}+\left(1-t_{i}\right)\hat{r}(x_{i},1)\bigg{)}-\frac{1}{n_{0}}\sum_{i=1}^{n}\left(1-t_{i}\right)\ell(y_{i},\cdot)$
$\displaystyle=\frac{1}{n_{0}}\sum_{i=1}^{n}\left(t_{i}\hat{w}(x_{i},1)-\left(1-t_{i}\right)\right)\left(\ell(y_{i},\cdot)-\hat{r}(x_{i},1)\right)$
$\displaystyle=\frac{1}{n_{0}}\sum_{i=1}^{n}\left(\frac{t_{i}-e(x_{i})}{e(x_{i})}\right)\left(\ell(y_{i},\cdot)-\hat{r}(x_{i},1)\right)$

Now we have that $\hat{r}(x_{i},1)=\mu_{Y\mid
X=x,T=1}=\mathbf{k}^{\top}_{1}(x)\mathbf{W}_{1}\mathbf{l}_{1}$ and we also
let:

$\displaystyle\beta(x_{i})=\frac{t-e(x_{i})}{e(x_{i})}$

Then we can write the full test stat as:

$\displaystyle\left\lVert\hat{\mu}_{\left\\{Y(1)\mid
T=0\right\\}}-\hat{\mu}_{\left\\{Y(0)\mid T=0\right\\}}\right\rVert^{2}$
$\displaystyle=\frac{1}{n^{2}_{0}}\sum_{i,j}\beta(x_{i})\beta(x_{j})\left\langle\ell(y_{i},\cdot)-\mathbf{k}^{\top}_{1}(x_{i})\mathbf{W}_{1}\mathbf{l}_{1},\ell(y_{j},\cdot)-\mathbf{k}^{\top}_{1}(x_{j})\mathbf{W}_{1}\mathbf{l}_{1}\right\rangle$
$\displaystyle=\frac{1}{n^{2}_{0}}\sum_{i,j}\beta(x_{i})\beta(x_{j})\left(\ell(y_{i},y_{j})-\mathbf{k}^{\top}_{1}(x_{i})\mathbf{W}_{1}\mathbf{l}_{1}(x_{j})-\mathbf{k}^{\top}_{1}(x_{i})\mathbf{W}_{1}\mathbf{l}_{1}(x_{i})+\phantom{a}\right.$
$\displaystyle\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\left.\phantom{a}\mathbin{+}\mathbf{k}^{\top}_{1}(x_{j})\mathbf{W}_{1}\mathbf{L}_{1,1}\mathbf{W}_{1}\mathbf{k}^{\top}_{1}(x_{j})\right),$

giving the full test statistic as:

$\displaystyle\mathbf{\beta}^{\top}\left(L(Y^{{\rm Te}},Y^{{\rm
Te}})-2K(X_{1}^{{\rm Tr}},X^{{\rm Te}})^{\top}\mathbf{W}_{1}L(Y_{1}^{{\rm
Tr}},Y^{{\rm Te}})+K(X_{1}^{{\rm Tr}},X^{{\rm
Te}})^{\top}\mathbf{W}_{1}L(Y_{1}^{{\rm Tr}},Y_{1}^{{\rm
Tr}})\mathbf{W}_{1}K(X_{1}^{{\rm Tr}},X^{{\rm Te}})\right)\mathbf{\beta}.$

##  Appendix C Simulation Details

For the simulation in Section 5.1 we select:

$\displaystyle\mathbf{a}$
$\displaystyle=[0.1,0.2,0.3,0.4,0.5,0.1,0.2,0.3,0.4]$
$\displaystyle\mathbf{b}$
$\displaystyle=[0.5,0.4,0.3,0.2,0.1,0.4,0.3,0.2,0.1]$ $\displaystyle\sigma$
$\displaystyle=0.2$ $\displaystyle\beta$ $\displaystyle=3,$

and for the second experiment we let:

$\displaystyle f_{0}(x)$ $\displaystyle=x_{1}$ $\displaystyle f_{1}(x)$
$\displaystyle=x_{1}^{2}$ $\displaystyle\alpha$ $\displaystyle=0.3$
$\displaystyle\sigma$ $\displaystyle=0.2.$

###  C.1 Fit Experiment Plots

Here we include the propensity plots for the simulated data in the fit plots:

We can see in the second plot with one sided overlap a much higher proportion
of datapoints are concentrated at extreme propensities.

###  C.2 Simulated Distributional Test Plots

Here we include some plots for the distribution in the experiments testing for
causal effects on simulated data:

And also the distribution over propensity scores: