프리트비 WxC: 날씨 및 기후를 위한 기반 모델

Prithvi WxC: Foundation Model for Weather and Climate

September 20, 2024

저자: 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

초록

AI 에뮬레이터가 HPC 시스템에서 실행되는 전통적인 수치 기상 예측 모델의 성능을 견줄 수 있다는 인식에 자극받아, 예보, 다운스케일링 또는 현재 예보와 같은 사용 사례를 다루는 대규모 AI 모델이 점점 더 많아지고 있습니다. AI 문헌의 병행 발전은 효과적으로 여러 가지 다른 사용 사례를 다루기 위해 조정될 수 있는 기초 모델에 초점을 맞추는 반면, 날씨 및 기후 분야의 발전은 주로 중기 예보에 중점을 둔 특정 사용 사례에 초점을 맞추고 있습니다. 저희는 Prithvi WxC를 소개함으로써 이 간극을 메웁니다. 이 모델은 160가지 변수를 사용하여 개발된 23억 개의 파라미터 기초 모델로, Modern-Era Retrospective Analysis for Research and Applications, Version 2 (MERRA-2)에서 가져온 데이터를 활용합니다. Prithvi WxC는 인코더-디코더 기반 아키텍처를 채택하며, 최근의 다양한 트랜스포머 모델에서 개념을 통합하여 입력 데이터의 지역 및 전역 의존성을 효과적으로 포착합니다. 이 모델은 다양한 해상도에서 다양한 위상의 날씨 현상을 모델링하기 위해 대규모 토큰 수용량을 갖출 수 있도록 설계되었습니다. 더불어, 이 모델은 마스크된 재구성과 예측의 패러다임을 결합한 혼합 목적으로 훈련되었습니다. 저희는 Autoregressive rollout forecasting, Downscaling, Gravity wave flux parameterization, 그리고 Extreme events estimation과 같은 일련의 어려운 하류 작업에서 이 모델을 테스트했습니다. 23억 개의 파라미터를 가진 사전 훈련된 모델과 관련된 파인 튜닝 워크플로우는 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.