Leveraging Gauge Freedom for Learning Non-Gradient Population Dynamics of Stochastic Systems

In existing works on population dynamics inference, there is a focus on flows arising from vector fields that are the gradients of scalar potentials. Among all admissible flows that are compatible with the population dynamic, gradient flows are optimal in a specific sense: they minimize kinetic energy. The selection of fields based on different criteria corresponds to a gauge freedom when determining population dynamics, which we leverage in this work. We propose Non-Gradient Inference Flows (NGIF), an algorithm to infer non-gradient population dynamics using a weak formulation of the continuity equation. This allows us to parameterize with general vector fields as well as choose other selection criteria beyond minimal energy. We demonstrate on a variety of low- and high-dimensional physics problems that this more general approach improves distributional accuracy over gradient-restricted baselines and better captures non-potential transport.

Existing methods often assume populations move in the simplest "downhill" way, like following the slope of a landscape. That choice is mathematically convenient and uses the least kinetic energy, but it may miss more complex motion. This work proposes NGIF, a method that allows populations to move in more general ways, not just downhill along gradients. It uses a flexible equation-based approach to infer these movements from data. In tests on physics problems, NGIF more accurately matched how populations actually evolved, especially when the motion involved swirling or non-downhill transport.