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

PDE-JEPA. Predictive Representation Learning of Latent Dynamics Modeling for Parametric PDEs. By Zhentao Tan and three others, from Zhejiang University. Published on the twenty-eighth of September. Its code repository is public but slight: four source files.

A physical trajectory is more than a sequence of snapshots. It also reveals how a system's states evolve under the conditions that govern it. Representation learning for parametric partial differential equations, though, has largely relied on reconstruction objectives, which emphasize recovering the physical fields that were observed. Reconstruction asks a model to reproduce what it saw, while a trajectory also carries how the system got there. The authors investigate predictive representation pretraining as an alternative, and find a mixed picture. Predictive representations do preserve rich physical information. That advantage alone, however, does not ensure accurate evolution of the fields.

PDE-JEPA is built on that observation, for the dynamics of parametric PDEs. It keeps predictive pretraining for what it does well, and adds two stages aimed at what it does not: one to adapt the learned space so it follows the dynamics more closely, and one to predict within that space with physical structure built in. Prediction supplies the physical information, and the later stages are there to turn it into accurate evolution of the fields.

There are three steps. First, an encoder is trained with masked-latent prediction, to capture the underlying regularities of PDE dynamics. Second, a geometry projector explicitly adapts that pretrained representation toward a more dynamics-aligned state space, by aligning the geometry of latent trajectories with the evolution geometry of the physical fields. Third, on top of this geometry-aligned latent space, a physics-structured latent predictor decomposes the dynamics into two components: an evolution that is independent of the parameters, and a response that depends on them.

Experiments cover nine widely used PDE benchmarks. In distribution, the framework outperforms existing state-of-the-art methods by an average of thirty-three point four percent. When extrapolating to governing parameters it has not seen, the average improvement is fifty-one point four percent. Both figures are averages, taken across the nine benchmarks.

The second number is larger than the first. The improvement grows rather than shrinks when extrapolating to unseen governing parameters. The recipe behind it runs in sequence: learn by prediction, align the latent geometry with how the physical fields evolve, then predict with the parameters' influence held apart from the evolution they do not affect. The decomposition is the step that names the parameters explicitly: one part of the dynamics is independent of them, and the other is a response to them.