NVIDIA はフーリエ モデルを使用して前例のない天気予報精度を実現

現代の数値天気予報 (NWP) は 1920 年代にまで遡ります。今日では、数値天気予報はいたるところで行われており、輸送、物流、農業、エネルギー生産などの主要分野における経済計画に貢献しています。正確な天気予報は、異常事態を事前に知らせることで、数え切れないほどの命を救います。

過去数十年にわたって天気予報の質は着実に向上してきました。最も初期の単一点動的モデリング数値天気予報は、1922 年に英国の科学者ルイス・フライ・リチャードソンによって計算尺と対数表を使用して計算され、6 時間の気象予報を計算するのに 6 週間かかりました。 1950 年代までに、初期の電子コンピュータによって予測速度が大幅に向上し、将来の天気を予測するのに十分な速さで予報を実行できるようになりました。計算能力の向上に加え、物理現象のより深い理解とより質の高い大気観測によって、小規模プロセスのパラメータ化が改善され、天気予報も向上します。

SOTA NWP モデルと比較すると、データ駆動型のディープラーニング モデルは桁違いに安価であるため、天気を予測するためのこのような方法の開発への関心が高まっています。多くの研究では、気候モデルの出力、一般循環モデル (GCM)、再解析製品、または気候モデルの出力と再解析製品の組み合わせに基づいてトレーニングされたデータ駆動型モデルを構築することにより、大規模な大気循環を予測しようと試みられてきました。データ駆動型モデルは、NWP モデルに存在するモデルバイアスを克服し、低い計算コストで確率予測とデータ同化のための大規模なアンサンブルを生成することで、天気予報を改善する大きな可能性を秘めています。

ほとんどのデータ駆動型気象モデルは、低解像度データ、典型的にはRasp and Thuerey [2021b]のように5.625°解像度、またはWeyn et al. [2020]のように2°解像度を使用してトレーニングされます。これらの事前の試みは、いくつかの粗粒度で低解像度の大気変数を予測する上で良好な結果を達成しました。しかし、粗大化プロセスにより、重要な微細物理情報が失われます。データ駆動型モデルが真に効果を発揮するには、現在の SOTA 数値気象モデル (約 0.1°) と同等かそれ以上の解像度で予報を生成する必要があります。

最近の論文では、NVIDIA、ローレンス・バークレー国立研究所、ミシガン大学アナーバー校、ライス大学などの研究者らが、フーリエベースのニューラルネットワーク予測モデル FourCastNet を開発しました。このモデルは、赤道付近で約 30×30 km の空間解像度、地球全体のグリッド サイズ 720×1440 ピクセルに相当する 0.25° の解像度で、主要な気象変数の地球規模のデータ駆動型予測を生成できます。これにより、初めて欧州中期予報センター（ECMWF）の高解像度統合予報システム（IFS）モデルとの直接比較が可能になりました。

論文の宛先:
出典：http://arxiv.org/pdf/2202.11214.pdf

FourCastNet は、ノード時間ベースで従来の NWP モデルよりも約 45,000 倍高速です。 FourCastNet の桁違いの高速化と、高解像度での前例のない精度を組み合わせることで、非常に低コストで非常に大規模なアンサンブル予測を生成することができます。 FourCastNet は、ハリケーン、大気河川、異常降水量などの大規模なアンサンブル予報を数秒で生成することで確率天気予報を大幅に改善し、よりタイムリーで情報に基づいた災害対応を可能にします。

さらに、FourCastNet の信頼性が高く、高速で安価な地表付近の風速予測により、陸上および洋上風力発電所の風力資源計画を改善できます。 FourCastNet をトレーニングするために必要なエネルギーは、IFS モデル (メンバー 50 個) を使用して 10 日間の予測を生成するために必要なエネルギーとほぼ等しくなります。ただし、一度トレーニングすると、FourCastNet は IFS モデルよりも予測を生成するのに必要なエネルギーが 12,000 倍少なくなります。研究者たちは、FourCastNet のトレーニングは一度だけでよく、その後の微調整のエネルギー消費はごくわずかになることを期待しています。

実装技術の面では、FourCastNetはフーリエ変換ベースのトークン混合法[Guibas et al., 2022]とViTバックボーン[Dosovitskiy et al., 2021]を使用しています。このアプローチは、解像度に依存しない方法で学習され、流体力学などの難しい偏微分方程式のモデリングに成功している最新のフーリエ ニューラル演算子に基づいています。さらに、彼らは ViT バックボーンを選択しました。これは、長距離の依存関係を適切にモデル化できるためです。 ViT とフーリエベースのトークン メソッドのハイブリッドにより、きめ細かい特徴を解決でき、解像度とデータセット サイズに合わせて適切にスケーリングできる SOTA 高解像度モデルが生成されます。研究者らは、このアプローチにより、これまでにない解像度で高忠実度のデータ駆動型モデルをトレーニングできるようになると述べています。

モデルのトレーニング

欧州中期予報センター (ECMWF) は、公開されている包括的なデータセット ERA5 を提供しており、本研究ではこれを使用して FourCastNet をトレーニングします。彼らは、（1）地表から10メートル上空の風速と（2）6時間の総降水量という2つの大気変数に焦点を当てましたが、この研究では、いくつかの異なる垂直高度でのジオポテンシャル高度、気温、風速、相対湿度、地表気圧や平均海面気圧などの地表付近の変数など、他のいくつかの変数も予測しました。

FourCastNet モデル

高解像度の予測を生成するために、この研究では適応型フーリエニューラル演算子 (AFNO) モデルを使用しました。このニューラル ネットワーク アーキテクチャは高解像度の入力用に設計されており、ディープラーニング (DL) の最新の主要な進歩を 1 つのモデルに統合しています。つまり、FourCastNet は、困難な PDE システムのモデリングで優れたパフォーマンスを発揮することが示されている Fourier Neural Operator (FNO) と強力な ViT バックボーンを組み合わせています。

現在、研究者らは計算の複雑さを軽減するために、いくつかの ViT バリアントを提案しています。しかし、AFNO モデルのユニークな点は、混合操作を連続的なグローバル畳み込みとして構築し、フーリエ領域で FFT を介して効率的に実装することで、空間およびチャネル次元にわたる依存関係の柔軟でスケーラブルなモデリングを可能にすることです。この設計により、空間混合の複雑さは O(N log N) に軽減されます。ここで、N は画像パッチまたはトークンの数です。このスケーリングにより、AFNO モデルは 0.25° 解像度の高解像度データだけでなく、より高い解像度での将来の研究にも適したものになります。

モデルの計算フローは次のとおりです。まず、720 × 1440 緯度経度グリッド上の入力変数が 2D パッチ グリッド (h × w) (パッチ サイズは p × p、たとえば p = 8) に投影され、各パッチは d 次元トークンとして表されます。次に、パッチ シーケンスは、位置エンコーディングとともに一連の AFNO レイヤーに送られます。

電車

この研究は主に地表の風速と降水量の予測に焦点を当てていますが、複雑な大気システムには、地球の表面から成層圏までの温度、地表気圧、湿度、水分含有量など、複数の変数間の強い非線形相互作用が関係しています。これらの相互作用をシミュレートするために、本研究では大気の瞬間的な状態を表すいくつかの変数（表 1）を選択しました。

この研究では、ERA5 データセットをトレーニング データセット、検証データセット、テスト データセットの 3 つのグループに分割します。トレーニング データセットは 1979 年から 2015 年までのデータで構成され、検証データセットには 2016 年と 2017 年のデータが含まれており、テスト データセットには 2018 年以降のデータが含まれています。

この研究では、テンソル X(k∆t) を使用して変数を表します。ここで、k は時間インデックス、∆t はトレーニング データセット内の連続するスナップショット間の時間間隔です。この研究では、ERA5 データセットを真の値と見なし、X_true(k∆t) を使用して真の値変数を表します。∆t は 6 時間に固定されています。トレーニングは、事前トレーニングと微調整の 2 つの段階で構成されます。

• 事前トレーニング段階では、トレーニングデータセットを使用して、AFNOモデルを教師あり方式でトレーニングし、X(k)からX(k + 1)へのマッピングを学習します。
• 微調整段階では、以前にトレーニングされたモデルから始めてモデルを最適化しました。つまり、モデルは最初に入力X(k)から出力X(k + 1)を生成し、次にモデルは出力X(k + 1)を入力として使用して出力X(k + 2)を生成しました。トレーニング損失は、X(k + 1)とX(k + 2)をトレーニングデータ内のそれぞれの真の値と比較することによって計算され、これら2つのトレーニング損失の合計がモデルの最適化に使用されます。

トレーニング プロセス全体は 64 個の Nvidia A100 GPU のクラスターで完了し、エンドツーエンドのトレーニングには約 16 時間かかります。

降水モデル

ERA5 再解析データセットの総降水量 (TP) は、雨や雪を通じて地球の表面に降り注ぐ液体および凍結した水の蓄積を表す変数です。この研究では総降水量を診断変数として使用し、それをp(kΔt)として表します。バックボーン モデルのトレーニングに使用されるデータセット (20 個の変数) には、変数総降水量が含まれていません。代わりに、研究者らは別の AFNO モデルをトレーニングし、バックボーン モデルの出力を使用して TP を診断しました (図 2(c) を参照)。このアプローチにより、降水モデリングの難しさと、大気の状態を予測するという一般的なタスクが切り離されます。さらに、一度トレーニングされた診断 TP モデルは、他の予測モデル (従来の NWP またはデータ駆動型予測) と組み合わせて使用​​することもできます。

バックボーン出力から降水量を診断するために使用されるモデルは、基本的な AFNO アーキテクチャと同じですが、追加の 2D 畳み込み層と ReLU アクティベーションが追加され、非負の降水量出力が強制されます。バックボーンモデルは 6 時間単位で予測を行うため、この研究では、6 時間で蓄積される総降水量を予測する診断降水モデルをトレーニングします。

実験結果

下の図 1 は、FourCastNet モデルが 0.25° 緯度経度の解像度で全球表面風速を予測する能力を定性的に示しています。このうち、地表風速の大きさは、風速の緯度成分と経度成分を用いて計算され、

。
下の図 1 は、FourCastNet を使用して生成された 96 時間先の世界の地表付近の風速予測を示しています。研究者らは、スーパー台風マンクット（2018年）や米国東海岸に向かう3つの命名されたサイクロン（フローレンス、アイザック、ヘレーネ）など、解像され正確に追跡された予報の主要な高解像度の詳細を強調しています。

このうち、図1(a)はモデルが初期化されたときの風速を示しています。図1(b)は、96時間先のモデル予測（上）と、その時点での対応する実際の風速（下）を示しています。ご覧のとおり、FourCastNet モデルは、優れた忠実度と正確な微細スケールの特徴により、96 時間先の風速を予測できます。

注目すべきことに、図 1 は、北緯約 10 度、西経 210 度の初期位置で発生し始めた、マンクットという名の超大型台風の予測される形成と進路を示しています。このモデルは、台風の激化と4日間の進路を驚くほど正確に定性的に予測しました。これは、FourCastNet モデルが台風などの気象現象の形成と移動軌道をうまく予測できることを示しています。

この研究では、サンプル外テストデータセットの初期条件を使用して FourCastNet モデルを初期化しました。この初期条件に基づいて、モデルは推論モードで16時間ステップにわたって自由に実行できます（下の図2（d）を参照）。

台風に比べると降水は断続的かつランダムであるため、降水予測は極めて難しいとされています。下の図 3 は、FourCastNet モデルが世界の総降水量を予測する能力を示しています。図3(a)は初期の降水状況を示し、図3(b)は36時間前のモデル予測と実際の状況の比較を示しています。この図は、米国西海岸の地域における降水量を例としており、FourCastNet モデルが極めて高い精度で地域的な降水量を解析し予測する能力を強調しています。研究者らは、FourCastNet が短期間の高解像度の降水特性を捉えるのに優れた性能を発揮し、それが異常気象の予測に非常に重要であることを観察しました。注目すべきは、ディープラーニング モデルが大規模な降水量予測にうまく使用されたのは今回が初めてであるということです。

次に、研究者らはハリケーンを予測するためのディープラーニングモデルの開発の潜在的な有用性をさらに調査しました。ハリケーンは、大きな破壊力を持つ異常気象です。計算コストが低く、迅速に利用できる大気警報モデルを構築することは、人命や財産の損失を軽減するために非常に重要ですが、そのようなモデルには、誤った予測を避けるためにより厳しい要件も必要です。下の図 4 は、FourCastNet が高速で移動するハリケーンを予測する強力な能力を示しています。

FourCastNet モデルは、台風や降水量に加えて、大気河川の形成と進化も予測できます。下の図 5 は、2018 年 4 月に北カリフォルニアの海岸沿いに最終的に上陸した大気河川の予測を示しています。

さらに、FourCastNet は定量的手法においても IFS よりも優れた大きな利点を持っています。 ACC と RMSE の両方の指標において、FourCastNet は IFS よりも短い時間 (約 48 時間以下) で優れたパフォーマンスを実現します。さらに、FourCastNet は IFS モデルの完全な変数セットの一部のみを使用するため、計算コストが大幅に低くなります。

アンサンブルは数値天気予報の重要な要素であり、大規模なアンサンブルを使用した FourCastNet のパフォーマンス向上は下の図に示されています。

興味のある読者は、研究の詳細について原著論文を読むことができます。