DOCSEARCHbyAlgolia
Submit Docs

OTT-JAX

repository·main·Indexed 20 days ago

https://github.com/ott-jax/ott

A JAX-powered library for solving optimal transport problems at scale on accelerators. It provides discrete solvers like Sinkhorn, advanced matching algorithms including Gromov-Wasserstein and Wasserstein barycenters, and neural network-based approaches for approximating optimal transport maps. The library includes tools for defining geometries (point clouds, grids, graphs), cost functions, and regularizers, as well as experimental support for multimarginal entropic optimal transport.

Tokens
30.1K
Snippets
74
Records
140
Agent score
71%

What's inside ott-jax

  1. Overview of OTT-JAX packages

    main

    The library is organized into several specialized modules:

    • ott.geometry: Defines the ground cost (e.g., using ott.geometry.pointcloud.PointCloud). Geometry objects describe the problem.
    • ott.problems: Defines the interaction between measures (e.g., linear_problem.LinearProblem for Kantorovich problems, or quadratic problems for Gromov-Wasserstein).
    • ott.solvers: Implements the algorithms to solve the problems defined in ott.problems.
    • ott.initializers: Provides strategies to initialize solvers, which is critical for convergence in non-convex problems.
    • ott.neural: Tools for parameterizing transport maps as neural networks (e.g., input convex neural networks or flow_matching).
    • ott.tools: Utilities to exploit OT solutions, such as computing Wasserstein distances, GMM approximations, or plotting tools.
    • ott.experimental: Contains immature APIs (e.g., ott.solvers.linear.mmsinkhorn for multi-marginal coupling) that may change.
    • ott.math: Low-level mathematical primitives (e.g., Legendre transform).
    • ott.utils: Miscellaneous helper functions.

  2. What is OTT-JAX?

    main

    OTT-JAX is a JAX-powered library designed to solve optimal transport problems at scale using hardware accelerators.

    Key capabilities include:

    • Discrete Solvers: Implements the Sinkhorn algorithm for matching point clouds, including optimizations like scheduling, momentum, acceleration, and low-rank extensions.
    • Advanced Problems: Supports Gromov-Wasserstein and Wasserstein barycenter computations.
    • Neural Approaches: Provides implementations of neural network methods that approximate optimal transport maps between measures.

  3. Overview of the ott.tools package

    main

    The ott.tools package provides high-level functions built on top of lower-level components like ott.solvers. It is designed to provide user-friendly APIs for common Optimal Transport (OT) tasks, including:

    • Unregularized OT: Calculating quantities like the Wasserstein distance for point clouds of the same size.
    • Sinkhorn Divergence: Implementations of Sinkhorn divergence and segmented versions.
    • Sliced Wasserstein Distance: Efficient distance computations using random projections.
    • Soft Sorting: Differentiable approximations to ranks, quantiles, and sorting.
    • Clustering: K-means implementations.
    • Gaussian Mixture Tools: Specialized tools for manipulating and comparing Gaussian mixtures using modified Wasserstein geometry.
    • Visualization: Plotting and transport animation utilities.

  4. Use the ott.neural module for neural optimal transport

    main
    The ott.neural module is used to parameterize optimal transport maps and couplings as neural networks. Unlike standard solvers in ott.solvers that output static vectors or matrices, neural methods in ott.neural learn representations that can generalize to new samples, allowing them to be evaluated on data points outside of the original training set. This module provides the necessary layers, models, and solvers to estimate these neural networks.

  5. Use neural networks in ott.neural.networks

    main

    The ott.neural.networks module provides various neural network architectures and potential functions used for parameterized optimal transport. Key components include:

    • ICNN: Input Convex Neural Networks (icnn.ICNN) and icnn.KeyNet.
    • Potentials: Base potential classes (potentials.BasePotential), MLP-based potentials (potentials.PotentialMLP), and standard MLPs (potentials.MLP).
    • Potential Training: potentials.PotentialTrainState for managing the state of potential training.

  6. Use velocity field networks in ott.neural.networks.velocity_field

    main

    The ott.neural.networks.velocity_field module provides architectures specifically designed for modeling velocity fields in transport problems:

    • MLP: Standard Multi-Layer Perceptrons via mlp.MLP.
    • UNet: U-Net architectures via unet.UNet.
    • EMA (Exponential Moving Average): Tools for applying EMA to network parameters, including ema.EMA (the stateful object), ema.init_ema for initialization, and ema.update_ema for updating weights.

  7. Use ott.neural.data for neural optimal transport

    main
    The ott.neural.data module provides specialized data loaders designed for neural optimal transport tasks. These loaders facilitate the streaming or batching of data required to train neural networks for computing optimal transport plans or potentials.

  8. Use specialized layers in ott.neural.networks.layers

    main

    The ott.neural.networks.layers module contains specialized layers and solvers for neural optimal transport:

    • Conjugate Solvers: Layers for handling Fenchel conjugates, including conjugate.FenchelConjugateSolver and conjugate.FenchelConjugateLBFGS. Results are returned as conjugate.ConjugateResults.
    • Positive Definiteness: Layers to ensure positive definiteness, such as posdef.PositiveDense and posdef.PosDefPotentials.

  9. Solve linear optimal transport problems with ott.problems.linear

    main

    The ott.problems.linear module is used to define and solve the simplest family of optimal transport problems, specifically the Kantorovich problem (linear optimal transport). It also supports objective functions that are sums of optimal transport costs, such as the two variants of Wasserstein barycenter problems.

    Available problem types include:

    • LinearProblem: The standard Kantorovich optimal transport problem.
    • SemidiscreteLinearProblem: A semidiscrete variant of the linear problem.
    • FixedBarycenterProblem: A Wasserstein barycenter problem where the barycenter is fixed.
    • FreeBarycenterProblem: A Wasserstein barycenter problem where the barycenter is a variable to be optimized.

  10. Use DualPotentials for dual optimal transport optimization

    main
    The DualPotentials class in ott.problems.linear.potentials provides functions that act as optimization variables for the dual optimal transport problem. These are used when solving the dual formulation of an OT problem rather than the primal formulation.

  11. Calculate Sinkhorn Divergence

    main

    The Sinkhorn divergence is a proxy for the Wasserstein distance between two samples. It is computed by taking the output of an entropy-regularized optimal transport (EOT) or low-rank optimal transport problem and applying a specific renormalization to ensure it behaves like a distance.

    Formula

    $$\text{SD}(\mu, \nu):= \Delta(\mu, \nu) - \tfrac12 \left(\Delta(\mu, \mu) + \Delta(\nu, \nu)\right)$$

    where Δ is the output of the regularized OT solver.

  12. Compute Barycenters with Barycenter Solvers

    main

    Barycenter solvers allow you to find a central measure (the barycenter) among a set of distributions. OTT provides different types of barycenter solvers:

    • FreeWassersteinBarycenter: For continuous barycenter computation.
    • FixedBarycenter: For discrete barycenter computation.