A Computational Framework for Solving Wasserstein Lagrangian Flows

Kirill Neklyudov, Rob Brekelmans, Alexander Tong, Lazar Atanackovic, Qiang Liu, Alireza Makhzani
Proceedings of the 41st International Conference on Machine Learning, PMLR 235:37461-37485, 2024.

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v235-neklyudov24a, title = {A Computational Framework for Solving {W}asserstein Lagrangian Flows}, author = {Neklyudov, Kirill and Brekelmans, Rob and Tong, Alexander and Atanackovic, Lazar and Liu, Qiang and Makhzani, Alireza}, booktitle = {Proceedings of the 41st International Conference on Machine Learning}, pages = {37461--37485}, year = {2024}, editor = {Salakhutdinov, Ruslan and Kolter, Zico and Heller, Katherine and Weller, Adrian and Oliver, Nuria and Scarlett, Jonathan and Berkenkamp, Felix}, volume = {235}, series = {Proceedings of Machine Learning Research}, month = {21--27 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v235/main/assets/neklyudov24a/neklyudov24a.pdf}, url = {https://proceedings.mlr.press/v235/neklyudov24a.html}, 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.} }
Endnote
%0 Conference Paper %T A Computational Framework for Solving Wasserstein Lagrangian Flows %A Kirill Neklyudov %A Rob Brekelmans %A Alexander Tong %A Lazar Atanackovic %A Qiang Liu %A Alireza Makhzani %B Proceedings of the 41st International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2024 %E Ruslan Salakhutdinov %E Zico Kolter %E Katherine Heller %E Adrian Weller %E Nuria Oliver %E Jonathan Scarlett %E Felix Berkenkamp %F pmlr-v235-neklyudov24a %I PMLR %P 37461--37485 %U https://proceedings.mlr.press/v235/neklyudov24a.html %V 235 %X 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.
APA
Neklyudov, K., Brekelmans, R., Tong, A., Atanackovic, L., Liu, Q. & Makhzani, A.. (2024). A Computational Framework for Solving Wasserstein Lagrangian Flows. Proceedings of the 41st International Conference on Machine Learning, in Proceedings of Machine Learning Research 235:37461-37485 Available from https://proceedings.mlr.press/v235/neklyudov24a.html.

Related Material