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

arXiv: https://arxiv.org/abs/2605.25107 · HTML: https://mlsanity.com/paper/2605.25107/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 32 references that are in the map.
- First posted on arXiv 24 May 2026.

## Builds on (most cited first, 21 of 21)

- [A survey of the Schr"odinger problem and some of its connections with optimal transport](https://mlsanity.com/paper/1308.0215/) (2013)
- [Inductive Moment Matching](https://mlsanity.com/paper/2503.07565/) (2025)
- [Action Matching: Learning Stochastic Dynamics from Samples](https://mlsanity.com/paper/2210.06662/) (2022)
- [Proximal Optimal Transport Modeling of Population Dynamics](https://mlsanity.com/paper/2106.06345/) (2021)
- [Metric Flow Matching for Smooth Interpolations on the Data Manifold](https://mlsanity.com/paper/2405.14780/) (2024)
- [On the Error of Random Fourier Features](https://mlsanity.com/paper/1506.02785/) (2015)
- [Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks](https://mlsanity.com/paper/2106.00774/) (2021)
- [Learning to correct spectral methods for simulating turbulent flows](https://mlsanity.com/paper/2207.00556/) (2022)
- [Deep Momentum Multi-Marginal Schr"odinger Bridge](https://mlsanity.com/paper/2303.01751/) (2023)
- [Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport](https://mlsanity.com/paper/2410.00844/) (2024)
- [Spectral Mixture Kernels for Multi-Output Gaussian Processes](https://mlsanity.com/paper/1709.01298/) (2017)
- [Multi-marginal Schr"odinger Bridges with Iterative Reference Refinement](https://mlsanity.com/paper/2408.06277/) (2024)
- [Learning diffusion at lightspeed](https://mlsanity.com/paper/2406.12616/) (2024)
- [Riemannian Metric Learning via Optimal Transport](https://mlsanity.com/paper/2205.09244/) (2022)
- [Curly Flow Matching for Learning Non-gradient Field Dynamics](https://mlsanity.com/paper/2510.26645/) (2025)
- [Modeling Cell Dynamics and Interactions with Unbalanced Mean Field Schr"odinger Bridge](https://mlsanity.com/paper/2505.11197/) (2025)
- [Trajectory Inference with Smooth Schr"odinger Bridges](https://mlsanity.com/paper/2503.00530/) (2025)
- [Variational Regularized Unbalanced Optimal Transport: Single Network, Least Action](https://mlsanity.com/paper/2505.11823/) (2025)
- [Parametric model reduction of mean-field and stochastic systems via higher-order action matching](https://mlsanity.com/paper/2410.12000/) (2024)
- [Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme](https://mlsanity.com/paper/2506.01502/) (2025)
- [DICE: Discrete inverse continuity equation for learning population dynamics](https://mlsanity.com/paper/2507.05107/) (2025)

## Built on by (most cited first, 1 of 1)

- [One-Step Generative Surrogate Models via Block-Triangular Joint Drifting](https://mlsanity.com/paper/2609.26435/) (2026)

## Abstract (arXiv)

Existing work on population dynamics inference often focuses on flows arising from vector fields that are the gradients of scalar potentials. Among all admissible flows that are compatible with the population dynamics, 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 general vector fields and choose other selection criteria beyond minimal kinetic 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.