[論文レビュー] ArchesWeather: An efficient AI weather forecasting model at 1.5° resolution

Guillaume Couairon, Christian Lessig+2 more|arXiv (Cornell University)|May 23, 2024

Meteorological Phenomena and SimulationsEarth and Planetary Sciences被引用数 3

ひとこと要約

ArchesWeather は、1.5°解像度で計算効率に優れたAI天気予報モデルであり、3D 局所的アテンションを新しいクロスレベルアテンション（CLA）機構に置き換えることで、数日分のGPU訓練予算で最先端の性能を達成した。1.4°の50メンバーのNeuralGCMアンサンブルを上回り、Pangu-Weather や Stormer と同等またはそれ以上の主要指標を満たす一方で、標準的なハードウェアでも低コストな推論を可能にした。

ABSTRACT

One of the guiding principles for designing AI-based weather forecasting systems is to embed physical constraints as inductive priors in the neural network architecture. A popular prior is locality, where the atmospheric data is processed with local neural interactions, like 3D convolutions or 3D local attention windows as in Pangu-Weather. On the other hand, some works have shown great success in weather forecasting without this locality principle, at the cost of a much higher parameter count. In this paper, we show that the 3D local processing in Pangu-Weather is computationally sub-optimal. We design ArchesWeather, a transformer model that combines 2D attention with a column-wise attention-based feature interaction module, and demonstrate that this design improves forecasting skill. ArchesWeather is trained at 1.5° resolution and 24h lead time, with a training budget of a few GPU-days and a lower inference cost than competing methods. An ensemble of four of our models shows better RMSE scores than the IFS HRES and is competitive with the 1.4° 50-members NeuralGCM ensemble for one to three days ahead forecasting. Our code and models are publicly available at https://github.com/gcouairon/ArchesWeather.

研究の動機と目的

• 1.5°解像度で高い精度を維持しつつ、計算効率に優れたAI天気予報モデルの設計。
• 3D 局所的アテンションが天気予報モデリングにおいて最適であるという仮定を挑戦し、その計算上の非効率性を示す。
• 圧力層間のグローバルな特徴相互作用を可能にする、新しいクロスレベルアテンション（CLA）機構を通じて、予報精度の向上。
• モデルサイズや訓練コストを増加させることなく、訓練コストと推論時間を削減する。
• ERA5データにおける分布シフトの影響を調査し、最近のデータでの微調整により性能を向上させる。

提案手法

• Pangu-Weatherをインspiredした地球特有の位置バイアスを備えた3D Swin U-Netトランスフォーマー・アーキテクチャを採用。3D 局所的アテンションを、層間特徴相互作用に特化した非局所的クロスレベルアテンション（CLA）層に置き換えている。
• CLA層により、圧力層全体にわたるグローバルアテンションが可能となり、隣接層に限定される3D 局所的アテンションの計算非効率性を克服した。
• 予測対象を $X_{t+ au}$ ではなく $X_{t+ au} - X_t$ にすることで、訓練の安定性と一般化性能を向上させた。
• 風成分（U, V, W）を入力特徴として含め、日付と月を条件として与えることで、時間的一般化性能を向上させた。
• 分布シフトの影響を軽減し、最近の大気状態での性能向上を図るため、2007–2018年のERA5データを用いた微調整フェーズを実施した。
• 複数の独立して訓練されたモデル（例：ArchesWeather-M4, ArchesWeather-L2）の出力を平均することでアンサンブル予測を生成し、精度を向上させた。

実験結果

リサーチクエスチョン

• RQ1 Pangu-Weatherのようなモデルにおける3D 局所的アテンションは、物理的妥当性は高いが、AI天気予報において計算的に非効率であるとされる。その仮定は妥当か？
• RQ2 圧力層間の非局所的アテンション機構により、モデルサイズや訓練コストを増加させることなく、予報精度を向上させられるか？
• RQ3 2007–2018年の最近のERA5データでの微調整により、歴史的データにおける分布シフトに起因する誤差が低減されるか？
• RQ4 1.5°解像度で訓練コストが数日分のGPUで済むモデルが、競争力のある性能を達成できるか？
• RQ5 24時間先の予報において、ArchesWeatherはNeuralGCM、Pangu-Weather、Stormerといった最先端モデルと比較してどの程度の性能を示すか？

主な発見

• ArchesWeather-M は2台のA100 GPUで約2日間の訓練で、IFS HRESを平均して -5.0% の相対RMSE改善（RRH）を達成し、1.4°の50メンバーのNeuralGCMアンサンブルを上回った。
• ArchesWeather-L2アンサンブル（2つのLモデル）は、Z500, T850, Q700, U850, V850 といった主要な上層大気変数において、1.4°の50メンバーのNeuralGCMアンサンブルよりも低いRMSEを達成した。
• CLAモジュールを導入することで、CLAを搭載しないベースラインモデルと比較し、Z500のRMSEが10.5%低減され、IFS HRESの性能に近づいた。
• 2007–2018年の最近のERA5データでの微調整により、Z500のRMSEが1.2%低減され、T2mのRMSEが0.001°C低減された（1979–2018年全期間で訓練した場合と比較）。
• 単一のA100 GPU上で1回の24時間予報の推論に約0.25秒で実行可能であり、低コストなデプロイが可能になった。
• 訓練コストは約10 V100日で、SphericalCNNの40倍、Stormerの256倍も少ないが、その精度を同等または上回った。