Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Latent Generative Solvers for Generalizable Long-Term Physics Simulation

Zituo Chen, Haixu Wu, Sili Deng

TL;DR

Latent Generative Solvers (LGS) address long-horizon forecasting for heterogeneous PDE systems by learning in a shared latent physics space via a pretrained P2VAE and evolving latent dynamics with a Flow Forcing Transformer trained through flow matching. An uncertainty knob and a flow-forcing mechanism stabilize autoregressive rollouts and align train/test conditioning, enabling robust, uncertainty-aware predictions across diverse dynamics and resolutions. Pretrained on about $2.5\times 10^6$ trajectories at $128^2$ across $12$ PDE families, LGS achieves competitive short-horizon accuracy with substantial reductions in rollout drift and up to $70\times$ lower FLOPs, while adapting efficiently to out-of-distribution $256^2$ Kolmogorov flow with limited finetuning. This framework provides a scalable, uncertainty-aware route to generalizable neural PDE surrogates for long-term forecasting and downstream scientific workflows.

Abstract

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.

Latent Generative Solvers for Generalizable Long-Term Physics Simulation

TL;DR

Latent Generative Solvers (LGS) address long-horizon forecasting for heterogeneous PDE systems by learning in a shared latent physics space via a pretrained P2VAE and evolving latent dynamics with a Flow Forcing Transformer trained through flow matching. An uncertainty knob and a flow-forcing mechanism stabilize autoregressive rollouts and align train/test conditioning, enabling robust, uncertainty-aware predictions across diverse dynamics and resolutions. Pretrained on about trajectories at across PDE families, LGS achieves competitive short-horizon accuracy with substantial reductions in rollout drift and up to lower FLOPs, while adapting efficiently to out-of-distribution Kolmogorov flow with limited finetuning. This framework provides a scalable, uncertainty-aware route to generalizable neural PDE surrogates for long-term forecasting and downstream scientific workflows.

Abstract

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 2.5M trajectories at 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} lower FLOPs than non-generative baselines, enabling scalable pretraining. We also show efficient adaptation to an out-of-distribution 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.
Paper Structure (59 sections, 7 theorems, 50 equations, 12 figures, 2 tables)

This paper contains 59 sections, 7 theorems, 50 equations, 12 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Related Work
  4. Neural PDE solvers.
  5. Probabilistic neural PDE solvers.
  6. Diffusion forcing for long-term stability.
  7. Method
  8. Learning Latent Physics
  9. Uncertainty-Aware Generative Solver
  10. Flow matching with an uncertainty knob.
  11. Flow forcing on the data manifold.
  12. Training and Inference
  13. Stage I: Latent Physics Encoding.
  14. Stage II: Latent Dynamics Learning.
  15. Inference and Rollout.
  16. ...and 44 more sections

Key Result

Proposition 3.2

Under Assumption ass:lipschitz_main, If $\delta_0=0$, then

Figures (12)

  • Figure 1: Overview of Latent Generative Solver. All physics states are leveraged into a unified latent space, where initial condition $\mathbf{x}_0$ are perturbed with noise to cover off-manifold self-rollout states, which are guided back to clean next state $\mathbf{x}_1$. The predicted states serve to generating system dynamics descriptor "context" which differentiate heterogeneous physical dynamics.
  • Figure 2: Model Architecture. (a) Essential components of the Flow Forcing Transformer, that first generate next state $\mathbf{x}_{s+1}$ given intermediate state $\hat{\mathbf{x}}_s$, diffusion time $t_s$, and physics context $c_s$. Then leveraging the generated state, a gated cross-attention unit that updates the physics condition to be $c_{s+1}$. (b) A set of FFTs to realize the next-state-prediction during training, where intermediate states are transformed from inputs and sequential FFTs autoregressively inherit context $c$ and make prediction at each step.
  • Figure 3: t-SNE visualization of inferred physics contexts $c$ for different PDE systems at the beginning and end of prediction. Each system forms a stable cluster and exhibits minimal drift over the diffusion/transport steps, indicating that context updates remain on the learned manifold.
  • Figure 4: (a) Average L2RE over 16 systems versus training throughput. Marker area denotes per-sample forward FLOPs. Latent-space models (LGS and ablations) have substantially lower FLOPs. Physics context and temporal pyramids improve throughput with limited loss of accuracy, while generative modeling and the uncertainty knob improve rollout accuracy. (b) Example 10-step rollout on FNO-v3 comparing LGS and U-AFNO. Errors accumulate and amplify into visible artifacts for the deterministic baseline (highlighted region). Error is visualized with a fixed range of $(-0.5, 0.5)$.
  • Figure 5: OOD adaptation on $256^2$ Kolmogorov flow. Left: 1-step predictions after finetuning. Right: 10-step autoregressive rollouts of the finetuned models. LGS better preserves coherent vorticity structures and exhibits reduced long-horizon drift.
  • ...and 7 more figures

Theorems & Definitions (11)

  • Proposition 3.2: Deterministic compounding error
  • Theorem 3.4: One-step probability-flow deviation
  • Lemma 1.2: Deterministic error recursion
  • proof
  • Proposition 1.3: Deterministic compounding bound
  • proof
  • Lemma 1.6: Rescaled Gronwall bound
  • proof
  • Lemma 1.7: Softening contracts the carried-over error
  • Theorem 1.9: Multi-step error recursion for flow matching with knob $k$
  • ...and 1 more