[論文レビュー] ClimaX: A foundation model for weather and climate

Tung Nguyen, J. Brandstetter+3 more|arXiv (Cornell University)|Jan 24, 2023

Meteorological Phenomena and Simulations被引用数 168

ひとこと要約

ClimaX は、異種の CMIP6 気候データで事前学習した柔軟なトランスフォーマーベースのファウンデーションモデルで、様々な気象・気候タスクに対する微調整を可能にし、優れた一般化能力と競争力のあるベンチマークを実現します。

ABSTRACT

Most state-of-the-art approaches for weather and climate modeling are based on physics-informed numerical models of the atmosphere. These approaches aim to model the non-linear dynamics and complex interactions between multiple variables, which are challenging to approximate. Additionally, many such numerical models are computationally intensive, especially when modeling the atmospheric phenomenon at a fine-grained spatial and temporal resolution. Recent data-driven approaches based on machine learning instead aim to directly solve a downstream forecasting or projection task by learning a data-driven functional mapping using deep neural networks. However, these networks are trained using curated and homogeneous climate datasets for specific spatiotemporal tasks, and thus lack the generality of numerical models. We develop and demonstrate ClimaX, a flexible and generalizable deep learning model for weather and climate science that can be trained using heterogeneous datasets spanning different variables, spatio-temporal coverage, and physical groundings. ClimaX extends the Transformer architecture with novel encoding and aggregation blocks that allow effective use of available compute while maintaining general utility. ClimaX is pre-trained with a self-supervised learning objective on climate datasets derived from CMIP6. The pre-trained ClimaX can then be fine-tuned to address a breadth of climate and weather tasks, including those that involve atmospheric variables and spatio-temporal scales unseen during pretraining. Compared to existing data-driven baselines, we show that this generality in ClimaX results in superior performance on benchmarks for weather forecasting and climate projections, even when pretrained at lower resolutions and compute budgets. The source code is available at https://github.com/microsoft/ClimaX.

研究の動機と目的

• タスク固有の制約を克服するため、気象と気候の一般的でデータ駆動型のファウンデーションモデルの創造を動機づける。
• 多様で物理情報を取り入れた気候データセット（CMIP6）を事前学習に活用して一般性を高める。
• 変数モダリティと不規則な時空カバレッジを可能にするアーキテクチャ上の革新を開発する。
• 予測、気候投影、ダウンスケーリングタスクの微調整の汎用性を実証する。
• データ量、モデルサイズ、解像度によるスケーリング挙動を評価し、今後の研究の指針とする。

提案手法

• ビジョン・トランスフォーマー（ViT）を二つの重要なモジュールとともに拡張する：変数トークン化（各入力変数を別々にトークン化）と変数集約（各空間位置でのクロス・アテンションを用いて変数を融合）。
• 任意の将来点にわたる時系列予測を導くために、位置とリードタイムの埋め込みを追加する。
• CMIP6 データで ClimaX を事前学習し、6h〜168hのリードタイム全域で現在の状態から未来の状態を予測するランダム化予測目的を設定する。
• 地球規模のグリッド化における球面積の違いを考慮するため、緯度加重平均二乗誤差を用いる。
• 事前学習中に見られた変数と見られない変数を含む下流タスクへの微調整を、モジュール式のファインチューニング戦略で実行する。
• 予測、気候投影、ダウンスケーリングのベンチマークで評価し、ニューラルベースラインおよびNWPベースラインと比較する。

実験結果

リサーチクエスチョン

• RQ1 単一の事前学習済み ClimaX モデルは、異種データを用いる複数の天気・気候タスクに適応できるか？
• RQ2 グローバル/地域予測、準季節的〜季節的予測、気候投影、ダウンスケーリングにおける ClimaX の性能は、専門的なベースラインと比較してどうか？
• RQ3 データセットサイズ、モデル容量、解像度が ClimaX の性能に与える影響（スケーリング法則）は？
• RQ4 未知の変数や変動する時空解像度へ ClimaX を最も効果的に転移させるファインチューニング戦略は何か？

主な発見

• 単一の事前学習済み ClimaX は、異なるリードタイム、解像度、および地域にわたる広範なタスクへ微調整可能。
• ClimaX は ClimateBench で最先端の性能を達成し、WeatherBench における運用 IFS とも競合する，モデレートな計算資源でも。
• スケーリング分析は、より多くの事前学習データ、より大きなモデル、より高い解像度の入力で性能が向上することを示す。
• アブレーション研究は、未知変数の取り扱いや共有コンポーネントを伴う複数タスクを含む効果的なファインチューニング戦略を支持する。
• 変数トークン化と集約のおかげで、データ覆域の不規則性や変数セットに対してロバスト性を維持する。