Problem
Research questions and friction points this paper is trying to address.
Predicts Optimal Transport distances accurately
Handles discrete measures of variable resolution
Improves Sinkhorn algorithm initialization
Universal Neural Optimal Transport
🤖 AI Summary
To address the high computational cost and poor generalization across arbitrary grid resolutions in entropy-regularized optimal transport (OT) between discrete measures, this paper proposes the first resolution-invariant, cost-function-adaptive neural OT framework that directly predicts both the OT distance and the transport plan. Innovatively, it incorporates the Fourier Neural Operator (FNO) into OT modeling to ensure resolution independence. A novel adversarial self-supervised guidance loss is introduced, with theoretical proof that it converges to the true OT loss. Experiments demonstrate that the method achieves high accuracy in predicting OT distances and transport plans across diverse datasets, faithfully reconstructs the geometric structure of the Wasserstein space, and—when used as initialization for the Sinkhorn algorithm—significantly outperforms existing baselines.
📝 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 of variable resolution 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 varying sizes. The network is trained adversarially using a second, generating network and a self-supervised bootstrapping loss. We theoretically justify the use of FNOs, prove that our generator is universal, and that minimizing the bootstrapping loss provably minimizes the ground truth loss. Through extensive experiments, we show that our network not only accurately predicts optimal transport 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, significantly outperforming existing approaches.
Innovation
Methods, ideas, or system contributions that make the work stand out.
Neural Optimal Transport
Fourier Neural Operators
Adversarial Training Framework
🔎 Similar Papers
No similar papers found.
