Fixed-Point Neural Optimal Transport without Implicit Differentiation
This work addresses training instability and computational overhead in neural optimal transport for practitioners needing scalable and stable transport map estimation.
The paper introduces a neural optimal transport method that uses a single network to parameterize a Kantorovich dual potential, reformulating the c-transform as a proximal fixed-point problem to avoid adversarial training and implicit differentiation. Experiments show strong transport accuracy with improved stability and efficiency across high-dimensional benchmarks and image translation tasks.
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.