Prithvi WxC: Foundation Model for Weather and Climate

TL;DR

Prithvi WxC is introduced, a 2.3 billion parameter foundation model developed using 160 variables from the Modern-Era Retrospective Analysis for Research and Applications, Version 2 (MERRA-2), incorporating concepts from various recent transformer models to effectively capture both regional and global dependencies in the input data.

Abstract

Triggered by the realization that AI emulators can rival the performance of traditional numerical weather prediction models running on HPC systems, there is now an increasing number of large AI models that address use cases such as forecasting, downscaling, or nowcasting. While the parallel developments in the AI literature focus on foundation models -- models that can be effectively tuned to address multiple, different use cases -- the developments on the weather and climate side largely focus on single-use cases with particular emphasis on mid-range forecasting. We close this gap by introducing Prithvi WxC, a 2.3 billion parameter foundation model developed using 160 variables from the Modern-Era Retrospective Analysis for Research and Applications, Version 2 (MERRA-2). Prithvi WxC employs an encoder-decoder-based architecture, incorporating concepts from various recent transformer models to effectively capture both regional and global dependencies in the input data. The model has been designed to accommodate large token counts to model weather phenomena in different topologies at fine resolutions. Furthermore, it is trained with a mixed objective that combines the paradigms of masked reconstruction with forecasting. We test the model on a set of challenging downstream tasks namely: Autoregressive rollout forecasting, Downscaling, Gravity wave flux parameterization, and Extreme events estimation. The pretrained model with 2.3 billion parameters, along with the associated fine-tuning workflows, has been publicly released as an open-source contribution via Hugging Face.

Figures (11)

• Figure 1: Prithvi WxC core architecture elements and masking scheme. For simplicity the figure ignores elements such as embedding and output layers as well as position encodings.
• Figure 2: Zero-shot reconstruction performance with Prithvi WxC. The first row shows "local" masking where we mask 95% of individual tokens. The second row shows "global" masking where we mask 75% of attention windows. At these extreme masking ratios some fine structure is lost in the reconstruction. See figure \ref{['fig:zero_shot_reconstruction']} for metrics.
• Figure 3: Zero-shot reconstruction performance of Prithvi WxC evaluated with 50, 60, 70, 80, 90, 95 and 99% masking. Note that the 6-hour ahead values are without any forecast tuning.
• Figure 4: Zero-shot forecasting performance of Prithvi WxC.
• Figure 5: (a) The track of Category 4 Hurricane Ida (2021) is shown from HURDAT, MERRA-FCN (FourCastNet model trained on MERRA-2 dataset), ERA-FCN (FourCastNet model trained on ERA5 dataset), and WxC models. All models were initialized at 00 UTC on 2021-08-27. The track errors of the models, compared to the HURDAT track, are 201.9 km for MERRA-FCN, 262.32 km for ERA-FCN, and 63.9 km for WxC, as noted in the legend. (b) A 5-day forecast of Mean Sea Level Pressure (MSLP) from MERRA-FCN, ERA-FCN, and WxC models. (c-e) Spatial distribution of Sea Level Pressure (SLP) for a 60-hour forecast (valid for 12 UTC on 2021-08-29). Among the models, the WxC model predicts the hurricane landfall most accurately in terms of both spatial location and timing, compared to the HURDAT reference.