Research map / Learning theory and algorithms / Optimal transport
Universal Neural Optimal Transport
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes UNOT, a Fourier-Neural-Operator-based solver that predicts entropic OT dual potentials for discrete measures of variable resolution, trained adversarially with a generator and a Sinkhorn-bootstrapped loss, and used as a fast Sinkhorn initialization.
Problem
Entropic optimal transport is computationally expensive to solve repeatedly, and prior neural OT solvers either train per OT instance or are limited to fixed-dimension measures and datasets.
Approach
Use discretization-invariant Fourier Neural Operators to predict the entropic dual potential for any pair of discrete measures (given a fixed cost), trained self-supervised against Sinkhorn-refined targets, with a provably universal generator (a residual ReLU network) producing training measures of varying resolution.
Key results
- Sinkhorn initialized with UNOT reaches 0.01 relative OT-distance error in about 3-7 iterations versus 16-80 with ones initialization (MNIST-LFW table).
- Speedups over default initialization up to 7.4x (CIFAR, squared Euclidean cost), with average 3.57x on 28x28 and 4.4x on 64x64 datasets.
- UNOT predicts OT distances with about 1-3% relative error and significantly outperforms Meta OT on all datasets except MNIST, where it almost matches it despite not training on MNIST.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Initialize Sinkhorn with UNOT instead of ones | Sinkhorn · MNIST | Sinkhorn iterations to 0.01 relative OT-distance error: 16±9 -> 3±5 iterations; speedup 1.25 | Helped |
| Initialize Sinkhorn with UNOT instead of ones | Sinkhorn · CIFAR | Sinkhorn iterations to 0.01 relative OT-distance error: 80±22 -> 3±6 iterations; speedup 7.4 | Helped |
| Initialize Sinkhorn with UNOT instead of ones | Sinkhorn · CIFAR-MNIST | Sinkhorn iterations to 0.01 relative OT-distance error: 32±15 -> 4±4 iterations; speedup 2.07 | Helped |
| Initialize Sinkhorn with UNOT instead of ones | Sinkhorn · LFW | Sinkhorn iterations to 0.01 relative OT-distance error: 78±20 -> 7±8 iterations; speedup 5 | Helped |
| Initialize Sinkhorn with UNOT instead of ones | Sinkhorn · Bear | Sinkhorn iterations to 0.01 relative OT-distance error: 41±16 -> 4±6 iterations; speedup 3.8 | Helped |
| Initialize Sinkhorn with UNOT instead of ones | Sinkhorn · LFW-Bear | Sinkhorn iterations to 0.01 relative OT-distance error: 53±18 -> 4±6 iterations; speedup 4.4 | Helped |
Surprising findings
- UNOT, without training on MNIST, almost matches Meta OT, which was explicitly trained on MNIST.
- Barycenters can be far apart even when transport distances are arbitrarily close, so accurate distances do not imply accurate barycenters.
Limitations
- Does not extrapolate well to measures with much higher resolution than training samples.
- Does not generalize to cost functions other than the training cost; a separate model is needed per cost.
- Restricted to grid-supported measures; non-uniform grids not handled.
Experimental setup
Datasets: MNIST CIFAR10 (grayscale) Google Quick, Draw! teddy bear Labeled Faces in the Wild (LFW) synthetic generator-produced measure pairs
Benchmarks: MNIST CIFAR10 (grayscale) CIFAR-MNIST Quick, Draw! teddy bear LFW LFW-Bear
Metrics: relative error on OT distance number of Sinkhorn iterations to 0.01 relative error wall-clock speedup barycenter quality geodesic approximation quality
Baselines: Sinkhorn with ones initialization Gaussian initialization Meta OT GeONet
Builds on: Fourier Neural Operator Sinkhorn algorithm GAN Spherical FNO Meta OT
Closest papers by meaning
- Neural Optimal Transport with General Cost Functionals (ICLR 2024 · similarity 0.87)
- Expectile Regularization for Fast and Accurate Training of Neural Optimal Transport (NeurIPS 2024 · similarity 0.87)
- Improving Neural Optimal Transport via Displacement Interpolation (ICLR 2025 · similarity 0.87)
- Progressive Entropic Optimal Transport Solvers (NeurIPS 2024 · similarity 0.87)
- Analyzing and Improving Optimal-Transport-based Adversarial Networks (ICLR 2024 · similarity 0.87)
- Geometry-Aware Optimal Transport: Fast Intrinsic Dimension and Wasserstein Distance Estimation (NeurIPS 2026 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.