Fixed-Point Neural Optimal Transport without Implicit Differentiation

TL;DR

This paper introduces a stable, efficient neural optimal transport method that avoids adversarial training by using a fixed-point formulation and does not require implicit differentiation, improving scalability and accuracy.

Contribution

It presents a novel fixed-point neural approach to optimal transport that simplifies training and enhances stability without implicit differentiation or multi-network architectures.

Findings

Achieves high transport accuracy on Gaussian benchmarks and image tasks.

Demonstrates improved training stability and computational efficiency.

Recovers both forward and backward transport maps effectively.

Abstract

We propose an implicit neural formulation of optimal transport that eliminates adversarial min--max optimization and multi-network architectures commonly used in existing approaches. Our key idea is to parameterize a single potential in the Kantorovich dual and reformulate the associated c-transform as a proximal fixed-point problem. This yields a stable single-network framework in which dual feasibility is enforced exactly through proximal optimality conditions rather than adversarial training. Despite the inner fixed-point computation, gradients can be computed without differentiating through the fixed-point iterations, enabling efficient training without requiring implicit differentiation. We further establish convergence of stochastic gradient descent. The resulting framework is efficient, scalable, and broadly applicable: it simultaneously recovers forward and backward transport maps and naturally extends to class-conditional settings. Experiments on high-dimensional Gaussian benchmarks, physical datasets, and image translation tasks demonstrate strong transport accuracy together with improved training stability and favorable computational and memory efficiency.