Prithvi WxC: Modelo Base para Tempo e Clima

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

Resumo

Impulsionado pela percepção de que os emuladores de IA podem rivalizar com o desempenho dos modelos tradicionais de previsão numérica do tempo em execução em sistemas de HPC, há agora um número crescente de grandes modelos de IA que abordam casos de uso como previsão, redução de escala ou previsão imediata. Enquanto os desenvolvimentos paralelos na literatura de IA se concentram em modelos fundamentais - modelos que podem ser efetivamente ajustados para lidar com múltiplos casos de uso diferentes - os desenvolvimentos no lado do clima e do tempo se concentram principalmente em casos de uso únicos com ênfase particular em previsões de médio prazo. Fechamos essa lacuna ao introduzir o Prithvi WxC, um modelo fundamental com 2,3 bilhões de parâmetros desenvolvido usando 160 variáveis da Análise Retrospectiva da Era Moderna para Pesquisa e Aplicações, Versão 2 (MERRA-2). O Prithvi WxC emprega uma arquitetura baseada em codificador-decodificador, incorporando conceitos de vários modelos de transformadores recentes para capturar efetivamente dependências regionais e globais nos dados de entrada. O modelo foi projetado para acomodar grandes contagens de tokens para modelar fenômenos climáticos em diferentes topologias em resoluções finas. Além disso, ele é treinado com um objetivo misto que combina os paradigmas de reconstrução mascarada com previsão. Testamos o modelo em um conjunto de tarefas desafiadoras subsequentes, a saber: previsão de rolagem autoregressiva, redução de escala, parametrização de fluxo de ondas gravitacionais e estimativa de eventos extremos. O modelo pré-treinado com 2,3 bilhões de parâmetros, juntamente com os fluxos de trabalho de ajuste fino associados, foi publicamente lançado como uma contribuição de código aberto 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.