Second-Order Models for Optimal Transport and Cubic Splines on the Wasserstein Space

Abstract

On the space of probability densities, we extend the Wasserstein geodesics to the case of higher-order interpolation such as cubic spline interpolation. After presenting the natural extension of cubic splines to the Wasserstein space, we propose a simpler approach based on the relaxation of the variational problem on the path space. We explore two different numerical approaches, one based on multimarginal optimal transport and entropic regularization and the other based on semi-discrete optimal transport.

Appendices

Proof of Theorem 1

The proof is a rewriting of the proof of [27, Theorem 1.33] when the initial and final spaces do not have the same dimension. In particular we prove that transport plans concentrated on a graph of a map $T : {\mathbb R}^d \rightarrow {\mathbb R}^p$ are dense into transport plans in ${\mathbb R}^d \times {\mathbb R}^p$ and deduce, taking $p= (n-1)d$, that for any continuous cost the multimarginal Kantorovich problem is the relaxation of the multimarginal Monge problem.

Theorem 2

Let $M={\mathbb R}^d$ and $c:M^n \rightarrow {\mathbb R}$ be a continuous cost function. Let $(\rho _i)_{i\in {1,\ldots ,n}}$ be n probability measures on M. We define the Monge Problem $(M_c)$ as

$$\begin{aligned} (M_c)= \inf { \int _{M} c \left( x,T_2(x),\ldots , T_n(x)\right) \rho _1,} \end{aligned}$$

over the set of map $\left. \rho _i, \, , i =2,\ldots ,n\right\}$. The Kantorovich problem $(K_c)$ is defined by

$$\begin{aligned} (K_c)= \inf { \int _{M^n} c \left( x_1,\ldots ,x_n \right) \pi \left( x_1,\ldots ,x_n \right) ,} \end{aligned}$$

over the set of plan , where $p_i$ is the projection of the $i^{\text {th}}$ factor. Then, if all $(\rho _i)_{i\in {1,\ldots ,n}}$ have compact support and $\rho _1$ is atomless, there holds $(M_c)=(K_c)$.

In order to prove Theorem 2 we first remark that [27, Corollary 1.29 and Theorem 1.32] have their multimarginal counterpart.

Lemma 2

Let $\mu \in \mathcal {P}({\mathbb R}^d)$ be atomless measure and $\nu \in \mathcal {P}({\mathbb R}^p)$, then there exists a transport map $T: {\mathbb R}^d \rightarrow {\mathbb R}^p$ such that $T_* \mu =\nu$.

Proof of Lemma 2

Let $\sigma _d : {\mathbb R}^d \rightarrow {\mathbb R}$ (resp $\sigma _p : {\mathbb R}^p \rightarrow {\mathbb R}$) be an injective Borel map with Borel inverse (see [27, Lemma 1.28] for instance for a very simple proof of existence in this case). Since $\mu$ is atomless $({\sigma _d})_*\mu$ is also atomless. Let $t: {\mathbb R}\rightarrow {\mathbb R}$ be the optimal transport map from $({\sigma _d})_*\mu$ to $({\sigma _{p}})_*\nu$ for the quadratic cost. $t_*\left( ({\sigma _d})_*\mu \right) = \left( {\sigma _{p}}\right) _*\nu$. Thus $T= \sigma _{p}^{-1} \circ t \circ \sigma _d$ is a map pushing forward $\mu$ to $\nu$. $\square$

Theorem 3

With the notation of Theorem 2, if the support of all $\rho _i$ are included in a compact domain then the set of plans $\varPi _T$ induced by a transport is dense, for the weak topology, in the set of plans $\varPi$ whenever $\rho _1$ is atomless.

Remark 7

Theorem 3 is in fact very general, one can consider MN to be only Polish spaces for instance. Then there exists invertible Borel maps from M (resp N) to [0, 1]. This is enough to obtain Lemma 2. Then one just need to consider a uniformly small partition of $\varOmega$ to prove the density Theorem 3.

Proof of Theorem 3

Again the proof is based on [27, Theorem 1.32]. In particular the strategy of the proof is to approach a transport plan by transport maps defined on small sets on which the measure is preserved.

We consider a compact domain $\varOmega = \varOmega _d \times \varOmega _p \in ({\mathbb R}^d \times {\mathbb R}^p)$ and $\pi \in \mathcal {P}(\varOmega _d \times \varOmega _p)$ such that $(p_{{\mathbb R}^d})_*(\pi )=\mu$ is atomless. For any m set a partition of $\varOmega _p$ (resp $\varOmega _q$) into (disjoint) sets $K_{i,m}$ (resp $L_{j,m}$) with diameter smaller than 1 / 2m. Then $C_{i,j,m} = K_{i,m}\times L_{j,m}$ is a partition of $\varOmega$ into sets with diameter smaller than 1 / m. Let $\pi _{i,m}$ be the restriction of $\pi$ on $K_{i,m}\times \varOmega _p$ and $\mu _{i,m} = (p_{{\mathbb R}^d})_*(\pi _{i,m})$ and $\nu _{i,m} = (p_{{\mathbb R}^d})_*(\pi _{i,m})$. Since $\mu$ is atomless $\mu _{i,m}=\mu _{|K_{i,m}}$ is also atomless and thanks to Lemma 2 there exists $t_{i,m}$ such that $(t_{i,m})_* \mu _{i,m}=\nu _{i,m}$. By definition

$$\begin{aligned} \pi [C_{i,j,m}]= & {} \pi _{i,m}[C_{i,j,m}]= \mu _{i,m}[K_{i,j}]\nu _{i,m}[L_{j,m}] \nonumber \\= & {} ({\mathrm{Id}},t_{i,m})_*(\mu _{i,m})([C_{i,j,m}])=({\mathrm{Id}},t_{m})_*(\mu )[C_{i,j,m}], \end{aligned}$$

(A.1)

where $t_m$ is define on $\varOmega$ by $t_{|K_{i,m}}=t_{i,m}$. In particular $(t_m)_*(\mu )=\nu$. Equation (A.1) and the definition of the partition sets $C_{i,j,m}$ implies that $(\mathrm {Id},t_{m})_*(\mu )$ weakly converges toward $\pi$ as $m +\infty$ (they give same masses to any set of the partition). See [Theorem 1.31]santambrogio2015optimal for instance. To finish the proof let us remark that we can set $p=d(n-1)$ then $\mu =\rho _1$ is atomless and $t_m:$${\mathbb R}^d \rightarrow {\mathbb R}^{d(n-1)}$ defines $(t_{2,n},\dots ,t_{n,m})$. $\square$

Proof of Theorem 2

The continuity of the cost c and the density Theorem 3 implies that $(K_c)\le (M_c)$. Since the converse is always true we have $(M_c)= (K_c)$.

Remark 8

Theorem 1 is a consequence of Theorem A since both the Monge and the Kantorovich (Definitions 1 and 2) problems reduces on $M^n$ with the spline cost which is continuous (see Corollaries 2 and 3).

Entropic Regularisation and Sinkhorn

1.1 Entropic Regularization and Sinkhorn Algorithm

The linear programming problems (5.7–5.10) is extremely costly to solve numerically and a natural strategy, which has received a lot of attention recently following the pioneering works of [9, 10] is to approximate these problems by strictly convex ones by adding an entropic penalization. It has been used with good results on a number of multimarginal optimal transport problems [1,2,3]. Here is a rapid and simplified description, see the references above for more details.

The regularized problem is

$$\begin{aligned} \min _{ T^\epsilon } \sum _{a,b} \{ C_{a,b} \, T^\epsilon _{a,b} + \epsilon \, T^\epsilon _{a,b} \, \log (T^\epsilon _{a,b}) \} . \end{aligned}$$

(B.1)

It is strictly convex. Denoting $u^k_{\alpha _{j_k}, \beta _{j_k}}$ the Lagrange multipliers of the k constraints (5.10), we obtain the optimality conditions:

$$\begin{aligned} T^\epsilon _{a,b} = K_{a,b} \, \varPi _{k =1}^N U^k_{j_k} , \end{aligned}$$

(B.2)

where

$$\begin{aligned} U^k_{j_k} = e^{ \frac{1}{\epsilon } u^k_{\alpha _{j_k}, \beta _{j_k}} } \quad K_{a,b} = e^{ - \frac{1}{\epsilon } C_{a,b} }. \end{aligned}$$

Equation (B.2) characterize the optimal tensor as a scaling of the Kernel K depending on the dual unknown $U^k$. Inserting this factorization into the constrains (5.10) the dual problem takes the form of the set of equations ( $\forall k \in [1,n]$)

$$\begin{aligned} U^k_{j_k} = \rho _{j_k} ( x_{\alpha _{j_k}, \beta _{j_k}} ) ( \sum _{a \setminus \{\alpha _{j_k} \} , \, b \setminus \{\beta _{j_k} \} } K_{a,b} \, \varPi _{ k^{\prime } \in \{1,\dots ,n\} \setminus k } \, U^{k^{\prime }}_{j_{k^{\prime }}} )^{-1} . \end{aligned}$$

(B.3)

Sinkhorn algorithm simply amounts to perform a Gauss–Seidel type iterative resolution of the system (B.3) and therefore consists in computing the sums on the right-hand side and then perform the (grid) point wise division.

1.2 Implementation

In dimension 2, each unknown $U_k$ has dimension $N_x^2$, the cost of one full Gauss Seidel cycle, i.e. on Sinkhorn iteration on all unknowns, will therefore be $n \times N_x^2 \times$ the cost to compute the tensor matrix products in the denominator of (B.3). Remember that n is the number of time steps with constraints and N the total number of time steps. The given tensor Kernel $K_{a,b}$ is a priori a large $N \times N_x \times N_x$ tensor with indices ${a,b} = {\alpha _1,\dots ,\alpha _N, \beta _1,\dots ,\beta _N}$. It can, however. advantageously be tensorized both along dimensions and also margins. First, using (5.4–5.8) we see that the Kernel is the product of smaller tensors

$$\begin{aligned} K_{a,b} = \varPi _{i=1,N-1} K^0_{i-1,i,i+1} , \,\, \text{ with } K^0_{i-1,i,i+1} := e^{ - \frac{1}{\epsilon \,\mathrm{d}\tau ^3 } \Vert x_{\alpha _{i+1},\beta _{i+1} }+ x_{\alpha _{i-1},\beta _{i-1} } -2 \, x_{\alpha _{i},\beta _{i} } \Vert ^2 } . \end{aligned}$$

Moreover as we chose to work on a cartesian grid at all time steps, $K^0$ tensorize again into

$$\begin{aligned} K^0_{i-1,i,i+1} = K^\alpha _{i-1,i,i+1} \, K^\beta _{i-1,i,i+1}\,\, \text{ with } K^\alpha _{i-1,i,i+1} := e^{ - \frac{h^2}{\epsilon \,\mathrm{d}\tau ^3 } \Vert \alpha _{i+1} + \alpha _{i-1} -2 \alpha _{i} \Vert ^2 }. \end{aligned}$$

Finally our large kernel $K_{a,b}$ can be represented a the product of $2\,(N-2)$ identical tensors of size $N_x \times N_x \times N_x$. Assuming a cubic cost $n^3$ for the multiplication of two $(n \times n)$ matrix, we see our algorithm is of order $O(N \, N_x^4)$ in dimension 2.