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

TL;DR

This work addresses spurious solutions in semi-dual neural optimal transport (SNOT) by proving a sufficient condition on the source measure that guarantees recovery of the true OT map from the max-min solution. When the condition fails, it introduces OTP, a smoothing-based method that learns the full OT plan (potentially stochastic) by gradually annealing the smoothed distribution back to the original, with convergence guarantees along subsequences. The approach yields accurate OT plans in failure cases and delivers state-of-the-art results in unpaired image-to-image translation, including stochastic colorization tasks where deterministic maps fall short. The results underscore the practical value of learning transport plans and provide a pathway to robust neural OT methods in settings with inherent stochasticity and complex support structures.

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 spurious 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 spurious 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.

Figures (10)

• Figure 1: Visualization of failure cases by comparing the Optimal Transport map (1st row) and the max-min solution (2nd row) of Semi-dual Neural OT in the failure cases. The source data $x \sim \mu$, target data $y \sim \nu$, and generated data $T(x)$ are represented in Blue, Orange, and Red. The max-min solution fails to recover the correct OT Map.
• Figure 2: Example of a stochastic transport map (OT Plan) task, e.g., MNIST-to-CMNIST colorization.
• Figure 3: Qualitative comparison between OTM (1st row) and our model (2nd row) on failure cases in Sec \ref{['sec:failure']}. The noised source sample $\Tilde{x}$ in Alg \ref{['alg:otp']} is denoted in Green. While OTM falls into spurious solutions and fails to generate the target distribution correctly, our OTP model successfully learns the OT Plan.
• Figure 4: Experimental results on a stochastic transport map application, i.e., MNIST-to-CMNIST translation.
• Figure 5: Visualization of failure cases with ICNN potential function $f_\phi(y) := \alpha \Vert y \Vert^2 - V_\phi(y)$. We compare the Optimal Transport map (a) and our OTP model (b) in the failure cases. The source data $x \sim \mu$, target data $y \sim \nu$, and generated data $T(x)$ are represented in Blue, Orange, and Red. As illustrated in the figure, OTM fails to transport source to the target, i.e. $T_\# \mu \neq \nu$. On the other hand, our model successfully transports source $\mu$ to target $\nu$.

Theorems & Definitions (19)

• Theorem 3.1
• Theorem 3.2
• Corollary 3.3
• Proposition 3.4: Informal
• Theorem 4.1
• Lemma 1.1: Theorem 5.10 in villani
• Definition 1.2: SC
• Definition 1.3: $\mathbf{H_\infty}$
• Remark 1.4
• Lemma 1.5: Thm. 10.28 and Thm. 10.42 in villani