Optimal transport inference

6.1 OTT

Optimal Transport Tools (OTT) (Cuturi et al. 2022), a toolbox for all things Wasserstein (documentation):

The goal of OTT is to provide sturdy, versatile and efficient optimal transport solvers, taking advantage of JAX features, such as JIT, auto-vectorization and implicit differentiation.

A typical OT problem has two ingredients: a pair of weight vectors a and b (one for each measure), with a ground cost matrix that is either directly given, or derived as the pairwise evaluation of a cost function on pairs of points taken from two measures. The main design choice in OTT comes from encapsulating the cost in a Geometry object, and [bundling] it with a few useful operations (notably kernel applications). The most common geometry is that of two clouds of vectors compared with the squared Euclidean distance, as illustrated in the example below:

import jax
import jax.numpy as jnp
from ott.tools import transport
# Samples two point clouds and their weights.
rngs = jax.random.split(jax.random.PRNGKey(0),4)
n, m, d = 12, 14, 2
x = jax.random.normal(rngs[0], (n,d)) + 1
y = jax.random.uniform(rngs[1], (m,d))
a = jax.random.uniform(rngs[2], (n,))
b = jax.random.uniform(rngs[3], (m,))
a, b = a / jnp.sum(a), b / jnp.sum(b)
# Computes the couplings via Sinkhorn algorithm.
ot = transport.solve(x, y, a=a, b=b)
P = ot.matrix

The call to sinkhorn above works out the optimal transport solution by storing its output. The transport matrix can be instantiated using those optimal solutions and the Geometry again. That transport matrix links each point from the first point cloud to one or more points from the second, as illustrated below.

To be more precise, the sinkhorn algorithm operates on the Geometry, taking into account weights a and b, to solve the OT problem, produce a named tuple that contains two optimal dual potentials f and g (vectors of the same size as a and b), the objective reg_ot_cost and a log of the errors of the algorithm as it converges, and a converged flag.

6.2 POT

POT: Python Optimal Transport (Rémi Flamary et al. 2021)

This open source Python library provides several solvers for optimization problems related to Optimal Transport for signal, image processing and machine learning.

Website and documentation: https://PythonOT.github.io/

Source Code (MIT): https://github.com/PythonOT/POT

POT provides the following generic OT solvers (links to examples):

• OT Network Simplex solver for the linear program/ Earth Movers Distance.
• Conditional gradient and Generalized conditional gradient for regularized OT.
• Entropic regularization OT solver with Sinkhorn Knopp Algorithm, stabilized version, greedy Sinkhorn and Screening Sinkhorn.
• Bregman projections for Wasserstein barycenter, convolutional barycenter and unmixing.
• Sinkhorn divergence and entropic regularization OT from empirical data.
• Debiased Sinkhorn barycenters Sinkhorn divergence barycenter
• Smooth optimal transport solvers (dual and semi-dual) for KL and squared L2 regularizations.
• Weak OT solver between empirical distributions
• Non regularized Wasserstein barycenters with LP solver (only small scale).
• Gromov-Wasserstein distances and GW barycenters (exact and regularized), differentiable using gradients from Graph Dictionary Learning
• Fused-Gromov-Wasserstein distances solver and FGW barycenters
• Stochastic solver and differentiable losses for Large-scale Optimal Transport (semi-dual problem and dual problem)
• Sampled solver of Gromov Wasserstein for large-scale problem with any loss functions
• Non regularized free support Wasserstein barycenters.
• One dimensional Unbalanced OT with KL relaxation and barycenter $$10,25$$. Also exact unbalanced OT with KL and quadratic regularization and the regularization path of UOT
• Partial Wasserstein and Gromov-Wasserstein (exact and entropic formulations).
• Sliced Wasserstein $$31,32$$ and Max-sliced Wasserstein that can be used for gradient flows.
• Graph Dictionary Learning solvers.
• Several backends for easy use of POT with Pytorch/jax/Numpy/Cupy/Tensorflow arrays.

POT provides the following Machine Learning related solvers:

• Optimal transport for domain adaptation with group lasso regularization, Laplacian regularization and semi supervised setting.
• Linear OT mapping and Joint OT mapping estimation.
• Wasserstein Discriminant Analysis (requires autograd + pymanopt).
• JCPOT algorithm for multi-source domain adaptation with target shift.

Some other examples are available in the documentation.

6.3 GeomLoss

• GeomLoss: Geometric Loss functions between sampled measures, images and volumes

The GeomLoss library provides efficient GPU implementations for:

• Kernel norms (also known as Maximum Mean Discrepancies).
• Hausdorff divergences, which are positive definite generalizations of the Chamfer-ICP loss and are analogous to log-likelihoods of Gaussian Mixture Models.
• Debiased Sinkhorn divergences, which are affordable yet positive and definite approximations of Optimal Transport (Wasserstein) distances.

It is hosted on GitHub and distributed under the permissive MIT license.

GeomLoss functions are available through the custom PyTorch layers SamplesLoss, ImagesLoss and VolumesLoss which allow you to work with weighted point clouds (of any dimension), density maps and volumetric segmentation masks.