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

2026-09-24

AI Summary

This study addresses the bottleneck of regenerating simulation data for surrogate models of reaction-diffusion equations when operators change. We propose RD-JEPA, a framework based on joint-embedding predictive architectures with self-supervised pretraining. By innovatively predicting future state representations in latent space rather than retraining on complete trajectories, it enables efficient few-shot transfer across systems. Experiments demonstrate that, using only minimal trajectory data, RD-JEPA significantly outperforms five supervised baselines and from-scratch trained models in both field error and spatial differential error. This work establishes a new paradigm for generalization and adaptation in partial differential equation surrogates.

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 $\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.

Problem

• reaction-diffusion equations
• few-trajectory transfer
• surrogate modeling
• partial differential equations
• data-efficient adaptation

Innovation

• Joint-Embedding Predictive Architecture
• Self-supervised pretraining
• Few-trajectory transfer
• Reaction-diffusion equations
• Latent representation prediction