Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
75,351 characters · 13 sections · 93 citation commands
Matching for causal effects via multimarginal unbalanced optimal transport
Identifying cause and effect is one of the primary goals of scientific research. The leading approaches to uncover causal effects are randomized controlled trials. Unfortunately, such trials are often practically infeasible on ethical grounds, might not be generalizable beyond the experimental setting due to lack of variation in the population, or simply have too few participants to generate robust results due to financial or logistical restrictions. An attractive alternative is to use observational data, which are ubiquitous, often readily available, and comprehensive.
The main challenge in using observational data for causal inference is the fact that assignment into treatment is not perfectly randomized. This implies that individuals assigned to different treatments may possess systematically different observable and unobservable covariates. Comparing the outcomes between individuals in different treatment groups may then yield a systematically biased estimator of the true causal effect. Matching methods are designed to balance the treatment samples in such a way that differences between the observed covariates of the groups are minimized. This allows the researcher to directly compare the balanced treatment groups for estimating the true causal effect under the assumption that the unobservable covariates of individuals are similar if their observed covariates are similar.
The most widely adopted matching approach in the setting of a binary treatment is matching on the propensity score rosenbaum1983central. Many adjustments and refinements have been proposed in the literature, see for instance austin2008critical, austin2011introduction and imbens2004nonparametric for reviews. The idea is to balance the samples of treatment and control groups by comparing individuals with similar propensity scores, i.e. the probability of receiving the treatment conditional on the observable covariates. Under the assumption that the observed covariates capture all the systematic heterogeneity of individuals, the so-called “unconfoundedness assumption” imbens2000role, imbens2015causal, the researcher can obtain an estimate of the respective causal effect by comparing the outcomes of individuals in treatment and control group with similar propensity scores. The utility of the propensity score derives from dimension reduction: instead of having to match individual covariates in a possibly high-dimensional space, one can match on a one-dimensional score.
However, propensity score matching is no panacea for unbalanced samples. First, as pointed out in imbens2000role, in multivalued treatment settings comparing individuals across groups with similar propensity scores is not equivalent to comparing individuals with similar covariates. This induces bias in the procedure and makes the propensity score “less well-suited” in a setting with more than two treatment arms imbens2000role, lopez2017estimation. Recent approaches address this issue at the cost of interpretability li2019propensity. Second, even in binary treatment settings, propensity score matching has been shown to introduce bias in finite samples king2019propensity. The reason lies in the fact that propensity score matching tries to generate a balanced sample that resembles a randomized trial. This does not imply that all corresponding covariates are matched exactly in finite samples, see king2019propensity for a detailed analysis. Third, in high-dimensional settings, one usually resorts to estimating the propensity score function based on parametric models, which introduces misspecification bias.
The main alternative to propensity score matching is to match the covariates directly. There are a multitude of matching approaches and refinements, see for instance abadie2004implementing, abadie2011bias, hainmueller2012entropy, lopez2017estimation, morucci2020adaptive, rosenbaum1989optimal, rosenbaum2020modern, rosenbaum2007minimum, zubizarreta2015stable and references therein. stuart2010matching provides a concise overview of the fundamental ideas. kallus2020generalized recently introduced a general framework for optimal matching that subsumes many approaches. The most widely used idea is to specify a distance, e.g. a weighted Euclidean distance or some discrepancy measure like the Kullback-Leibler divergence, and proceed via some form of $k$-nearest neighbor matching abadie2004implementing, rubin1973matching, stuart2010matching.
Direct matching comes with computational and statistical challenges. From a computational perspective, matching covariates in high-dimensional settings is more complex than matching on a one-dimensional score, see for instance rosenbaum1989optimal for a classical analysis. From a statistical perspective, most optimal matching approaches, and especially $k$-nearest neighbors matching approaches or variants thereof, are non-smooth approaches due to the strict cut-off at $k$ neighbors. Therefore, they often provide optimal matchings with poor statistical properties. In particular, the obtained estimators for causal effects are often asymptotically biased and need not converge at the parametric rate as shown in abadie2006large in the case of $k$-nearest neighbor approaches. abadie2011bias introduce a correction for these estimators based on linear regression. Analogous rigorous statistical results for similar matching algorithms have been absent in the literature, especially in the multivalued treatment setting as argued in lopez2017estimation.
There is hence a need for a direct matching method that can provide interpretable weights of the optimal matches across several treatment arms while possessing good computational and statistical properties.
This article introduces a natural matching method based on unbalanced optimal transport. It is an alternative to existing matching approaches in terms of how it weights the matches. It is similar in spirit to the nearest-neighbor matching idea, with the important difference that the number of good matches and corresponding weights are chosen endogenously. The idea of optimal transportation is to construct a coupling $\gamma$ between the distributions of the covariate distributions of the respective treatment and control groups. This coupling introduces a natural weighting scheme by considering the respective conditional measures $\gamma(A|x_1)$ for each treatment group, where $A$ is the set of control individuals for the given treated individual with covariate $X_1=x_1$.
This weighting is what sets the proposed method apart from other matching approaches, especially the classical $k$-nearest neighbor approach, where one chooses a fixed number $k$ of neighbors and weights them uniformly, a non-smooth approach for given data abadie2006large. By contrast, the proposed method chooses a natural weighting $\gamma(\cdot |x_1)$ of the matched control individuals that takes into account the cost of matching. One commonly chooses distances on the covariate space as the cost function, for instance the Euclidean distance or weighted counterparts. Then the proposed method naturally re-weights the matched control units in terms of their distance from the treated unit. This provides a natural and smooth weighting procedure and is the main reason why our method enjoys good statistical performance, as we prove in section (ref) below.
Moreover, and equally importantly, the proposed method is based on unbalanced optimal transport chizat2019unbalanced, chizat2018scaling, sejourne2019sinkhorn, liero2018optimal, which enables it to endogenously choose optimal matchings and drop observations in either treatment- or control group that do not have a good match in the other group. This is different from the $k$-nearest neighbor approach, where one matches every observation in one group and only drops observations from the other group that are bad matches. The fact that our proposed method can drop bad potential matches from both groups makes it competitive in terms of bias in finite samples with respect to $k$-nearest neighbor approaches.
The concepts of unbalanced and partial optimal transport were introduced to allow for partial matches without the constraint from classical optimal transport that requires researchers to match every individual in the dataset caffarelli2010free, figalli2010optimal, kitagawa2015multi. The utility of the unbalanced optimal transport approach stems from the fact that it relies on entropy regularization to achieve this and can hence be implemented efficiently in practice via versions of the standard Iterative Proportion Fitting Procedure (IPFP) deming1940least, also called the Sinkhorn algorithm sinkhorn1964relationship. This procedure can straightforwardly be extended to the multiple-treatment setting, which we analyze below. In this sense, the proposed method addresses the challenges outlined in the introduction: it (i) provides optimal matchings in finite samples, (ii) gives interpretable weights based on covariate distance, (iii) is computationally efficient with (iv) good statistical properties, and (v) can even provide a joint optimal matching in multivalued treatment settings. We demonstrate the excellent finite sample properties in a simple simulation exercise in section (ref).
We provide the statistical properties of the optimal matching obtained via the proposed procedure in section (ref). These results are valid when there are different numbers of individuals in different treatment arms, which is crucial in practical settings. From a mathematical perspective, it complements the analytical and statistical results in sejourne2019sinkhorn by showing that the optimal weights of the matches obtained via this procedure converge at the parametric rate for a fixed penalty term. It is built on a central limit theorem for the potentials of the dual problem of the unbalanced optimal transport problem. This result also complements other existing statistical results for the Sinkhorn alogorithm for classical entropy-regularized balanced optimal transport problems from harchaoui2020asymptotics, who require that the number of individuals in both treatment arms are the same. However, they derive an expression for the asymptotic variance of the estimator via Gaussian chaos decompositions, which we do not do.
In concurrent work to ours, dunipace2021optimal provides a matching estimator that is also based on optimal transportation. Their setting is a prediction problem, which is different from ours. They use optimal transport matching to impute the outcomes and the corresponding effects. In contrast, we use optimal transportation to jointly balance and re-weight several samples that only receive one type of treatment. In particular, a main result of this paper is to introduce a balancing of the groups via the conditional measure $\gamma(\cdot|x)$ obtained by the joint coupling $\gamma$. Our proposed method hence redefines the estimation for causal effects in a way that is different from existing approaches.
This section briefly introduces the mathematical background on optimal transport needed for the proposed matching estimator. For further information we refer to santambrogio2015optimal and villani2003topics,villani2008optimal for introductions to general optimal transportation theory and to peyre2019computational for a review of the divergence-regularized optimal transportation and unbalanced optimal transport approach we consider.
For two given marginal distributions $\mu_0$ and $\mu_1$ of characteristics of groups to be matched, defined on spaces $\mathcal{X}$ and $\mathcal{Y}$ respectively, the classical Kantorovich problem villani2003topics consists in finding an optimal transport plan, i.e. a probabilistic matching between the group whose characteristics are distributed as $\mu_0$ and the group whose characteristics are distributed as $\mu_1$. This is a joint coupling $\gamma$ in the space of all couplings $\Pi(\mu_0,\mu_1)$ of these two marginals in such a way that the average cost of transportation of mass from $\mathcal{X}$ to $\mathcal{Y}$ is minimized. Formally, the Kantorovich problem takes the form
where $\Pi(\mu,\nu)$ is the set of all (Borel-) probability measures on $\mathcal{X}\times\mathcal{Y}$ whose marginals are $\mu$ and $\nu$, respectively, and where $c:\mathcal{X}\times\mathcal{Y}\to\mathbb{R}$ is a cost function encoding the difficulty of moving mass from $\mathcal{X}$ to $\mathcal{Y}$. In many cases of interest the cost function is chosen to be a metric on the underlying space, such as the Euclidean distance, i.e. \[c(x,y)\equiv |x-y|^2.\] A distance function such as the squared Euclidean distance is the most natural choice for a cost function, but our method applies to any cost function that is infinitely often differentiable. This can be useful in settings where the researchers want to take into account more information on the cost of comparing individuals than only the distance of their covariates. The Kantorovich problem is a convex relaxation of the Monge problem, which requires that the optimal transportation is achieved via a deterministic mapping $T$, that is, it takes the form
where $T_\# P$ denotes the pushforward measure of $P$ via $T:\mathcal{X}\to\mathcal{Y}$.
This binary setting can straightforwardly be extended to the multimarginal setting by searching for a joint coupling over several marginals $\{\mu_j\}_{j=1,\ldots, J}$ gangbo1998optimal:
where we denote $\mathcal{X} = \mathcal{X}_1,\ldots,\mathcal{X}_J$ and $x\coloneqq (x_1,\ldots, x_J)$. Note that in the multimarginal setting we index the measures $\mu_j$ starting from $1$, in contrast to the binary setting. This is to signify that in a multimarginal setting one does not necessarily have a true control group.
In many cases in the binary setting, for instance when $\mu_0$ possesses a density with respect to Lebesgue measure and the cost function is strictly convex, $\gamma$ is supported on the graph of a function. In this case one obtains a deterministic matching between individuals in $\mu_0$ and individuals in $\mu_1$ that takes the form of a function, i.e. every individual with characteristic $x\in\mathcal{X}$ is mapped to at most one individual with characteristic $y\in\mathcal{Y}$. A special case of this is Brenier's classical theorem on the structure of optimal solutions of the Monge-Kantorovich problem in the case of the squared Euclidean distance as a cost function brenier1991polar. Similar results have been shown in the multimarginal setting, see agueh2011barycenters, carlier2010matching, gangbo1998optimal, pass2015multi.
The classical optimal transport problem is not well-suited for the purpose of matching for causal effects for two reasons. First, its computational and statistical properties are poor. In particular, in practical settings, when the actual distributions $\mu_0$ and $\mu_1$ are unknown and have to be estimated from the observable characteristics of individuals, estimating the optimal matching $\gamma$ suffers from the curse of dimensionality forrow2019statistical, gunsilius2021convergence, hutter2021minimax, manole2021plugin, weed2019sharp. This curse of dimensionality is also computational altschuler2021hardness. Second, classical optimal transport imposes the strong constraint that all elements in the sample have to be matched. From a causal inference perspective, this induces bias because not all individuals will have a close enough match, and matching individuals that are too different induces bias.
This is why we rely on unbalanced optimal transport chizat2018scaling, chizat2019unbalanced, sejourne2019sinkhorn, liero2018optimal, which solves the partial optimal transport problem via divergence regularization. A special case is entropy-regularization of the classical balanced optimal transport problem, which was introduced in galichon2010matching in economics and in cuturi2013sinkhorn in the machine learning literature and is one of the main ways to compute approximations to optimal transportation matchings in practice peyre2019computational.
To obtain a multimarginal matching between the distributions $\mu_1,\ldots,\mu_J$ of characteristics of individuals, we solve the following divergence-regularized optimal transport problem for some $\varepsilon> 0$ and a cost function $c(x)$ which we will assume throughout is smooth, i.e. $c(x)\equiv c(x_1,\ldots,x_J)\in C^\infty\left(\mathcal{X}\right)$ sejourne2019sinkhorn:
Here, $\mathcal{M}^+\left(\mathcal{X}\right)$ denotes the set of all non-negative finite measures on $\mathcal{X}\subset \left(\mathbb{R}^d\right)^J$, $\bigotimes_j\mu_j(x) = \mu_1(x_1)\otimes\cdots\otimes\mu_J(x_J)$ denotes the independence coupling of the marginal measures $\mu_1,\ldots,\mu_J$, \[KL(\gamma||\bigotimes_j\mu_j)\coloneqq
\] denotes the Kullback-Leibler divergence kullback1951information, kullback1959information, $\pi_j\gamma$ denotes the projection onto the $j$-th marginal of the coupling $\gamma$, and $D_\varphi$ denotes the $\phi$-divergence (or Csisz\`ar-divergence) ,csiszar1967information, csiszar1975divergence \[D_\phi\left(\mu||\nu\right)\coloneqq \int_{\mathcal{X}} \phi\left(\frac{\mathop{}\!\mathrm{d}\mu}{\mathop{}\!\mathrm{d}\nu}\right)\mathop{}\!\mathrm{d}\nu + \phi'_\infty \int_{\mathcal{X}}\mathop{}\!\mathrm{d}\mu^\perp,\] where $\phi:[0,\infty)\to[0,\infty]$ is an entropy function, which is assumed to be convex, lower semi-continuous, and $\phi(1)=0$. $\phi'_\infty$ is the recession constant and $\mu^\perp$ is the singular part of $\mu$ with respect to $\nu$ in the Lebesgue decomposition. By assumption, if $D_\phi\left(\mu||\nu\right)<+\infty$ then $\mu\ll \nu$.
We phrased the unbalanced optimal transport problem as a multimarginal problem, but in most practical settings it will be used in the binary setting, which is of the form
The main difference to the classical optimal transport problem are (i) that the optimal coupling $\gamma$ does not need to be a probability measure and (ii) the addition of the $\phi$-divergence terms, which allow for the creation and destruction of mass and do not enforce the constraint from classical optimal transport that the corresponding measures $\mu$ and $\nu$ need to have the same mass. In practice, this means that computing optimal couplings via (ref) or its multimarginal analogue (ref) will provide partial optimal matches where individuals that do not have a close enough match are automatically discarded. This is what provides a small bias in finite samples.
From a mathematical perspective, the unbalanced problem is a generalization of the entropy-regularized optimal transport problem: setting $D_\varphi = +\infty$ if $\pi_j\gamma \neq \mu_j$ almost everywhere and $0$ otherwise and fixing $\varepsilon>0$ reduces our setting to this balanced problem. The regularity properties of this classical balanced entropy regularized optimal transport problem have been analyzed thoroughly, see for instance carlier2020differential, carlier2021linear, genevay2019sample, di2020optimal for recent results and references therein. Note that as $\varepsilon\to0$ the optimal matching $\gamma^*$ solving (ref) approximates the optimal matching obtained via the classical Kantorovich problem (ref) under mild regularity assumptions nutz2021entropic, eckstein2021quantitative.
The utility of the unbalanced optimal transport formulation in our matching setup stems from its dual representation, which follows from the classical Fenchel-Rockafellar theorem chizat2018scaling, sejourne2019sinkhorn:
where $C(\mathcal{X})$ is the space of all continuous functions on $\mathcal{X}$.
The reason for the utility is that the first-order optimality conditions to (ref) can be derived by Moreau's decomposition in Banach spaces combettes2013moreau, which leads to the following generalized Sinkhorn system sejourne2019sinkhorn:
where \[\text{aprox}_{\phi^*}^\varepsilon(p)\coloneqq \operatorname*{\arg\!\min}_{q\in\mathbb{R}}\varepsilon\exp\left(\frac{p-q}{\varepsilon}\right) +\phi^*(q) \qquad\text{for all $p\in\mathbb{R}$}\] is the anisotropic approximation operator combettes2013moreau, sejourne2019sinkhorn, and where \[\phi^*(q)\coloneqq \sup_{p\geq 0} pq-\phi(p)\] is the classical Legendre-Fenchel transform of $\phi(p)$ rockafellar1997convex. We write as a shorthand
The generalized Sinkhorn system (ref) can be solved by a generalization of the classical Sinkhorn algorithm/IPFP deming1940least, sinkhorn1964relationship as in sejourne2019sinkhorn, which is computationally efficient in practice by recursively solving for the sequences $\left\{\varphi_j^m\right\}_{m\in\mathbb{N}\cup\{0\}}$, $j=1\ldots,J$ defined via
for all $j=1,\ldots, J$. Note that in general the optimal potentials $\{\varphi^*_j\}_{j=1}^J$ need not be unique. However, under the assumption that $\phi^*$ is strictly convex, the dual functional (ref) becomes strictly convex in $(\varphi_1,\ldots,\varphi_J)$, which implies a unique solution sejourne2019sinkhorn. We uphold this assumption for our statistical results below (Assumption (ref)).
In the two-marginal setting, sejourne2019sinkhorn showed that this approach converges to the solution of (ref) as $m\to\infty$ for Lipschitz cost functions under some regularity assumptions on the function $\phi^*$ and a normalization of the potentials, which we also uphold (Assumption (ref)). The same argument using compactness also holds in the multimarginal setting on compact supports with Lipschitz cost functions. Addressing related questions, di2020optimal and carlier2021linear proved (linear) convergence of the sequence $\varphi_j^m$ to the optimal $\varphi^*\equiv(\varphi^*_1,\ldots,\varphi^*_J)$ as $m\to\infty$ in the special case of the multimarginal extension of the classical Sinkhorn algorithm for different degrees of smoothness of the cost function, extending convergence results from the two-marginal case for the classical Sinkhorn algorithm chen2016entropic, franklin1989scaling, ruschendorf1995convergence.
This dual approach computes the optimal potentials $\varphi^*$, and only up to the additive constants $\lambda_{\varphi_j}$, which are fixed by normalization. From these optimal potentials, one can generate the optimal matching $\gamma^*$ via
i.e. the optimal potentials are used to generate the optimal matching $\gamma$ as an exponential tilting of the independence coupling of the marginals $\mu_j$. $\gamma^*$ and $\varphi^*_j$ all depend on $\varepsilon$, but we make this dependence implicit for notational convenience.
It is the IPFP that makes solving the entropy regularized problem so convenient. The above procedure (ref) is written out for general measures $\mu_j$. In practice, where the measures $\mu_j$ are unknown and have to be estimated from observations, efficient implementations have been proposed, see peyre2019computational and references therein. Of particular relevance for high-dimensional matching scenarios are mini-batch approaches that break the problem into smaller parts fatras2021unbalanced.
For the multimarginal case introduced above, the optimal probabilistic matching $\gamma^*$ of the participants is joint over all $J$ treatment arms. In particular, for each individual $n_j$ in treatment arm $j$ it provides the optimally weighted matchings in each other treatment arm $i\neq j$. The additional information obtained from the full multimarginal matching comes with an increase in computational cost. That is, one needs to specify (and store in memory) the squared distances between all possible points of all $J$ marginal measures. This becomes increasingly costly with the number of arms $J$, even for the most efficient implementations pham2020unbalanced. In addition, to obtain a tensor joint distribution that can provide a mapping close to a curve in multi-dimensional space, instead of a more uniformly spread-out mass, one needs to specify a penalty term $\varepsilon>0$ that is much smaller than that in the pairwise setting, which not only can make it difficult for the algorithm to converge, but also bring stability issues.
Therefore, we recommend a pairwise comparisons approach to reduce the computational burden in very high-dimensional settings with many treatment arms. Still, we phrase and prove all results in the general multimarginal setting.
This section introduces matching via optimal transportation as a natural framework to minimize bias in weighting approaches for estimating average treatment effects. We rely on the potential outcome notation rubin1974estimating. Let $Y(j)$ be the potential outcome under the causal inference framework for treatment $j\in \{1,...,J\}$, $T$ be the treatment assignment, and random variable $D$ supported on $\{\tau\in \{0,1\}^J: \sum_{j=1}^J \tau_j=1\}$ be the binary vector of the treatment assignment: $T=j$ if and only if $D(j)=1$. $X\in\mathbb{R}^d$ is the vector of observed covariates.
We make the following set of assumptions, the first three standard in the literature abadie2006large, imbens2000role, imbens2004nonparametric.
Parts 1 and 2 of Assumption (ref) are standard. The weak unconfoundedness assumption formalizes the idea of matching on observables: it guarantees that there are no more systematic differences between the unobservable covariates of the groups once the observable characteristics have been accounted for. The weak unconfoundedness assumption was introduced in imbens2000role and is slightly weaker than the strong unconfoundedness assumption introduced in rosenbaum1983central since it is locally constrained for every realization of the treatment assignment. The overlap assumption (Part 2 of Assumption (ref)) guarantees that one can find appropriate matches in each group. If this were not the case, the estimate of the causal effect would be biased. It is a fundamental assumption for any matching approach stuart2010matching. Part 3 of Assumption (ref) states that all treatment groups are independently sampled; a similar assumption is made in classical matching approaches abadie2006large. Part 4 introduces regularity assumptions that simplify the derivations of our statistical results; in particular, the boundedness assumption can be relaxed. In general, these regularity assumptions are weaker than most existing assumptions in the literature, as we do not specify any parametric restrictions on the relationship between $Y$ and $X$ like linearity or convexity and only require additive separability analogously to abadie2006large.
Under part 3 in Assumption (ref), we can denote $X_j := X|D(j)=1$, $X_j\sim \mu_j$ and $\gamma$ as the joint coupling over $\prod_{j=1}^J\mathcal{X}_j$ with the marginal distributions $(\mu_1,...,\mu_J)$ obtained via the proposed matching procedure. The expected counterfactual outcome $\mathbb{E}[Y(t)]$ for treatment $T=t$ can then be written as $\mathbb{E}[Y(t)] = \mathbb{E}_{X}\left[ \mathbb{E}(Y(t)|D(t)=1,X) \right] = \mathbb{E}_X[y(t,X)]$.
The idea of the counterfactual notation is that one can only observe one counterfactual outcome $Y(j)$ in realized samples. This is the “fundamental problem of causal inference” holland1986statistics: for each individual $i=1,\ldots, N$ the researcher only observes one potential outcome $Y_i(j)$ for the given treatment $T=j$ that the individual was assigned to. In order to compute the causal effect of a treatment, the researcher should compare $Y_i(j)$ to $Y_i(k)$, $k\neq j$ for each $i$, which is impossible. The idea of matching on covariates is hence to impute the counterfactual outcomes $Y_i(k)$ for individual $i$ by considering the outcomes of other individuals that received a different treatment $k$. The unconfoundedness assumption implies that individuals with similar observable covariates also possess similar unobservable characteristics, which makes matching on observable covariates a feasible approach.
This leads to the following definition of the potential outcome imputed by our proposed method.
The above definition is designed for the multiple treatment setting. For binary treatments, the expression (ref) is the empirical analogue of \[\mathbb{E}\left[\hat{Y}(1)\right] = \mathbb{E}\left[\mathbb{E}\left[\hat{Y}(1)|T=1,X\right] + \mathbb{E}\left[\hat{Y}(1)|T=0,X\right]\right].\] It is derived by replacing the population measures $\mu_j$ by empirical measures. From this, the average treatment effect or average treatment effects on the treated can be obtained. In the same manner, we can provide estimators of average treatment effects of subgroups like the average treatment effect on the treated, or conditional average treatment effects. We now provide the asymptotic statistical properties of this estimator.
We now provide asymptotic statistical properties of the estimator $\hat{\gamma}_N$ of the optimal matching $\gamma$ obtained by solving (ref) and using the characterization (ref). As mentioned, estimating the weights $\gamma$ via the classical optimal transport problem is known to suffer from a statistical curse of dimensionality. However, we show below that for fixed $\varepsilon>0$, the unbalanced optimal transport problem does not suffer from this curse and obeys a central limit theorem at the parametric rate. This makes it convenient for practical applications, especially in high-dimensional settings.
The results in these section complement the analytical and statistical results in sejourne2019sinkhorn, who inter alia proved that the statistical estimator of the value function of the unbalanced optimal transport problem is consistent at the parametric rate. We show that even the empirical process of the optimizer, i.e. the optimal matching, converges at the parametric rate to a tight limit process with finite (but potentially degenerate) variance if we replace the population measures $\mu_j$ by their empirical counterparts. This result also complements the recently introduced asymptotic linearity result in harchaoui2020asymptotics for the classical balanced entropy-regularized optimal transport problem. There, the authors focus on the two-treatment setting (but their results can be extended to the multimarginal setting) under the assumption that both treatment arms contain the same number of observations. In contrast to our result, they also derive expressions for the asymptotic variance via Gaussian chaos expansions. klatt2020empirical, also in the binary setting, show asymptotic normality of the optimal matching of the classical entropy regularized optimal transport matching and validity of the bootstrap in the setting where all distributions are purely discrete. So far, analogous results for the unbalanced optimal problem have been absent and the results in this section are designed to take a first step to fill this gap.
To set the notation, we define the empirical measures
for $j=1,\ldots, J$, where $\delta_{X_{n_j}}$ denotes the Dirac measure at the observation $X_{n_j}$. This allows for differing sample sizes $N_j$ in each cohort.
Throughout, we assume that we observe independent and identically distributed draws $X_{n_j}$ from the respect distributions $\mu_j$. For notational convenience we will index estimators of functional relationships $f$ that depend on all $j$ elements as $\hat{f}_N$, where it is implicitly understood that $N\equiv \bigcup_{j=1}^J N_j$ is the overall number of data points across all groups. We require the following assumptions for our statistical analysis.
To ensure the regularity and uniqueness of the optimal potentials in the problem, we assume more structure on the function $\phi$ in the definition of the divergence $D_\phi$. The following assumption is satisfied by many functions used for divergences, such as the Kullback-Leibler- and power entropies sejourne2019sinkhorn. For other forms of divergences such as the total variation or even the classical balanced entropy regularized optimal transport problem, uniqueness of the potentials is no longer guaranteed sejourne2019sinkhorn, but similar results can be obtained by different proof methods.
The strong regularity on the cost function is important as the regularity transfers onto the optimal potentials. The smoothness of the potentials is crucial for a parametric rate of convergence of the optimal couplings $\gamma^*_{N}$, as it provides a Donsker class for all dimensions $d$ wellner2013weak. A standard example for a cost function satisfying Assumption (ref) is the squared Euclidean distance function, which is the most natural for a causal matching approach.
sejourne2019sinkhorn have shown that the generalized Sinkhorn algorithm solves the unbalanced optimal transport problem and that the potentials vary continuously with the input measures. We repeat the results for the consistency and convergence of the generalized Sinkhorn algorithm from sejourne2019sinkhorn here in our setting. The proofs are analogous to theirs, as their arguments can be directly extended to the multimarginal setting.
We now provide the main results of this section: a statistical analysis of the asymptotic properties of the optimal matching generated by the generalized Sinkhorn algorithm. To simplify the subsequent statements, we define the operator $T(\varphi)\coloneqq (T_1(\varphi),\ldots, T_J(\varphi))$ with \[T_j(\varphi)(x_j) = \varphi_j(x_j) - \text{aprox}_{\phi^*}^\varepsilon \left(\varepsilon\log\int_{\mathcal{X}_{-j}}\exp\left(\frac{\sum_{i\neq j}\varphi_i(x_i)-c(x)}{\varepsilon}\right)\mathop{}\!\mathrm{d}\bigotimes_{i\neq j}\mu_i(x_{-j})\right),\] which corresponds to the generalized Sinkhorn setting (ref). We take its domain to be $\prod_{j=1}^J C(\mathcal{X}_j)$, which is a Banach space under the supremum norm \[\|\varphi\|_{\infty} \coloneqq \sum_{j=1}^J \|\varphi_j\|_{\infty} = \sum_{j=1}^J\sup_{x_j\in\mathcal{X}_j} \left\lvert \varphi_j(x_j)\right\rvert. \] We also denote $\hat T_N(\varphi) \equiv \left(\hat{T}_{N_1}(\varphi),\ldots,\hat{T}_{N_J}(\varphi)\right)$ with \[\hat{T}_{N_j}(\varphi) = \varphi_j(x_j) - \text{aprox}_{\phi^*}^\varepsilon \left(\varepsilon\log\int_{\mathcal{X}_{-j}}\exp\left(\frac{\sum_{i\neq j}\varphi_i(x_i)-c(x)}{\varepsilon}\right)\mathop{}\!\mathrm{d}\bigotimes_{i\neq j}\hat{\mu}_{N_i}(x_{-j})\right).\]
As a shorthand, and following sejourne2019sinkhorn, we write \[T_j(\varphi)(x_j) = \varphi_j(x_j) - \text{aprox}_{\phi^*}^\varepsilon\left(-\text{Smin}_j^\varepsilon\left(c-\sum_{i\neq j}\varphi_i\right)\right),\] where \[\text{Smin}_j^\varepsilon\left(c-\sum_{i\neq j}\varphi_i\right) \coloneqq -\varepsilon\log\int_{\mathcal{X}_{-j}}\exp\left(\frac{\sum_{i\neq j}\varphi_i-c}{\varepsilon}\right)\mathop{}\!\mathrm{d}\bigotimes_{i\neq j}\mu_i\] is the “soft-min” operator. Based on these definitions we have the following result for the asymptotic distribution of the optimal empirical potentials $\hat{\varphi}^m_{N_j}$ obtained in the $m$-th iteration of the IPFP (ref).
The number of iterations $m$ needs to be chosen such that \[\hat{T}_{N_j}(\hat{\varphi}^m_N) = o_P(N^{-\frac{1}{2}}),\] where $o_P$ is the standard Landau symbol and we always understand statements about probabilities in terms of outer probability wellner2013weak. Theorem (ref) shows that the potential functions obtained via the IPFP possess excellent statistical properties. In particular, the proof of Theorem (ref) also shows that a standard bootstrap procedure, i.e. a resampling procedure of the observable data with replacement, provides correct results in the sense that the bootstrapped estimator $\tilde{\varphi}_N^{m}$ satisfies \[\left(\sqrt{N_1}\left(\tilde{\varphi}^m_{N_1} - \hat{\varphi}^m_{N_1}\right),\ldots, \sqrt{N_J}\left(\tilde{\varphi}^m_{N_J}-\hat{\varphi}^m_{N_J}\right)\right) \] weakly converges to a tight Gaussian limit process conditional on the observed data like the standard empirical process \[\left(\sqrt{N_1}\left(\hat{\varphi}^m_{N_1} - \varphi_1^*\right),\ldots, \sqrt{N_J}\left(\hat{\varphi}^m_{N_J}-\varphi_J^*\right)\right).\] This follows from Theorem 13.4 in kosorok2008introduction. We refer to kosorok2008introduction and chapter 3.6 in wellner2013weak for a detailed analysis of statistical properties of different bootstrap procedures.
Theorem (ref) is the main step in the asymptotic properties of the optimal matching $\gamma^*_\varepsilon$. Recall that the optimal matching $\gamma^*_\varepsilon$ can be derived from the potentials $\varphi^*$ via \[ \gamma^*_\varepsilon(A) = \int_{A}\exp\left(\frac{\sum_{J=1}^J\varphi^*_j(x_j) - c(x)}{\varepsilon}\right)\mathop{}\!\mathrm{d}\bigotimes_j\mu_j(x)\] for any measurable set $A$, i.e. \[K^*(x)\coloneqq \exp\left(\frac{\sum_{J=1}^J\varphi^*_j(x_j) - c(x)}{\varepsilon}\right)\] is the Radon-Nikodym derivative of $\gamma^*_\varepsilon$ with respect to $\bigotimes_j\mu_j$. The asymptotic behavior of the Radon-Nikodym derivative is straightforward to obtain via the delta method wellner2013weak as in the following proposition.
Corollary (ref) shows that the exponential tilt imposed by $K^*(x)$ also converges at the parametric rate to a Gaussian process. Hence, if all measures $\mu_j$ are known, this implies that the optimal coupling converges at the parametric rate to a Gaussian process. Moreover, the Delta method for bootstrapped measures wellner2013weak implies that the standard bootstrap procedure works in this case.
Corollary (ref) provides the statistical properties of the optimal matching $\gamma^*_\varepsilon$ in the case where $\mu_j$ are known. In general, this is not the case and we need to take into account the estimation of the measures $\mu_j$. The following theorem shows that we still have convergence to a tight mean-zero limit process in this case. Also, the classical bootstrap is still valid.
Corollary (ref) and Corollary (ref) show that the optimal balancing possesses converges at the parametric rate to the limit matching. Note that Corollary (ref) is weaker than the central limit result in harchaoui2020asymptotics for the classical balanced entropy regularized optimal transport problem, who provide a Gaussian limit and compute the variance for the classical balanced entropy regularized optimal transport problem. We conjecture that the limit process in Corollary (ref) will be Gaussian in general, but we have not managed to show this. Moreover, the variance of the limit process can be trivial, in which case one would need a higher-order expansion of the variance functional similar to harchaoui2020asymptotics in the classical balanced case. On the other hand, our result allows for different sample sizes in the different treatment arms, while harchaoui2020asymptotics rely on a weighted Monge representation for their results, which requires the same number of observations in each treatment arm.
The tight limit process $G$ is achieved for bounded Donsker classes $\mathcal{F}$, i.e. functional classes that are “small enough” to be estimable with a relatively small number of observations. This is a standard regularity requirement in the theory of empirical processes. Smooth functions, indicator functions and many other functions fall under this category. For a thorough introduction to Donsker classes we refer to wellner2013weak. Corollary (ref) implies the following statistical properties for the induced estimator of the average treatment effect between two groups.
Note that in Part B of Corollary (ref), under the binary treatment setting, where the optimal coupling $\mathop{}\!\mathrm{d}\gamma_\varepsilon(s,x) = K_\varepsilon(s,x) \mathop{}\!\mathrm{d}\mu_0(s)\mathop{}\!\mathrm{d}\mu_1(x)$, the bias between $\mathbb{E}_\epsilon[Y(0)|T=1]$ and the true counterfactual $\mathbb{E}[Y(0)|T=1]$ is
where $Y_i(0)=f_0(X_i)+\eta_i$, $\eta_i$ is an independent noise term as in Assumption (ref) and where we have made the dependence of $K$ on epsilon explicit.
Based on this expression, it is easy to see that if the divergence terms $D_\phi$ vanish quicker than $\varepsilon$ as the latter goes to zero (for instance, when scaling the divergence terms with a parameter $\rho = o(\varepsilon)$), then $K_\varepsilon(s,x)$ converges to a Dirac delta distribution $\delta_x(s)$, which sets the bias to zero. The explanation for this provides further insight into matching approaches: by setting the bias to zero in the limit one only keeps perfect matches. If the distributions of characteristics between groups differ even in the limit, this entails discarding (potentially a significant amount of) mass even in the limiting distribution. This obviously increases the variance of potential estimates, potentially dramatically so.
Classical optimal transport has a balancing constraint, which enforces a match of all the mass in both distributions. This leads to matches that are not perfect and introduces potentially severe bias; on the other hand, the balancing constraint allows the method to use all data points, which in general provides a low variance.
Unbalanced optimal transport is a middle ground between the two: by introducing a penalty term for the balancing constraint from classical optimal transport via the divergence terms, it allows for the balancing constraint to not be satisfied perfectly, which allows for discardment of bad matches. On the other hand, the balancing constraint enforces the use of most of the available data, which in general yields a lower variance. Therefore, from a causal inference perspective, the unbalanced optimal transport approach balances bias and variance: the more the balancing constraint of the classical optimal transport problem is enforced, the more bad matches are created introducing bias into the solution while reducing the variance, and vice versa.
The competitive statistical finite-sample properties of the proposed method, especially compared to the standard $k$-nearest neighbor approach abadie2004implementing, can already be appreciated in simple two-dimensional simulation studies. We propose two simulations, one with a region of overlapping individuals from the two groups, and the other one with the treated sample concentrated within the support of the control sample. Figure (ref) provides an illustration of the covariate distributions for the two testing cases. To compare the method to the existing $k$-nearest neighbor approach (KNN), we compare the performance of unbalanced optimal transport with inverse propensity score weighting (IPW), and KNN when $k=1$ and when $k=3$, in terms of the ATE and ATT estimation accuracy.
Denote the 2-dimensional covariate $\mathbf{X} = (X_1,X_2)$. Throughout the simulation studies, we sample the outcome from the following distributions.
In addition, we sample $n_0=1000$ and $n_1=100$ individuals for the control and treated group respectively.
\noindentCase 1.
For the first simulation, we sample the treated and control groups from two different mixture normal distributions with 2 components. The treated and control groups only have a proportion of covariates in the overlapping area.
\noindentCase 2. For the second simulation, the treated group is concentrated within the region of the control group.
To compare the performance of unbalanced OT with KNN and IPW, 100 bootstrap samples are drawn in each case, the estimated ATE and ATT are computed for each method. For unbalanced OT, the simulation study includes a varying range of the penalty parameter $\epsilon$ from 0.001 to 0.05. In addition, we also compare the estimated standard deviation for ATT by unbalanced OT and KNN methods. Since IPW is a weighting-based method, whereas KNN and OT are imputation-based methods, the standard deviation for IPW is not included. The unadjusted ATE, defined as the observed sample difference between $Y(1)$ and $Y(0)$, is computed as a baseline estimation. Table (ref) gives the summary results of different methods under case 1 and 2. For each case, to evaluate the estimation accuracy, we compute the mean of the absolute difference between the estimates and the true values across all bootstrap samples, denoted as ATE_diff, ATT_diff, and ATT_sd_diff respectively in Table (ref).
Based on the results in Table (ref), we have the following observations:
To further illustrate the weight distribution in the unbalanced OT map, Figure (ref) provides an enlarged covariates scatter plot, where the top 3 pairs that are assigned with the largest weights are marked. Different from KNN(k=1), OT assigns large weights from the treated “1” to both controls “1” and “3”.
Overall, we recommend using unbalanced OT matching for estimating ATT when the overlap assumption is likely to be satisfied at least for one group, like in case 2, where the treated group is concentrated within the controls.
We provide an application to the classical Lalonde dataset. The data are obtained from \url{https://users.nber.org/ rdehejia/data/.nswdata2.html} and are slightly larger and more comprehensive than the version used in dehejia1999causal or imbens2015matching. The dataset contains covariates and earning outcomes for men who were randomly selected into the Nationally Supported Work (NSW) program. The main outcome of interest is earnings in $1978$, re78. We match on 7 covariates: for continuous variables, age ranges from 17 to 55, education is the years of schooling, ranging from 3 to 16, re75 indicates earnings in 1975, ranging from 0 to 37431.7; for the discrete variables, black, \texttt{hispanic} equal to 1 if an individual is Black or Hispanic, respectively,\texttt{married} is the indicator variable for martial status, \texttt{nodegree} equal to 1 if an individual is a high school dropout. We standardize \texttt{age}, \texttt{education} and \texttt{re75}. Among the 722 individuals in this dataset, 297 are in the treated group. Table (ref) provides summary statistics of the covariates stratified by treatment groups.
Table (ref) provides the estimates for ATE and ATT under different methods. The penalty parameter is set to $10^{-3}$ for OT matching. From Table (ref), especially the ATT estimates, KNN(k=1) is almost half the value of OT and KNN(k=3) estimates. There are 2 possible explanations: (i) the noise in the outcome is large, in which case KNN(k=3) and OT results would be more trustworthy, since they are both based on weighted averages of nearby outcomes; (ii) the nearby controls of each treated individual are not very close in distance due to the lack of enough control observations, in which case KNN(k=1) estimate is more trustworthy since it only selects the nearest control as the match. We believe that the first explanation is more likely because the covariate distributions to be matched are mostly binary or discrete, which does not introduce a severe curse of dimensionality, despite the relatively small number of observations.
The estimates of our method and the $3$-nearest neighbor method are qualitatively similar, even though we estimate a slightly higher average treatment effect of $\$760.74$ and a slightly lower average treatment effect on the treated group of $\$828.34$. The sizable difference in outcome between the $3$-nearest and $1$-nearest neighbor estimator displays a challenge for a method that has a discrete penalty term: it is not quite clear at face-value which is the less biased estimate because there are no intermediate values, compared to the case of a continuous penalty term as in our method, which can better illustrate the sensitivity of the outome with respect to the penalty much more clearly. Our additional estimate implies that the less biased estimate among the $k$-nearest neighbor approaches is the $3$-nearest neighbor estimate.
We propose a matching approach for causal inference in observational studies based on unbalanced multimarginal optimal transportation. It can be understood as an alternative to the $k$-nearest neighbor method where the number of neighbors and optimal weights are chosen endogenously. The user instead specifies a penalty term $\varepsilon$. This proposed method possesses desirable properties: the optimal weights for the matches are automatically adjusted to the function that captures the cost of matching, which is usually chosen to be the squared Euclidean distance. Moreover, the method discards bad matches from either group. Together, these two properties lead to demonstrably competitive estimates in terms of bias and variance in finite samples compared to existing methods, such as the classical $k$-nearest neighbor- or propensity score methods. An application to the classical Lalonde dataset lalonde1986evaluating shows that our method provides estimates that are qualitatively similar to those from the $3$-nearest neighbor approach.
In addition, we prove that, for fixed penalty term $\varepsilon$, the potential functions related to the optimal matching converge at the parametric rate to the true potential functions; moreover, the classical bootstrap method can be used to obtain confidence regions for the potential functions of the optimal matching. So, while a positive penalty term $\varepsilon$ induces some asymptotic bias in the estimator, this bias is still competitive with respect to those of existing methods, such as variants of the $k$-nearest neighbor matching approach.
The method also possesses good computational properties with a strong foundation in computational optimal transport theory peyre2019computational. This allows the applied researcher to leverage the latest advances in these areas for matching in high-dimensional spaces. A simple simulation exercise demonstrates the method's competitive finite sample properties in comparison to classical $k$-nearest neighbor approaches. Finally, the proposed method allows researchers to estimate a joint matching over several treatment arms, which could be interesting in answering general non-linear causal questions in this setting.
From a mathematical perspective, we provide a statistical analysis of the potential functions of the optimal matching in the generalized Sinkhorn algorithm for a fixed penalty term $\varepsilon$. This complements existing statistical and mathematical results for the value function of the unbalancd optimal transport problem sejourne2019sinkhorn, as well as results for the classical balanced entropy regularized optimal transport problem carlier2021linear, carlier2020differential, di2020optimal, eckstein2021quantitative, nutz2021entropic, genevay2019sample, klatt2020empirical, pooladian2021entropic, especially that in harchaoui2020asymptotics. A next step in this direction is a finer analysis of the variance functional of the (potential functions of the) optimal matching in the unbalanced optimal transport case, also in the setting where the penalty term is allowed to vanish, similar to the contributions in the classical balanced setting. Finally, another interesting question is to analyze which matches are chosen by the method and which observations are discarded as a function of the penalty term. This could yield insights to optimize the matching approach even further and position unbalanced optimal transportation to potentially be used as a general method for nonparametric regression approaches.