プリトヴィ 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文献の並行する発展は、複数の異なるユースケースに効果的に調整できるファウンデーションモデルに焦点を当てていますが、気象と気候側の発展は、中期予測を特に重視した単一のユースケースに焦点を当てています。こうしたギャップを埋めるために、160の変数を使用して開発された23億のパラメータを持つファウンデーションモデルであるPrithvi WxCを紹介します。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.