ott
ott copied to clipboard
Optimal transport tools implemented with the JAX framework, to get differentiable, parallel and jit-able computations.
![logo](https://raw.githubusercontent.com/ott-jax/ott/main/docs/_static/images/logoOTT.png)
Optimal Transport Tools (OTT).
See full documentation.
What is OTT-JAX?
Optimal transport theory can be loosely described as the branch of mathematics and optimization that studies matching problems: given two sets of points, how to find (given some prior information, typically a cost function) a good way to associate bijectively every point in the first set with another in the second. A typical matching problem arises, for instance, when sorting numbers (when sorting, one associates to numbers [3.1, -4.2, -18, 5.4] the ranks [3, 2, 1, 4] that reorder them in increasing fashion) or when matching submitted papers with reviewers at ML conferences!
These problems are easy to describe yet hard to solve. Indeed, while matching optimally two sets of n points using a pairwise cost can be solved with the Hungarian algorithm, this requires an effort that scales as $n^3$. Additionally, one may run into various issues. For instance, the two sets might have different sizes and a relevant matching cost function might not be given (it remains an open problem to set a score to qualify how well a reviewer is likely to be to check a paper!). Fortunately, optimal transport theory has made decisive progresses since the Hungarian algorithm was proposed, and can count on many efficient algorithms and extensions that can handle such situations.
OTT-JAX is a toolbox providing sturdy, scalable and efficient solvers for those problems. OTT builds upon the JAX framework. Some of JAX features include JIT, auto-vectorization and implicit differentiation.
The typical ingredient in OT problems consists of two discrete measures (an efficient way to encode weighted sets of points), and a cost function comparing points. By default, OTT assumes that the measures are supported on vectors, and that these vectors are compared with a squared Euclidean distance, as given below:
Example
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. That transoprt matrix links each point from the first point cloud to one or more points from the second, as illustrated below.
Overall description of source code
Currently implements the following classes and functions:
-
The geometry folder describes tools that to encapsulate the essential ingredients of OT problems: measures and cost functions.
-
The
CostFn
class in costs.py and its descendants define cost functions between points. A few simple costs are considered,Euclidean
between vectors, andBures
, between a pair of mean vector and covariance (p.d.) matrix. -
The
Geometry
class in geometry.py describes a cost structure between two measures. That cost structure is accessed through various member functions, either used when running the Sinkhorn algorithm (typically kernel multiplications, or log-sum-exp row/column-wise application) or after (to apply the OT matrix to a vector).-
In its generic
Geometry
implementation, as in geometry.py, an object can be initialized with either acost_matrix
along with anepsilon
regularization parameter (or scheduler), or with akernel_matrix
. -
If one wishes to compute OT between two weighted point clouds $x=\left(x_1, \ldots, x_n\right)$ and $y=\left(y_1, \ldots, y_m\right)$ endowed with a cost function (e.g. Euclidean) $c$, the
PointCloud
class in pointcloud.py can be used to define the corresponding cost and kernel matrices $C_{i j}=c\left(x_{i}, y_{j}\right)$ and $K_{i j}=\exp\left(-c\left(x_{i}, y_{j}\right) / \epsilon\right)$. When $n$ and $m$ are very large, this geometry can be instantiated with abatch_size
parameter, to avoid storing the cost and/or kernel matrices, to recompute instead these matrices on the fly as needed,batch_size
lines at a time, at each application. -
Simlarly, if all measures to be considered are supported on a separable grid (e.g. $\{1, \ldots, n\}^{d}$), and the cost is separable along all axis, i.e. the cost between two points on that grid is equal to the sum of (possibly $d$ different) cost functions evaluated on each of the $d$ pairs of coordinates, then the application of the kernel is much simplified, both in log space or on the histograms themselves. This particular case is exploited in the
Grid
geometry in grid.py which can be instantiated as a hypercube using agrid_size
parameter, or directly through grid locations inx
. -
LRCGeometry
, low-rank cost geometries, of which aPointCloud
endowed with a squared-Euclidean distance is a particular example, can efficiently carry apply their cost to another matrix. This is leveraged in particular in the low-rank Sinkhorn (and Gromov-Wasserstein) solvers.
-
-
-
In the core folder,
-
The
sinkhorn
function in sinkhorn.py is a wrapper around theSinkhorn
solver class, running the Sinkhorn algorithm, with the aim of solving approximately one or various optimal transport problems in parallel. An OT problem is defined by aGeometry
object, and a pair $\left(a, b\right)$ (or batch thereof) of histograms. The function's outputs are stored in aSinkhornOutput
named t-uple, containing potentials, regularized OT cost, sequence of errors and a convergence flag. Such outputs (with the exception of errors and convergence flag) can be differentiated w.r.t. any of the three inputs(Geometry, a, b)
either through backprop or implicit differentiation of the optimality conditions of the optimal potentialsf
andg
. -
A later addition in sinkhorn_lr.py is focused on the
LRSinkhorn
solver class, which is able to solve OT problems at larger scales using an explicit factorization of couplings as being low-rank. -
In discrete_barycenter.py: implementation of discrete Wasserstein barycenters : given $N$ histograms all supported on the same
Geometry
, compute a barycenter of theses measures, using an algorithm by Janati et al. (2020). -
In continuous_barycenter.py: implementation of continuous Wasserstein barycenters : given
probability measures described as points which can be compared with an arbitrary cost function, compute a barycenter of theses measures, supported at most $k$ points on using an algorithm by Cuturi and Doucet (2014).
-
In gromov_wasserstein.py: implementation of two Gromov-Wasserstein solvers (both entropy-regularized and low-rank) to compare two measured-metric spaces, here encoded as a pair of
Geometry
objects,geom_xx
,geom_xy
along with weightsa
andb
. Additional options include using a fused term by specifyinggeom_xy
.
-
-
In the tools folder,
-
In soft_sort.py: implementation of soft-sorting operators, notably soft-quantile transforms
-
The
sinkhorn_divergence
function in sinkhorn_divergence.py, implements the unbalanced formulation of the Sinkhorn divergence, a variant of the Wasserstein distance that uses regularization and is computed by centering the output ofsinkhorn
when comparing two measures. -
The
Transport
class in sinkhorn_divergence.py, provides a simple wrapper to thesinkhorn
function defined above when the user is primarily interested in computing and storing an OT matrix. -
The gaussian_mixture folder provides novel tools to compare and estimate GMMs with an OT perspective.
-
Citation
If you have found this work useful, please consider citing this reference:
@article{cuturi2022optimal,
title={Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein},
author={Cuturi, Marco and Meng-Papaxanthos, Laetitia and Tian, Yingtao and Bunne, Charlotte and Davis, Geoff and Teboul, Olivier},
journal={arXiv preprint arXiv:2201.12324},
year={2022}
}