[Paper Review] Prithvi WxC: Foundation Model for Weather and Climate

Johannes Schmude, Sujit Roy+27 more|arXiv (Cornell University)|Sep 20, 2024

Geophysics and Gravity Measurements7 citations

TL;DR

Prithvi WxC is a 2.3B-parameter transformer-based foundation model trained on 160 MERRA-2 atmospheric variables, designed for forecasting, downscaling, gravity wave parameterization, and extreme-events estimation, with open-source release on Hugging Face.

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.

Motivation & Objective

• Motivate the need for a foundation-model approach in weather and climate spanning global and regional tasks.
• Develop a scalable encoder-decoder transformer architecture capable of handling large token counts and diverse topologies.
• Train with a mixed objective combining masked reconstruction and forecasting grounded in climatology.
• Evaluate zero-shot reconstruction and forecasting performance, plus downstream tasks like downscaling and gravity-wave parameterization.
• Publicly release pretrained weights and fine-tuning workflows to enable broad usage in the community.

Proposed method

• Encoder-decoder transformer with local and global attention across windows, enabling long-range dependencies.
• 2D vision transformer with 50% input masking and a flexible tokenization scheme (windows and tokens per window).
• Pretraining objective that blends masked reconstruction with climatology-based forecasting, including a normalized residual formulation (Equation 2).
• Training on 160 MERRA-2 variables (1980-2019) with static inputs, climatology C_t, and lead-time parameters δt and δτ.
• Two-phase pretraining: initial phase with masking and random lead times, followed by fine-tuning for forecasting with adjusted losses and autoregressive rollout.
• Downstream fine-tuning workflows for downscaling and gravity-wave parameterization.

Experimental results

Research questions

• RQ1 Can a single foundation model handle global and regional weather-climate tasks with zero-shot capabilities?
• RQ2 What pretraining objective best supports both reconstruction and forecasting across variable lead times?
• RQ3 How does a large WxC model perform on autoregressive forecasting, downscaling, and hurricane-track prediction compared to task-specific models?
• RQ4 What are effective strategies for integrating climatology and static features into a unified pretraining objective?
• RQ5 What are the practical fine-tuning considerations for downstream tasks (downscaling, gravity-wave parameterization, extreme events)?

Key findings

• Zero-shot reconstruction works with high masking (as low as 5-25% remaining data) under local or global masking schemes.
• Zero-shot forecasting performs strongest at short lead times (6-12 hours) with reasonable performance up to ~66 hours, declining thereafter relative to some baselines.
• In hurricane-track forecasting, Prithvi WxC achieves superior track accuracy and landfall timing compared with FourCastNet variants trained on MERRA-2 or ERA5 data (example: Ida, 2021; mean track error 63.9 km vs 201.9–262.3 km).
• Downscaling experiments show substantial gains over interpolation baselines (MERRA-2 T2m: spatial RMSE 3.22→0.73; temporal RMSE 2.46→0.64; correlation 0.89→0.98).
• CORDEX downscaling to 12.5 km resolution achieves similar improvements (spatial RMSE 1.89→0.44; temporal RMSE 1.14→0.37; correlation 0.99→1.00).
• The model is released publicly on Hugging Face with pretrained weights and finetuning workflows for downstream tasks.