Generative Modeling of Stochastic Dynamics for Long-Time Evolution
Authors: Yang-yang Tan, Jinyang Li, Lingxiao Wang
Organizations: Institute for Physics of Intelligence, Graduate School of Science, The University of Tokyo, Bunkyo-ku, Tokyo 113-0033, Japan · RIKEN Center for Interdisciplinary Theoretical and Mathematical Sciences (iTHEMS), Wako, Saitama 351-0198, Japan · KEK Theory Center, Institute of Particle and Nuclear Studies · Graduate University for Advanced Studies (SOKENDAI), Oho 1-1, Tsukuba, Ibaraki 305-0801, Japan
Exact stochastic equations for non-equilibrium dynamics are rarely accessible. We show that the long-time evolution of stochastic dynamics can be predicted from configuration pairs at a fixed short time lag, without knowledge of the equation of motion. Generative diffusion models learn the finite-time transition kernel from these pairs, and iterating it propagates the dynamics far beyond the training lag. For two-dimensional Model B, the diffusive dynamics of a conserved order parameter, the learned kernels reproduce dynamic critical scaling and self-similar t1/3 coarsening. Agreement with direct simulations persists on lattices twice the largest training size and for initial ensembles absent from training. For driven colloids in a periodic optical potential, ten minutes of measured trajectories suffice to predict the particle current and mean passage time over the next twenty minutes within experimental uncertainty. Short-time observations thus contain the information needed to predict emergent non-equilibrium dynamics at much longer times.
Figures & tables
Figure 1: Conditional transition kernel. The upper row shows the physical trajectory from ϕ0 to ϕt and onward in steps of δt . In the lower row, the present field ϕt conditions a score network that guides reverse diffusion from Gaussian noise at s=1 to an increment Δϕs=0∼pθ(Δϕ∣ϕt) . Adding this increment to ϕt yields ϕt+δt . Repeated sampling generates ϕt+nδt . The snapshots illustrate coarsening deep in the broken phase.
Figure 2: Long-time prediction of the learned kernel from the diffusion model (DM) against a high-precision stochastic equation solver (SOSRA [ 22 ] ) for Model B at the critical point ( κ=0.34034 , λ=1 , L=32 , and δt=10 ). (a) Field configurations at t=0 , 10 , 50 , 100 , 500 , and 2000 for equilibrium HMC (upper pair) and non-equilibrium Gaussian (lower pair, with the spatial mean of the paired HMC field) initial ensembles. Within each pair, SOSRA and DM are shown above and below, respectively, and start from the same field at t=0 . (b) Equal-time power G1(t) of the smallest diagonal mode. Dotted lines with bands show SOSRA, and points show DM. (c) Mean-field drift ⟨∣M(t)−M(0)∣⟩ . Both dynamics conserve M to single-precision accuracy. Line and marker styles are as in (b).
Figure 3: Critical relaxation at κ=0.34034 , λ=1 , and L=32 . The diagonal-mode correlations collapse when plotted against x=Δt/τn . Colored markers show the DM results, with error bars denoting standard errors. The dotted black curve and the dashed blue curve show fits of Eq. 11 to the SOSRA and DM data, respectively. The gray dash-dotted curve is e−x . The inset shows the relaxation times versus lattice momentum, yielding effective exponents z=3.792(19) for the DM and z=3.812(20) for SOSRA.
Figure 4: Coarsening deep in the broken phase ( κ=0.85 , λ=1 ). The kernel was trained only on synthetic field configurations at L≤64 and applied at L=128 from a disordered initial state. (a) DM configurations at t=10 , 100 , and 1000 . (b) Collapse of the sign-field correlation C(r,t) for t=100 – 1000 . The insets compare the domain size L(t) , defined by the first zero of C(r,t) , with the direct SOSRA result and show their ratio. The black dotted line shows SOSRA, and blue open circles show DM. The gray dashed line denotes the t1/3 growth law. The SOSRA and DM ensembles contain 128 and 32 trajectories, respectively.
Wall time (s)
Method
Eq.
Noneq.
DPM-Solver-v3 (This work)
3.24
3.26
SOSRA (This work)
1.30
3.28
SOSRA (StochasticDiffEq.jl) [ 22 , 51 ]
12.36
30.67
ShARK (Diffrax) [ 52 , 53 ]
5.56
17.76
EM (Diffrax) [ 54 , 53 ]
19.45
38.75
Table 1: Wall time for evolving 512 configurations on an L=64 lattice from t=0 to 100 , for the equilibrium (Eq.) HMC ensemble and the non-equilibrium (Noneq.) quench ensemble of Fig. 4 . Details are given in Sec. S6.1 and Sec. S6.2 .
\fnum@figure : Conserved transport during phase ordering on an L=128 lattice at κ=0.85 and λ=1 . All panels show the same enlarged region of one SOSRA trajectory. (a,b) Field ϕ at t=100 and 300 . The dotted line in (b) marks ϕ(t=100)=0 . (c) Colors show the field change Δϕ=ϕ(300)−ϕ(100) . Selected streamlines show the direction of the integrated current J=∫100300jdt , including the noise contribution in Eq. S8 . Very weak currents are omitted from the streamline overlay. The fixed box Ω encloses a shrinking domain. Fields and current are smoothed with the same periodic Gaussian of standard deviation two lattice spacings, preserving the discrete continuity relation Δϕ+∇⋅J=0 .
\fnum@figure : Time-origin averaged mode correlations ρn(Δt) of the first four diagonal modes on the equilibrium ensemble of Fig. 2 . Dotted lines with bands show SOSRA, and open markers show the DM kernel trained on Gaussian fields. All time origins from the t=2000 trajectories are included. Bands and error bars are trajectory jackknife standard errors over 1024 trajectories.
\fnum@figure : Effective decay rate Γeff=−dlnρn/dx of the DM scaling function against 1/x=τn/Δt for the data of Fig. 3 . Open markers show the four resolved diagonal modes. The dashed curve is a separate tail fit 1/(2x)+λ0 for x≥2 . The dash-dotted curve is the logarithmic derivative of the fit in Eq. 11 . The dotted lines mark a pure exponential decay, Γeff=1 , the short-time rate r fixed by the static amplitudes through Eq. S17 , and the long-time rate λ0 from the tail fit.
\fnum@figure : Dynamic scaling function of the structure factor during coarsening. Structure factor of the sign field for the L=128 DM trajectories of Fig. 4 , rescaled by the domain size L(t) as in Eq. S22 . The dashed line is the Porod law k^−3 with fitted amplitude. Error bars are standard errors over the 32 trajectories.
\fnum@figure : Reverse-sampling resolution at L=64 . Relaxation times are extracted by joining the statistically resolved segments of time-origin-averaged correlations on output grids with spacings δt=0.5 up to t=50 and δt=10 up to t=1000 . We compare SOSRA, the learned kernels sampled with 2000 EM steps, and the same kernels sampled with DPM-Solver-v3 using 32 NFE per physical transition. Only the commonly resolved modes n=2,…,8 are shown. The dotted line for SOSRA and dashed lines for the learned kernels are power-law fits to τn∝k^n−z . Error bars and the parenthesized uncertainties of z are trajectory-bootstrap standard errors.
\fnum@figure : Time-origin-averaged diagonal-mode correlations ρn(Δt) on the equilibrium ensemble at L=32 . Panels (a)–(d) show n=1,…,4 for kernels trained at L=32 and at L=8,16,32 , with training lag δt=10 . Panels (e)–(h) show n=5,…,8 for both kernels with δt=0.5 . Dotted lines without markers show SOSRA on the corresponding time grid, squares the kernel trained at L=32 , and circles the kernel trained at L=8,16,32 . For the kernels, dashed and solid lines denote δt=10 and 0.5 . The upper row is shown up to Δt=710 . In the lower row, the kernel trained at L=32 is shown up to Δt=5.5 , and the other curves up to Δt=10 . Every point averages at least 30 time origins. Error bars are leave-one-trajectory-out jackknife standard errors over 512 trajectories for training at L=8,16,32 and over 1024 trajectories for training at L=32 .
\fnum@figure : One-step diagonal-mode transfer tests at L=32 and L=64 . Panels (a) and (b) show ρn(δt) at L=32 and 64 , respectively. Each correlation is evaluated at the single-step separation Δt=δt . The L=32 initial fields are drawn from the training set, while L=64 was excluded from training. Each panel compares δt=0.5 and 10 for SOSRA and the two diffusion kernels. Dotted lines without markers show SOSRA at both time separations. For the kernels, solid lines denote δt=0.5 and dashed lines δt=10 . Squares show the kernel trained only at L=32 , and circles show the kernel trained on mixed L=8,16,32 lattices. Error bars on the kernel results are leave-one-trajectory-out jackknife standard errors over 512 trajectories.
Method
Step or NFE
Time (s)
S/V
ϕ2
ϕ4
ϵ
DPM-Solver-v3 (This work)
32 NFE per 10
3.26
−0.078
+0.098
+0.183
0.128
SOSRA (This work)
0.0025
3.28
−0.075
−0.061
−0.157
0.107
SOSRA (StochasticDiffEq.jl) [ 22 , 51 ]
0.0025
30.67
−0.146
−0.021
−0.102
0.104
ShARK (Diffrax) [ 52 , 53 ]
0.00125
17.76
+0.007
−0.052
−0.130
0.081
EM (Diffrax) [ 54 , 53 ]
0.0002
38.75
+0.209
−0.047
−0.025
0.125
EM (StochasticDiffEq.jl) [ 54 , 51 ]
0.0002
256.44
+0.163
−0.005
+0.049
0.098
Table S1: Accuracy at t=100 and wall time for the non-equilibrium benchmark. Observable entries are signed relative differences from the h=0.0001 SOSRA reference in percent. The last column is Eq. ( S24 ). All rows use one NVIDIA GH200.
Undriven
Driven
Data
DM
Langevin
Data
DM
Langevin
P(Δn=−1) per step
0.0157
0.0154
0.0161
0.0728
0.0735
0.0763
P(Δn=+1) per step
0.0168
0.0170
0.0161
0.0047
0.0047
0.0043
vˉ [ μ m/s]
0.0005(1)
0.0006
0.0000
−0.0282(3)
−0.0284
−0.0298
Deff,x [ μ m 2 /s]
0.0152
0.0163
0.0163
0.0430
0.0396
0.0408
Passage time [s]
610(15)
576
576
168(4)
166.5
157.3
Table S2: Measured and predicted colloid transport at 10 s resolution. The comparison covers 50 min without driving and 20 min under driving, starting at recording time 10 min. Parentheses give standard errors over measured tracks. Model means use 32 replicas per initial position, with sampling errors omitted. Passage times include only trajectories reaching either xi±λ without driving, or xi−λ under driving, within the observation window.
\fnum@figure : Colloid transport without driving (upper row) and under driving (lower row). Symbols and lines as in Fig. End Matter . (a,b,d,e) Mean-square displacement along x and y over the time Δt elapsed since the common origin at t=10 min. The gray line shows free diffusion 2DΔt . (c,f) Mean transverse position against recording time. Shading shows the standard error of the mean displacement from t=0 . The dotted line marks the end of training. The later transverse drift in the driven data is absent from both models.
\fnum@figure : Distributions of the net displacement Δx from t=10 min after 1 , 5 , 20 , and 50 min of undriven motion (upper row) and after 1 , 5 , 10 , and 20 min of driven motion (lower row). Panels (d) and (h) repeat Fig. End Matter (b,f). Bins have width λ/4 and are centered at integer multiples of λ/4 . Symbols and lines as in Fig. End Matter .
\fnum@figure : Periodic density and transition-path times of the colloid trajectories, undriven (upper row) and driven (lower row). (a,e) −lnp(x) , the negative logarithm of the position probability density folded onto one spatial period. The coordinate x measures distance from the barrier top along the drive direction. Gray open circles show experimental data from the training window ( t=0 – 10 min). Black open circles show data from the subsequent prediction window ( t=10 – 60 min in (a) and t=10 – 30 min in (e)). Blue solid and orange dashed lines show the DM and Langevin predictions over the corresponding prediction window. At equilibrium, −lnp(x) equals U(x)/(kBT) up to an additive constant. This identification does not hold under driving. Dotted lines mark the milestones A and B at 0.5 and 3.0μ m. (b–d,f–h) Transition-path times between A and B at 10 s resolution. Black open circles and gray open squares show the data for A→B and B→A . Full and faint lines show DM and the Langevin reference for the two directions. (b,f) All paths. (c,g) Paths passing the well bottom between the milestones. (d,h) Paths passing the barrier top. Without driving, 56% of the A→B paths and 60% of the B→A paths pass the well bottom. Under driving these fractions are 99% and 2% .
Hydrodynamic models of stochastic particle systems represented by coarse-grained stochastic partial differential equations (SPDE), such as the regularized Dean-Kawasaki (DK) equation, do not accurately capture the short-time system dynamics that is dominated by non-Markovian effects, and low particle density regimes where the distributions are highly non-Gaussian. We develop a generative flow matching method that directly models the probability distribution of fluxes from particle simulations that explicitly incorporates non-Markovian and non-Gaussian effects. As a demonstration, we use this method to simulate the Kramers first passage time problem for a system of non-interacting Brownian particles. We show the model accurately captures the short-time behavior and provides better predictions of the statistical moments of the number density when compared against the solution of the Markovian baseline, regularized DK equation.
Bhargav Sriram Siddani, John B. Bell, Alejandro L. Garcia +1
Lawrence Berkeley National Laboratory · San Jose State University
Neural stochastic differential equations (SDEs) have emerged as powerful tools for learning noisy or stochastic dynamics directly from data; however, existing approaches largely assume uncoupled and continuous noise, limiting their applicability to realistic stochastic drivers, and often scale poorly in time, requiring expensive autoregressive training. To address these limitations, we propose Neural Kolmogorov Equations (NKEs), a deterministic, infinite-dimensional reformulation of Neural SDEs based on the Kolmogorov Forward equation, transforming the learning problem from modelling individual stochastic trajectories to modelling the evolution of probability densities. NKEs learn general Lévy-type stochastic forcing directly through the operator structure of the KFE, and enable parallel-in-time training via a Lagrangian Galerkin projection and operator splitting. We evaluate NKEs on several stochastic benchmarks, including systems with coupled noise and jump processes, and verify that NKEs provide flexible models that accurately recover deterministic and stochastic dynamics with competitive predictive accuracy and improved training efficiency. Code and pretrained models will be released.
We propose a latent score-based generative AI framework for learning stochastic, non-local closure models and constitutive laws in nonlinear dynamical systems of computational mechanics. This work addresses a key challenge of modeling complex multiscale dynamical systems without a clear scale separation, for which numerically resolving all scales is prohibitively expensive, e.g., for engineering turbulent flows. While classical closure modeling methods leverage domain knowledge to approximate subgrid-scale phenomena, their deterministic and local assumptions can be too restrictive in regimes lacking a clear scale separation. Recent developments of diffusion-based stochastic models have shown promise in the context of closure modeling, but their prohibitive computational inference cost limits practical applications in many real-world settings. This work addresses this limitation by jointly training convolutional autoencoders with conditional diffusion models in latent space, significantly reducing the dimensionality of the sampling process while preserving essential physical characteristics. Numerical results demonstrate that the joint training approach helps discover a proper latent space that not only guarantees small reconstruction errors but also ensures good performance of the diffusion model in the latent space. When integrated into numerical simulations, the proposed stochastic modeling framework via latent conditional diffusion models achieves significant computational acceleration while maintaining comparable predictive accuracy to standard diffusion models in physical space.
Xinghao Dong, Huchen Yang, Jin-Long Wu
Department of Mechanical Engineering, University of Wisconsin-Madison, Madison, WI, 53706, USA