A Transformer-Based Surrogate for Distributed Groundwater Dynamics with Zero-Shot Cross-Basin Transfer and Precipitation-Regime Data Augmentation

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

Large-scale physically based groundwater models generate spatially distributed simulations, but their cost limits the regional datasets available for training artificial-intelligence-based surrogates. Here, we developed and systematically evaluated a groundwater-adapted PredFormer surrogate for regional ParFlow-CLM simulations. The framework uses temporal-first factorized attention, pressure-residual prediction, prescribed precipitation-minus-evapotranspiration (PME) forcing, and static hydrogeological attributes. The model was developed and evaluated over the 34,000 km² upper Beijiang River basin (UBRB), with zero-shot transfer tested over the 24,000 km² Hejiang basin. Across 12 independent 720 h test segments, PredFormer maintained decimeter-scale water-table-depth errors. Upper-tail errors were associated with precipitation statistics, whereas abrupt wet–dry transitions coincided with temporal error fluctuations. PME reduced the overall mean absolute error from 0.151 to 0.099 m. Depth-stratified analysis showed that PME was most beneficial in shallow-water-table areas, whereas static attributes contributed most at intermediate depths and little in deep areas. PredRNN was more accurate initially, but PredFormer performed better at extended lead times, reducing the across-segment mean error at 720 h from 0.232 to 0.152 m. Without target-basin training, fine-tuning, or recalculation of standardization statistics, zero-shot transfer produced an overall mean error of 0.106 m in Hejiang, close to 0.099 m in the UBRB. Precipitation-regime data augmentation produced modest overall changes but clearer benefits when added samples matched prediction conditions; in a precipitation-active segment, the mean error decreased from 0.141 to 0.095 m. Together, these results show how groundwater state, training-data coverage, and cross-basin transfer shape the performance of Transformer-based groundwater surrogates.