Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Neural Optimal Transport with Lagrangian Costs

Aram-Alexandre Pooladian, Carles Domingo-Enrich, Ricky T. Q. Chen, Brandon Amos

TL;DR

This paper extends optimal transport to general costs defined by a least-action principle, enabling transport under obstacles and non-Euclidean geometries. It develops Neural Lagrangian Optimal Transport (NLOT) by neurally parameterizing the Kantorovich potential, amortizing the $c$-transform and path computations via cubic splines, and outputting deterministic transport maps and geodesic-like paths without ODE solvers. A second thrust learns Riemannian metrics from sequential measure pairs, casting metric learning as a bilevel optimization that leverages the LOT dual to adjust the underlying geometry for minimal transport cost. Empirically, NLOT outperforms baselines on low-dimensional synthetic tasks with barriers and circular geometries, and the metric-learning setup successfully recovers target geodesics, demonstrating practical impact for physics-informed transport and geometry discovery. The work opens avenues for unbalanced OT, multi-marginal extensions, and statistical analysis of non-standard transport costs.

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.

Neural Optimal Transport with Lagrangian Costs

TL;DR

This paper extends optimal transport to general costs defined by a least-action principle, enabling transport under obstacles and non-Euclidean geometries. It develops Neural Lagrangian Optimal Transport (NLOT) by neurally parameterizing the Kantorovich potential, amortizing the -transform and path computations via cubic splines, and outputting deterministic transport maps and geodesic-like paths without ODE solvers. A second thrust learns Riemannian metrics from sequential measure pairs, casting metric learning as a bilevel optimization that leverages the LOT dual to adjust the underlying geometry for minimal transport cost. Empirically, NLOT outperforms baselines on low-dimensional synthetic tasks with barriers and circular geometries, and the metric-learning setup successfully recovers target geodesics, demonstrating practical impact for physics-informed transport and geometry discovery. The work opens avenues for unbalanced OT, multi-marginal extensions, and statistical analysis of non-standard transport costs.

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.
Paper Structure (26 sections, 38 equations, 4 figures, 4 tables)

This paper contains 26 sections, 38 equations, 4 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Background on optimal transport
  3. Kantorovich primal-dual problems and optimal transport mappings
  4. Lagrangian Optimal Transport (LOT)
  5. Lagrangian OT between two measures via neural networks
  6. Experiments
  7. Related work
  8. Lagrangian Schrödinger bridges.
  9. Estimation and applications of optimal transport maps under the squared-Euclidean cost.
  10. Estimation of optimal transport maps for other notions of cost.
  11. Metric learning with NLOT
  12. Experiments
  13. Related work
  14. Conclusion
  15. Details on the use of splines for geodesics and paths
  16. ...and 11 more sections

Figures (4)

  • Figure 1: Optimal transport paths with Lagrangian costs on the obstacles setting from liu2022deep and circular geometry from scarvelis2022riemannian.
  • Figure 2: Paths on the Neural Lagrangian Schrödinger Bridge (NLSB) datasets koshizuka2022neural.
  • Figure 3: We successfully recover the metrics on the settings from scarvelis2022riemannian.
  • Figure 4: Our transport geodesics are able to reconstruct continuous versions of the original data that can predict the movement of individual particles given only samples from the first measure.

Theorems & Definitions (4)

  • Remark 1
  • Example 1: Euclidean distances, cf. benamou2000computational
  • Example 2: Obstacles and other potential functions
  • Example 3: Squared geodesic distances on Riemannian manifolds