Geodesic Causal Inference

Consider a uniquely geodesic metric space $(\Omega,d)$ that is endowed with a probability measure $P$. For any $\alpha,\beta\in\Omega$, the unique geodesic connecting $\alpha$ and $\beta$ is a curve $\gamma_{\alpha,\beta}:[0,1]\mapsto\Omega$ such that $d(\gamma_{\alpha,\beta}(s),\gamma_{\alpha,\beta}(t))=d(\alpha,\beta)|t-s|$ for $s,t\in[0,1]$. For $\alpha,\beta,\zeta\in\Omega$, define the following simple operations on geodesics,

The space $(\Omega,d)$ is further assumed to be geodesically complete (Gallier and Quaintance,, 2020), meaning that every geodesic can be extended to a geodesic defined on $\mathbb{R}$. Specifically, for any $\alpha,\beta\in\Omega$, the geodesic $\gamma_{\alpha,\beta}:[0,1]\mapsto\Omega$ can be extended to $\bar{\gamma}_{\alpha,\beta}:\mathbb{R}\mapsto\Omega$. The restriction of $\bar{\gamma}_{\alpha,\beta}$ to $[0,1]$ coincides with $\gamma_{\alpha,\beta}$, i.e., $\bar{\gamma}_{\alpha,\beta}|_{[0,1]}=\gamma_{\alpha,\beta}$. We can then define a scalar multiplication for geodesics $\gamma_{\alpha,\beta}$ with a factor $\rho\in\mathbb{R}$ by

In practical applications, the relevant metric space often constitutes a subset $(\mathcal{M},d)$ of $(\Omega,d)$, which is typically closed and convex. Prominent instances of such spaces are given in Examples 2.1–2.4 below. We use these same example spaces to demonstrate the practical performance of the proposed geodesic treatment effect in simulations and real-world data analyses. In Appendix B we demonstrate that these examples satisfy the pertinent assumptions outlined in Sections 3 and 4; see Propositions B.1–B.4. Geodesic extensions are confined to operate within the closed subset $\mathcal{M}\subset\Omega$, in some cases necessitating suitable modifications to ensure that they do not cross the boundary (Zhu and Müller,, 2023).

Considering the forward extension of a geodesic, assume that the geodesic $\gamma_{\alpha,\beta}$ extends to the boundary point $\zeta$ and denote the extended geodesic by $\gamma_{\alpha,\zeta}:[0,1]\mapsto\mathcal{M}$. For $\rho>1$, we then define a scalar multiplication by

where

It is straightforward to verify that the start point of $\rho\odot\gamma_{\alpha,\beta}$ is always $\alpha$, while the end point is $\beta$ for $\rho=1$ and $\zeta$ for $\rho=\infty$; see Figure 1 for an illustration. Similar modifications apply to the reverse extension where $\rho<-1$. To simplify the notation, we write the end point of $\rho\odot\gamma_{\alpha,\beta}$ as $\gamma_{\alpha,\beta}(\rho)$.

In Examples 2.1 and 2.2, for any two matrices $\alpha,\beta$ in the space, the geodesic is the line segment connecting $\alpha$ and $\beta$, i.e., $\gamma_{\alpha,\beta}(t)=\alpha+(\beta-\alpha)t$. In Example 2.3, for any two points $\alpha,\beta\in\mathcal{S}_{+}^{d-1}$, the geodesic connecting $\alpha$ and $\beta$ is

where $\theta=\arccos(\alpha^{\prime}\beta)$ represents the angle between $\alpha$ and $\beta$.

Write $\tau_{\#}\mu$ for the pushforward measure of $\mu$ by the transport $\tau$. In Example 2.4, for any two distributions $\alpha,\beta\in\mathcal{W}$, the geodesic between them is given by McCann’s interpolant (McCann,, 1997),

where $\mathrm{id}$, $F_{\alpha}$, $F_{\beta}^{-1}$ denote the identity map, cumulative distribution function of $\alpha$, and quantile function of $\beta$, respectively.

Let $(\mathcal{M},d)$ be a uniquely extendable geodesic metric space. Units $i=1,\dots,n$ are either given a treatment or not treated. For the $i$-th unit, we observe a treatment indicator variable $T_{i}$, where $T_{i}=1$ if the $i$-th unit is treated and $T_{i}=0$ otherwise, along with an outcome

where $Y_{i}(0),Y_{i}(1)\in\mathcal{M}$ are potential outcomes for $T_{i}=0$ and $1$, respectively. Additionally, we observe Euclidean covariates $X_{i}\in\mathbb{R}^{p}$. We assume that $\{Y_{i},T_{i},X_{i}\}_{i=1}^{n}$ is an i.i.d. sample from a joint distribution $P$, where a generic random variable distributed according to $P$ is $(Y,T,X)\in\mathcal{M}\times\{0,1\}\times\mathcal{X}$ and we assume that all conditional distributions are well-defined; expectations and conditional expectations in the following are with respect to $P$. We assume that $\mathcal{X}$ is a compact subset of $\mathbb{R}^{p}$ and write $Y=Y(t)$ if $T=t\in\{0,1\}$.

A standard way to quantify causal effects for the usual situation where outcomes are located in Euclidean spaces is to target the difference of the potential outcomes. However, when outcomes take values in a general metric space this is not an option anymore, due to the absence of linear structure in general metric spaces (see, e.g., Example 2.3), which means that the notion of a difference does not exist. This requires to extend the conventional notion of causal effects for the case of random objects in $\mathcal{M}$, and to address this challenge we propose here to replace the conventional difference by the geodesic that connects the Fréchet means of the potential outcomes.

Define the Fréchet mean of a random element $A\in\mathcal{M}$ and the conditional Fréchet mean of $A\in\mathcal{M}$ given $Z\in\mathbb{R}^{p}$ as

respectively. Assuming the minimizers are unique and well-defined, we introduce the following notions of causal effects in geodesic metric spaces.

In Euclidean spaces, the geodesic individual treatment effect and geodesic average treatment effect (GATE) reduce to the difference of the potential outcomes for the $i$-th unit, $Y_{i}(1)-Y_{i}(0)$, and the difference of the expected potential outcomes, $\operatorname{E}[Y(1)]-\operatorname{E}[Y(0)]$, respectively. These can be interpreted as the quantities that represent the shortest distance and directional information from $Y_{i}(0)$ to $Y_{i}(1)$ or from $\operatorname{E}[Y(0)]$ to $\operatorname{E}[Y(1)]$ with respect to the Euclidean metric. Since a geodesic connecting points $\alpha$ and $\beta$ in a geodesic metric space $(\mathcal{M},d)$ embodies both the shortest distance and directional information from $\alpha$ to $\beta$ with respect to the metric $d$, Definition 2.1 emerges as a natural extension of quantifying treatment effects in geodesic metric spaces.

To construct doubly robust estimators in metric spaces, one needs to extend the notion of doubly robust estimation of causal effects to the case of random objects in a geodesic metric space. In analogy to the Euclidean case, key ingredients are propensity scores $p(x)=\operatorname{P}(T=1|X=x)$ and conditional Fréchet means for the potential outcomes $m_{t}(x)=\operatorname{E}_{\oplus}[Y(t)|X=x]$ for $t\in\{0,1\}$. We impose the following conditions.

Assumptions 3.1 (ii) and (iii) are standard overlap and unconfoundedness (or ignorability) conditions in the context of causal inference. Assumption 3.1 (ii) says that for any $x\in\mathcal{X}$, there is a positive probability that there are units in both the treatment and control groups in any arbitrarily small neighborhood around $x$. Assumptions 3.1 (iii) says that conditionally on $X$, the treatment is randomly assigned from the potential outcomes. In the Euclidean case, this setup leads to various estimation methods for the average treatment effect, including inverse probability weighting, outcome regression, and doubly robust estimation methods.

Our goal here is to extend the doubly robust estimation to geodesic metric spaces, aiming at GATE (geodesic average treatment effect). To this end, we need to introduce the notion of a (Fréchet) regression model of $Y(t)$ on $X$ in geodesic metric spaces. More precisely, we impose the following conditions, where $\operatorname{E}_{\oplus}$ is as in (2.1).

The random perturbation map $\mathcal{P}_{t}$ is a generalization of additive noise $\varepsilon$ in the Euclidean case that is needed here because the addition operation is not defined, and generalizes the error assumption $\operatorname{E}[\varepsilon|X]=0$. Specifically, Assumption 3.2 (ii) is equivalent to $\operatorname{E}_{\oplus}[Y(t)|X]=m_{t}(X)$ so that $m_{t}$ becomes the Fréchet regression function of $Y(t)$ on $X$ (Petersen and Müller,, 2019). In the following, we refer to $m_{t}$ as the outcome regression function. Note that our conditions on the perturbation maps are generalizations of those considered in Chen and Müller, (2022).

To obtain a doubly robust representation of the GATE, we additionally introduce an expectation operation for geodesics. For random objects $A,B\in\mathcal{M}$ and Euclidean random variables $\kappa\in\mathbb{R}$ and $X\in\mathbb{R}^{p}$, define

These definitions of geodesic expectations $\operatorname{E}_{G}[\cdot]$, $\operatorname{E}_{G}[\cdot|X]$, and $\operatorname{E}_{G,X}[\cdot]$ are consistent with the corresponding (conditional) expectations in the Euclidean case if $\kappa$ and $\{A,B\}$ are conditionally independent given $X$. Indeed, if $\kappa$ and $\{A,B\}$ are conditionally independent given $X$ and $A,B$ are Euclidean, it is easy to see that $\operatorname{E}_{G,X}[\kappa\odot\gamma_{A,B}]=\operatorname{E}_{G}[\kappa\odot\gamma_{A,B}]$ by the law of iterated expectations.

Based on the above assumptions and notions, a doubly robust representation of the GATE for geodesic metric spaces is as follows.

This proposition shows that the GATE is a natural generalization of the average treatment effect for Euclidean data to the case of geodesic metric spaces (see Remark 3.1 below). Furthermore, we note the double robustness property of the GATE estimator, i.e., either $e(x)$ or $\mu_{t}(x)$ for $t\in\{0,1\}$ could be misspecified (but not both) while still achieving consistent estimation of the GATE. Since the GATE is defined via the geodesic between the Fréchet means for the potential outcomes, we first establish the representation from an arbitrary starting point $\mu$ in (3.1). Although the representation of the GATE in (3.2) also involves an arbitrary starting point $\mu$, the doubly robust, cross-fitting, and outcome regression GATE estimators presented below do not require to choose $\mu$; however, implementing the inverse probability weighting estimator requires choice of a starting point $\mu$ that can be arbitrary (see Remark 3.3 below for further details). We illustrate in Figure 2 the doubly robust representation of the GATE when the outcome regression model or alternatively the propensity score model is correctly specified.

Building upon Proposition 3.1, we can construct a sample based doubly robust estimator for the GATE. Note that the right-hand side of (3.2) can be rewritten as $\operatorname{E}_{G,X}[\gamma_{A_{0},A_{1}}]$ with

which corresponds to the representation in (3.1). By moving to the sample counterpart, a doubly robust estimator for the GATE is obtained as $\gamma_{\widehat{\Theta}_{0}^{(\mathrm{DR})},\widehat{\Theta}_{1}^{(\mathrm{DR})}}$, where

Here $e(x;\varphi)$ is a parametric model for the propensity score $p(x)$, $\widehat{\varphi}$ is an estimator of $\varphi$, and $\widehat{\mu}_{t}(x)$ is an estimator of the outcome regression function $m_{t}(x)$. In our real data analysis, we estimate the propensity score by the logistic regression model and the outcome regression function by global Fréchet regression (Petersen and Müller,, 2019):

where $I_{t}=\{1\leq i\leq n:T_{i}=t\}$, $N_{t}$ is the sample size of $I_{t}$, $\bar{X}=n^{-1}\sum_{i=1}^{n}X_{i}$, and $\widehat{\Sigma}=n^{-1}\sum_{i=1}^{n}(X_{i}-\bar{X})(X_{i}-\bar{X})^{\prime}$.

To examine the asymptotic properties of the proposed doubly robust and cross-fitting estimators, we impose the following assumptions on the models for the propensity score and outcome regression functions. Let $\mathcal{M}_{e}=\{e(x;\varphi):x\in\mathcal{X},\varphi\in\Phi\}$ be a class of parametric models for the propensity score $p(x)$, where $\Phi\subset\mathbb{R}^{p}$ is a compact set. For $t\in\{0,1\}$, let $\widehat{\mu}_{t}(\cdot)$ be an estimator for the outcome regression function $m_{t}(\cdot)$.

Assumptions 4.1 (i) and (ii) imply the uniform convergence rate for the propensity score estimator, $\sup_{x\in\mathcal{X}}|e(x;\widehat{\varphi})-e(x;\varphi_{*})|=O_{p}(\varrho_{n})$, regardless of whether the propensity score model is misspecified or not. These assumptions are typically satisfied with $\varrho_{n}=n^{-1/2}$ under mild regularity conditions. Assumption 4.1 (iii) concerns the outcome regression. If one employs the global Fréchet regression model (3.6) for estimating $m_{t}$, then under some regularity conditions, this assumption is satisfied with $r_{n}=n^{-\alpha^{\prime}}$ for some $\alpha^{\prime}>1/2$ (see Theorem 1 in Petersen and Müller, (2019)) and

where $\Sigma=\operatorname{Var}(X)$, again regardless of whether the model is misspecified or not for the conditional Fréchet mean, which is the true target. Alternatively, one can employ local Fréchet regression as outcome regression method, which is more flexible but comes with slower rates of convergence and is subject to the curse of dimensionality (see Theorem 1 in Chen and Müller, (2022)).

To establish the consistency of $\widehat{\Theta}_{t}^{(\mathrm{DR})}$ and $\widehat{\Theta}_{t}^{(\mathrm{CF})}$ in (3.2) and (3.7), we additionally require that the geodesics in $\mathcal{M}$ satisfy the following assumption.

This is a Lipschitz-type condition to control the criterion function (3.2) that involves the geodesics $\gamma_{\mu(X_{i}),Y_{i}}$. When the space $\Omega$ is a CAT($0$) space (i.e., a geodesic metric space that is globally non-positively curved in the sense of Alexandrov), this assumption is automatically satisfied. Many metric spaces of statistical interest are CAT($0$) spaces, such as the Wasserstein space of univariate distributions, the space of networks expressed as graph Laplacians with the (power) Frobenius metric, covariance/correlation matrices with the affine-invariant metric (Thanwerdas and Pennec,, 2023), the Cholesky metric (Lin,, 2019) and various other metrics (Pigoli et al.,, 2014), as well as phylogenetic tree space with the BHV metric (Billera et al.,, 2001; Lin and Müller,, 2021). Assumption 4.2 also holds for some relevant positively curved spaces such as open hemispheres or orthants of spheres that are relevant for compositional data with the Riemannian metric; these spaces are not CAT(0), see Example 2.3.

For $t\in\{0,1\}$, denote by $\Theta_{t}^{(\mathrm{DR})}:=\mathop{\rm arg~{}min}\limits_{\nu\in\mathcal{M}}Q_{t}(\nu;\mu_{t},\varphi_{*})$ the minimizer of the population criterion function, where

To guarantee the identification of the object of interest $\operatorname{E}_{\oplus}[Y(t)]$ by the doubly robust estimation approach, we add the following conditions.

Assumption 4.3 (i) is a standard separation condition to achieve the consistency of the M-estimator (see, e.g., Chapter 3.2 in van der Vaart and Wellner, (2023)). Note that the existence of the minimizers follows immediately if $\mathcal{M}$ is compact. Assumption 4.3 (ii) requires that at least one of the models for the propensity score and the outcome regression is correctly specified.

The following result provides the consistency of the doubly robust and cross-fitting estimators $\widehat{\Theta}_{t}^{(\mathrm{OR})}$ in (3.2) and $\widehat{\Theta}_{t}^{(\mathrm{CF})}$ in (3.7).

To obtain rates of convergence for estimators $\widehat{\Theta}_{t}^{(\mathrm{DR})}$ and $\widehat{\Theta}_{t}^{(\mathrm{CF})}$ in the metric space $(\Omega,d)$ we need additional entropy conditions to quantify the complexity of the space. Let $B_{\delta}(\omega)$ be the ball of radius $\delta$ centered at $\omega\in\Omega$ and $N(\varepsilon,B_{\delta}(\omega),d)$ its covering number using balls of size $\varepsilon$.