[Paper Review] A Computational Framework for Solving Wasserstein Lagrangian Flows

Kirill Neklyudov, Rob Brekelmans+4 more|arXiv (Cornell University)|Oct 16, 2023

Model Reduction and Neural Networks4 citations

TL;DR

This paper proposes a unified deep learning framework for solving Wasserstein Lagrangian flows by leveraging dual formulations of Lagrangian action functionals, enabling efficient optimization without trajectory backpropagation or access to optimal couplings. The method outperforms prior approaches in single-cell trajectory inference by flexibly incorporating diverse prior knowledge through customizable kinetic and potential energy terms.

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.

Motivation & Objective

• To address the ill-posed problem of trajectory inference from sparse, marginal observations in single-cell RNA-seq data.
• To unify diverse optimal transport variants—such as Schrödinger bridge, unbalanced OT, and physically constrained OT—under a single computational framework.
• To enable flexible incorporation of prior knowledge into dynamics via customizable kinetic and potential energy terms in the Lagrangian formulation.
• To develop a differentiable, scalable method that avoids simulating or backpropagating through individual particle trajectories.
• To demonstrate superior performance in single-cell trajectory inference and image generation by leveraging dual optimization.

Proposed method

• The framework formulates trajectory inference as minimizing a Lagrangian action functional composed of kinetic energy (Wasserstein-2 or Wasserstein-Fisher-Rao geometry) and potential energy (e.g., entropy or spatial potentials).
• It employs a dual formulation of the Lagrangian to derive tractable optimization objectives, avoiding the need to simulate or differentiate through individual particle trajectories.
• The method uses deep neural networks to parameterize the vector field $ s_t(x, \theta) $ and the density path $ \rho_t(x, \eta) $, trained via dual objectives to match marginal constraints.
• The optimization is performed using a dual objective that enables end-to-end training without explicit trajectory rollout or coupling estimation.
• For image generation, the method enables single-step sampling via the Hamilton-Jacobi condition, where zero acceleration implies $ X_1 = X_0 + \nabla s_0(X_0) $.
• The framework supports multiple configurations by combining energy terms, allowing inductive biases for specific biological or physical constraints.

Experimental results

Research questions

• RQ1 Can a unified deep learning framework solve diverse optimal transport problems, including unbalanced OT, Schrödinger bridge, and physically constrained OT, under a single formulation?
• RQ2 How can prior knowledge about cell dynamics be flexibly encoded in a variational framework for trajectory inference?
• RQ3 Can dual optimization of the Lagrangian avoid the computational burden of trajectory backpropagation while maintaining accuracy?
• RQ4 To what extent does the framework improve trajectory inference performance in single-cell RNA-seq data compared to existing methods?
• RQ5 Can the framework enable single-step image generation via optimal transport maps without iterative ODE solving?

Key findings

• The proposed framework outperforms previous methods in single-cell trajectory inference tasks, particularly when incorporating prior biological knowledge through tailored Lagrangians.
• The method achieves state-of-the-art performance on low- and high-dimensional single-cell datasets, including Embryoid Body, CITE-seq, and Multiome data.
• For image generation, the framework enables single-step sampling via the Hamilton-Jacobi condition, matching the quality of iterative ODE solvers with only one function evaluation.
• The dual optimization approach avoids trajectory backpropagation and coupling estimation, significantly reducing computational cost.
• The framework successfully generalizes across multiple optimal transport variants, including unbalanced OT and Schrödinger bridge, by adjusting the Lagrangian components.
• Experiments on MNIST and CIFAR-10 show that single-step generation via the learned vector field produces high-fidelity samples comparable to iterative diffusion models.