[論文レビュー] Earthformer: Exploring Space-Time Transformers for Earth System Forecasting

Zhihan Gao, Xingjian Shi+5 more|arXiv (Cornell University)|Jul 12, 2022

Computational Physics and Python Applications被引用数 101

ひとこと要約

EarthformerはCuboid Attentionスペース-タイム transformerを提案し、地球系予測において、降水予測のnowcasting (SEVIR) と ENSO SST 予測 (ICAR-ENSO) で最先端の結果を達成する。ローカルのキューブ自己注意と少数のグローバルベクトルの組み合わせを実現する。

ABSTRACT

Conventionally, Earth system (e.g., weather and climate) forecasting relies on numerical simulation with complex physical models and are hence both expensive in computation and demanding on domain expertise. With the explosive growth of the spatiotemporal Earth observation data in the past decade, data-driven models that apply Deep Learning (DL) are demonstrating impressive potential for various Earth system forecasting tasks. The Transformer as an emerging DL architecture, despite its broad success in other domains, has limited adoption in this area. In this paper, we propose Earthformer, a space-time Transformer for Earth system forecasting. Earthformer is based on a generic, flexible and efficient space-time attention block, named Cuboid Attention. The idea is to decompose the data into cuboids and apply cuboid-level self-attention in parallel. These cuboids are further connected with a collection of global vectors. We conduct experiments on the MovingMNIST dataset and a newly proposed chaotic N-body MNIST dataset to verify the effectiveness of cuboid attention and figure out the best design of Earthformer. Experiments on two real-world benchmarks about precipitation nowcasting and El Nino/Southern Oscillation (ENSO) forecasting show Earthformer achieves state-of-the-art performance. Code is available: https://github.com/amazon-science/earth-forecasting-transformer .

研究の動機と目的

• 地理時空観測データを活用したデータ駆動型DL手法で、地球系の変動予測を改善する動機づけ。
• 高次元の地球データに適したスケーラブルな時空Transformerアーキテクチャを開発する。
• 局所的なアテンションとグローバル情報共有を可能にする汎用的なcuboid attention機構を設計・評価する。
• 粗から細への予測を可能にする階層的エンコーダ-デコーダを導入する。
• 最適なEarthformer構成を見つけるため、パターン・グローバルベクトル・階層構造などの設計選択を検討する。

提案手法

• 入力テンソルを非重畳のキューブに分解し、各キューブ内で自己注意を並列に適用するCuboid Attentionを導入する。
• (bT, bH, bW)サイズのキューブ内で動作させることで、アテンションの計算量を概ねO(THW · bT bH bW)に削減する。
• すべてのキューブに対してアテンションを行い、別個のグローバルアテンションステップで更新されるグローバルベクトルGを導入する。
• 粗から細の予測を可能にする初期のダウンサンプリング/アップサンプリングを備えた階層的エンコーダ-デコーダを付随させる。
• learned positional embeddingsから直接Kステップ予測を生成する非自回帰デコードを使用する（Appendix Cで自回帰変種を評価）。
• 複数のcuboid attentionパターン（例：Axial、Divided Space-Time、Video-Swin、Spatial Local-Global）を探索し、グローバルベクトルを追加して性能利得を評価する。

実験結果

リサーチクエスチョン

• RQ1 Cuboid-attentionを備えた時空トランスフォーマーは、地球系予測タスクで最先端の性能に匹敵するか、またはそれを超えるか。
• RQ2 実世界データに対して最適な予測精度と効率性を生むcuboid attentionの構成（サイズ・戦略・シフト）およびアーキテクチャの選択は何か。
• RQ3 グローバルベクトルと階層設計は、合成データと実データの両方で一貫した改善をもたらすか。
• RQ4 降水nowcastingとENSO予測において、Earthformerは既知のベースラインと比較して精度と計算コストの両面でどうか。
• RQ5 Axialタイプのcuboidパターンとグローバルベクトルは、タスクを跨いで頑健で効率的な選択か。

主な発見

• Axialパターンとグローバルベクトルの組み合わせが、合成実験と実世界ベンチマークの総合的な性能で最良を示した。
• グローバルベクトルは計算オーバーヘッドをほとんど増やさず、一貫した性能向上を提供する。
• 階層的エンコーダ-デコーダはFLOPSを大幅に増やすことなく予測精度を向上させる。
• EarthformerはSEVIR降水nowcastingとICAR-ENSO ENSO予測データセットで最先端の結果を達成する。
• Earthformerは提案設計を用いると、UNet、ConvLSTM、PredRNN、PhyDNet、E3D-LSTM、Rainformerなどのベースラインを上回る。
• 合成テストでは、Axial + グローバルパターンを用いたEarthformerが長距離相互作用とダイナミクスのモデリングで強力であることを示した。