RAPID-Hydrology: Pretraining Hydro-Meteorological Representations With Diffusion for Enhanced Causal Prediction
Abstract
Accurate basin-scale prediction of hydro-meteorological variables is central to water-resource management and environmental monitoring, but target-basin observations are often insufficient for training high-capacity sequence models from scratch. RAPID-Hydrology extends the TimeDART diffusion-autoregressive backbone into a large-source/small-target hydrological transfer framework. It combines diffusion-based representation pretraining with a relative-position causal Transformer tailored to temporally ordered hydro-meteorological forecasting. HydroDiffuser first pretrains the encoder on a large-scale national hydro-meteorological dataset. The diffusion decoder is then discarded, and the learned encoder is transferred and fine-tuned with the limited local observations available for the Danjiangkou Reservoir basin. This design converts knowledge learned from broad climatic and spatial variability into deterministic 365-step prediction for a data-scarce target basin without diffusion sampling at inference. Experiments on precipitation, mean temperature, and wind speed show that RAPID-hydrology consistently outperforms CNN-BiLSTM-MTS, Transformer, and lag-aware causal Transformer. Using the 12-month average precipitation results, RAPID-hydrology reduces mean squared error (mse), mean absolute error (MAE), and root mse (RMSE) by 98.32%, 95.77%, and 86.08%, respectively, relative to the strongest evaluated from-scratch comparator, CNN-BiLSTM-MTS. The improvements are also supported by symmetric mean absolute percentage error (sMAPE), mean squared skill score (MSSS), anomaly correlation coefficient (ACC), spatial skill maps, and prediction-residual diagnostics. These results demonstrate that large-scale hydro-meteorological pretraining can support accurate long-horizon prediction when local basin observations are limited. The core implementation is available at https://github.com/Kilimetarlop/RAPID-Hydrology