Model predicts future states to improve adapting reaction diffusion solvers
RD-JEPA: Predictive latent pretraining for few-trajectory transfer across reaction--diffusion equations
Artificial IntelligenceMachine Learning
Summary
Partial differential equations describe how things like chemicals spread and react over time, but solving them for new scenarios can take lots of computing power. The authors show a method called RD-JEPA that learns from several reaction-diffusion systems and then adapts quickly to new ones using only a few example simulations. Their approach predicts hidden future states to make better guesses about how the system evolves, improving accuracy with less data. This could help build faster simulation tools for related problems without needing to start from scratch.
What this means in practice
- •For computational physicists: Improve simulation tool accuracy by adapting pretrained models to new reaction-diffusion systems with minimal additional data.
- •For software engineers building simulation platforms: Incorporate pretrained predictive latent embeddings to speed up surrogate model training for time-dependent physics problems.
Authors
Chenhao Si, Ming Yan
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.