A Fast AI Surrogate for Coastal Ocean
Circulation Models

Existing machine learning techniques for numerical simulations include physics-informed neural networks (PINNs) and neural operator learning. A PINN trains a neural network to solve a partial different equation (PDE) through the automatic differentiation [25]. PINNs have been tested in simulating coastal flood and wave [26, 27], but their training is challenging due to the complex physics-informed regularization terms [28] and the huge computational cost for solving the PDE [29]. Neural operator learning aims to train a neural network surrogate as the solver of a family of PDE instances [30]. The surrogate takes the initial or boundary conditions and predicts the solution function. Existing surrogate models include deep convolutional neural networks [31], graph neural operators [32, 33], Fourier neural operators [34, 35], DeepONet [30], NodeFormer [36, 37] and vision transformers [38, 21]. Some simple methods have been tested for individual coastal problems (e.g., storm surges) [39, 40]. There are also methods for irregular spatial points or unstructured meshes in continuous space through implicit neural representation [41, 42, 43] or geometric deformation [44], or fixed-graph transformation [45], but their neural networks are limited by only taking each sample’s spatial coordinates without explicitly capturing spatial structure. There also exists a rich literature that particularly focuses on machine learning for continuous fluid dynamics (CFD) and turbulence [46], e.g., accelerating direct numerical simulations [47], improving turbulence closure modeling [48], and enhancing reduced-order models [49, 50]. However, existing research mostly focuses on idealized settings, with few works addressing real-world complex coastal circulation in estuaries.

Another body of closely related research is AI surrogates for global weather forecasting. For instance, GraphCast [19] utilizes Graph Neural Networks (GNNs) for medium-range weather forecasts, presenting a more efficient approach compared to traditional weather simulation systems. It predicts a resolution of 0.25 degrees in longitude and latitude, equivalent to approximately 28km $\times$ 28km at the equator. Similarly, FourCastNet [20] uses the Fourier Neural Operator. The Pangu model [21] employs a unique Earth-specific transformer model architecture. Despite the significant advancements of AI surrogates in global weather forecasting, these existing works focus on coarse-scale global forecasting. They cannot be directly applied to regional simulations like ROMS, since ROMS takes not only the initial condition but also the regional boundary conditions of future temporal snapshots. Such boundary conditions are not required in global forecasting models.

Regional Ocean Modeling System (ROMS) is a semi-explicit ocean model that solves the hydrostatic primitive equations using a Finite Volume scheme. ROMS integrates the Primitive Equations in time, starting at an initial condition and producing, at each time step, updated values of the physical variables (free surface, temperature, salinity, and velocities). It uses a non-uniform, structured spatial discretization (3D Arakawa-C grid) composed of various sigma-coordinate layers, the first of which follows the terrain, the last follows the free surface, and the rest occupies intermediate levels. Time integration is done by decomposing the 3D fields into their barotropic (depth-averaged) and baroclinic (the residual) parts. Further details of the forcing in the governing equations of ROMS can be found in [51, 52, 53].

ROMS is computationally intensive due to the need to discretize numerous state variables across high-resolution spatial grids and store intermediate time levels and halo cells for parallel processing. ROMS divides the simulation area into horizontally rectangular zones, assigning each zone to a CPU core. As the number of subdivisions increases, the practicality and efficiency of the parallelization diminish. Furthermore, smaller subdivisions lead to larger MPI messages because of the need to share boundary cells among the cores.

Existing works explore efficient scaling of ROMS by deploying it across multiple nodes on HPC clusters and commercial clouds, leveraging parallelization techniques using MPI [8], containerization for reproducibility and portability [22, 23], and I/O parallelism to accelerate ROMS execution [24]. Table I summarizes existing efforts in optimizing ROMS simulation on HPC clusters. The table compares the simulation overhead, which represents the computational time required to run the simulation, across different solutions utilizing various computational and memory resources, such as CPUs, GPUs, and RAM. Additionally, the table presents the simulation workload details, including the simulation size (i.e., the mesh size) and the simulation time horizon, which refers to the time duration of coastal circulation simulations. As can be seen in the table, the existing works [24, 22, 8, 23] exhibit significantly higher computational overheads. For example, in the first row [8], running ROMS for 3-day forecasting on a $1520\times 1088\times 30$ with 256 CPU cores took 19,915 seconds. This result is consistent with our benchmarking results of running ROMS simulation on a $898\times 598\times 12$ mesh with 512 cores, which took $9,908$ seconds (half of the time cost of [8]). In contrast, our AI surrogate solution demonstrates remarkable efficiency in optimizing ROMS simulations. By leveraging a combination of 16 AMD EPYC 7742 Processors and one NVIDIA A100 GPU, AI surrogate achieves a simulation overhead of just 22 seconds for 12-day forecasts of coastal ocean circulation with a mesh size of 898$\times$598$\times$12, which is near 450$\times$ faster than traditional MPI ROMS with the same simulation size. This indicates a significant potential for performance improvement by adapting AI surrogates in ROMS simulations.

The overall workflow is shown in Figure 1. Historical coastal simulations from ROMS are pre-processed and used to train an AI surrogate. Once trained, the AI surrogate can make forecasts by taking the initial condition of physical variables at time $t_{0}$ and the boundary conditions of those variables from time $t_{1}$ to time $t_{T}$ as inputs. The surrogate predicts the interior values of those variables at time $t_{1}$ to $t_{T}$ as outputs. Some variables, such as free surface elevation, are 2D, while others, such as circulation velocities, are 3D. A physics-based verification module checks if the surrogate’s forecast satisfies the physical constraints, such as the conservation of mass. If the forecasts meet these requirements, the forecasting continues with the surrogate generating the next episode of snapshots. Otherwise, the pipeline temporarily switches back to the original ROMS model to run the forecast.

For long-term forecasting, we employ a dual-model approach, i.e., training two models with the same architecture but at different temporal resolutions. Specifically, one model makes forecasts over 12 days with 24 snapshots (one snapshot every 12 hours), while the other model makes forecasts over 12 hours with 24 snapshots (one snapshot every half an hour). To generate 12-day forecasts with half-hour intervals (576 snapshots), we combine the two models, i.e., first running the coarse model to generate a 12-day forecast with a snapshot every half day and then using each coarse snapshot as the initial condition to generate 24 half-hour snapshots via the second model.

The dataset for this study is derived from our decade-long (2008-2017) Regional Ocean Modeling System (ROMS) simulation in Charlotte Harbor. This comprehensive dataset includes temporal snapshots every half-hour, featuring over ten physical variables such as circulation velocities (horizontal and vertical), temperature, salinity, and free surface elevation, alongside several external force variables, including surface wind components, air pressure, relative humidity, solar radiation, and rainfall rate. The spatial domain is represented by a non-uniform 3D mesh of $898\times 598\times 12$ with a higher resolution near river channels and inlets to capture the complex land-water interactions. For simplicity, we focus on the tidal wave propagation process in the estuary, which involves a subset of four physical variables, including circulation velocities along two horizontal directions ($u,v$) and the vertical direction ($w$), as well as the free surface elevation ($\zeta$).

We use simulations in 2011 for training and validation (9:1) and simulations in 2012 for testing. To train the 12-hour forecasting model, we augment the 2011 dataset using a sliding window of 24 time steps (12 hours) with a stride of 6. This results in 2,588 training instances, 288 validation instances, and 720 test instances. Similarly, to train the 12-day model, we resample the 2011 data with 12-hour intervals, generating 2,509 training instances and 279 validation instances. The volumes of the two training datasets are 2.6 TBs and 2.5 TBs, respectively. The 2012 data was partitioned into non-overlapping windows for testing in both cases.

In the original ROMS simulation data, the current velocity variables are located on the sides of cells, while the other variables are located in cell centers. To accommodate neural network training, we use linear interpolation to resample all variables to cell centers. Additionally, zero-padding is used to adjust the mesh size to $900\times 600\times 12$. The original ROMS simulations are based on FP64 for numerical stability, but the data is converted to FP16 for model training to enable faster computation and reduced memory usage. Numerical models like ROMS require higher precision due to accumulated rounding errors in numerous iterations to solve complex equations, while AI models can capture patterns from data without iterative operations. Further studies are needed to see the effect of such data compression on the accuracy performance of AI surrogates. All variables are normalized using z-score normalization based on the mean and standard deviation from the 2011 data.

Our AI surrogate model is built upon the 4D Swin Transformer model [54, 55, 56]. It has an encoder-decoder architecture, as shown in Figure 2. The model starts by taking the initial conditions at time $t=0$ and boundary conditions from $t=1$ to $T$ across all simulation variables (both 3D variables $u,v,w$ and 2D variable $\zeta$ on the mesh surface). To efficiently process this high-dimensional data with millions of mesh cells, we first employ patch partitioning and embedding for 3D and 2D variables. Patch partitioning divides the 3D/2D spatial domain into smaller blocks, effectively reducing the spatial resolution. Patch embedding transforms these blocks into the same latent representation space, facilitating the concatenation of a 2D variable (free surface elevation $\zeta$) and 3D variables (current velocity $u,v,w$) along the depth dimension. After considering an additional temporal dimension, the 4D data was then fed into a 4D Swin Transformer encoder. It leverages the 4D spatiotemporal attention mechanism to learn complex dependency patterns across both space and time with advanced window self-attention and shifted-window self-attention layers for computational efficiency, simultaneously downsampling the data with patch merging operations to capture the multi-scale information. The decoder employs multiple layers of transposed convolutions, with batch normalization [57] and Gaussian Error Linear Units (GELU) [58] activations to gradually upsample the encoded data, restoring it to its original resolution. A skip connection operation similar to the U-Net architecture [59, 60] is introduced to fuse fine-grained information directly from the corresponding encoder layers. Finally, the model splits the 3D and 2D variables after upsampling and recovers them to the original spatial and temporal shape.

Patch embedding. To efficiently handle high-dimensional data with millions of mesh cells, we first employ patch partitioning and embedding, dividing the 3D/2D spatial domain into smaller segments and embedding these segments into the same embedding space. Choosing the appropriate patch size is crucial: smaller patches can capture fine-grained information more effectively but increase the number of patches, thereby raising the computational cost of the following attention mechanisms. Thus, a balance must be made between model accuracy and training complexity. In this study, 3D variables ($u,v,w$) form a $3\times H\times W\times D\times T$ tensor, while the 2D variable ($\zeta$) contains a $1\times H\times W\times T$ volume, where $H,W,D$, and $T$ are height, width, depth, and time respectively. These variables are partitioned into patches and mapped into a $C$-dimensional latent space. For the 3D variables, the patch size is $P_{H}\times P_{W}\times P_{D}$, resulting an embedded shape of $C\times\frac{H}{P_{H}}\times\frac{W}{P_{W}}\times\frac{D}{P_{D}}\times T$. For the 2D surface variable, the patch size is $P_{H}\times P_{W}$, so the embedded data have a shape of $C\times\frac{H}{P_{H}}\times\frac{W}{P_{W}}\times T$, where $C$ is the initial embedding dimension. These embedded volumes are then concatenated along the depth dimension and passed to the encoder.

Swin transformer block.

To effectively process high-dimensional spatiotemporal data in coastal simulations, we adopt the 4D Swin Transformer architecture. This approach is well-suited for the task because it can capture complex spatial and temporal dependencies while mitigating the computational challenges associated with traditional attention mechanisms.

Multi-head self-attention (MSA) [61] has revolutionized modeling capabilities by allowing the network to focus on different parts of the input simultaneously. It achieves this by computing attention scores using queries ($Q$), keys ($K$), and values ($V$), derived from the input data

where $d$ is the embedding dimension. MSA extends this concept by splitting the attention mechanism across multiple heads, each computing attention independently:

where head_i = Attention$(QW_{i}^{Q},KW_{i}^{K},VW_{i}^{V})$, $W^{O}$, $W_{i}^{Q}$, $W_{i}^{K}$, $W_{i}^{V}$ are learnable parameter matrices, and $h$ is the number of heads.

While MSA is powerful, it has an $O(n^{2})$ computational complexity to the number of patches, making it unsuitable for high-dimensional simulation data due to prohibitive computational cost. The Swin Transformer addresses this challenge by partitioning the input into smaller, non-overlapping windows and applying attention within these local windows.

The core of the Swin Transformer model is the window multi-head self-attention (W-MSA) layer. Given $H^{\prime}\times W^{\prime}\times D^{\prime}\times T$ input patches, the patches are partitioned by a window of size $M_{H}\times M_{W}\times M_{D}\times M_{T}$, resulting in $\frac{H^{\prime}}{M_{H}}\times\frac{W^{\prime}}{M_{W}}\times\frac{D^{\prime}}{M_{D}}\times\frac{T}{M_{T}}$ non-overlapping windows. While W-MSA improves efficiency, it does not allow for information exchange between windows. To address this, we use shifted window multi-head self-attention (SW-MSA) layer which shifts the windows by a set number of patches in the subsequent layer to enable cross-window connections for a more comprehensive understanding of the data. Although the number of windows increases in SW-MSA, the actual computation cost remains similar by leveraging the cyclic-shifting batch [54]. An example of detailed operations of 4D W-MSA and SW-MSA is shown in Figure 3(a). Combining the W-MSA layer and the SW-MSA layer, two successive 4D Swin Transformer blocks, as shown in Figure 3(b), are computed as the following:

in which (S)W-MSA, LN, and MLP denote the (Shifted) Window Multi-head Self-Attention, Layer Normalization [62], and Multi-Layer Perceptron module, respectively. Moreover, $\hat{z}^{l}$ and $z^{l}$ denote the output features of the (S)W-MSA module and the following MLP module for block $l$, respectively.

Patch merging. Patch merging is used to downsample the data, enabling a hierarchical feature extraction structure while reducing the computational complexity of subsequent layers.

Following prior works [55, 63, 56], the patch merging step performs on the three spatial dimensions (height, width, depth) but not for the temporal dimension. During the patch merging operation, the spatial dimensions are reduced by half, and the channel size is doubled as compensation. As shown in Figure 4, consider a tensor with dimensions $C^{\prime}\times H^{\prime}\times W^{\prime}\times D^{\prime}\times T$. During patch merging, this tensor is reshaped into dimensions $8C^{\prime}\times\frac{H^{\prime}}{2}\times\frac{W^{\prime}}{2}\times\frac{D^{\prime}}{2}\times T$, where $2\times 2\times 2$ spatially neighboring patches are concatenated along the channel dimension. Then, each channel in the resulting tensor is projected onto a $2C^{\prime}$ dimensional space by applying a single fully connected layer, resulting $2C^{\prime}\times\frac{H^{\prime}}{2}\times\frac{W^{\prime}}{2}\times\frac{D^{\prime}}{2}\times T$ in total.

Positional encoding. To accurately represent the spatiotemporal distribution of data on non-uniform grids, we use absolute positional encoding, which explicitly captures the positions of patches in both space and time. In line with [64], we separately add positional embeddings for the spatial and temporal dimensions. Specifically, given an input tensor with dimensions of $C^{\prime}\times H^{\prime}\times W^{\prime}\times D^{\prime}\times T$, we define spatial and temporal positional embedding tensors with dimensions of $C^{\prime}\times H^{\prime}\times W^{\prime}\times D^{\prime}\times 1$ and $C^{\prime}\times 1\times 1\times 1\times T$, respectively. These tensors are then added to the input tensor using broadcasting.

Decoder. After the encoding process, the decoder reconstructs the spatial and temporal features from their embedded space, enabling the generation of high-resolution outputs. It starts by applying a series of 3D transposed convolutional layers with batch normalization and Gaussian Error Linear Unit (GELU) activations to upsample the embedded tensor back to its original dimensions. At this stage, a skip connection similar to the U-Net architecture [59, 60] is introduced, allowing the decoder to fuse spatial information directly from the corresponding encoder layers, which helps preserve fine-grained spatial details and improves the overall reconstruction quality. Following this step, the decoder splits the 3D and 2D patches and applies different patch recovery methods, which contain a 3D/2D transposed convolution, batch normalization, and GELU activation, with $1\times 1$ 3D/2D convolutions to reconstruct the spatiotemporal features and align the features with their original shape.

AI surrogate model training on high-resolution spatiotemporal simulations requires intensive I/O to load a large volume of training data into GPU memory. To fully utilize the parallelism of GPUs and achieve efficient training with large datasets, it is essential to identify and address performance bottlenecks. We analyze the throughput of each stage in the training pipeline, including training sample loading, training sample processing, and model parameter updating, to determine the limiting factor. Table II shows results.

Table II lists the memory requirement for each stage in the training pipeline when we process one sample with mesh size as 900$\times$600$\times$12 using mixed precision training. We identify two bottlenecks for AI surrogate training. (1) Low bandwidth of SSD slows down training sample loading from SSD to higher memory hierarchy. Loading one sample from SSD to CPU memory causes a 5.5-second latency on the critical execution path. (2) Training batch size is limited by the capacity of GPU since the activation tensor generated in each forward pass consumes most GPU memory. Given the NVIDIA A100 GPU memory capacity limitation of 80GB, we can only use per GPU training batch size as 1 for AI surrogate training, which significantly impacts the convergence efficiency.

We propose two optimization techniques to address the above performance bottlenecks by (1) reducing I/O strain by leveraging OS-level caching; and (2) reducing peak memory consumption of training with customized activation checkpoint. By addressing the I/O and memory usage bottlenecks head-on, we pave the way for more scalable, efficient, and effective AI surrogate training.

Leveraging OS cache for I/O strain reduction. An additional optimization was introduced via the PyTorch DataLoader’s prefetch technique. By increasing the number of subprocesses dedicated to data loading, we significantly improved the efficiency of moving data from SSD storage to CPU RAM. This adjustment ensures that data is pre-fetched and ready for processing in parallel with ongoing computations, effectively hiding I/O overhead during model training. Such an approach is crucial in maximizing computational throughput.

Further, to optimize the transfer of data from CPU to GPU, we employed the use of pinned memory and non-blocking I/O operations. Pinned memory, or page-locked memory, facilitates faster data transfer rates by preventing the paging of memory, while non-blocking I/O operations allow for concurrent execution of data transfers and computation. These modifications substantially accelerate the data pipeline, mitigating potential bottlenecks between CPU and GPU.

Reduce GPU peak memory consumption through activation checkpointing. A pivotal aspect of our optimization strategy involves the implementation of checkpointing techniques to effectively manage GPU memory usage. By selectively storing the activations of the Swin Transformer layers, we can reconstruct the intermediate states as needed, rather than maintaining them in memory throughout the training process. To minimize the reconstruction overhead while minimizing the memory usage for activations, we store the activations of the shifted-window multi-head self-attention (SW-MSA) layers but discard the activations of other layers. This is because the SW-MSA layer is the most computationally expensive, but the activations of the SW-MSA layer and the remaining layers have the same size. This technique not only reduces the per-sample GPU memory footprint but also enables us to increase the per-GPU batch size, effectively doubling the computational workload without necessitating additional GPU memory resources.

One crucial aspect of the AI surrogate’s inference is result verification. A verification module grounded in fundamental physical laws represents a major innovation in our AI surrogate system. By ensuring that the surrogate’s predictions adhere to mass conservation law, we establish a rigorous framework for evaluating and improving the quality of our forecasting results. This verification process, coupled with a dynamic feedback loop involving both our AI surrogate and ROMS, underscores our commitment to developing a predictive system that is both scientifically robust and computationally efficient.

Verification to ensure adherence to physical laws. At the core of our verification strategy is the mass conservation law of water. To satisfy this law, the divergence of circulation velocities into a water column (water inflow) should be equal to the time derivative of free surface elevation (increase of volume). For a horizontal domain $\Omega$ of contour $\Gamma$, the mass conservation law can be expressed as Equation 4,

where $h$ is the depth (it is a constant for the same study area), $\mathbf{u}$ is the horizontal velocity vector (a vector of $u,v$), $\mathbf{n}$ is the outward-pointing unit normal vector to the boundary, and the horizontal domain $\Omega$ is a horizontal grid cell in our case. This physical law serves as a critical metric for assessing the quality of predictions generated by our AI surrogate. By integrating this verification step, we can evaluate whether the surrogate’s outputs are physically consistent. Specifically, we compute the residual value by subtracting both sides of the equation and taking the absolute value:

We compute this value for all cells in one forecast. If the average value is less than a setting threshold, then the prediction is considered physics-consistent and can pass the verification.

Combine AI and ROMS for correction and improvement. The verification module checks the AI surrogate’s predictions against established conservation criteria. For instances in which the AI surrogate’s results fail to meet the verification criteria, indicating a deviation from expected physical behaviors, our system is designed to revert to the Regional Ocean Modeling System (ROMS) for direct simulation. This fallback mechanism ensures that our predictions remain anchored to robust, physics-based simulations, maintaining the reliability of our forecasts. A key factor that will impact the scaling performance of our pipeline is the passing rate of the AI surrogate in physics verification. A high passing rate will lead to a smaller chance of reverting back to the original ROMS simulation and thus achieving a better speed-up.

Our computational experiments are executed on an HPC cluster with 140 NVIDIA DGX-2 nodes. Each node is equipped with eight NVIDIA A100 GPUs (each A100 has 80GB of GPU memory), alongside 2 AMD EPYC 7742 64-Core Processors (128 CPU cores for each node) and 2010GB of CPU memory. The GPUs within each node are intricately linked through Nvlink technology, facilitating efficient intra-node communication. The cluster employs a sophisticated network of 10x HDR InfiniBand (IB) interfaces for inter-node data transfer, designed to ensure seamless, non-blocking communication across nodes. In the training phase of our AI surrogate model, we use up to 32 A100 GPUs within four DGX-2 nodes. For inference, our model’s requirements are significantly reduced, necessitating only a single A100 GPU complemented by 16 CPU cores and 250 GB of memory, highlighting the model’s efficiency.

In our experiments, the AI surrogate model is designed to generate predictions based on initial conditions and boundary conditions for the next 24 time steps. With two different time intervals (0.5-hour and 12-hour), the model can produce results spanning 12 hours and 12 days, both at a 0.5-hour resolution. Table III presents the average MAE and RMSE for prediction performance on the 2012 test data.

The AI surrogate model handles both 3D and 2D variables. For 3D variables, we use a patch size of $5\times 5\times 4$ for 3D variables and $5\times 5$ for 2D variables, with an initial patch embedding size of 24. The backbone consists of 4D Swin Transformer operations, applied iteratively three times, and employed attention heads of 3, 6, and 12, respectively. The initial window attention operation utilizes a window size of (4, 4, 2, 2) (for 4 dimensions) while subsequently adopting a window size of (2, 2, 2, 2).

In alignment with the optimization strategies outlined in Section III-D, we implement pin memory and non-blocking I/O to streamline data transfer between CPU and GPU, reducing bottlenecks and improving computational efficiency. During training, we use a batch size of 2 per GPU and employ 6 worker processes with a prefetch factor of 2 to ensure continuous data loading. Both fine- and coarse-resolution models are trained for 30 epochs, with training times of 58 and 56 minutes, respectively.

For inference on a single A100 GPU, a 12-hour forecast takes 0.888 seconds. For the 12-day forecast, which involves running the coarse-resolution model followed by 24 iterations of the fine-resolution model, the total inference time is 22.2 seconds.

To demonstrate the forecast capabilities of our AI surrogate system, we present the visualization of spatial distributions of the $u,v$ and $\zeta$ variables in Figure 5. The vertical velocity $w$ is always very small and is very close to 0 for most areas, so we do not show the visualization of $w$ here. We also show the temporal variation of ROMS’s and AI surrogate’s output on free surface elevation ($\zeta$) at three distinct locations from 12 p.m. January 2, 2012, to 12 p.m. January 14, 2012, as shown in Figure 6. These visualizations underscore the system’s proficiency in capturing the complex spatial and temporal dynamics inherent in the ROMS simulation.

We assess the performance of the AI surrogate-based workflow. First, we quantify the AI surrogate inference quality by comparing water mass residuals calculated from AI surrogate results with various thresholds. We set the threshold of passing the verification to 3e-4 m/s, 3.5e-4 m/s, 4e-4 m/s, 4.5e-4 m/s, 5e-4 m/s, and 5.5e-4 m/s and tested on the 2012 dataset. Figure 7 shows these results. As the water mass residual threshold increases, the acceptance rate of AI surrogate inference results also increases. This is because higher thresholds allow for results that deviate more from strict adherence to the water mass conservation law while still being considered acceptable. Water mass residuals of simulated results smaller than 5.0e-4 m/s are typically considered acceptable by oceanographers. Our analysis shows that over 99% of AI surrogate inference results meet this criterion, indicating that they are of high quality.

We further quantify the end-to-end performance of the AI surrogate-based workflow for coastal ocean circulation simulation. Figure 8 illustrates the execution time of the workflow at various verification thresholds, comprising AI surrogate computation time plus ROMS execution time for cases where AI surrogate inference fails verification. The verification time can be ignored. Specifically, the workflow switches back to ROMS for computation after detecting low-quality results from the AI surrogate. Notably, as discussed in Section III-A, we employ dual-model inference for long-term forecasting. Since the verification is very lightweight, we check inference quality after every AI surrogate inference, enabling early error detection during the calculation. Evidence of this is seen in the varying AI surrogate execution times in figure 8.

The AI surrogate-based workflow demonstrates significant efficiency gains, as shown in figure 8. It achieves up to a 450$\times$ speedup compared to traditional methods. Even with a strict verification threshold (i.e., water mass residual smaller than 3e-4 m/s), the workflow still outperforms traditional ROMS simulation by 1.8$\times$. Furthermore, using the threshold typically accepted by oceanographers, the AI surrogate-based workflow achieves up to a 275$\times$ speedup. These results highlight the substantial computational efficiency gained through the AI surrogate approach while maintaining simulation quality acceptable to the oceanographic community.

We conduct a sensitivity analysis in patching size for the AI surrogate efficiency. As a key hyperparameter within the AI model architecture, patch size critically impacts the system’s computational efficiency and predictive accuracy. We explore the effects of three patch sizes—5 (utilized in our AI surrogate system), 15, and 25 horizontally. The results, as depicted in Table IV, reveal an effect of patch size on the model’s complexity, primarily reflected through the parameter count. Larger patch sizes would decrease the number of patches in attention computations, thereby reducing parameters, but will increase the decoder’s parameter count because of the transpose convolutional upsampling operations. The smallest patch size assessed, 5, although has the largest parameter count mostly achieved the lowest MAE and RMSE for the 4 variables, underscoring its ability to capture spatiotemporal dynamics with finer resolution. Conversely, enlarging the patch size led to a marginal increase in inference time.

We conduct an ablation study to evaluate the impact of various system optimization strategies on the AI surrogate training process. We test the performance on 1 GPU with 16 CPU cores and 250 GB RAM. Figure 9 illustrates the training throughput associated with our proposed method and its counterparts, each lacking one of the optimizations: pin memory, activation checkpointing, and data prefetching. The analysis indicates that removing pin memory results in the most substantial decrease in throughput, underscoring its critical role in efficient data transfer from CPU to GPU. The omission of activation checkpointing also increases time costs, suggesting their roles in enhancing computational efficiency through effective memory management. During model training, data prefetching is used to load data in advance, this results show that prefetching significantly reduces idle time by overlapping data loading with model computation, thus improving training efficiency.

Figure 10 illustrates the exploration of the weak scalability in the training AI surrogate model, presenting the training throughput as the system proportionally scales with the number of GPUs. Due to the large training instance size, an A100 GPU with 80 GB memory can initially train only one instance per GPU. By enabling activation checkpointing, we increase the training batch size to 2 per GPU, achieving higher training throughput by fully utilizing GPU computational capacity. The implementation of activation checkpointing not only allows for a larger batch size but also demonstrates improved memory efficiency. This enables simultaneous processing of more data, potentially leading to faster convergence and better utilization of available hardware. These improvements pave the way for AI surrogate training with even larger simulation sizes.