[Paper Review] Optimal transport mapping via input convex neural networks

Ashok Vardhan Makkuva, Amirhossein Taghvaei+2 more|arXiv (Cornell University)|Aug 28, 2019

Stochastic Gradient Optimization Techniques44 references39 citations

TL;DR

The paper presents a minimax framework using input convex neural networks (ICNNs) to learn the optimal transport map between distributions under the 2-Wasserstein distance, representing the transport map as the gradient of a convex function.

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.

Motivation & Objective

• Motivate learning an optimal transport map between distributions from samples without regularization biases.
• Propose a minimax formulation that convexifies the dual problem and avoids constrained projections.
• Leverage ICNNs to parameterize convex functions and their conjugates to recover the transport map as a gradient.
• Establish consistency and stability of the learned map under the proposed framework.
• Demonstrate applicability to high-dimensional and real-world datasets as a deep generative modeling tool.

Proposed method

• Formulate the 2-Wasserstein dual problem and convexify constraints to obtain a minimax objective (Equation 5).
• Reparametrize using convex function f and its conjugate to express W2^2(P,Q) via a single convex function (Theorem 3.3).
• Parameterize convex f with input convex neural networks (ICNNs) to ensure convexity in the input.
• Represent the transport map as the gradient of a convex function (g's gradient in the minimax setup).
• Solve the resulting minimax problem with stochastic optimization (Adam), enforcing nonnegative weights in f’s ICNN and optional regularization in g for stability (Equation 9).
• Provide theoretical guarantees: consistency (Theorem 3.3) and stability bound (Theorem 3.6).

Experimental results

Research questions

• RQ1 Can we learn the optimal transport map T* from Q to P using samples without regularization that biases the primal problem?
• RQ2 Can a minimax formulation over convex functions yield the exact 2-Wasserstein transport map when parameterized by ICNNs?
• RQ3 Does representing the transport map as the gradient of a convex function enable discontinuous transport mappings and sharper boundaries between supports?
• RQ4 How does the proposed method perform in high-dimensional and real-world datasets compared to standard OT-based or GAN-based approaches?
• RQ5 Is the learned transport map robust to initialization and training dynamics, and can stability bounds be established?

Key findings

• The proposed minimax ICNN-based framework yields transport maps that visually and qualitatively align with optimal transport between simple and complex distributions.
• The method demonstrates robustness to initialization, unlike W1-LP and W2GAN baselines which show initialization sensitivity.
• The gradient-of-convex-function representation enables discontinuous transport mappings, enabling sharp boundaries between disconnected supports.
• Experiments on high-dimensional data (including Gaussian-to-Gaussian, Gaussian-to-mixture, and MNIST-based tasks) show the approach can scale to complex distributions with reasonable accuracy.
• Quantitative results on a 728-dimensional Gaussian transport (Table 1) report mean-transport errors that increase with target shift magnitude (alpha), with specific reported values for mean transport error and relative error.
• Theoretical contributions include a consistency result (Theorem 3.3) and a stability bound (Theorem 3.6) for the learned transport map.