A Call to Lagrangian Action: Learning Population Mechanics from Temporal Snapshots

TL;DR

This paper introduces Wasserstein Lagrangian Mechanics, a novel framework for modeling population dynamics that captures complex behaviors like periodicity, and presents WLM, an algorithm to learn these second-order dynamics from data.

Contribution

It formalizes Wasserstein Lagrangian Mechanics and develops WLM, the first method to learn second-order population dynamics without predefined Lagrangians.

Findings

WLM outperforms existing methods in forecasting and interpolation.

It captures complex dynamics such as vortex behavior and flocking.

The framework unifies classical and quantum mechanics with population dynamics.

Abstract

The population dynamics of molecules, cells, and organisms are governed by a number of unknown forces. In the last decade, population dynamics have predominantly been modeled with Wasserstein gradient flows. However, since gradient flows minimize free energy, they fail to capture important dynamical properties, such as periodicity. In this work, we propose a change in perspective by considering dynamics that minimize a population-level action under a damped Wasserstein Lagrangian. By deriving the corresponding Hamiltonian equations of motion, we formalize Wasserstein Lagrangian Mechanics, a structured class of second-order dynamics that encompasses classical mechanics, quantum mechanics, and gradient flows. We then propose WLM as the first algorithm that learns these second-order dynamics from observed marginals, without specifying the Lagrangian. By directly learning the population mechanics, WLM can both forecast and interpolate unseen marginals, and outperforms existing gradient flow and flow matching methods across a wide range of dynamics, including vortex dynamics, embryonic development, and flocking.