Neural Optimal Transport with Lagrangian Costs (2406.00288v1)

Abstract: We investigate the optimal transport problem between probability measures when the underlying cost function is understood to satisfy a least action principle, also known as a Lagrangian cost. These generalizations are useful when connecting observations from a physical system where the transport dynamics are influenced by the geometry of the system, such as obstacles (e.g., incorporating barrier functions in the Lagrangian), and allows practitioners to incorporate a priori knowledge of the underlying system such as non-Euclidean geometries (e.g., paths must be circular). Our contributions are of computational interest, where we demonstrate the ability to efficiently compute geodesics and amortize spline-based paths, which has not been done before, even in low dimensional problems. Unlike prior work, we also output the resulting Lagrangian optimal transport map without requiring an ODE solver. We demonstrate the effectiveness of our formulation on low-dimensional examples taken from prior work. The source code to reproduce our experiments is available at https://github.com/facebookresearch/lagrangian-ot.

• The paper presents a neural network method that efficiently computes Lagrangian optimal transport maps by solving the dual Kantorovich problem.
• It leverages amortized optimization and cubic splines to represent complex transport paths in non-Euclidean geometries with improved computational efficiency.
• The framework learns unknown metrics from sequential probability measures, expanding its applications to advanced metric learning and physics simulations.

Overview of Neural Optimal Transport with Lagrangian Costs

This paper investigates the optimal transport problem using Lagrangian costs to advance the computational approximation of transport maps and paths in geometric contexts. Traditional approaches often utilize the squared-Euclidean cost, however, this approach integrates a broader class of Lagrangian costs, allowing for more nuanced modeling aligned with physical systems and non-Euclidean geometries.

Contribution

The paper presents methodologies to compute Lagrangian optimal transport (LOT) maps efficiently, without relying on an ODE solver. The approach leverages neural networks to parameterize and solve the dual Kantorovich problem, considering Lagrangian costs. Two main applications are explored:

1. Optimal Transport Between Measures: It utilizes known Lagrangian or metric information to compute optimal transport paths and mappings. The neural network models the Kantorovich potential, optimizing it through gradient descent methods.
2. Metric Learning from Data: When the underlying metric is unknown, the framework learns this metric from sequential pairs of probability measures. The learned metric is used to derive optimal transport maps, offering a novel approach to non-Euclidean geometry learning.

Methodology

The computational approach introduced involves the use of amortized optimization and cubic splines:

• Amortized Optimization: This technique helps in approximating the $c$-transform effectively, and along with gradient techniques, offers a robust mechanism for learning transport maps.
• Spline-based Path Representation: Cubic splines are used to parameterize optimal paths, computationally solving the transport problem in complex geometries.

Numerical Experiments

Evaluations demonstrate this approach's effectiveness by comparing to existing strategies like Neural Lagrangian Schrödinger Bridges. Key insights include:

• Successful representation of non-standard transport paths in potential-dominated and obstructions-laden environments.
• Improved accuracy in approximating target distributions, with substantial gains in computational efficiency.

Implications and Future Research

The paper's advancements pave the way for more comprehensive modeling of transport phenomena affected by complex geometric and physical constraints. From a theoretical perspective, the findings provide insights into bridging computational methodology with classical mechanics theories, such as those formulated by Lagrangian mechanics.

Looking forward, potential research directions include:

• Statistical Analysis of Novel Costs: Extending the statistical frameworks to cover the general class of Lagrangian costs.
• Unbalanced and Multi-Marginal Transport Problems: Applying the proposed method to unbalanced transport settings or those involving multiple marginals.
• Further Integration with Data Science Applications: Including genomic profiling, physics simulations, and other domains requiring precise geometric modeling.

In conclusion, this work builds a solid foundation for enhancing optimal transport computations beyond traditional settings, incorporating the nuances of Lagrangian costs into the computational field effectively.