PDE-JEPA

RESEARCH PREDICTIVE PHYSICAL DYNAMICS

PDE-JEPA.

Predictive Representation Learning of Latent Dynamics Modeling for Parametric PDEs

Learning physical states that are informative,
evolvable, and ready to generalize.

Zhentao TanJianrong ZhangRuijie QuanYi Yang

Zhejiang University

Read the paper Code Watch the demos
PHYSICS BEYOND THE TRAINING RANGEFIG. 03
Numerical ground-truth vorticity field under an out-of-distribution viscosityNumerical ground truth
PDE-JEPA prediction of the same out-of-distribution vorticity fieldPDE-JEPA prediction
Same initial field. Unseen viscosity.
Vorticity at the final rollout step.
ν = 3.4545 × 10−5STEP 29
PREDICTIVE REPRESENTATIONS → PHYSICAL DYNAMICSDiscover the research
33.4%

Mean ID improvement

51.4%

Mean OOD improvement

9

Parametric PDE benchmarks

8/9

Best in-distribution results

TRAJECTORY DEMOS

Watch the dynamics unfold.

Explore ground truth, model output, and absolute error side by side across nine PDE datasets.

132videos · 116 trajectories

Vorticity

Model output

ν = 1.818e-05

Out-of-distribution (OOD)
Sample 00000 · 30 frames · 6 s
Open MP4 ↗Download ↓
8 videos in this setting
01 / THE IDEA

A useful representation
must also be
easy to evolve.

Predictive learning captures rich physical information. But an informative latent space alone does not guarantee accurate long-horizon forecasting.

PDE-JEPA connects representation learning with the structure of physical evolution: first learning predictive features, then aligning their trajectory geometry, and finally evolving them with a predictor that explicitly separates shared dynamics from parameter-dependent responses.

OBSERVATION 01

Information is there.

Frozen JEPA features support stronger local-state and parameter probes than reconstruction-based features on the evaluated PDEs.

OBSERVATION 02

Evolution needs structure.

Vanilla JEPA still gives higher in-distribution rollout errors on Wave-2D and Vorticity. Better probes do not automatically mean better forecasts.

Observations from Figure 1 and Section 1 of the manuscript ↗

02 / THE METHOD

From predictive features
to physical evolution.

Three stages. One goal: a latent state space
suited to forecasting and extrapolation.

01PREDICTIVE PRETRAINING

Learn the state.

Predict masked latent targets to learn physical regularities. Freeze the encoder and encode each downstream state frame by frame.

zt = E(ut)
02PHYSICS-ALIGNED GEOMETRY · PAG

Align its geometry.

A lightweight residual projector aligns latent trajectory directions with physical evolution. An anchor preserves pretrained information.

qt = zt + Gϕ(zt)
03PHYSICS-STRUCTURED PREDICTION · PSP

Evolve with structure.

Separate shared evolution from parameter-dependent responses, then integrate the learned vector field with a fixed-step RK4 solver.

dq/dt = Cθ(q) + Σj rj(ξ) Dθ(j)(q)
PDE-JEPA pipeline. Left: physics-aligned latent geometry with a frozen encoder, trainable geometry projector and predictor, and dynamics, geometry and anchor losses. Right: a physics-structured latent predictor combines a parameter-independent vector field with weighted parameter-dependent fields and evolves the latent state through an ODE solver.
Figure 2. Physics-aligned latent geometry (left) and the physics-structured latent predictor (right). View full size ↗

The structured components provide a physics-inspired inductive bias; they are not required to recover the exact analytical PDE operators.

03 / THE RESULTS

Better rollouts.
Beyond familiar parameters.

Transport, diffusion, waves, reaction–diffusion,
and fluid dynamics. Relative L2 error; lower is better.

Against the strongest reported baseline
8 OF 9 BENCHMARKS

Lowest ID error on eight benchmarks; second on Advection.

In-distribution rollout results from Table 1, compared with the strongest baseline for each benchmark.
PDE benchmarkBest baselineBaseline errorPDE-JEPAImprovement
AdvectionCoDA0.00680.0074−8.8%
BurgersLE-PDE0.08690.042850.7%
HeatPoseidon-T0.09330.027470.6%
Wave-BPoseidon-T0.10930.035068.0%
CombinedCAPE0.00850.007412.9%
Wave-2DZebra0.20700.114044.9%
VorticityLNS0.05920.034841.2%
HeterNSUniSolver0.00980.00899.2%
Gray–ScottUniSolver0.03230.028412.1%

Source: Table 1 ↗. Values and relative improvements are reproduced as reported. A negative improvement indicates a higher error.

The headline averages, 33.4% ID and 51.4% OOD, are reported in the abstract. See the manuscript for full baseline comparisons, parameter ranges, and evaluation protocols.

Each component matters.

OOD rollout error decreases as geometry alignment and structured prediction are added.

MODULE ABLATION · TABLE 2

Vorticity

Vanilla JEPA0.491
+ PAG0.397
+ PAG + PSP0.288

Wave-2D

Vanilla JEPA0.502
+ PAG0.321
+ PAG + PSP0.157

Same initial condition. A different viscosity.

The geometry-aligned and physics-structured models follow the numerical solution more closely under the unseen parameter.

Vorticity at final step 29 for matched initial conditions. Top row: training-reference viscosity. Bottom row: unseen viscosity. Columns compare numerical ground truth, vanilla JEPA, geometry-aligned predictions, physics-structured predictions, and their absolute errors.
Figure 3. Training-reference parameter (top) and OOD parameter (bottom), with shared field and error color scales. View full size ↗

EXPLORE THE COMPLETE STUDY

PDE-JEPA

Methods, benchmark details, and additional analyses.

Code on GitHub · Tanpig-X/PDE-JEPA
Read the full paper