Systematic Evaluation of a Transformer-Based Surrogate for Distributed Groundwater Dynamics: Long-Horizon Prediction, Dynamic–Static Fusion, and Data Augmentation

Hui Huang, Aoqi Sun, Hoang Tran, Yongjiu Dai, Chen Yang

Transformers are increasingly reshaping artificial-intelligence-based Earth system prediction, yet their suitability for distributed groundwater surrogate modeling remains poorly understood. Groundwater dynamics combine strong temporal persistence with spatial interactions governed by heterogeneous hydrogeological conditions. These characteristics make attention-based architectures potentially well suited to representing the long-range spatiotemporal dependencies of groundwater systems, rather than merely providing an alternative to recurrent models. Here, we systematically evaluated a groundwater-adapted PredFormer as a surrogate for regional ParFlow-CLM simulations. The adapted model combined temporal-first factorized attention with pressure-residual prediction and blockwise autoregressive rollout. We assessed its long-horizon prediction behavior over 720 h, the contributions of prescribed P−ET (PME) forcing and static attributes, alternative static-fusion strategies, its performance relative to PredRNN, and precipitation-regime-specific data augmentation generated by perturbing precipitation forcing. Over the horizon, the adapted PredFormer maintained decimeter-scale median water table depth (WTD) errors, although error magnitude varied among WTD regimes. Larger upper-tail spatial errors were associated with precipitation amount and variability, whereas temporal fluctuations in the error trajectories were more closely associated with abrupt wet–dry transitions. Adding prescribed PME produced the largest and most consistent reduction in prediction error, although the benefits of dynamic and static inputs varied systematically with WTD. Post-temporal–spatial cross-attention performed best among the static-fusion strategies. PredRNN was more accurate initially, but PredFormer performed better at longer lead times and in most segments. Data augmentation produced modest overall changes but clear regime-dependent benefits when the added samples matched the precipitation conditions encountered during prediction. These findings show that the advantages of Transformer-based groundwater surrogates emerge primarily over extended autoregressive horizons and depend on groundwater state, input composition, information-fusion strategy, and training-sample coverage. The results provide process-informed guidance for developing computationally efficient groundwater surrogates for regional prediction and water-resources management, while also highlighting their longer-term potential to provide AI-enabled representations of groundwater dynamics within Earth system models.