Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Improving Neural Optimal Transport via Displacement Interpolation

Jaemoo Choi, Yongxin Chen, Jaewoong Choi

TL;DR

The paper addresses instability in neural optimal transport map learning by leveraging displacement interpolation from dynamic OT theory. It derives a dual formulation for displacement interpolation and links time-evolving interpolation potentials through a Hamilton–Jacobi–Bellman condition, enabling training over the full interpolation trajectory. The proposed method, DIOTM, parametrizes intermediate maps via boundary maps and introduces an HJB-based regularizer to achieve stable, accurate OT map estimation, demonstrated on synthetic maps and image-to-image translation benchmarks with strong performance and stability. Limitations include the quadratic cost restriction and the need for bidirectional transport maps, suggesting avenues for extending the framework to more general costs and applications.

Abstract

Optimal Transport (OT) theory investigates the cost-minimizing transport map that moves a source distribution to a target distribution. Recently, several approaches have emerged for learning the optimal transport map for a given cost function using neural networks. We refer to these approaches as the OT Map. OT Map provides a powerful tool for diverse machine learning tasks, such as generative modeling and unpaired image-to-image translation. However, existing methods that utilize max-min optimization often experience training instability and sensitivity to hyperparameters. In this paper, we propose a novel method to improve stability and achieve a better approximation of the OT Map by exploiting displacement interpolation, dubbed Displacement Interpolation Optimal Transport Model (DIOTM). We derive the dual formulation of displacement interpolation at specific time $t$ and prove how these dual problems are related across time. This result allows us to utilize the entire trajectory of displacement interpolation in learning the OT Map. Our method improves the training stability and achieves superior results in estimating optimal transport maps. We demonstrate that DIOTM outperforms existing OT-based models on image-to-image translation tasks.

Improving Neural Optimal Transport via Displacement Interpolation

TL;DR

The paper addresses instability in neural optimal transport map learning by leveraging displacement interpolation from dynamic OT theory. It derives a dual formulation for displacement interpolation and links time-evolving interpolation potentials through a Hamilton–Jacobi–Bellman condition, enabling training over the full interpolation trajectory. The proposed method, DIOTM, parametrizes intermediate maps via boundary maps and introduces an HJB-based regularizer to achieve stable, accurate OT map estimation, demonstrated on synthetic maps and image-to-image translation benchmarks with strong performance and stability. Limitations include the quadratic cost restriction and the need for bidirectional transport maps, suggesting avenues for extending the framework to more general costs and applications.

Abstract

Optimal Transport (OT) theory investigates the cost-minimizing transport map that moves a source distribution to a target distribution. Recently, several approaches have emerged for learning the optimal transport map for a given cost function using neural networks. We refer to these approaches as the OT Map. OT Map provides a powerful tool for diverse machine learning tasks, such as generative modeling and unpaired image-to-image translation. However, existing methods that utilize max-min optimization often experience training instability and sensitivity to hyperparameters. In this paper, we propose a novel method to improve stability and achieve a better approximation of the OT Map by exploiting displacement interpolation, dubbed Displacement Interpolation Optimal Transport Model (DIOTM). We derive the dual formulation of displacement interpolation at specific time and prove how these dual problems are related across time. This result allows us to utilize the entire trajectory of displacement interpolation in learning the OT Map. Our method improves the training stability and achieves superior results in estimating optimal transport maps. We demonstrate that DIOTM outperforms existing OT-based models on image-to-image translation tasks.
Paper Structure (48 sections, 5 theorems, 33 equations, 15 figures, 7 tables, 1 algorithm)

This paper contains 48 sections, 5 theorems, 33 equations, 15 figures, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Notations and Assumptions
  3. Background
  4. Optimal Transport
  5. Dynamic Optimal Transport and Displacement Interpolation
  6. Method
  7. Dual Formulation of DI and the Relationship between Interpolation Potential Functions
  8. Dual Problem of Displacement Interpolant
  9. Relationship between Interpolation Potential Functions
  10. Displacement Interpolation Optimal Transport Model
  11. Parametrization of the DI Dual Problem
  12. Algorithm
  13. Related Work
  14. Experiments
  15. Optimal Transport Map Evaluation on Synthetic Datasets
  16. ...and 33 more sections

Key Result

Theorem 3.1

Given the assumptions in Appendix appen:proofs, for a given $t \in (0,1)$, the minimization problem $\inf_{\rho} \mathcal{L}_{DI} (t, \rho)$ (Eq. eq:displacement_opt) is equivalent to the following dual problem: where the supremum is taken over two continuous potential functions $f_{1, t} : \mathcal{Y} \rightarrow \mathbb{R}$ and $f_{2,t} : \mathcal{X} \rightarrow \mathbb{R}$, which satisfy $(1-t

Figures (15)

  • Figure 1: Visualization of transport maps $T$ on synthetic datasets. The transport map is visualized as a black line connecting each source sample $x$ to its corresponding generated data $T(x)$.
  • Figure 2: Image-to-Image translation results of DIOTM for Wild $\rightarrow$ Cat ($64 \times 64$) on AFHQ afhq. The left figure shows the source images and the right figure shows the corresponding translated images by DIOTM.
  • Figure 2: Image-to-Image translation benchmark results compared to existing neural (entropic) optimal transport models. $\dagger$ indicates the results conducted by ourselves. DSBM scores are taken from asbmSB-flow.
  • Figure 3: Image-to-Image translation results of DIOTM for Male $\rightarrow$ Female ($128 \times 128$) on CelebA celeba.
  • Figure 4: Ablation study on the regularizer hyperparameter $\lambda$.
  • ...and 10 more figures

Theorems & Definitions (7)

  • Theorem 3.1
  • Corollary 3.2
  • Theorem 3.3
  • Theorem A.1
  • proof
  • Theorem A.2
  • proof