Prithvi WxC: Modelo Base para Clima y Tiempo atmosférico

Prithvi WxC: Foundation Model for Weather and Climate

September 20, 2024

Autores: Johannes Schmude, Sujit Roy, Will Trojak, Johannes Jakubik, Daniel Salles Civitarese, Shraddha Singh, Julian Kuehnert, Kumar Ankur, Aman Gupta, Christopher E Phillips, Romeo Kienzler, Daniela Szwarcman, Vishal Gaur, Rajat Shinde, Rohit Lal, Arlindo Da Silva, Jorge Luis Guevara Diaz, Anne Jones, Simon Pfreundschuh, Amy Lin, Aditi Sheshadri, Udaysankar Nair, Valentine Anantharaj, Hendrik Hamann, Campbell Watson, Manil Maskey, Tsengdar J Lee, Juan Bernabe Moreno, Rahul Ramachandran

cs.AI

Resumen

Impulsado por la realización de que los emuladores de IA pueden rivalizar con el rendimiento de los modelos tradicionales de predicción numérica del clima que se ejecutan en sistemas HPC, ahora hay un número creciente de grandes modelos de IA que abordan casos de uso como pronósticos, escalamiento descendente o pronóstico inmediato. Mientras que los desarrollos paralelos en la literatura de IA se centran en modelos fundamentales, es decir, modelos que pueden ajustarse eficazmente para abordar múltiples casos de uso diferentes, los desarrollos en el ámbito del clima y clima se centran principalmente en casos de uso únicos con un énfasis particular en el pronóstico a medio plazo. Cerramos esta brecha presentando Prithvi WxC, un modelo fundamental de 2.3 mil millones de parámetros desarrollado utilizando 160 variables del Análisis Retrospectivo de la Era Moderna para Investigación y Aplicaciones, Versión 2 (MERRA-2). Prithvi WxC emplea una arquitectura basada en codificador-decodificador, incorporando conceptos de varios modelos transformadores recientes para capturar eficazmente dependencias regionales y globales en los datos de entrada. El modelo ha sido diseñado para acomodar un gran número de tokens para modelar fenómenos climáticos en diferentes topologías a resoluciones finas. Además, se entrena con un objetivo mixto que combina los paradigmas de reconstrucción enmascarada con pronósticos. Probamos el modelo en un conjunto de desafiantes tareas secundarias, a saber: Pronóstico de despliegue autoregresivo, Escalamiento descendente, Parametrización de flujo de ondas gravitatorias y Estimación de eventos extremos. El modelo preentrenado con 2.3 mil millones de parámetros, junto con los flujos de trabajo de ajuste fino asociados, se ha lanzado públicamente como una contribución de código abierto a través de Hugging Face.

English

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.