Overcoming Spurious Solutions in Semi-Dual Neural Optimal Transport: A Smoothing Approach for Learning the Optimal Transport Plan

TL;DR

This paper introduces OTP, a novel method that learns both the optimal transport map and plan to overcome fake solutions in neural OT, improving accuracy and applicability in tasks like image translation and colorization.

Contribution

The paper proposes OTP, a new approach that jointly learns the OT map and plan, addressing fake solutions in semi-dual neural OT and enabling stochastic transport learning.

Findings

OTP outperforms existing models in image-to-image translation.

OTP successfully learns stochastic transport maps in one-to-many tasks.

The method eliminates fake solutions under certain distribution assumptions.

Abstract

We address the convergence problem in learning the Optimal Transport (OT) map, where the OT Map refers to a map from one distribution to another while minimizing the transport cost. Semi-dual Neural OT, a widely used approach for learning OT Maps with neural networks, often generates fake solutions that fail to transfer one distribution to another accurately. We identify a sufficient condition under which the max-min solution of Semi-dual Neural OT recovers the true OT Map. Moreover, to address cases when this sufficient condition is not satisfied, we propose a novel method, OTP, which learns both the OT Map and the Optimal Transport Plan, representing the optimal coupling between two distributions. Under sharp assumptions on the distributions, we prove that our model eliminates the fake solution issue and correctly solves the OT problem. Our experiments show that the OTP model recovers the optimal transport map where existing methods fail and outperforms current OT-based models in image-to-image translation tasks. Notably, the OTP model can learn stochastic transport maps when deterministic OT Maps do not exist, such as one-to-many tasks like colorization.