[論文レビュー] Prithvi WxC: Foundation Model for Weather and Climate

Johannes Schmude, Sujit Roy+27 more|arXiv (Cornell University)|Sep 20, 2024

Geophysics and Gravity Measurements被引用数 7

ひとこと要約

tldr: Prithvi WxC は、160 MERRA-2 大気変数上で訓練された2.3Bパラメータのトランスフォーマー型基盤モデルで、予測、降下スケーリング、重力波パラメータ化、極端イベント推定を目的とし、Hugging Face でオープンソース公開されている。

ABSTRACT

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.

研究の動機と目的

• 全球および地域のタスクを横断する気象・気候分野において、基盤モデルアプローチの必要性を喚起する。
• 大規模なトークン数と多様なトポロジーに対応できるスケーラブルなエンコーダ-デコーダー変換器アーキテクチャを開発する。
• 気候学に基づく予測とマスキング再構成を組み合わせた混合目的で訓練する。
• ゼロショット再構成と予測性能、および降下スケーリングや重力波パラメータ化などの下流タスクを評価する。
• 事前訓練済みウェイトとファインチューニングワークフローを公開して、コミュニティでの広範な活用を可能にする。

提案手法

• ウィンドウ全体で局所注意とグローバル注意を組み合わせたエンコーダ-デコーダー変換器で、長距離依存を可能にする。
• 入力マスキング50%と柔軟なトークン化方式（ウィンドウと1ウィンドウあたりのトークン数）を備えた2Dビジョントランスフォーマー。
• 正規化残差形式（式2）を含む、マスキング再構成と気候学に基づく予測を組み合わせた事前訓練目的。
• 静的入力、気候学C_t、およびリードタイムパラメータδtとδτを用いて、160個のMERRA-2変数（1980-2019）で訓練。
• 2段階の事前訓練：マスキングとランダムリードタイムを用いた初期フェーズ、続く予測の微調整フェーズで損失を調整し自己回帰ロールアウトを行う。
• 降下スケーリングと重力波パラメータ化のための下流微調整ワークフロー。

実験結果

リサーチクエスチョン

• RQ1 単一の基盤モデルがゼロショット機能を備えて全球・地域の気象・気候タスクを扱えるか？
• RQ2 可変リードタイムを横断する再構成と予測の両方を最も効果的に支援する事前訓練目的は何か？
• RQ3 タスク固有モデルと比較して、大規模WxCモデルは自己回帰予測、降下スケーリング、ハリケーン軌跡予測でどの程度性能を示すか？
• RQ4 気候学と静的特徴を統合した統一された事前訓練目的の効果的戦略は何か？
• RQ5 下流タスク（降下スケーリング、重力波パラメータ化、極端事象）における実践的なファインチューニングの考慮事項は何か？

主な発見

• ゼロショット再構成は局所マスキングまたはグローバルマスキングの下で、マスキング率が高くても（残存データが5-25%程度と低くても）機能する。
• ゼロショット予測は短リードタイム（6-12時間）で最も強く、約66時間まで妥当な性能を示し、それ以降は基準モデルと比較して低下する。
• ハリケーン軌跡予測では、Prithvi WxCはMERRA-2またはERA5データで訓練された FourCastNet 系列と比較して軌跡精度と上陸時刻の予測精度が優れている（ Ida, 2021 の例; 平均軌跡誤差63.9km 対 201.9–262.3km）。
• 降下スケーリング実験は、補間ベースラインより大幅な利得を示す（MERRA-2 T2m: 空間RMSE 3.22→0.73; 時間RMSE 2.46→0.64; 相関0.89→0.98）。
• CORDEXの12.5km解像度へのダウンスケーリングは同様の改善を達成（空間RMSE 1.89→0.44; 時間RMSE 1.14→0.37; 相関0.99→1.00）。
• モデルは公開され、Hugging Face上で事前訓練済みウェイトと下流タスクのファインチューニングワークフローが提供されている。