T ANGENTIAL W ASSERSTEIN P ROJECTIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
We develop a notion of projections between sets of probability measures using the geometric properties of the 2 -Wasserstein space.
In contrast to existing methods, it is designed for multivariate probability measures that need not be regular, is computationally efﬁcient to implement via a linear regression, and provides a unique solution in general.
The idea is to work on tangent cones of the Wasserstein space using generalized geodesics.
Its structure and computational properties make the method applicable in a variety of settings where probability measures need not be regular, from causal inference to the analysis of object data.
An application to estimating causal effects yields a generalization of the synthetic controls method for systems with general heterogeneity described via multivariate probability measures, something that has been out of reach of existing approaches.
I NTRODUCTION
The concept of projections, that is, approximating a target quantity of interest by an optimally weighted combination of other quantities, is of fundamental relevance in learning theory and statistics.
Projections are generally deﬁned between random variables in appropriately deﬁned linear spaces (e.g. van der Vaart, 2000, chapter 11).
In modern statistics and machine learning applications, the objects of interest are often probability measures themselves.
Examples range from object- and functional data (e.g. Marron & Alonso, 2014) to causal inference with individual heterogeneity (e.g.
Athey & Imbens, 2015).
A notion of projection between sets of probability measures should be applicable between any set of general probability measures, replicate geometric properties of the target measure, and possess good computational and statistical properties. We introduce such a notion of projection between sets of general probability measures supported on Euclidean spaces.
It provides a unique solution to the projection problem under mild conditions.
To achieve this, we work in the 2 -Wasserstein space, that is, the set of all probability measures with ﬁnite second moments equipped with the 2 -Wasserstein distance.
Importantly, we focus on the multivariate setting, i.e. we consider the Wasserstein space over some
Euclidean space R d , denoted by W 2 , where the dimension d can be high.
The multivariate setting poses challenges from a mathematical, computational, and statistical perspective.
In particular, Wis a positively curved metric space for d > 1 (e.g. Ambrosio et al., 2008, Kloeckner, 2010).
Moreover, the 2 -Wasserstein distance between two probability measures is deﬁned as the value function of the Monge-Kantorovich optimal transportation problem (Villani, 2003, chapter 2), which does not have a closed-form solution in multivariate settings.
This is coupled with a well-known statistical curse of dimensionality for general measures (Ajtai et al., 1984, Dudley, 1969, Fournier & Guillin, 2015, Talagrand, 1992; 1994, Weed & Bach, 2019).
1. E XISTING APPROACHES
These challenges have impeded the development of a method of projections between general and potentially high-dimensional probability measures.
A focus so far has been on the univariate and low-dimensional setting.
In particular, Chen et al.
, Ghodrati & Panaretos (2022), and Pegoraro & Beraha (2021) introduced frameworks for distribution-on-distribution regressions in the univariate setting for object data.
Bigot et al.
, Cazelles et al.
(2017) developed principal component analyses on the space of univariate probability measures using geodesics on the Wasserstein space.
The most closely related works to ours are Bonneel et al.
(2016), Mérigot et al.
(2020), and Werenski et al.
The ﬁrst develops a regression approach in barycentric coordinates with applications in computer graphics as well as color and shape transport problems.
Their method is deﬁned directly on W 2 and requires solving a computationally costly bilevel optimization problem, which does not necessarily yield global solutions.
The second introduces a linearization of the 2 -Wasserstein space by lifting it to a L 2 -space anchored at measure that is absolutely continuous with respect to Lebesgue measure.
This approach relies on the existence of optimal transport maps between this absolutely continuous “anchor” distribution and other distributions and hence only deﬁnes tangent spaces at absolutely continuous measures.
The third works on a tangential structure based on “Karcher means” (Karcher, 2014, Zemel & Panaretos, 2019), which is more restrictive still.
This implies that their method requires all involved measures to be absolutely continuous measures with densities that are bounded away from zero, with the target measure lying in the convex hull of the control measures.
1. O UR CONTRIBUTION
In contrast to the existing approaches, our method is applicable for general probability measures, allows for the target measure to be outside the generalized geodesic convex hull of the control measures, can be implemented by a standard constrained linear regression, and provides a global—and in many cases unique—solution.
The proposed method transforms the projection problem on the positively curved Wasserstein space into a linear optimization problem in the geometric tangent cone, which can be implemented via a linear regression.
This problem takes the form of a deformable template (Boissard et al., 2015, Yuille, 1991), which connects our approach to this literature.
Our method can be implemented in three steps: (i) obtain the general tangent cone structure at the target measure, (ii) construct a tangent space from the tangent cone via barycentric projections if it does not exist, and (iii) perform a linear regression to carry out the projection in the tangent space.
This implementation of the projection approach via linear regression is computationally efﬁcient, in particular compared to the existing methods in Bonneel et al.
and Werenski et al.
The challenging part of the implementation is lifting the problem to the tangential structure: this requires computing the corresponding optimal transport plans between the target and each measure used in the projection.
Many methods have been developed for this, see for instance Benamou &
Brenier (2000), Jacobs & Léger (2020), Makkuva et al.
(2020), Peyré & Cuturi (2019), Ruthotto et al.
(2020) and references therein.
Other alternatives compute approximations of the optimal transport plans via regularized optimal transport problems (Peyré & Cuturi, 2019), such as entropy regularized optimal transport (Galichon & Salanié, 2010, Cuturi, 2013).
The proposed projection approach is compatible with any such method, therefore its complexity scales with that of estimating optimal transport plans.
We provide results for the statistical consistency when estimating the measures via their empirical counterparts in practice.
To demonstrate the efﬁciency and utility of the proposed method, we apply our method in different settings and compare it to existing benchmarks such as Werenski et al. Furthermore, we extend the classical synthetic control estimator (Abadie & Gardeazabal, 2003, Abadie et al., 2010) to settings with observed individual heterogeneity in multivariate outcomes.
The synthetic controls estimator is a projection approach, where one tries to predict an aggregate outcome of a treated unit by an optimal convex combination of control units and to use the weights of this optimal combination to construct the counterfactual state of the treated unit had it not received treatment.
The novelty of our application is that it lets us perform the synthetic control method on the joint distribution of several outcomes, which complements the recently introduced method in Gunsilius (2022) designed for univariate outcomes. The possibility to project entire probability measures allows us to disentangle treatment heterogeneity at the treatment unit level.
The possibility of working with general probability measures is key in this setting, as many outcomes of interest are not regular. We illustrate this by applying our method to estimate the effects of a Medicaid expansion policy in Montana, where we
M ETHODOLOGY 2. T HE 2 -W ASSERSTEIN SPACE W 2 ( R d )
The 2 -Wasserstein Distance For probability measures P X , P Y ∈ P ( R d ) with supports X , Y ⊆ R d , respectively, the 2 -Wasserstein distance W 2 ( P X , P Y ) is deﬁned as
[Equation]
Here, | · | denotes the Euclidean norm on R d and
[Equation]
is the set of all couplings of P X and P Y .
The maps π 1 and π 2 are the projections onto the ﬁrst and second coordinate, respectively, and T # P denotes the pushforward measure of P via T , i.e. for any measurable
A ⊆ Y , T # P ( A ) ≡ P ( T − 1 ( A )) .
An optimal coupling γ ∈ Γ( P X , P Y ) solving the optimal transport problem equation 2.1 is an optimal transport plan .
By Prokhorov’s theorem, a solution always exists in our setting.
When P X is regular, i.e. when it does not give mass to sets of lower Hausdorff dimension in its support, then the optimal transport plan γ solving equation 2. is unique and takes the form γ
= (Id ×∇ ϕ )
# P X , where Id is the identity map on R d and ∇ ϕ
( x ) is the gradient of some convex function.
This result is known as Brenier’s theorem (Brenier, 1991, McCann, 1997, Villani, 2003, Theorem 2.12).
By deﬁnition, all measures that possess a density with respect to Lebesgue measure are regular.
Our main contribution is to allow for general probability measures, where only optimal transport plans but no maps exist.
The 2 -Wasserstein Space The 2 -Wasserstein space W 2 (cid:44) ( P 2 ( R d ) , W 2 ) is the metric space deﬁned on the set P 2 ( R d ) of all probability measures with ﬁnite second moments supported on R d , with the 2 -Wasserstein distance as the metric.
It is a geodesically complete space in the sense that between any two measures P, P (cid:48) ∈ W 2 , one can deﬁne a geodesic P t : [0 , 1] → W 2 via the interpolation P t (cid:44) ( π t ) # γ , where γ is an optimal transport plan and π t : R d × R d
→ R d is deﬁned through π t ( x, y ) (cid:44) (1 − t ) x
+ ty (Ambrosio et al., 2008, McCann, 1997).
Using this, it can be shown that W 2 is a positively curved metric space d > 1 (Ambrosio et al., 2008, Theorem 7.3.2) and ﬂat for d = 1 (Kloeckner, 2010), where curvature is deﬁned in the sense of Aleksandrov (1951).
This difference in the curvature properties is the main reason for why the multivariate setting requires different approaches compared to the established results for measures on the real line.
2. T ANGENT CONE STRUCTURE ON W
We exploit a tangential structure that can be deﬁned for general measures on W 2 (Ambrosio et al., 2008, Otto, 2001, Takatsu & Yokota, 2012).
In particular, it allows us to circumvent solving a bilevel optimization problem as the one considered in Bonneel et al.
(2016), whose statement we have included in the appendix.
The tangential structure relies on the fact that geodesics P t in W 2 are linear in the transport plans ( π t ) # γ .
This implies a geometric tangent cone structure at each measure P ∈ W that can be deﬁned as the closure in P 2 ( R d ) of the set G ( P
[Equation]
(cid:44)
[Equation]
P [Equation] [Equation] d × R d ) : ( π 1 ) # γ = P, ( π 1 , π 1 + επ 2 ) # γ is optimal for some ε > 0 (cid:111) with respect to the local distance
[Equation]
where γ 12 and γ 13 are couplings between P and some other measures P 2 and P 3 , respectively, and Γ 1 ( γ 12 , γ 13 ) is the set of all 3 -couplings
γ 123 such that the projection of γ 123 onto the ﬁrst two elements is γ 12 and the projection onto the ﬁrst and third element is γ
(Ambrosio et al., 2008,
Appendix 12).
We can then deﬁne the exponential map at P with respect to some tangent element γ ∈ G ( P ) by
[Equation]
This tangent cone can be constructed at every P ∈ W , irrespective of its support; in particular, we do not assume that the corresponding measures are regular, i.e., give mass to subsets of R d of lower
Hausdorff dimension.
In the case where P is regular the tangent cone structure reduces to a tangent space (Ambrosio et al., 2008, Theorem 8.5.1).
This tangent space structure has been exploited in Mérigot et al.
and Werenski et al.
, and we include the results for our projection approach in this special case in Appendix A.
2. T ANGENTIAL W ASSERSTEIN PROJECTIONS
Our main contribution is to deﬁne a projection approach between general probability measures, where the target need not be regular.
To deﬁne this notion of projection, we need to ﬁrst deﬁne an appropriate notion of a geodesic convex hull.
The novelty here is that we deﬁne this notion via generalized geodesics (Ambrosio et al., 2008, section 9.2) centered at the target measure P 0 .
For this, we extend the deﬁnition of W P to the multimarginal setting, by deﬁning, for given couplings γ
[Equation]
where Γ 1 ( γ 01 , . .
, γ 0 J ) ⊆
Γ( P 0 , P 1 , . .
, P J ) is the set of all ( J + 1) -couplings γ such that the projection of γ onto the ﬁrst- and j -th element is γ
0 j .
Note that this deﬁnition is similar to the multimarginal deﬁnition of the 2 -Wasserstein barycenter (Agueh & Carlier, 2011, Gangbo & ´ Swi˛ech, 1998), but “centered” at P 0 .
Based on this, we deﬁne the generalized geodesic convex hull of measures { P j } j ∈ (cid:74) J (cid:75) with respect to the measure P 0 as
[Equation]
γ solves W 2 P 0 ; λ ( γ 01 , . . .
γ 0 J ) , γ
j is optimal in Γ( P 0 , P j ) ∀
[Equation]
∈
[Equation]
A direct application of our tangential projection idea would lead us to solving
[Equation]
which would be a computationally prohibitive bilevel optimization problem similar to the one in
Bonneel et al.
[-In such settings, the corresponding optimal transport problems equation 2.1 between the target and the respective control measures P j is only achieved via optimal transport plans γ 
 j , not maps ∇ ϕ j . In contrast-]We therefore rely on barycentric projections to reduce the general cone structure to a regular tangent space which we denote by T P 0 W 2 (Ambrosio et al., 2008).
In this structure the projection problem equation 2.5 is replaced by
[Equation]
[-where γ 
-] j − Id (cid:1) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
L 2 ( P 0 ) , with b γ 0 j ( x 1 ) (cid:44) (cid:90)
R d x 2 d γ 0 j,x 1 ( x 2 ) (2.6) denoting the barycentric projections of optimal transport plans γ
j between P 0 and P j .
Here, γ xdenotes the disintegration of the optimal transport plan γ with respect to P 0 .
This approach is a natural extension of the regular setting to general probability measures for two reasons.
First, if the optimal transport plans γ 0 j are actually induced by some optimal transport map, then b γ 0 j reduces to this optimal transport map; in this case the general tangent cone G ( P 0 ) reduces to the regular tangent cone T P 0 W 2 (Ambrosio et al., 2008, Theorem 12.4.4).
[-∇ γ j , then b γ 
-]Second, by the
T P 0 W W P P
P P P π b 1 − Id π (cid:102) Co P deﬁnition of b γ and disintegrations in conjunction with Jensen’s inequality it holds for all λ ∈ ∆ J that
[Equation]
This implies that for general P 0 we can also deﬁne a convex hull based on barycentric projections, which is of the form
[Equation]
[Equation]
Furthermore, the contraction property equation 2.7 implies that Co P 0 ⊆
(cid:102)
Co P 0 , with equality when all transport plans are achieved via maps ∇ ϕ j .
Using these deﬁnitions, the following deﬁnes our notion of projection for general P 0 and shows that it projects onto (cid:102)
Co P 0 .
Proposition 2.1.
Consider a general target measure P 0 and a set { P j } j ∈ (cid:74) J (cid:75) of general control measures.
Construct the measure (cid:101) P π as
[Equation]
where the optimal weights λ ∗ ∈ ∆ J are obtained by solving equation 2.6 and γ
j are optimal plans transporting P 0 to P j , respectively.
Then for given optimal plans γ 0 j , (cid:101) P π is the unique metric projection of P 0 onto (cid:102)
Co P 0 (cid:16)(cid:8) P j (cid:9) J j =1 (cid:17) .
The optimal plans γ [-
-]0 j transporting P 0 to P j need not be unique if P j lies outside the cut locus of P 0 , i.e., when there is more than one optimal way to transport P 0 onto P j .
However, the projection for ﬁxed γ
j is always unique by virtue of the linear regression.
S TATISTICAL PROPERTIES OF THE WEIGHTS AND PROJECTION
We now provide statistical consistency results for our method when the corresponding measures { P j } j ∈ (cid:74) J (cid:75) are estimated from data.
We consider the case where the measures P j are replaced by their empirical counterparts
[Equation]
for every measurable set A in the Borel σ -algebra on R d , where δ x ( A ) is the Dirac measure and (cid:0)
X 1 j , . . .
, X N j ,j (cid:1) is an independent and identically distributed set of random variables whose distribution is P j .
We explicitly allow for different sample sizes (cid:83)
Jj =0 N j = N for the different measures.
To save on notation we write (cid:98) ϕ
N j ≡ (cid:98) ϕ
j , (cid:98) b
j ≡ (cid:98) b γ 0 j ,N j and (cid:98) γ
j ≡ (cid:98) γ N j ,N 0 in the following.
Proposition 3.1 (Consistency of the optimal weights) .
Let (cid:8) P N j (cid:9) J j =0 be the empirical measures corresponding to the data (cid:0) X 1 j , . . .
, X N j j (cid:1)
J j =0 which are independent and identical draws from
P j , respectively, and are supported on some common latent probability space (Ω , A , P ) .
Assume all P j have ﬁnite second moments.
As N j → ∞ for all j ∈ (cid:74) J (cid:75) , the corresponding optimal weights (cid:98) λ ∗ N = (cid:16)(cid:98) λ ∗ N 1 , . .
, (cid:98) λ ∗ N J (
cid:17) ∈ ∆ J obtained via (cid:98) λ ∗ N (cid:44) arg min λ ∈ ∆ J (cid:13) (cid:13) (cid:13)(cid:13)(cid:13)(cid:13)
J (cid:88)
{+j = λ j (cid:16) (cid:98) b 
 j − Id (cid:17) (cid:13) (cid:13) (cid:13)(cid:13)(cid:13)(cid:13) [Equation]+}
[Equation]
satisfy
[Equation]
for all ε > 0 , where λ ∗ solve equation 2.6.
This consistency result directly implies consistency of the optimal weights in case the optimal transport problems between P N 0 and each P N j are achieved by optimal transport maps ∇ (cid:98) ϕ
N j .
We also have a consistency result for the empirical counterparts (cid:101) P π,N of the optimal projection (cid:101) P π .
Corollary 3.1 (Consistency of the optimal projections) .
In the setting of Proposition 3.1, the estimated projections (cid:101) P π,N converge weakly in probability to the projection (cid:101) P π as N j → ∞ for all j ∈ (cid:74) J (cid:75) .
Proposition 3.1 and Corollary 3.1 hold in all generality and without any assumptions on the corresponding measures P j , except that they possess ﬁnite second moments.
To get stronger results, for instance parametric rates of convergences, one needs to make strong regularity assumptions on the measures P j .
Without these, the rate of convergence of optimal transport maps in terms of expected square loss is as slow as n − 2 /d
(Hütter & Rigollet, 2021).
Under such additional regularity conditions, the results for the asymptotic properties are standard, because the proposed method reduces to a classical semiparametric estimation problem, as the weights λ j are ﬁnite-dimensional.
I LLUSTRATIONS 4. M IXTURES OF G AUSSIANS
We consider mixtures of Gaussian in dimension d = 10 .
We draw from the following Gaussians:
[Equation]
where µ 0 =
. , 10] , µ 1 = . , 50] , µ 2 = , 200] , µ 3 = [ − 50 , − 50 , . .
− 50] and Σ
= Id 10 +0 .
8 Id − 10 , with Id − 10 the 10 × 10 matrix with zeros on the main diagonal and ones on all off-diagonal terms.
We then deﬁne the following mixtures: Y 0 as target, and Y 1 , Y 2 , and Y 3 as controls, where
[Equation]
Each sample is 10000 points.
The estimated optimal weights are λ ∗ =
[0 . 4329 , 0 . 4002 , 0 . 1669] with corresponding projection (cid:101) Y 0 .
With only 3 control units, it is not possible to perfectly replicate the entire target distribution.
Still, in Figure 2, the optimal projection approximates Y 0 reasonably well.
Moreover, the weights are non-sparse in this case, indicating that the target P Y 0 lies inside the geodesic convex hull of the control measures.
In many real-world applications we observe sparse optimal weights; see, for instance, Section 4.3, and our application to synthetic controls in Section 5.
This implies that in these settings the target lies outside the geodesic convex hull of the controls and is projected onto one of the faces.
4. I MAGE E XPERIMENT : MNIST
We compare our results to those from the experiment in Section 4.3 of Werenski et al.
We follow the experimental procedure described therein, taking as experimental data the MNIST dataset of 28 × 28 pixel images of hand-written digits (LeCun, 1998).
We show comparison to the test case with image occlusion.
We treat the normalized matrix as probability measures supported on a 28 × 28 grid.
Figure 3 shows our results.
We are able to more clearly replicate the edges and contours of the target image, compared to both the Euclidean projection and the method described in
Werenski et al.
Moreover, our method manages to replicate the overall shape of the speciﬁc handwritten number closer than the other methods; in particular, it is the only method that correctly replicates the horizontal bar at the bottom of this particular handwritten “ 4 ”.
target image.
To examine the general properties of how our method obtains the optimal weights, we provide an application on replicating a target image of an object using images of the same object taken from different angles. We use the Lego Bricks dataset available from Kaggle , which contains approximately 12,700 images of 16 different Lego bricks in RGBA format. Our method manages to replicate the target block rather well, while only using the information of control units that look sufﬁciently like the target.
In particular, in replication, our method does not use information from any image of the underside of the Lego brick.
In contrast, the Euclidean projection does not provide the correct rotation in the replication, and suffers from the standard blur induced by using a mixture of images. [Figure]  A PPLICATION TO CAUSAL INFERENCE VIA SYNTHETIC CONTROLS
When analyzing the causal effect of treatment on a unit, such as that of public policies or medical interventions, there is often no comparable control unit that can capture the treated unit’s underlying characteristics.
The classical synthetic controls method (Abadie & Gardeazabal, 2003, Abadie et al., 2010) aims to create a suitable control unit by replicating the pre-treatment outcome trends of the treated unit, using some optimally chosen set of control units.
This is achieved by projecting the observed characteristics of the target unit onto the convex hull deﬁned by the characteristics of control units in the pre-treatment periods. The optimal weights obtained by this projection, therefore, describe how much each control unit contributes to the target unit’s counterfactual outcome in the post-treatment period (Abadie, 2021).
We apply our notion of projections to extend the classical synthetic control method to work on joint measures of several outcomes, which allows to disentangle heterogeneous treatment effects and complements the univariate method introduced in Gunsilius (2022).
As demonstration, we study the effect of health insurance coverage following state-level Medicaid expansion in Montana in 2016 . The variables of interest are Medicaid coverage, employment status, log wages, and log hours worked.
For control units, we use the twelve states for which such expansion has never occurred; these are: Alabama, Florida, Georgia, Kansas, Mississippi, North Carolina, South Carolina, South
Dakota, Tennessee, Texas, Wisconsin, Wyoming. Additional information can be found in Appendix
C. We estimate “synthetic Montana”, i.e. Montana had it not adopted Medicaid expansion, by estimating the optimal weights λ ∗ using data from 2010 to 2016, and solving equation 2.6 over the joint distribution of the four outcomes over the time period from 2010 to 2016, which generates measures in d = 28 dimensions.
We note that we estimate one set of optimal weights—speciﬁcally, one for each control state—over the entire time period. We then estimate the counterfactual joint distribution using data from 2017 to 2020, by using the optimal weights λ ∗ and computing the weighted barycenter (Agueh & Carlier, 2011) of the control states using these weights.
Details of sample selection and estimating “synthetic Montana” are described in Appendix C.
The results of the general causal effect of the Medicaid expansion policy in Montana averaged over the years 2017 − 2020 are illustrated in Figure 5.
Consistent with ﬁndings in Courtemanche et al.
(2017), Mazurenko et al.
(2018), we ﬁnd signiﬁcant ﬁrst- and second order effects of Medicaid expansion, which are summarized in the top row and the bottom row of Figure 5, respectively.
“Synthetic Montana” has much lower proportion of individuals insured under Medicaid, suggesting that expanding Medicaid eligibility directly affects the extensive margin of Medicaid enrollment.
The disemployment effect is less pronounced in comparison to the enrollment effect we estimated, but nonetheless positive and nontrivial, consistent with the ﬁndings in, e.g., Peng et al.
(2020), but inconsistent with those in, e.g., Gooptu et al.
We also ﬁnd positive second-order effects, summarized in the bottom row of Figure 5.
Additional details are in Appendix C.
[Figure]
C ONCLUSION
We have developed a projection method between sets of probability measures supported on R d based on the tangent cone structure of the 2-Wasserstein space.
Our method seeks to best approximate some general target measure using some chosen set of control measures.
In particular, it provides a global (and in most cases unique) optimal solution.
Our application to evaluating the ﬁrst- and second-order effects of Medicaid expansion in Montana via an extension of the synthetic controls estimator (Abadie & Gardeazabal, 2003, Abadie et al., 2010) demonstrates the method’s utility in allowing for a method that is applicable for general probability measures.
The method still works without restricting optimal weights to be in the unit simplex, which would allow for extrapolation beyond the convex hull of the control units, providing a notion of tangential regression.
It can also be extended to a continuum of measures, using established consistency results of barycenters (e.g. Le Gouic & Loubes, 2017).
A W ASSERSTEIN BARYCENTERS AND THE SPECIAL CASE OF A REGULAR TARGET MEASURE
The natural approach to deﬁning projections on W 2 is to work on the manifold directly. As mentioned in the main text, this leads to a bilevel optimization problem, based on the notion of barycenters in Wasserstein space (Agueh & Carlier, 2011, Carlier & Ekeland, 2010): [Equation]
With this deﬁnition, and assuming that the barycenter ¯ P ( λ ) is unique for given λ , the bilevel projection problem reads:
[Equation]
where
[Equation]
A version of this approach is used in Bonneel et al.
(2016) to deﬁne a notion of regression between probability measures in low dimensions.
The challenges here are mathematical and computational.
Importantly, the optimal weights λ ∗ need not be unique.
This is not an issue for the applications considered in Bonneel et al.
(2016), like color transport; however, it is important in statistical settings when the weights convey information used in further procedures, like causal inference via synthetic controls, where the optimal weights are used to introduced a counterfactual outcome of a treated unit had it not been treated (Abadie & Gardeazabal, 2003, Abadie et al., 2010, Abadie, 2021).
Moreover, the bi-level optimization structure makes solving the problem prohibitively costly in higher dimensions.
Bonneel et al.
(2016) introduce a gradient descent approach based on an entropy-regularized analogue of W 2 (Cuturi, 2013, Peyré & Cuturi, 2019) that can be implemented in low-dimensional settings.
Other
approaches like Werenski et al.
(2022) introduce a tangential approach, but under strong assumptions on the involved measures: they need to be absolutely continuous with densities bounded away from zero on their support, and in particular the target measure must be known to lie inside the convex hull of the other measures.
A starting point for this is to consider a characterization of the barycenter ¯ P ( λ ) for ﬁxed weights of a set { P j } j ∈ (cid:74) J (cid:75) in regular tangent spaces.
Agueh & Carlier (2011, Equation (3.10)) show that if at least one of the measures is absolutely continuous with respect to Lebesgue measure, then ¯ P ( λ ) can be characterized via
[Equation]
where { ˜ ϕ j } j ∈ (cid:74) J (cid:75) are the optimal transport maps from the barycenter to the respective measure P j , i.e. ( ˜ ϕ j )
# ¯ P ( λ )
=
P j .
Each term of the summand in equation A.2 is an element in T ¯ P ( λ ) W 2 ( R d ) by construction.
More generally, the condition equation A.2 is a sufﬁcient condition for ¯ P ( λ ) to be a “Karcher mean” (Karcher, 2014) in W 2 (Zemel & Panaretos, 2019).
In fact, a “Karcher mean” of a set of measures { P j } j ∈ (cid:74) J (cid:75) is deﬁned as the gradient of the Fréchet functional in W 2 and is characterized through equation A.2 holding ¯ P ( λ ) -almost everywhere.
equation A.2 is a stronger condition because it is assumed to hold at every point in the support of ¯ P ( λ ) , not just almost every point.
Álvarez-Esteban et al.
(2016) use this characterization to introduce a ﬁxed-point approach to compute Wasserstein barycenters, and Werenski et al.
(2022) use this structure to introduce a replication approach for absolutely continuous measures whose densities are bounded away from zero and whose target measure lies inside the convex hull of the control measures.
Related is the recent deﬁnition of weak barycenters in Cazelles et al.
(2021), where the authors replace the optimal transport maps from the classical optimal transport problem by the weak optimal transport problem introduced in Gozlan et al.
Heuristically, this characterization is that of a deformable template .
A measure P is a deformable template if there exists a set of deformations (cid:8)
ψ j (cid:9)
j =1 ,...,J such that ψ j # P = P j , in a way that their weighted average is “as close to the identity” as possible.
In our setting ψ j ≡ ∇ ϕ j − Id (Anderes et al., 2015, Boissard et al., 2015, Yuille, 1991).
In our setting of interest, our tangential projection reduces to
[Equation]
[Table]
[Table]
0 ) , (A.3) where ∇ ϕ j are the optimal transport maps between the target P 0 and the control measures P j , j ∈ (cid:74) J (cid:75) .
In contrast to Werenski et al. (2022) the target measure does not need to lie inside the convex hull of the other measures.
Based on these deﬁnitions we can show that our approach is a projection of the target P 0 onto
Co P 0 (cid:16) { P j } Jj =1 (cid:17) in the case where P 0 is regular.
Proposition A.1.
Consider a regular target measure P 0 and a set { P j } j ∈ (cid:74) J (cid:75) of general control measures.
Construct the measure P π as
[Equation]
where the optimal weights λ ∗ ∈ ∆ J are obtained by solving equation A.3 and ∇ ϕ
j are the optimal maps transporting P 0 to P j , respectively.
Then P π is the unique metric projection of P 0 onto
Co P 0 (cid:16)(cid:8) P j (cid:9) J j =1 (cid:17) .
B P ROOFS
Proof of Proposition A.1.
Deﬁne the following closed and convex subset C ⊆
L 2 ( P 0 ) for ﬁxed optimal transportation maps between P 0 and P j , denoted ∇ ϕ j :
[Equation]
for some λ ∈ ∆
[Equation]
  .
Recall that the transport maps ∇ ϕ j exist since P 0 is regular.
Using C , we can rewrite equation A. as
[Equation]
which by deﬁnition is the metric projection of Id onto C .
Since C is a non-empty closed and convex subset of the Hilbert space L 2 ( P 0 ) , this metric projection exists and is unique (Aliprantis & Border, 1999, Theorem 6.53).
Moreover, if Id ∈ C , then π C
= Id ; otherwise, π C ∈ ∂ C , where ∂ C is the boundary of C (Aliprantis & Border, 1999, Lemma 6.54).
Since P 0 is regular, the exponential map is continuous.
In fact, for every j (cid:54) =
k ,
[Equation]
In other words, the distance between P j and P k in W 2 ( R d ) is smaller than that between corresponding elements ∇ ϕ j , ∇ ϕ k in the tangent space.
This implies continuity of the exponential map.
Furthermore, in this regular setting, the exponential map sends convex sets in T P 0 W 2 to generalized geodesically convex sets in W 2 .
Mechanically, for any two (scaled) elements t ( ∇ ϕ j − Id) and s ( ∇ ϕ k − Id) in T P 0 W 2 , and any ρ ∈ [0 , 1] , [Equation]
[Equation]
where (cid:101) ρ (cid:44) ρt + (1 − ρ ) s .
This is a generalized geodesic connecting P j and P k , via the optimal transport map between them and P
(Ambrosio et al., 2008, section 9.2).
The same argument holds when extending generalized geodesics to generalized barycenters by taking convex combination of more measures than a binary interpolation with respect to ρ .
Mechanically, for any λ ∈ ∆ J and t j > 0 for all j ∈ (cid:74) J (cid:75) ,
[Equation]
where (cid:101)
ρ J (cid:44) (cid:80) Jj =
λ j t j .
This proves the exponential map is generalized geodesically convex.
From above it follows that P π (cid:44) exp P
( π C ) is either in the interior of C , which is the case if Id ∈ C , or on its boundary: since π C ∈ ∂C , exp P
( π C ) ∈ exp P 0 ( ∂C ) .
By continuity of the exponential map it follows that exp P 0 ( ∂C ) =
∂ exp P 0 ( C ) .
Combining all steps above show that P π is a geodesic metric projection of P 0 onto the geodesic convex hull of (cid:8) P j (cid:9) J j =1 .
[Equation]
Proof of Proposition 2.1.
The result follows from the same argument as the proof of Proposition A.1.
Theorem 12.4.4 in Ambrosio et al.
(2008) shows that T P 0 W 2 is the image of the barycentric projection of measures in the general tangent cone:
b γ ( x ) is an optimal transport map if γ is an optimal transport plan.
But the exponential map satisﬁes
[Equation]
This implies that
[Equation]
since the convex combination of elements in the subgradients of convex functions lie in the subgradient of a convex function (provided the subgradient of each convex function is nonempty, which is the case here).
Then the continuity and generalized convexity of the exponential map for elements in the regular tangent space T P 0 W 2 implies the result.
[Equation]
Proof of Proposition 3.1.
We split the proof into two parts.
In the ﬁrst part we prove the convergence in probability of the family of objective functions equation 3.1 to their population counterparts equation 2.6 if the empirical measures P N j converge weakly in probability to the population measures P j .
In the second step we use the fact that (cid:98) λ ∗ is a classical semiparametric estimator (Andrews, 1994, Newey & McFadden, 1994) to derive the convergence of the weights.
[Table]
We hence want to show that
[Table]
where (cid:86) j N j ≡ min { N 0 , . . . , N J } .
We split the result into two parts.
The ﬁrst part shows that
[Table]
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) d P 0 ( x 0 ) .
In the second part we use the L 2 ( P 0 ) convergence of the barycentric projections to prove that the limit exists and coincides with the limit inferior.
For the ﬁrst part, we have
[Equation]
where (cid:98) γ N ( x 0 , x 1 , . .
, x J ) is a measure that solves
[Equation]
d γ : γ[-
-]
[Equation]
{+[Equation] [Equation]   , (cid:98) γ 
+} j are the optimal couplings between P N 0 and (cid:101) P N j (cid:44) (cid:16)(cid:98) b 0 j (cid:17) # P N 0 .
Since all measures are deﬁned on the complete and separable space R d , and by assumption of ﬁnite second moments, i.e.
[Equation]
it holds that each sequence (cid:98) γ
j is tight by Ulam’s theorem (Dudley, 2018, Theorem 7.1.4).
Using the fact that λ ∈ ∆ J and (cid:98) γ N ∈ Γ 1 ( (cid:98) γ 01 , . .
, (cid:98) γ 0 J ) , applying Jensen’s inequality gives us
[Equation]
which implies that (cid:98) γ N is tight.
By Prokhorov’s theorem, there exists a subsequence (cid:98) γ N k that weakly converges to a limit measure γ .
Therefore, by the continuity of the map ( x 0 , x j ) (cid:55)→ (cid:80) j λ j x j − x 0 , it follows from classical convergence results (Ambrosio et al., 2008, Lemma 5.1.12(d))
that
[Table]
=
[Table]
[Table]
[Equation]
[Table]
[Equation]
= [Table] [Table]
, x J ) .
Furthermore, by the same argument via Jensen’s inequality, i.e., (cid:90) [Equation] d
[Equation]
J + (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
[Equation]
[Equation]
(cid:88)
[Equation]
[Table]
[Equation]
[Table] d [Equation] , it follows that the limit γ ∈ Γ 1 ( γ 01 , . . .
, γ 0 J ) .
Now note that by the deﬁnition of disintegration it follows that (Ambrosio et al., 2008, Lemma 5.
[Equation]
where
[Equation]
and
[Equation]
are the disintegrations of γ and γ
j with respect to P 0 , respectively. Therefore, we have [Equation] (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
(cid:88)
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
γ
[Equation]
[Equation]
 
[Equation]

[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
d P
[Equation]
[Equation]
) = (cid:90)
R d (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
J (cid:88) j
=
[Equation]
λ j (cid:90) ( R d
[Equation]
J
[Equation]
j
[Equation]
[Equation]
[Equation]
−
[Equation]
0 (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
d
[Equation]
0 ( x 0 )
= (cid:90) R d (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) [Equation]
(cid:88) [Equation] (cid:90) R
[Equation]
( x [Equation] 0 (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
P 0 ( x 0 ) =
(cid:90)
R d (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) J (cid:88) j = [Equation]
λ
[Equation]
[Equation] ( x [Equation] 0 (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) d P 0 ( x 0 ) , where the third lines follows from Jensen’s inequality and the ﬁfth line from γ x∈ Γ (cid:16) γ 1 | x 0 , . .
. , γ J | x 0 (cid:17) .
This shows the ﬁrst part.
For the second part we use the fact that each barycentric projection (cid:98) b 0 j ( x 1 ) is an optimal transport map between P N 0 and (cid:101) P N j if (cid:98) γ
j is an optimal transport plan between P N 0 and P N j , which follows from Theorem 12.4.4 in Ambrosio et al.
As before, we know that (cid:16)(cid:98) b 0 j (cid:17) # P N 0 is a tight sequence that converges to some (cid:101) P j .
By deﬁnition and the fact that ˆ b
j is the gradient of a convex function between P N 0 and (cid:101) P N j , (cid:98) b
j is the unique optimal transport map between P N 0 and (cid:101) P N j for all N j and all j .
Since the measures P j have ﬁnite second moments by assumption, we have
[Equation]
[Equation]
where the last equality follows from the tightness of (cid:98) γ 0 j , as shown earlier.
Therefore, by standard stability results for optimal transport maps (Segers, 2022, Panaretos & Zemel, 2020), it holds that (cid:98) b
j converges uniformly on every compact subset K ⊆ R d in the support of the limit measure (cid:101) P j , that is
[Equation]
where v j is the optimal transport map between P 0 and (cid:101) P j .
We now show that v j = b 0 j P 0 -almost everywhere.
From the local uniform convergence, we can then derive “strong L 2 -convergence”
(Ambrosio et al., 2008, Deﬁnition 5.4.3) of the potentials:
[Equation]
[Table]
the ﬁrst term converges to zero by Hölder’s inequality and the local uniform convergence of the optimal transport maps from above.
The second term satisﬁes
[Equation]
But since P 0 has ﬁnite second moments, it holds that this term also converges to zero.
Based on this we can show that (cid:98) γ 0 j ≡ (cid:16) Id , (cid:98) b 0 j (cid:17) converge weakly to γ
j ≡ (cid:0) I
d , v j (cid:1) .
Indeed, if γ 0 j is a limit point of the sequence (cid:98) γ
j , it holds that
[Equation]
Disintegrating the left-hand side with respect to P 0 , and applying Jensen’s inequality, gives
[Equation]
[Equation]
that is,
[Equation]
But since v j is an optimal transport map between P 0 and (cid:101) P j by deﬁnition, it holds that
[Equation]
which implies that equality holds and we have that
[Equation]
which implies that b
j = v j P 0 -almost everywhere.
We have hence shown that (cid:16) Id , (cid:98) b 0 j (cid:17) # P Nconverges weakly to (cid:0) Id , b 0 j (cid:1) # P 0 for all j , where the barycentric projection b
j is the optimal transport map between P 0 and (cid:101)
P j (e.g. Villani, 2003, Theorem 2.12.(iii)).
Moreover, we have shown “strong L 2 -convergence” of the barycentric projections in terms of Definition 5.4.3 in Ambrosio et al.
Since this holds for all j , it also holds for their convex combination for ﬁxed weights λ ∈ ∆ J .
Putting everything together, we then have that
[Equation]
Since all observable measures P j are empirical measures, they converge weakly in probability (Varadarajan, 1958), which implies that
[Equation]

[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
(cid:88)
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
2 ( P 0 ) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
ε  
[Equation]
for all ε > 0 .
This shows convergence in probability of the objective function for ﬁxed λ .
Step 2: Convergence of the optimal weights (cid:98) λ ∗ N
The convergence of the optimal weights now follows from standard consistency results in semiparametric estimation.
In particular, the objective functions are all convex for any λ ∈ R J , which implies that they converge uniformly on any compact set (Rockafellar, 1970, Theorem 10.8), so the objective function converges uniformly on ∆ J .
Now a standard consistency result like Theorem 2.1 in Newey & McFadden (1994) then implies that
[Equation]
which is what we wanted to show.
Note that the result can also be shown if we allow the weights λ to be negative, i.e., if we only require that (cid:80)
Jj =
λ j = 1 .
In this case, the fact that the objective functions are convex and coercive implies that an optimal λ ∗ will be achieved at the interior of the extended Euclidean space, from which consistency follows by Theorem 2.7 in Newey & McFadden (1994).
[Equation]
Proof of Corollary 3.1.
We want to show that (cid:16)(cid:80)
Jj =
1 ˆ λ ∗
N j ˆ b
j (cid:17) # P N 0 converges weakly in probability to (cid:16)(cid:80)
Jj =1 λ ∗ j b 0 j (cid:17) # P 0 , where ˆ λ ∗ N (cid:44) (cid:16) ˆ λ ∗ N 1 , . . .
, ˆ λ ∗ N J (cid:17) are the optimal weights obtained in equation 3.1 and equation 2.6, respectively.
The result follows by applying the extended continuous mapping theorem (van der Vaart & Wellner, 2013, Theorem 1.11.1) as follows.
As shown in the proof of Proposition 3.1 we have “strong L 2 -convergence” of the maps (cid:80)
Jj =1 ˆ λ ∗ N j ˆ b
j − Id to (cid:80)
Jj =1 λ ∗ j b
j − Id .
Therefore, by Theorem 5.4.4 (iii) in Ambrosio et al.
(2008), it holds that
[Equation]
for any continuous function such that | f ( x 0 ) | (cid:54) C 1 + C 2 | x 0 − x 0 | 2 for all x 0 in the support of P 0 , where C 1 , C 2 < + ∞ are some constants and x 0 in some element in the support of P
(Ambrosio et al., 2008, equation (5.1.21)).
In particular, this holds for any bounded and continuous function f , which implies that
[Equation]
for any bounded and continuous function, which implies that (cid:16)(cid:80)
[Equation]
(cid:17) # P N 0 converges weakly to (cid:16)(cid:80)
Jj =1 λ ∗ j b 0 j (cid:17) # P 0 if P N j converge weakly to P j , j ∈ (cid:74) J (cid:75) .
Now we apply the extended continuous mapping theorem (van der Vaart & Wellner, 2013, Theorem 1.11.1).
Equip P 2 ( R d ) with any metric (cid:101) d ( · , · ) that metrizes weak convergence.
We deﬁne the maps g : × J j =0 (cid:16) P 2 ( R d ) , (cid:101) d (cid:17)
j → (cid:16) P 2 ( R d ) , (cid:101) d (cid:17) by
[Equation]
and analogously for their empirical counterparts g N .
Note that g and g N are non-random functions if the measures P j and P N j are non-random themselves for all j ∈ (cid:74) J (cid:75) .
Moreover, by deﬁnition, g and g N are continuous maps because (cid:80)
Jj =1 λ ∗ j b
j are gradients of convex functions, which are continuous P 0 -almost everywhere; the same thing holds for their empirical counterparts.
Now from what we have shown above and in Proposition 3.1, it holds that
[Equation]
as P N j converge weakly to P j .
Since { P N j } Jj =1 here instead are the only random elements in × J j =0 (cid:16) P 2 ( R d ) , (cid:101) d (cid:17) j , the extended continuous mapping theorem implies that
[Equation]
which is what we wanted to show.
[Equation]
[-We sample N 
-]
C D ETAILS OF M EDICAID E XPANSION A PPLICATION
We use the ACS data with harmonized variables made available by IPUMS, a uniﬁed source of
Census and survey data collected around the world.
The data is at the household-person-year level.
For our application, we select the household head and the spouse as our unit of analysis.
The continuous outcomes are adjusted using the person-level sample weights available in the data.
We adopt the following sample restriction criteria: we included individuals
• of working age, i.e. between ages 18 and• who has no missing outcomes (for those listed in the main text) • who has no top-coded responses • who are either household heads or their spouses
The sample size breakdown by states are follows:
[Table]
We randomly select N = 1500 observations from each unit for estimating λ ∗ .
In the Python implementation, we face a challenge where if the entries of the target and control data are large enough, equation A.3 becomes too large for CVXPY to compute an optimal solution.
Therefore, we introduce a stabilizing constant to prevent this.
This stabilizing constant is determined by the mean value and dimensions of the target distribution, and the number of controls.
The weights we obtained are sparse and are displayed in Table 2.
[Table]
We check whether the obtained weights are ﬁt for creating synthetic Montana by examining if they well-approximate actual Montana in the pre-treatment period.
As seen in Figures 6 and 7, our projection is similar to the actual data.
[Figure] [Figure]
Once we obtain the optimal weights λ ∗ , we estimate the counterfactual outcomes of interest for the four years after Medicaid expansion in Montana (namely, between 2017 and 2020).
This involves solving equation A.1 with λ ∗ obtained from the pre-intervention period.
Implementation-wise, we computed the free-support barycenter, using the POT package; this does not ﬁx the support of the barycenter a priori , and allows it to be different from those of the control distributions.
We plot the densities and distributions of the counterfactual outcomes in Figure 5 of the main text.
To perform inference on the estimated causal effect, we use a placebo permutation test in analogy to Abadie et al.
(2010), Gunsilius (2022).
The idea is to repeatedly apply the procedure described above to each control unit, pretending in turn each control unit is the treated unit.
Post-intervention, if an actual treatment effect only appears in the treatment unit (Montana, in this application), then the estimated effect for the actual treatment unit should be among the largest.
[Figure]
We plot the 2-Wasserstein distance between the treated, joint distribution of all outcomes and the pre-/post-intervention optimal projection (i.e. equation A.1 with λ ∗ ).
We present two sets of results in Figure 8: in panel (A), the optimal projection is computed using λ ∗ estimated using all years in the pre-intervention period; in panel (B), the λ ∗ used is constructed from taking a simple average of weights estimated in each year of the pre-intervention period.
Our results suggest that the estimated causal effect is valid in the post-intervention period, as we consistently observe the largest difference coming from Montana, especially from 2017-2019.
The effect is less pronounced in 2020, however.
To accompany Figure 8, we also compute p -values, which we denote by and deﬁne as p t (cid:44) r
( d 1 t )
J +1 , where d 1 t is the 2-Wasserstein distance from the optimal projection to actual distribution when the target unit is Montana, r ( d 1 t ) is the rank of d 1 t amongst d jt s at given time t , and J is the number of control units.
Results are described in Table 3.
A smaller p t value indicates larger treatment effect.
We observed that r ( d 1 t ) = 1 for 2018 and 2019, implying a nontrivial effect of the Medicaid expansion in Montana during these years.
The values are p t are comparably higher in 2017 and 2020, which we attribute to the fact that it was the ﬁrst year of the policy implementation, and the
COVID-19 pandemic, respectively.
[Table]