Latent Generative Solvers for Generalizable Long-Term Physics Simulation

We study long-horizon surrogate simulation across heterogeneous PDE systems. We introduce Latent Generative Solvers (LGS), a two-stage framework that (i) maps diverse PDE states into a shared latent physics space with a pretrained VAE, and (ii) learns probabilistic latent dynamics with a Transformer trained by flow matching. Our key mechanism is an uncertainty knob that perturbs latent inputs during training and inference, teaching the solver to correct off-manifold rollout drift and stabilizing autoregressive prediction. We further use flow forcing to update a system descriptor (context) from model-generated trajectories, aligning train/test conditioning and improving long-term stability. We pretrain on a curated corpus of $\sim$2.5M trajectories at $128^2$ resolution spanning 12 PDE families. LGS matches strong deterministic neural-operator baselines on short horizons while substantially reducing rollout drift on long horizons. Learning in latent space plus efficient architectural choices yields up to \textbf{70$\times$} lower FLOPs than non-generative baselines, enabling scalable pretraining. We also show efficient adaptation to an out-of-distribution $256^2$ Kolmogorov flow dataset under limited finetuning budgets. Overall, LGS provides a practical route toward generalizable, uncertainty-aware neural PDE solvers that are more reliable for long-term forecasting and downstream scientific workflows.

💡 Research Summary

The paper introduces Latent Generative Solvers (LGS), a two‑stage framework for long‑horizon surrogate simulation across heterogeneous partial differential equation (PDE) families. First, a pretrained variational auto‑encoder (P2V‑AE) compresses high‑resolution PDE fields (128²) from twelve distinct families into a shared low‑dimensional latent physics space, abstracting away resolution and discretization while preserving dynamical structure. Second, a Transformer‑based Flow‑Forcing Transformer (FFT) learns probabilistic latent dynamics using flow‑matching. A key “uncertainty knob” k adds isotropic Gaussian noise to the latent input during training and inference, forcing the model to denoise off‑manifold states and thereby learn an auto‑correcting transition. The model also updates a physics context vector c via gated cross‑attention on its own generated predictions (flow forcing), aligning train‑time and test‑time distributions and mitigating exposure bias.

The authors provide a theoretical comparison: deterministic autoregressive operators accumulate error geometrically under Lipschitz dynamics, whereas flow‑matching with softened inputs does not propagate errors across steps. To handle long histories efficiently, a temporal pyramid (PFFT) downsamples early latent tokens.

Experiments on a curated dataset of ~2.5 M trajectories (five‑step length) show that LGS matches state‑of‑the‑art deterministic neural operators (FNO, U‑AFNO, CNextU‑Net) on short horizons, while reducing long‑term L2 error by 30‑45 % on rollouts of 50‑100 steps. Because dynamics are learned in latent space, FLOP consumption is roughly 70× lower than non‑generative baselines, enabling scalable pre‑training. Fine‑tuning on an out‑of‑distribution 256² Kolmogorov flow dataset with ≤ 5 % of parameters still yields competitive performance.

Overall, LGS demonstrates that a unified latent physics representation combined with uncertainty‑aware flow‑matching yields a practical, scalable, and more reliable neural PDE solver capable of probabilistic long‑term forecasting and uncertainty quantification across diverse physical systems.