Capturing non-Markovian dynamics in non-equilibrium stochastic systems using flow matching

This paper introduces a generative flow matching method that accurately captures non-Markovian and non-Gaussian effects in short-time stochastic particle dynamics, outperforming traditional regularized Dean-Kawasaki models in predicting statistical moments and first passage times.

The Gist

Imagine you are trying to predict how a crowd of people moves through a hallway.

If the hallway is packed with thousands of people, you can use a simple rule: "The crowd flows like water." This is easy to calculate and works well for long periods. In the scientific world, this is called the regularized Dean-Kawasaki (DK) equation. It treats the crowd as a smooth, continuous fluid.

However, what happens if the hallway is mostly empty, with only a few people wandering around? Or what if you want to know exactly what happens in the first few seconds after the doors open?

The "water" rule breaks down.

1. The "Few People" Problem: When there are very few people, the crowd isn't smooth; it's jagged and unpredictable. The "water" model might even predict negative numbers of people (which is impossible), just like a weather model might predict negative rain.
2. The "Memory" Problem: The "water" model assumes that where people are right now is all that matters. It forgets the past. But in reality, if a person just turned left, they are less likely to turn left again immediately. They have "memory." The old model ignores this, leading to wrong predictions about how the crowd spreads out quickly.

The New Solution: "Flow Matching"

The authors of this paper built a new, smarter way to simulate these crowds using a technique called Flow Matching. Think of this not as a rigid rulebook, but as a highly trained AI coach.

Instead of guessing how the crowd moves, the AI coach watches millions of real simulations of individual particles (like watching individual people walk). It learns two tricky things that the old "water" model missed:

• Non-Gaussian (The "Jagged" Shape): It learns that when there are few particles, the movement isn't a smooth bell curve; it has wild, unpredictable spikes.
• Non-Markovian (The "Memory"): It learns that the future depends on the past. It remembers the history of where the particles have been to predict where they will go next.

The Experiment: The "Kramers" Challenge

To test their new AI coach, the researchers set up a specific challenge called the Kramers first passage time problem.

Imagine a ball (or a particle) sitting in a valley (a low point). There is a hill in the middle, and another valley on the other side. The goal is to see how long it takes for the ball to roll over the hill and land in the new valley.

• The Setup: They simulated 5,120 different scenarios with 100 "cells" (small sections of the hallway).
• The Comparison: They ran the simulation three ways:
  1. The Gold Standard: Tracking every single particle individually (very accurate, but slow).
  2. The Old Way: The "water" model (DK equation).
  3. The New Way: Their AI "Flow Matching" model.

What They Found

1. The Old Way Failed Early: The "water" model (DK) was okay at predicting the average number of people in the new valley, but it was terrible at showing the actual movement. It created "ghosts" (negative numbers of particles) and missed the chaotic, jagged nature of the early movement.
2. The New Way Won: The AI model, especially the one that remembered the past (the non-Markovian version), perfectly captured the short-term chaos. It predicted the "higher-order statistics" (the weird, jagged details of the crowd) much better than the old model.
3. The Catch: The new AI model is very good at the beginning (short time). However, as time goes on, it starts to drift away from the truth, just like a GPS that gets slightly lost after a long drive.

The Bottom Line

This paper doesn't claim to solve every physics problem. It specifically shows that for systems with few particles and short timeframes, the old "smooth fluid" math is too simple.

By using Flow Matching, they created a model that acts like a smart observer who remembers the past and understands that small crowds are messy, not smooth. This allows for much more accurate predictions of how these systems behave in their critical early moments, which is something the old equations couldn't do.

Note: The authors mention this method is currently slower than tracking individual particles for simple systems, but they believe it will be much faster and more efficient for complex systems where particles interact with each other over long distances (like in chemistry or biology), where the old methods get stuck in computational traffic jams.

Technical Summary

Technical Summary: Capturing Non-Markovian Dynamics in Non-Equilibrium Stochastic Systems Using Flow Matching

Problem Statement
Hydrodynamic models based on coarse-grained stochastic partial differential equations (SPDEs), such as the regularized Dean-Kawasaki (DK) equation, are widely used to describe stochastic particle systems. However, these models face significant limitations in two specific regimes:

1. Low particle density: In regions with few particles, the number density distributions become highly non-Gaussian, which the regularized DK equation fails to capture accurately.
2. Short-time dynamics: The evolution of such diffusive, stochastic systems is heavily influenced by memory effects (non-Markovian dynamics), where initial conditions impact early-time trajectories. Classical coarse-graining, including the regularized DK formulation, typically loses this memory, treating fluxes as uncorrelated in space and time.

While the regularized DK equation predicts long-time statistics correctly for systems with high particle counts per cell, it produces nonphysical negative densities and inaccurate higher-order statistical moments during early-time evolution and in low-density regions.

Methodology
To address these shortcomings, the authors propose a data-driven generative approach using Flow Matching (FM) to model the probability distribution of particle fluxes directly from particle simulations.

• Target Variable: The method models the net particle current, $J(x, t)$, crossing a computational cell face. Unlike the Gaussian, uncorrelated noise used in the DK equation, the true current $J$ exhibits non-Gaussian tails and non-Markovian dependencies on the system's history.
• Flow Matching Framework: The authors construct a probability path $p_\tau$ from a known source distribution (a Student's $t$-distribution to capture heavy tails) to the target distribution of the particle current. A neural network velocity field, $u_\theta^\tau$, is trained to map samples from the source to the target by solving an ordinary differential equation (ODE).
• Incorporating Non-Markovian Effects: To capture memory, the velocity field is conditioned on the history of the local state from $k$ previous time steps. The input vector includes the number of particles in the left ($N_L$) and right ($N_R$) cells, and the current ($J$), for time steps $t, t-\Delta t, \dots, t-k\Delta t$.
  • For the Markovian limit ($k=0$), a DeepONet-based architecture is used.
  • For non-Markovian cases ($k>0$), a Transformer-based architecture is employed to handle the temporal sequence.
• Symmetry Enforcement: The model enforces statistical reflection symmetry, ensuring that the probability density of a current $J$ given particle counts $(N_L, N_R)$ is identical to that of $-J$ given $(N_R, N_L)$. This is implemented by reconstructing the velocity field as an antisymmetric combination of the network output.

Key Contributions

1. Generative Modeling of Fluxes: The paper introduces a flow matching framework that directly learns the non-Gaussian and non-Markovian probability distribution of stochastic fluxes, bypassing the need for ad-hoc regularization of noise.
2. Explicit Memory Integration: The method successfully integrates historical state information into the generative model, demonstrating that including memory ($k>0$) is critical for accurate short-time predictions.
3. Architectural Adaptation: The study demonstrates the utility of switching between DeepONet and Transformer architectures depending on whether the model requires historical context.

Results
The method was evaluated on the Kramers first-passage time problem for a system of non-interacting Brownian particles in a double-well potential ($V(x) = (x-\alpha)^2(x-\beta)^2$). The system was simulated using three approaches: direct Brownian random-walkers (ground truth), the regularized DK equation, and the proposed Flow Matching (FM) model (tested with both $k=0$ and $k=10$).

• Negative Densities: The regularized DK model frequently produced nonphysical negative density values, whereas the FM models maintained physical validity.
• Statistical Accuracy: While the DK model predicted mean particle densities reasonably well, it failed to capture higher-order statistical moments. The non-Markovian FM model ($k=10$) significantly outperformed both the DK model and the Markovian FM model ($k=0$) in capturing spatial particle distributions and higher-order moments during early-time evolution.
• Computational Cost: The non-Markovian FM model was computationally more expensive than direct Brownian dynamics (averaging 9 seconds per time step on CPUs) but scaled well up to 100 cells per CPU.

Significance and Claims
The authors claim that their generative flow matching approach provides a more accurate representation of short-time dynamics in stochastic particle systems compared to traditional Markovian SPDEs like the regularized DK equation. By explicitly modeling non-Gaussian and non-Markovian effects, the method reproduces higher-order statistics that are lost in standard coarse-graining.

The paper modestly notes that while the current methodology is computationally more intensive than non-interacting Brownian simulations, it is expected to offer significant efficiency gains for systems with complex, long-ranged interactions. In such systems, traditional particle simulations suffer from costly neighbor-list construction and strict time-step constraints, areas where a learned coarse-grained SPDE could potentially reduce computational overhead while maintaining accuracy. The authors acknowledge that for the specific non-interacting case studied, the model predictions begin to deviate from particle simulations at later times, with errors increasing over time.