cs.AISep 24, 2026

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

Authors: Chenhao Si, Ming Yan

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

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

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 ℓ2\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.

Figures & tables

Figure 1: a , During pretraining, a context encoder represents four observed fields and an exponential-moving-average target encoder represents a future field. A lead-time-conditioned predictor estimates the target representation from context; the loss is evaluated in latent space without reconstructing the future field. b , For downstream forecasting, the source-pretrained encoder is frozen and the predictor is adapted jointly with a dense decoder using KK target-system trajectories. The same checkpoint is evaluated on source equations and on governing equations excluded from pretraining, and it produces direct forecasts at h=1,…,5h=1,\ldots,5 .
Figure 2: Few-trajectory forecasting on the source reaction–diffusion systems. a , Direct forecasts from the prespecified representative downstream run for a FitzHugh–Nagumo test trajectory after adaptation with K=5K=5 trajectories. The two field components, uu and vv , are shown at forecast horizons h=1h=1 , 33 and 55 . For each component and horizon, columns show the reference field, the RD-JEPA prediction and the pointwise absolute error. Absolute-error colour limits are rescaled independently at each horizon, so magnitudes should be read from the corresponding colour bars rather than compared by colour intensity. The displayed trajectory was selected randomly from the held-out test set. b , Cross-equation qualitative comparison from the same representative run at K=5K=5 and h=5h=5 for Gray–Scott, Ginzburg–Landau, Schnakenberg and Brusselator. Columns show the reference solution, RD-JEPA, the no-predictive-latent control and an equation-specific FNO trained from scratch using the same adaptation budget. Fields within each row share a common colour scale. c , Mean relative L2L^{2} error at K=5K=5 and h=5h=5 , reported separately for each source equation and as an equal-weight mean across the five equations ( Overall ). Each run-level equation score is averaged over the same 300300 held-out test trajectories. Markers and error bars show the mean and sample standard deviation across three paired support-selection runs, respectively. The Overall score is first averaged across equations within each run. Lower values indicate better performance, and the horizontal axis is logarithmic.
Figure 3: Dependence of source-system forecasting accuracy on adaptation budget and forecast horizon. a , Mean relative L2L^{2} error as a function of forecast horizon h∈{1,…,5}h\in\{1,\ldots,5\} for adaptation budgets K=5K=5 , 1010 and 2020 . RD-JEPA is compared with the no-predictive-latent control (orange) and an equation-specific FNO (purple). Each run-level value is an equal-weight mean across the five source equations after averaging over 300300 fixed test trajectories per equation. Points and error bars show the mean and sample standard deviation across three paired support-selection runs. The vertical axis is logarithmic. b , Equation-level few-trajectory scaling at h=5h=5 . Relative L2L^{2} error is plotted against KK for each source equation and for the equal-weight mean across equations. Points and error bars again show the mean and sample standard deviation across the three runs. Both axes are logarithmic. c , Summary across all adaptation budgets and forecast horizons. The upper heatmap reports the across-run mean relative L2L^{2} error of RD-JEPA. The lower heatmap reports its percentage reduction relative to FNO, 100(1−errRD\mbox−JEPAerrFNO)100\left(1-\frac{\mathrm{err}_{\mathrm{RD\mbox{-}JEPA}}}{\mathrm{err}_{\mathrm{FNO}}}\right) . Positive values indicate lower error for RD-JEPA. Equation- and split-resolved across-run summaries are reported in Supplementary Tables S6 – S12 , with individual run-level values in Supplementary Data 1.
Figure 4: Representative one-trajectory forecasts on governing equations excluded from pretraining. a , Ground-truth and RD-JEPA forecasts of the uu field for Lambda–Omega, Barkley and Oregonator after adaptation with one target trajectory, shown at forecast horizons h=1h=1 , 33 and 55 . Ground truth and prediction share a common colour scale within each equation. The dashed line in each h=5h=5 ground-truth field marks the spatial transect used in panel b. b , Values of the uu field along the indicated transects at h=5h=5 , plotted against normalized spatial coordinate. Black solid lines denote ground truth, blue solid lines RD-JEPA and orange dashed lines the no-predictive-latent control. Aggregate three-run results over all 300300 held-out test trajectories are reported in Fig. 5 and Supplementary Table S13 .
Figure 5: Few-trajectory forecasting on governing equations excluded from pretraining. a , Mean relative L2L^{2} error versus forecast horizon h=1,…,5h=1,\ldots,5 for Lambda–Omega, Barkley and Oregonator (columns) after adaptation with K=1K=1 , 55 and 1010 trajectories (rows). RD-JEPA (blue) is compared with the no-predictive-latent control (orange) and FNO (purple). Each run-level value is averaged over the same 300300 held-out test trajectories. Points and error bars show the mean and sample standard deviation across three paired support-selection runs; vertical axes are logarithmic. b , Percentage reduction in horizon-averaged relative L2L^{2} error at K=1K=1 , computed within each matched run as 100(1−errRD\mbox−JEPAerrbaseline)100\left(1-\frac{\mathrm{err}_{\mathrm{RD\mbox{-}JEPA}}}{\mathrm{err}_{\mathrm{baseline}}}\right) , relative to FNO (purple) and the control (orange). Bars and error bars show the mean and sample standard deviation across the three paired run-level reductions. c , Corresponding paired reduction in horizon-averaged gradient L1L^{1} error. Positive values indicate lower error for RD-JEPA. The horizon-averaged full-model–control comparison is reported in Supplementary Table S13 , and the horizon-resolved run-level values underlying the figure are provided in Supplementary Data 2.
  Relative L2L^{2} ↓\downarrow (×10−2)(\times 10^{-2})Gradient L1L^{1} ↓\downarrow (×10−2)(\times 10^{-2})
PDEModelK=1K=1K=5K=5K=10K=10K=1K=1K=5K=5K=10K=10
Lambda–OmegaRD-JEPA12.24±2.58\mathbf{12.24\pm 2.58}9.85±0.79\mathbf{9.85\pm 0.79}8.38±1.46\mathbf{8.38\pm 1.46}0.68±0.12\mathbf{0.68\pm 0.12}0.52±0.05\mathbf{0.52\pm 0.05}0.45±0.08\mathbf{0.45\pm 0.08}
 FNO50.53±13.0650.53\pm 13.0625.48±1.0925.48\pm 1.0920.29±1.5920.29\pm 1.593.87±1.023.87\pm 1.021.89±0.071.89\pm 0.071.50±0.081.50\pm 0.08
 LNO34.56±8.1534.56\pm 8.1521.07±1.2621.07\pm 1.2616.65±2.1616.65\pm 2.162.51±0.552.51\pm 0.551.57±0.111.57\pm 0.111.37±0.161.37\pm 0.16
 ReViT22.64±4.1122.64\pm 4.1115.35±0.7215.35\pm 0.7213.61±1.3713.61\pm 1.372.85±0.222.85\pm 0.223.24±0.143.24\pm 0.143.04±0.133.04\pm 0.13
 RieszNO19.20±4.7219.20\pm 4.7213.41±0.8813.41\pm 0.888.47±1.148.47\pm 1.141.26±0.281.26\pm 0.280.88±0.060.88\pm 0.060.47±0.060.47\pm 0.06
Table 1: Forecasting accuracy on governing equations excluded from pretraining. KK denotes the number of independent target-system trajectories used for adaptation. Each run-level score is obtained by first averaging each trajectory’s error over h=1,…,5h=1,\ldots,5 and then averaging across the same 300300 held-out test trajectories. For every model, values are the mean ±\pm sample standard deviation across three paired support-selection runs. Values are multiplied by 100100 , and lower values are better. Boldface marks the lowest numerical mean for each equation and adaptation budget; no formal hypothesis tests were performed. The paired horizon-averaged comparison between RD-JEPA and the no-predictive-latent control is reported in Supplementary Table S13 .
SystemDomain and gridIntegrator and spatial discretizationInternal Δt\Delta tSaved window
Gray–Scott[−1,1)2[-1,1)^{2} , 128×128128\times 128Explicit Euler; second-order periodic central differences10−410^{-4}t∈[0.05,1.00]t\in[0.05,1.00]
FitzHugh–NagumoPeriodic 128×128128\times 128 lattice, Δx=1\Delta x=1RK4; fourth-order periodic finite difference Laplacian2×10−32\times 10^{-3}Δtsave=1\Delta t_{\mathrm{save}}=1
Brusselator[−1,1)2[-1,1)^{2} , 128×128128\times 128Explicit Euler; second-order periodic central differences10−310^{-3}t∈[2,4]t\in[2,4]
Complex Ginzburg–Landau[−1,1)2[-1,1)^{2} , 128×128128\times 128Explicit Euler; second-order periodic central differences10−310^{-3}t∈[2,10]t\in[2,10]
Schnakenberg[−1,1)2[-1,1)^{2} , 128×128128\times 128Explicit Euler; second-order periodic central differences10−310^{-3}t∈[2,4]t\in[2,4]
Lambda–Omega[−1,1)2[-1,1)^{2} , 128×128128\times 128Explicit Euler; second-order periodic central differences2×10−32\times 10^{-3}t∈[3,12]t\in[3,12]
Table S1: Numerical discretization and temporal sampling in the primary periodic benchmark. Each retained trajectory contains 40 saved two-channel fields.
Target typeProbabilitySelection of time offsets
Future block0.50H={h}\mathcal{H}=\{h\} , where hh is sampled uniformly from {1,…,8}\{1,\ldots,8\} .
Future tube0.40H\mathcal{H} contains four distinct offsets sampled uniformly without replacement from {1,…,8}\{1,\ldots,8\} . The same spatial region R\mathcal{R} is used at each selected offset.
Same time0.10H={0}\mathcal{H}=\{0\} , targeting the latest context frame xtx_{t} . The patches in R\mathcal{R} are masked in the online encoder input.
Table S2: Target selection during predictive pretraining. All three target types use the same spatial-region sampling procedure to obtain R\mathcal{R} . The selected offset set H\mathcal{H} defines the candidate set C=R×H\mathcal{C}=\mathcal{R}\times\mathcal{H} , from which M=256M=256 query pairs are subsampled for loss computation.
ComponentInputOutput
Online encoder4×128×128×24\times 128\times 128\times 2256×512256\times 512
Target encoder1×128×128×21\times 128\times 128\times 2256×512256\times 512
Latent predictor256×512256\times 512 context tokens and 256256 query pairs (i,h)(i,h)256×512256\times 512 predicted tokens
Forecast decoder256×512256\times 512 predicted tokens and the latest 128×128×2128\times 128\times 2 context field128×128×2128\times 128\times 2
Table S3: RD-JEPA input and output dimensions and parameter counts. Shapes omit the batch dimension. The decoder output corresponds to one forecast offset. Parameter counts are reported in millions (M). Totals are computed before rounding the displayed component counts. The pretraining total includes the EMA target encoder.
SettingValue
Shared settings
OptimizerAdamW, (β1,β2)=(0.9,0.95)(\beta_{1},\beta_{2})=(0.9,0.95)
Batch size4
Gradient clippingGlobal norm 1.0
Predictive pretraining
Trainable modulesOnline encoder and predictor
Table S4: Optimization settings for predictive pretraining and downstream adaptation. Each training iteration uses one batch. The downstream settings apply to RD-JEPA adaptation runs in the primary benchmark, for both source systems and systems excluded from pretraining.
SettingRD-JEPANo predictive latentExternal baselines
Training updates5,0005,0005,000
Batch size444
OptimizerAdamWAdamWAdamW
Learning rate10−510^{-5} predictor; 2×10−42\times 10^{-4} decoder10−410^{-4} conditioner; 2×10−42\times 10^{-4} decoder2×10−42\times 10^{-4}
Weight decay10−410^{-4}10−410^{-4}10−410^{-4}
Adam betas(0.9,0.95)(0.9,0.95)(0.9,0.95)(0.9,0.95)(0.9,0.95)(0.9,0.95)
Table S5: Downstream optimization settings for the primary benchmark. The external-baseline column applies to FNO, LNO, RieszNO, ReViT and CNextU-Net.
PDEModelh=1h=1h=2h=2h=3h=3h=4h=4h=5h=5Avg.
Gray–ScottRD-JEPA3.56±1.21\mathbf{3.56}\pm\mathbf{1.21}4.66±0.97\mathbf{4.66}\pm\mathbf{0.97}5.74±0.84\mathbf{5.74}\pm\mathbf{0.84}6.83±0.75\mathbf{6.83}\pm\mathbf{0.75}7.99±0.66\mathbf{7.99}\pm\mathbf{0.66}5.76±0.87\mathbf{5.76}\pm\mathbf{0.87}
No predictive latent6.54±1.686.54\pm 1.6810.13±2.3810.13\pm 2.3813.40±2.9713.40\pm 2.9716.35±3.1916.35\pm 3.1919.00±2.8519.00\pm 2.8513.08±2.5813.08\pm 2.58
FNO6.10±0.496.10\pm 0.497.30±0.227.30\pm 0.229.36±0.569.36\pm 0.5611.55±0.9311.55\pm 0.9313.71±1.2913.71\pm 1.299.60±0.469.60\pm 0.46
FitzHugh–NagumoRD-JEPA4.11±1.38\mathbf{4.11}\pm\mathbf{1.38}4.59±1.34\mathbf{4.59}\pm\mathbf{1.34}5.36±1.20\mathbf{5.36}\pm\mathbf{1.20}6.43±1.08\mathbf{6.43}\pm\mathbf{1.08}7.75±1.02\mathbf{7.75}\pm\mathbf{1.02}5.65±1.19\mathbf{5.65}\pm\mathbf{1.19}
No predictive latent12.00±3.3012.00\pm 3.3013.31±3.4813.31\pm 3.4815.86±3.8415.86\pm 3.8419.09±3.9419.09\pm 3.9423.00±4.2623.00\pm 4.2616.65±3.7416.65\pm 3.74
FNO21.99±0.6521.99\pm 0.6521.49±0.8821.49\pm 0.8821.79±1.1021.79\pm 1.1022.64±1.3422.64\pm 1.3424.08±1.6024.08\pm 1.6022.40±1.1122.40\pm 1.11
Table S6: Complete relative L2L^{2} results on the five source systems at K=5K=5 . Entries are 100×100\times the error and report the mean ±\pm sample s.d. across three paired support-set selections. Within each run, errors are first averaged over all valid windows within each trajectory and then over the same 300 test trajectories (100 from each test regime). The Avg. column averages h=1,…,5h=1,\ldots,5 within each run before across-run aggregation. Lower values are better. Boldface marks the lowest numerical mean within each equation and horizon; no formal hypothesis tests were performed.
PDEModelh=1h=1h=2h=2h=3h=3h=4h=4h=5h=5Avg.
Gray–ScottRD-JEPA2.08±0.88\mathbf{2.08}\pm\mathbf{0.88}2.96±1.07\mathbf{2.96}\pm\mathbf{1.07}3.91±1.25\mathbf{3.91}\pm\mathbf{1.25}4.85±1.38\mathbf{4.85}\pm\mathbf{1.38}5.84±1.39\mathbf{5.84}\pm\mathbf{1.39}3.93±1.19\mathbf{3.93}\pm\mathbf{1.19}
No predictive latent4.89±2.134.89\pm 2.137.95±3.767.95\pm 3.7610.96±4.8910.96\pm 4.8913.84±5.6613.84\pm 5.6616.24±5.9416.24\pm 5.9410.78±4.4710.78\pm 4.47
FNO4.74±0.864.74\pm 0.865.77±1.045.77\pm 1.047.40±1.177.40\pm 1.179.13±1.219.13\pm 1.2110.90±1.2010.90\pm 1.207.59±1.097.59\pm 1.09
FitzHugh–NagumoRD-JEPA2.64±0.56\mathbf{2.64}\pm\mathbf{0.56}3.37±0.59\mathbf{3.37}\pm\mathbf{0.59}4.23±0.80\mathbf{4.23}\pm\mathbf{0.80}5.27±1.04\mathbf{5.27}\pm\mathbf{1.04}6.49±1.17\mathbf{6.49}\pm\mathbf{1.17}4.40±0.83\mathbf{4.40}\pm\mathbf{0.83}
No predictive latent7.39±1.787.39\pm 1.789.19±1.819.19\pm 1.8111.85±2.2711.85\pm 2.2714.93±3.0614.93\pm 3.0618.40±3.7018.40\pm 3.7012.35±2.5012.35\pm 2.50
FNO17.07±2.4217.07\pm 2.4216.65±2.2116.65\pm 2.2116.78±2.0016.78\pm 2.0017.33±1.7817.33\pm 1.7818.46±1.5818.46\pm 1.5817.26±1.9917.26\pm 1.99
Table S7: Complete relative L2L^{2} results on the five source systems at K=10K=10 . Entries are 100×100\times the error and report the mean ±\pm sample s.d. across three paired support-set selections. Within each run, errors are first averaged over all valid windows within each trajectory and then over the same 300 test trajectories (100 from each test regime). The Avg. column averages h=1,…,5h=1,\ldots,5 within each run before across-run aggregation. Lower values are better. Boldface marks the lowest numerical mean within each equation and horizon; no formal hypothesis tests were performed.
PDEModelh=1h=1h=2h=2h=3h=3h=4h=4h=5h=5Avg.
Gray–ScottRD-JEPA1.42±0.04\mathbf{1.42}\pm\mathbf{0.04}1.87±0.09\mathbf{1.87}\pm\mathbf{0.09}2.48±0.18\mathbf{2.48}\pm\mathbf{0.18}3.19±0.26\mathbf{3.19}\pm\mathbf{0.26}4.00±0.36\mathbf{4.00}\pm\mathbf{0.36}2.59±0.17\mathbf{2.59}\pm\mathbf{0.17}
No predictive latent2.43±0.242.43\pm 0.243.97±0.493.97\pm 0.495.53±0.755.53\pm 0.757.19±1.007.19\pm 1.008.98±1.158.98\pm 1.155.62±0.715.62\pm 0.71
FNO3.36±0.533.36\pm 0.534.24±0.644.24\pm 0.645.65±0.725.65\pm 0.727.25±0.817.25\pm 0.818.91±0.878.91\pm 0.875.88±0.715.88\pm 0.71
FitzHugh–NagumoRD-JEPA2.24±0.32\mathbf{2.24}\pm\mathbf{0.32}2.55±0.08\mathbf{2.55}\pm\mathbf{0.08}3.12±0.11\mathbf{3.12}\pm\mathbf{0.11}3.90±0.13\mathbf{3.90}\pm\mathbf{0.13}4.84±0.15\mathbf{4.84}\pm\mathbf{0.15}3.33±0.11\mathbf{3.33}\pm\mathbf{0.11}
No predictive latent7.64±0.407.64\pm 0.408.74±0.348.74\pm 0.3411.16±0.6111.16\pm 0.6113.96±0.7213.96\pm 0.7217.20±0.5817.20\pm 0.5811.74±0.5011.74\pm 0.50
FNO13.95±1.4313.95\pm 1.4313.61±1.2313.61\pm 1.2313.78±0.9913.78\pm 0.9914.52±0.9114.52\pm 0.9115.61±0.7915.61\pm 0.7914.29±1.0714.29\pm 1.07
Table S8: Complete relative L2L^{2} results on the five source systems at K=20K=20 . Entries are 100×100\times the error and report the mean ±\pm sample s.d. across three paired support-set selections. Within each run, errors are first averaged over all valid windows within each trajectory and then over the same 300 test trajectories (100 from each test regime). The Avg. column averages h=1,…,5h=1,\ldots,5 within each run before across-run aggregation. Lower values are better. Boldface marks the lowest numerical mean within each equation and horizon; no formal hypothesis tests were performed.
PDEModelh=1h=1h=2h=2h=3h=3h=4h=4h=5h=5Avg.
Gray–ScottRD-JEPA0.35±0.10\mathbf{0.35}\pm\mathbf{0.10}0.53±0.12\mathbf{0.53}\pm\mathbf{0.12}0.70±0.12\mathbf{0.70}\pm\mathbf{0.12}0.86±0.11\mathbf{0.86}\pm\mathbf{0.11}1.02±0.09\mathbf{1.02}\pm\mathbf{0.09}0.69±0.11\mathbf{0.69}\pm\mathbf{0.11}
No predictive latent0.72±0.200.72\pm 0.201.21±0.321.21\pm 0.321.63±0.391.63\pm 0.392.00±0.412.00\pm 0.412.32±0.362.32\pm 0.361.57±0.331.57\pm 0.33
FNO0.69±0.020.69\pm 0.020.88±0.030.88\pm 0.031.16±0.051.16\pm 0.051.44±0.061.44\pm 0.061.72±0.061.72\pm 0.061.18±0.041.18\pm 0.04
FitzHugh–NagumoRD-JEPA0.16±0.04\mathbf{0.16}\pm\mathbf{0.04}0.19±0.04\mathbf{0.19}\pm\mathbf{0.04}0.23±0.04\mathbf{0.23}\pm\mathbf{0.04}0.28±0.03\mathbf{0.28}\pm\mathbf{0.03}0.34±0.03\mathbf{0.34}\pm\mathbf{0.03}0.24±0.04\mathbf{0.24}\pm\mathbf{0.04}
No predictive latent0.44±0.110.44\pm 0.110.53±0.120.53\pm 0.120.65±0.130.65\pm 0.130.80±0.140.80\pm 0.140.96±0.160.96\pm 0.160.68±0.130.68\pm 0.13
FNO1.16±0.051.16\pm 0.051.18±0.051.18\pm 0.051.23±0.061.23\pm 0.061.29±0.061.29\pm 0.061.38±0.071.38\pm 0.071.25±0.061.25\pm 0.06
Table S9: Complete gradient L1L^{1} results on the five source systems at K=5K=5 . Entries are 100×100\times the error and report the mean ±\pm sample s.d. across three paired support-set selections. Within each run, errors are first averaged over all valid windows within each trajectory and then over the same 300 test trajectories (100 from each test regime). The Avg. column averages h=1,…,5h=1,\ldots,5 within each run before across-run aggregation. Lower values are better. Boldface marks the lowest numerical mean within each equation and horizon; no formal hypothesis tests were performed.
PDEModelh=1h=1h=2h=2h=3h=3h=4h=4h=5h=5Avg.
Gray–ScottRD-JEPA0.21±0.07\mathbf{0.21}\pm\mathbf{0.07}0.35±0.11\mathbf{0.35}\pm\mathbf{0.11}0.49±0.14\mathbf{0.49}\pm\mathbf{0.14}0.63±0.17\mathbf{0.63}\pm\mathbf{0.17}0.77±0.17\mathbf{0.77}\pm\mathbf{0.17}0.49±0.13\mathbf{0.49}\pm\mathbf{0.13}
No predictive latent0.52±0.250.52\pm 0.250.93±0.470.93\pm 0.471.33±0.621.33\pm 0.621.68±0.721.68\pm 0.721.99±0.741.99\pm 0.741.29±0.561.29\pm 0.56
FNO0.53±0.030.53\pm 0.030.70±0.080.70\pm 0.080.96±0.100.96\pm 0.101.23±0.121.23\pm 0.121.50±0.121.50\pm 0.120.98±0.090.98\pm 0.09
FitzHugh–NagumoRD-JEPA0.12±0.02\mathbf{0.12}\pm\mathbf{0.02}0.15±0.02\mathbf{0.15}\pm\mathbf{0.02}0.19±0.02\mathbf{0.19}\pm\mathbf{0.02}0.24±0.03\mathbf{0.24}\pm\mathbf{0.03}0.29±0.03\mathbf{0.29}\pm\mathbf{0.03}0.20±0.03\mathbf{0.20}\pm\mathbf{0.03}
No predictive latent0.33±0.030.33\pm 0.030.41±0.040.41\pm 0.040.53±0.060.53\pm 0.060.66±0.080.66\pm 0.080.80±0.100.80\pm 0.100.55±0.060.55\pm 0.06
FNO0.86±0.040.86\pm 0.040.87±0.030.87\pm 0.030.90±0.020.90\pm 0.020.95±0.020.95\pm 0.021.02±0.021.02\pm 0.020.92±0.020.92\pm 0.02
Table S10: Complete gradient L1L^{1} results on the five source systems at K=10K=10 . Entries are 100×100\times the error and report the mean ±\pm sample s.d. across three paired support-set selections. Within each run, errors are first averaged over all valid windows within each trajectory and then over the same 300 test trajectories (100 from each test regime). The Avg. column averages h=1,…,5h=1,\ldots,5 within each run before across-run aggregation. Lower values are better. Boldface marks the lowest numerical mean within each equation and horizon; no formal hypothesis tests were performed.
PDEModelh=1h=1h=2h=2h=3h=3h=4h=4h=5h=5Avg.
Gray–ScottRD-JEPA0.16±0.03\mathbf{0.16}\pm\mathbf{0.03}0.23±0.03\mathbf{0.23}\pm\mathbf{0.03}0.32±0.04\mathbf{0.32}\pm\mathbf{0.04}0.41±0.05\mathbf{0.41}\pm\mathbf{0.05}0.52±0.06\mathbf{0.52}\pm\mathbf{0.06}0.33±0.04\mathbf{0.33}\pm\mathbf{0.04}
No predictive latent0.29±0.020.29\pm 0.020.48±0.070.48\pm 0.070.69±0.100.69\pm 0.100.92±0.130.92\pm 0.131.18±0.161.18\pm 0.160.71±0.100.71\pm 0.10
FNO0.40±0.020.40\pm 0.020.53±0.040.53\pm 0.040.76±0.070.76\pm 0.071.01±0.101.01\pm 0.101.27±0.111.27\pm 0.110.80±0.070.80\pm 0.07
FitzHugh–NagumoRD-JEPA0.10±0.01\mathbf{0.10}\pm\mathbf{0.01}0.12±0.00\mathbf{0.12}\pm\mathbf{0.00}0.15±0.01\mathbf{0.15}\pm\mathbf{0.01}0.19±0.01\mathbf{0.19}\pm\mathbf{0.01}0.24±0.01\mathbf{0.24}\pm\mathbf{0.01}0.16±0.01\mathbf{0.16}\pm\mathbf{0.01}
No predictive latent0.35±0.020.35\pm 0.020.41±0.010.41\pm 0.010.51±0.020.51\pm 0.020.63±0.030.63\pm 0.030.76±0.030.76\pm 0.030.53±0.020.53\pm 0.02
FNO0.67±0.030.67\pm 0.030.67±0.030.67\pm 0.030.70±0.030.70\pm 0.030.74±0.030.74\pm 0.030.80±0.030.80\pm 0.030.72±0.030.72\pm 0.03
Table S11: Complete gradient L1L^{1} results on the five source systems at K=20K=20 . Entries are 100×100\times the error and report the mean ±\pm sample s.d. across three paired support-set selections. Within each run, errors are first averaged over all valid windows within each trajectory and then over the same 300 test trajectories (100 from each test regime). The Avg. column averages h=1,…,5h=1,\ldots,5 within each run before across-run aggregation. Lower values are better. Boldface marks the lowest numerical mean within each equation and horizon; no formal hypothesis tests were performed.
  Relative L2L^{2}Gradient L1L^{1}
ModelKKIDCoeff.-OODIC-OODIDCoeff.-OODIC-OOD
RD-JEPA56.61±1.06\mathbf{6.61\pm 1.06}6.89±1.14\mathbf{6.89\pm 1.14}6.88±1.23\mathbf{6.88\pm 1.23}0.79±0.10\mathbf{0.79\pm 0.10}0.85±0.10\mathbf{0.85\pm 0.10}0.86±0.09\mathbf{0.86\pm 0.09}
104.69±0.36\mathbf{4.69\pm 0.36}4.87±0.45\mathbf{4.87\pm 0.45}4.84±0.37\mathbf{4.84\pm 0.37}0.60±0.07\mathbf{0.60\pm 0.07}0.64±0.09\mathbf{0.64\pm 0.09}0.65±0.07\mathbf{0.65\pm 0.07}
203.44±0.10\mathbf{3.44\pm 0.10}3.52±0.11\mathbf{3.52\pm 0.11}3.64±0.11\mathbf{3.64\pm 0.11}0.45±0.02\mathbf{0.45\pm 0.02}0.46±0.01\mathbf{0.46\pm 0.01}0.48±0.01\mathbf{0.48\pm 0.01}
No predictive latent520.68±4.5520.68\pm 4.5521.73±5.3421.73\pm 5.3422.88±6.0422.88\pm 6.042.29±0.152.29\pm 0.152.41±0.142.41\pm 0.142.51±0.192.51\pm 0.19
1012.48±1.9612.48\pm 1.9613.28±2.8113.28\pm 2.8114.28±4.7214.28\pm 4.721.27±0.141.27\pm 0.141.34±0.211.34\pm 0.211.38±0.211.38\pm 0.21
Table S12: Split-resolved source-system forecasting. Each entry reports the mean ±\pm sample s.d. across three paired support-set selections. Within each run, errors are averaged over the indicated 100-trajectory split and then given equal weight across the five source equations and five forecast horizons. Values are multiplied by 100; lower values are better. Boldface marks the lowest numerical mean within each split and adaptation budget; no formal hypothesis tests were performed.
  Relative L2L^{2}Gradient L1L^{1}
PDEKKRD-JEPANo predictive latentReduction (%)RD-JEPANo predictive latentReduction (%)
Lambda–Omega112.24±2.58\mathbf{12.24\pm 2.58}23.06±5.2823.06\pm 5.2846.78±1.2246.78\pm 1.220.68±0.12\mathbf{0.68\pm 0.12}1.24±0.281.24\pm 0.2845.11±2.1745.11\pm 2.17
59.85±0.79\mathbf{9.85\pm 0.79}18.73±1.3218.73\pm 1.3247.38±3.4247.38\pm 3.420.52±0.05\mathbf{0.52\pm 0.05}0.97±0.050.97\pm 0.0546.39±4.6946.39\pm 4.69
108.38±1.46\mathbf{8.38\pm 1.46}17.84±1.6817.84\pm 1.6853.16±5.4653.16\pm 5.460.45±0.08\mathbf{0.45\pm 0.08}0.94±0.100.94\pm 0.1052.01±4.7652.01\pm 4.76
Barkley110.53±2.55\mathbf{10.53\pm 2.55}18.97±4.0418.97\pm 4.0444.72±3.0444.72\pm 3.041.09±0.27\mathbf{1.09\pm 0.27}1.97±0.421.97\pm 0.4244.60±2.6444.60\pm 2.64
56.90±0.24\mathbf{6.90\pm 0.24}12.97±0.8512.97\pm 0.8546.76±1.7646.76\pm 1.760.74±0.04\mathbf{0.74\pm 0.04}1.37±0.121.37\pm 0.1246.35±2.4646.35\pm 2.46
Table S13: Paired comparison with the no-predictive-latent control on governing equations excluded from pretraining. Error entries are 100×100\times the error and report the mean ±\pm sample s.d. across three paired support-set selections. Within each selection, errors are first averaged over the same 300 test trajectories and then with equal weight over the five direct forecast horizons. Reduction entries are percentages computed within each paired selection as 100(1−ERD\mbox−JEPA/Econtrol)100(1-E_{\mathrm{RD\mbox{-}JEPA}}/E_{\mathrm{control}}) and then summarized across the three paired values. Positive reductions favour RD-JEPA. Lower error values are better; boldface marks the lower numerical mean in each paired comparison; no formal hypothesis tests were performed.
  Relative L2L^{2} ↓\downarrow (×10−2)(\times 10^{-2})Gradient L1L^{1} ↓\downarrow (×10−2)(\times 10^{-2})
PDEModelK=1K=1K=5K=5K=10K=10K=1K=1K=5K=5K=10K=10
Lambda–OmegaRD-JEPA12.24±2.58\mathbf{12.24\pm 2.58}9.85±0.79\mathbf{9.85\pm 0.79}8.38±1.46\mathbf{8.38\pm 1.46}0.68±0.12\mathbf{0.68\pm 0.12}0.52±0.05\mathbf{0.52\pm 0.05}0.45±0.08\mathbf{0.45\pm 0.08}
Architecture-matched scratch23.42±4.7323.42\pm 4.7320.14±0.9120.14\pm 0.9119.11±1.7919.11\pm 1.791.27±0.261.27\pm 0.261.05±0.041.05\pm 0.041.00±0.121.00\pm 0.12
BarkleyRD-JEPA10.53±2.55\mathbf{10.53\pm 2.55}6.90±0.24\mathbf{6.90\pm 0.24}5.23±0.32\mathbf{5.23\pm 0.32}1.09±0.27\mathbf{1.09\pm 0.27}0.74±0.04\mathbf{0.74\pm 0.04}0.58±0.04\mathbf{0.58\pm 0.04}
Architecture-matched scratch19.23±3.4819.23\pm 3.4813.03±0.4813.03\pm 0.4810.01±1.1210.01\pm 1.122.01±0.382.01\pm 0.381.39±0.061.39\pm 0.061.10±0.121.10\pm 0.12
OregonatorRD-JEPA11.36±1.34\mathbf{11.36\pm 1.34}4.83±0.79\mathbf{4.83\pm 0.79}4.05±0.47\mathbf{4.05\pm 0.47}0.47±0.05\mathbf{0.47\pm 0.05}0.26±0.06\mathbf{0.26\pm 0.06}0.22±0.02\mathbf{0.22\pm 0.02}
Table S14: Architecture-matched control trained without predictive pretraining on governing equations excluded from pretraining. Each entry is 100×100\times the error and reports the mean ±\pm sample s.d. across three paired support-set selections. For each selection, errors are first averaged over the same 300 test trajectories and then over the five direct forecast horizons. The model-training random-number-generator state is fixed, and both models use the same support trajectories within each selection. Lower values are better. Boldface marks the lower numerical mean in each paired comparison; no formal hypothesis tests were performed.
  Relative L2L^{2} ↓\downarrow (×10−2)(\times 10^{-2})Gradient L1L^{1} ↓\downarrow (×10−2)(\times 10^{-2})
Boundary conditionModelK=1K=1K=5K=5K=10K=10K=1K=1K=5K=5K=10K=10
NeumannRD-JEPA10.79±1.42\mathbf{10.79\pm 1.42}7.33±1.30\mathbf{7.33\pm 1.30}5.49±0.46\mathbf{5.49\pm 0.46}0.68±0.10\mathbf{0.68\pm 0.10}0.48±0.08\mathbf{0.48\pm 0.08}0.38±0.03\mathbf{0.38\pm 0.03}
No predictive latent26.88±5.5826.88\pm 5.5817.78±3.2417.78\pm 3.2410.83±0.9110.83\pm 0.911.69±0.331.69\pm 0.331.16±0.181.16\pm 0.180.75±0.090.75\pm 0.09
FNO30.89±1.5030.89\pm 1.5019.15±2.3219.15\pm 2.3215.02±0.7615.02\pm 0.762.02±0.092.02\pm 0.091.21±0.131.21\pm 0.130.97±0.040.97\pm 0.04
RobinRD-JEPA11.13±1.35\mathbf{11.13\pm 1.35}7.21±1.26\mathbf{7.21\pm 1.26}5.57±0.39\mathbf{5.57\pm 0.39}0.69±0.10\mathbf{0.69\pm 0.10}0.47±0.08\mathbf{0.47\pm 0.08}0.39±0.02\mathbf{0.39\pm 0.02}
No predictive latent25.61±4.4625.61\pm 4.4617.91±3.1317.91\pm 3.1310.79±2.6710.79\pm 2.671.61±0.281.61\pm 0.281.18±0.171.18\pm 0.170.76±0.190.76\pm 0.19
Table S15: Few-trajectory Barkley forecasting under non-periodic boundary conditions. Each entry is 100×100\times the error and reports the mean ±\pm sample s.d. across three paired support-set selections. For each selection, errors are first averaged over the same 300 test trajectories and then over the five direct forecast horizons. The model-training random-number-generator state is fixed, and all models use the same support trajectories within each selection. Lower values are better. Boldface marks the lowest mean for each boundary condition, metric and adaptation budget.

Explore similar work

CardsList

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

    Sep 28, 2026Zhentao Tan, Jianrong Zhang, Ruijie Quan +1Latent PredictionPartial Differential Equations

    Zhejiang University.

  2. Recovering Physical Dynamics from Discrete Observations via Intrinsic Differential Consistency

    May 8, 2026Yuxiang Luo, Andrew PerraultNeural Ordinary Differential EquationsVelocity Field

    Department of Computer Science and Engineering The Ohio State University

  3. Learnable composition for neural operators

    Sep 2, 2026Zituo Chen, Baiming Zhang, Sili DengNeural OperatorsNeural Fields

    MIT MechE · University of Toronto