24AI Data-Driven Modeling for Medium-Range Weather Prediction Using Different Architectures

Artificial intelligence (AI) is transforming the landscape of numerical weather prediction (NWP), offering performance that now rivals and even surpasses traditional physically based models used by leading institutions such as ECMWF and NOAA. This presentation provides an in-depth exploration of our research into the application of AI in medium-range weather prediction (MWP). Our primary goal is to establish a robust testbed to understand the key factors driving the success of data-driven forecasts and to assess the infrastructure required for training state-of-the-art AI models in this domain.

Our work leverages the ERA5 reanalysis dataset at multiple resolutions—5.6, 2.8, and 1.4 degrees. The central architecture employed in this study is the squeeze-and-excitation residual neural network (SE-ResNet), chosen for its balance of simplicity and effectiveness. We extend our exploration to include dilated-convolution layers and more advanced architectures such as Vision Transformers. Additionally, we investigate foundational models that can be trained on diverse datasets, offering greater flexibility in data-driven modeling.

Training is conducted on NOAA's Hera cluster, utilizing P100 GPUs. Notably, training a 30-block SE-ResNet with 192 filters on the coarsest resolution dataset (5.6 degrees) is completed within just three days on a single P100 GPU, achieving performance comparable to the best models in WeatherBench version 1. As the resolution increases to 2.8 and 1.4 degrees, the model’s performance improves significantly, approaching the state of the art. Specifically, the model trained on 2.8-degree resolution outperforms FourCastNet and is competitive with the ClimaX foundational model. Following the methodology of the Stormer model, our models are trained to predict multiple lead times (6, 12, and 24 hours) concurrently. During inference, we combine these base lead times in an auto-regressive manner to extend predictions out to 10-15 days.

In this talk, we will delve into the implementation details and discuss various advanced concepts from the literature, including multi-resolution training, randomized iterative forecasting, and multi-lead time training. We will also address the challenges of scaling up training on older-generation GPUs in the Hera cluster and share practical workarounds to enhance scalability. This presentation aims to provide both theoretical insights and practical guidance for the future of AI-driven medium-range weather prediction.