[논문 리뷰] Earthformer: Exploring Space-Time Transformers for Earth System Forecasting

Zhihan Gao, Xingjian Shi외 5명|arXiv (Cornell University)|2022. 07. 12.

Computational Physics and Python Applications인용 수 101

한 줄 요약

Earthformer는 Cuboid Attention 공간-시간 트랜스포머를 이용해 지구 시스템 예측을 제시하며, 로컬 큐보이드 자기-주목과 소수의 글로벌 벡터를 결합하여 강수 예보(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 방법을 통해 지구 시스템 변동성의 예측 성능을 향상시키는 것을 목표로 한다.
• 고차원 지구 데이터에 적합한 확장 가능한 공간-시간 트랜스포머 아키텍처를 개발한다.
• 전역 정보 공유를 통한 효율적인 로컬 어텐션을 가능하게 하는 일반적인 큐보이드 어텐션 메커니즘을 설계하고 평가한다.
• 거친-정밀 예측을 위한 계층적 인코더-디코더의 효과를 입증한다.
• 최적의 Earthformer 구성을 식별하기 위한 디자인 선택(패턴, 글로벌 벡터, 계층 구조)을 조사한다.

제안 방법

• 입력 텐서를 비중첩 큐보이드로 분해하고 각 큐보이드 내부에서 병렬로 self-attention을 적용하는 Cuboid Attention을 도입한다.
• 크기가 (bT, bH, bW)인 큐보이드 안에서 연산하여 주의 집중 복잡도를 대략 O(THW · bT bH bW)로 감소시킨다.
• 모든 큐보이드에 어텐션하는 글로벌 벡터 모음 G를 도입하고 별도의 글로벌 어텐션 단계에서 업데이트한다.
• 거친에서 미세한 예측을 가능하게 하려는 초기 다운샘플링/업샘플링을 갖춘 계층적 인코더-디코더를 부착한다.
• 학습된 위치 임베딩에서 직접 K단계 예측을 생성하기 위해 비자회귀(non-auto-regressive) 디코딩을 사용한다(Appendix C에서 자회귀 변형을 평가).
• 다양한 큐보이드 어텐션 패턴(예: Axial, Divided Space-Time, Video-Swin, Spatial Local-Global)을 탐색하고 글로벌 벡터를 추가하여 성능 이점을 평가한다.

실험 결과

연구 질문

• RQ1 큐보이드-어텐션이 적용된 공간-시간 트랜스포머가 지구 시스템 예측 작업에서 최첨단 성능에 필적하거나 이를 능가할 수 있는가?
• RQ2 실세계 지구 데이터에 대해 어떤 큐보이드 어텐션 구성(크기, 전략, 시프트)과 아키텍처 선택(글로벌 벡터, 계층 구조)이 최상의 예측 정확도와 효율성을 제공하는가?
• RQ3 글로벌 벡터와 계층적 설계가 합성 데이터와 실제 데이터 모두에서 일관된 향상을 제공하는가?
• RQ4 강수 예보(nowcasting)와 ENSO 예측에서 Earthformer가 정확도와 계산 비용 측면에서 기존 기준선과 어떻게 비교되는가?
• RQ5 Axial-type 큐보이드 패턴과 글로벌 벡터가 다양한 작업에서 강건하고 효율적인 선택인가?

주요 결과

• Axial 패턴과 글로벌 벡터가 합성 실험과 실제 벤치마크에서 전체적으로 가장 우수한 성능을 제공한다.
• 글로벌 벡터는 계산 오버헤드가 거의 없으면서도 일관된 성능 향상을 제공한다.
• 계층적 인코더-디코더는 FLOPS를 크게 증가시키지 않으면서 예측 정확도를 향상시킨다.
• Earthformer는 SEVIR 강수 NOWCASTING 및 ICAR-ENSO ENSO 예측 데이터셋에서 최첨단 결과를 달성한다.
• 제안된 설계를 사용할 때 Earthformer는 보고된 작업에서 UNet, ConvLSTM, PredRNN, PhyDNet, E3D-LSTM, Rainformer와 같은 기준선보다 성능이 우수하다.
• 합성 테스트에서 Axial + global 패턴을 사용하는 Earthformer는 장거리 상호작용과 동역학의 강력한 모델링을 보여준다.