PDE-JEPA: Predictive Representation Learning of Latent Dynamics Modeling for Parametric PDEs

2026-09-28

AI Summary

This study addresses the challenge that reconstructive representation learning struggles to accurately capture physical field evolution dynamics in parametric partial differential equations (PDEs). To this end, we propose PDE-JEPA, a framework that introduces predictive pretraining to this domain for the first time. Methodologically, through masked latent variable prediction and geometric projection alignment, our approach innovatively decouples parameter-independent evolution from parameter-dependent responses, constructing a dynamics-aligned latent space. Evaluated across nine benchmarks, the proposed method achieves average performance improvements of 33.4% on in-distribution tasks and 51.4% on unseen parameter extrapolation, significantly enhancing model generalization capabilities.

Abstract

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}.

Problem

• parametric PDEs
• representation learning
• latent dynamics modeling
• predictive representation
• field evolution

Innovation

• Predictive Representation Learning
• Masked Latent Prediction
• Geometry Projector
• Physics-Structured Latent Predictor
• Parametric PDEs