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 investigates the optimal transport problem between probability measures when the underlying cost function is governed by a *least action principle*, also known as a *Lagrangian* cost. The authors propose a computational framework to enforce transport with more general costs that can incorporate geometries such as obstacles and non-Euclidean spaces. They demonstrate the ability to efficiently compute geodesics and amortize spline-based paths, which has not been achieved before, even in low-dimensional problems. Unlike prior work, they output the *Lagrangian optimal transport map* without requiring an ODE solver. The effectiveness of their formulation is demonstrated on low-dimensional examples from prior work, and the source code is available at <https://github.com/facebookresearch/lagrangian-ot>. The paper focuses on solving the optimal transport problem when the cost of displacement is governed by a *least action principle*. The authors aim to compute the Lagrangian optimal transport maps and the resulting paths for these maps. They address two main challenges: computing the $c$-transform and evaluating the Lagrangian cost. To overcome these challenges, they use amortized optimization to approximate the $c$-transform and parameterize paths using cubic splines. They then apply gradient descent to optimize the Kantorovich dual problem and learn the Lagrangian cost and paths. The authors present experiments on two types of optimal transport problems: transport between a pair of probability measures and transport between consecutive pairs of measures. They demonstrate the effectiveness of their method on various examples, including obstacles and circular geometries. The results show that their method outperforms existing baselines in terms of marginal distribution accuracy and trajectory quality. The paper discusses related work, including Lagrangian Schrödinger bridges, estimation and applications of optimal transport maps under the squared-Euclidean cost, and estimation of optimal transport maps for other notions of cost. It also highlights the differences between their work and other methods, emphasizing the self-contained and unregularized nature of their approach.This paper investigates the optimal transport problem between probability measures when the underlying cost function is governed by a *least action principle*, also known as a *Lagrangian* cost. The authors propose a computational framework to enforce transport with more general costs that can incorporate geometries such as obstacles and non-Euclidean spaces. They demonstrate the ability to efficiently compute geodesics and amortize spline-based paths, which has not been achieved before, even in low-dimensional problems. Unlike prior work, they output the *Lagrangian optimal transport map* without requiring an ODE solver. The effectiveness of their formulation is demonstrated on low-dimensional examples from prior work, and the source code is available at <https://github.com/facebookresearch/lagrangian-ot>. The paper focuses on solving the optimal transport problem when the cost of displacement is governed by a *least action principle*. The authors aim to compute the Lagrangian optimal transport maps and the resulting paths for these maps. They address two main challenges: computing the $c$-transform and evaluating the Lagrangian cost. To overcome these challenges, they use amortized optimization to approximate the $c$-transform and parameterize paths using cubic splines. They then apply gradient descent to optimize the Kantorovich dual problem and learn the Lagrangian cost and paths. The authors present experiments on two types of optimal transport problems: transport between a pair of probability measures and transport between consecutive pairs of measures. They demonstrate the effectiveness of their method on various examples, including obstacles and circular geometries. The results show that their method outperforms existing baselines in terms of marginal distribution accuracy and trajectory quality. The paper discusses related work, including Lagrangian Schrödinger bridges, estimation and applications of optimal transport maps under the squared-Euclidean cost, and estimation of optimal transport maps for other notions of cost. It also highlights the differences between their work and other methods, emphasizing the self-contained and unregularized nature of their approach.
Reach us at info@study.space