Low-Rank Optimal Transport with FRLC in OTT-JAX

This repository provides a tutorial-style implementation in OTT-JAX of “Low-Rank Optimal Transport through Factor Relaxation with Latent Coupling” by Peter Halmos, Xinhao Liu, Julian Gold, and Benjamin J. Raphael.

The goal is to explain the key ideas behind Factor Relaxation with Latent Coupling (FRLC) and show how they can be implemented in an OTT-JAX-based workflow.

Official references:

• Paper: Low-Rank Optimal Transport through Factor Relaxation with Latent Coupling
• Official code: raphael-group/FRLC

Overview

FRLC solves for a low-rank optimal transport coupling
$\mathbf{P} \in \Pi_r(\mathbf{a}, \mathbf{b})$ that aligns a source distribution $\mathbf{a}$ with a target distribution $\mathbf{b}$ while minimizing the Wasserstein cost:

$$ \mathbf{P}^{*} = \arg\min_{\mathbf{P} \in \Pi_r(\mathbf{a}, \mathbf{b})} \langle \mathbf{C}, \mathbf{P} \rangle_F. $$

Instead of forming the full coupling matrix directly, FRLC returns a factorized representation:

$$ \mathbf{P} = \mathbf{Q} \text{diag}(1 / \mathbf{g}_Q) \mathbf{T} \text{diag}(1 / \mathbf{g}_R) \mathbf{R}^{T}. $$

Here, $\mathbf{Q}$ and $\mathbf{R}$ are sub-couplings, while $\mathbf{T}$ is a latent coupling between low-dimensional latent clusters. This can be interpreted as a coarsening of the full transport problem from observed data points to latent groups:

$$ (\mathbf{P}, \mathbf{a}, \mathbf{b}) \mapsto (\mathbf{T}, \mathbf{g}_Q, \mathbf{g}_R). $$

Unlike diagonal low-rank OT approaches, FRLC allows the latent coupling $\mathbf{T}$ to be non-diagonal and non-square, making it possible to represent richer hidden transport structures.

Experiments and usage

The repository contains two notebooks organized around an explanation of our implemantion and simulated experiments.

The experiments include:

• a comparison between FRLC and LOT on a dataset composed of $1000$ points sampled from a two-moons distribution and $1000$ points sampled from eight 2D Gaussian densities;
• an interpretation experiment for the latent coupling on a dataset composed of $1000$ samples from 2D Gaussians centered at the 5th roots of unity and $1000$ samples from 2D Gaussians centered at the 10th roots of unity.

FRLC_demo.ipynb

This notebook runs the implemented FRLC_solver/ directly and focuses on the experiments and visualizations.

FRLC_ott_jax_tuto.ipynb

This notebook provides a step-by-step OTT-JAX tutorial explaining the implementation of the FRLC solver and provide the experiments and visualizations.

Code structure

├── FRLC_demo.ipynb
├── FRLC_ott_jax_tuto.ipynb
├── FRLC_solver/
│   ├── __init__.py
│   ├── gradients_computation.py
│   ├── initialization.py
│   ├── plotting.py
│   ├── sinkhorn_projection.py
│   ├── solver.py
│   └── utils.py
├── LICENSE
├── pyproject.toml
├── README.md
└── uv.lock

FRLC_solver/

Contains the implementation of the FRLC solver used in the notebooks.

• solver.py: main FRLC solver.
• initialization.py: initialization of the sub-couplings.
• sinkhorn_projection.py: balanced and semi-relaxed Sinkhorn projection operators implemented with OTT-JAX.
• gradients_computation.py: gradient computations for the factor and latent coupling updates.
• utils.py: utility functions used by the solver.
• plotting.py: plotting and visualization helpers.
• __init__.py: package initialization file.

Installation

This project uses uv for environment and dependency management.

To install the dependencies, run:

uv sync

Then open the notebooks from the resulting environment.

If you do not use uv, you can inspect pyproject.toml and install the required dependencies manually.