Source-linked AI summary
Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein
Marco Cuturi, Laetitia Meng-Papaxanthos, Yingtao Tian, Charlotte Bunne, Geoff Davis, Olivier Teboul
TL;DR
Optimal transport faces computational and statistical challenges, alongside difficulties differentiating optimal transport matrices. OTT organizes linear and quadratic transport problems into a toolbox with solvers and extensions, and includes examples of advanced applications.
Problem
Linear transport has super-cubic worst-case bounds, quadratic transport is NP-hard in a stated case, empirical measures face a curse of dimensionality, and optimal transport matrices can be poorly behaved under differentiation.
Source: §1 · p. 3Method
Results
Takeaways & Limitations
The documentation includes advanced-application notebooks that can run online in Colab without setup.
Source: §3 · p. 4
Abstract
from arXiv · showhide
Optimal transport tools (OTT-JAX) is a Python toolbox that can solve optimal transport problems between point clouds and histograms. The toolbox builds on various JAX features, such as automatic and custom reverse mode differentiation, vectorization, just-in-time compilation and accelerators support. The toolbox covers elementary computations, such as the resolution of the regularized OT problem, and more advanced extensions, such as barycenters, Gromov-Wasserstein, low-rank solvers, estimation of convex maps, differentiable generalizations of quantiles and ranks, and approximate OT between Gaussian mixtures. The toolbox code is available at \texttt{https://github.com/ott-jax/ott}
1. Background and Motivation: Optimal Transport in a Nutshell
The section frames optimal transport as matching weighted point sets by linear or pairwise-structure objectives, then identifies computational, statistical, and differentiability challenges that motivate OTT. It also outlines continuous couplings and regularized or low-rank approaches to these problems.
1. Background and Motivation: Optimal Transport in a Nutshell
- When a cross-domain ground cost is hard to specify, the quadratic objective compares within-domain costs and favors assignments close to an isometry.The section notes this setting can arise when the two spaces differ.Source: §1 · p. 2
- OTT focuses on linear and quadratic transport objectives over transportation matrices with prescribed marginal sums.The linear objective uses a point-to-point cost, while the quadratic objective compares within-dataset costs.Source: §1 · p. 2
- Continuous formulations optimize over couplings with marginals µ and ν, and discrete measures recover the earlier objectives.The section describes entropic regularization as helping with scalability, statistical aspects, and differentiability, while low-rank Sinkhorn imposes a low-rank structure on the transport matrix.Source: §1 · pp. 2–3
- Exact linear OT has super-cubic worst-case bounds, quadratic OT is NP-hard in a stated case, and empirical-measure estimation faces a curse of dimensionality.The section also notes that optimal transport matrices can be almost everywhere insensitive to small input-point changes.Source: §1 · p. 3
2. Implementation
OTT separates transport geometry and problem definitions from solver implementations, while providing components for linear and quadratic transport and several extensions.
2. Implementation
- Geometry objects encapsulate cost-matrix properties to avoid storing full cost matrices when that is unnecessary.For squared-Euclidean point-cloud costs, the matrix rank is at most d + 2; grid costs can be applied in O(dn^(d+1)) operations.Source: §2 · p. 3
- Problem classes combine a geometry with probability weights, and Sinkhorn or low-rank Sinkhorn solvers address linear transport problems.Gromov solvers address the quadratic problem through iterated calls to linear solvers.Source: §2 · p. 4
- The toolbox computes regularized barycenters between discrete measures when the barycenter support is fixed.Source: §2 · p. 4
- Additional components cover Gaussian-mixture transport, input-convex neural networks, and differentiable soft sorting.The Gaussian-mixture distance can be used to fit pairs of mixtures with a map between them.Source: §2 · p. 4
3. Applications
The applications section describes a typical point-cloud transport workflow and illustrates advanced applications combining algorithms with differentiability.
3. Applications
- A typical workflow builds point-cloud geometry, defines marginal weights and a linear problem, then applies Sinkhorn or low-rank Sinkhorn to obtain a transport matrix.The geometry can use a custom cost function.Source: §3 · p. 4
- Figure 1 illustrates a grid-geometry voxel barycenter, a spiral-to-Swiss-roll Gromov-Wasserstein match, and soft sorting with varying ε.The soft-sorted array uses different ε values to control differentiability.Source: §3 · p. 4, Figure 1 · p. 4