ClimaX: A foundation model for weather and climate

論文読み会の資料です．

(A slide for the paper-reading activity at my company, written in Japanese.)

Transcript

1. ClimaX: A foundation model for weather and climate (arXiv:2301.10343) 2023/2/15 @cocomoff 1 レポジトリも最近公開 https://github.com/microsoft/ClimaX
2. この資料は？ • ClimaXの全容（の雰囲気）を理解したいので、全体を読んだときのメモをスライドにまとめたもの • ClimaXを一言で説明するなら、ViTの上に作成した ↓タスク向けの基盤モデル 2 図:https://www.microsoft.com/en-us/research/group/autonomous-systems-group-robotics/articles/introducing-climax-the- first-foundation-model-for-weather-and-climate/
3. 余談 • Googleも気象データ系の何か出しそう？（ナツコミ↓DeepMindとかもやってたらしいが…）。 3
4. アブスト (1 out of 2) • Most state-of-the-art approaches for weather and climate modeling are based on physics informed numerical models of the atmosphere. • These approaches aim to model the non-linear dynamics and complex interactions between multiple variables, which are challenging to approximate. • Additionally, many such numerical models are computationally intensive, especially when modeling the atmospheric phenomenon at a fine-grained spatial and temporal resolution. • Recent data-driven approaches based on machine learning instead aim to directly solve a downstream forecasting or projection task by learning a data-driven functional mapping using deep neural networks. • However, these networks are trained using curated and homogeneous climate datasets for specific spatiotemporal tasks, and thus lack the generality of numerical models. 4 大気などの観測量の非線形な関係をモデル化すると気象タスクに使えそう。 だけど計算が大変（特に粒度の細かいモデル) なのでDeepでポンしたい。 Deepでポンで結構いい感じに解けるけど、Deepの学習を適当にやると汎用性がないかもしれない。
5. アブスト (2 out of 2) • We develop and demonstrate ClimaX, a flexible and generalizable deep learning model for weather and climate science that can be trained using heterogeneous datasets spanning different variables, spatio-temporal coverage, and physical groundings. • ClimaX extends the Transformer architecture with novel encoding and aggregation blocks that allow effective use of available compute while maintaining general utility. • ClimaX is pre-trained with a self-supervised learning objective on climate datasets derived from CMIP6. • The pre-trained ClimaX can then be fine-tuned to address a breadth of climate and weather tasks, including those that involve atmospheric variables and spatio-temporal scales unseen during pretraining. • Compared to existing data-driven baselines, we show that this generality in ClimaX results in superior performance on benchmarks for weather forecasting and climate projections, even when pretrained at lower resolutions and compute budgets. 5 今流行りのTransformer+事前学習で頑張ったら出来た😄
6. 論文の構成 • §1 Introduction • §2 Background and Related Work ◦ §2.1 Data sources (CMIP6, ERA5) ◦ §2.2 Tasks (→ corresponding to §4) • §3 Approach ◦ §3.1 Input representation ◦ §3.2 Model architecture (tokenization, aggregation, transformer) ◦ §3.3 Datasets (pre-training, fine-tuning, and evaluation) ◦ §3.4 Training (pre-training, fine-tuning) • §4 Experiments ◦ §4.1 Neural baselines ◦ §4.2 Forecasting (global, regional, sub-seasonal) ◦ §4.3 Climate projection ◦ §4.4 Climate model downscaling ◦ §4.5 Scaling laws analysis ◦ §4.6 Ablation studies • §5 Discussion and Future Work • Appendix A〜E (これが§1-§5と同じぐらいの分量😭) 6 §4.2と§4.3の赤字は イントロで言及されてたタスク §2はデータとタスクの紹介 なので読み飛ばし （関連研究が書いてあるから チラ見する価値あり？)
7. §2 Background and Related Work • モデルなど ◦ ESM (Earth System Models) ◦ GCM (General Circulation Models) ▪ ノーベル賞2021年の言及あり ◦ NWP (Numerical Weather Prediction) • データソース ◦ CMIP (Coupled Model Intercomparison Project) ▪ CMIP6: 49 groups, hundreds of models ◦ ERA5: ECMWFの再解析データ ▪ 40 years (1979-2018) 0.25°×0.25° resolution grid • 721 x 1440 grid points ▪ hourly intervals with different climate variables (37 altitude levels) • タスク ◦ Global / regional forecasting, sub-seasonal to seasonal (S2S) prediction ◦ Climate projection ◦ Downscaling 7 ※CMIP6:第6期結合モデル相互比較プロジェクト Pythonによる気象データサイエンス:https://www.dpac.dpri.kyoto-u.ac.jp/enomoto/pymetds/Reanalysis.html
8. §2.1 Foundation models, §2.2 Tasks 8 • 基盤モデル (Foundation models) ◦ 最近流行っているやつ (NLPのBERT/GPT/PaLM、画像映像のCLIP/Florence/BEiT) ▪ ブログによれば「大量かつ多様なデータで訓練され、多様な下流タスクに適応（ファイン チューニングなど）できるモデル」 ◦ 特徴と呼ばれているもの ▪ emergence (≒ DLの発展から、高度な特徴量を学習できた/転移学習の発展も) ▪ homogenization (≒ 手法の均質化、全部Transformerで解ける！みたいなやつ) • 気象タスクではどうなの？ → ない。関連領域 (衛星画像解析など) はされてる ◦ YL20: 衛星画像系列のクラス分類のための自己教師あり事前学習 + Transformer ◦ Con+22: SatMae (NeurIPS2022) ↑ + Transformer + 事前学習 ◦ Ree+22a: Scale-MAE ◦ MAE: Masked Auto Encoder ▪ 最近流行りの隠した (マスク)したやつを当てに行く ▪ ViT (後で出てくる) の事前学習で注目されているらしい AI開発の新たなパラダイム「基盤モデル」とは: https://blog.recruit.co.jp/data/articles/foundation_models/ ファウンデーションモデル（基盤モデル）とは何か？　～あらたなパラダイムシフトのはじまり～: https://www.macnica.co.jp/business/ai/blog/142451/ ← この辺は衛星画像向け。 特に読まなくてもいい
9. 補足: Transformer • NLP界隈で生まれたモデル ◦ 背景: 長い時系列 (文書) の処理はスケールしづらい ◦ 逐次伝搬していくブロック自体に問題がある • Transformer ◦ 逐次伝搬していくブロックをAttention+NNに 置き換えて頑張って学習する ◦ 系列の情報 (h 0 -> h 1 -> …) が消えそうなので 位置情報をデータ側に付与する (positional encoding) ▪ Attention に置き換えた分を補う 9 Transformer メタサーベイ https://www.slideshare.net/cvpaperchallenge/transformer-247407256 深層学習界の大前提Transformerの論文解説！ https://qiita.com/omiita/items/07e69aef6c156d23c538
10. 補足: Vision Transformer • NLP (単語列) のTransformerを画像に転用したい 10 画像認識の大革命。AI界で話題爆発中の「Vision Transformer」を解説！ https://qiita.com/omiita/items/0049ade809c4817670d7
11. §3 Approach • 図2の雰囲気を理解したい 11 Transformer すごい埋め込みしてくれる 最強のニューラルネット 正解データ (⊿t後) 時刻情報 位置情報 いい感じのロス いろんな変数の情報を全 部埋め込んだすごいやつ
12. §3.1 Input Representation • 基本的な考え方 ◦ 格子に区切った領域について物理量などを予測したい ◦ 様々な気象データを画像として入力し、将来を予測する ▪ 入力: V × H × W • V: 入力変数 (≒ 重力ポテンシャル、気温、CO 2 /SO 2 の観測量…) • H, W: グリッドで切った領域 ▪ 出力: V’ × H’ × W’ (実際はV’ × H × W @Appendix) ◦ (§3.2冒頭のコメント) このタスクが「image-to-image translation」 っぽいと思う人がいる かもしれない (Vチャネル→V’チャネル) → なので画像系のモデルと相性が良さそう。例は UNet/ResNet/ViT。ただし気象がより複雑そうな系なので、ViTを活用して更に工夫する。 • 考えている対象 ◦ 2パターン ▪ 5.625° (32×64 grid points; 64=360/5.625) ▪ 1.40625° (128×256 grid points) ◦ 意味 ▪ 全球 (the entire globe) or 一部 (a specific region such as North America) 12
13. §3.2.1 Variable tokenizer • 図2の左側: Variable tokenizer ◦ いろいろな観測変数をトークン化する ▪ 後段で使っているViTの世界で大事と知られているらしい • RGB画像でうまくいくのは知られているけど 気象データでうまくいくかは知られていない！ ▪ V×H×Wを(H/p)×(W/p)=h×wにパッチ化した • V×p2 の大きさのパッチ（よく分からないけど) ▪ データセットによって変数Vが異なるので必要みたいな 説明が書いてあるけどあまりよく分からない (可変だから?) 13 こんな感じの物理量? (推測) t850: temperature at 850hpa u500: zonal wind speed at 500 meters from the ground q950: specific humidity at 950 meters from the ground アニメーション (Googleのブログ): https://ai.googleblog.com/2021/12/improving-vision-transformer-efficiency.html ViTのTokenizerが良いことの解説など: 自己蒸留でCVのBERTを目指すiBOT: https://ai-scholar.tech/articles/transformer/iBOT パッチの列 埋め込みベクトル↓ (p×p → Conv → D次元)
14. §3.2.2 Variable aggregator • 図2の真ん中: Variable aggregationの雰囲気を理解したい • 考えていること ◦ V×h×wの次元が大きくなりそう ▪ attention layerで最悪 O(n2) ◦ パッチ化したやつが来るので物理的な関係性が忘れられそう • 図4 集約するイメージ図 ◦ 同じ場所のデータを集約 （Position-based agg.) ▪ このような考え方はAttentionの世界では cross-attentionと呼ばれるmoduleらしい？ ◦ Appendixに実装が書いてないので詳細は不明 ▪ 素直に読むと、位置情報 → IDにした補足情報 14 V変数の情報がまとめ られたすごいやつ 時間 (Lead Time) 位置 (Position) も全部 ベクトルにして突っ込む (back-propできる) Transformer メモリ的に重いみたいな 話はたくさんある
15. github:ClimaX/src/arch.py; network • 作成 15 PatchEmbedは VisionTransformerの 実装からそのまま来てる var -> idに埋め込む PyTorchのParameterになってる 変数埋め込みの部分 位置埋め込み
16. github:ClimaX/src/arch.py; forward • 順伝搬 16 各変数をトークン化してくっつけた 変数のVar IDをパラメータにした やつを加えた (ただの足し算) 変数集約 トランスフォーマーに突っ込んだ 位置情報のパラメータを足した 予測時間の情報も足した 集約意味不明でワロタ😄
17. §3.3 Datasets / pre-training • CMIP6の様々なデータは基盤モデルの事前学習に良さそう 😄 • データセット毎に変数（観測した物理量） が異なるので混ぜるの難しそう 😭 • 事前学習で使ったデータ ◦ 5つのdata sources: MPI-ESM、TaiESM、AWI-ESM、HAMMOZ、CMCC ▪ global projections of climate scenarios from 1850 to 2015 ▪ time delta 6 hours ◦ 解像度が異なるので、対象の粒度 (5.626°/1.40625°) でregridding (≒ 補完) した • 例えばMPI-ESMの例 (↓) 17 regriddingでぐぐったら出てきた図
18. §3.3 Datasets / fine-tuning • 基本的な考え方: 事前学習した特徴量をベースに、本当に解きたいタスクに向けて学習する • Fine-tuningに使うデータはERA5 (regriddingする) • 実験評価方法は Rasp and Thuerey 2021 に従う • 詳細は（論文では）後ほど出てくる 18 • WeatherBench [Rasp+2020] ◦ Rasp et al., WeatherBench: a benchmark data set for data-driven weather forecasting, J. of Advances in Modeling Earth Systems, 2020 ◦ （おそらく）元が0.25°のERA5データをregridして、5.625°/2.8125°/1.40625°にしている Appendix C.2 ERA5 ERA5の詳細: https://confluence.ecmwf.int/display/CKB/ERA5%3A+data+documentation
19. • 後段タスクのfine-tuning ◦ 後 (個別のタスク)で説明すると論文に書いてあった §3.4 Training • 予測タスク: ある時刻 t の気象のスナップショット X t (これが V × H × W) から、 lead time ⊿t 後 の状況 X t+⊿t を予測する ◦ 学習時は ⊿t を 6時間〜168時間 (= 1週間) でランダムに事前学習する ▪ ⊿t も情報に入っている • スカラ量 → D次元ベクトル (埋め込み) • 事前学習 (Fig. 2) ◦ ランダムに⊿t〜U[6, 168]を取る ◦ 答え X t+⊿t をサンプルし、ロス関数を評価してbackprop ◦ 誤差関数: 重み付けされた平均二乗誤差 19 地球が球体なので緯度lat(i)で 縦方向Hの重み付けするらしい
20. §4 Experiments • §4 Experiments ◦ §4.1 Neural baselines ◦ §4.2 Forecasting (global, regional, sub-seasonal) ◦ §4.3 Climate projection ◦ §4.4 Climate model downscaling ◦ §4.5 Scaling laws analysis ◦ §4.6 Ablation studies 20 • Global forecasting ◦ IFS [Wed+2015]: current gold standard in weather forecasting • 他 (baselineがないようなタスク): 普通の画像系のNN (image-to-image) を試す ◦ UNet [RFB2015, GB2022], ResNet [He+ 2016], CNNs ▪ ResNetはWeatherBenchの論文と同じアーキテクチャ §4.1 比較手法
21. §4.2.1 Global Forecasting | タスク・設定 • ⊿t 時間後を予測するタスク ◦ 入力: 48 variables: 6 atmospheric variables (7 vertical levels)、3 surface variables、3 constant fileds ◦ 出力: Z500 (geopotential at 500hPa)、T850 (temperature at 850hPa)、T2m (temperature at 2 meters from the ground)、U10 (zonal wind speed at 10 meters from the ground) ▪ Z500/T850: (中規模な)NWPモデルのスタンダードな予測対象 ▪ T2m/U10: より人間の活動にとって大事そうな量として選択 ◦ 予測するlead time ⊿t: 6 hours、{1, 3, 5, 7} days、2 weeks、1 month • Fine-tuning: 全部の予測対象と⊿tで別々にfine-tuningした（他の方法もある） 21
22. §4.2.1 Global Forecasting | 結果(1) • ERA5 (5.625°) の結果 22 画像系の手法は 超長期予想できる
23. §4.2.1 Global Forecasting | 結果(2) • ERA5 (1.40625°) の結果 23 IFSと比べて だいぶいい (10日以上先)
24. §4.2.1 Global Forecasting | 付録 • 他手法 (FCN=FourCastNet/PW=PanguWeather/GC=GraphCast)と比較 24
25. §4.2.2 Regional Forecasting • Regionalな気持ち: 全球は無理でも部分的なら当てられそう ◦ ここではNorth Americaのトークンだけが入力されるようにfine-tuning ◦ 過去にされていないタスクなので、CNN・画像系・事前学習なしのClimaX (Cli-ViT)と比較 ◦ 解像度補完のよく知られた手法を使ったもの ClimaX-pos-interp 25
26. §4.2.3 S2S prodiction • Sub-seasonal to seasonal (S2S) prediction ◦ forecasting at a time range between 2 weeks and 2 months [VR,2018] ◦ competition: https://s2s-ai-challenge.github.io/ • あまり注目されていないタスクらしい（けどタスクの意味がよく分からない） 26
27. §4.3 Climate Projection | モデル • ClimateBench [WP+2022] (climate projection taskのベンチマーク) で評価する • （タスクの？)ゴール ◦ 予測対象: annual mean global distributions of surface temperature / diurnal temperature range / precipitation (降水量) / 90th percentile of precipitation ◦ 入力: CO 2 / SO 2 / BC / CH 4 • どういう意味があるタスクか: 「気温のモデリング」や「過去から未来を予測する」のではなく、例えば 「CO 2 があるレベルだと平均気温がどうなる？」みたいなことを考える (CNのシナリオ的な解釈？) 27 ClimaXと言いつつ中身変わってる！ 事前学習で見てない(出力も)。 基盤モデルならいけるやろ感。 このタスクで新しく差し込む 事前学習したやつの流用 Fine-tuningの種別: ClimaX frozen /ClimaX 10年分の予想を貯めて ある年の予想に活かすらしい（謎）
28. §4.3 Climate Projection | 実験結果 • ClimateBench [WP+2022] の結果と比較した 28 Fine-tuningの種別: ClimaX frozen /ClimaX
29. §4.4 Climate model downscaling • Downscalingの目的: coarse gridsで計算されるので高精度化したい、バイアスを減らしたい • 評価のためにやったこと ◦ CMIP6とERA5を使ってデータセットを作った (coarse → higher resolultion 予測) ◦ MPI-ESM (CMIP6 5.626°)を入力にして、ERA5 (1.40625°)を予測するように学習 29 著者コメント • 最近のダウンスケーリング手法 [Van+17,Rod+18,Hoh+20,VKG19,LG D20] はConv-basedだけど比較してない。 • 最新の手法 YNet [LGD20] はデータが足り なくて使えない (high-resolution auxiliary dataがない）
30. §4.5 Scaling laws analysis | データ・モデル • 背景: Transformerはスケール則の議論がされている。参考文献; 言語 [Kap+20, Hof+22]、画像 [Zha+22a]、マルチモーダルタスク [Hen+20b, Hen+21, Ree+22b] ◦ ブログから引用; Scaling Lawを一言で説明するなら「Transformerの性能はたった3つの変 数のべき乗則に支配されている」というものです。 （画像のイメージ➘） • ClimaXの気象予測に関する調査(左) ◦ x軸: 事前学習のデータ量 ◦ y軸: RMSE (3-day予測タスク) ◦ 埋め込み次元 (128 to 1024) 30 OpenAIが発見したScaling Lawの秘密: https://deeplearning.hatenablog.com/entry/scaling_law データ増えると減っていくみたい なことが書いてあるけど…？ （それはそう） 小さいモデルはダメそう みたいなことも書いてある
31. §4.5 Scaling laws analysis | 解像度 • 前ページ: データサイズ・モデルサイズ大事そう • このページ: （事前学習に使う？）解像度の方 (5.625°、1.40625°) 31 高解像度なデータ (1.40625°)の方を 食べさせたほうが良い みたいなことを主張してる気がする
32. §4.6 Ablation Studies • いろんなClimaXの亜種で性能比較した 32 ClimaX-all-vars: fine-tuningするときに 1変数 or 全変数でやるか 予算限られてるなら全部 一緒に予測するタスクにしたらいい ClimaX-iter: 6時間ずつ予測するのを繰り返す or 将来全部一気に予測するか 事前学習であまりやってない(?)ので 変な使い方するとすぐに精度が死ぬ ClimaX-cont: 説明読んだけど分からなかった😭 (事前学習のときの⊿t関係の話?)
33. §5 Discussion and Future Work • 力尽きました 33