A Dual-Branch Coupled Fourier Neural Operator for High-Resolution Multi-Phase Flow Modeling in Porous Media

Abstract

This paper investigates a physics-informed surrogate modeling framework for multi-phase flow in porous media based on the Fourier Neural Operator. Traditional numerical simulators, though accurate, suffer from severe computational bottlenecks due to fine-grid discretizations and the iterative solution of highly nonlinear partial differential equations. By parameterizing the kernel integral directly in Fourier space, the operator provides a discretization-invariant mapping between function spaces, enabling efficient spectral convolutions. We introduce a Dual-Branch Adaptive Fourier Neural Operator with a shared Fourier encoder and two decoders: a saturation branch that uses an inverse Fourier transform followed by a multilayer perceptron and a pressure branch that uses a convolutional decoder. Temporal information is injected via Time2Vec embeddings and a causal temporal transformer, conditioning each forward pass on step index and time step to maintain consistent dynamics across horizons. Physics-informed losses couple data fidelity with residuals from mass conservation and Darcy pressure, enforcing the governing constraints in Fourier space; truncated spectral kernels promote generalization across meshes without retraining. On SPE10-style heterogeneities, the model shifts the infinity-norm error mass into the ${10}^{-2}$ to ${10}^{-1}$ band during early transients and sustains lower errors during pseudo-steady state. In zero-shot three-dimensional coarse-to-fine upscaling from $30\times 110\times 5$ to $60\times 220\times 5$, it attains $R^2=0.90$, RMSE = $4.4\times {10}^{-2}$, and MAE = $3.2\times {10}^{-2}$, with more than 90% of voxels below five percent absolute error across five unseen layers, while the end-to-end pipeline runs about three times faster than a full-order fine-grid solve and preserves water-flood fronts and channel connectivity. Benchmarking against established baselines indicates a scalable, high-fidelity alternative for high-resolution multi-phase flow simulation in porous media.

1. Introduction

Multi-phase flow in porous media underpins critical applications in subsurface engineering—from groundwater remediation to carbon sequestration. The governing equations, nonlinear partial differential equations (PDEs) that capture the intricate interactions among fluid phases and the heterogeneous porous structure, are notoriously difficult to solve numerically. Traditional solvers, whether implemented in commercial software or open-source tools, require extremely fine discretizations to capture the essential dynamics, which in turn leads to prohibitive computational costs and long turnaround times, especially when conducting multi-parameter studies or producing high-resolution predictions.

While deep learning has revolutionized fields such as imaging, speech recognition, and natural language processing [1], the application of machine learning techniques that explicitly incorporate physical laws is relatively new. Early contributions such as the physics-informed neural networks (PINNs) introduced [2] and later refined in works like those of [3,4] have demonstrated that embedding governing equations into the training process can guide the solution toward physically consistent behavior. Yet, these methods have often yielded modest gains when compared with the dramatic successes observed in other areas of deep learning.

A particularly transformative approach comes from the emerging field of operator learning, which seeks to approximate mappings between infinite-dimensional function spaces. In this context, the Fourier Neural Operator (FNO) introduced [5] and extended [6] represents a paradigm shift. By reformulating the kernel integral in Fourier space and truncating the high-frequency components, FNO achieves spectral convolution that is both computationally efficient and inherently suited to enforcing physical constraints. Crucially, its invariance with respect to discretization allows the learned operator to generalize seamlessly across arbitrary grids without requiring re-training—a property of immense value for high-resolution simulations.

In porous media flow, where multi-scale heterogeneities and strong nonlinearities prevail, integrating physics-informed loss terms into the FNO framework (as demonstrated in physics-informed neural operators [7,8]) has proven essential. This integration ensures that the surrogate model rigorously adheres to the underlying PDEs, thereby reducing the computational effort needed to achieve accurate predictions. For example, implementations on platforms such as NVIDIA Modulus have shown that FNO-based approaches can reduce simulation times by up to three orders of magnitude while preserving accuracy across different discretizations [9]. More recently, several studies have adapted FNO-style architectures specifically to multiphase flow in porous media and related subsurface applications, including enhanced multiphase FNO variants, residual-corrected neural operators for nonlinear PDE inverse problems, and operator-based surrogates for field-scale reservoir simulations [10,11,12,13].

However, existing applications of Fourier Neural Operators to porous-media flow still face important limitations. Most studies focus on single-phase or weakly coupled settings and treat only one state variable (typically pressure or saturation), thereby neglecting the strong two-way coupling that characterizes multiphase displacement. Purely data-driven FNOs also tend to oversmooth sharp saturation fronts and underestimate channelized flow in highly heterogeneous formations such as SPE10, particularly during early transients. In many implementations, time is handled as a simple scalar feature or as an additional spatial coordinate, which can degrade accuracy over long horizons and lead to drift between transient and pseudo-steady regimes. Moreover, existing FNO surrogates rarely enforce PDE residuals during training, so saturation and pressure predictions may violate mass balance and Darcy’s law, reducing robustness when extrapolating across injection scenarios.

This work presents an advanced framework for applying FNO to the physics-informed simulation of multi-phase flow in porous media. Building on a dual-branch formulation (DB-AFNO), the physics-based loss functions are integrated with different temporal encodings (e.g., Time2Vec, TFT/TST) to honor the governing PDE constraints in both saturation and pressure fields. Our comprehensive computational analysis demonstrates that DB-AFNO outperforms conventional, purely data-driven FNO approaches across diverse tracer transport configurations, reducing errors during early transient phases and maintaining accuracy in later pseudo-steady states. In particular, DB-AFNO with Time2Vec and TFT/TST concentrates the majority of saturation predictions in the lower error ranges, preserving critical flow structures that conventional FNO often oversmooth. Blind testing on top and bottom layers of the 10th SPE Comparative Solution Project [14] highlights the framework’s potential for scalable, high-resolution simulations without incurring the computational costs of traditional full-order solvers.

The uniqueness of this approach lies in explicitly coupling the two PDE variables using a multitask, time-aware architecture. While conventional FNOs typically address a single field, our model retains the core spectral convolution operations and extends them via key modifications. First, a shared spectral encoder predicts both pressure and saturation, extracting spatial features that naturally encode interdependencies. Second, two distinct time encoders—one for saturation and one for pressure—transform the scalar time t into high-dimensional representations that capture differing temporal sensitivities while preserving mutual spatial information. Finally, after decoding via separate branches, the outputs are fused through an additional convolutional layer that enforces consistency between the predictions through cross-information exchange.

2. Methods

This section presents a technical description of the problem using the dynamical state-space framework. The focus of this formulation is on the incompressible and immiscible displacement of two phases within a porous media. More recently, the same principles have been extended to carbon sequestration efforts, specifically in the displacement of brine by ${\mathrm{CO}}_2$ [15]. In this work, we adopt a Darcy-scale, incompressible and immiscible two-phase formulation on a fixed Cartesian grid. The two phases are denoted wetting and non-wetting, with saturations $S_w$ and $S_{nw}$ and corresponding pressures $P_w$ and $P_{nw}$. The physical domain $\mathrm{\Omega }$ is discretized into N control volumes on a structured grid, and the primary unknowns collected in the state vector are the cell-wise saturation and pressure fields. The fluid system is isothermal, with constant viscosities and densities, and rock properties (porosity and absolute permeability) are time-invariant and defined per grid cell. Flow is governed by Darcy’s law under a single-continuum representation, without explicit geomechanical deformation, fracture propagation, or matrix–fracture transfer terms, and without capillary-pressure hysteresis. Gravity is included through a constant acceleration vector, while molecular diffusion and dispersion are neglected so that transport is driven by advection and sources/sinks only. Under these assumptions, the saturation constraint reduces to $S_w+S_{nw}=1$, and a monotone capillary-pressure law $P_c\left(S_w\right)$ together with Corey-type relative permeability functions $k_{r\psi }\left(S_{\psi }\right)$ closes the system. The governing equations fully describe the displacement process considered here. Gas-phase transport, solution-gas liberation, and evaporation/condensation processes are neglected; consequently, no gas saturation or gas pressure variables are introduced into the state vector, and all volumetric source terms represent injection and production of the two liquid phases only. A concise summary of the formulation is provided in

Table 1. State-Space Variables for Multiphase Flow in Porous Media.

Component	Description/Equation
State vector $\mathit{x}$	$\mathit{x}=[S_{w1},S_{nw1},\dots ,S_{wN},S_{nwN}]$
Parameter vector $\mathit{p}$	$\mathit{p}=[\varphi ,K,{\mu }_w,{\mu }_{nw},{\rho }_w,{\rho }_{nw},P_{cap},\mathit{M}]$
Excitation vector $\mathit{u}\left(t\right)$	$\mathit{u}\left(t\right)=[q_{w1}\left(t\right),q_{nw1}\left(t\right),\dots ,q_{wN}\left(t\right),q_{nwN}\left(t\right)]$
Output vector $\mathit{y}$	$\mathit{y}=[S_{w1},S_{nw1},P_1,\dots ,S_{wN},S_{nwN},P_N]$

,

Table 2. Governing Equations for Multiphase Flow in Porous Media.

Component	Description/Equation
Mass conservation	${\displaystyle \frac{\partial \left(\varphi S_{\psi }\right)}{\partial t}+\nabla ·\left({\mathbf{v}}_{\psi }\right)=q_{\psi },\psi \in \{w,nw\}}$
Darcy’s law	${\displaystyle {\mathbf{v}}_{\psi }=-\frac{K{k}_{r\psi }}{{\mu }_{\psi }}\left(\nabla P_{\psi }-{\rho }_{\psi }\mathbf{g}\right),\psi \in \{w,nw\}}$
Saturation constraint	$S_w+S_{nw}=1$
Capillary pressure	$P_{c\psi }=P_{\psi }-P_w,\psi \in \left\{nw\right\}$
Relative permeabilities	${\displaystyle {\lambda }_w\left(s\right)=\frac{{\left(s^{*}\right)}^2}{{\mu }_w},{\lambda }_{nw}\left(s\right)=\frac{{(1-s^{*})}^2}{{\mu }_{nw}},s^{*}=\frac{s-s_{wc}}{1-s_{nwr}-s_{wc}}}$
Effective mobility	$k_{r\psi }=f\left(S_{\psi }\right),\psi \in \{w,nw\}$

and

Table 3. Discrete Saturation Dynamics and Coupling Terms.

Component	Description/Equation
Saturation Dynamics	${\displaystyle \frac{\partial S_{\psi }}{\partial t}=-\nabla ·\left(\frac{K{k}_{r\psi }}{{\mu }_{\psi }}\left(\nabla P_{\psi }-{\rho }_{\psi }\mathbf{g}\right)\right)+\frac{q_{\psi }}{\varphi }}$
	${\displaystyle \mathrm{or}\varphi \frac{\partial S_{\psi }}{\partial t}+\nabla ·\left(f_{\psi }\left(s\right)\mathbf{v}\right)=\frac{q_{\psi }}{{\rho }_{\psi }}}$
Transmissibility, T	${\displaystyle T_{ij}=2\left|A_{ij}\right|{\left(\frac{\Delta x_i}{{\lambda }_{i,ij}}+\frac{\Delta x_j}{{\lambda }_{j,ij}}\right)}^{-1}}$
Pressure Equation	Solve $TP=q$ (Two-Point Flux Approx.)
Coupling	${\displaystyle \frac{\partial S_{\psi }}{\partial t}=diag(f_{w1},\dots ,f_{wN})A+[f_{i1},\dots ,f_{iN}],}$
	where A captures velocity interactions and source scaling by local $PV$.

Notes: In the above table, $S_{wi}$ and $S_{nwi}$ denote the saturations/concentration of the wetting and non-wetting phases in the i-th grid block, respectively; $\varphi$ is the porosity; K represents permeability; ${\mu }_w$ and ${\mu }_{nw}$ are the viscosities of the wetting and non-wetting phases; ${\rho }_w$ and ${\rho }_{nw}$ are their densities; $P_{cap}$ is the capillary pressure; and $\mathit{M}$ includes any additional parameters. The excitation vector $\mathit{u}\left(t\right)$ consists of injection/production rates $q_{wi}\left(t\right)$ and $q_{nwi}\left(t\right)$. The output vector $\mathit{y}$ pairs saturation and pressure measurements. Mass conservation ($q_{\psi }$) and Darcy’s law (with velocity ${\mathit{v}}_{\psi }$) govern flow, while the saturation constraint and capillary pressure ensure physical consistency. Relative permeabilities ${\lambda }_w\left(s\right)$ and ${\lambda }_{nw}\left(s\right)$ (with normalized saturation $s^{*}$) and effective mobility $k_{r\psi }=f\left(S_{\psi }\right)$ model fluid transport properties. The transmissibility matrix $\mathit{T}$ (computed via TPFA) couples the pressure field $\mathit{P}$ to the source $\mathit{q}$, and the pressure-to-saturation coupling (via matrix $\mathit{A}$) reflects the influence of the velocity field and local pore volume ($PV$) on saturation evolution.

.

The above formulation was coded in Python using Pytorch and Cuda to allow for interoperable transfer between the simulator and the neural operator architecture. In the following sections, we will present the numerical methods used for the simulation part along with the neural operator architectures.

2.1. Model Description

In this study, we used the SPE10 Model 2 dataset [14], an open source benchmark known to evaluate up-scaling methods for aquifer flow and tracer transport. This is used as an analogue to an underground water aquifer for tracer transport. This data set is composed of a Brent sequence mapped on a Cartesian grid with dimensions of $60\times 220\times 85$, resulting in 1,122,000 cells. The model encompasses two distinct formations: the shallow-marine Tarbert formation occupying the upper 35 layers, characterized by relatively smooth permeability, and the fluvial Upper-Ness formation constituting the lower 50 layers, noted for its significant permeability variations spanning 8 to 12 orders of magnitude. These formations present markedly different permeability structures, as highlighted in Figure 1, where the model is inverted to emphasize the heterogeneity within the Upper-Ness formation. This complex permeability and porosity distribution poses a rigorous challenge, making it an ideal test case for operator learning.

2.2. Numerical Methods

We have implemented a framework that incorporates an adaptive trapezoidal integration scheme. The Generalized Conjugate Residual (GCR) method serves as a robust solver for linear systems within each Newton iteration, providing efficient handling of sparse matrix structures and enhancing the convergence characteristics of the overall method. The Newton’s algorithm, integrated with the GCR solver, is designed to leverage its fast convergence properties, significantly reducing the computational overhead associated with direct methods like LU decomposition, while maintaining the capacity for handling nonlinear systems effectively. This integration of the adaptive trapezoidal method with iterative linear solvers is crucial in our simulations, striking an optimal balance between computational efficiency and the accuracy of transient dynamic responses in complex multiphase flows.

There are two encountered main challenges in this work: (1) building a stable model to act as a base for our neural operator implementation and (2) ensuring that the model is computationally efficient.

To circumvent these challenges, we have implemented several advanced numerical techniques and optimization methods that enhance both the stability and efficiency of the model. These techniques include the adaptive dynamic stepping scheme, the use of GCR with a banded LU preconditioner, and the employment of automatic differentiation for efficient Jacobian computation.

2.2.1. Preconditioned Iterative Solvers

Iterative methods were deemed to be more suitable for our approach due to the ability to predetermine the accuracy needed for phase saturations. However, iterative methods, specifically GCR, may not be efficient for sparse matrices out-of-the-box. Therefore, using a preconditioner is important to achieve fast convergence while having control over the desired accuracy.

Banded LU/ILU preconditioners are particularly suitable for systems where the matrix A exhibits a banded structure [16]. This approach is advantageous for large sparse systems with a banded pattern, as it significantly reduces the fill-in compared to a full LU decomposition, thereby preserving the computational efficiency.

$$\kappa \left(U^{-1}{L}^{-1}A\right)<\kappa \left(A\right),\mathrm{for}\mathrm{banded}A.$$

(1)

Alternatively, when the Jacobian matrix A is symmetric and positive-definite, a Cholesky decomposition can be used for preconditioning [17]. The preconditioned system is solved via

$$L^{-T}{L}^{-1}Ax=L^{-T}{L}^{-1}b,$$

(2)

Hence, multi-grid preconditioners are also common for fluid flow problems, however, Python does not support GPU direct implement similar to PyAMG on the CPU therefore it was not included in our benchmark. The choice between banded LU and Cholesky preconditioning depends on the structure and properties of the Jacobian matrix, Cholesky preferred for symmetric positive-definite systems because of its computational efficiency. In our model, we found that ILU is more stable for our Neumann boundary condition pressure solver.

2.2.2. Adaptive Dynamic Stepping:

In the adaptive time-stepping scheme implemented within the simulation framework, the trapezoidal method is utilized to preserve numerical stability without the need for the implementation of CFL, particularly crucial in the context of stiff systems encountered in multiphase flow simulations. The integration process begins with a prediction step that takes advantage of the explicit Forward Euler method to approximate saturations at the subsequent timestep, denoted as $S_{\mathrm{predictor}}$. This preliminary prediction is calculated according to the following equation:

$$S_{\mathrm{predictor}}=S_{\mathrm{current}}+\Delta t·{\left({\displaystyle \frac{dS}{dt}}\right)}_{\mathrm{current}}$$

where $S_{\mathrm{current}}$ represents the current saturation state, $\Delta t$ signifies the current timestep, and ${\left({\displaystyle \frac{dS}{dt}}\right)}_{\mathrm{current}}$ is the rate of change in saturation evaluated in the current state. This predictor step serves as an initial approximation for the subsequent implicit corrector step, which refines this estimate to enhance stability and accuracy. The corrector phase is carried out using an implicit trapezoidal integration, where the corrected saturation $S_{\mathrm{next}}$ is determined by resolving the nonlinear equation:

$$S_{\mathrm{next}}-S_{\mathrm{current}}-{\displaystyle \frac{\Delta t}{2}}\left({\left({\displaystyle \frac{dS}{dt}}\right)}_{\mathrm{current}}+{\left({\displaystyle \frac{dS}{dt}}\right)}_{\mathrm{next}}\right)=0$$

Here, ${\left({\displaystyle \frac{dS}{dt}}\right)}_{\mathrm{next}}$ denotes the rate of saturation change evaluated in the next predicted state. This equation is solved iteratively using the Newton–Raphson method.

To dynamically adjust the timestep $\Delta t$, an error estimation mechanism is incorporated that compares the norm of the discrepancy between the predictor and corrector solutions against a predefined tolerance. This adaptive time-step control is vital for managing computational efficiency and accuracy throughout varying dynamics of the modeled system. If the error estimate exceeds the tolerance threshold, $\Delta t$ is reduced to enhance resolution; conversely, if the error is significantly below the tolerance, $\Delta t$ is increased to expedite computations without sacrificing significant accuracy.In the context of dynamically regulating the time step $\Delta t$ within numerical simulations, a sophisticated approach is required to ensure stability and precision under varying conditions. The reduction factor for adjusting $\Delta t$ can be mathematically defined using an exponential function, providing a smooth and continuous regulation mechanism. The reduction factor R for the time step $\Delta t$ is given by the following equation:

$$R(\Delta t)=\alpha +(\beta -\alpha )exp(-\gamma \Delta t)$$

where $\alpha$ is the minimum reduction factor, $\beta$ is the maximum reduction factor and $\gamma$ is a scaling parameter controlling the steepness of the exponential decay.

In this formulation, as $\Delta t$ increases, the reduction factor R approaches the lower bound $\alpha$, whereas as $\Delta t$ decreases, R approaches the upper bound $\beta$. To dynamically adjust the time step $\Delta t$, the new time step $\Delta t_{\mathrm{new}}$ is defined as

$$\Delta t_{\mathrm{new}}=max\left(\Delta t\times R(\Delta t),\Delta t_{\min }\right),$$

where $\Delta t_{\min }$ is the minimum allowable step size, preventing excessively small steps that would be computationally inefficient. The stability constraints for $\Delta t$ are derived from the spectral properties of the system’s Jacobian $\mathit{J}$, using the inverse of its maximum eigenvalue, ${\displaystyle \frac{1}{{\lambda }_{\max }\left(\mathit{J}\right)}}$, and a more conservative bound, ${\displaystyle \frac{0.1}{{\lambda }_{\max }\left(\mathit{J}\right)}}$. These limits ensure that timestep updates remain within a range that preserves the numerical stability of the integration method for multiphase flow. All pseudocode listings for the above algorithms, Algorithms A1–A4, are provided in Appendix A.

2.3. Neural Operator Architecture

We use a DB–AFNO neural operator that fuses spectral, spatial, and physics-based information to predict mean saturation $\widehat{S}\in {\mathbb{R}}^{N_x\times N_y}$ and variance ${\sigma }^2\in {\mathbb{R}}^{N_x\times N_y}$. Static geological properties—porosity $\varphi \left(\mathbf{x}\right)$, permeability $k\left(\mathbf{x}\right)$—initial states $P_{\mathrm{initial}}$, $S_{\mathrm{initial}}$, and phase viscosities ${\mu }_w$, ${\mu }_{nw}$ are combined with temporal encoding via Time2Vec and a causal transformer that conditions each step on the time index and $\Delta t$. Fourier layers capture the coupled space–time response, and training minimizes a composite loss that blends data misfit with PDE residuals enforcing mass conservation and Darcy flow. In this way, the network advances the state in time and returns high-resolution, physically consistent pressure P and saturation S fields at arbitrary query times, as summarized in Figure 2.

2.3.1. Data Preparation

The data preparation pipeline for the physics-informed Fourier Neural Operator framework integrates simulation outputs—including grid parameters, K fields, well configurations, initial saturation distributions, and fluid properties—stored in serialized files. The raw data is reformatted for a consistent spatial representation; for example, the K tensor is converted to an effective scalar field via magnitude computation. Temporal dynamics are captured using a multi-frequency positional encoding with sine and cosine functions over logarithmically spaced scales, which is then spatially broadcast and concatenated with static input channels to form a comprehensive feature tensor. Global statistical measures (mean and standard deviation) are computed over the entire dataset and applied to both input and target fields to standardize the data, reducing numerical variability and aiding convergence during training. Early time steps are replicated to better represent rapid transients. Over 1000 simulation runs at half fidelity were performed, with the data partitioned 80% for training and 20% for validation, plus additional cases reserved for blind testing. These runs span a wide range of heterogeneous SPE10 Model 2 realizations, as well as idealized cases with spatially uniform permeability and porosity that are included explicitly in the training set. Across all runs, we vary initial water saturations, viscosity ratios, well locations, and injection schedules, resulting in substantial variability in flood-front shapes, breakthrough times, and recovery factors. A statistical overview of the key petrophysical variables is provided in Figure 3, and descriptive statistics are listed in

Table 4. Statistical summary of petrophysical properties and fluid parameters. For each parameter, N is the number of samples, $\mu$ is the mean, $\sigma$ is the standard deviation, and Min/Max denote the range of values.

Parameter	Unit	N	$\mathit{\mu }$	$\mathit{\sigma }$	Min/Max
$K_x$	mD	$2.89\times {10}^6$	$3.86\times {10}^2$	$1.16\times {10}^3$	$5.08\times {10}^{-3}/3.98\times {10}^4$
$\varphi$	-	$2.89\times {10}^6$	$1.92\times {10}^{-1}$	$8.60\times {10}^{-2}$	$1.00\times {10}^{-1}/5.00\times {10}^{-1}$
$S_{\mathrm{w},\mathrm{init}}$	-	$2.89\times {10}^6$	$1.97\times {10}^{-1}$	$5.81\times {10}^{-2}$	$1.00\times {10}^{-1}/3.00\times {10}^{-1}$
${\mu }_{\mathrm{w}}$	cP	$8.75\times {10}^2$	$3.48\times {10}^{-1}$	$8.61\times {10}^{-2}$	$2.00\times {10}^{-1}/5.00\times {10}^{-1}$
${\mu }_{\mathrm{nw}}$	cP	$8.75\times {10}^2$	$2.97$	$5.62\times {10}^{-1}$	$2.00/4.00$

, demonstrating that the training set covers a broad portion of the oil–water displacement behavior within this benchmark setting. This process is implemented using a modular Python framework built on PyTorch (version 2.9.1, CUDA 12.6), supporting efficient file handling, dynamic sample generation, and mini-batch processing for large-scale data, with the simulations executed on hardware platforms equipped with RTX Titan GPUs.

2.3.2. Pure Data-Driven Approach

The model processes inputs $\mathit{X}$ via:

$${\mathit{Z}}_0=\mathrm{InputProj}\left(\mathit{X}\oplus {\mathcal{F}}^{-1}(\mathit{W}\odot \mathcal{F}\left(\mathit{G}\right))\right),$$

where $\mathcal{F}$ and ${\mathcal{F}}^{-1}$ denote Fourier and inverse Fourier transforms, $\mathit{W}$ are learnable spectral weights, and $\mathit{G}$ represents learned Fourier features. Subsequent layers combine spectral, spatial, and physics-guided paths:

$${\mathit{Z}}_{l+1}=\mathrm{GELU}\left(\sum _{k=1}^2{w}_k^{\left(l\right)}{\mathcal{P}}_k\left({\mathit{Z}}_l\right)\right)+{\mathit{Z}}_l,l\in \{1,...,L\},$$

where ${\mathcal{P}}_1$, and ${\mathcal{P}}_2$ correspond to spectral convolution and spatial depthwise convolution, respectively.

2.3.3. Dual Output Approach

In this physics-informed multi-task framework, the model jointly predicts transient pressure (P) and saturation (S) fields governed by two-phase flow dynamics.

$$\left(K,\mathrm{wells}\left(q\right),S_{\mathrm{i}},P_{\mathrm{i}},{\mu }_w,{\mu }_{nw},\varphi ,t\right)⟼\left(\widehat{S}(x,t),\widehat{P}(x,t)\right).$$

The composite loss function is designed to enforce both data fidelity and the underlying physical constraints through a dual-task formulation. It is defined as

$$\mathcal{L}=\underset{\mathrm{Data}\mathrm{terms}}{\underbrace{\parallel \widehat{S}-S_{\mathrm{true}}{\parallel }^2+{\parallel \widehat{P}-P_{\mathrm{true}}\parallel }^2}}+\alpha \underset{\mathrm{PDE}\mathrm{residuals}}{\underbrace{\left({\mathcal{R}}_S^2+{\mathcal{R}}_P^2\right)}},$$

where $\alpha$ is a weighting hyperparameter. The PDE residuals are expressed as

$${\mathcal{R}}_S={\displaystyle \frac{\partial \widehat{S}}{\partial t}}+\nabla ·\left(f\left(\widehat{S}\right)\widehat{\mathit{v}}\right),$$

$${\mathcal{R}}_P=\nabla ·\left({\lambda }_t\left(\widehat{S}\right)K\nabla \widehat{P}\right)-q.$$

$$\widehat{\mathit{v}}=-{\displaystyle \frac{{\lambda }_t\left(\widehat{S}\right)K}{\varphi }}\nabla \widehat{P},$$

In this formulation, pressure gradients drive the propagation of the saturation front via the nonlinear interaction expressed by the product $f\left(S\right){\lambda }_t\left(S\right)$, while porosity $\varphi$ modulates the effective velocity in the aquifer. This dual-output approach, therefore, not only minimizes discrepancies between the predicted and true fields but also rigorously enforces the physical laws governing two-phase flow in porous media.

For this coupling, handling the temporal part becomes a necessity. The temporal part of the PDE was tackled by incorporating a temporal representation into the input features with the goal of capturing both short-term variations and long-term trends. In our work, we consider three complementary approaches:

Temporal Sinusoidal Encoding

Temporal sinusoidal encoding relies on fixed periodic functions to map a scalar time t into a high-dimensional vector. At time step t, the encoding is defined as

$${\mathit{P}}_t=\left[sin\left({\displaystyle \frac{t}{{\lambda }_1}}\right),cos\left({\displaystyle \frac{t}{{\lambda }_1}}\right),\dots ,sin\left({\displaystyle \frac{t}{{\lambda }_{d/2}}}\right),cos\left({\displaystyle \frac{t}{{\lambda }_{d/2}}}\right)\right],$$

where ${\lambda }_i={10}^{4i/d}$ and d is the embedding dimension. This encoding is then concatenated with the static input features ${\mathit{x}}_t$ (e.g., permeability, well locations) and the previous saturation $S_{t-1}$ to form

$${\mathit{z}}_t=\left[S_{t-1},{\mathit{x}}_t,{\mathit{P}}_t\right].$$

This approach, introduced in the Transformer architecture [18], remains a popular and effective means to provide temporal context.

Time2Vec

Time2Vec is a learnable temporal encoding method that transforms the scalar time t into a vector composed of one linear component and several periodic components. The encoding is given by

$$\mathit{T}\left(t\right)=\left[a_{0}t+b_0,sin\left(a_{1}t+b_1\right),\dots ,sin\left(a_{d-1}t+b_{d-1}\right)\right],$$

where $a_i$ and $b_i$ are learnable parameters. By learning these parameters from data, Time2Vec is capable of capturing complex, nonlinear temporal patterns. This method has demonstrated advantages over fixed encodings in various time-series applications [19].

Transformer-Based Temporal Encoding

Transformer-based temporal encoding leverages self-attention mechanisms to generate context-aware representations of time. Dedicated time-series transformer architectures—such as the Temporal Fusion Transformer (TFT) [20] and the Time-series-specific Transformer family [21,22,23]—often incorporate gating mechanisms and variable selection networks to enhance stability and interpretability. In this approach, the scalar-time input is first projected into a high-dimensional space and then processed as a sequence via a Transformer encoder. The resulting outputs are aggregated (e.g., by mean pooling) to yield a dynamic embedding that effectively captures both local and global temporal dependencies.

The model predicts both saturation ($\widehat{S}$) and pressure ($\widehat{P}$) fields via separate decoding branches. Recognizing that the pressure signal is typically lower in magnitude, we apply a learnable amplification factor ${\gamma }_P$ (initialized at 2.0) to the pressure output:

$${\widehat{P}}_{\mathrm{amp}}={\gamma }_P\widehat{P}.$$

These outputs are then fused using a cross-gated mechanism that implements bidirectional modulation: the pressure branch provides context for refining saturation, and the saturation branch provides context for refining pressure. Concretely, we first construct gating maps by passing each branch through a shallow convolutional block followed by a sigmoid nonlinearity, so that the gates lie in $[0,1]$ and can act as spatially and temporally varying masks. The fused outputs are computed as

$${\widehat{S}}_{\mathrm{fused}}=\widehat{S}+\alpha \sigma \left({\mathcal{C}}_1\left({\widehat{P}}_{\mathrm{amp}}\right)\right)\odot {\widehat{P}}_{\mathrm{amp}},$$

$${\widehat{P}}_{\mathrm{fused}}={\widehat{P}}_{\mathrm{amp}}+\beta \sigma \left({\mathcal{C}}_2\left(\widehat{S}\right)\right)\odot \widehat{S},$$

where $\alpha$ and $\beta$ are learnable scalar (or channel-wise) parameters that control the strength of cross-branch contributions, $\sigma (·)$ denotes the sigmoid function, and ⊙ is element-wise multiplication. In this formulation, $\sigma \left({\mathcal{C}}_1\left({\widehat{P}}_{\mathrm{amp}}\right)\right)$ acts as a data-dependent gate that selectively injects pressure information into the saturation field, while $\sigma \left({\mathcal{C}}_2\left(\widehat{S}\right)\right)$ plays the analogous role in the pressure branch. This cross-gated fusion allows the network to dynamically downweight noisy or less informative regions and to emphasize locations where pressure and saturation must remain tightly coupled (e.g., near displacement fronts and high-permeability channels). As a result, the model can better balance the relative contributions of pressure and saturation, leading to final predictions that are more consistent with the underlying two-phase flow physics and less prone to spurious decoupling between the two fields. final predictions adhere to the underlying physics of two-phase flow. Figure 4 shows a dual-branch, time-aware FNO architecture that splits shared spectral features into separate saturation and pressure paths.

2.4. Training and Optimization

All models were implemented in PyTorch and trained on a single NVIDIA GPU. Inputs were standardized per channel using statistics computed once on the training set; saturation and pressure targets were standardized separately. Losses were computed in standardized space, and predictions were de-standardized for reporting. The network uses per-channel normalization in the spectral trunk and BatchNorm in the convolutional decoder; dropout is applied only in the pressure decoder ($p=0.2$). A learned 64-dimensional temporal/positional embedding of the scalar time/step is concatenated to the static inputs and broadcast over the spatial grid. No geometric augmentation was used; to emphasize early transients we oversampled early steps ($t\le 10$) by a factor of 2.

Optimization used AdamW (initial learning rate $5\times {10}^{-3}$, weight decay ${10}^{-2}$). A cosine-annealing schedule with warm restarts (CosineAnnealingWarmRestarts; $T_0=100$, $T_{\mathrm{mult}}=2$, ${\eta }_{min}={10}^{-8}$) was stepped per optimizer update and not reset across curriculum phases. Mini-batch size was 64, with global-norm gradient clipping at 2. The data term was Smooth-$L_1$ (Huber) with $\beta =0.01$, applied to saturation and pressure with equal weights. A physics-residual term (conservation and Darcy consistency) was weighted by ${\lambda }_{\mathrm{phys}}\left(e\right)=exp(-e/25)$, where e is the global epoch index; a total-variation penalty on the pressure head was available but set to zero in the main runs.

Training followed a staged curriculum over the simulation horizon. In the first phase, the network was trained only on the earliest portion of each trajectory (up to time index 25), and in subsequent phases the maximum time index was progressively increased to 100, 500, and finally the full trajectory (∞), using 50, 50, 50, and 500 epochs per phase, respectively. This schedule forces the model to first capture early-time transients and sharp displacement fronts before fitting the later-time, quasi-steady-state behavior over the entire horizon. A concise summary of the training and configuration settings is provided in

Table 5. Training and configuration settings.

Block	Item	Setting
Data	Normalization	z-score per channel; separate $S,P$
Model	Temporal/Coord emb.	64-D time/step; grid $(x,y,(z\left)\right)$
FNO	Width/Layers	$W=64$, $L=8$
FNO	Activation/Residual	GELU; skips
FNO	Paths	Spectral conv + depthwise conv
FNO	Fourier modes (curric.)	$(2,2,2)\to (6,6,3)\to (18,18,8)$
Decoders	S/P	S: iFFT→MLP; P: Conv, dropout $p=0.2$
Fusion	Gating/Amp.	Cross-gated conv + $\sigma$; ${\gamma }_P=2.0$
Optim.	Batch/Optimizer	64; AdamW ($\mathrm{lr}=5\times {10}^{-3}$, wd ${10}^{-2}$)
Optim.	LR schedule	Cosine restarts ($T_0=100$, $T_{\mathrm{mult}}=2$, ${\eta }_{min}={10}^{-8}$)
Optim.	Grad clip	Global norm $=2$
Losses	Data/Physics	Huber ($\beta =0.01$) on $S,P$; ${\lambda }_{\mathrm{phys}}\left(e\right)=exp(-e/25)$
Protocol	Sampling/Split	Early $t\le 10\times 2$; 80/20 train/val
Protocol	Curriculum/CKPT	$t_{max}\in \{25,100,500,\infty \}$; epochs $\{50,50,50,500\}$; best val

.

3. Results

3.1. Numerical Validation

For numerical validation, we examine both single-phase and multi-phase flow scenarios. In the single-phase case, the formation is modeled as a homogeneous $64\times 64\times 1$ grid with unit porosity and viscosities, and irreducible saturations $\left(s_w,s_{nw}\right)$ set to zero, allowing full pore space utilization. Injection occurs at one corner and production at the opposite. Unlike Aarnes et al. [24], who normalize the injection rate to one pore volume of water, our non-normalized simulation reproduces the expected physics, as shown in Figure 5 (A vs. B). For the multi-phase case, we use SPE10 Model 2 with an incompressible two-phase system $\left({\mu }_w=0.3\mathrm{cP},{\mu }_{nw}=3.0\mathrm{cP},s_{wc}=s_{nwr}=0.2\right)$. Figure 5 presents breakthrough saturation distributions for the top (A1 vs. B1) and bottom layers (A2 vs. B2). Our model matches the reference, with minor discrepancies attributable to numerical smoothing differences.

3.2. 2D FNO Results

We compare (i) a purely data-driven FNO baseline and (ii) a DB-AFNO with different temporal encodings (Positional Encoding, Time2Vec, TFT/TST). Numerical simulations serve as ground truth. We evaluate absolute error fields and spatio-temporal accuracy over time, including blind tests on the top and bottom layers of SPE10 (Figure 6).

As shown in Figure 6, DB-AFNO with Time2Vec and TFT/TST yields lower early-time errors ($t^{*}<0.25$), with more samples in the ${10}^{-2}$–${10}^{-1}$ range. For later times ($t^{*}\ge 0.25$), Time2Vec and TFT/TST remain close; Positional Encoding improves over the purely data-driven FNO but trails the former two. In Figure 7, each column corresponds to a different injection time, and each row compares the ground-truth simulator, the conventional FNO baseline, and DB-AFNO. The baseline FNO visibly diffuses the displacement front, particularly along high-permeability channels, and in several time slices the advancing water front becomes disconnected or breaks into isolated patches. In contrast, DB-AFNO preserves the continuity of the main displacement front and better reproduces the fingering patterns and channelized advance observed in the simulator. This is most evident near breakthrough, where DB-AFNO aligns both the front position and the saturation contrast across the front, whereas the baseline front arrives too early and exhibits a smeared transition zone. These qualitative differences are consistent with the error distributions in Figure 6, where DB-AFNO has a higher proportion of predictions in the ${10}^{-2}$–${10}^{-1}$ error band at both early and late times.

3.3. Zero-Shot Super Resolution

DB-AFNO is trained once—as an upscaler—on a single coarse/fine pair from the upper five layers of SPE10: input lattice $30\times 110\times 5$, target lattice $60\times 220\times 5$. After convergence, weights are frozen. At deployment, the coarse simulator supplies $\left(S^{\mathrm{coarse}},P^{\mathrm{coarse}}\right)$, which the operator upsamples

$${\mathcal{G}}_{\theta }:\left(S^{\mathrm{coarse}},P^{\mathrm{coarse}},K,\varphi ,{\mu }_w,{\mu }_{nw},q,t\right)⟼\left(\widehat{S},\widehat{P}\right){|}_{{\mathrm{\Omega }}_{60\times 220\times 5}}.$$

The FNO backbone includes a Cartesian coordinate-grid embedding (normalized $(x,y,z)$ appended per voxel) for absolute spatial context on any mesh. Training uses an incremental multi-resolution curriculum: start with $2\times 2\times 2$ retained Fourier modes on a $2\times$ subsampled grid, then raise spectral/spatial resolutions $\left[(2,2,2)\to (6,6,3)\to (18,18,8)\right]$ once validation-loss gradients plateau. In this zero-shot setup the full model (including gating and temporal channels) converges in 100 epochs, with ≈1.5 GPU-min per case on a single RTX Titan.

With >90% of voxels below a $5\%$ absolute-error threshold. Figure 8 shows 3-D renderings of $\widehat{S}$ and $|S-\widehat{S}|$; fronts, channel connectivity, and high-saturation ridges are preserved at $8\times$ volumetric up-resolution. Figure 9 provides layer-wise scatter plots at the final time step; similar behavior holds at other times. Results remain sensitive to the coarse grid: using $15\times 55\times 5$ did not reproduce the same accuracy.

4. Discussion

Across the 2-D benchmarks, DB-AFNO consistently outperforms a purely data-driven FNO. Error distributions in ${\ell }^{\infty }$ norm show that Time2Vec- and TFT/TST-encoded variants concentrate early-time errors ($t^{*}<0.25$) in the ${10}^{-2}$–${10}^{-1}$ band, whereas Positional Encoding and the conventional FNO exhibit heavier tails (Figure 6). Near pseudo–steady state ($t^{*}\ge 0.25$), Time2Vec and TFT/TST remain closely matched and continue to dominate Positional Encoding. Visual comparisons corroborate these statistics: the baseline FNO produces over-smoothed fields with disconnected saturation artifacts, while DB-AFNO retains sharp, physically plausible fronts and spatial connectivity (Figure 7). These gains arise from (i) the explicit coupling of pressure and saturation through dual branches, (ii) temporal conditioning that disambiguates early transients, and (iii) physics-informed losses that penalize mass-balance and Darcy residuals.

For zero-shot super-resolution, a train-low, infer-high regimen enables coarse-to-fine mapping without retraining. A single coarse/fine pair $30\times 110\times 5\to 60\times 220\times 5$ is sufficient to learn mesh-agnostic upscaling, aided by a Cartesian coordinate-grid embedding for absolute spatial context and an incremental multi-resolution curriculum that increases retained Fourier modes and spatial resolution only after validation plateaus. This curriculum halves wall-clock training and reduces GPU memory by ∼40% relative to single-shot high-resolution training, while the deployed pipeline (coarse TPFA + one forward pass) achieves ∼3× speedup over full-order fine-grid simulation (

Table 6. Runtime benchmark for zero-shot upscaling vs. full-order TPFA on five unseen layers.

Step	Wall-Time	Comment
Full-order TPFA on $60\times 220\times 5$	≈15 min	Ground truth
Coarse TPFA on $30\times 110\times 5$	3 min	Generates inputs
DB-AFNO zero-shot upscale	$1.5$ min	Single forward pass
Total	$4.5$ min	∼3 × faster

Note: All runs on a single RTX Titan; identical simulator settings except grid resolution.

). On five unseen layers, we obtain $\overline{R^2}=0.90$, $\overline{\mathrm{RMSE}}=4.4\times {10}^{-2}$, and $\overline{\mathrm{MAE}}=3.2\times {10}^{-2}$, with >90% of voxels below $5\%$ absolute error. Three-dimensional renderings confirm preservation of water-flood fronts, channel continuity, and high-saturation ridges at $8\times$ volumetric up-resolution (Figure 8), and layer-wise scatter plots indicate tight agreement at the final time step with similar behavior at other times (Figure 9).

A key technical limitation of the zero-shot path is its dependence on the fidelity of the coarse solution: the mapping presumes that the coarse grid resolves the dominant transport features. When the input was degraded to $15\times 55\times 5$, the same accuracy was not recovered, indicating a lower bound on coarse resolution for reliable upscale inference. Minor discrepancies in the validation figures are consistent with numerical smoothing choices rather than architectural deficiencies.

Overall, the results indicate that (i) enforcing coupled physics with dual-branch decoding and temporal encodings improves early-time accuracy and suppresses over-smoothing; (ii) the operator generalizes to unseen layers within SPE10 when conditioned on physically meaningful inputs; and (iii) zero-shot super-resolution provides accurate fine-grid fields at a fraction of the computational cost, provided the coarse simulation captures the essential flow physics.

The DB-AFNO surrogates presented here are trained entirely on incompressible, immiscible two-phase displacements, and thus are most reliable for formations and operating conditions with similar heterogeneity statistics, viscosity ratios, and flow regimes; all results in this study pertain to such systems with no explicit gas phase or compositional effects. When the underlying physics changes significantly—e.g., strong capillary hysteresis, gas evolution or gas caps, or fracture-dominated flow—retraining or transfer learning on simulations from the new regime would be required. While the dual-branch architecture is, in principle, extensible to three-phase or compositional flow (e.g., via an additional output branch and corresponding physics residuals), this lies beyond the scope of the present work and is left for future investigation.

5. Conclusions

• The proposed DB-AFNO framework achieves a significant reduction in computational cost compared to traditional full-order solvers while preserving high-resolution prediction accuracy, as demonstrated on the SPE10 benchmark.
• The dual-branch architecture effectively couples pressure and saturation fields, ensuring that the predicted outputs remain physically consistent and adhere closely to the governing PDE constraints.
• The incorporation of advanced temporal encoding methods (e.g., Time2Vec, TFT/TST) enhances the model’s ability to capture dynamic transient behaviors, resulting in lower validation errors during both early transient and pseudo-steady states.
• By integrating physics-informed loss functions, the model not only minimizes data-driven prediction errors but also enforces adherence to the underlying flow physics, thereby improving overall reliability.
• The overall architecture demonstrates scalability to higher-resolution simulations, given that it is trained on similar lower-fidelity cases.
• Despite its high-fidelity performance within the training distribution, DB-AFNO-based approaches may fall out-of-distribution when the underlying physics deviates from those encountered during training.
• The scalability of DB-AFNO was found to be similar to that of FNO-based models. It is limited in scenarios involving sudden or drastic changes in governing physical processes, indicating the need for adaptive re-training or architectural modifications when the physics evolve beyond the original training regime.
• Zero-shot super-resolution cuts the end-to-end turnaround for a 60 × 220 × 5 grid from ∼15 min (full TPFA) to 4.5 min by upscaling a coarse $30\times 110\times 5$ run with a single forward pass, achieving $\overline{R^2}=0.90$ and $\overline{\mathrm{RMSE}}=4.4\times {10}^{-2}$ without additional fine-grid training.