Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

A Computational Framework for Solving Wasserstein Lagrangian Flows

Kirill Neklyudov, Rob Brekelmans, Alexander Tong, Lazar Atanackovic, Qiang Liu, Alireza Makhzani

TL;DR

This work provides a unified dual-formulation framework for Wasserstein Lagrangian Flows, recasting trajectory inference as action minimization on the space of probability densities with kinetic and potential energies. By exploiting duality and linearizable objectives, it learns a cotangent vector field without simulating trajectories or requiring optimal couplings, enabling flexible incorporation of priors such as Schrödinger bridges, unbalanced OT, and physically constrained OT. The framework parameterizes the dual variables and density paths with neural networks and generative models, using Action Matching and Wasserstein gradient updates to optimize density trajectories that satisfy endpoint marginals. Empirically, it demonstrates competitive and improved performance on single-cell trajectory inference tasks, highlighting the practical impact of incorporating prior dynamics into marginal interpolation. The approach offers a principled, scalable route to diverse OT variants while preserving data-consistent marginals, with potential applicability to quantum, biological, and social science trajectory problems.

Abstract

The dynamical formulation of the optimal transport can be extended through various choices of the underlying geometry (kinetic energy), and the regularization of density paths (potential energy). These combinations yield different variational problems (Lagrangians), encompassing many variations of the optimal transport problem such as the Schrödinger bridge, unbalanced optimal transport, and optimal transport with physical constraints, among others. In general, the optimal density path is unknown, and solving these variational problems can be computationally challenging. We propose a novel deep learning based framework approaching all of these problems from a unified perspective. Leveraging the dual formulation of the Lagrangians, our method does not require simulating or backpropagating through the trajectories of the learned dynamics, and does not need access to optimal couplings. We showcase the versatility of the proposed framework by outperforming previous approaches for the single-cell trajectory inference, where incorporating prior knowledge into the dynamics is crucial for correct predictions.

A Computational Framework for Solving Wasserstein Lagrangian Flows

TL;DR

This work provides a unified dual-formulation framework for Wasserstein Lagrangian Flows, recasting trajectory inference as action minimization on the space of probability densities with kinetic and potential energies. By exploiting duality and linearizable objectives, it learns a cotangent vector field without simulating trajectories or requiring optimal couplings, enabling flexible incorporation of priors such as Schrödinger bridges, unbalanced OT, and physically constrained OT. The framework parameterizes the dual variables and density paths with neural networks and generative models, using Action Matching and Wasserstein gradient updates to optimize density trajectories that satisfy endpoint marginals. Empirically, it demonstrates competitive and improved performance on single-cell trajectory inference tasks, highlighting the practical impact of incorporating prior dynamics into marginal interpolation. The approach offers a principled, scalable route to diverse OT variants while preserving data-consistent marginals, with potential applicability to quantum, biological, and social science trajectory problems.

Abstract

The dynamical formulation of the optimal transport can be extended through various choices of the underlying geometry (kinetic energy), and the regularization of density paths (potential energy). These combinations yield different variational problems (Lagrangians), encompassing many variations of the optimal transport problem such as the Schrödinger bridge, unbalanced optimal transport, and optimal transport with physical constraints, among others. In general, the optimal density path is unknown, and solving these variational problems can be computationally challenging. We propose a novel deep learning based framework approaching all of these problems from a unified perspective. Leveraging the dual formulation of the Lagrangians, our method does not require simulating or backpropagating through the trajectories of the learned dynamics, and does not need access to optimal couplings. We showcase the versatility of the proposed framework by outperforming previous approaches for the single-cell trajectory inference, where incorporating prior knowledge into the dynamics is crucial for correct predictions.
Paper Structure (28 sections, 3 theorems, 39 equations, 2 figures, 3 tables, 1 algorithm)

This paper contains 28 sections, 3 theorems, 39 equations, 2 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Wasserstein-2 Geometry
  4. Wasserstein Fisher-Rao Geometry
  5. Action Matching
  6. Wasserstein Lagrangian Flows
  7. Wasserstein Lagrangian and Hamiltonian Flows
  8. Inner Optimization: Evaluating ${\mathcal{A}}_{{\mathcal{L}}}[\rho_t]$ using Action Matching
  9. Outer Optimization over Constrained Distributional Paths
  10. Computational Approach for Solving Wasserstein Lagrangian Flows
  11. Linearizable Objectives
  12. Parameterization and Optimization
  13. Wasserstein Gradient Approach
  14. Discontinuous Interpolation
  15. Examples of Wasserstein Lagrangian Flows
  16. ...and 13 more sections

Key Result

Theorem 1

For a Lagrangian $\mathcal{L}[\rho_t, \dot{\rho}_t,t]$ which is lsc and strictly convex in $\dot{\rho}_t$, the optimization where, for $s_t \in T^*_{\rho_t}{\mathcal{P}}$, the Hamiltonian ${\mathcal{H}}[\rho_t, s_t, t]$ is the Legendre transform of $\mathcal{L}[\rho_t, \dot{\rho}_t,t]$ (Eq eq:hamiltonian). The action ${\mathcal{A}}_{{\mathcal{L}}}[\rho_t]$ of a given curve is the solution to the

Figures (2)

  • Figure 1: Our Wasserstein Lagrangian Flows are action-minimizing curves for various choices of Lagrangian $\mathcal{L}_i[\rho_t, \dot{\rho}_t,t]$ on the space of densities, which each translate to particular state-space dynamics. Toy examples of dynamics resulting from various potential or kinetic energy terms are given in (a)-(d). We may also constrain Wasserstein Lagrangian flows match intermediate data marginals $\rho_{t_i} = \mu_{t_i}$ and combine energy terms to define a suitable notion of interpolation between given $\mu_{t_i}$.
  • Figure 2: For different definitions of Lagrangian $\mathcal{L}[\rho_t, \dot{\rho}_t,t]$ or Hamiltonian $\mathcal{H}[\rho_t, s_t,t]$ on the space of densities, we obtain different action functionals $\mathcal{A}[\rho_t]$. Here, we show state-space velocities and optimal density paths for the $W_2$ geometry and problem. (a) The action functional for each curve can be evaluated using Action Matching (inner optimization in \ref{['prop:dual_obj']}), which is performed in the state-space. (b,c) Minimization of the action functional (outer optimization in \ref{['prop:dual_obj']}) is performed on the space of densities satisfying two endpoint constraints and possible intermediate constraints.

Theorems & Definitions (9)

  • Theorem 1
  • Definition 3.1: (Dual) Linearizability
  • Proposition 1
  • Example 4.1: $W_2$ Optimal Transport
  • Example 4.2: Unbalanced Optimal Transport
  • Example 4.3: Physically-Constrained Optimal Transport
  • Example 4.4: Schrödinger Bridge
  • Theorem 1
  • proof