Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark

Primal formulation

• A transport map may fail to move P to Q when mass must split, whereas a transport plan allows such splitting through a measure with marginals P and Q.When the optimal plan has the form induced by a map, that map also minimizes the Monge formulation.Source: §1 · p. 2

Dual formulation

• The optimal dual potential f* recovers the optimal map as T*(x) = x − ∇f*(x).The section also links c-concavity of optimal potentials to convex functions and Brenier’s theorem [6].Source: §1 · pp. 2–3

Solving optimal transport problems

• OT solvers address three distinct tasks: evaluating W2, computing an optimal map or plan, and estimating the W2 gradient used to update generative models.The gradient task arises when a generated measure Pα is compared with data measure Q.Source: §1 · p. 3

Quantitative evaluation of OT solvers

• Continuous OT evaluation is difficult because the existing discrete benchmark does not extend to continuous OT, and solvers are often tested on self-generated examples or through downstream models.Direct map metrics are scarce because few continuous measure pairs have known ground-truth maps.Source: §1 · p. 3
• Downstream metrics such as GAN quality do not isolate solver quality because they also depend on model components unrelated to OT.The section gives Fréchet Inception distance as an example of an indirect metric.Source: §1 · p. 3

Evaluation scope

• The paper quantitatively evaluates dual-form continuous solvers, while noting that primal-form solvers typically require more complex optimization and careful hyperparameter search.The benchmark could test other continuous solvers that compute a map or potential gradient.Source: §2 · p. 4

Continuous dual solvers

• Existing continuous dual solvers fit parametric potentials whose gradients yield approximate OT maps and the derivatives used to update generative models.The paper summarizes these solvers in Table 1.Source: §2 · p. 4, Table 1 · p. 4

Inner minimization

• The common computational challenge is the inner minimization used to evaluate f^c(y); tMM-B restricts the search to the current P minibatch, overestimating the minimum.This restricted search makes the approach fast, but the minimum is taken over only a subset of the full domain.Source: §2 · p. 4

ICNN and minimax solvers

• tMM-v1 uses an ICNN potential to make the inner problem convex, but solving it to high precision is computationally costly.tMMs instead uses neural networks for the potential and inner minimizer, trained by stochastic gradient ascent/descent.Source: §2 · pp. 4–5

Gradient deviation

• Potential fit does not guarantee gradient fit: even small L2 potential error can coexist with arbitrarily large gradient error, an issue formally addressed only for three ICNN-based solvers.The named solvers are tMM-v1, tMM-v2, and tW2.Source: §2 · p. 5

Reversed solvers

• tMM, tMM-v2, and tW2 can recover forward and inverse maps; swapping P and Q during training defines reversed variants, and reversed tMM performs better in generative modeling.The solvers are asymmetric in P and Q.Source: §2 · p. 5

Key idea

• For any continuous P with sample access and known differentiable convex ψ, the pair (P, ∇ψ#P) has the analytically known OT map ∇ψ.Samples from the second measure are generated by pushing samples from P through ∇ψ.Source: §3 · p. 5

Arbitrary pairs

• For arbitrary P and Q, the method approximates a map with tW2 and benchmarks the modified pair (P, ∇ψθ#P), which may slightly favor ICNN-based methods.Other solvers could supply the approximation if ψθ is convex.Source: §3 · p. 5

Extensions

• Combining known convex functions with a convex monotone function can create further convex potentials and therefore additional benchmark pairs.The passage gives nonnegative scaling and summation as examples.Source: §3 · p. 5

Inversion

• When ∇ψθ is bijective, the inverse transport map exists and can be computed by solving a convex problem.All ICNN gradients used in the paper are stated to be bijective.Source: §3 · p. 5

4 Benchmark Details and Results

• The authors implement their benchmark in PyTorch, provide pre-trained transport maps for all benchmark pairs, and make the code publicly available at the stated GitHub repository.Experiments used 4 GTX 1080ti GPUs and required about 100 hours of computation per GPU.Source: §4 · p. 5

4.1 Datasets

• The high-dimensional benchmarks use a three-Gaussian source and two ten-Gaussian targets, averaging two fitted ICNN transport potentials to form each output measure.The study also builds CelebA64 pairs from generative-model checkpoints, adding Gaussian noise to make measures absolutely continuous.Source: §4.1 · p. 6

4.2 Metrics and Baselines

• L2-UVP measures map error relative to target variance, with values near 0% indicating a close approximation and values at least 100% indicating a far-from-optimal map.The constant baseline has L2-UVP of 100%.Source: §4.2 · p. 6

4.3 Evaluation of Solvers on High-dimensional Benchmark Pairs

• In higher dimensions, only tMMv1, tMM, tMMv2, tW2 and their reversed versions produce reasonable results; several maximin solvers are difficult to optimize.All solvers perform well at D = 2, while some diverge during optimization at higher dimensions.Source: §4.3 · p. 7, Table 2 · p. 8
• As dimension increases, tQC, tLS, and tMM-B have sharply rising L2-UVP; tMM-B only notably outperforms the linear baseline, while tQC and tLS retain high cosine similarity.The reported explanations include biased inner-problem estimation for tMM-B, regularization bias for tLS, and mini-batch discrete OT bias for tQC.Source: §4.3 · p. 7, Table 2 · p. 8

4.4 Evaluation of Solvers in CelebA 64 ˆ 64 Images Benchmark Pairs

• On CelebA image benchmarks, tQCs and tMM-Bs show extreme bias and nearly orthogonal derivative estimates, whereas tMMs, reversed tMM, and tW2 recover maps well.The tMMs map is slightly noisier than reversed tMM's, illustrating gradient deviation.Source: §4.4 · p. 8, Table 3 · p. 9

4.5 Evaluation of Solvers in Generative Modeling of CelebA 64 ˆ 64 Faces

• tQCs generates well despite recovering a map far from identity, while tMMs produces blurry images linked to gradient deviation; reversed tMM improves map quality but remains blurry.The fitted map from reversed tMM produces images comparable to tQCs.Source: §4.5 · pp. 9–10

5 Conclusion

• Bad OT solvers can yield good generative performance, and simply reversing asymmetric solvers can affect performance dramatically.The benchmark is intended to support evaluation of continuous OT solvers; a generic method for producing new benchmark pairs is also described.Source: §5 · p. 10
• The benchmark relies on ICNN-gradient maps and absolutely continuous measures, and the representational power of those maps may be insufficient.The authors also say it remains unclear whether analogous constructions work for other costs such as W1.Source: §5 · p. 10
• The benchmark could narrow future solver evaluation to its datasets, so the paper also describes a generic method for producing new pairs.The authors hope the benchmark improves the quality and reusability of OT-related research.Source: §5 · p. 10

A Benchmark Pairs Details

• Appendix A.1 discusses high-dimensional benchmark pairs, and Appendix A.2 is devoted to CelebA 64 × 64 image benchmark pairs.Source: A Benchmark Pairs Details · p. 14

A.1 High-dimensional Benchmark Pairs

• In each dimension, the procedure fixes random Gaussian mixtures P, Q1, and Q2 to create a benchmark.The random seeds are hard-coded in the code.Source: §A.1 · p. 14
• The mixture-generation procedure selects random component means from an M-dimensional grid and normalizes the final mixture to axis-wise variance 1.The component means are selected so that no pair shares a coordinate; the component covariance matrices are initialized from random row directions.Source: §A.1 · p. 14
• The resulting Gaussian mixtures have the same marginals across dimensions, as illustrated by a projection onto the first two dimensions.Source: §A.1 · p. 14, Figure 6 · p. 14

A.2 CelebA 64 × 64 Images Benchmark Pairs

• Three generative models are trained on aligned CelebA64 faces using a 128-dimensional latent Gaussian and WGAN-QC with a ResNet generator.For the first two trials, checkpoints at 1,000, 5,000, and 10,000 iterations produce measures; the third uses its 50,000-iteration final checkpoint.Source: §A.2 · p. 14
• White Normal noise with axis-wise σ = 0.01 is added to generator outputs to make each measure absolutely continuous.Source: §A.2 · p. 14
• The generated measures are used to construct image benchmark pairs through the pipeline described in M4.1.The pipeline is visualized in Figure 2.Source: §A.2 · p. 14

• Appendix B.1 discusses experimental neural network architectures, while Appendix B.2 provides the other training hyperparameters.Source: B Experimental Details · p. 15

B.1 Neural Network Architectures

• Networks are pretrained so ∇ψθ(x) ≈ x and Hω(y) ≈ y, a strategy reported to make optimization more stable.Source: §B.1 · p. 15
• DenseICNN uses input-quadratic skip connections and can be made input-convex by constraining selected weights to be non-negative.The experiments impose this constraint only for tMMv1, tMMv2, and tW2 solvers; DenseICNN (U) has unrestricted weights.Source: §B.1 · p. 15
• For image experiments, tMMv2 and tW2 use ConvICNN64 for both potentials, with an added quadratic term to make its gradient bijective.The added term uses β = 10^-4.Source: §B.1 · p. 15, Figure 7 · p. 16
• Conv2D-CQ blocks output tensors of input-quadratic functions, while PosConv2D layers constrain weights other than biases to be non-negative.GroupChannelSumPool splits channels into groups and sums each group's channel maps.Source: §B.1 · pp. 15–16, Figure 8 · p. 16

B.2 Hyperparameters and Implementation Details

• The evaluation required implementing most solvers from scratch and extracting discriminator components for methods otherwise used only in GANs.The authors used Adam with default hyperparameters except the learning rate, and official code for tQCs and tW2.Source: §B.2 · p. 16

B.2.1 High-dimensional Benchmark Pairs

• For the high-dimensional benchmark, tQCs uses batch size 64 because its per-step linear-program solve makes larger batches computationally infeasible.tW2 uses the same hyperparameters for benchmark training and evaluation.Source: §B.2.1 · p. 16

B.2.2 CelebA 64 × 64 Images Benchmark Pairs

• The CelebA image benchmark's solver hyperparameters are listed in Table 6.Source: §B.2.2 · p. 16, Table 6 · p. 17

B.2.3 CelebA 64 × 64 Images Generation Experiment

• Image generation alternates generator updates with OT-solver updates; tQCs uses one solver update per generator update, while other methods use ten.The generator learning rate is 3 × 10^-4, with 50,000 generator iterations.Source: §B.2.3 · p. 17
• For the generation experiment, the generator gradient uses the current OT potential, while tMM:R parameterizes the forward OT map instead.Source: §B.2.3 · p. 17