Optimal transport mapping via input convex neural networks (1908.10962v2)

Abstract: In this paper, we present a novel and principled approach to learn the optimal transport between two distributions, from samples. Guided by the optimal transport theory, we learn the optimal Kantorovich potential which induces the optimal transport map. This involves learning two convex functions, by solving a novel minimax optimization. Building upon recent advances in the field of input convex neural networks, we propose a new framework where the gradient of one convex function represents the optimal transport mapping. Numerical experiments confirm that we learn the optimal transport mapping. This approach ensures that the transport mapping we find is optimal independent of how we initialize the neural networks. Further, target distributions from a discontinuous support can be easily captured, as gradient of a convex function naturally models a {\em discontinuous} transport mapping.

• The paper introduces an ICNN-based minimax optimization framework that accurately estimates optimal transport maps.
• The approach guarantees theoretical consistency and stability by bridging optimal transport theory with practical neural network applications.
• Numerical experiments on synthetic and real-world data demonstrate the method's capability to capture complex, discontinuous transport mappings.

Optimal Transport Mapping via Input Convex Neural Networks

This academic work introduces an innovative approach to solving the problem of optimal transport between two probability distributions using machine learning techniques, particularly focusing on deep neural networks characterized by convexity properties. The authors present a novel method leveraging input convex neural networks (ICNNs) to estimate the optimal transport map, which is a pivotal concept in various applications like machine learning, domain adaptation, and more. Their method is grounded in optimal transport theory, proposing a principled framework that addresses some of the fundamental challenges in this domain.

Key Components and Contributions

The paper explores employing ICNNs, a class of neural networks ensuring the neural function is convex in its inputs. This architectural choice underpins the ability of the ICNNs to inherently satisfy certain mathematical constraints pivotal in optimal transport problems. Here's a succinct overview of the paper's primary contributions and methodologies:

1. Minimax Optimization Framework: The authors propose a minimax optimization problem where the aim is to find a pair of convex functions that estimate the optimal transport map. This is a deviation from regularization-based approaches, which often introduce a penalty term but consequently complicate optimization dynamics.
2. Theoretical Guarantees: Through rigorous mathematical formulations and proofs, the paper establishes the consistency and stability of the proposed minimax optimization approach. Theorems supporting these claims indicate that the learned map approaches the true optimal transport map as the solution converges.
3. Numerical Experiments: The paper discusses various numerical experiments that validate the accuracy and robustness of the learned transport map. These experiments, including both synthetic and real-world datasets, underscore the ability of the proposed ICNN framework to capture complex transport mappings with discontinuous characteristics. The experiments illustrate that this approach can effectively tackle high-dimensional problems.
4. Discontinuous Mappings: A significant advantage of their approach is the natural handling of discontinuous mappings. Neural networks generally struggle with representing discontinuities due to their inherent continuity but using the gradient of a convex neural network enables the model to approximate such functions more faithfully.
5. Implications for Deep Generative Models: By utilizing ICNNs, the authors outline a methodology that could be directly applied to train deep generative models like generative adversarial networks (GANs). When trained in this optimal transport framework, these models show improved stability and robustness compared to traditional approaches.

Implications and Future Directions

In synthesizing these results, the authors make several critical observations on the implications of their findings:

• Robust Learning with ICNNs: The work highlights how utilizing ICNNs offers robustness in learning optimal transport maps. By representing the transport map as a gradient of a convex function, the method can circumvent issues related to traditional neural architectures, particularly the sensitivity to initial conditions and difficulties in capturing discontinuous mappings.
• Scalability Challenges and Prospects: While the method shows promise, scaling it to larger datasets remains a challenge. Future work might investigate more sophisticated architectures or alternative optimization strategies to address these issues. Also, incorporating the concept of ICNNs into broader machine learning tasks holds potential for extensive exploration and practical applications.
• Theoretical and Computational Balance: The combination of theoretical solidification and computational application presented in this paper paves the way for a deeper understanding and development of machine learning models capable of leveraging the rich framework provided by optimal transport theory.

In conclusion, this paper makes significant strides in bridging theoretical concepts in optimal transport with implementable machine learning solutions. By situating their work within a broader context of neural network-based approaches to distribution mapping, the authors contribute meaningfully to both the theoretical understanding and practical implementation of optimal transport maps, encouraging further exploration and development in this promising intersection of mathematics and artificial intelligence.