Neural Optimal Transport with Lagrangian Costs

Neural Optimal Transport with Lagrangian Costs

1 Jun 2024 | Aram-Alexandre Pooladian, Carles Domingo-Enrich, Ricky Tian Qi Chen, Brandon Amos
This paper introduces a framework for computing optimal transport (OT) under Lagrangian costs, which generalize the traditional squared-Euclidean cost by incorporating geometric constraints and non-Euclidean structures. The proposed method, Neural Lagrangian Optimal Transport (NLOT), uses neural networks to approximate the optimal transport maps and paths, enabling efficient computation of geodesics and amortized spline-based paths. Unlike prior approaches, it avoids the need for ODE solvers to output the transport map, making it computationally efficient. The method is demonstrated on low-dimensional examples, including transport with obstacles and circular geometries, showing improved performance over existing baselines. The framework is also extended to learn Riemannian metrics from sequential probability measures, allowing for the recovery of underlying geometries without prior knowledge of the metric. The approach is validated on synthetic data and real-world scenarios, achieving accurate transport maps and paths. The method is generalizable to various cost functions, including those with potential energy terms and position-dependent costs, and is applicable to a wide range of domains such as computational biology, computer vision, and physics. The results highlight the effectiveness of the proposed framework in handling complex transport problems with geometric constraints.This paper introduces a framework for computing optimal transport (OT) under Lagrangian costs, which generalize the traditional squared-Euclidean cost by incorporating geometric constraints and non-Euclidean structures. The proposed method, Neural Lagrangian Optimal Transport (NLOT), uses neural networks to approximate the optimal transport maps and paths, enabling efficient computation of geodesics and amortized spline-based paths. Unlike prior approaches, it avoids the need for ODE solvers to output the transport map, making it computationally efficient. The method is demonstrated on low-dimensional examples, including transport with obstacles and circular geometries, showing improved performance over existing baselines. The framework is also extended to learn Riemannian metrics from sequential probability measures, allowing for the recovery of underlying geometries without prior knowledge of the metric. The approach is validated on synthetic data and real-world scenarios, achieving accurate transport maps and paths. The method is generalizable to various cost functions, including those with potential energy terms and position-dependent costs, and is applicable to a wide range of domains such as computational biology, computer vision, and physics. The results highlight the effectiveness of the proposed framework in handling complex transport problems with geometric constraints.
Reach us at info@study.space