Rates of Estimation of Optimal Transport Maps using Plug-in Estimators via  Barycentric Projections

1 Introduction . Given two random variables X ∼ µ and Y ∼ ν , where µ , ν are probability measures on Rd , d ≥ 1 , the problem of finding a “ nice '' map T0 ( · ) such that T0 ( X ) ∼ ν has numerous applications in machine learning such as domain adaptation and data integration [ 34 , 35 , 38 , 48 , 61 , 112 ] , dimension reduction [ 12 , 66 , 90 ] , generative models [ 60 , 81 , 88 , 110 ] , to name a few . Of particular interest is the case when T0 ( · ) is obtained by minimizing a cost function , a line of work initiated by Gaspard Monge [ 97 ] in 1781 ( see ( 1.1 ) below ) , in which case T0 ( · ) is termed an optimal transport ( OT ) map and has applications in shape matching/transfer problems [ 29 , 47 , 107 , 121 ] , Bayesian statistics [ 46 , 75 , 80 , 108 ] , econometrics [ 15 , 28 , 45 , 50 , 54 ] , nonparametric statistical inference [ 39– 41 , 113 , 114 ] ; also see [ 111 , 128 , 129 ] for book-length treatments on the subject . In this paper , we will focus on the OT map obtained using the standard squared Euclidean cost function , i.e. , T0 : = argmin T : T # µ=ν E‖X − T ( X ) ‖2 , ( 1.1 ) where T # µ = ν means T ( X ) ∼ ν for X ∼ µ . The estimation of T0 has attracted a lot of interest in recent years due to its myriad applications ( as stated above ) and interesting geometrical properties ( see [ 19 , 56 , 91 ] and Definition 1.1 below ) . In practice , the main hurdle in constructing estimators for T0 is that the explicit forms of the measures µ , ν are unknown ; instead only random samples X1 , . . . , Xm ∼ µ and Y1 , . . . , Yn ∼ ν 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . are available . A natural strategy in this scenario is to estimate T0 using T̃m , n , where T̃m , n is computed as in ( 1.1 ) with µ and ν replaced by µ̃m and ν̃n which are empirical approximations of µ and ν based on X1 , . . . , Xm and Y1 , . . . , Yn respectively ( see Definition 1.2 ) . Such estimators are often called plug-in estimators and have been used extensively ; see [ 7 , 30 , 67 , 93 , 94 , 102 , 116 ] . The main goal of this paper is to study the rates of convergence of general plug-in estimators of T0 under a unified framework . We show that when µ̃m and ν̃n are chosen as µ̂m and ν̂n respectively , where µ̂m and ν̂n are the standard empirical distributions supported on m and n atoms , i.e. , µ̂m : = 1 m m∑ i=1 δXi and ν̂n : = 1 n n∑ j=1 δYj , ( 1.2 ) T̃m , n ( appropriately defined using Definition 1.2 ) converges at a rate of m−2/d + n−2/d for d ≥ 4 in the sense of ( 1.8 ) . This rate happens to be minimax optimal under minimal smoothness assumptions ( see [ 72 , Theorem 6 ] ) but suffers from the curse of dimensionality . We next show that , if µ and ν are known to admit sufficiently smooth densities , it is possible to apply kernel or wavelet based smoothing techniques on µ̂m and ν̂n to obtain plug-in estimators that mitigate the aforementioned curse of dimensionality . Our next contribution pertains to the estimation of W 22 ( µ , ν ) ( the squared Wasserstein distance ) , see ( 1.3 ) below , a quantity of independent interest in statistics and machine learning with applications in structured prediction [ 51 , 89 ] , image analysis [ 18 , 59 ] , nonparametric testing [ 16 , 106 ] , generative modeling [ 10 , 96 ] , etc . In this paper , we also obtain rates of convergence for plug-in estimators W 22 ( µ̃m , ν̃n ) of W 2 2 ( µ , ν ) . We show that kernel smoothing µ̂m and ν̂n can be used to obtain plug-in estimators of W 22 ( µ , ν ) that mitigate the curse of dimensionality as opposed to a direct plug-in approach using µ̂m and ν̂n ( as used in [ 30 , Theorem 2 ] ) . This provides an answer to the open question of estimating W 22 ( µ , ν ) when µ , ν admit smooth densities laid out in [ 30 ] . 1.1 Background on optimal transport . In this section , we present some basic concepts and results associated with the OT problem that will play a crucial role in the sequel . Let Pac ( Rd ) denote the set of all Lebesgue absolutely continuous probability measures on Rd andP2 ( Rd ) be the set of probability measures with finite second moments . Then the 2-Wasserstein distance ( squared ) between µ , ν ∈ P2 ( Rd ) is defined as : W 22 ( µ , ν ) : = min π∈Π ( µ , ν ) ∫ ‖x− y‖2 dπ ( x , y ) , ( 1.3 ) where Π ( µ , ν ) is the set of probability measures on Rd×Rd with marginals µ and ν . The optimization problem in ( 1.3 ) is often called the Kantorovich relaxation ( see [ 76 , 77 ] ) of the optimization problem in ( 1.1 ) . The existence of a minimizer in ( 1.3 ) follows from [ 129 , Theorem 4.1 ] . Proposition 1.1 ( Brenier-McCann polar factorization theorem , see [ 91 , 128 ] ) . Suppose µ ∈ Pac ( Rd ) . Then there exists a µ-a.e . ( almost everywhere ) unique function T0 ( · ) : Rd → Rd , which is the gradient of a real-valued d-variate convex function , say ϕ0 ( · ) : Rd → R , such that T0 # µ = ν . Further , the distribution defined as π ( A×B ) = µ ( A ∩ ( T0 ) −1 ( B ) ) for all Borel sets A , B ⊆ Rd is the unique minimizer in ( 1.3 ) provided µ , ν ∈ P2 ( Rd ) . Definition 1.1 ( OT map and potential function ) . The function T0 : Rd → Rd in Proposition 1.1 which satisfies T0 # µ = ν will be called the OT map from µ to ν . A convex function ϕ0 ( · ) in Proposition 1.1 satisfying∇ϕ0 = T0 will be termed an OT potential . The next and final important ingredient is the alternate dual representation of ( 1.3 ) which gives : 1 2 W 22 ( µ , ν ) = 1 2 ∫ ‖x‖2 dµ ( x ) + 1 2 ∫ ‖y‖2 dν ( y ) −min f∈F Sµ , ν ( f ) , where ( 1.4 ) Sµ , ν ( f ) = ∫ f dµ+ ∫ f∗ dν . ( 1.5 ) Here F denotes the space of convex functions on Rd which are also elements of L1 ( µ ) and f∗ ( · ) is the standard Legendre-Fenchel dual defined as : f∗ ( x ) : = sup y∈Rd [ y > x− f ( y ) ] , for x ∈ dom ( f ) . ( 1.6 ) 1.2 Estimating OT map via barycentric projection . Recall the setting from the Introduction . Let µ̃m , ν̃n ∈ P2 ( Rd ) . Here µ̃m , ν̃n need not be absolutely continuous and can be very general . Intuitively , µ̃m and ν̃n can be viewed as some empirical approximation of µ and ν respectively . Example 1.2 ( Simple choices of µ̃m and ν̃n ) . Let X1 , . . . , Xm i.i.d.∼ µ and Y1 , . . . , Yn i.i.d.∼ ν ; in which case a natural choice would be to set µ̃m = µ̂m and ν̃n = ν̂n where µ̂m and ν̂n are the empirical distributions of X1 , . . . , Xm and Y1 , . . . , Yn respectively , as defined in ( 1.2 ) . This is the standard choice adopted in the discrete-discrete Kantorovich relaxation ; see [ 104 , Section 2.3 ] . Another popular choice is µ̃m = µ̂m , ν̃n = ν or µ̃m = µ , ν̃n = ν̂n . This is the semi-discrete Kantorovich problem and is popular when one of the measures is fully specified ; see [ 26 , 55 ] . A natural way to estimate T0 ( · ) , as defined in ( 1.1 ) , would be to approximate it using the OT map from µ̃m to ν̃n . However as µ̃m and ν̃n may not be elements of Pac ( Rd ) , Proposition 1.1 does not apply and an OT map may not exist from µ̃m to ν̃n . Such is the case in Example 1.2 in the discrete-discrete case when m 6= n. To circumvent this issue , we leverage the notion of barycentric projections ( see [ 3 , Definition 5.4.2 ] ) defined below : Definition 1.2 ( Barycentric projection ) . Define the set Γ̃min : = argmin π∈Π ( µ̃m , ν̃n ) ∫ ‖x− y‖2 dπ ( x , y ) . The optimization problem above is the plug-in analog of the optimization problem on the right hand side of ( 1.3 ) . Given any γ ∈ Γ̃min , define the barycentric projection of γ as the conditional mean of y given x under γ , i.e. , T̃m , n ( x ) ≡ T̃ γm , n ( x ) : = ∫ y y dγ ( x , y ) ∫ y dγ ( x , y ) , for x ∈ supp ( µ̃m ) . ( 1.7 ) In general , Γ̃min need not be a singleton which is why we index the barycentric projection T̃ γm , n ( · ) by γ ∈ Γ̃min . Note that T̃ γm , n ( · ) need not be a transport map ; however , if an OT map exists then it must be equal to T̃ γm , n ( · ) ( µ̃m-a.e. ) . Our goal is to obtain stochastic upper bounds for sup γ∈Γ̃min ∫ ∥∥T̃ γm , n ( x ) − T0 ( x ) ∥∥2 dµ̃m ( x ) . ( 1.8 ) In addition , our proof techniques also yield rates of convergence for∣∣W 22 ( µ̃m , ν̃n ) −W 22 ( µ , ν ) ∣∣ . ( 1.9 ) In this paper , we will focus on d ≥ 2 . Due to the canonical ordering of R , the case d = 1 can be handled easily using the classical Hungarian embedding theorem [ 82 ] . 1.3 Contributions . 1 . We provide a new and flexible stability estimate Theorem 2.1 which yields a unified approach to obtaining rates of convergence for general plug-in estimators of the OT map T0 ( · ) . Unlike existing stability estimates , Theorem 2.1 holds for the barycentric projection ( which is the same as the OT map when it exists ) and does not require any smoothness assumptions on µ̃m , ν̃n or T̃ γm , n ( · ) ; also see Remark 2.1 for a comparison with the existing literature . 2 . In Sections 2.1 and 2.2 , we use Theorem 2.1 to bound ( 1.8 ) and ( 1.9 ) : • In Section 2.1 , we show that in both the discrete-discrete and semi-discrete Kantorovich relaxation problems ( see Example 1.2 ) , the rate of convergence of ( 1.8 ) is m−2/d + n−2/d for d ≥ 4 when T0 is assumed to be Lipschitz ( see Theorem 2.2 ) , which is the minimax rate ( see [ 72 , Theorem 6 ] ) . To the best of our knowledge , rates of convergence for these natural estimators weren ’ t previously established in the literature . • In Section 2.2 and Appendix A , we show that the curse of dimensionality in the above rates can be mitigated provided µ and ν admit ( uniform ) Sobolev smooth densities ( see Section 2.2 ) or Besov smooth densities ( see Appendix A ) . In Section 2.2 , our plug-in estimator is obtained by choosing µ̃m ( and ν̃n ) as the convolution of µ̂m ( and ν̂n ) and a smooth kernel with an appropriate bandwidth . Under this choice , the rate of convergence in ( 1.8 ) is m− ( s+2 d ∧ 1 2 ) + n− ( s+2 d ∧ 1 2 ) , where s denotes the degree of Sobolev smoothness ( see Theorem 2.5 ) . Clearly , if 2 ( s + 2 ) ≥ d , the rate of convergence becomes dimension-free and mitigates the curse of dimensionality . We also show the same rates of convergence mentioned above hold for ( 1.9 ) ( see e.g. , Proposition 2.6 ) which makes a strong case in favor of incorporating smoothness in the construction of plug-in estimators as was conjectured in [ 30 ] . In Appendix A , our plug-in estimator is obtained using natural wavelet based density estimators . The rate of convergence in ( 1.8 ) turns out to be n− 1+s d+2s where s denotes the degree of Besov smoothness ( see Theorem A.1 ) . Note that by choosing s large enough , the exponent in the rate can be made arbitrarily close to 1/2 , thereby reducing the curse of dimensionality . 3 . In Section 2.3 , we use a discretization technique from [ 131 ] to construct discrete approximations to the smoothed µ̃m and ν̃n from the previous paragraph that in turn yield computable plug-in estimators for T0 ( provided one can sample from µ̃m and ν̃n ) that also achieve the same statistical guarantees as the smoothed plug-in estimator from Section 2.2 ( see Theorem 2.7 ) . However the number of atoms required in the discretizations and correspondingly the computational complexity increases with the degree of smoothness ; this highlights a statistical and computational trade-off . 4 . We provide implications of our results in popular applications of OT such as estimating the barycenter of two multivariate probability distributions ( see Theorem B.1 in Appendix B.1 ) and in nonparametric independence testing ( see Theorem B.3 in Appendix B.2 ) .

Summary: This paper studies the problem of estimating an optimal transport map between two probability distributions on Euclidean space, given IID samples from each of the distributions. The main result (Theorem 2.1) gives a general upper bound, in terms of the dual representation of the 2-Wasserstein distance. This is used to give upper bounds, in probability, on the rate of decay of the estimation risk both for the plugging in empirical distributions in the absence of smoothness assumptions and for kernel estimates under smoothness assumptions. Finally, implications for two applications (Wasserstein barycenter estimation and nonparametric independence testing) are discussed.