Universal Neural Optimal Transport

TL;DR

UNOT introduces a neural network framework based on Fourier Neural Operators that accurately predicts optimal transport distances and plans across various domains, offering significant speedups and geometric insights.

Contribution

The paper presents UNOT, a universal neural network framework for OT that is discretization-invariant and capable of generalizing across different resolutions and domains.

Findings

UNOT accurately predicts OT distances and plans.

UNOT captures the geometry of Wasserstein space.

UNOT accelerates Sinkhorn algorithm by up to 7.4 times.

Abstract

Optimal Transport (OT) problems are a cornerstone of many applications, but solving them is computationally expensive. To address this problem, we propose UNOT (Universal Neural Optimal Transport), a novel framework capable of accurately predicting (entropic) OT distances and plans between discrete measures for a given cost function. UNOT builds on Fourier Neural Operators, a universal class of neural networks that map between function spaces and that are discretization-invariant, which enables our network to process measures of variable resolutions. The network is trained adversarially using a second, generating network and a self-supervised bootstrapping loss. We ground UNOT in an extensive theoretical framework. Through experiments on Euclidean and non-Euclidean domains, we show that our network not only accurately predicts OT distances and plans across a wide range of datasets, but also captures the geometry of the Wasserstein space correctly. Furthermore, we show that our network can be used as a state-of-the-art initialization for the Sinkhorn algorithm with speedups of up to $7.4\times$, significantly outperforming existing approaches.