Research map / Probabilistic and scientific modeling / Variational and generative sampling
A Computational Framework for Solving Wasserstein Lagrangian Flows
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a unified dual-based deep learning framework for Wasserstein Lagrangian Flows (OT, unbalanced OT, Schrödinger bridge, physically-constrained OT) that avoids simulating trajectories and optimal couplings, and applies it to single-cell trajectory inference.
Problem
Solving Lagrangian action minimization over density paths (optimal transport variants such as Schrödinger bridge, unbalanced OT, and physically constrained OT) with marginal constraints, which is computationally hard, e.g., for trajectory inference from single-cell snapshots.
Approach
Derive a dual (Hamiltonian) objective that is linear in the density when the Lagrangian is dual-linearizable, so only samples from a parameterized density path are needed; parameterize the cotangent vector s_t with a neural net and the path with a generative interpolation model, optimized as a min-max problem (plus Wasserstein gradient steps), without simulating ODEs/SDEs or needing optimal coupli
Key results
- WLF-(UOT + potential, λ=1) achieves the best W1 on Cite (34.160) and Multi (36.131) dims in the 50-dim PCA setting, outperforming OT-CFM (38.756 on Cite dim=50).
- On the EB train/test split, WLF-SB and WLF-(OT + potential) reach mean W1 0.73, improving on the best prior baseline NLSB (0.78) and OT-CFM (0.77).
- On melanoma data, WLF-OT (no growth) achieves MMD 5.04e-3 versus 1.75e-2 for uDSB (w/growth).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add the physical potential to WLF-OT | WLF-OT · EB, Cite, and Multi high-dim PCA leave-one-out | Wasserstein-1: 0.814 ± 0.002 -> 0.651 ± 0.002; 38.253 ± 0.071 -> 36.167 ± 0.031; 47.736 ± 0.110 -> 38.743 ± 0.060; 44.769 ± 0.054 -> 42 | Helped |
| Add the physical potential to WLF-UOT | WLF-UOT (λ = 1) · EB, Cite, and Multi high-dim PCA leave-one-out | Wasserstein-1: 0.800 ± 0.002 -> 0.634 ± 0.001; 37.035 ± 0.079 -> 34.160 ± 0.041; 45.903 ± 0.161 -> 36.131 ± 0.023; 43.530 ± 0.067 -> 41 | Helped |
| Use unbalanced OT instead of balanced OT | WLF · High-dim PCA leave-one-out | Wasserstein-1: 0.814 ± 0.002 -> 0.800 ± 0.002; 38.253 ± 0.071 -> 37.035 ± 0.079; 47.736 ± 0.110 -> 45.903 ± 0.161; 44.769 ± 0.054 -> 43 | Mixed |
| Use the Schrödinger Bridge objective instead of WLF-OT | WLF · High-dim PCA leave-one-out | Wasserstein-1: 0.814 ± 0.002 -> 0.816 ± 7.7e-4; 38.253 ± 0.071 -> 39.240 ± 0.068; 47.736 ± 0.110 -> 47.788 ± 0.111; 44.769 ± 0.054 -> 4 | Mixed |
| Use UOT instead of OT | WLF · EB 5-dim train/test split mean | Wasserstein-1: 0.74 -> 0.77 | Hurt |
| Add the physical potential to WLF-OT | WLF-OT · EB 5-dim train/test split mean | Wasserstein-1: 0.74 -> 0.73 | Helped |
| Add the physical potential to WLF-UOT | WLF-UOT (λ = 0.1) · EB 5-dim train/test split mean | Wasserstein-1: 0.77 -> 0.76 | Helped |
| Use WLF-SB instead of WLF-OT | WLF · EB 5-dim train/test split mean | Wasserstein-1: 0.74 -> 0.73 | Helped |
| Use WLF-OT instead of uDSB (no growth) | Melanoma treatment model · Melanoma treatment | MMD: 1.86e-2 -> 5.04e-3 | Helped |
| Use WLF-UOT instead of uDSB (w/growth) | Melanoma treatment model · Melanoma treatment | W2: 6.11 -> 5.01 | Helped |
Surprising findings
- The exact OT solver is not clearly better than the learned WLF-OT in some settings; WLF-UOT with λ=1 beats exact OT on Cite dim=5 and dim=100 leave-one-out results.
Limitations
- Schrödinger equation is not dual-linearizable, so the framework does not directly cover it
- Physical potential in experiments is a simple linear-acceleration estimate from finite differences of marginal means
- Discontinuous interpolation indicator needed for disconnected optimal paths
Experimental setup
Datasets: Embryoid body (EB) scRNA-seq CITE-seq Multiome melanoma treatment data
Benchmarks: EB scRNA-seq leave-one-out CITE-seq leave-one-out Multiome leave-one-out melanoma treatment (uDSB setting)
Metrics: Wasserstein-1 distance MMD Wasserstein-2 distance
Baselines: OT-CFM I-CFM [SF]^2 M-Geo [SF]^2 M-Exact Neural SDE TrajectoryNet IPF (GP) IPF (NN) SB-FBSDE NLSB uDSB exact OT
Builds on: Action Matching Wasserstein Hamiltonian flows Benamou-Brenier formulation Wasserstein-Fisher-Rao metric OT-CFM
Closest papers by meaning
- Wasserstein residuals: Learning Gradient Flows from Population Dynamics (NeurIPS 2026 · similarity 0.86)
- Transport meets Variational Inference: Controlled Monte Carlo Diffusions (ICLR 2024 · similarity 0.86)
- Progressive Entropic Optimal Transport Solvers (NeurIPS 2024 · similarity 0.86)
- Computing Optimal Transport Maps and Wasserstein Barycenters Using Conditional Normalizing Flows (ICML 2025 · similarity 0.86)
- Generalized Wasserstein Flow Matching: Transport Plans, Everywhere, All at Once (NeurIPS 2026 · similarity 0.86)
- Learned Lagrangian Models of PDEs via Euler–Lagrange Residual Minimization (NeurIPS 2026 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.