PaperlayerView exact revision on arXiv

Source-linked AI summary

Optimal transport mapping via input convex neural networks

Ashok Vardhan Makkuva, Amirhossein Taghvaei, Sewoong Oh, Jason D. Lee

arXiv:1908.10962v2cs.LGstat.ML

TL;DR

The paper addresses the challenge of learning optimal transport maps from samples without relying on regularization or discrete approximations. It uses a minimax optimization over input convex neural networks, with a convex-function gradient representing the map. The authors report exact optimal-map recovery under exact optimization, initialization-independent learning, and representation of discontinuous mappings.

  • Problem

    Learning optimal transport maps from samples is difficult because high-dimensional methods can be intractable and regularized dual formulations may bias the learned transport.

  • Method

    The method restricts dual functions to input convex neural networks and solves a novel minimax optimization whose transport map is the gradient of a convex function.

  • Results

    Theorem 3.3 states that exact solution of the proposed objective yields the exact optimal transport map, while experiments qualitatively recover optimal maps on synthetic datasets.

  • Takeaways & Limitations

    The framework can learn maps independently of neural-network initialization and represent discontinuous maps for target distributions with disconnected supports.

  • Takeaways & Limitations

    Further performance improvements in high-dimensional experiments require careful hyperparameter tuning, which the authors identify as ongoing work.

Abstract

from arXiv · show

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.

1. Introduction

Optimal transport provides a cost-minimizing inductive bias for mapping distributions, but existing approaches face scalability and constraint-enforcement challenges. The paper addresses these challenges with a convex-neural-network minimax framework that learns optimal, initialization-independent maps and can represent discontinuities.

  • Optimal transport seeks mappings that minimize mass-transport cost, but finding such maps is challenging for high-dimensional distributions.
  • Quantization-based approaches for general distributions become intractable in the high-dimensional settings encountered in modern applications.
  • Sampled Kantorovich constraints can be imposed as regularizers, but the resulting biases hinder learning the true optimal transport.
  • Restricting the search to convex functions eliminates the Kantorovich constraints and yields a novel minimax formulation.
  • The proposed training framework is reported to recover the optimal map independently of initialization and represent sharp boundaries for disconnected target supports.
  • Input convex neural networks support a scalable framework in which the gradient of a convex function represents the transport map.

2. Background on optimal transport

The paper reviews Monge and Kantorovich formulations of quadratic-cost optimal transport, including the existence challenges of deterministic maps. Regularized dual methods offer tractability but alter the original constrained problem, motivating a non-regularized alternative.

  • Monge’s problem transports Q to P using a map that minimizes quadratic cost, when such a deterministic map exists.
  • An optimal transport map may not exist; for example, no map can transport a Dirac distribution to a non-Dirac distribution.
  • Kantorovich’s relaxation optimizes over couplings with prescribed marginals, and its optimal value is the squared 2-Wasserstein distance.
  • The Kantorovich dual formulation uses function constraints that are difficult to maintain during stochastic gradient updates.
  • Entropic or quadratic regularization makes the dual unconstrained and enables numerical optimization, but changes the original optimal-transport problem.
  • The paper instead solves the dual without regularization using a minimax formulation over input convex neural networks.

3. A novel minimax formulation to learn optimal transport

The paper replaces the difficult constrained Kantorovich dual with a convex-function minimax formulation, then parameterizes the functions with ICNNs to learn the optimal transport map from samples.

  • 3. A novel minimax formulation to learn optimal transport: The standard Kantorovich dual is difficult because its constraints require function evaluations at every point in the domain.Sampling constraints as regularizers can introduce biases that hinder learning the true optimal transport.
  • 3. A novel minimax formulation to learn optimal transport: The alternative formulation eliminates the distance constraints by restricting the search to convex functions.This convexification enables integration with neural architectures that guarantee convexity.
  • 3. A novel minimax formulation to learn optimal transport: The formulation relies on Brenier’s theorem: when Q has a density, the optimal coupling is unique and the gradient of a convex conjugate solves the Monge problem.The paper therefore assumes throughout that Q admits a density in R^d.
  • 3. A novel minimax formulation to learn optimal transport: For the functional V_P,Q(f,g), an optimal pair (f_0,g_0) exists whose gradient ∇g_0 is the optimal transport map from Q to P.The functional is defined using expectations under P and Q and the term ⟨Y,∇g(Y)⟩−f(∇g(Y)).
  • 3. A novel minimax formulation to learn optimal transport: The method learns the optimal transport map ∇g from Q to P by solving a minimax optimization over convex functions.Under the paper’s density assumption on Q, the optimal pair includes a function whose gradient is the optimal map.
  • 3.1. Minimax optimization over ICNNs: ICNNs approximate convex functions with scalar-valued neural networks, allowing the proposed minimax optimization to be estimated from samples.The paper uses ICNN parameterizations for the convex functions and proposes an associated training algorithm.
  • 3.1. Minimax optimization over ICNNs: Unlike semi-dual optimization that repeatedly solves inner problems for f*, the proposed minimax formulation represents the conjugate with an ICNN for scalable learning.The earlier approach is described as slow and memory intensive on large datasets.
  • 3.2. Stability analysis of the learned transport map: Theorem 3.3 gives exact recovery when the minimax objective is solved exactly, while the stability analysis bounds approximation error using ε_1 and ε_2.The paper also presents a training framework in which the generator is modeled as a gradient of a convex function.

4. Experiments

The experiments train ICNN-based convex potentials with a minimax procedure and evaluate optimal transport on synthetic, high-dimensional, and MNIST-related distributions. The learned maps qualitatively capture optimal couplings, remain independent of initialization, represent discontinuities, and scale to high-dimensional settings, although further performance improvements require hyperparameter tuning.

  • Training methodology: The proposed ICNN minimax procedure learns transport maps by parameterizing convex functions and optimizing their empirical objective with Adam.Convexity is enforced strictly for f, while g uses a data-independent regularizer that empirically accelerates convergence.
  • Two-dimensional experiments: On Checkerboard and eight-Gaussian data, the method qualitatively learns optimal maps that symmetrically split continuous mass toward the nearest target support.The experiments compare against Barycentric-OT, W1-LP, and W2GAN.
  • Robustness: Unlike W1-LP and W2GAN, the learned transport map is independent of random initialization rather than selecting arbitrary maps.The comparison uses different random trials on the Checkerboard example.
  • Discontinuous transport maps: The gradient-based representation naturally models discontinuous transport maps, avoiding spurious probability mass connecting disconnected target supports.Continuous neural-network transport maps produce trailing dots or other connecting mass in the compared approaches.
  • High-dimensional experiments: In high-dimensional experiments, the method maps a 784-dimensional Gaussian to a four-component low-dimensional mixture and maps remaining components to zero.For MNIST embeddings, visually similar digits are coupled, including 1 →9, 2 →8, and 4 →9.
  • High-dimensional experiments: The experiments provide proof of concept for scaling to high-dimensional and real-world data, while further performance improvements require careful hyperparameter tuning.The authors describe hyperparameter tuning as ongoing work.

5. Conclusion

The paper presents a minimax framework for learning optimal transport under the W2 metric, replacing regularized dual constraints with ICNN-based convex representations and gradient transport maps.

  • The framework learns optimal transport under the W2 metric through a novel minimax formulation.
  • Unlike regularization-based approaches, ICNN representations automatically satisfy the dual Kantorovich constraint.
  • Expressing the transport map as the gradient of a convex function allows discontinuous maps to be represented.

A. Proof of Theorem 3.3

The proof establishes the key variational identity by exploiting convexity and almost-everywhere differentiability under the source distribution.

  • The proof defines V_f(g) as an expectation involving ∇g and shows that sup_g∈CVX(Q) V_f(g) = E_Q[f*(Y)].
  • The argument uses differentiability of convex functions g and f* at Q-almost every point.
  • The inequality is attained by choosing g = f*.
  • Finite expectation implies that the effective domain has full Q-measure, while convexity yields differentiability except on a Q-null set.

B. Proof of Theorem 3.6

The proof controls the estimator’s optimality gaps using bounds derived from strong convexity, convex conjugacy, and the triangle inequality.

  • The stability proof begins from bounds and applies the triangle inequality.
  • Strong convexity of f is used to bound properties of its convex conjugate f*.
  • The resulting inequality controls the optimality gap ϵ1(f, g).
  • The proof separately analyzes the second optimality gap ϵ2(f).
  • The argument concludes the second bound using expressions involving f, f*, and their gradients.

C.1. Two-dimensional experiments

The two-dimensional experiments use synthetic Checkerboard and eight-Gaussian datasets, ICNN parameterizations, established baselines, and detailed reproducibility settings.

  • The experiments use Checkerboard and Mixture of eight Gaussians synthetic datasets.
  • For Checkerboard, the source and target distributions are formed by adding uniform square noise to distinct finite mixtures of points.
  • Algorithm 1 parameterizes both convex functions f and g with ICNNs using leaky-ReLU-based convex activations.
  • The baselines are Barycentric-OT, W1-LP, and W2GAN, using publicly available implementations and default hyperparameters.
  • Checkerboard training uses hidden size m = 64, L = 4 layers, batch size M = 1024, learning rate 10^-4, K = 10 generator iterations, and T = 10^5 total iterations.

C.2. High dimensional experiments

The experiments evaluate optimal transport learning across Gaussian, high-dimensional mixture, and MNIST settings, using input convex neural networks with specified training configurations.

  • Gaussian experiments: The Gaussian-to-Gaussian experiment uses 784-dimensional source and target distributions differing by mean shifts α ∈ {1, 5, 10}.Both distributions have identity covariance, with target mean µ = α(1, . . . , 1)⊤.
  • Mixture experiment: The high-dimensional Gaussian-to-mixture experiment transports a 784-dimensional standard Gaussian to a mixture of four Gaussians.The mixture components differ in their first two coordinates and have diagonal covariance concentrated on those coordinates.
  • MNIST experiments: The MNIST-to-MNIST experiment learns transport between latent embeddings of digit groups {0, 1, 2, 3, 4} and {5, 6, 7, 8, 9}.The embeddings are obtained using a pretrained variational autoencoder, and the ICNNs operate in 16 dimensions.
  • MNIST experiments: The Gaussian-to-MNIST experiment uses the same pretrained VAE and hyperparameter settings as the MNIST experiment, changing only the batch size to 64.The source is Gaussian while the target is represented through MNIST latent embeddings.

D. Further discussion of related work

Related work includes semi-dual, adversarial, primal, and projection-based approaches to approximating optimal transport, with limitations involving discretization, computational cost, or restricted optimality.

  • Semi-dual approaches: Classical semi-dual optimization methods discretize the space and require explicit probability densities, limiting applicability to high-dimensional real-world problems.These methods derive a functional derivative and optimize the objective with gradient descent.
  • Semi-dual approaches: Semi-discrete methods assume one discrete marginal and one continuous marginal with compact convex support, connecting function learning to the Alexandrov problem.The discrete marginal is supported on N points, while the continuous marginal has density on a compact convex domain.
  • Semi-dual approaches: Statistical analyses establish minimax convergence rates for certain regular distributions using wavelet-based representations of the Kantorovich potential.The associated procedure represents the function with wavelet basis functions up to a selected order.
  • Alternative approaches: Adversarial dual approaches replace the optimal-transport constraint with a quadratic penalty term to approximate the transport map.This approach remains within the dual formulation but uses regularization rather than enforcing the constraint directly.
  • Alternative approaches: Primal approaches learn an optimal coupling with two generators and two discriminators, making them computationally expensive despite handling cases where an optimal map may not exist.The generators map Gaussian variables to the data space, while discriminators enforce marginal constraints.
  • Alternative approaches: Projection-based methods optimize transport only on a subspace, while selecting an optimal projection is computationally challenging on the Grassmann manifold.Their approximation is inspired by sliced Wasserstein methods but does not optimize transport over the entire space.
  • ICNN-based approaches: Modeling the convex conjugate with an ICNN and adding an inverse-gradient penalty makes the optimization non-convex even with unrestricted function capacity.The penalty enforces ∇g as the inverse of ∇f and thereby encourages g = f∗.
Loading 1908.10962v2…