PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Chenhao Si, Ming Yan

arXiv:2609.29403v1cs.AIcs.LG

TL;DR

Learning surrogates across changing reaction operators raises the question of whether pretraining can yield useful dynamics representations when target trajectories are scarce. RD-JEPA predicts future-state representations during pretraining, then adapts to held-out reaction laws; its reported transfer ordering persists across budgets, horizons, and specialized surrogate architectures.

  • Problem

    Whether pretraining across reaction operators yields reusable dynamics representations when target-system trajectories are scarce is the central question.

    Source: Predictive latent pretraining targets… · p. 2

  • Method

    RD-JEPA predicts future-state representations from trajectory context during pretraining, then freezes the pretrained online encoder and adapts the predictor with a newly initialized full-field decoder.

    Source: Predictive latent pretraining targets… · p. 2

  • Results

    RD-JEPA's representation remained effective after adaptation to previously unseen reaction laws, and the mean-error ordering persisted across adaptation budgets, forecast horizons, and specialized surrogate architectures.

    Source: Discussion · pp. 4–5

  • Takeaways & Limitations

    A simulation corpus from related equations can be amortized into one representation and repeatedly specialized when target-system trajectories are scarce.

    Source: Discussion · p. 7

    The evaluation covers deterministic two-component systems at one two-dimensional resolution and five direct-forecast horizons; longer rollouts, additional resolutions, and repeated pretraining runs remain uncharacterized.

    Source: Discussion · pp. 6–7

Abstract

from arXiv · show

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.

Results

RD-JEPA was pretrained on five reaction–diffusion systems and adapted to three held-out systems whose reaction laws and trajectories were excluded from pretraining. It achieved lower reported errors across target equations and adaptation settings, with additional results on source-system reuse and boundary-condition shifts.

Predictive-latent pretraining improves source-system forecasting with limited adaptation data

  • Across source systems, RD-JEPA had the smallest mean relative discrete ℓ2 error at every evaluated budget–horizon configuration, while mean error rose with horizon and fell with more adaptation trajectories.At K = 5 and h = 5, it also had the smallest reported mean relative discrete ℓ2 field error on each of the five equations and their equal-weight average.Source: Predictive-latent pretraining improves… · p. 3
  • RD-JEPA had the lowest mean relative L2 error against FNO and the no-predictive-latent control for every target equation, adaptation budget, and forecast horizon.This ordering held after adaptation with one trajectory and through h = 5.Source: Predictive-latent representations… · p. 4
  • Removing the trajectory-specific encoder–predictor pathway while retaining the physical-context decoder route and lead-time conditioning left the control with higher mean field and gradient errors than RD-JEPA.This comparison supports a contribution from the complete predictive-latent representation pathway to few-trajectory transfer.Source: Predictive-latent representations… · p. 4
  • Across all nine target equation–budget settings, RD-JEPA had the lowest mean on both metrics among five supervised surrogates, with relative L2 essentially indistinguishable from RieszNO for Lambda–Omega at K = 10.The other eight settings had lower mean relative L2 error than the strongest external supervised baseline.Source: Cross-equation transfer persists… · p. 4
  • RD-JEPA had lower mean relative L2 and gradient L1 errors than the architecture-matched scratch model in all nine target-equation–budget settings.The comparison used matched support sets, downstream objectives, and update budgets.Source: Cross-equation transfer persists… · p. 4
  • For Barkley with homogeneous Neumann, Robin, and Dirichlet conditions, RD-JEPA had the lowest mean relative L2 and gradient L1 errors at K = 1, 5, and 10.This result covers the evaluated combined shift in boundary conditions and diffusion discretization, domain, resolution, and simulation protocol.Source: Cross-equation transfer persists… · p. 4

Discussion

The discussion presents predictive-latent pretraining as a reusable dynamics representation that remained effective after adaptation to unseen reaction laws. It frames representation-space prediction as a distinct cross-system pretraining target while delimiting the evidence to the evaluated setting.

Discussion

  • RD-JEPA's mean-error ordering persisted across adaptation budgets, forecast horizons, and specialized surrogate architectures after transfer to previously unseen reaction laws.The discussion describes the learned representation as remaining effective after adaptation to unseen reaction laws.Source: Discussion · pp. 4–5
  • Where prior cross-system approaches predict or reconstruct fields [27,28,33], RD-JEPA predicts future representations from trajectory context, testing a different locus for shared dynamical information.The discussion identifies direct comparisons among predictive, reconstructive, and contrastive objectives as a route for characterizing their relative transfer properties.Source: Discussion · p. 6
  • Diffusion-like and reaction-like latent pathways provide an inductive bias, not an exact operator splitting, and require no equation identifier, coefficient, or symbolic governing equation.The reported controls assess the complete source-pretrained pathway; component-wise analyses could further resolve how its pathways and objective interact.Source: Discussion · p. 6
  • The evaluation covers deterministic two-component systems at one 2D resolution and five direct forecast horizons, using few-trajectory adaptation rather than zero-shot inference.Variability estimates cover support-set selection for one pretraining checkpoint and one fixed training random-number-generator state; longer rollouts, more resolutions, and repeated pretraining runs remain needed to characterize beyond this setting.Source: Discussion · pp. 6–7
  • Within these limits, future-state prediction in representation space is presented as a viable route to domain-specific reusable dynamics models.The discussion describes amortizing a simulation corpus from related equations into one representation that can be specialized when target-system trajectories are scarce.Source: Discussion · p. 7

Methods

RD-JEPA was pretrained across five reaction–diffusion systems and evaluated on systems excluded from pretraining using complete-trajectory adaptation. Its representation-learning and forecasting stages use distinct objectives and evaluation protocols.

Reaction–diffusion datasets and experimental design

  • Pretraining used Gray–Scott, FitzHugh–Nagumo, Brusselator, complex Ginzburg–Landau, and Schnakenberg; Lambda–Omega, Barkley, and Oregonator were held out at the equation level.Each equation had 1,000 trajectories for pretraining or adaptation and 300 test trajectories; K counts complete trajectories, with three prespecified support-set selections.Source: Reaction–diffusion datasets and… · p. 8
  • A separate Barkley stress test evaluated homogeneous Neumann, Robin, and Dirichlet boundaries; the predictor retained its pretrained periodic neighbor map without a boundary-condition label or explicit target-domain operator.Each boundary-specific dataset used 20 adaptation trajectories and 300 test trajectories, with K ∈ {1,5,10}; all models used 5,000 optimization steps.Source: Reaction–diffusion datasets and… · p. 9, Joint-embedding predictive pretraining · p. 10, Few-trajectory adaptation and… · p. 10
  • One model uses four consecutive states to predict future-state representations, with diffusion-like and reaction-like latent updates conditioned on forecast horizon.A future-field target encoder follows the online encoder through an exponential moving average; the online encoder and predictor use no equation identity, coefficients, or symbolic equation.Source: Joint-embedding predictive pretraining · pp. 9–10
  • Pretraining sampled future targets up to eight stored-frame intervals ahead and used future-block, future-tube, and same-time masks to optimize a representation-space objective.Training lasted 200,000 optimization steps, and no field decoder was attached during pretraining.Source: Joint-embedding predictive pretraining · p. 10
  • For forecasting, the online encoder was frozen, the predictor was fine-tuned with a new dense decoder, and each requested horizon was predicted directly rather than autoregressively.RD-JEPA and the no-predictive-latent control shared a supervised multi-term objective; supervised baselines received the same support trajectories, context, and target horizons. The scratch control used the matched architecture and update budget without pretrained tensors or the JEPA objective.Source: Few-trajectory adaptation and… · p. 10
  • Evaluation used relative field error and spatial-gradient error; the latter averages absolute adjacent first differences across channels and valid neighboring grid pairs, without wrap-around.Both metrics were computed on the original field scale.Source: Evaluation and statistical aggregation · p. 10
  • Horizon-averaged scores weight h = 1,...,5 equally, and reported means and standard deviations summarize three paired support-set selections, not independent pretraining runs.All valid temporal windows were evaluated, with errors averaged within trajectories and then equally across 300 test trajectories; no formal hypothesis tests were performed.Source: Table 1 · p. 9, Evaluation and statistical aggregation · pp. 10–11

Data availability

The paper states that quantitative data and supporting materials will be deposited in a public archival repository.

Data availability

  • Supplementary Data 1–5 are to include run-level values, reported means and standard deviations, quantitative and qualitative display inputs, dataset splits, support-set indices, and plotting inputs.The planned deposit also includes simulation and evaluation code, parameter files, and recorded random seeds; its persistent repository identifier will be added before publication.Source: Data availability · p. 11

Code availability

The paper describes planned code access for editors and reviewers and public archiving upon publication.

Code availability

  • The code release is to cover simulation, predictive pretraining, adaptation, baseline training, evaluation, and figure generation, with environment specifications and result-to-source-data links.The anonymous repository is to be archived publicly upon publication, with a persistent version identifier added before publication.Source: Code availability · p. 11

Supplementary Information for

The supplied material gives the RD-JEPA paper title and an institutional line.

Supplementary Information for

Supplementary Note 1: Reaction–diffusion systems and dataset construction

The supplement specifies an eight-system reaction–diffusion benchmark, separating five pretraining systems from three fully held-out transfer systems, and details its numerical and boundary-condition datasets.

S1.2 Source systems used for predictive pretraining

S1.3 Held-out systems for equation-level transfer

S1.1 Benchmark composition, evaluation regimes, and notation

  • Each primary trajectory contains 40 two-channel fields on a 128×128 periodic grid, with 1,000 source-pool trajectories, 50 validation trajectories, and 300 test trajectories per source system.Source-system test sets include in-distribution, coefficient-out-of-distribution, and initial-condition-out-of-distribution regimes.Source: S1.1 Benchmark composition, evaluation… · p. 16

S1.4 Numerical integration and trajectory screening

  • The benchmark uses equation-specific integration schemes, summarized with temporal sampling in Table S1, and screens trajectories for invalid, divergent, nearly static, or unresolved dynamics.The numerical schemes differ across systems, including finite-difference time stepping, Runge–Kutta, and Strang splitting.Source: S1.4 Numerical integration and… · p. 18, Table S1 · p. 18

S1.5 Barkley benchmark with non-periodic boundary conditions

  • Three additional Barkley datasets test adaptation beyond periodic boundaries, each with 20 adaptation, five validation, and 300 test trajectories; only the activator requires a spatial boundary operator because Dv = 0.Source: S1.5 Barkley benchmark with… · p. 18
  • Boundary-specific Barkley datasets share matched generating distributions but are not paired trajectory-by-trajectory because screening was performed separately for each condition.Source: S1.5 Barkley benchmark with… · p. 19

Supplementary Note 2: RD-JEPA architecture and predictive latent learning

RD-JEPA separates predictive latent pretraining from full-field forecasting adaptation: it predicts future target representations from observed context, then uses predicted tokens to produce field forecasts.

S2.1 Overall two-stage RD-JEPA formulation

  • During predictive pretraining, an online encoder represents observed trajectory context, a target encoder represents future fields, and a predictor learns future target tokens at selected patches and time offsets.The target encoder tracks the online encoder through an exponential moving average.Source: S2.1 Overall two-stage RD-JEPA… · p. 20
  • During downstream adaptation, the encoder is frozen while the predictor and decoder are trained to turn predicted latent tokens and the latest observed field into future-field predictions.Source: S2.1 Overall two-stage RD-JEPA… · p. 20
  • The online encoder and predictor are trained jointly across all five source systems to match normalized predicted and target tokens.Source: S2.1 Overall two-stage RD-JEPA… · p. 20
  • Forecasts for horizons h ∈ {1,2,3,4,5} are produced directly from the same observed context rather than recursively feeding predictions back into the model.Source: S2.1 Overall two-stage RD-JEPA… · p. 20

S2.2 Online encoder Eθ: spatiotemporal tokenization and multiscale spatial encoding

This section specifies RD-JEPA’s encoding and prediction components for reaction–diffusion contexts, together with the few-trajectory adaptation protocol and downstream forecasting objectives. It also points to summaries of the architecture, parameter counts, and optimization settings.

S2.2 Online encoder Eθ: spatiotemporal tokenization and multiscale spatial encoding

  • The online encoder maps four 128×128 fields to a 16×16 latent grid through patch tokenization, per-position temporal aggregation, and five-stage multiscale spatial encoding.Each spatial position’s four frame tokens are aggregated by cross-attention; the spatial encoder then exchanges information across positions using attention and spectral–local mixing.Source: S2.2 Online encoder Eθ: spatiotemporal… · pp. 21–22

S2.3 Target sampling and EMA target encoding

  • Pretraining samples rectangular target regions and selects 256 patch-offset pairs per sample for latent prediction.The target encoder processes complete, unmasked target frames without gradient updates and is updated from the online encoder by an EMA rule.Source: S2.3 Target sampling and EMA target… · pp. 22–23

S2.4 Lead-time-conditioned predictor with diffusion-inspired and reaction-inspired branches

  • The predictor processes requested patch-offset queries through eight blocks that combine query and context attention with diffusion-inspired and reaction-inspired updates.These branches operate on learned latent representations, not as numerical discretizations of physical operators; the same periodic neighbor sets are used even in Barkley experiments with non-periodic physical boundaries.Source: S2.4 Lead-time-conditioned predictor… · pp. 23–24

S2.5 Predictive latent objective and pretraining

  • The predictive latent objective normalizes predicted and target vectors before averaging squared differences, and pretraining runs for 200,000 optimization iterations.The target encoder receives no gradients or optimizer updates; its EMA coefficient follows a cosine schedule from 0.996 to 0.99995.Source: S2.5 Predictive latent objective and… · pp. 24–25

S2.6 Dense forecasting and adaptation from few trajectories

  • For dense downstream forecasts, the pretrained encoder is held fixed while the predictor and randomly initialized decoder are fine-tuned, with five future fields predicted directly from the same four-frame context.The decoder combines the predicted latent grid with a multiscale feature pyramid from the latest context field to reconstruct each two-channel 128×128 field.Source: S2.6 Dense forecasting and adaptation… · p. 25
  • Each adaptation run uses a system-, trajectory-budget-, and support-set-specific predictor-decoder pair and 5,000 optimization steps; held-out systems use K = 1, 5, or 10 trajectories.For pretrained source systems, the evaluated budgets are K = 5, 10, or 20.Source: S2.6 Dense forecasting and adaptation… · p. 26
  • The downstream objective weights five forecast offsets equally; spatial first differences use adjacent values without grid-spacing normalization or boundary wrap-around, and directional averages are computed separately then added.The 2D real-input Fourier transform is applied independently to each field component; at 128 × 128 resolution it returns 128 × 65 coefficients, whose magnitudes receive equal weight without conjugate-frequency multiplicity weights.Source: S2.6 Dense forecasting and adaptation… · p. 26

S2.7 Architecture and optimization summary

Supplementary Note 3: Experimental protocols, baselines and statistics

The protocols define few-trajectory forecasting with matched support sets and fixed evaluation windows, compare supervised surrogates and two controls, and report field and gradient errors with uncertainty across support-set selections.

S3.1 Forecasting tasks, support selection and evaluation

  • Held-out systems use K = 1, 5, or 10 support trajectories, while source systems use K = 5, 10, or 20, with matched selections across methods.Primary benchmarks evaluate the same 300 test trajectories and 32 valid windows per trajectory; the boundary-condition experiment selects supports separately by boundary condition.Source: S3.1 Forecasting tasks, support… · p. 28

S3.2 Comparison models and controls

  • The supervised comparators are FNO, LNO, RieszNO, ReViT, and CNextU-Net, spanning neural operators, a transformer, and a convolutional encoder–decoder.Each is trained separately for each equation, support budget, and support-set selection.Source: S3.2 Comparison models and controls · p. 28
  • The no-predictive-latent control removes the sample-dependent predictive pathway, while a scratch control uses RD-JEPA’s architecture with randomly initialized modules.The scratch model jointly updates its encoder, predictor, and decoder, unlike RD-JEPA adaptation, which freezes the online encoder.Source: S3.2 Comparison models and controls · pp. 28–29

S3.3 Training objectives and optimization settings

  • All models use the same weighted forecasting loss and equal weighting across five lead times; primary-benchmark models are evaluated after 5,000 updates without validation-based checkpoint selection.RD-JEPA uses fixed pretrained parameters across runs, and the reported variation measures support-set sensitivity rather than variability across training or pretraining seeds.Source: S3.3 Training objectives and… · p. 29

S3.4 Evaluation metrics and statistical aggregation

  • Relative L2 and gradient L1 are the principal metrics; gradient errors use adjacent spatial first differences without wrap-around and average across channels and valid neighbors.MSE is retained as a secondary diagnostic.Source: S3.4 Evaluation metrics and… · pp. 29–30
  • Reported standard deviations are computed across three run-level means, so they measure support-set sensitivity rather than variation across test trajectories.Each run-level value summarizes the same fixed 300 test trajectories.Source: S3.4 Evaluation metrics and… · p. 30
  • Percentage reductions are averaged across paired support-set selections, with displayed centers and error bars giving the mean and sample standard deviation of those paired reductions.This is a mean of matched run-level reductions, not a ratio computed after pooling runs.Source: S3.4 Evaluation metrics and… · p. 30

Supplementary Note 4: Complete forecasting results on source systems

Across the five source equations, RD-JEPA has the lowest mean error among the three compared models for every evaluated system, support-set size, and forecast lead time. The supplement details the error averaging and reports relative-field, spatial-difference, and split-resolved results.

S4.1 Evaluation settings and error averaging

  • RD-JEPA has the lowest mean error on both metrics in all 75 source-system, support-set-size, and forecast-lead combinations among RD-JEPA, the no-predictive-latent control, and FNO.Errors are averaged over valid windows and 300 test trajectories; reported means and sample standard deviations span three support-set selections.Source: S4.1 Evaluation settings and error… · p. 31

S4.2 Complete relative-field-error results

S4.3 Errors in spatial first differences

S4.4 Split-resolved forecasting results

Supplementary Note 5: Paired predictive-latent control on held-out governing equations

The supplement compares RD-JEPA with an independently retrained control on governing equations excluded from pretraining, using matched adaptation and test trajectories. It describes horizon-averaged errors and paired percentage reductions, without reproducing the main-text comparison with other baselines.

S5.1 Evaluation coverage, data mapping and interpretation

  • On held-out governing equations, the independently retrained control removes the trajectory-specific JEPA encoder–predictor pathway while retaining the dense decoder and physical-context pathway.Matched support sets and test trajectories make this a comparison of the complete predictive-latent pathway, not an inference-time zeroing intervention.Source: S5.1 Evaluation coverage, data mapping… · p. 35

S5.2 Horizon-averaged paired comparison

  • Errors are averaged over the 300 test trajectories and five direct forecast horizons; paired percentage reductions are calculated per support-set selection, with positive values favoring RD-JEPA.The summaries report means and sample standard deviations across three paired selections; no formal hypothesis tests were performed.Source: S5.2 Horizon-averaged paired comparison · p. 35, Table S13 · p. 35

Supplementary Note 6: Matched control and boundary-condition stress test

On held-out equations, RD-JEPA outperforms an architecture-matched scratch model across the tested adaptation budgets. It also has the lowest reported errors in the evaluated Barkley boundary-condition settings.

S6.1 Architecture-matched training without predictive pretraining

  • Across three target equations and all three adaptation budgets, RD-JEPA has lower mean relative L2 and gradient L1 errors than the architecture-matched scratch model.Under matched support sets and update budgets, the ordering indicates that source pretraining supplies reusable information not recovered from the limited target trajectories alone.Source: S6.1 Architecture-matched training… · p. 36, Table S14 · p. 36

S6.2 Few-trajectory adaptation under non-periodic boundary conditions

  • RD-JEPA has the lowest mean relative L2 and gradient L1 error in all nine Barkley boundary-condition–budget settings: Neumann, Robin, and Dirichlet, at K = 1, 5, and 10.This extends transfer to a combined boundary-condition and discretization shift within Barkley at the evaluated equation, domain, resolution, and simulator; the predictor retains its pretrained periodic latent-neighbour map and receives no explicit boundary-condition label.Source: S6.2 Few-trajectory adaptation under… · p. 36, Table S15 · p. 37

Related papers

Loading 2609.29403v1…