GradNetOT: Learning Optimal Transport Maps with GradNets

2025/07/17 by Shreyas Chaudhari, Chaudhari, Shreyas, Srinivasa Pranav +3
Computer Science · #Advanced Neural Network Applications #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning and Algorithms #Optimization and Search Problems

paper · pdf · doi:10.48550/arxiv.2507.13191

openalex publication_date 2025/07/17 · openalex created_date 2025/10/18 · openalex updated_date 2026/07/28

Abstract

Monotone gradient functions play a central role in solving the Monge formulation of the optimal transport (OT) problem, which arises in modern applications ranging from fluid dynamics to robot swarm control. When the transport cost is the squared Euclidean distance, Brenier's theorem guarantees that the unique optimal transport map satisfies a Monge-Ampère equation and is the gradient of a convex function. In [arXiv:2301.10862] [arXiv:2404.07361], we proposed Monotone Gradient Networks (mGradNets), neural networks that directly parameterize the space of monotone gradient maps. In this work, we leverage mGradNets to directly learn the optimal transport mapping by minimizing a training loss function defined using the Monge-Ampère equation. We empirically show that the structural bias of mGradNets facilitates the learning of optimal transport maps across both image morphing tasks and high-dimensional OT problems.

Citations

• Monge-Ampère gravity, optimal transport theory and their link to the Galileons
• High-dimensional Optimal Density Control with Wasserstein Metric Matching
• Learning Gradients of Convex Functions with Monotone Gradient Networks
• Supervised Training of Conditional Monge Maps
• Input Convex Gradient Networks
• Riemannian Convex Potential Maps
• Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2\n Benchmark
• Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization
• Optimal transport mapping via input convex neural networks
• Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations
• Input Convex Neural Networks
• The Monge-Ampère equation and its link to optimal transportation
• Gradient Networks
• Polar factorization and monotone rearrangement of vector‐valued functions
• A computational fluid mechanics solution to the Monge-Kantorovich mass transfer problem

Related

• Optimal Transport for Machine Learners
• The Monge Gap: A Regularizer to Learn All Transport Maps
• Learning Gradients of Convex Functions with Monotone Gradient Networks
• Learning to Transport with Neural Networks
• Optimal transport mapping via input convex neural networks