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
   

   View Slide

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/
   

   View Slide

3. 余談
   ● Googleも気象データ系の何か出しそう？（ナツコミ↓DeepMindとかもやってたらしいが…）。
   3
   

   View Slide

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の学習を適当にやると汎用性がないかもしれない。
   

   View Slide

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+事前学習で頑張ったら出来た😄
   

   View Slide

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はデータとタスクの紹介
   なので読み飛ばし
   （関連研究が書いてあるから
   チラ見する価値あり？)
   

   View Slide

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
   

   View Slide

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/
   ← この辺は衛星画像向け。
   特に読まなくてもいい
   

   View Slide

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
   

   View Slide

10. 補足: Vision Transformer
    ● NLP (単語列) のTransformerを画像に転用したい
    10
    画像認識の大革命。AI界で話題爆発中の「Vision Transformer」を解説！ https://qiita.com/omiita/items/0049ade809c4817670d7
    

    View Slide

11. §3 Approach
    ● 図2の雰囲気を理解したい
    11
    Transformer
    すごい埋め込みしてくれる
    最強のニューラルネット
    正解データ
    (⊿t後)
    時刻情報
    位置情報
    いい感じのロス
    いろんな変数の情報を全
    部埋め込んだすごいやつ
    

    View Slide

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
    

    View Slide

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次元)
    

    View Slide

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
    メモリ的に重いみたいな
    話はたくさんある
    

    View Slide

15. github:ClimaX/src/arch.py; network
    ● 作成
    15
    PatchEmbedは
    VisionTransformerの
    実装からそのまま来てる
    var -> idに埋め込む
    PyTorchのParameterになってる
    変数埋め込みの部分
    位置埋め込み
    

    View Slide

16. github:ClimaX/src/arch.py; forward
    ● 順伝搬
    16
    各変数をトークン化してくっつけた
    変数のVar IDをパラメータにした
    やつを加えた (ただの足し算)
    変数集約
    トランスフォーマーに突っ込んだ
    位置情報のパラメータを足した
    予測時間の情報も足した
    集約意味不明でワロタ😄
    

    View Slide

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でぐぐったら出てきた図
    

    View Slide

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
    

    View Slide

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の重み付けするらしい
    

    View Slide

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 比較手法
    

    View Slide

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
    

    View Slide

22. §4.2.1 Global Forecasting | 結果(1)
    ● ERA5 (5.625°) の結果
    22
    画像系の手法は
    超長期予想できる
    

    View Slide

23. §4.2.1 Global Forecasting | 結果(2)
    ● ERA5 (1.40625°) の結果
    23
    IFSと比べて
    だいぶいい
    (10日以上先)
    

    View Slide

24. §4.2.1 Global Forecasting | 付録
    ● 他手法 (FCN=FourCastNet/PW=PanguWeather/GC=GraphCast)と比較
    24
    

    View Slide

25. §4.2.2 Regional Forecasting
    ● Regionalな気持ち: 全球は無理でも部分的なら当てられそう
    ○ ここではNorth Americaのトークンだけが入力されるようにfine-tuning
    ○ 過去にされていないタスクなので、CNN・画像系・事前学習なしのClimaX (Cli-ViT)と比較
    ○ 解像度補完のよく知られた手法を使ったもの ClimaX-pos-interp
    25
    

    View Slide

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
    

    View Slide

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年分の予想を貯めて
    ある年の予想に活かすらしい（謎）
    

    View Slide

28. §4.3 Climate Projection | 実験結果
    ● ClimateBench [WP+2022] の結果と比較した
    28
    Fine-tuningの種別:
    ClimaX
    frozen
    /ClimaX
    

    View Slide

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がない）
    

    View Slide

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
    データ増えると減っていくみたい
    なことが書いてあるけど…？
    （それはそう）
    小さいモデルはダメそう
    みたいなことも書いてある
    

    View Slide

31. §4.5 Scaling laws analysis | 解像度
    ● 前ページ: データサイズ・モデルサイズ大事そう
    ● このページ: （事前学習に使う？）解像度の方 (5.625°、1.40625°)
    31
    高解像度なデータ (1.40625°)の方を
    食べさせたほうが良い
    みたいなことを主張してる気がする
    

    View Slide

32. §4.6 Ablation Studies
    ● いろんなClimaXの亜種で性能比較した
    32
    ClimaX-all-vars:
    fine-tuningするときに
    1変数 or 全変数でやるか
    予算限られてるなら全部
    一緒に予測するタスクにしたらいい
    ClimaX-iter:
    6時間ずつ予測するのを繰り返す
    or 将来全部一気に予測するか
    事前学習であまりやってない(?)ので
    変な使い方するとすぐに精度が死ぬ
    ClimaX-cont:
    説明読んだけど分からなかった😭
    (事前学習のときの⊿t関係の話?)
    

    View Slide

33. §5 Discussion and Future Work
    ● 力尽きました
    33
    

    View Slide