[논문 리뷰] W-MAE: Pre-trained weather model with masked autoencoder for multi-variable weather forecasting

Xin Man, Chenghong Zhang외 3명|arXiv (Cornell University)|2023. 04. 18.

Hydrological Forecasting Using AI인용 수 7

한 줄 요약

W-MAE는 ERA5 데이터에 대한 자체 감독 마스킹 자동 인코더 사전 학습으로 공간 기상 구조를 학습한 뒤, 다변량 및 강수 예측에 대해 미세 조정하여 강건성을 향상시키고 빠른 앙상블 예측을 가능하게 한다.

ABSTRACT

Weather forecasting is a long-standing computational challenge with direct societal and economic impacts. This task involves a large amount of continuous data collection and exhibits rich spatiotemporal dependencies over long periods, making it highly suitable for deep learning models. In this paper, we apply pre-training techniques to weather forecasting and propose W-MAE, a Weather model with Masked AutoEncoder pre-training for weather forecasting. W-MAE is pre-trained in a self-supervised manner to reconstruct spatial correlations within meteorological variables. On the temporal scale, we fine-tune the pre-trained W-MAE to predict the future states of meteorological variables, thereby modeling the temporal dependencies present in weather data. We conduct our experiments using the fifth-generation ECMWF Reanalysis (ERA5) data, with samples selected every six hours. Experimental results show that our W-MAE framework offers three key benefits: 1) when predicting the future state of meteorological variables, the utilization of our pre-trained W-MAE can effectively alleviate the problem of cumulative errors in prediction, maintaining stable performance in the short-to-medium term; 2) when predicting diagnostic variables (e.g., total precipitation), our model exhibits significant performance advantages over FourCastNet; 3) Our task-agnostic pre-training schema can be easily integrated with various task-specific models. When our pre-training framework is applied to FourCastNet, it yields an average 20% performance improvement in Anomaly Correlation Coefficient (ACC).

연구 동기 및 목표

• 대량의 비레이블 기상 데이터를 활용하기 위한 기상 예측의 자체 감독 사전 학습 도입.
• 마샬드 자동 인코더 사전 학습으로 공간 의존성을 비전 트랜스포머 백본에서 모델링.
• 사전 학습된 모델을 다변량 예측 및 강수 예측에 대해 미세 조정하고 강건성 및 효율성을 평가.
• 작업 비특정 사전 학습이 하류 기상 모델을 향상시키고 빠른 앙상블 예측을 가능하게 할 수 있음을 보여준다.

제안 방법

• ERA5 샘플을 20채널 기상 이미지로 취급하고 패치로 나눈다.
• 마스크된 패치를 재구성하기 위해 높은 마스킹 비율의 MAE 유사 사전 학습을 적용한다.
• 언마스크된 패치를 처리하기 위해 ViT 인코더를 사용하고 픽셀 재구성을 위한 수정된 2D 패치 기반 디코더를 사용한다.
• autoregressively 미래 상태를 예측하여 다변수 예측에 대해 사전 학습된 모델을 미세 조정한다.
• 강수 및 기타 변수에 대해 FourCastNet와 비교 평가하고 추론 속도와 메모리를 분석한다.

실험 결과

연구 질문

• RQ1 작업 비특정 자체 감독 사전 학습이 단기에서 중기 기상 예측의 안정성과 정확성을 향상시키는가?
• RQ2 W-MAE가 다변수 및 강수 예측 작업에서 기존 모델(FourCastNet 등)을 능가할 수 있는가?
• RQ3 특히 앙상블 예측에서 사전 학습이 추론 속도 및 메모리 사용에 어떤 영향을 미치는가?
• RQ4 사전 학습 프레임워크가 다른 작업 특화 기상 모델로 쉽게 전이될 수 있는가?

주요 결과

• W-MAE는 저층 바람 및 500 hPa 지오포텐셜 높이와 같은 변수에 대해 명시적 오차 누적 최적화 없이도 10일 예측에서 안정된 성능을 제공한다.
• W-MAE는 강수 예측에서 FourCastNet을 크게 능가한다.
• 사전 학습을 FourCastNet에 적용하면 강수 예측에서 약 20%의 개선과 더 빠른 수렴을 보인다.
• 개인화된 W-MAE 디코더는 계산 및 메모리 오버헤드를 줄여 에포크 당 1.2x 학습 속도와 바닐라 MAE 대비 최대 8%의 메모리 절감을 달성한다.
• 사전 학습 프레임워크는 다른 작업 특화 모델과 통합되어 다수의 예측 작업에 공통 기반으로 활용될 수 있다.