Adaptive Physics Transformer with Fused Global-Local Attention for Subsurface Energy Systems

TL;DR

The paper tackles the computational cost and rigidity of full-physics subsurface simulations by introducing the Adaptive Physics Transformer (APT), a mesh-agnostic neural operator that fuses global attention with local graph-based encoding to capture multi-scale subsurface dynamics. By directly learning from adaptive meshes and enabling cross-dataset pretraining, APT achieves state-of-the-art accuracy across irregular, Cartesian, and adaptive grids, while delivering substantial speedups over traditional simulators. Key contributions include the fused Global Perceiver and Local GNO encoder design, a gated fusion mechanism, direct learning from dynamic meshes, and demonstrated robustness in out-of-distribution and cross-dataset settings, positioning APT as a scalable backbone for subsurface foundation models. The practical impact lies in enabling fast, accurate, and transferable surrogate models for carbon storage, geothermal energy, and hydrocarbon systems, thereby accelerating decision-making and energy-transition deployment.

Abstract

The Earth's subsurface is a cornerstone of modern society, providing essential energy resources like hydrocarbons, geothermal, and minerals while serving as the primary reservoir for $CO_2$ sequestration. However, full physics numerical simulations of these systems are notoriously computationally expensive due to geological heterogeneity, high resolution requirements, and the tight coupling of physical processes with distinct propagation time scales. Here we propose the \textbf{Adaptive Physics Transformer} (APT), a geometry-, mesh-, and physics-agnostic neural operator that explicitly addresses these challenges. APT fuses a graph-based encoder to extract high-resolution local heterogeneous features with a global attention mechanism to resolve long-range physical impacts. Our results demonstrate that APT outperforms state-of-the-art architectures in subsurface tasks across both regular and irregular grids with robust super-resolution capabilities. Notably, APT is the first architecture that directly learns from adaptive mesh refinement simulations. We also demonstrate APT's capability for cross-dataset learning, positioning it as a robust and scalable backbone for large-scale subsurface foundation model development.

Figures (14)

• Figure 1: Architectural Overview. The fused encoder combines a (a) Global Perceiver Encoder that projects input features onto supernode queries via cross-attention with a (b) Local GNO Encoder that aggregates neighborhood information through radius graph pooling. A (c) Gated Fusion Mechanism adaptively combines global ($\mathbf{v}_{\text{attn}}$) and local ($\mathbf{v}_{\text{gno}}$) representations via a learned gate $G \in [0,1]^d$, followed by DiT blocks and Perceiver pooling to produce fixed-size latent tokens. The overall pipeline (d) APT architecture maps input fields $a(x)$ through the fused encoder $\mathcal{E}$, latent dynamics approximator $\mathcal{A}$ with temporal modulation, and Perceiver decoder $\mathcal{D}$ to generate output fields $z(x,t)$ at arbitrary query locations.
• Figure 2: Benchmark datasets for evaluating APT across diverse subsurface applications. (a) to (d) each displays input parameter fields (e.g., permeability) alongside the temporal evolution of output fields (e.g., saturation, pressure, temperature). (a) 2D Geologic Carbon Storage (GCS) utilizing an irregular mesh to resolve complex fault geometries. (b) 2D Hydrocarbon extraction on a Cartesian mesh with varying well configurations. (c) 2D Wastewater injection system modeled on a radial mesh. (d) 3D Basin-scale GCS featuring a million-cell nested semi-adaptive mesh with local grid refinement (LGR), supporting both Gaussian and channelized permeability geomodels. (e) 3D Aquifer Thermal Energy Storage (ATES) employing dynamic mesh optimization to track moving thermal fronts.
• Figure 3: Visualization of prediction error and state evolution across different subsurface applications. Rows compare APT against baseline models (MGN-LSTM, U-FNO, and FNO) at multiple time steps ($t$). (a) Saturation error maps ($\delta S_g$) for the 2D irregular mesh geologic carbon storage dataset, showing APT's stability near complex fault geometries. (b) Water saturation error maps ($\delta S_w$) for the 2D Cartesian hydrocarbon dataset. (c) Spatiotemporal temperature evolution for the 3D dynamic mesh ATES dataset, where APT maintains higher fidelity to the thermal front compared to FNO. Errors and physical quantities are normalized within each sample for comparative visualization.
• Figure 4: Cross-dataset training setup. (a) Gaussian permeability fields defined on nested grids with 4 levels of LGRs. (b) Channelized permeability fields LGR defined with 3 levels of LGRs.
• Figure 5: Heterogeneous permeability realizations with two fixed impermeable faults and one injection well for three cases. The well coordinates for each case are shown at the top, with insets displaying an enlarged view of the well vicinity.