ICML 2025

Universal Neural Optimal Transport
Jonathan Geuter 1 2 * Gregor Kornhardt 3 * Ingimar Tomasson 3 * Vaios Laschos 4
Abstract
Optimal Transport (OT) problems are a corner-
stone 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 net-
work to process measures of variable resolutions.
The network is trained adversarially using a sec-
ond, generating network and a self-supervised
bootstrapping loss. We ground UNOT in an ex-
tensive theoretical framework. Through experi-
ments 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×, significantly outper-
forming existing approaches.
1. Introduction
Optimal Transport (Villani, 2009; Peyr´e & Cuturi, 2019)
plays an increasing role in various areas in machine learn-
ing, such as domain adaptation (Courty et al., 2017), single-
cell genomics (Schiebinger et al., 2019), imitation learn-
ing (Dadashi et al., 2020), imaging (Schmitz et al., 2018),
dataset adaptation (Alvarez-Melis & Fusi, 2021), and signal
processing (Kolouri et al., 2017). Oftentimes, an entropic
*Equal contribution
1Harvard John A. Paulson School of En-
gineering and Applied Sciences 2Kempner Institute at Harvard
University 3Department of Mathematics, Technische Universit¨at
Berlin, Germany 4Weierstrass Institute, Berlin, Germany. Corre-
spondence to: Jonathan Geuter...