Learning to Transport with Neural Networks

TL;DR

This paper compares various neural network-based methods for learning optimal transport maps between probability distributions, introducing a novel approach using dynamic flows and supervised learning reductions.

Contribution

It presents a new approach that reduces optimal transport to supervised learning and compares it with heuristic and mathematically justified methods.

Findings

The novel dynamic flow-based method performs competitively.

Mathematically justified approaches show stronger theoretical guarantees.

Heuristic methods are faster but less accurate.

Abstract

We compare several approaches to learn an Optimal Map, represented as a neural network, between probability distributions. The approaches fall into two categories: ``Heuristics'' and approaches with a more sound mathematical justification, motivated by the dual of the Kantorovitch problem. Among the algorithms we consider a novel approach involving dynamic flows and reductions of Optimal Transport to supervised learning.

Tables

Table 1. Table 4.1. Performance of “heuristic” flows

model name	${\epsilon }_2$	$\sigma \left({\epsilon }_2\right)$	$t_{\text{min}}$	$T$
covariance	0.44	0.010	1000	5000
discr_1	0.76	0.124	2000	100000
discr_5	0.63	0.183	800	40000
exp	0.31	0.019	9000	30000
sym_discr_1	0.28	0.023	21000	50000
sym_discr_5	0.26	0.039	21000	50000
tp_covariance	0.44	0.012	1000	5000
tp_discr_1	0.46	0.020	7000	50000
tp_discr_5	0.46	0.017	5400	30000
tp_exp	0.17	0.013	27000	30000
tp_sym_discr_1	0.29	0.014	36000	50000
tp_sym_discr_5	0.23	0.020	23000	50000

Table 2. Table 4.2. Performance of “adversarial” models

model name	${\epsilon }_2$	$\sigma \left({\epsilon }_2\right)$	$t_{\text{min}}$	$T$
adv_l0.1_2	0.70	0.010	22000	50000
adv_l1_10	0.49	0.013	4000	50000
adv_l1_2	0.37	0.037	39000	50000
adv_l1_2_clip_0.01	0.37	0.037	39000	50000
adv_l10_2	0.62	0.104	2000	50000
adv_l100_2	0.71	0.188	9000	50000

Table 3. Table 4.3. Performance of models based on Regularized duals

model name	${\epsilon }_2$	$\sigma \left({\epsilon }_2\right)$	$t_{\text{min}}$	$T$
seguy_ent_mean_0.1	0.15	0.021	9200	10000
seguy_ent_sum_0.1	0.17	0.016	9200	10000
seguy_l2_mean_0.1	0.27	0.012	4500	5000
seguy_l2_sum_0.1	0.15	0.019	4500	5000
supervised_dual_0.05	0.21	0.093	13800	30600
supervised_dual_0.1	0.18	0.082	14000	20800
supervised_map_iters_1000_0.05	0.16	0.096	24000	51000
supervised_map_iters_200_0.05	0.18	0.084	8000	50000
supervised_prob	0.29	0.030	3000	51000

Table 4. Table 4.4. Timing statistics for models

model name	secs per step	secs to $t_{\text{min}}$
tp_exp	0.00456	123.04
supervised_map_iters_1000_0.05	0.00303	72.60
seguy_l2_sum_0.1 (dual)	0.00795	35.78
seguy_l2_sum_0.1 (map)	0.01192	59.60
seguy_ent_mean_0.1 (dual)	0.00806	74.18
seguy_ent_mean_0.1 (map)	0.01245	124.48
supervised_dual_0.1 (dual)	0.00482	67.48
supervised_dual_0.1 (map)	0.00607	60.70

Equations

$${\mathcal{P}}_0$$

$$\pi ={\int }_{\mathcal{X}}d\mu (x){\int }_{{\mathrm{pr}}_{\mathcal{Y}}^{-1}(x)}{\pi }_x,$$

$${\mathcal{M}}_0$$

$$\pi \stackrel{\text{heur}}{\to }T$$

$${\mathcal{D}}_0$$

$${\mathcal{P}}_{\epsilon }$$

$${\mathcal{D}}_{\epsilon }$$

$$d\pi (x,y)=\exp \left(\frac{u(x)+v(y)-c(x,y)}{\epsilon }\right)d\mu \otimes \nu (x,y).$$

$${\mathcal{P}}_{\epsilon }^{l_2}$$

$${\mathcal{D}}_{\epsilon }^{l_2}$$

$$d\pi (x,y)=\frac{(u(x)+v(y)-c(x,y){)}_{+}}{2\epsilon }d\mu \otimes \nu (x,y).$$

$$\frac{dw(t)}{dt}=-{\nabla }_{w(t)}\mathcal{L}(T_{w(t)},\mu ,\nu ).$$

$${\int }_{\mathcal{X}}f(T_{w(t)}(x))d\mu (x)={\int }_{\mathcal{Y}}f(y)d\nu (y).$$

$$\mathcal{L}(T_{w(t)\#}\mu ,\nu )=\sum _{k=1}^K{\left|{\int }_{\mathcal{X}}{f}_k(T_{w(t)}(x))d\mu (x)-{\int }_{\mathcal{Y}}{f}_k(y)d\nu (y)\right|}^2.$$

$$\mathcal{L}(\{T_{w(t)}(X_i){\}}_{i=1}^{b_X},\{Y_j{\}}_{j=1}^{b_Y})=\sum _{k=1}^K{\left|\frac{1}{b_X}\sum _{i=1}^{b_X}{f}_k(T_{w(t)}(X_i))-\frac{1}{b_Y}\sum _{j=1}^{b_Y}{f}_k(Y_j)\right|}^2.$$

$$f_k(y)=\exp \left(\frac{-d(y,z_k{)}^2}{{\sigma }^2}\right).$$

$${\mathrm{disc}}_k(\{T_{w(t)}(X_i){\}}_{i=1}^{b_X},\{Y_j{\}}_{j=1}^{b_Y})$$

$${\mathrm{symdisc}}_k(\{T_{w(t)}(X_i){\}}_{i=1}^{b_X},\{Y_j{\}}_{j=1}^{b_Y})$$

$$\mathcal{L}(T_{w(t)},\mu ,\nu )={\int }_{\mathcal{X}}c(x,T_{w(t)}(x))d\mu (x)+\mathcal{L}(T_{w(t)\#}\mu ,\nu ).$$

$$Adv$$

$$\sum _{i=1}^{b_X}|u_{\theta }(X_i)-{\hat{u}}_i|+\sum _{j=1}^{b_Y}|v_{\eta }(Y_j)-{\hat{v}}_j|.$$

$${\epsilon }_2={\int }_{\mathcal{X}}\Vert T_w(x)-T_{\text{opt}}{\Vert }_2^{2}d\mu (x).$$

$${\epsilon }_2=\frac{1}{B}\sum _{i=1}^B\Vert T_w(X_i)-T_{\text{opt}}(X_i){\Vert }^2.$$

Full text

1. 1 Introduction
2. 2 Review of Optimal Transport
3. 3 Learning to Transport
4. 4 Experiments

1. Introduction

Many problems in Machine Learning require comparing probability distributions. Applications of computing distances between probability histograms include bag-of-words for natural language processing [KSKW15], color and shape processing in computer vision [SdGP*+*15], and regularization of generative models [ACB17, PvF*+*18]

While many distances between probability distributions have been studied and used in the literature, the Optimal Transport distance [PC19, San15] has come to hold a prominent place. One main reason for its success is its flexibility as it can be parametrized using a cost chosen depending on the application domain. Another reason is that solving the Optimal Transport problem allows to find a transformation, an Optimal Map, between two probability distributions which minimizes a given loss function.

In particular, while computing a distance allows to say how much two histograms are close to each other, computing a map allows to solve additional problems, for example sampling from intractable distributions using an optimal map to transform a tractable distribution into an intractable one [TT16, SDF*+*18], flow tracing in fluid dynamics and medical diagnostics [OZM*+*16, TT16], and database aggregation between data collected in different laboratories [TT16]. In my work in E-Commerce I have additionally used Optimal Transport to design quasi-experiments [KS12, Net] in situations where adoption of a given product cannot be properly randomized between treatment and control. In many cases using Optimal Transport has allowed me to build comparable treatment and control groups increasing the statistical power to measure the effect of interventions.

Finally, there is increasing work in Optimal Transport in general metrics spaces and with non-smooth costs [LA08, LV09] where the lack of regularity and high-dimensionality makes traditional approaches based on the Monge-Ampere equation unfeasible.

In this work we make a comparison of different approaches to learn optimal transport using Neural Networks. We look at a data-driven formulation to solve the problem exploiting a class of Machine Learning models that have gained prominence in the last two decades with their successes in Natural Language Processing, Computer Vision and Variational Inference.

1.1. Previous work

It his hard to make proper justice to the amount of literature on computational Optimal Transport and its applications to Machine Learning. We pick a few previous works related to this paper and refer the reader to [PC19]. In [TT16] Trigila and Tabak studied algorithms to find the Optimal Map using dynamic flows. Their approach is data-driven and motivated by problems where the probability histograms live in high-dimensional spaces.

In [Cut13] Cuturi showed that Optimal Transport distances, after adding an entropic regularization term, can be computed efficiently using the celebrated Sinkhorn’s algorithm [SK67]. The work [Cut13] is a landmark in making optimal transportation distances attractive to the Machine Learning community. Moreover, in [GCPB16] the authors showed how to scale up the computation of distances using online learning and stochastic gradient descent.

Optimal transport distances have soon attracted the interest of people working on Generative Models as the optimal transport distance can be used as a regularization term. For example in [ACB17] the authors introduced Wasserstein GANs, which are GANs where the adversarial network computes an $l_{1}$-optimal transport distance, and show the superior performance of Wasserstein GANs over traditional GANs. Moreover they also provided a mathematical justification for preferring optimal transport to other regularization terms, a line of research further developed in [AB17]. Another recent work using the Sinkhorn algorithm is the paper [PvF*+*18] introducing the so-called Sinkhorn autoencoders.

The first work we are aware of that uses Neural Network to compute transport maps is [SDF*+*18] where the authors show applications to generative models and domain adaptation.

1.2. Contributions

The contribution of this paper is mainly comparing several approaches to learn the optimal transport map using neural networks. From the existing literature we take two approaches employed in [SDF*+*18], namely solving the dual transport problem with entropic / $l_{2}$-regularization. Novel approaches that we consider are using neural networks in dynamic flows motivated by [TT16], adversarial training for the optimal map motivated by the literature on generative models [ACB17, GPAM*+*14, ZML16], and supervised learning, which is motivated by the assumption that discrete / semi-discrete problems can be used to train statistical models which are capable of generalization.

1.3. Code Links

We try to make this work reproducible providing the code and data in the GitHub repository: https://github.com/salayatana66/learn_to_transport_code. In the following we just refer to this repository as the GitRepo.

Please note that the experiments/code there is not an exhaustive representation of the experiments we tried out. Quite a bit of work went into engineering each experiment and decide which parameters to test. This was broken down into intermediate experiments that resulted in the final scripts. Concretely, in Subsection 3.7 we discuss the importance of proper initialization. We learned this while experimenting on learning linear maps between Gaussians in different dimension, code which we exclude from the GitRepo. Similarly, in Subsection 3.6 we discuss the issue of normalization while learning potentials in a supervised way. Making that work took a bit of effort, also in light of numerical instability of the Sinkhorn iterations for small values of the regularization. The final solution is based on going to log-space, compare [PC19, Subsec. 4.4].

2. Review of Optimal Transport

2.1. Notation

We now recall the formulation of the Optimal Transport Problem. Let $\mathcal{X}$ and $\mathcal{Y}$ denote two metric spaces and let ${\rm Prob}(\mathcal{X}\times\mathcal{Y})$ denote the set of Radon probability measures on $\mathcal{X}\times\mathcal{Y}$, and for each metric space $\mathcal{Z}$ let $C({\mathcal{Z}})$ denote the set of real valued continuous functions defined on $\mathcal{Z}$. Let ${\rm pr}_{\mathcal{\mathcal{X}}}$, ${\rm pr}_{\mathcal{\mathcal{Y}}}$ denote the projections of $\mathcal{X}\times\mathcal{Y}$ on the first and second factor, respectively. In the following we fix a probability measure $\mu$ on $\mathcal{X}$ and another one $\nu$ on $\mathcal{Y}$. In the optimal transport problem we want to move $\mu$ into $\nu$ minimizing some kind of cost; therefore, to each pair of points $(x,y)\in\mathcal{X}\times\mathcal{Y}$ we associate a transport cost $c(x,y)$, and, without loss of generality, we can assume that $c\in C({\mathcal{X}\times\mathcal{Y}})$ and $c$ is non-negative.

2.2. The Kantorovich Formulation

Let ${\rm Plan}(\mu,\nu)$ denote the set of those $\pi\in{\rm Prob}(\mathcal{X}\times\mathcal{Y})$ such that the marginals on $\mathcal{X}$ and $\mathcal{Y}$ give back $\mu$, $\nu$: ${{\rm pr}_{\mathcal{\mathcal{X}}}}_{\#}\pi=\mu$, ${{\rm pr}_{\mathcal{\mathcal{Y}}}}_{\#}\pi=\nu$. In the Kantorovich formulation the optimal transport problem becomes:

[TABLE]

the infimum value of the objective in ($\mathcal{P}_{0}$) will be denoted by ${{\rm W}}_{0}$. We will refer to problem ($\mathcal{P}_{0}$) as the un-regularized primal.

In cases of common interest $\mathcal{X}=\mathcal{Y}={\mathbb{R}}^{d}$ and the cost $c$ is the square of the Euclidean distance; in that case ${{\rm W}}_{0}^{1/2}$ is called $l_{2}$-Wasserstein distance between the probability distributions $\mu$ and $\nu$.

2.3. The Monge formulation

Having found a minimizer $\pi$ for ($\mathcal{P}_{0}$) one can use the Disintegration Theorem [San15, Box 2.2] to represent $\pi$ as:

[TABLE]

where $x\mapsto\pi_{x}$ is a map associating to each $x\in\mathcal{X}$ in the support of $\mu$ a family of probability measures in ${\rm Prob}(\mathcal{X}\times\mathcal{Y})$ supported on the fiber of $x$ under the projection. The probability $\pi_{x}$ can be interpreted as a transport plan for a unit of mass at $x$ towards points of $Y$ in the support of $\nu$. In the Monge formulation of the transport problem one requires that $\pi_{x}$ is concentrated at a single point $T(x)$ of $\mathcal{Y}$; equivalently there is a (measurable )map $T:\mathcal{X}\to\mathcal{Y}$ (note there is need to define $T$ only on the support of $\mathcal{X}$) such that $({{\rm Id}},T)_{\#}\mu=\pi$. One can also directly formulate the Monge Problem: let ${\rm TM}(\mu,\nu)$ denote the set of those measurable maps $T:\mathcal{X}\to\mathcal{Y}$ with $T_{\#}\mu=\nu$; then the problem becomes:

[TABLE]

The convexity and linearity of ($\mathcal{P}_{0}$) implies the existence of a minimizer $\pi$ under mild conditions [San15, Theorem 1.4]. On the other hand in the Monge Problem it is not always the case that ${\rm TM}(\mu,\nu)$ is nonempty or that the infima of the objectives of ($\mathcal{P}_{0}$) and ($\mathcal{M}_{0}$) are the same; we refer the reader to [San15, Secs 1.3–1.5] for more information. A map $T\in{\rm TM}(\mu,\nu)$ is a transport map and an optimal transport map is a solution to ($\mathcal{M}_{0}$).

In this paper we are mainly concerned with finding an approximation to an optimal $T$ and we will deliberately neglect issues about the existence and regularity of such a $T$. We just observe that given a plan $\pi\in{\rm Prob}(\mathcal{X}\times\mathcal{Y})$ and mild assumptions on $c(x,y)$ (e.g. convexity or being a distance) there is a simple heuristic to associate to it a transport map: for each $x$ in the support of $\mu$ one defines $T(x)$ as

[TABLE]

for example in the case of the quadratic Euclidean cost one maps $x$ to the barycenter of $\pi_{x}$.

2.4. The dual problem

The problem ($\mathcal{P}_{0}$) is a convex problem; its dual can be formulated in terms of a pair of potentials $u\in C({\mathcal{X}})$, $v\in C({\mathcal{Y}})$. We let ${\rm Pot}(c)$ denote the set of those pairs $(u,v)\in C({\mathcal{X}})\times C({\mathcal{Y}})$ such that for each $(x,y)$ one has $u(x)+v(y)\leq c(x,y)$. The dual of ($\mathcal{P}_{0}$) is then:

[TABLE]

One might show [San15, Theorems 1.37,1.38] that given an optimal $(u,v)$ for ($\mathcal{D}_{0}$) an optimal $\pi$ for ($\mathcal{P}_{0}$) can be found using the disintegration formula: $\pi_{x}$ can be taken as the uniform probability distribution on those $y$ such that $u(x)+v(y)=c(x,y)$.

2.5. Regularization

The problem ($\mathcal{P}_{0}$) is not strongly convex; adding a regularization term it can be turned into a strongly convex one improving its convergence properties. Let $\mu\otimes\nu\in{\rm Prob}(\mathcal{X}\times\mathcal{Y})$ denote the product of $\mu$ and $\nu$ and for $\pi\in{\rm Plan}(\mu,\nu)$ let $\frac{d\pi}{d\mu\otimes\nu}$ be the Radon-Nikodym of $\pi$ with respect to $\mu\otimes\nu$. If $\pi$ is not absolutely continuous with respect to $\mu\otimes\nu$, in Real Analysis it is common practice to restrict the Radon Nikodym derivative to the part of $\pi$ which is absolutely continuous with respect to $\mu\otimes\nu$. However in Optimal Transport it is a common practice to set the derivative to $+\infty$ in such a case. As we will then consider regularizations which are $+\infty$ in the case of $\pi$ not being absolutely continuous with respect to $\mu\otimes\nu$, optimal plans for ($\mathcal{P}_{\varepsilon}$) will be in general “diffused” or “smoothed” by the regularization. In fact, if the optimal plan $\pi$ for ($\mathcal{P}_{0}$) is induced by an optimal transport map $T$ for ($\mathcal{M}_{0}$), then $\pi$ is generally singular with respect to $\mu\otimes\nu$.

For $\varepsilon>0$ the entropic regularization of ($\mathcal{P}_{0}$) is

[TABLE]

The dual of ($\mathcal{P}_{\varepsilon}$) is the unconstrained problem:

[TABLE]

Given a solution $(u,v)$ for ($\mathcal{D}_{\varepsilon}$) on obtains a solution $\pi$ for ($\mathcal{P}_{\varepsilon}$) setting:

[TABLE]

An alternative regularization is the $l_{2}$-one which gives the problems:

[TABLE]

[TABLE]

and the analogue of (2) becomes

[TABLE]

2.6. Discretization

When $\mu$ and $\nu$ are discrete measures all problems ($\mathcal{P}_{0}$), ($\mathcal{M}_{0}$), ($\mathcal{D}_{0}$) and their regularized versions become discrete optimization problems. For problem ($\mathcal{D}_{\varepsilon}$) there is an efficient computation algorithm, the Sinkhorn algorithm [SK67], with a convergence rate $O(\varepsilon^{-3}\log n)$, $n$ being an upper bound on the number of points at which $\mu$, $\nu$ are supported. The Sinkhorn algorithm is a matrix rescaling procedure that can be efficiently parallelized on GPUs [Cut13] and in the limit $\varepsilon\to 0$ one can guarantee to recover a solution of the unregularized problem. However, in practice too small values of $\varepsilon$ lead to rounding errors and slow convergence; here we have not used values smaller than $0.005$.

3. Learning to Transport

Figure 3.1 summarizes the approaches that we examine to learn a transport map $T$ from $\mu$ to $\nu$ where $T$ is represented by a Neural Network. We can classify the approaches in two categories:

1. (1)

Heuristics: the mathematical justification is not so solid, but the approach is motivated by intuition and might work reasonably well in practice. In the Machine Learning literature it is common to resort to such approaches when problems are hard to treat or correctly formulate (e.g. Variational Inference, Variational AutoEncoders, Generative Adversarial Neural Networks). 2. (2)

Solving a regularized dual: these are based on solving ($\mathcal{D}_{\varepsilon}$) or ($\mathcal{D}^{l_{2}}_{\varepsilon}$). If $\mu$ and $\nu$ are discrete the theory is classical; when $\mu$ or $\nu$ are continuous and the potentials $u$ and $v$ are represented via neural networks we do not know of rigorous convergence guarantees, but methods can work well in practice.

3.1. Representations

What is a good way to represent mathematical objects like $\mu$, $\nu$, $T$, $u$ or $v$? The answer will likely depend on the setting. For example $\mu$ and $\nu$ can be discretized on a fixed mesh and then one needs to keep track of the value of maps and functions only on a fixed set of points. As observed in [TT16] mesh-based approaches do not scale well in the dimension and for Big Data applications [TT16] advocates having access to random samples drawn from $\mu$ and $\nu$. For a discussion of discretization and Wasserstein distances look at [WB17].

From a Machine Learning perspective it is also convenient to store $T$ as a model, hence in a compact form. Besides storage efficiency, one also gains in generalization power. For example assume taht $T$ has been learned on discrete samples; now such a $T$ can be also evaluated on new samples. Here we focus on representing $T$ (or $u$, $v$ or $\pi$ depending on the circumstances) as neural networks. The first reason is that even with just two hidden layers one can represent rich spaces of non linear functions. The second reason is availability of Open Source libraries to train Neural Networks, here we will use PyTorch.

3.2. Flows

A first set of heuristics to find a transport map $T$ considers flows, i.e. dynamically evolving trajectories from points sampled from $\mu$ to points sampled from $\nu$. We loosely follow [TT16] to outline the main idea and, just for the moment, we let $\mathcal{X}$, $\mathcal{Y}$ be subsets of some Euclidean space ${\mathbb{R}}^{d}$ and assume that $\mu$ and $\nu$ are continuous with respect to the Lebesgue measure. In this case ($\mathcal{M}_{0}$) an ($\mathcal{P}_{0}$) will have the same infimum value and there will be an optimal map inducing an optimal plan. In [TT16] the authors try to reach the optimal $T$ using a flow $\{T_{t}\}_{t}$ starting at the identity. Under the aforementioned regularity assumptions they are able to combine the Monge and Kantorovich formulations by updating at each step $T_{t}$ and the potentials $u$ and $v$. This is made possible by linearizing the Monge Ampere equation with different approaches that they investigate. To fix the ideas, a simple linearization approach consists in using a linear perturbation of the identity (equivalently assuming that $u$ and $v$ are updated via second order polynomials), which result in $\{T_{t}\}$ being represented as a composition of linear maps. A drawback of this approach is that to compute $T$ at new unseen points one needs to unroll the chain of compositions.

Therefore, here we propose to parametrize $T$ as a neural network $T_{w}$, where $w$ are the weights initialized at a point $w_{0}$ such that $T_{w_{0}}$ is the identity map. Having defined a Lagrangian $\mathcal{L}(T_{w(t)},\mu,\nu)$ we can simply evolve $w(t)$ by gradient descent:

[TABLE]

In reality updates are done stochastically drawing batches of samples from $\mu$ and $\nu$. Note however, that while we still use flows, we are no longer making regularity assumptions on the measure $\mu$, $\nu$ or the spaces $\mathcal{X}$ and $\mathcal{Y}$.

3.3. Flow Potentials

Let us first forget about the cost part and look for flows that would induce $T_{w(t)\#}\mu\to\nu$ as $t\to\infty$. Here we look for Lagrangians of the form $\mathcal{L}(T_{w(t)\#}\mu,\nu)$. If one had $T_{w(t)\#}\mu=\nu$ then for each $f\in C({\mathcal{Y}})$ one would have:

[TABLE]

Therefore if we choose a sufficiently rich family $\{f_{k}\}_{k=1}^{K}\subset C({\mathcal{Y}})$ we might hope of flowing $\mu$ into $\nu$ by using a Lagrangian of the form:

[TABLE]

In reality, at each gradient descent iteration, one has access to batches $\{X_{i}\}_{i=1}^{b_{X}}\sim\mu$, $\{Y_{j}\}_{j=1}^{b_{Y}}\sim\nu$ and the Lagrangian becomes:

[TABLE]

Motivated by [TT16] in our experiments we choose quadratic polynomials to match means and covariances, or concentrated Gaussians with centers lying on a grid $\{z_{k}\}_{k=1}^{K}$:

[TABLE]

For a more geometric approach we could try to directly match the “shapes” of $\{X_{i}\}_{i=1}^{b_{X}}$, $\{Y_{j}\}_{j=1}^{b_{Y}}$ in a Hausdorff sense [PC19, Subsec 10.6.1]. Specifically assume that $\mathcal{X}$ and $\mathcal{Y}$ are subsets of a metric space $\mathcal{Z}$ with distance $d$. Fix $t$ and let $Y_{j_{l}(i)}$ denote an $l$-th closest point of $\{Y_{j}\}_{j=1}^{b_{Y}}$ to $T_{w(t)}(X_{i})$ (i.e. sort the points in ascending order of distance from $T_{w(t)}(X_{i})$, breaking ties arbitrarily, and take the $l$-th point). Similarly let $T_{w(t)}(X_{i_{l}(j)})$ denote an $l$-th closest point of $\{T_{w(t)}(X_{i})\}_{i=1}^{b_{X}}$ to $Y_{j}$. We can then define a discrepancy @ k and a symmetric discrepancy @ k:

[TABLE]

Note that ${\rm disc}_{{k}}$ will tend to move the mass of $\mu$ inside the support of $\nu$; while ${\rm symdisc}_{{k}}$ will tend to move the mass of $\mu$ inside the support of $\nu$ while matching the shape of $\nu$, thus avoiding collapsing $\mu$ to a measure singular with respect to $\nu$.

Finally, to bring back the optimal transport problem ($\mathcal{M}_{0}$) we simply introduce a term depending on $c$ which acts as a regularization for the flow associated with $\mathcal{L}(T_{w(t)\#}\mu,\nu)$:

[TABLE]

3.4. Adversarial Training

Adversarial training formulates the learning problem as a min-max:

[TABLE]

where the neural network $\theta\mapsto f_{\theta}$ parametrizes a function in $C({\mathcal{Y}})$ to penalize choices of $w$ such that $T_{w\,\#}\mu\neq\nu$. In a typical experiment in Euclidean space one would initialize $T_{w}$ to the identity and $f_{\theta}$ to the zero function.

The theoretical justification for ($Adv$) is that the inner $\max$ over $\theta$ would give $+\infty$ if $T_{w\,\#}\mu\neq\nu$. As in the case of GANs [GPAM*+*14] in practice one cannot really train a min-max because of the vanishing gradients problem [AB17]. In a typical application for each minimization step updating $w$ there is a loop of a fixed number of steps updating $\theta$.

In [ACB17] it is observed that if one could force $f_{\theta}$ to be $L$-Lipschitz for some $L>0$, then the penalization term in ($Adv$) involving $f_{\theta}$ would converge to a multiple of the $l_{1}$-Wasserstein distance between the measures $T_{w\,\#}\mu$ and $\nu$. In WGANs the authors propose a heuristic to keep $f_{\theta}$ Lipschitz by using gradient clipping. In our experiments we did not benefit much from gradient clipping and we are not quite sure whether the min-max training would yield a penalization term close to a Wasserstein distance.

3.5. Regularized duals after [SDF*+*18]

In [SDF*+*18] the authors propose to parametrize $u$ and $v$ in ($\mathcal{D}_{\varepsilon}$) or ($\mathcal{D}^{l_{2}}_{\varepsilon}$) as neural networks $u_{\theta}$, $v_{\eta}$. They use stochastic gradient descent to learn an optimal $(u_{\theta^{*}},v_{\eta^{*}})$ and can thus generate a transport plan using (2) or (3). Finally they also parametrize the transport map $T$ as a neural network $T_{w}$ and use the heuristic ($\pi\xrightarrow{\textrm{heur}}T$) to learn the parameter $w$.

Despite lack of rigorous guarantees on optimizing the regularized dual(s) using neural networks, this approach solves scalability issues on large data-sets. Assume that $\mu$ (resp. $\nu$) is represented by storing observed data in a data-set $\mathbf{X}$ (resp. $\mathbf{Y}$). Then the cost $c$ can be stored as a matrix which grows as $\#\mathbf{X}\times\#\mathbf{Y}$, making its storage cost prohibitive. The advantage of [SDF*+*18] is that one needs only to compute $c$ on each batch, incurring a cost growing as $b_{X}\times b_{Y}$. Nevertheless, it might be the case that using heuristics like in [GM17] one might be able to sparsfiy $c$ reducing its storage cost. However, it seems like that scaling [GM17] to large datasets would require distributing the computation across several nodes in a cluster.

To summarize, the approach of [SDF*+*18] is very effective as large data-sets can be streamed to a single machine incurring a fixed cost in memory usage. Note that the dimensionality of the parameters $\theta$, $\eta$, $w$ will grow with the dimensions of the spaces $\mathcal{X}$, $\mathcal{Y}$ so there is still an effect of the dimensionality on the scalability of this approach. Moreover, it might also be the case that in high dimensions large batch sizes might be needed to properly train the neural networks [WB17]. We think investigation of the interrelation between converge, batch sizes and dimensionality of $\mathcal{X}$, $\mathcal{Y}$ is an interesting topic for further research.

3.6. Supervised learning

An alternative approach to [SDF*+*18] is to use ($\mathcal{D}_{\varepsilon}$) as a “source of truth” to generate a stream of training data. One would then train the neural networks using this data, reducing optimal transport to a supervised learning task.

To fix the ideas, let us go back to learning $u_{\theta}$, $v_{\eta}$; on each batch we can use the Sinkhorn Algorithm [SK67] to find optimal $\{\hat{u}_{i}\}_{i=1}^{b_{X}}$, $\{\hat{v}_{j}\}_{j=1}^{b_{Y}}$ and then iterate to minimize a loss like:

[TABLE]

An advantage of this approach is that Sinkhorn iterations, especially on small batches, are much faster to converge than stochastic gradient descent for ($\mathcal{D}_{\varepsilon}$). Moreover, learning the source of truth for large data-sets might be carried out in parallel on different nodes in a cluster; finally, the training data would be streamed to a single node to optimize the parameters $\theta$, $\eta$.

In the case in which the training data is still fitted at the time of training $u_{\theta}$, $v_{\eta}$, a further advantage of this approach is that one can supply $\{u_{\theta}(X_{i})\}_{i=1}^{b_{X}}$, $\{v_{\eta}(X_{j})\}_{j=1}^{b_{Y}}$ as starting values for the Sinkhorn iterations. A disadvantage of this approach involving potentials is that each solution $\{\hat{u}_{i}\}_{i=1}^{b_{X}}$, $\{\hat{v}_{j}\}_{j=1}^{b_{Y}}$ is only well-defined up to an additive constant, i.e. $\{\hat{u}_{i}-C\}_{i=1}^{b_{X}}$, $\{\hat{v}_{j}+C\}_{j=1}^{b_{Y}}$ would still yield a solution. In our experiments we overcome this issue enforcing some kind of normalization.

There is nothing special about learning $u_{\theta}$, $v_{\eta}$: one can apply this approach to directly solve on the batches for an optimal plan $\hat{\pi}_{i,j}$ or transport map $\hat{T}$ and then train neural networks $\pi_{w}$ or $T_{w}$ to approximate $\hat{\pi}_{i,j}$, $\hat{T}$. In our experiments we also explore these approaches. Learning $\hat{T}$ is particularly favorable as one avoids the additional step in [SDF*+*18] of using the potentials to learn the map.

Finally, in this approach there is nothing special about Sinkhorn iterations or ($\mathcal{D}_{\varepsilon}$). One can use the reduction of optimal transport to supervised learning anytime there is a good black-box approach to learn potentials, plans or maps on the minibatches.

3.7. Initializations and general metric spaces

In our experiments we focus on the cases in which $\mathcal{X}$ and $\mathcal{Y}$ are subsets of ${\mathbb{R}}^{d}$ and $c$ is the standard quadratic cost. In that case we initialize potentials to [math], plans to the product measure and maps to the identity. At the start of the project we made simple tests on using neural networks to learn a transport map which is linear. We found that even in $1$-dimensions initialization to the identity can be important. For example, if $T_{w_{0}}$ has negative determinant, e.g. in the case in which $T_{w_{0}}$ reverses some space directions, we find that during training updates to the parameters $w_{0}$ will tend to keep the determinant negative.

In the case of $\mathcal{X}$, $\mathcal{Y}$ being general metric spaces (e.g. Carnot groups or graphs) one can always initialize potentials to [math] or plans to the product measure. However, initialization / representation of the transport map will depend on the metric space. We think this is also an interesting area for further research.

4. Experiments

For details on the commands we used to run the experiments we refer to the GitRepo, in file Experiments.md. Performance metrics and pictures can be found in evaluation_metrics/ and snapshot_images_and_movies/.

4.1. The Dataset

The crucial decision we have taken at the beginning of our work has been to limit our tests to a single dataset while exploring a variety of algorithms. We have deliberately decided to keep the dataset as simple as possible; however we have avoided trying to learn a linear transport map.

We have thus opted for planar measures where the transport map is not smooth as it requires splitting the domain of $\mu$ into pieces. Specifically, $\mu$ is the uniform distribution on the unit ball; $\nu$ is a uniform distribution supported on $4$ balls of radius $\frac{1}{2}$ and centers at the points $(\pm 1,\pm 1)$. See Figure 4.1 for a visualization.

4.2. Performance

In order to assess the performance of each algorithm we compute the mean squared distance between the transport map $T_{w}$ and the “optimal” one $T_{\textrm{opt}}$:

[TABLE]

To define $T_{\textrm{opt}}$ we take samples $\{X_{i}\}_{i=1}^{B}\sim\mu$, $\{Y_{j}\}_{j=1}^{B}\sim\nu$ ($B=1000$ in our experiments) and compute a “ground truth” $T_{\textrm{opt}}$ using the Sinkhorn algorithm with a low value of the entropic regularization (i.e. we set $\varepsilon=10^{-2}$ in ($\mathcal{P}_{\varepsilon}$) in our experiments). We thus compute:

[TABLE]

We stress that our experiment are dynamic, i.e. $T_{w}$ is the final point of a “flow” $\{T_{w(t)}\}_{t}$ starting at $T_{w(0)}$ being the identity map (at least when restricted to the support of $\mu$). We thus take a fixed number $S$ ($S=50$ in our experiments) of snapshots of $T_{w(t)}$ across the training iterations $t\in\{0,\cdots,T-1\}$. In this way we are both able to estimate the rate of convergence and the stability of the training process.

A non trivial thing to account for is that different algorithms run for a different number $T$ of iterations. In our plots we account for this normalizing the steps with $t\mapsto t/T$ to compare the convergence rate across the iterations. Once we single out promising algorithms we are able to dig more into the training times using statistics that we log at periodic intervals in Tensorboard.

Finally, for each of the $S$ timesteps we save an image of the map $T_{w(s)}$, see for example figure 4.2. At the end of training we compose the images into a movie (see GitRepo) to get an idea of the flow and also inspect visually the “quality” of the final map.

Even though visual inspection is not as objective as using an evaluation metric, it allows us to compare maps that have similar error $\varepsilon_{2}$ but different properties. For example we will see cases where $\varepsilon_{2}$ is similar between two different algorithms but in once case the final map is too much diffused around the support to $\nu$, while in the second case the final map is squashed inside the support of $\nu$. We will see also cases where $\varepsilon_{2}$ is comparable to the one of a map far from the optimal transport map, but nevertheless on visual inspection the final map looks quite “reasonable”.

While writing the training scripts we found qualitative methods helpful in debugging errors. Indeed at the start of this project we were just looking at metrics in Tensorboard, but quickly realized that we were not getting enough insights into the behavior of the “flow” $\{T_{w(t)}\}_{t}$. Obviously we are also helped by our choice of a lower dimensional dataset.

We think a possible area of further research is to design metrics that make these objective inspections more quantitative. We find that in practice just looking at the transport cost is not enough. Good inspection metrics should take into account not just how far or close $T_{w(t)\#}\mu$ is from the support of $\nu$, but also the relative shapes of the measures.

4.3. Comparison of Heuristics

Here is a summary of experiments we ran using the heuristics in Subsection 3.3:

• •

covariance: We try to force matching of second order momenta.

• •

exp: We use a grid of centers and use a family of Gaussian bumps in (6).

• •

discr_N: We use discrepancy @ k.

• •

sym_discr_N: We use symmetric discrepancy @ k.

When we add tp_ in front of an experiment name we use the cost as a regularization term, compare (9).

Results are reported in Table 4.1; $\varepsilon_{2}$ is the minimum $\varepsilon_{2}$ across the $S$ snapshot iterations; the $\varepsilon_{2}$ is realized at iteration $t_{\textrm{min}}$ out of the total $T$ iterations. The standard deviation $\sigma(\varepsilon_{2})$ of $\varepsilon_{2}$ is computed on the iterations in the $S$ snapshots that occur after $t_{\textrm{min}}$.

Except for covariance all methods benefit by adding the cost regularization. The final transport yielded by covariance is quite diffused, see Figure 4.3, and does not improve substantially after 10% of the iterations, see Figure 4.4.

For the flows discr_N the error $\varepsilon_{2}$ increases with the number of iterations, see Figure 4.5. We observe a quick chaotic push to collapse $T_{\#}\mu$ inside the support of $\nu$, see Figure 4.6. Adding regularization improves these flows but the results remain disappointing, not improving over covariance and not making substantial progress after the first 20% of the iterations, see Figure 4.7. Despite adding regularization we still observe a sharp collapsing tendency, see Figure 4.8.

The flow sym_discr benefits a little bit by adding regularization with tp_sym_discr_5 reaching the minimal $\varepsilon_{2}$. After the first 50% of iterations there is not a real improvement, see Figure 4.9. On visual inspection this flow is quite successful in mapping $\mu$ into the support of $\nu$ while keeping the shapes of $T_{\#}\mu$ and $\nu$ comparable, see Figure 4.10.

The flows based on exp give the best results and greatly benefit from adding the regularization. The algorithm exp reaches the minimum after about 30% of the iterations while tp_exp does not improve substantially after 60% of the iterations, see Figure 4.11. Both flows are quick to map $\mu$ into a measure matching the geometry of $\nu$, see Figures 4.12 and 4.13.

4.4. Comparison of Adversarial Models

Table 4.2 summarizes the results of trying to learn a transport map following the “adversarial” approach of Subsection 3.4. The name of each model refers to training parameters: for example naming a model adv_l10.0_3_clip_0.01 means that the adversarial term (i.e. the one involving integration over $f_{\theta}$) in ($Adv$) is “boosted” by multiplying $f_{\theta}$ by a factor $\lambda=10$; moreover the “3” means that for each gradient descent step for the parameters $w$, 2 gradient ascent steps are performed for the parameters $\theta$; finally if “clip” is used, one is specifying the parameter to clip the gradients, e.g. in our case to $0.01$.

In Figure 4.14 we can inspect the convergence rate of the adversarial models. We see that the “out-of-the-box” parameter $\lambda=1$ seems to yield the best results and gradient clipping does not seems to help. Large values of $\lambda$ make the training unstable; for example with $\lambda=100$ there is a steady increase in $\varepsilon_{2}$ after about 25% of the training iterations. Inspection of frames shows a lot of chaotic behavior in this case, compare Figure 4.15.

Also increasing the number of iterations of the adversarial network leaves it stuck at a point with a higher value of $\varepsilon_{2}$. The best model is adv_l1_2 which reaches the minimal $\varepsilon_{2}$ at around 75% of the iterations. However, the final results are quite disappointing, both in terms of $\varepsilon_{2}$, which is higher than the one obtained simply with the exp flow, and of the final map as $T_{w(T)\#}\mu$ lies into a subset of the support of $\nu$, see Figure 4.16. Generally speaking, we found adversarial training to be under-performing and very sensitive to parameter specification.

In Figure 4.17 we compare the adversarial and the cost component of the loss during the training. We observe that at the beginning of the training (say $t<10k$) there is a first phase displaying a quick increase in the adversarial component which is then followed by a saturation phase in which the adversarial component stays constant while the cost component increases. The first phase corresponds to the adversarial loss starting to discriminate between $\mu$ and $\nu$ and the second phase to the cost increasing as $\mu$ starts to move towards $\nu$.

In the training regime (say $t\geq 10k$ and $t<30k$) the adversarial loss continues to grow driving a further rearrangement of $T_{w(t)\#}\mu$ inside the support of $\nu$. Finally, the last phase (say $t\geq 30k$) is dominated by a degree of reduction of the adversarial loss leading to a collapse of $T_{w(t)\#}\mu$ inside the support of $\nu$. This phenomenon points out that at a certain point the adversarial network starts to fail at discriminating between $T_{w(t)\#}\mu$ and $\nu$. Thus, for the rest of the paper we set the adversarial models aside.

4.5. Comparison of Approaches based on duals & supervised learning

In table 4.3 we can see a comparison of the performances of the algorithms based on the dual methods of Subsection 3.5 and supervised learning of Subsection 3.6. The models with a name seguy_.* can use either the entropic or the $l_{2}$ regularization. Note that the name can contain either “mean” or “sum” depending on how the regularization is aggregated across the batch. In fact, we sum the potentials $\{u_{i}\}_{i=1}^{B}$, $\{v_{j}\}_{j=1}^{B}$ on the batch; however, the regularization is a matrix $\{R_{i,j}\}_{i,j=1}^{B}$ and the mathematical formulation of duality [GCPB16, e.g. Sec. 3] would suggest that we need to sum on one dimension (say $i$) and take the mean on the other one (say $j$). However the equations in the original paper [SDF*+*18, Alg. 1] suggest to use a double sum for the regularization term. While taking the sum degrades a bit the entropic model, it has a positive effect on the $l_{2}$-regularized model, which otherwise suffers from a collapsing phenomenon, see Figure 4.18. Overall seguy_ent_mean and seguy_l2_sum have comparable performances, see Figures 4.19 and 4.20, with the latter requiring less iterations.

In Figure 4.21 we can see the convergence rate for the models using a supervised approach. Except for supervised_prob all these models have a higher variance $\sigma(\varepsilon_{2})$ but, except for supervised_prob, the final performance is comparable to the one of seguy_ent_mean and seguy_l2_sum.

In the model supervised_prob we try to first learn a transport plan and then use the heuristic ($\pi\xrightarrow{\textrm{heur}}T$) to learn a transport map. In practice this model performs poorly and the $\varepsilon_{2}$ starts to grow; on a qualitative inspection we observe that the first plans tend to distort the geometry of $T_{\#}\mu$ relatively to that of $\nu$, while towards the end of the training $T_{\#}\mu$ becomes too diffused, see Figure 4.22. Finally, we also point out that the training of this model is considerably slower than for the others as it requires fitting a neural network for the plan against a $B\times B$ matrix (while for potentials or maps we fit against vectors of size $B$).

Both the methods based on the supervised dual or the supervised map perform well. The high variance seems to be an artifact of the iterations using supervised data and can be reduced in practice by using a validation rule which decides when to stop the fitting of the map $T$. Concretely, periodically one evaluates the performance of $T_{w(t)}$ against a “ground-truth” $T_{\textrm{opt}}$ to decide when to stop the training of $T$. Note that even with the high variance the maps tend to stay qualitatively closed to $T_{\textrm{opt}}$, i.e. we do not observe phenomena like collapsing or diffusion, see Figure 4.23. Finally, we point out that the approach based on directly learning $T$ is faster while both the models seguy_.* and supervised_dual_.* require a second step to fit $T$.

4.6. Timings

We now discuss the time consumption for the models that we find more promising, see Table 4.4. We ran our experiments on a laptop with a $4$-Core Intel (R) Core (TM) i5-5257U@2.70 GHz CPU. We used a docker image of ubuntu:bionic and version 1.0 of PyTorch on Python 3.6.

In Table 4.4 the seconds per step are computed using the wall clock time. The seconds per step are then used to estimate the time to $t_{\textrm{min}}$, which we recall was the “best” of the iterations in the $S$ snapshots. For models requiring two steps, namely training networks to optimize a dual problem and use the potentials to learn a transport map, we report two rows of timings, one for the dual and another for the map. While secs per step is an objective metric, secs to $t_{\textrm{min}}$ must be taken with a grain of salt as it depends on the parameters used to run the corresponding script.

In general, total timings are relatively comparable. However, we see the advantage of using models like tp_exp or supervised_map_iters_1000_0.05 that learn directly the transport map. Using Sinkhorn iterations also gives a speed up compared to using the dual problem. For example in seguy_l2_sum_0.1 about $7.95\times 10^{-3}$ seconds are required for each step of the dual problem, compared to the $4.82\times 10^{-3}$ required by supervised_dual_0.1.

4.7. Conclusion

We have compared a variety of approaches to find an optimal map between probability distributions.

We find that, despite different theoretical/heuristic justifications, some algorithms yield similar good optimal maps. Specifically, we find flows using local Gaussian bumps, supervised learning approaches learning potentials or directly a transport map and the dual formulation of [SDF*+*18] to yield good results. In terms of time consumption, algorithms learning directly the transport map and using Sinkhorn’s iterations are more favorable.

On the other hand, we also find other approaches to under-perform or being unstable. In particular, flows using the covariance loss or the **discr @ ** $N$ seem to yield poor maps. Approaches using adversarial training yield poor maps and are also unstable to train.