Optimal Transport Tools (OTT)
Contents
Optimal Transport Tools (OTT)#
Introduction#
OTT is a JAX package that bundles a few utilities to compute,
and differentiate as needed, the solution to optimal transport (OT) problems, taken in a fairly wide sense.
For instance, OTT can of course compute Wasserstein (or Gromov-Wasserstein) distances between
weighted clouds of points (or histograms) in a wide variety of scenarios,
but also estimate Monge maps, Wasserstein barycenters, and help with simpler tasks
such as differentiable approximations to ranking or even clustering.
To achieve this, OTT rests on two families of tools:
The first family consists in discrete solvers computing transport between point clouds,
using the Sinkhorn [Cuturi, 2013] and low-rank Sinkhorn [Scetbon et al., 2021] algorithms,
and moving up towards Gromov-Wasserstein [Mémoli, 2011, Peyré et al., 2016];
the second family consists in continuous solvers, using suitable neural architectures [Amos et al., 2017] coupled
with SGD type estimators [Korotin et al., 2021, Makkuva et al., 2020].
Installation#
Install OTT from PyPI as:
pip install ott-jax
or with conda via conda-forge as:
conda install -c conda-forge ott-jax
Design Choices#
OTT is designed with the following choices:
Take advantage whenever possible of JAX features, such as Just-in-time (JIT) compilation, auto-vectorization (VMAP) and both automatic but most importantly implicit differentiation.
Split geometry from OT solvers in the discrete case: We argue that there should be one, and one implementation only, of every major OT algorithm (Sinkhorn, Gromov-Wasserstein, barycenters, etc…), regardless of the geometric setup that is considered. To give a concrete example, any speedups one may benefit from by using a specific cost (e.g. Sinkhorn being faster when run on a separable cost on histograms supported on a separable grid [Solomon et al., 2015]) should not require a separate reimplementation of a Sinkhorn routine.
As a consequence, and to minimize code copy/pasting, use as often as possible object hierarchies, and interleave outer solvers (such as quadratic, aka Gromov-Wasserstein solvers) with inner solvers (e.g. Low-Rank Sinkhorn). This choice ensures that speedups achieved at lower computation levels (e.g. low-rank factorization of squared Euclidean distances) propagate seamlessly and automatically in higher level calls (e.g. updates in Gromov-Wasserstein), without requiring any attention from the user.
Packages#
ott.geometry package contains classes to instantiate objects that describe two point clouds paired with a cost function. Geometry objects are used to describe OT problems, handled by solvers in the ott.solvers package.
ott.tools package provides an interface to exploit OT solutions, as produced by solvers in the ott.solvers package. Such tasks include computing approximations to Wasserstein distances [Genevay et al., 2018, Séjourné et al., 2019], approximating OT between GMMs, or computing differentiable sort and quantile operations [Cuturi et al., 2019].
Examples
API
References