🤖 machine learning

Two-Parameter Flows for Learning Population Dynamics of Physical Systems

This paper introduces a scalable method for learning high-dimensional physical system dynamics from unlabeled samples by utilizing two-parameter flows to derive unique, regular physics-time velocities from base-to-marginal transports, thereby enabling the modeling of non-gradient phenomena like rotation without requiring trajectory data or per-step optimal transport couplings.

Original authors: Paul Schwerdtner, Tobias Blickhan, Benjamin Peherstorfer

Published 2026-05-27

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Paul Schwerdtner, Tobias Blickhan, Benjamin Peherstorfer

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Predicting the Crowd, Not the Individual

Imagine you are trying to predict how a massive crowd of people moves through a city over time.

The Old Way (Sample Dynamics): You try to track every single person's exact footsteps. If the crowd is chaotic (like a mosh pit or a swarm of bees), this is impossible. People bump into each other, change direction randomly, and their individual paths are messy and unpredictable. Trying to write a rule for "Person A's exact path" fails because the data is too noisy.
The New Way (Population Dynamics): Instead of tracking individuals, you look at the shape of the crowd. You ask: "Where is the dense part of the crowd at 9:00 AM? Where is it at 10:00 AM?" You want to learn the rules that move the entire blob of people from one shape to another, without worrying about who is exactly where inside that blob.

This paper introduces a new method to learn these "crowd movement rules" for complex physical systems (like fluids or particles) without needing to know the exact path of every single particle.

The Problem: The "Shortest Path" Trap

Many existing methods try to find the most efficient way to move the crowd from Point A to Point B. In math, this is called Optimal Transport. Imagine you have a pile of sand (the starting crowd) and you want to spread it into a new shape (the ending crowd). The "Optimal Transport" method tries to move every grain of sand the shortest possible distance.

The Catch:

It's too expensive: Calculating the shortest path for millions of grains of sand takes a massive amount of computer power.
It's too rigid: Sometimes, the crowd doesn't just move forward; it swirls, rotates, or spins. The "shortest path" method forces everything to move in straight lines or smooth gradients. It struggles to explain rotation (like a spinning vortex in water) because rotation isn't the "shortest" way to get from A to B.

The Solution: The "Two-Parameter Flow"

The authors propose a clever two-step trick using something they call Two-Parameter Flows. Think of this as having two different clocks:

Sampling Time (The "Vertical" Clock): This is the time it takes to build a model.
Physics Time (The "Horizontal" Clock): This is the actual time the physical system evolves (seconds, minutes, years).

Step 1: The Vertical Lift (Building the Map)

First, the model learns how to transform a simple, boring starting shape (like a perfect circle of dots) into the complex shapes the crowd takes at different times.

Analogy: Imagine you have a lump of clay (the base distribution). You learn how to stretch and mold this clay into a specific shape (the crowd at 9:00 AM), then another shape (the crowd at 10:00 AM), and so on.
The Trick: You don't worry about how the clay moves between 9:00 and 10:00 yet. You just learn the "recipe" to turn the lump of clay into the 9:00 shape, and the recipe to turn it into the 10:00 shape. This is easy and fast because modern AI tools are already very good at this.

Step 2: The Horizontal Slide (Finding the Real Motion)

Once you have the "recipes" for all the shapes, the model looks at how the same piece of clay moves as you change the time.

Analogy: Imagine you have a video where the clay morphs from a circle to a square, then to a triangle. If you watch the video frame-by-frame, you can see exactly how a specific speck of dust on the clay moves as the shape changes.
The Result: By connecting these "recipes" together, the model figures out the velocity (speed and direction) of the crowd moving through Physics Time.

Why This Is Better

It Handles Spin: Because the model isn't forced to find the "shortest path" (Optimal Transport), it can naturally learn that the crowd is spinning. If the data shows a vortex, the model learns a swirling motion. This is crucial for fluids and particles that rotate.
It's Fast: It avoids the heavy math of calculating the shortest path for every single particle. It uses standard, fast AI tools for the first step and then a simple regression (curve fitting) for the second step.
It Works on Huge Systems: The authors tested this on a system with over 10,000 dimensions (imagine a fluid simulation with millions of tiny points). Previous methods crashed or took forever on systems this big. This method handled it smoothly.

Real-World Examples from the Paper

The Spinning Vortex: In a simulation of a swirling fluid, the old methods tried to force the swirl to be a straight-line movement, resulting in jagged, unrealistic motion. The new method correctly learned the smooth, circular spin.
Particle Instabilities: In a physics problem involving charged particles (Vlasov-Poisson), the particles form complex, rotating filaments. The new method captured these fine, rotating details, while methods that force "straight-line" movement failed to see the structure.

Summary

The paper introduces a way to learn how complex groups of things move over time by first learning how to reshape a simple starting point into the target shapes, and then figuring out the motion that connects those shapes. This approach is faster, handles spinning/rotating systems naturally, and works on massive, high-dimensional problems that other methods can't solve.

Technical Summary: Two-Parameter Flows for Learning Population Dynamics of Physical Systems

Problem Statement
The paper addresses the challenge of learning the dynamics of high-dimensional probability densities over time using unlabeled samples, without access to trajectory information. The authors distinguish between sample dynamics (fitting models to minimize point-wise discrepancies along individual trajectories) and population dynamics (learning the evolution of the underlying probability law). In stochastic or chaotic systems, learning sample dynamics is often ill-posed or uninformative because individual trajectories may be non-differentiable or highly sensitive to initial conditions. Instead, the goal is to infer a velocity field $u(x, t)$ that satisfies the continuity equation $\partial_t \rho + \nabla \cdot (\rho u) = 0$ , ensuring that the flow induced by $u$ transports the initial distribution $\rho(0)$ to the observed time marginals $(\rho(t))_{t \in [0, T]}$ .

A key difficulty is that admissible velocity fields satisfying the continuity equation are generally not unique. Existing approaches often select a specific representative, such as the minimal kinetic energy field (which is a gradient flow), via Optimal Transport (OT) or Schrödinger Bridges. However, these methods face significant limitations:

Computational Cost: Computing OT couplings scales poorly (often cubically) with data size, making them intractable for high-dimensional problems.
Modeling Restriction: Enforcing a gradient structure ( $u = \nabla \phi$ ) excludes non-gradient phenomena like rotational or circulating flows, which are common in physical systems.
Regularity Issues: Minimal energy fields can exhibit sharp gradients or oscillations, particularly with finite samples, leading to unstable dynamics.

Methodology: Two-Parameter Flows
The authors propose Two-Parameter Flows (TPF), a framework that decouples the learning of distribution transport from the extraction of physics-time dynamics. The method introduces two time variables:

Sampling time ( $s \in [0, 1]$ ): The parameter along which a base distribution $\nu$ is transported to a specific time marginal $\rho(t)$ .
Physics time ( $t \in [0, T]$ ): The physical time evolution of the system.

The approach proceeds in three steps:

Learning Sampling-Time Transports: A standard conditional flow matching model (or stochastic interpolant) is trained to learn a velocity field $v(x, t, s)$ that transports the base distribution $\nu$ to each marginal $\rho(t)$ . This defines a two-parameter flow map $\Phi(a, t, s)$ such that $\Phi(\cdot, t, 1)_{\#} \nu = \rho(t)$ . This step leverages scalable, mature generative modeling tools and does not require trajectory data.
Deriving Physics-Time Velocity: Once $\Phi$ is fixed, the physics-time velocity field $u$ is uniquely determined by the consistency condition that mixed partial derivatives commute ( $\partial_t \partial_s \Phi = \partial_s \partial_t \Phi$ ). This leads to a partial differential equation relating $u$ and $v$ :
$\partial_s u + v \cdot \nabla u - u \cdot \nabla v - \partial_t v = 0$
The authors prove that if $v$ is sufficiently regular (e.g., $C^2$ ), $u$ is uniquely defined and inherits regularity properties from $v$ .
Extraction via Regression: To obtain an explicit, usable model for $u$ , the authors generate synthetic "coherent" trajectories. By fixing a noise sample $a \sim \nu$ and varying $t$ , they compute $\hat{x}(t) = \Phi(a, t, 1)$ . These synthetic trajectories are coupled across physics time (unlike the original independent marginals). A neural network $u_\theta$ is then trained via standard regression to match the time derivatives of these synthetic trajectories.

Key Contributions

Scalable Pipeline: The method reduces population dynamics inference to a standard conditional flow training step followed by a regression step. This avoids per-step optimal transport coupling computations, enabling application to high-dimensional systems (demonstrated up to $d > 10^4$ ).
Regularity Guarantees: The authors prove that the extracted physics-time velocity $u$ inherits the regularity of the learned sampling-time transport $v$ . This provides a theoretical guarantee that well-posed, smooth dynamics can be inferred if the base-to-marginal transport is regular.
Admissibility Beyond Gradient Flows: Unlike OT-based methods that enforce a minimal kinetic energy (gradient) structure, TPF allows for admissible non-gradient dynamics. This enables the model to naturally capture rotational and circulating physics phenomena that gradient-based approaches might miss or represent with irregular fields.
Uniqueness: The paper proves that fixing the transport from the base to all marginals uniquely selects a single physics-time velocity field among the many admissible ones.

Experimental Results
The authors evaluate TPF on several benchmarks:

Evolving Gaussian Mixture: TPF produces smoother, more regular dynamics compared to piecewise OT couplings, avoiding the sharp gradients observed in OT-based trajectories.
Barotropic Flow (Turbulence): In a high-dimensional ( $128 \times 128$ ) vorticity transport problem, TPF successfully captures the slow self-organization of vortices (population-level statistics like enstrophy) while filtering out fast, fine-scale advective motions. This results in a kinetic energy roughly two orders of magnitude lower than physical trajectories (indicating the removal of noise) and orders of magnitude faster inference times compared to numerical solvers. In contrast, MSE-fitted models failed to capture the vortex-merging behavior.
Kolmogorov Flow: TPF-generated samples accurately match the energy spectrum decay rate ( $\omega^{-3}$ ), demonstrating physical fidelity.
Vlasov-Poisson Instabilities: In particle instability benchmarks (two-stream and bump-on-tail), TPF outperforms gradient-based methods (HOAM, DICE) in capturing fine phase-space filamentation and rotational structures, achieving lower Wasserstein errors.

Significance and Claims
The paper claims that Two-Parameter Flows provide a principled and scalable route to learning population dynamics. The significance lies in:

Decoupling: Separating the difficult problem of matching marginals (solved by scalable flow matching) from the extraction of dynamics.
Flexibility: Avoiding the restrictive assumption of minimal kinetic energy/gradient flows, thereby accommodating rotational physics without sacrificing regularity.
Practicality: Enabling the inference of dynamics in high-dimensional regimes where OT-based methods are computationally prohibitive.

The authors modestly note that the inferred dynamics inherit the inductive bias of the chosen base-to-marginal transport. While the method guarantees regularity and marginal matching, future work is needed to determine what specific structures in the transport are required to satisfy additional criteria like minimal curl or low Lipschitz constants. The approach does not assume access to drift or diffusion coefficients, relying solely on unlabeled samples of time marginals.