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.

Tutorials:

• Point clouds
• Grid geometry

Benchmarks:

• OTT vs. POT
• Sinkhorn in all flavors
• Low-Rank Sinkhorn

Advanced Applications:

• Sinkhorn Barycenters
• Gromov Wasserstein
• Low-Rank Gromov-Wasserstein
• Sinkhorn Divergence Hessians
• Soft Sorting
• Single-cell genomics
• Fairness regularizers
• Learning the Kantorovich Dual using Input Convex Neural Networks

Public API: ott packages

• ott.geometry package
• ott.core package
• ott.tools package

Indices and tables

• Index

• Module Index

• Search Page

References

1

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

2

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

3

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

4

10. 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

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

7

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