Search papers, labs, and topics across Lattice.
This paper introduces PDE-JEPA, a novel predictive representation learning framework for parametric PDEs, moving beyond traditional reconstruction-based objectives. By first pretraining an encoder with masked-latent prediction to capture dynamics and then employing a geometry projector to align latent trajectory geometry with physical field evolution, PDE-JEPA explicitly models the underlying dynamics. The resulting framework decomposes latent dynamics into parameter-independent and parameter-dependent components, achieving significant out-of-distribution generalization improvements over state-of-the-art methods.
Predictive pretraining for parametric PDEs unlocks a 51.4% generalization improvement on unseen governing parameters by explicitly modeling latent dynamics geometry.
Physical trajectories contain more than snapshots of a system: they also reveal how its states evolve under governing conditions. However, representation learning for parametric partial differential equations (PDEs) has largely relied on reconstruction-based objectives that emphasize recovering observed physical fields. In this paper, we investigate predictive representation pretraining as an alternative to reconstruction-based learning. We find that predictive representations preserve rich physical information, yet this advantage alone does not ensure accurate field evolution. Based on these observations, we introduce PDE-JEPA for parametric PDE dynamics. Specifically, we first train an encoder using a masked-latent prediction to capture the underlying regularities of PDE dynamics. To explicitly adapt the pretrained representation toward a more dynamics-aligned state space, we then introduce a geometry projector that aligns latent trajectory geometry with the evolution geometry of physical fields. Finally, building on this geometry-aligned latent space, we further develop a physics-structured latent predictor that decomposes the dynamics into parameter-independent evolution and parameter-dependent response components. Extensive experiments on nine widely used PDE benchmarks demonstrate that our framework outperforms existing state-of-the-art methods by an average of 33.4\% in-distribution, while achieving an average improvement of 51.4\% when extrapolating to unseen governing parameters. The project page is available \href{https://tanpig-x.github.io/PDE-JEPA/}{here}.
