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 a sturdy and versatile implementation of the Sinkhorn algorithm 1 2. This implementation takes advantage of several JAX features, such as Just-in-time (JIT) compilation, auto-vectorization (VMAP), and both automatic and/or implicit differentiation. Basic tutorial snippets are provided below, along with a more advanced example to process single-cell genomics data 3.

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 4). A geometry, along with weight vectors a and b, describes an OT problem. core algorithms can solve such problems, using subroutines provided by geometries;

  • core contains the Sinkhorn algorithm, the main workhorse to solve OT in this package, as well as variants that can be used to compute 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 5 6, or computing differentiable sort and quantile operations 7.

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. Schiebinger et al., Optimal-Transport Analysis of Single-Cell Gene Expression Identifies Developmental Trajectories in Reprogramming, Cell 176, 928–943.

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

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

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

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