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

2026-05-08

AI Summary

Existing Wasserstein gradient flow methods struggle to capture complex population dynamics involving periodic or higher-order behaviors. This work proposes the Wasserstein Lagrangian Mechanics (WLM) framework, which models second-order collective dynamics by minimizing a damped Wasserstein Lagrangian action. WLM establishes, for the first time, a unified theoretical foundation encompassing classical mechanics, quantum mechanics, and gradient flows. The proposed algorithm does not require a predefined Lagrangian form; instead, it learns the Lagrangian end-to-end directly from temporal snapshots, leveraging unsupervised marginal distribution learning and flow matching techniques. Experiments demonstrate that WLM significantly outperforms existing approaches in tasks such as vortex dynamics, embryonic development, and collective aggregation, exhibiting superior capabilities in distribution interpolation and future-state prediction.

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.

Problem

• population dynamics
• Wasserstein gradient flows
• periodicity
• Lagrangian mechanics
• temporal snapshots

Innovation

• Wasserstein Lagrangian Mechanics
• population dynamics
• second-order dynamics
• action minimization
• marginal matching