Fourier Neural Operators and Transformers for Extreme Weather and Climate Prediction
, Senior AI Developer Technologist, AI-HPC, NVIDIA
, Director of ML Research, NVIDIA
Predicting extreme weather events in a warming world at fine scales is a grand challenge facing climate scientists. Policymakers and society at large depend on reliable predictions to plan for the disastrous impact of climate change and develop effective adaptation strategies. Deep learning (DL) offers novel methods that are potentially more accurate and orders of magnitude faster than traditional weather and climate models for predicting extreme events. The Fourier Neural Operator (FNO), a novel deep-learning method, has shown promising results for predicting complex systems, such as spatio-temporal chaos, turbulence, and weather phenomena. The power of deep learning is best realized at scale — when models are trained on very large datasets (hundreds of terabytes or more) utilizing thousands of GPUs, with data and model parallelism. Scaling the FNO model requires innovations using efficient token mixing in the paradigm of self-attention in vision transformers, which led to developing Adaptive FNO (AFNO). Further, embedding physics into FNO, the so-called PINO approach, improves predictive skill and increases physical consistency and fidelity of the model. We'll describe our results on extreme weather prediction using AFNO and PINO at scale, and highlight how scaling, optimization, and performance challenges were addressed on three large supercomputing systems: Selene, Perlmutter, and Jureca-DC. Our results show promise for large-scale DL potentially competing with state-of-the-art numerical weather prediction. We'll conclude with results on the science: accuracy, fidelity, and acceleration in predicting the complex spatio-temporal evolution of critical extreme weather events: hurricanes, extreme precipitation, heat waves, and cold waves. The work presented here is a significant step toward building a reliable high-fidelity, high-resolution digital twin of Earth for extremes, which is widely recognized as grand challenge in the climate community.