Hierarchical U-net with re-parameterization technique for spatio-temporal weather forecasting

Abstract

Due to the considerable computational demands of physics-based numerical weather prediction, especially when modeling fine-grained spatio-temporal atmospheric phenomena, deep learning methods offer an advantageous approach by leveraging specialized computing devices to accelerate training and significantly reduce computational costs. Consequently, the application of deep learning methods has presented a novel solution in the field of weather forecasting. In this context, we introduce a groundbreaking deep learning-based weather prediction architecture known as Hierarchical U-Net (HU-Net) with re-parameterization techniques. The HU-Net comprises two essential components: a feature extraction module and a U-Net module with re-parameterization techniques. The feature extraction module consists of two branches. First, the global pattern extraction employs adaptive Fourier neural operators and self-attention, well-known for capturing long-term dependencies in the data. Second, the local pattern extraction utilizes convolution operations as fundamental building blocks, highly proficient in modeling local correlations. Moreover, a feature fusion block dynamically combines dual-scale information. The U-Net module adopts RepBlock with re-parameterization techniques as the fundamental building block, enabling efficient and rapid inference. In extensive experiments carried out on the large-scale weather benchmark dataset WeatherBench at a resolution of 1.40625$^\circ$, the results demonstrate that our proposed HU-Net outperforms other baseline models in both prediction accuracy and inference time.