Most representation learning for PDE dynamics focuses on reconstructing fields, but reconstructive objectives alone can miss the temporal regularities needed for accurate long-horizon forecasts. This work flips that emphasis: it first trains predictive latent features that capture how states evolve, then aligns those features with physical evolution geometry, and finally evolves them using a physics-structured predictor that separates shared dynamics from parameter-specific responses. The result is better in-distribution accuracy and stronger extrapolation to unseen parameters across diverse PDE families.
Key Findings
- Predictive pretraining (masked-latent prediction) produces latent features that contain richer physical information than reconstruction-based encoders, enabling stronger local probes for state and parameter inference.
- A lightweight geometry projector aligns latent trajectory directions with the true evolution geometry of physical fields, adapting pretrained features into a space better suited for forecasting.
- A physics-structured predictor decomposes dynamics into a parameter-independent vector field plus parameter-dependent response components and integrates them with an RK4 solver, providing an inductive bias helpful for parametric PDEs.
- Empirically, the framework yields an average in-distribution error reduction of 33.4% and an average extrapolation improvement of 51.4% across nine PDE benchmarks (transport, diffusion, waves, reaction–diffusion, fluids), with particularly large gains on Burgers, Heat, Wave-B, Wave-2D and Vorticity benchmarks.
Who it's for and trade-offs
Great fit if you need latent-space PDE surrogates that generalize across governing parameters and you value long-horizon rollouts with physics-inspired inductive structure. It is suitable for researchers building learned solvers, surrogate models for design/optimization, or benchmarks of latent dynamics. Look elsewhere if you need explicit, interpretable analytical PDE operators (this approach favors learned latent dynamics with structured components) or if you require minimal architecture complexity—PDE-JEPA adds pretraining plus projector and structured predictor stages that increase pipeline complexity compared with plain reconstruction-based models.