[Paper Review] Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein

Marco Cuturi, Laetitia Meng-Papaxanthos+4 more|arXiv (Cornell University)|Jan 28, 2022

Asphalt Pavement Performance Evaluation18 citations

TL;DR

OTT-JAX is a JAX-based Python toolbox that enables efficient, differentiable optimal transport computations using entropic regularization and low-rank approximations. It supports linear and quadratic OT problems, barycenters, Gromov-Wasserstein, and Gaussian mixture matching, with scalable solvers and automatic differentiation for machine learning applications.

ABSTRACT

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 exttt{https://github.com/ott-jax/ott}

Motivation & Objective

• To address the computational and differentiability challenges in optimal transport (OT) for large-scale and differentiable machine learning applications.
• To provide a unified, high-performance framework for solving regularized OT problems between point clouds, histograms, and measures.
• To enable differentiable OT computations via JAX’s automatic differentiation and JIT compilation, supporting end-to-end training in deep learning pipelines.
• To extend OT capabilities beyond standard Wasserstein distances to include barycenters, Gromov-Wasserstein, and soft-sorting operations.
• To support efficient computation through low-rank approximations and geometry-aware cost computation without explicit matrix storage.

Proposed method

• Leverages JAX’s automatic differentiation and JIT compilation to enable differentiable OT solvers with high performance on CPU and TPU/GPU.
• Implements entropic regularization via the Sinkhorn algorithm to smooth the optimal transport plan and enable efficient, differentiable optimization.
• Introduces low-rank Sinkhorn solvers that reduce memory and computation costs by approximating the transport matrix with rank-r factors.
• Uses geometry classes to compute cost matrices implicitly, avoiding explicit storage—e.g., for point clouds via kernelized operations or grid-based structures.
• Supports Gromov-Wasserstein via iterative linearization, solving quadratic OT problems by reducing them to sequences of linear OT problems.
• Integrates input-convex neural networks (ICNNs) for learning convex maps and soft-sorting operations via differentiable OT-based reparameterizations.

Experimental results

Research questions

• RQ1 How can optimal transport problems be solved efficiently and differentiably at scale using modern deep learning frameworks?
• RQ2 To what extent can low-rank approximations and geometry-aware computation reduce memory and time complexity in OT?
• RQ3 Can differentiable OT be used to learn structured representations such as barycenters or maps between distributions?
• RQ4 How does implicit differentiation through the optimal transport plan improve training stability in differentiable ML pipelines?
• RQ5 Can advanced OT variants like Gromov-Wasserstein and soft-sorting be efficiently implemented with full differentiability and scalability?

Key findings

• OTT-JAX enables differentiable optimal transport with full support for JAX’s automatic differentiation, allowing backpropagation through the transport plan.
• The low-rank Sinkhorn solver achieves significant memory and computational savings, especially for large-scale problems, without sacrificing accuracy.
• The toolbox supports end-to-end differentiable computation of barycenters, Gromov-Wasserstein distances, and soft-sorted arrays.
• Geometry classes allow implicit computation of cost matrices (e.g., for point clouds or grids), avoiding explicit storage and enabling O(dn^{d+1}) operations on grids.
• The implementation supports efficient computation of Wasserstein-like distances between Gaussian mixtures using a differentiable approximation from Delon and Desolneux (2020).
• The toolbox is production-ready and used in advanced applications such as iso-barycenter computation on medical imaging data and shape matching between complex manifolds like spirals and Swiss rolls.