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

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

1. Problem Statement

The paper addresses the challenge of inferring population-level dynamics for stochastic systems where individual trajectories are unobserved or intrinsically unpredictable (e.g., chaotic or turbulent systems). The input consists of independent samples drawn from time-marginal distributions $\rho_t$ at discrete time points $t \in [0, T]$. The objective is to learn a time-dependent velocity field $u_t: X \to \mathbb{R}^d$ such that the flow induced by $u_t$ transports the initial distribution $\rho_0$ to the observed sequence of marginals $(\rho_t)_{t \in [0, T]}$.

A fundamental difficulty in this setting is that the continuity equation, $\partial_t \rho_t + \nabla \cdot (\rho_t u_t) = 0$, provides only an underdetermined constraint on the velocity field. The equation constrains only the $\rho_t$-weighted divergence of the flow, leaving the divergence-free (solenoidal) components of the velocity field unconstrained. This "gauge freedom" implies that infinitely many velocity fields can produce the same evolution of time marginals.

Existing methodologies often resolve this non-uniqueness by hard-coding a specific gauge choice: minimizing kinetic energy. This selection principle, rooted in optimal transport theory, forces the velocity field to be a gradient of a scalar potential ($u_t = \nabla s_t$). While optimal in terms of kinetic energy, this gradient restriction imposes an inductive bias that can be detrimental when the underlying dynamics involve significant rotational components (e.g., Hamiltonian systems or incompressible fluids). Enforcing a gradient structure on rotational dynamics can result in highly oscillatory, pathological potentials that are difficult to represent and learn, leading to poor distributional accuracy.

2. Methodology: Non-Gradient Inference Flows (NGIF)

The authors propose Non-Gradient Inference Flows (NGIF), an algorithm that decouples the task of matching time marginals from the task of selecting a specific velocity field.

Weak Continuity Equation Formulation:
Instead of enforcing the continuity equation pointwise, NGIF utilizes a weak (moment-based) formulation. For a set of test functions $\Phi$, the method requires that the rate of change of the expected value of any test function $\phi \in \Phi$ under the marginal $\rho_t$ matches the transport predicted by the candidate velocity field:
$$\frac{d}{dt} \mathbb{E}_{x \sim \rho_t}[\phi(x)] = \mathbb{E}_{x \sim \rho_t}[\nabla \phi(x) \cdot u_t(x)]$$
This identity is derived by integrating the continuity equation against test functions and applying integration by parts.

Empirical Loss and Random Fourier Features:
To operationalize this constraint, the authors employ random Fourier feature test functions. Frequencies $\omega$ are sampled from a distribution $\nu$ (using multiscale bandwidths to capture structures at different resolutions), generating test functions of the form $\sin(\omega^\top x)$ and $\cos(\omega^\top x)$.
The empirical loss minimizes the residual of the weak continuity equation over these test functions. Crucially, the time derivatives of the moments are estimated via smoothing splines fitted to the empirical moments of the training data, allowing for pre-computation and variance reduction.

Explicit Gauge Regularization:
Rather than restricting the hypothesis class to gradient fields, NGIF introduces an explicit gauge regularizer $G(u)$ to the optimization objective:
$$\min_{u} \mathcal{L}(u) + \lambda G(u)$$
where $\mathcal{L}(u)$ is the weak continuity residual loss. This allows the user to select from a broad family of structural biases. The paper proposes three specific regularizers:

1. Kinetic Energy ($G_{KIN}$): Minimizes $\int \|u_t\|^2 \rho_t$, recovering the standard gradient-based approach but as a soft penalty rather than a hard constraint.
2. Curl Penalty ($G_{CURL}$): Penalizes the antisymmetric part of the Jacobian, discouraging rotation and biasing toward gradient-like dynamics without enforcing a potential parametrization.
3. Divergence Penalty ($G_{DIV}$): Penalizes $\nabla \cdot u_t$, promoting approximately incompressible flows. This is particularly useful for fluid-like transport where the true dynamics are divergence-free.

By using a neural network to parameterize $u_t$ directly (rather than a potential $s_t$), the method avoids the computational complexities of backpropagating through spatial derivatives of a potential and allows for the learning of non-gradient fields.

3. Key Contributions

The paper outlines three primary contributions:

1. Decoupling of Marginal Matching and Gauge Fixing: The authors formulate population dynamics inference using weak continuity constraints, separating the matching of time marginals from the selection of the velocity field structure. Non-uniqueness is addressed via an explicit, user-chosen gauge regularizer.
2. Scalable Weak Constraints: The method instantiates the weak constraints using random Fourier feature test functions, yielding a computationally efficient and scalable loss function.
3. Non-Gradient Learning: The approach enables the learning of non-gradient velocity fields, which improves representability and stability. It demonstrates that relaxing the gradient constraint leads to more accurate population dynamics inference compared to gradient-restricted baselines.

4. Experimental Results

The authors evaluate NGIF on a variety of low- and high-dimensional physics problems:

• Gigli's Example (Rotational Dynamics): In a setup where Gaussian components rotate on a circle, gradient-based methods fail to capture the simple rotational motion, instead producing highly oscillatory potentials. NGIF with a curl regularizer accurately recovers the smooth rotational velocity field with only a marginal increase in kinetic energy compared to the optimal gradient solution.
• Tracer Particles in Turbulent Flow: In a Kolmogorov-forced turbulent flow (where the ground truth is divergence-free), NGIF with a divergence regularizer ($G_{DIV}$) produces particle trajectories and density reconstructions that closely match the ground truth. Gradient-based baselines (Action Matching, DICE, HOAM) perform significantly worse.
• Vlasov-Poisson Instabilities: On two-stream, bump-on-tail, and strong Landau damping problems, NGIF variants outperform state-of-the-art gradient-based methods (AM, HOAM, DICE) in predicting the evolution of electric energy. The divergence regularizer allows for richer particle movement, improving the accuracy of the final distribution.
• High-Dimensional Turbulence: In a 128x128 vorticity field simulation, NGIF applied to a latent space accurately captures the decay of enstrophy (a statistical measure of the flow), demonstrating that the learned flow reproduces salient population-level statistics even when individual trajectories are chaotic and unrecoverable.

5. Significance and Claims

The paper claims that population dynamics inference is inherently underdetermined, and the choice of gauge is a modeling decision rather than a search for a universally correct solution. The significance of NGIF lies in its ability to:

• Avoid Pathological Targets: By not hard-coding the minimal kinetic energy/gradient gauge, the method avoids the need to learn highly oscillatory potentials that arise when forcing gradient structures onto rotational dynamics.
• Incorporate Physical Priors: It provides a modular framework to incorporate specific physical knowledge (e.g., incompressibility) directly through the choice of gauge regularizer, without requiring reference flows or sample-level velocity data.
• Improve Distributional Accuracy: The experiments suggest that for systems with non-potential transport, non-gradient approaches yield superior distributional accuracy compared to gradient-restricted baselines.

The authors conclude that while the weak continuity equation constrains only the flux divergence, the explicit selection of a gauge regularizer allows for the learning of dynamical models that are better aligned with the underlying physics of the system, particularly in regimes where trajectory data is unavailable or uninformative.