Prithvi WxC : Modèle de base pour la météo et le climat

Prithvi WxC: Foundation Model for Weather and Climate

September 20, 2024

Auteurs: 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

Résumé

Suite à la réalisation que les émulateurs d'IA peuvent rivaliser avec les performances des modèles traditionnels de prévision numérique du temps fonctionnant sur des systèmes HPC, on observe maintenant un nombre croissant de grands modèles d'IA qui abordent des cas d'utilisation tels que la prévision, le rééchantillonnage ou la prévision immédiate. Alors que les développements parallèles dans la littérature sur l'IA se concentrent sur les modèles fondamentaux - des modèles qui peuvent être efficacement ajustés pour aborder plusieurs cas d'utilisation différents - les développements du côté météorologique et climatique se concentrent largement sur des cas d'utilisation uniques avec un accent particulier sur la prévision à moyen terme. Nous comblons cette lacune en introduisant Prithvi WxC, un modèle fondamental de 2,3 milliards de paramètres développé en utilisant 160 variables de l'Analyse Rétrospective de l'Ère Moderne pour la Recherche et les Applications, Version 2 (MERRA-2). Prithvi WxC utilise une architecture basée sur un encodeur-décodeur, incorporant des concepts de divers modèles de transformateurs récents pour capturer efficacement les dépendances régionales et mondiales dans les données d'entrée. Le modèle a été conçu pour prendre en charge de grands nombres de jetons pour modéliser les phénomènes météorologiques dans différentes topologies à des résolutions fines. De plus, il est entraîné avec un objectif mixte qui combine les paradigmes de la reconstruction masquée avec la prévision. Nous testons le modèle sur un ensemble de tâches aval difficiles, à savoir : la prévision de déploiement autorégressif, le rééchantillonnage, la paramétrisation du flux d'ondes gravitationnelles et l'estimation des événements extrêmes. Le modèle pré-entraîné avec 2,3 milliards de paramètres, ainsi que les flux de travail de fine-tuning associés, ont été publiquement publiés en tant que contribution open-source via 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.