Google Deepmindの気象予報AI「GraphCast」

2024年10月11日 物理気候ゼミ Google Deepmindの気象予報AI：GraphCast 松田 拓巳 京都大学理学部 物理気候学研究室 B4 0

アジェンダ ❶ GraphCastとは何か？ ❷ GraphCastのモデル詳細 ❸ 精度比較：GraphCast vs. 数値予報モデル 1

GraphCastとは何か？ GraphCastは評価項目1246個/1380で数値予報モデルを上回る、データ駆動型の全球予報モデル 評価項目 ( 5変数 13高度 + 4地上変数 ) (u, v, t, q) (1000〜50hPa) (2t, 10u, 10v, msl) (東西風速) (南北風速) 20step = 1380個 (12h〜10days) (ジオポテンシャル) (気温) このうち90.3％で 数値予報モデルHRES よりも良い精度を出した (比湿) (地上気温) (地上東西風速) (地上南北風速) (平均海面気圧) 学習データ ERA5 (0.25 解像度) の約40年分のデータ（1979-01-01〜2022-01-10） 学習時 32GBのCloud TPU v4を32台並列で使用：約4週間 推論時 32GBのCloud TPU v4を1台使用：1分未満 (10日先までの予報) 計算コスト 2

余談 クラウドで計算すると、どれくらいお金がかかるのだろうか…？ Google Cloud の料金 (論文記載のTPUv4) • 1回の学習で どれくらい かかる？ 1回の推論で どれくらい かかる？ 割引価格の「3年間のコミットメント」の場合でも… $1.4490 • 24時間 28日 32台 150円/$ = 約470万円😱 32台 150円/$ = 約1,040万円😱😱 オンデマンド価格の場合は… $3.2200 24時間 28日 • さらにこれが実験回数分だけかかってくる…😱😱😱 • TPU1台を1分間動かせば良いので… $1.4490(〜$3.2200) • 1/60時間 1台 150円/$ = 約3.62円〜8.05円☺ 短時間＆低コストなので大規模なアンサンブル回せそう！ 3

GraphCastとは何か？ 現在＆6時間前の大気場を与えると、6時間後の未来の大気場を予測するモデル GraphCastの予報の枠組み • 入力：現在時刻と6時間前の大気状態 • 出力：6時間後の大気状態 • これを再帰的に繰り返すことで何日 も先の予報ができる →Δt=6hの時間積分をAIでやるイメージ なぜ入力は2時刻分？ * Supplementary materialsの3.1参照 • 1時刻分の場合よりも精度が良かった • 3時刻分にすると精度は確かに上がるが、メモリ使用量増加に見合った精度向上ではなかった 4

GraphCastとは何か？ " "#$ " "%$ 1グリッドが持つ特徴量は「大気物理量𝒙"#$ , 𝒙 ，外部変数𝒇 , 𝒇 , 𝒇 ! ! ! ! ! ，定数𝒄! 」の474次元ベクトル 大気物理量 (from ERA5) 𝒙"! , 𝒙"#$ ! 外部変数 " "%$ 𝒇"#$ ! , 𝒇! , 𝒇! (解析的に計算できる) 定数 𝒄! • 時刻𝑡、グリッド𝑖における1hPa〜1000hPaの37等圧面での変数 → 気温・東西風速・南北風速・鉛直風速・ジオポテンシャル・比湿 • 時刻𝑡、グリッド𝑖における地上変数 → 気温・東西風速・南北風速・平均海面気圧・積算降水量(6h) (6 37+5) 2 = 454次元 • 時刻𝑡、グリッド𝑖における大気上端の太陽放射量 • 現地時間をsin, cosにして0〜1に押し込めた値 → $ & sin &'" &( +1 , $ & cos &'" &( +1 5 3 = 15次元 • 1年の進行度をsin, cosで表して0〜1に押し込めた値 → 春分の日を0 、夏至を90 、として…？(論文には詳細なかったので推測) • グリッド𝑖が海か陸か • グリッド𝑖の標高 • グリッド𝑖の緯度のcos, 経度のsin, cos 5次元 5

アジェンダ ❶ GraphCastとは何か？ ❷ GraphCastのモデル詳細 ❸ 精度比較：GraphCast vs. 数値予報モデル 6

GraphCastのモデル詳細 GraphCastでは、グラフニューラルネットワークを採用している グラフとは？ ノード(頂点)とエッジ(辺)で構成されるデータ構造 GraphCastで使われているグラフ→ ノード エッジ グラフのデータ構造 グラフニューラルネットワーク(GNN) ノードとエッジがそれぞれベクトルを保持している ノードやエッジのベクトルを、ニューラルネット ワークで更新していく 1 2 3 𝒆*→& ∈ ℝ) 𝒗$ ∈ ℝ) ! * GraphCastのエッジは双方向的(bidirectional)で、どちらの向きでも同じ値 ※図はテキトー… 7

GraphCastのモデル詳細 メッシュは、𝑀! 〜 𝑀" の多面体から構成される「multi-mesh」という構造を使用 GraphCastのグラフ構造 「multi-mesh」＋「緯度経度格子の各グリッドノード」の組み合わせ* 以降”メッシュ”や”mesh”と呼ぶ 以降”グリッド”や”grid”と呼ぶ * グリッドノードとmulti-meshのノードもエッジで繋がっている multi-meshとは • 下図の7つの多面体𝑀, 〜 𝑀- を組み合わせたグラフ。「𝑀! のノード集合 ⊂ 𝑀 . のノード集合」の関係。(𝑖 < 𝑗) • この構造によってさまざまな空間スケールの現象を表現できるようになっている？ 正20面体 𝑀!%$ ：𝑀! の各面を4つの三角形で分割した多面体 8

GraphCastのモデル詳細 グリッドの情報をメッシュにのせる→メッシュ上で情報伝播→グリッドに情報を戻す グリッドの情報をメッシュにのせる メッシュ上で情報伝播 グリッドに情報を戻す meshノードから一定距離内にある gridノードから情報を集約し、 meshに情報をのせる mesh上でエッジに沿った情報伝播が 行われ、ノードやエッジが持っている ベクトルが更新されていく gridノードから一定距離内にある meshノードから情報を集約し、 gridに情報を戻す 9

GraphCastのモデル詳細 【前提知識】MLP (Multi Layer Perceptron)は典型的なニューラルネットワークの一つ 単層パーセプトロンとは アフィン変換と非線形関数(=活性化関数という)を組み合わせたモデル。 図示すると左のようになる。 𝐘 = 𝑆𝐿𝑃 𝑿 ≡ 𝜎(𝑾𝑿 + 𝒃) 𝑥8 ℎ8 𝑦8 ℎ7 𝑦7 𝑥7 𝑥6 MLP=多層パーセプトロンとは 単層パーセプトロンを何層も重ねたモデル。 𝐘 = 𝑀𝐿𝑃 𝑿 ≡ 𝜎 … 𝜎 𝑾" 𝜎 𝑾# 𝑿 + 𝒃# + 𝒃# … ちなみに：理論的には、3層MLPで任意の関数を任意の精度で近似できる（万能近似定理） 10

GraphCastのモデル詳細 Encoder：変数をMLPで変換し、エッジに沿って情報を集約＆更新 グリッドの情報をメッシュにのせる meshノードから一定距離内にある gridノードから情報を集約し、 meshに情報をのせる ❷ 初期特徴量をMLPで変換する ❶ ノードとエッジを初期化する Meshノード 緯度経度を表す 3次元ベクトル* を持っている Meshノード ",$%&' 𝒗" ) ! = MLP(𝒗! 𝒆*+" = MLP(𝒆",$%&' ) ( "→( $ ( "→( $ 長さや向きを表す 4次元ベクトル** を持っている ! P.4で説明した 474次元ベクトル を持っている # ! # 𝒗*! = MLP(𝒗*,$%&' ) ! Gridノード Gridノード Gridノード ❸ ノードとエッジの値を1回だけ更新する " 𝒗" ! ← 𝒗! + MLP 𝒗" ! , , (#$,(%$ *+" 𝒆*+" ← 𝒆 + MLP " $ ( →( ( "→( $ ! # ! # Meshノード 𝒆*+" " (! →(#$ 𝒆*+" , 𝒗* , 𝒗" / ( "→( $ . ! # : = 𝒆#$% ! " →! $ ! & # 𝒗*! ← 𝒗*! + MLP(𝒗*! ) * この3次元ベクトルは、具体的には：[cos(緯度), sin(経度), cos(経度)] ** この4次元ベクトルは、具体的には：[エッジ長, エッジの方向ベクトル(3D)] Gridノード 11

GraphCastのモデル詳細
Processor：メッシュ上でエッジに沿って情報を集約＆更新
メッシュ上で情報伝播
mesh上でエッジに沿った情報伝播が
行われ、ノードやエッジが持っている
ベクトルが更新されていく

❶ エッジを初期化する。

❷ 初期特徴量をMLPで変換する

Encoderで
処理済みの値を
持っているので
そのまま

",$%&'
𝒆"
( $→( $ = MLP 𝒆( $→( $
!

#

!

#

長さや向きを表す
4次元ベクトル*
を持っている

❸下記のようなノードとエッジの更新を16回繰り返す

"
𝒗"
! , , 𝒆( $→( $

"
𝒗"
! ← 𝒗! + MLP

!

(#$,(%$

"
𝒆"
( $→( $ ← 𝒆( $→( $ + MLP
!

#

!

#

-

#

" "
𝒆"
( $→( $ , 𝒗. , 𝒗/
!

#

: = 𝒆"
( $→( $
!

* この4次元ベクトルは、具体的には：[エッジ長, エッジの方向ベクトル(3D)]

-

#

12

GraphCastのモデル詳細 Decoder：各グリッドセルに近いメッシュノードから情報を集約し、グリッドに値を戻す グリッドに情報を戻す ❶ ノードとエッジの更新を1回だけ行う。 Meshノード gridノードから一定距離内にある meshノードから情報を集約し、 gridに情報を戻す 𝒆"+* ← 𝒆"+* + MLP ( $→( " ( $→( " ! # ! # 𝒆"+* , 𝒗" , 𝒗*/ ( $→( " . ! # : = 𝒆"+* ( $→( " ! 𝒗*! ← 𝒗*! + MLP 𝒗*! , , 𝒆"+* ( $→( " (#",(%" ! - # - # Gridノード ❷ グリッドノードの値をMLPに通すと、現在時刻との差分が得られる。 𝑋 9:8 = 𝑋 9 + 𝑌 9 → モデル自体は残差を予測している 𝒚*! = MLP(𝒗*! ) Gridノード 13

GraphCastのモデル詳細 Lossは重みつきMSEを使用。最大12ステップ先までの予測のLoss全体を使って誤差逆伝播 損失関数 通常の二乗和誤差 backpropagation 通時的 初期値 T=0 GraphCast GraphCast (6h後を予報) (6h後を予報) …… GraphCast (6h後を予報) 誤差逆伝播 (BPTT) 予報値 T=6h ※「GraphCast」はすべて同一のモデル 予報値 T=12h 予報値 T=72h loss 14

アジェンダ ❶ GraphCastとは何か？ ❷ GraphCastのモデル詳細 ❸ 精度比較：GraphCast vs. 数値予報モデル 15

精度比較：GraphCast vs. 数値予報モデル ベンチマーク「WeatherBench」「ECMWF Scores Cards」に概ね則って評価 対象期間 2018年1月1日〜12月31日 評価項目 ( 5変数 13高度 + 4地上変数 ) (u, v, t, q) (1000〜50hPa) (2t, 10u, 10v, msl) 20step = 1380個 (12h〜10days) GraphCastは総降水量も予報するが、評価項目には含まれていな い。論文中では「ERA5の降水量データには既知のバイアスがあるた め、評価の際には総降水量は除外されている」と言い訳している。 主な 評価指標 • RMSE (Root Mean Square Error)：二乗平均平方根誤差 • ACC (Anomaly Correlation Coefficient)：実況値の偏差と予測値の偏差の相関係数 • Normalized Difference of RMSE (ACC)：2つのモデルのRMSE(ACC)の相対差 16

精度比較：GraphCast vs. 数値予報モデル GraphCastの予報精度は数値予報モデルHRESを上回る◎ ※縦破線の左右でHRES(数値予報モデル)の初期値時刻が異なっている（左側：06z/18z，右側：00z/12z）。HRESが異なる初期時刻で異なる予報期間を提供しているため。 ・初期時刻06z/18z：予報期間は3.75日まで ・初期時刻00z/12z：予報期間は10日まで 17

精度比較：GraphCast vs. 数値予報モデル GraphCastの予報は、成層圏が苦手… 図の見方 X軸：予報時間，Y軸：高度 🟦：GraphCastの方がRMSEが低い (東西風速) 🟥：HRESの方がRMSEが低い (南北風速) (ジオポテンシャル) (気温) (比湿) (地上気温) (地上東西風速) (地上南北風速) (平均海面気圧) 考えられる原因 • (see 7.2.2) 損失関数の重みつけ 上層大気ほど、損失関数の重みを小さく設定しているため、学習が十分に行えていない？ • ERA5とHRESのデータの違い ERA5とHRESのデータセット間で成層圏における変数の予測可能性の違いが影響？ 論文では結局 「よくわからない」 と言っている 18

精度比較：GraphCast vs. 数値予報モデル 低緯度帯の予報精度はGraphCastが優位◎、成層圏では数値予報モデルに負ける 図の見方 X軸：緯度，Y軸：高度，黒線：対流圏界面 🟦：GraphCastの方がRMSEが低い 🟥：HRESの方がRMSEが低い • 低緯度帯で数値予報モデルよりも高精度 • 数値予報モデルに負けている部分はほぼ成層圏 に対応していそう 19

精度比較：GraphCast vs. 数値予報モデル GraphCastの予報は、標高が高い地域が苦手… 図の見方 X軸：地表のジオポテンシャル高度，Y軸：HRESとのRMSEの相対差分(負の値→GraphCastが◎)，色の濃さ：予報時間の違い 0より小さいなら GraphCastが良い 20

精度比較：GraphCast vs. 数値予報モデル 台風の進路予測精度は、GraphCastは数値予報モデルよりも少し良い◯（勢力には言及無し） 進路からのズレ平均 同じ台風ペアにおけるズレの差 HRESの進路 実際の進路 GraphCastの進路 21

精度比較：GraphCast vs. 数値予報モデル 大気の川の予測精度は、GraphCastは数値予報モデルよりも良い◎ 大気の川の予測精度をどう評価するか？ 予報したq, u, vの値を使って、右式で定義される 大気の川の強さ(IVT)の実況とのズレを評価 22

精度比較：GraphCast vs. 数値予報モデル GraphCastの異常気温予測精度は、短期だと数値予報モデルに劣るが、中期ではやや優位 “異常気温”の定義 上位2%＝異常高温 下位2%＝異常低温 気温 図の見方 X軸：再現率，Y軸：適合率 曲線の下の面積が大きいほど異常気温を予報できている Precision Area Under the PR Curve Recall 23

精度比較：GraphCast vs. 数値予報モデル 小スケール現象の短期予報の場合を除いて、GraphCastはHRESよりも好成績だった • 予報期間が2day〜 ：スケール問わずGraphCastはHRESよりも良い！ • 予報期間が12hour〜1day：200〜2000kmのスケールではHRESの方が良い 24

まとめ GraphCastはグラフニューラルネットワークベースのモデルで、HRESを上回る性能を持つ GraphCastとは何か？ • 40年分のERA5データを学習しており、1分以内で10日先までの全球予報が高解像度(0.25 )で可能 GraphCastのモデル詳細 • 多様な空間スケールを表現可能にするために「multi-mesh」と呼ばれるグラフを使ったGNNである 精度比較：GraphCast vs. 数値予報モデル • 90.3%の評価変数で、数値予報モデルHRESよりも高精度を達成した • 成層圏大気や、標高の高い地点の予報は苦手である • 台風や大気の川、異常気温などの極端現象も予報することができる 25