Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein (2201.12324v1)

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

• The paper presents OTT-JAX, a robust framework that leverages JAX’s automatic differentiation and just-in-time compilation for efficient optimal transport computations.
• It efficiently computes regularized OT, barycenters, and Gromov-Wasserstein calculations using entropic regularization and low-rank approximations.
• The toolbox addresses scalability and differentiability challenges, enabling practical applications in machine learning, data science, and advanced computational settings.

Optimal Transport Tools (OTT): A JAX Toolbox for All Things Wasserstein

Optimal Transport Tools (OTT-JAX) presents a robust framework designed for solving optimal transport (OT) problems involving point clouds and histograms. This toolbox capitalizes on the computational capabilities of JAX, incorporating features like automatic differentiation and just-in-time compilation, to facilitate efficient problem-solving across various computational environments.

Core Contributions

The primary focus of OTT-JAX lies in addressing fundamental and advanced OT computations. These include solving the regularized OT problem, computing barycenters, Gromov-Wasserstein calculations, and exploring low-rank solutions. Additionally, the toolbox delivers capabilities for estimating convex maps, obtaining differentiable generalizations of quantiles and ranks, and approximating OT between Gaussian mixtures.

Mathematical Background

Optimal transport involves assigning pairs of points from two distinct datasets optimally. The assignment aims to minimize or handle discrepancies, typically formulated through an objective function. In certain scenarios, where prior information regarding cost functions is lacking, OTT-JAX facilitates more generalized approaches using discrepancies between candidate cost functions.

The toolbox tackles linear and quadratic OT problems through relaxed permutations represented in transportation polytope matrices. Mathematically, these are described by two primary objectives:

• The linear OT problem, $\mathcal{L}_c(\mu, \nu)$, focuses on minimizing the transportation cost over matrices.
• The quadratic OT problem, $$\mathcal{Q}_{c_\mathcal{X}, c_\mathcal{Y}}(\mu, \nu)$$, extends this to capture more nuanced isometry characteristics.

Computational Challenges and Solutions

OTT-JAX addresses several key challenges in solving OT problems:

1. Scalability: High-dimensional datasets often make exact OT solutions computationally prohibitive. Using entropic regularization and low-rank approximations significantly alleviates computational burdens, enhancing scalability.
2. Curse of Dimensionality: Transforming empirical measures representing sample datasets, particularly in high dimensions, is tackled by leveraging statistical insights from entropic techniques.
3. Differentiability: While differentiating objective function values remains straightforward, obtaining derivatives of optimal transport plans is non-trivial. The toolbox addresses this through methods like implicit differentiation and low-rank Sinkhorn.

Implementation and Features

OTT-JAX is structured into distinct modules, each serving critical roles:

• Geometry Module: Abstracts mathematical properties of cost matrices, offering memory-efficient implementations.
• Core Module: Contains problem definitions and solvers, including Entropic and Low-Rank Sinkhorn solvers for linear problems, and Gromov solvers for quadratic challenges.
• Tools Module: Offers additional utilities like soft sorting, which underscore the toolbox's versatile applicability.

Applications and Implications

OTT-JAX's functionality extends to diverse applications. From computing regularized barycenters in morphometry to aligning geometric structures with Gromov-Wasserstein frameworks, the toolbox shows adaptability in merging traditional computations with differentiability insights. This capacity opens avenues in machine learning, particularly for applications requiring gradient-based optimization techniques.

Consequently, OTT-JAX forms a significant resource not only for theoretical explorations in mathematics and computer science but also for practical applications in data science and AI. The toolbox’s ability to efficiently manage computational and statistical requirements suggests potential for continual integration into more complex machine learning models and algorithms.

Future Directions

The development of OTT-JAX prompts further inquiry into refining approximations in large-scale and high-dimensional tasks. By expanding the underlying frameworks to incorporate novel regularization techniques or integrating with other scalable computational libraries, there is considerable potential for advancing its current capabilities. Additionally, exploring deeper connections between optimal transport and neural network methodologies may reveal further insights, propelling advancements in AI research.

In summary, OTT-JAX stands as an impactful contribution to computational optimal transport, offering a comprehensive and technically sophisticated toolbox for researchers and practitioners aiming to explore or utilize Wasserstein metrics within their work.