UniDynamic: Unified Surrogate Modeling for PDE and SPDE via Flow Matching

Pretraining shared models across deterministic and stochastic PDEs requires accommodating both diverse dynamics and different levels of predictive uncertainty. With sufficient observations, deterministic dynamics yield a unique continuation, whereas unobserved future stochastic forcing induces a conditional distribution of possible futures. A controlled Kolmogorov-flow experiment reveals complementary strengths of point prediction for deterministic accuracy and conditional generation for stochastic distribution matching. We therefore introduce UniDynamic, a flow-matching framework that jointly learns both modes across deterministic and stochastic PDEs. A shared backbone, conditioned on multiscale windowed log-signature history representations, supports one-pass prediction and iterative sampling. We construct a paired benchmark with repeated-future DNS ensembles for evaluating conditional moments, marginal distributions, and spatial fluctuations. UniDynamic achieves competitive accuracy across 11 deterministic PDE subsets. On the three stochastic PDEs, it attains the lowest lead-20 marginal energy distance among the compared methods. Furthermore, pretraining improves few-shot performance on unseen equations, while adding deterministic training data improves marginal-distribution accuracy on stochastic Kolmogorov flow and reaction–diffusion. These results support a unified flow-matching framework that combines accurate point prediction, conditional generation, and transfer across equations and regimes.