Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

PI-JEPA: Label-Free Surrogate Pretraining for Coupled Multiphysics Simulation via Operator-Split Latent Prediction

Brandon Yee, Pairie Koh

Abstract

Reservoir simulation workflows face a fundamental data asymmetry: input parameter fields (geostatistical permeability realizations, porosity distributions) are free to generate in arbitrary quantities, yet existing neural operator surrogates require large corpora of expensive labeled simulation trajectories and cannot exploit this unlabeled structure. We introduce \textbf{PI-JEPA} (Physics-Informed Joint Embedding Predictive Architecture), a surrogate pretraining framework that trains \emph{without any completed PDE solves}, using masked latent prediction on unlabeled parameter fields under per-sub-operator PDE residual regularization. The predictor bank is structurally aligned with the Lie--Trotter operator-splitting decomposition of the governing equations, dedicating a separate physics-constrained latent module to each sub-process (pressure, saturation transport, reaction), enabling fine-tuning with as few as 100 labeled simulation runs. On single-phase Darcy flow, PI-JEPA achieves $1.9\times$ lower error than FNO and $2.4\times$ lower error than DeepONet at $N_\ell{=}100$, with 24\% improvement over supervised-only training at $N_\ell{=}500$, demonstrating that label-free surrogate pretraining substantially reduces the simulation budget required for multiphysics surrogate deployment.

PI-JEPA: Label-Free Surrogate Pretraining for Coupled Multiphysics Simulation via Operator-Split Latent Prediction

Abstract

Reservoir simulation workflows face a fundamental data asymmetry: input parameter fields (geostatistical permeability realizations, porosity distributions) are free to generate in arbitrary quantities, yet existing neural operator surrogates require large corpora of expensive labeled simulation trajectories and cannot exploit this unlabeled structure. We introduce \textbf{PI-JEPA} (Physics-Informed Joint Embedding Predictive Architecture), a surrogate pretraining framework that trains \emph{without any completed PDE solves}, using masked latent prediction on unlabeled parameter fields under per-sub-operator PDE residual regularization. The predictor bank is structurally aligned with the Lie--Trotter operator-splitting decomposition of the governing equations, dedicating a separate physics-constrained latent module to each sub-process (pressure, saturation transport, reaction), enabling fine-tuning with as few as 100 labeled simulation runs. On single-phase Darcy flow, PI-JEPA achieves lower error than FNO and lower error than DeepONet at , with 24\% improvement over supervised-only training at , demonstrating that label-free surrogate pretraining substantially reduces the simulation budget required for multiphysics surrogate deployment.

Paper Structure

This paper contains 55 sections, 1 theorem, 16 equations, 5 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Neural Operators for PDE Surrogates
  4. Physics-Informed Neural Networks and Operators
  5. Self-Supervised Pretraining for Simulation Surrogates
  6. Graph-Based Simulation
  7. Latent-Space Reduced-Order Modeling for PDEs
  8. Surrogate Modeling for Subsurface Flow and Reactive Transport
  9. Surrogate Modeling Framework
  10. Problem Formulation
  11. Two-Phase Darcy Flow.
  12. Reactive Transport.
  13. The PI-JEPA Framework
  14. Fourier-Enhanced Encoder Architecture
  15. Patch Tokenization.
  16. ...and 40 more sections

Key Result

Proposition 1

Suppose the encoder $f_\theta$ is represented by a matrix $\Phi \in \mathbb{R}^{d \times n}$ with $\Phi\Phi^\top = I_d$, and each sub-operator $A_k$ has effective rank at most $r_k$ in the sense that $A_k = \Phi^\top B_k \Phi + E_k$ with $\|E_k\|_F \leq \delta_k$. Given $N_u$ unlabeled transition pa When $d \ll n$ and the projection error $\sum_k \delta_k^2$ is small (i.e., the encoder captures th

Figures (5)

  • Figure 1: PI-JEPA architecture overview. The solution field $\mathbf{u}(\mathbf{x},t)$ is partitioned into context and target patch sets. A context encoder $f_\theta$ and an EMA target encoder $f_\xi$ map these to latent codes $\mathbf{z}_c$ and $\mathbf{z}_t$, respectively. A bank of $K$ latent predictors $\{g_{\phi_k}\}$ predicts the target embeddings $\hat{\mathbf{z}}_t^{(k)}$ aligned to each sub-operator in the physical splitting. The total loss $\mathcal{L}$ combines a predictive term, a PDE residual physics term, and a covariance regularizer. The EMA update ensures the target encoder provides stable, slowly-evolving learning targets throughout self-supervised pretraining.
  • Figure 2: Operator splitting correspondence. The numerical Lie--Trotter splitting (top row) decomposes each timestep into sequential physical sub-operators $\mathcal{L}_1,\ldots,\mathcal{L}_K$. PI-JEPA's latent predictor bank (bottom row) mirrors this structure: predictor $g_{\phi_k}$ advances the latent state through the $k$-th sub-step, and a per-sub-operator PDE residual loss $\mathcal{L}_\mathrm{phys}^{(k)}$ regularizes each prediction. The dashed arrows indicate this physics supervision coupling. Illustrated here for $K=2$ (pressure + saturation), the scheme extends to $K=3$ for reactive transport by adding a reaction predictor.
  • Figure 3: Spatiotemporal block masking strategy. Context patches (blue) are selected from a contiguous subregion at time $t$. Target patches (orange) form a spatially displaced block at the subsequent timestep $t + \Delta t$. The predictor must anticipate the latent representation of the target region, implicitly learning the causal dynamics---advection, diffusion, or reaction---linking context to target. Unmasked patches (gray) contribute to the target encoder representation but are not used in the predictive loss.
  • Figure 4: Data efficiency on single-phase Darcy flow ($64\times64$). PI-JEPA (blue) outperforms FNO (red) and DeepONet (orange) for $N_\ell \leq 100$. The gap between PI-JEPA and the scratch baseline (gray, dashed) quantifies the benefit of self-supervised pretraining. FNO surpasses PI-JEPA at $N_\ell \geq 250$ where its spectral inductive bias is fully exploited. All results averaged over 3 seeds.
  • Figure :

Theorems & Definitions (4)

  • Remark 1: Target Encoder Stability
  • Proposition 1: Sample complexity reduction
  • proof
  • Remark 2