Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark
An ICNN benchmark shows that tW2s best recovers high-dimensional W2 maps, yet OT accuracy does not guarantee better generation.
Key Findings
Methodology
The paper constructs continuous distributions with known quadratic-cost OT maps using input-convex neural networks (ICNNs). For a convex ψ, Brenier’s theorem makes ∇ψ the optimal map from P to Q=∇ψ#P. The authors build Gaussian-mixture and 64×64 CelebA benchmarks, then evaluate tLS, tMM, tMM-B, tQCs, tMMv1, tMMv2, and tW2s.
Key Results
- On the D=256 Gaussian-mixture benchmark, tW2s reaches 2.7% L2-UVP and cosine similarity 1.00; tMM-B reaches 22.5% and 0.93, the linear baseline 67.4% and 0.77, and tQCs 88.2% and 0.66.
- On CelebA64, tW2s obtains L2-UVP of 1.7%, 0.5%, and 0.25% on Early, Mid, and Late benchmarks, with cosine scores 0.99, 0.95, and 0.93. tMM-B yields 45.9%, 46.1%, and 47.7%.
- Accurate map recovery does not ensure superior generation. ICNN solvers recover OT maps faithfully, but the paper shows that downstream performance can favor biased methods such as tQCs, because generator gradients and full OT maps are different objectives.
Significance
This work turns the long-standing lack of ground truth for continuous OT into a reproducible benchmarking problem. It shows that reporting a GAN metric or an estimated Wasserstein value does not establish that a solver recovered the underlying transport structure. The finding matters for generative modeling, domain adaptation, and image translation, and provides a practical protocol for evaluating maps, gradients, and high-dimensional robustness separately.
Technical Contribution
The central construction combines Brenier’s theorem with ICNNs: Q=∇ψ#P is sampled from P, while the exact map T*=∇ψ remains known. Averaging convex potentials creates richer multimodal benchmarks. The paper distinguishes map error from the generator-relevant derivative ∇f*=id−T*, and evaluates both through L2-UVP and cosine similarity across dimensions 2–256 and CelebA64 images.
Novelty
Unlike prior benchmarks based on discrete, low-dimensional measures with limited support, this is a systematic continuous W2 framework that scales to image spaces and provides analytically known maps. Its deeper novelty is conceptual: it separates accurate distance estimation, accurate map recovery, accurate training gradients, and useful downstream behavior—four notions often conflated in neural OT research.
Limitations
- Some benchmark potentials are first fitted with tW2s, so the protocol may favor ICNN-based methods; approximate construction error remains unless ψ is analytically specified.
- The study focuses on continuous dual solvers rather than exhaustively covering primal, diffusion, or sampling-based approaches. Around 100 GPU-hours also make broad reproduction expensive.
Future Work
Future work should include non-ICNN and real-data benchmarks, primal solvers, and calibrated correction for regularization, minibatch bias, and maximin instability. A unified protocol should jointly predict map fidelity, gradient quality, optimization cost, and downstream utility rather than ranking solvers by one metric.
AI Executive Summary
Optimal transport is increasingly embedded in generative modeling, domain adaptation, and image translation, yet neural solvers for continuous distributions have rarely been tested against known answers. Existing studies often report FID or an estimated Wasserstein value, measurements that can hide errors in the OT component itself. This paper asks a basic but consequential question: do neural W2 solvers actually recover optimal transport maps?
Its answer comes from input-convex neural networks. For an absolutely continuous source P and convex potential ψ, Brenier’s theorem guarantees that ∇ψ is the true quadratic-cost optimal map to Q=∇ψ#P. The authors use this fact to build Gaussian-mixture benchmarks from dimension 2 to 256 and continuous 64×64 CelebA image benchmarks. They compare tLS, tMM, tMM-B, tQCs, tMMv1, tMMv2, and tW2s using map error, L2-UVP, and gradient cosine similarity.
The results expose a sharp high-dimensional divide. At D=256, tW2s achieves 2.7% L2-UVP and cosine 1.00, whereas tQCs reaches 88.2% and 0.66, and tMM-B 22.5% and 0.93. On CelebA, tW2s remains highly accurate, with UVP between 0.25% and 1.7%. Yet map accuracy does not automatically produce the best image generation. The broader lesson is methodological: W2 estimation, map recovery, generator gradients, and downstream quality must be evaluated as related but distinct objectives.
Deep Analysis
Background
Continuous OT replaces discretization with neural potentials, kernel expansions, or learned maps, enabling large-scale applications such as WGANs and domain adaptation. W2 offers geometric structure and stronger theory than many alternatives, but existing solvers are usually tested on self-generated examples or indirectly through GAN performance. The discrete benchmark of prior work does not extend to continuous high-dimensional measures.
Core Problem
The paper separates three tasks: evaluating W2², recovering the optimal map or plan, and estimating ∇W2² for generator updates. A potential can be numerically close while its gradient is inaccurate, a failure called gradient deviation. Regularization, minibatch restrictions, and approximate inner minimization introduce additional bias, especially as dimension grows.
Innovation
First, ICNNs and Brenier’s theorem provide continuous benchmarks with known maps. Second, averaging convex potentials creates multimodal, high-dimensional targets. Third, the construction extends from Gaussian mixtures to CelebA64 images. Fourth, L2-UVP and cosine similarity separately measure map and gradient quality, while downstream generation tests whether direct OT accuracy transfers to applications.
Methodology
- �� Define W2²(P,Q)=minπ∫||x−y||²dπ and use T*=∇ψ for convex ψ.
- �� Sample a three-component Gaussian mixture P and fit ICNN maps ∇ψ1,∇ψ2 toward two ten-component targets using tW2s.
- �� Form Q=(1/2)(∇ψ1+∇ψ2)#P and repeat for D=2,4,8,16,32,64,128,256.
- �� For CelebA64, use WGAN-QC checkpoints to obtain Early, Mid, and Late continuous measures, then construct targets with ConvICNN64.
- �� Compare tLS, tMM, tMM-B, tQCs, tMMv1, tMMv2, tW2s, plus identity, constant, and linear baselines.
- �� Evaluate 214 source samples using L2-UVP and gradient cosine similarity, then train a 128-dimensional latent CelebA generator with each solver.
Experiments
The high-dimensional benchmark uses random Gaussian mixtures; the image benchmark uses aligned CelebA64 faces. The constant map has UVP=100%, while the linear baseline is the closed-form OT map after Gaussian approximation. DenseICNN, ConvICNN64, ResNet, and U-Net architectures are used. Experiments run on four GTX 1080Ti GPUs and require roughly 100 GPU-hours.
Results
At D=2, most methods approach the truth, but errors diverge with dimension. At D=256, tW2s obtains 2.7% UVP and cosine 1.00; tMM-B 22.5% and 0.93; tLS 54.7% and 0.81; tQCs 88.2% and 0.66; identity reaches 153% UVP. On CelebA, tMM-B and tQCs are severely biased, while tW2s, tMM, and reversed tMM recover visually and quantitatively accurate maps.
Applications
The benchmark can select W2 solvers for image translation, domain adaptation, and WGAN-style training, and can diagnose whether a discriminator supplies a meaningful generator direction. Before deployment, practitioners should report UVP, cosine similarity, convergence, and FID on matched data modalities rather than relying on one downstream score.
Limitations & Outlook
Some ψ functions are approximated with tW2s, potentially favoring ICNN methods, and Gaussian noise is added to image outputs to ensure absolute continuity. The study does not cover every primal solver. Maximin methods can diverge under imperfect hyperparameters, while tMMv1 is computationally expensive because of its inner convex optimization. Thus the results reveal structural failure modes, not a universal final ranking.
Plain Language Accessible to non-experts
Imagine a moving company. Distribution P is a warehouse full of packages, and Q is the arrangement required in a new warehouse. Optimal transport asks the movers to deliver every package while minimizing total travel distance. Neural OT solvers are teams of automated movers: they may create a new warehouse that looks roughly correct, even if many packages took unnecessarily long routes.
The authors build a warehouse with an answer key. An ICNN defines a special moving map whose gradient is mathematically guaranteed to be the shortest valid route. Because the destination is generated from that map, the researchers know exactly where every package should go. They can therefore compare routes directly instead of merely checking whether the final warehouse looks similar.
In low dimensions, most teams perform well. In high-dimensional image-like spaces, some fail badly: at dimension 256, tW2s has only 2.7% route error, while tQCs has 88.2%. But another surprise appears: the team with the most accurate routes is not always best for a later automated production task. The production system needs the correct direction for its next adjustment, not necessarily a perfect record of every route. OT evaluation must therefore inspect routes, directions, and final outcomes separately.
ELI14 Explained like you're 14
Think of a game where blurry faces must be upgraded into clear faces. An OT solver tells the generator how features should move: eyes, hair, lighting, and other details. A message can sound reasonable while still pointing in the wrong direction, causing the generator to improve slowly or produce odd results.
The authors create a challenge level with a known answer. ICNNs produce a special map that mathematics guarantees is the shortest way to rearrange one collection into another. The researchers then test several neural “moving robots,” including tW2s, tMM, tQCs, and tMM-B, on Gaussian mixtures and CelebA64 face images.
The score gap is huge. In 256 dimensions, tW2s has only 2.7% map error, while tQCs has 88.2%. On the face benchmark, tW2s again follows the correct map closely, whereas some other methods nearly lose the route entirely.
But here comes the twist: the best map does not always create the best pictures. Training mainly needs the right direction for the next update, not just a complete travel plan. It is like basketball: knowing where the ball should end up does not automatically tell you how hard to push it right now. The paper’s lesson is simple—check the map, the direction, and the final score!
Glossary
Optimal transport
A framework for moving one probability distribution into another at minimum cost. Here the cost is squared Euclidean distance.
Defines the maps, plans, and W2 objective studied throughout the paper.
Wasserstein-2 distance
A distributional distance induced by quadratic transport cost, denoted W2²(P,Q). It has both primal transport-plan and dual-potential formulations.
The central quantity and training loss evaluated by the solvers.
Input-convex neural network (ICNN)
A neural network constrained to be convex in its input and therefore suitable for representing convex potentials. Its gradient naturally has the structure required by Brenier transport.
Used to create ground-truth benchmarks and in tW2s and related solvers.
Brenier map
For an absolutely continuous source under quadratic cost, the optimal transport map is the gradient of a convex function, T*=∇ψ. This converts a hard unknown map into a known construction.
Provides the theoretical basis of the benchmark generator.
L2-UVP
A variance-normalized percentage measuring squared error between an estimated and true transport map. Values near 0 are accurate; 100% corresponds to the constant baseline.
Primary direct metric for map recovery.
Gradient deviation
A small error in a potential does not guarantee a small error in its gradient, because optimization may never constrain derivatives directly.
Explains why potential-based solvers can have inaccurate maps or generator gradients.
Open Questions Unanswered questions from this research
- 1 How to build continuous W2 benchmarks on realistic image manifolds without relying on any candidate solver remains unresolved.
- 2 The relationship among map fidelity, gradient cosine similarity, and FID is not stable; a general theory explaining their mismatch is still needed.
- 3 It remains unclear whether adaptive optimization or unbiased minibatch estimators can systematically prevent maximin divergence and high-dimensional bias.
Applications
Immediate Applications
W2 solver validation
Research teams can run Gaussian-mixture and CelebA64 tests before deployment, reporting L2-UVP, cosine similarity, and convergence. This exposes high-dimensional bias that FID alone can conceal.
Image translation and domain adaptation
Practitioners can inspect ICNN-style transport maps for structural consistency and direction quality when transferring styles or domains. Adequate samples and substantial training compute are required.
Long-term Vision
Auditable generative systems
Future generators could report three audit layers—map fidelity, training-gradient quality, and downstream utility—creating reproducible distribution-matching standards and reducing black-box claims.
Abstract
Despite the recent popularity of neural network-based solvers for optimal transport (OT), there is no standard quantitative way to evaluate their performance. In this paper, we address this issue for quadratic-cost transport -- specifically, computation of the Wasserstein-2 distance, a commonly-used formulation of optimal transport in machine learning. To overcome the challenge of computing ground truth transport maps between continuous measures needed to assess these solvers, we use input-convex neural networks (ICNN) to construct pairs of measures whose ground truth OT maps can be obtained analytically. This strategy yields pairs of continuous benchmark measures in high-dimensional spaces such as spaces of images. We thoroughly evaluate existing optimal transport solvers using these benchmark measures. Even though these solvers perform well in downstream tasks, many do not faithfully recover optimal transport maps. To investigate the cause of this discrepancy, we further test the solvers in a setting of image generation. Our study reveals crucial limitations of existing solvers and shows that increased OT accuracy does not necessarily correlate to better results downstream.