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

Zhentao Tan; Jianrong Zhang; Ruijie Quan; Yi Yang

cs.AISep 28, 2026

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

Authors: Zhentao Tan, Jianrong Zhang, Ruijie Quan, Yi Yang

Organizations: Zhejiang University.

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

Figures & tables

Figure 1: Overview of our key observations. Predictive representation (JEPA) learns physically informative representations, but informativeness alone does not ensure effective forecasting of parametric PDEs. Our approach, PDE-JEPA, improves PDE forecasting and extrapolation through geometry alignment and physics-structured latent prediction.

Figure 2: Overview of PDE-JEPA. (a) We first adapt the frozen predictive representation initialized by JEPA with a light weight geometry projector that aligns latent trajectory geometry with physical evolution. (b) We then evolve the aligned latent states using a physics-structured latent predictor, which decomposes the dynamics into parameter-independent and parameter-dependent components and integrates their combined dynamics with an ODE solver.

Method	Advect	Burgers	Heat	Wave-B	Combined	Wave-2D	Vorticity	HeterNS	GS
Parametric Solvers
FNO	0.0390	0.4972	0.3449	0.9819	0.0374	0.9913	0.1411	0.0210	0.0547
CAPE	0.0094	0.2230	0.2130	0.9780	0.0085	–	–	–	–
CoDA	0.0068	0.5460	0.7670	1.0200	0.0120	0.7770	0.6780	–	–
GEPS	0.0930	0.3989	0.6641	0.6317	0.0097	0.5138	0.0821	0.1032	0.0332
In-Context Solvers

Table 1: ID rollout performance across PDE benchmarks. All results are reported in Relative

L^{2}

error, lower is better. Best results are bold and second-best results are underlined .

	Vorticity		Wave-2D
Model	ID	OOD	ID	OOD
Vanilla	.086	.491	.363	.502
+ PAG	.040	.397	.143	.321
+ PSP	.034	.288	.114	.157

Table 2: Module ablation on rollout.

	Vorticity		Wave-2D
Model	ID	OOD	ID	OOD
Vanilla	.086	.491	.363	.502
+ PAG	.040	.397	.143	.321
+ PSP	.034	.288	.114	.157

Table 2: Module ablation on rollout.

Dataset	Angle MAE ( ∘ )	Sym. Acc. MAE	Lag-1 Cos. MAE
Vorticity	$29.1\!\rightarrow\!\mathbf{6.3}$ (-78%)	$.46\!\rightarrow\!\mathbf{.10}$ (-77%)	$.36\!\rightarrow\!\mathbf{.05}$ (-84%)
Wave-2D	$35.5\!\rightarrow\!\mathbf{10.1}$ (-71%)	$.46\!\rightarrow\!\mathbf{.14}$ (-68%)	$.50\!\rightarrow\!\mathbf{.14}$ (-74%)
Burgers	$41.4\!\rightarrow\!\mathbf{20.8}$ (-49%)	$.73\!\rightarrow\!\mathbf{.37}$ (-49%)	$.43\!\rightarrow\!\mathbf{.17}$ (-59%)
GS	$28.8\!\rightarrow\!\mathbf{20.4}$ (-29%)	$.39\!\rightarrow\!\mathbf{.27}$ (-30%)	$.39\!\rightarrow\!\mathbf{.26}$ (-32%)

Table 3: Representation geometry before and after alignment.

Method	Combined	Wave-2D	Vorticity	HeterNS Visc./Force	GS
UniSolver	0.038	1.003	0.923	0.037 / 0.105	0.1636
Poseidon-T	0.146	1.511	0.665	0.560/0.821	0.083
LNS	0.1667	0.610	0.481	0.610/0.932	0.146
ENMA	0.243	1.151	0.467	1.501/2.271	0.134
Zebra	–	0.680	0.320	–	–
Ours	0.008	0.157	0.288	0.011 / 0.103	0.033

Table 4: OOD rollout performance across PDE benchmarks. Relative

L^{2}

error is reported, lower is better. Best results are bold and second-best results are underlined .

Figure 3: Visualization on Vorticity under matched initial conditions on final timestep. Prediction errors are shown for the vanilla, geometry-aligned and physics-structured models. More visualization can be found in Appendix .

		Field		Latent
Dataset	Model	Cos. $\uparrow$	Amp. $\rightarrow 1$	Cos. $\uparrow$	Amp. $\rightarrow 1$
Vorticity	PAG	.626	.846	.913	1.142
	+ PSP	.708	.987	.924	1.043
Wave-2D	PAG	.958	.979	.939	.990
	+ PSP	.986	.996	.947	1.002

Table 5: Parameter-response consistency under matched ICs.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Shape per trajectory	Train	Val.	Test	OOD
Advection	$140\times 256$	12,000	120	120	–
Burgers	$250\times 256$	12,000	120	120	–
Heat	$250\times 256$	12,000	120	120	–
Wave-B	$250\times 256$	12,000	120	120	–
Combined Equation	$140\times 256$	12,000	120	120	120
Vorticity	$1\times 128\times 128\times 30$	12,000	1,200	1,200	120

Appendix

Table 6: Datasets overview. Counts denote trajectory number. A dash means that no separate OOD split is specified for the audited dataset version. One-dimensional shapes use

(T,X)

; two-dimensional shapes use

(C,X,Y,T)

Dataset	Parameter	In-D	Out-D
Combined	$\alpha$	$\mathcal{U}([0,1])$	$[1.0682,1.7581]$
	$\beta$	$\mathcal{U}([0,0.4])$	same as In-D
	$\gamma$	$\mathcal{U}([0,1])$	same as In-D
Vorticity	$\nu$	$[10^{-3},10^{-2}]$	$[10^{-5},10^{-4}]$
HeterNS	$\nu$	$\{10^{-5},5{\times}10^{-5},10^{-4},5{\times}10^{-4},10^{-3}\}$	$\begin{gathered}\text{Interp.: }(\{2,3,4,6,7,8,9\}{\times}10^{-5})\\ {}\cup(\{2,3,4,6,7,8,9\}{\times}10^{-4})\\ \text{Extra.: }\{2,3,4,5,6,7,8,9\}{\times}10^{-3}\end{gathered}$
HeterNS	$m$	$\{1,2,3\}$	$\begin{gathered}\{1,2,3\}\;(\text{viscosity OOD})\\ \{0.5,1.5,2.5,3.5\}\;(\text{forcing OOD})\end{gathered}$

Appendix

Table 7: In-distribution (In-D) and out-of-distribution (Out-D) parameter settings for each dataset.

Figure 4: Advection. Transport speed $\beta=2.01597$ , within the training support $[0,4]$ , with periodic boundaries on a domain of length 128. The stored segment spans $t\simeq 44.18$ to $100$ .

Figure 5: Burgers. Viscosity $\beta=0.072296$ , within the training support $[10^{-3},5]$ , with periodic boundaries on a domain of length 16. The forcing is $F(t,x)=\sum_{j=1}^{5}A_{j}\sin(\omega_{j}t+2\pi\ell_{j}x/16+\phi_{j})$ ; its coefficients $A_{j}$ , $\omega_{j}$ , $\ell_{j}$ , and $\phi_{j}$ are listed beneath.

Figure 6: Heat. Diffusivity $\beta=0.0653548$ , within the training support $[10^{-3},5]$ . The trajectory solves $\partial_{t}u=\beta\partial_{xx}u+F(t,x)$ on a periodic interval of length 16, with the five-mode forcing coefficients shown beneath.

Figure 7: Wave-B. Wave speed $c=2$ , a left Dirichlet and a right Neumann boundary, and initial pulse center $s=-0.31821$ . The initial displacement and velocity are $u_{0}(x)=e^{-(x-s)^{2}}$ and $v_{0}(x)=-2c(x-s)u_{0}(x)$ . The plot shows $t\simeq 0$ – $12.05$ after restoring forward time from the reversed storage order, retaining the original nonuniform Chebyshev grid.

Figure 8: Combined Equation. ID: $(\alpha,\beta,\gamma)=(0.436745,0.185805,0.633549)$ ; OOD: $(\alpha,\beta,\gamma)=(1.652660,0.214112,0.414353)$ . The OOD nonlinear-transport coefficient exceeds the ID interval $[0,1]$ , while $\beta$ and $\gamma$ remain within their ID ranges $[0,0.4]$ and $[0,1]$ . Both trajectories use periodic boundaries, no external forcing, and the same saved coordinates on $x\in[0,16)$ and $t\in[0,10]$ . Their initial fields and all three coefficients differ, so every trajectories have different initial conditions.

Figure 9: Vorticity. ID: $\nu=0.00593069$ ; OOD: $\nu=5.90909\times 10^{-5}$ , below the training support $[10^{-3},10^{-2}]$ . ID and OOD both use initial-condition, and their first saved vorticity fields are exactly equal. Rows compare the same five saved-frame positions, with periodic boundaries and a shared symmetric color scale.

Figure 10: HeterNS. ID: $\nu\simeq 10^{-4}$ ; OOD: $\nu\simeq 3\times 10^{-3}$ , above the training range $[10^{-5},10^{-3}]$ . Both use $F(x,y)=0.1[\sin(2\pi(x+y))+\cos(2\pi(x+y))]$ .

Figure 11: Wave-2D:. ID: $(c,k)=(398.98990,43.93939)$ ; OOD: $(c,k)=(525,55)$ , outside the training ranges $c\in[100,500]$ and $k\in[0,50]$ .

Figure 12: Gray–Scott. ID: $(f,k)=(0.0364141,0.0627879)$ ; OOD: $(f,k)=(0.0459899,0.0585152)$ , outside the training ranges $f\in[0.023,0.045]$ and $k\in[0.059,0.064]$ .

Explore similar work

RD-JEPA: Predictive latent pretraining for few-trajectory transfer across reaction--diffusion equations
Sep 24, 2026Chenhao Si, Ming YanLatent Prediction Surrogate Models
School of Data Science, The Chinese University of Hong Kong, Shenzhen, Shenzhen, China
Disentangled Latent Dynamics Manifold Fusion for Solving Parameterized PDEs
Mar 13, 2026Zhangyong Liang, Huanhuan GaoNeural Partial Differential Equation Solvers Partial Differential Equations
National Center for Applied Mathematics, Tianjin University, Tianjin, 300072, PR China · School of Mechanical and Aerospace Engineering, Jilin University, Changchun, 130025, China
Learning Interpretable PDE Representations for Generative Reconstructions with Structured Sparsity
Apr 26, 2026Valerie Tsao, Nathaniel Chaney, Manolis VeveakisLatent Diffusion Model Spatiotemporal Fields
Department of Civil & Environmental Engineering, Duke University, NC, USA.

Sep 24, 2026, cs.AIJ/K move · Enter open · S save

RD-JEPA: Predictive latent pretraining for few-trajectory transfer across reaction--diffusion equations

Chenhao Si, Ming Yan

School of Data Science, The Chinese University of Hong Kong, Shenzhen, Shenzhen, China

Latent Prediction Surrogate Models Joint-Embedding Predictive Architectures Latent Trajectories S-Jepa Self-Supervised Learning

Open paper AskPDF arXiv

Learning surrogates for time-dependent partial differential equations often requires a new simulation corpus when the governing operator changes. We introduce RD-JEPA, a joint-embedding predictive architecture for self-supervised pretraining on reaction-diffusion trajectories. A single model is pretrained on five parameterized systems and then adapted to three held-out systems whose reaction operators and trajectories are excluded from pretraining. Using one, five, or ten complete trajectories from a held-out system, RD-JEPA achieves lower mean relative discrete

\ell^2

field error and mean absolute spatial first-difference error than five supervised surrogate baselines, an independently trained control that removes the trajectory-dependent predictive latent pathway, and an architecture-matched model trained from scratch. Within the evaluated equations, output resolution, forecast horizons, and choices of adaptation trajectories, the results indicate that prediction of future-state representations can support data-efficient adaptation across related reaction-diffusion systems.

Abstract

Explore similar work

RD-JEPA: Predictive latent pretraining for few-trajectory transfer across reaction--diffusion equations

Disentangled Latent Dynamics Manifold Fusion for Solving Parameterized PDEs

Learning Interpretable PDE Representations for Generative Reconstructions with Structured Sparsity