[論文レビュー] Scaling transformer neural networks for skillful and reliable medium-range weather forecasting

Tung Nguyen, Rohan Shah+7 more|arXiv (Cornell University)|Dec 6, 2023

Meteorological Phenomena and Simulations被引用数 19

ひとこと要約

天気に特化した埋め込み、乱択ダイナミクス予測、圧力重み付き損失を備えたシンプルなTransformerベースのモデルは、WeatherBench 2で競争力のある短期予測と、7日を超える長期予測で優れた性能を示し、データと計算資源がはるかに少なくて済む。

ABSTRACT

Weather forecasting is a fundamental problem for anticipating and mitigating the impacts of climate change. Recently, data-driven approaches for weather forecasting based on deep learning have shown great promise, achieving accuracies that are competitive with operational systems. However, those methods often employ complex, customized architectures without sufficient ablation analysis, making it difficult to understand what truly contributes to their success. Here we introduce Stormer, a simple transformer model that achieves state-of-the-art performance on weather forecasting with minimal changes to the standard transformer backbone. We identify the key components of Stormer through careful empirical analyses, including weather-specific embedding, randomized dynamics forecast, and pressure-weighted loss. At the core of Stormer is a randomized forecasting objective that trains the model to forecast the weather dynamics over varying time intervals. During inference, this allows us to produce multiple forecasts for a target lead time and combine them to obtain better forecast accuracy. On WeatherBench 2, Stormer performs competitively at short to medium-range forecasts and outperforms current methods beyond 7 days, while requiring orders-of-magnitude less training data and compute. Additionally, we demonstrate Stormer's favorable scaling properties, showing consistent improvements in forecast accuracy with increases in model size and training tokens. Code and checkpoints are available at https://github.com/tung-nd/stormer.

研究の動機と目的

• よりシンプルでスケーラブルなデータ駆動型アプローチを中期天気予測へ動機づける。
• パフォーマンスを推進する主要なアーキテクチャとトレーニング要素を特定する。
• 適切なトレーニングレシピを用いれば標準的なTransformerが複雑なモデルと同等または上回ることを示す。
• モデルサイズとデータのスケーリングが有利に働くことを示し、最先端のベースラインと比較する。

提案手法

• 天気特化の埋め込みを備え、変数をトークン化しクロスアテンションで集約する標準的なTransformerバックボーンを使用する。
• ランダムな間隔に渡る乱択ダイナミクス予測の目的で訓練し、天気のダイナミクス Δδtを予測する。
• near-surface variables を強調するために大気圧で損失に重みを付ける。
• 長期予測を改善するためのマルチステップ微調整 regime を採用する。
• 複数の区間ベースのロールアウトを組み合わせて推論する（best m in n または均質戦略）。
• WeatherBench 2 ERA5データで 1–14 day leads を評価し、Pangu-Weather、GraphCast、 climatology と比較する。

実験結果

リサーチクエスチョン

• RQ1 専門化された埋め込みとトレーニングレシピを備えた単純な Transformer は WeatherBench 2 で競争力のある短期予測と長期予測で優位性を示せるか？
• RQ2 乱択間隔予測と圧力重み付き損失はリードタイム全体で予測精度を有意に改善するか？
• RQ3 モデルサイズ、パッチサイズ、トレーニングトークン数は性能とスケーラビリティにどのような影響を与えるか？
• RQ4 長いリードタイムでのロールアウト誤差を減らすにはマルチステップ微調整が不可欠か？
• RQ5 提案手法はデータと計算効率の点で最先端のディープラーニングベースラインと比べてどうか？

主な発見

• モデルは1–7日予測で競合精度を達成し、7日を超えるとベースラインを上回る。
• 訓練はPangu-WeatherおよびGraphCastよりも桁違いに少ないデータと計算資源で済む。
• 乱択ダイナミクス予測は複数区間ロールアウトを追加の計算なしに可能にして精度を向上させる。
• 圧力重み付き損失とダイナミクス予測はこれらの成分を欠くモデルより優れている。
• より大きなモデルとより多いトレーニングトークンで性能が向上し、より小さなパッチサイズが利益をもたらす。