Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark (2106.01954v2)

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.

• The paper introduces a standardized benchmark using ICNNs and analytical OT maps to assess solver performance on quadratic-cost problems.
• It reveals that high Wasserstein-2 accuracy does not guarantee better generative outcomes, prompting a reevaluation of evaluation metrics.
• The research establishes a reproducible framework for continuous OT in high-dimensional settings, guiding future neural solver development.

Evaluation of Neural Optimal Transport Solvers: A Continuous Wasserstein-2 Benchmark

The paper "Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark" by Korotin et al. addresses a critical gap in the evaluation of neural network-based solvers for optimal transport (OT). Despite the increasing application of these solvers in machine learning, particularly in generative modeling and domain adaptation, a standardized methodology for assessing their performance has been lacking. Their contribution is a benchmark for evaluating solvers focused on quadratic-cost transport, specifically the Wasserstein-2 ($\mathbb{W}_2$) distance.

Methodology and Contributions

The authors provide a comprehensive evaluation framework using input-convex neural networks (ICNN) to construct benchmark pairs of measures with analytically-known optimal transport maps. This allows for direct calculation of the ground truth, facilitating quantitative comparisons of solver performance. Their benchmark addresses continuous distributions in high-dimensional spaces, such as image data, making it particularly relevant for modern deep learning applications.

The paper details several existing OT solvers, highlighting their distinct methodologies and optimization procedures. Among these are approaches that leverage dual forms of the OT problem, parameterizing functions to approximate optimal potentials. The solvers differ notably in how they handle the computational challenges inherent in estimating continuous OT maps, with some using mini-batch strategies while others incorporate input convexity constraints to aid in optimization.

Experimental Evaluation

Korotin et al. performed extensive experiments to evaluate the solvers on their benchmark sets, including Gaussian mixtures and CelebA64 image data. Their findings indicate that while some solvers exhibit moderate errors even in low-dimensional settings, those using parameterization techniques with ICNNs, which maintain input convexity, tend to perform best across dimensions. Surprisingly, they find that solvers achieving high $\mathbb{W}_{2}$ accuracy do not necessarily yield superior performance in generative modeling tasks, suggesting that these tasks may benefit from alternative metrics or regularization strategies.

Implications and Future Directions

The insights gained from this paper have significant implications for the development and application of OT solvers in machine learning. The revelation that higher OT accuracy does not directly translate to better downstream performance challenges the community to reassess the evaluation metrics used in solver development. It also opens opportunities for exploring other formulations that might better capture the intricacies of generative tasks.

Furthermore, this benchmark establishes a pathway towards more reproducible and transparent research in continuous OT, encouraging the exploration of new neural architectures and optimization methods that can overcome current limitations. The authors suggest that while ICNNs currently represent a strong approach, future research might uncover novel avenues for leveraging the structure and dynamics of high-dimensional data in OT settings.

Overall, this paper provides a valuable resource and set of tools for the ongoing exploration of neural networks in optimal transport, proposing a rigorous standard for future research and development in the field.