[논문 리뷰] Prithvi WxC: Foundation Model for Weather and Climate

Johannes Schmude, Sujit Roy외 27명|arXiv (Cornell University)|2024. 09. 20.

Geophysics and Gravity Measurements인용 수 7

한 줄 요약

tldr: Prithvi WxC는 2.3B-parameter 트랜스포머 기반 파운데이션 모델로, 160 MERRA-2 대기 변수에 대해 학습되어 예보, 다운스케일링, 중력파 매개화, 극한 이벤트 추정에 사용되며, Hugging Face에서 오픈 소스로 공개됩니다.

ABSTRACT

Triggered by the realization that AI emulators can rival the performance of traditional numerical weather prediction models running on HPC systems, there is now an increasing number of large AI models that address use cases such as forecasting, downscaling, or nowcasting. While the parallel developments in the AI literature focus on foundation models -- models that can be effectively tuned to address multiple, different use cases -- the developments on the weather and climate side largely focus on single-use cases with particular emphasis on mid-range forecasting. We close this gap by introducing Prithvi WxC, a 2.3 billion parameter foundation model developed using 160 variables from the Modern-Era Retrospective Analysis for Research and Applications, Version 2 (MERRA-2). Prithvi WxC employs an encoder-decoder-based architecture, incorporating concepts from various recent transformer models to effectively capture both regional and global dependencies in the input data. The model has been designed to accommodate large token counts to model weather phenomena in different topologies at fine resolutions. Furthermore, it is trained with a mixed objective that combines the paradigms of masked reconstruction with forecasting. We test the model on a set of challenging downstream tasks namely: Autoregressive rollout forecasting, Downscaling, Gravity wave flux parameterization, and Extreme events estimation. The pretrained model with 2.3 billion parameters, along with the associated fine-tuning workflows, has been publicly released as an open-source contribution via Hugging Face.

연구 동기 및 목표

• 전 세계 및 지역 작업에 걸친 기상 및 기후 분야에서 파운데이션 모델 접근의 필요성을 제시한다.
• 대규모 토큰 수와 다양한 토폴로지를 처리할 수 있는 확장 가능한 인코더-디코더 트랜스포머 아키텍처를 개발한다.
• 기후학에 기반한 예측과 함께 마스킹 재구성과를 결합한 혼합 목표로 학습한다.
• 제로샷 재구성과 예측 성능을 평가하고, 다운스케일링 및 중력파 매개화와 같은 다운스트림 작업을 평가한다.
• 커뮤니티의 광범위한 활용을 가능하게 하기 위해 사전 학습 가중치와 파인튜닝 워크플로를 공개적으로 공개한다.

제안 방법

• 윈도우를 가로지르는 로컬 및 글로벌 어텐션을 갖춘 인코더-디코더 트랜스포머로, 장거리 의존성을 가능하게 한다.
• 입력 마스킹 50%를 사용하는 2D 비전 트랜스포머와 윈도우 및 윈도우당 토큰 수가 유연한 토큰화 체계.
• 정규화된 잔차 형식화(식 2)를 포함하여 마스킹 재구성과를 기후학 기반 예측과 혼합한 사전 학습 목표.
• 정적 입력, 기후학 C_t 및 선두 시간 매개변수 δt 및 δτ와 함께 160 MERRA-2 변수(1980-2019)로 학습한다.
• 마스킹 및 임의의 선두 시간과 함께 초기 단계의 이중 단계 사전 학습, 이어 예측용 교정 손실 및 자기회귀 롤아웃으로 파인튜닝한다.
• 다운스트림 다운스케일링 및 중력파 매개화용 파인튜닝 워크플로를 제공한다.

실험 결과

연구 질문

• RQ1 하나의 파운데이션 모델이 제로샷 기능으로 전세계 및 지역 기상-기후 작업을 처리할 수 있는가?
• RQ2 가변 선두 시간에 걸친 재구성 및 예측 모두를 가장 잘 지원하는 사전 학습 목표는 무엇인가?
• RQ3 대형 WxC 모델이 자회귀 예측, 다운스케일링 및 허리케인 트랙 예측에서 태스크-스펙IFIC 모델과 비교해 어떤 성능을 보이는가?
• RQ4 기후학과 정적 특징을 통합된 사전 학습 목표로 통합하기 위한 효과적인 전략은 무엇인가?
• RQ5 다운스트림 작업(다운스케일링, 중력파 매개화, 극한 이벤트)에 대한 실용적인 파인튜닝 고려사항은 무엇인가?

주요 결과

• 제로샷 재구성은 로컬 또는 글로벌 마스킹 하에서 남은 데이터가 5-25%로도 높은 마스킹 상태에서 작동한다.
• 제로샷 예측은 짧은 선두 시간(6-12시간)에서 가장 강하게 작동하며, 약 66시간까지 합리적인 성능을 보이고 그 이후는 일부 기준선 대비 감소한다.
• 허리케인 트랙 예측에서 Prithvi WxC는 MERRA-2 또는 ERA5 데이터로 학습된 FourCastNet 변형에 비해 트랙 정확도와 상륙 시점을 우수하게 달성한다(예: Ida, 2021; 평균 트랙 오차 63.9 km 대 201.9–262.3 km).
• 다운스케일링 실험은 보간 기준선에 비해 상당한 개선을 보인다(MERRA-2 T2m: 공간 RMSE 3.22→0.73; 시간 RMSE 2.46→0.64; 상관 0.89→0.98).
• CORDEX 다운스케일링을 12.5 km 해상도로 수행하면 유사한 개선이 나타난다(공간 RMSE 1.89→0.44; 시간 RMSE 1.14→0.37; 상관 0.99→1.00).
• 모델은 Hugging Face에서 공개적으로 사전 학습 가중치 및 다운스트림 작업용 파인튜닝 워크플로와 함께 공개된다.