Optimal Transport Tools (OTT) documentation#

Code hosted on Github. To install, clone that repo or simply run pip install ott-jax.

Intro#

OTT is a JAX package that bundles a few utilities to compute and differentiate the solution to optimal transport problems. OTT can help you compute Wasserstein distances between weighted clouds of points (or histograms), using a cost (e.g. a distance) between individual points.

To that end OTT uses various implementation of the Sinkhorn algorithm 1 2 3. These implementation take advantage of several JAX features, such as Just-in-time (JIT) compilation, auto-vectorization (VMAP) and both automatic and/or implicit differentiation. A few tutorials are provided below, along with different use-cases, notably for single-cell genomics data 4.

Packages#

There are currently three packages, geometry, core and tools, playing the following roles:

  • geometry defines classes that describe two point clouds paired with a cost function (simpler geometries are also implemented, such as that defined by points supported on a multi-dimensional grids with a separable cost 5). The design choice in OTT is to state that cost functions and algorithms should operate independently: if a particular cost function allows for faster computations (e.g. squared-Euclidean distance when comparing grids), this should not be taken advantage of at the level of optimizers, but at the level of the problems description. Geometry objects are therefore only considered as arguments to describe OT problem handled in core, using subroutines provided by geometries;

  • core help define first an OT problem (linear, quadratic, barycenters). These problems are then solved using Sinkhorn algorithm and its variants, the main workhorse to solve OT in this package, as well as variants that can comppute Gromov-Wasserstein distances or barycenters of several measures;

  • tools provides an interface to exploit OT solutions, as produced by core functions. Such tasks include instantiating OT matrices, computing approximations to Wasserstein distances 6 7, or computing differentiable sort and quantile operations 8.

Indices and tables#

References#

1
  1. Cuturi, Sinkhorn Distances: Lightspeed Computation of Optimal Transport, NIPS’13.

2
  1. Peyré, M. Cuturi, Computational Optimal Transport, FNT in ML, 2019.

3
  1. Scetbon et al., Low-Rank Sinkhorn Factorization , ICML 2021.

4
  1. Schiebinger et al., Optimal-Transport Analysis of Single-Cell Gene Expression Identifies Developmental Trajectories in Reprogramming, Cell 176, 928–943.

5
  1. Solomon et al, Convolutional Wasserstein distances: efficient optimal transportation on geometric domains, ACM ToG, SIGGRAPH’15.

6
  1. Genevay et al., Learning Generative Models with Sinkhorn Divergences, AISTATS’18.

7
  1. Séjourné et al., Sinkhorn Divergences for Unbalanced Optimal Transport, arXiv:1910.12958.

8
  1. Cuturi et al. Differentiable Ranking and Sorting using Optimal Transport, NeurIPS’19.