[論文レビュー] W-MAE: Pre-trained weather model with masked autoencoder for multi-variable weather forecasting

Xin Man, Chenghong Zhang+3 more|arXiv (Cornell University)|Apr 18, 2023

Hydrological Forecasting Using AI被引用数 7

ひとこと要約

W-MAE は ERA5 データを用いた自己監視型マスクドオートエンコーダ pre-training を利用して空間的気象構造を学習し、マルチ変数および降水予測のファインチューニングを行い、頑健性を向上させ、フォークのエンサンブル準備が整った迅速な予測を可能にする。

ABSTRACT

Weather forecasting is a long-standing computational challenge with direct societal and economic impacts. This task involves a large amount of continuous data collection and exhibits rich spatiotemporal dependencies over long periods, making it highly suitable for deep learning models. In this paper, we apply pre-training techniques to weather forecasting and propose W-MAE, a Weather model with Masked AutoEncoder pre-training for weather forecasting. W-MAE is pre-trained in a self-supervised manner to reconstruct spatial correlations within meteorological variables. On the temporal scale, we fine-tune the pre-trained W-MAE to predict the future states of meteorological variables, thereby modeling the temporal dependencies present in weather data. We conduct our experiments using the fifth-generation ECMWF Reanalysis (ERA5) data, with samples selected every six hours. Experimental results show that our W-MAE framework offers three key benefits: 1) when predicting the future state of meteorological variables, the utilization of our pre-trained W-MAE can effectively alleviate the problem of cumulative errors in prediction, maintaining stable performance in the short-to-medium term; 2) when predicting diagnostic variables (e.g., total precipitation), our model exhibits significant performance advantages over FourCastNet; 3) Our task-agnostic pre-training schema can be easily integrated with various task-specific models. When our pre-training framework is applied to FourCastNet, it yields an average 20% performance improvement in Anomaly Correlation Coefficient (ACC).

研究の動機と目的

• Introduce self-supervised pre-training for weather forecasting to leverage large unlabeled meteorological data.
• Model spatial dependencies in meteorological fields using a Vision Transformer backbone with masked autoencoder pre-training.
• Fine-tune the pre-trained model for multi-variate and precipitation forecasting and assess robustness and efficiency.
• Demonstrate that task-agnostic pre-training can enhance downstream weather models and enable fast ensemble forecasting.

提案手法

• ERA5 サンプルを 20 チャンネルの天気画像として扱い、パッチに分割する。
• MAE に類似した pre-training を高いマスキング比で適用し、マスクされたパッチを再構成する。
• unmasked パッチを処理する ViT エンコーダと、画素を再構成する修正版の二次元パッチベースのデコーダを用いる。
• 事前学習済みモデルをマルチ変数予測のためにファインチューニングし、将来状態を自己回帰的に予測する。
• FourCastNet と比較して降水量および他の変数を評価し、推論速度とメモリを分析する。

実験結果

リサーチクエスチョン

• RQ1 タスク非特異的で自己監視型の pre-training は短期〜中期の天気予報の安定性と精度を改善するか。
• RQ2 W-MAE はマルチ変量および降水予測タスクで既存モデル（例：FourCastNet）を上回るか。
• RQ3 特にエンサンブル予測における pre-training は推論速度とメモリ使用にどのように影響するか。
• RQ4 pre-training フレームワークは他のタスク特化型天気モデルへ容易に移植可能か。

主な発見

• W-MAE は低層風や 500 hPa 干渉高度などの変数で、明示的な誤差蓄積最適化なしに安定した 10 日予測性能を提供する。
• W-MAE は降水予測で FourCastNet を大幅に上回る。
• pre-training を FourCastNet に適用すると降水予測で約 20% の改善と収束の速さが向上する。
• カスタマイズされた W-MAE デコーダは計算とメモリのオーバーヘッドを削減し、1 エポックあたりのトレーニング速度を 1.2 倍、ベースの MAE に対して最大 8% のメモリ節約を達成する。
• pre-training フレームワークは他のタスク特化型モデルと統合可能で、複数の予測タスクの共通基盤として機能する。