Two-Parameter Flows for Learning Population Dynamics of Physical Systems

Abstract: This work addresses the problem of learning the dynamics of high-dimensional probability densities over time using unlabeled samples, without assuming access to trajectory information. We introduce two-parameter flows that learn only sampling-time transports from a base distribution to each marginal and then extract a physics-time velocity by regressing on coupled synthetic trajectories. We prove that the resulting physics-time dynamics are unique and inherit regularity from the sampling-time transports. Because we can build on standard, well-developed conditional flow matching techniques for learning the base-to-marginal transports, our approach scales to high dimensions and avoids per-step optimal-transport couplings, while allowing admissible non-gradient dynamics that can naturally explain rotational or circulating physics phenomena.

• The paper introduces a two-parameter flows (TPF) framework that decouples base-to-marginal flows from physics-time regression to efficiently capture high-dimensional population dynamics.
• It rigorously establishes the uniqueness and regularity of the induced physics-time velocity field, overcoming limitations of gradient-restricted optimal transport methods.
• Empirical evaluations demonstrate that TPFs achieve superior regularity and computational efficiency in modeling complex phenomena such as turbulent flows and plasma instabilities.

Two-Parameter Flows for Learning Population Dynamics of Physical Systems

Problem Statement and Motivation

The paper introduces the problem of inferring the time evolution of high-dimensional probability distributions in physical systems, given only samples from time marginals and without trajectory-level supervision. Traditional approaches to model population dynamics often focus on optimal-transport-based minimal kinetic energy vector fields, which are computationally demanding and, by construction, restricted to gradient fields. This gradient restriction limits the representation of physical phenomena that exhibit rotational or circulation dynamics. Further, optimal-transport approaches become infeasible in high dimensions due to the cubic scaling in the number of samples. The introduced work seeks to overcome these limitations and provide a scalable, flexible method for learning population-level dynamics.

Methodology: Two-Parameter Flows Framework

The core contribution is the development of two-parameter flows (TPFs). TPFs decompose the modeling and learning process into two axes:

1. Sampling-time Direction (Base-to-Marginal Transports): The approach first learns transports from a fixed base distribution $\nu$ to each target time marginal $\rho(t)$, using established conditional flow matching techniques or stochastic interpolants.
2. Physics-time Direction (Dynamical Evolution): Once such base-to-marginal transports are fixed, the method exploits the induced structure to uniquely extract a physics-time velocity field via a regression step over coupled synthetic trajectories.

This two-stage approach has two main theoretical consequences:

• Uniqueness and Regularity: The induced physics-time velocity is unique (given the regularity of the flow) and obeys the continuity equation, thus matching the desired time marginals. Regularity of the physics-time velocity is provably inherited from the regularity of the learned base-to-marginal transports.
• Beyond Gradient Flows: Since the framework does not rely on minimal kinetic energy or gradient vector fields, TPFs can naturally model rotational or non-gradient dynamics prevalent in physical systems, addressing a key limitation of OT-based approaches.

Figure 1: Snapshots from a time curve $t \mapsto \rho(t)$, comparing piecewise OT transport (top row) to TPF dynamics (bottom row), showing improved regularity with TPF especially at later times.

Theoretical Properties and Relations to Optimal Transport

A fundamental insight in the work is the decoupling of distributional matching from explicit optimal transport coupling. The flow matching component focuses on learning a conditional velocity $v(x, t, s)$ in the 'sampling time' direction. Once this is established, the physics-time velocity $u$ is determined uniquely via a compatibility (Schwarz) condition, ensuring commutativity of mixed derivatives (sampling time $s$ and physics time $t$).

The paper rigorously proves several key results:

• Proposition of Uniqueness: If the learned flow is sufficiently regular, the physics-time velocity $u$ is uniquely determined via differentiation with respect to $t$, and it satisfies the continuity equation for $\rho(t)$.
• Regularity Propagation: An explicit, closed-form for $\rho(t)$0 is presented, showing that regularity in $\rho(t)$1 (learned in the first stage) directly transfers to $\rho(t)$2 (used for temporal dynamics), albeit with a potential loss of derivative order.
• Beyond OT: General Admissible Fields: Unlike classic OT-based methods (which recover the minimal kinetic energy/generating vector field in some special cases, e.g., commuting covariance matrices in evolving Gaussians), TPFs can represent rotational, non-gradient flows, including those phenomena where OT yields highly oscillatory or irregular velocity fields.

Practical Algorithm and Implementation

The algorithmic pipeline follows three main steps:

1. Train Base-to-Marginal Flows: Conditional flow matching is employed to learn $\rho(t)$3 to transport a base distribution to each $\rho(t)$4.
2. Generate Coupled Synthetic Trajectories: By integrating in sampling-time for each fixed noise realization, coherent trajectories across $\rho(t)$5 are obtained, coupling samples in physics time.
3. Regression for Physics-Time Velocity: A neural network $\rho(t)$6 is fit using regression to finite-difference time derivatives along synthesized trajectories. This yields a velocity field that governs the evolution of $\rho(t)$7.

This approach eliminates the need for expensive per-interval optimal transport coupling and can be scaled to high-dimensional dynamical systems—orders of magnitude larger than what has previously been reported in the literature.

Empirical Results and Numerical Evaluation

The method's efficacy is extensively validated on several systems:

• Evolving Gaussian Mixtures: TPF-generated sample trajectories exhibit superior regularity compared to those obtained via piecewise OT, which become increasingly non-smooth over time.
• Barotropic and Kolmogorov Flows: In turbulent 2D vorticity equations with state dimension $\rho(t)$8, TPF models successfully capture population-level phenomena such as vortex merging and the energy spectrum scaling ($\rho(t)$9 decay), while being orders of magnitude faster in inference than the underlying numerical simulators. MSE-based models that fit individual trajectories fail to reproduce the high-level distributional statistics.
• Vlasov-Poisson Particle Instabilities: On benchmarks for two-stream and bump-on-tail instabilities in the Vlasov–Poisson system, TPF models outperform or match methods such as action matching, DICE, and OT-based approaches in both $t \mapsto \rho(t)$0 error and the preservation of fine structure in phase space.

Implications, Limitations, and Future Directions

The TPF methodology broadens the class of admissible learned population dynamics, enabling representation of non-gradient, rotational, and smooth temporal evolutions where OT is restrictive or computationally prohibitive. The proven regularity guarantees and computational scalability make it appealing for high-dimensional physical systems encountered in turbulence, plasma physics, and other computational science domains.

However, the inductive bias of the induced population dynamics inherits from the initial base-to-marginal transport, which, while regular, may not always produce velocity fields with desirable additional properties (e.g., minimal kinetic energy, minimal curl, or specific Lipschitz constants). Future work may address the systematic enforcement of these constraints at the level of the initial transport, or via regularized training objectives. Moreover, interpretability and physical consistency of the learned velocity field $t \mapsto \rho(t)$1 warrant further investigation, especially for applications in scientific modeling and control.

Conclusion

Two-parameter flows provide a principled, scalable framework for inferring the time evolution of high-dimensional populations in physical systems, overcoming key limitations associated with gradient and OT-based approaches. With strong theoretical backing and empirical demonstration across diverse regimes, TPFs expand the toolbox for physical surrogate modeling and open avenues for further development in population-dynamics learning.