FourCastNet3 (FCN3) は、NVIDIA Earth-2 による最新の AI グローバル気象予測システムです。FCN3 は、確率予測の精度、計算効率、スペクトル忠実度、アンサンブルの較正精度、そして亜季節スケールでの安定性という、これまでにない組み合わせを実現しています。中期予報における FCN3 の予測精度は、GenCast などの最先端の機械学習モデルに匹敵し、IFS-ENS などの従来の数値予報システムを上回っています。
0.25 度、6 時間間隔の解像度で 60 日分の FCN3 シミュレーションを実行しても、NVIDIA H100 Tensor コア GPU 1 基でわずか 4 分以内に計算が完了します。これは、GenCast の約 8 倍、IFS-ENS の約 60 倍の高速化に相当します。
また、優れた較正精度とスペクトル忠実度を備えており、60 日という長期リード タイムにおいても、アンサンブル各メンバーが現実的なスペクトル特性を維持しています。FCN3 は、中期から亜季節スケールにおける大規模アンサンブルを用いたデータ駆動型気象予測の分野で、大きな飛躍を示しています。
FCN3 アーキテクチャ
FourCastNet3 は、球面信号処理の基本原理に基づいた 完全畳み込み型の球面ニューラル オペレーター アーキテクチャを採用しています (図 2 を参照)。FourCastNet2 が球面フーリエ ニューラル オペレーターに基づいているのに対し、 FCN3 は局所的な球面畳み込みとスペクトル畳み込みの両方を組み合わせて使用しています。
これらの畳み込みはモルレー ウェーブレットを用いてパラメーター化されており、離散‐連続群畳み込みの枠組みで定式化されています。この手法により、局所的な大気現象に適した異方的で局所化されたフィルターを実現できる一方で、NVIDIA CUDA によるカスタム実装によって計算効率も保証されています。
FCN3 は、球面上の拡散過程によって進化が制御される潜在ノイズ変数を用いて、各予測ステップごとに確率性を導入しています。この隠れマルコフ形式により、アンサンブル メンバーを 1 ステップで効率的に生成できるようになります。これは、拡散モデル ベースの手法に対する大きな利点です。FCN3 はアンサンブルとして同時に学習されており、 空間領域およびスペクトル領域における連続順位確率スコア (CRPS) を組み合わせた複合損失関数を最小化するように訓練されています。この手法により、FCN3 は大気中の確率的プロセスに内在する空間的な相関関係を正しく学習できるようになります。
機械学習モデルのスケール拡大は、競争力のある精度を達成するためにしばしば重要ですが、データ駆動型気象モデルにおいては、そのスケールの影響がこれまで検証されていません。FCN3 は、その計算的な野心という点で異例の存在です。FCN3 のスケーリングのために、従来の数値気象モデルで用いられる領域分割の概念に着想を得た、新しいモデル並列化のパラダイムを導入しました。
この手法により、モデルを複数のデバイス間で分割することで、トレーニング中により大きなモデルを VRAM 内に収めることができ、同時に各デバイスあたりのディスク I/O も削減できます。これを可能にするために、NVIDIA Collective Communications Library (NCCL) を使用して畳み込みなどの空間操作を分散方式で実装します。この技術を用いることで、FCN3 は最大 1,024 基の GPU 上で、ドメイン並列、バッチ並列、アンサンブル並列を同時に活用して学習されています。トレーニング コードをご覧ください。
FourCastNet3 は、物理ベースの最良のアンサンブル モデルである IFS-ENS を上回り、予測精度の点では GenCast に匹敵します (図 3 参照)。NVIDIA H100 GPU 1 基で、FCN3 は 6 時間ごとの時間解像度と 0.25 度の空間解像度で 15 日間の予報を、わずか 1 分で生成します。これは、GenCast に比べて約 8 倍、IFS-ENS に比べて約 60 倍の高速化に相当します。
FCN3 の確率的アンサンブルは、一貫して 1 に近いスプレッド スキル比を示しており、予測された不確実性が実際の大気変動と密接に一致する、十分にキャリブレーションされた予報であることを示しています。ランク ヒストグラムおよびその他の診断結果から、アンサンブル メンバーが実際の観測値と同等に扱えることが確認されており、FCN3 の予測が信頼性と妥当性を兼ね備えていることが裏付けられています。
重要なのは、FCN3 があらゆるスケールにおいて大気のスペクトル特性を保持し、実際の気象パターンに見られるエネルギー カスケードや鋭い構造を、最大 60 日先という長期リード タイムでも忠実に再現している点です。多くの機械学習モデルが時間の経過とともに高周波成分をぼやかしたり、ノイズの多いアーティファクトに崩れていくのとは異なり、FCN3 は安定した物理的に現実的なスペクトルを維持し、亜季節スケールに至るまで正確で鮮明かつ物理的整合性のある予報を可能にしています。
これは図 4 に示されており、2020 年 2 月 11 日を初期値とした FCN3 による 850hPa 風速の予測を描いています。この日は、ストーム デニスがヨーロッパに上陸する直前の時期にあたります。FCN3 は、風速強度の大きさおよびそのスケールごとの変動を正確に捉えており、それは各予測結果における角度方向パワー スペクトル密度が忠実に再現されていることからも確認できます。この精度は、30 日 (720 時間) 以上の長期シミュレーションにおいても維持されています。
FourCastNet3 を始める
完全にトレーニングされた FourCastNet3 チェックポイントは NVIDIA NGC で利用できます。
FCN3 推論は、Earth2Studio で簡単に実行できます。単一の 4 メンバー アンサンブル推論を実行するなら、次のコードを実行できます。
from earth2studio.models.px import FCN3
from earth2studio.data import NCAR_ERA5
from earth2studio.io import NetCDF4Backend
from earth2studio.perturbation import Zero
from earth2studio.run import ensemble as run
import numpy as np
# load default package
model = FCN3.load_model(FCN3.load_default_package())
# determine output variables
out_vars = ["u10m", "v10m", "t2m", "msl", "tcwv"]
# data source initial condition
ds = NCAR_ERA5()
io = NetCDF4Backend("fcn3_ensemble.nc", backend_kwargs={"mode": "w"})
# no perturbation required due to hidden Markov formulation of FCN3
perturbation = Zero()
# invoke inference with 4 ensemble members
run(time=["2024-09-24"],
nsteps=16,
nensemble=4,
prognostic=model,
data=ds,
io=io,
perturbation=perturbation,
batch_size=1,
output_coords={"variable": np.array(out_vars)},
)
この推論の結果は図 5 に示されています。FCN3 を最適な性能で動作させるためには、カスタム CUDA 拡張を有効にした torch-harmonicsをインストールし、推論時には bf16 形式で自動混合精度を使用することを推奨します (これは Earth2Studio でのデフォルト設定です)。カスタム FCN3 推論を実行したり、自分でトレーニングしたりする場合は、makani でコードを見つけることができます。
FCN3 の詳細を見る
FourCastNet3 の関連情報はこちらでご覧いただけます。
- arXiv で公開されている FourCastNet3 論文
- FourCastNet3 チェックポイント
- GitHub で配布されている FCN3 のトレーニング/推論コード
- torch-harmonics による球面信号処理
- 推論のための Earth2Studio
- Physicsnemo
- arXiv で公開されている球面フーリエ ニューラル演算子論文
- 技術ブログ: 地球の大気を球面フーリエ ニューラル演算子でモデル化する
- FourCastNet2(SFNO) チェックポイント
著者一覧
Boris Bonev (NVIDIA)、Thorsten Kurth (NVIDIA)、Ankur Mahesh (LBNL)、Mauro Bisson (NVIDIA)、Karthik Kashinath (NVIDIA)、Anima Anandkumar (Caltech)、William D. Collins (LBNL)、Mike Pritchard (NVIDIA)、Alex Keller (NVIDIA)
