[논문 리뷰] ArchesWeather: An efficient AI weather forecasting model at 1.5° resolution

Guillaume Couairon, Christian Lessig외 2명|arXiv (Cornell University)|2024. 05. 23.

Meteorological Phenomena and SimulationsEarth and Planetary Sciences인용 수 3

한 줄 요약

ArchesWeather는 1.5° 해상도에서 계산 효율성이 뛰어난 인공지능 기상 예측 모델로, 3D 국소 주의 메커니즘을 새로운 크로스레벨 주의(Cross-Level Attention, CLA) 구조로 대체하여 몇 일치 GPU 훈련 예산으로도 최신 기준 성능을 달성한다. 이 모델은 1.4° 50개 구성의 NeuralGCM 앙상블을 능가하고, Pangu-Weather 및 Stormer와 동등하거나 뛰어난 핵심 지표를 확보하면서도 표준 하드웨어에서 저비용 추론을 가능하게 한다.

ABSTRACT

One of the guiding principles for designing AI-based weather forecasting systems is to embed physical constraints as inductive priors in the neural network architecture. A popular prior is locality, where the atmospheric data is processed with local neural interactions, like 3D convolutions or 3D local attention windows as in Pangu-Weather. On the other hand, some works have shown great success in weather forecasting without this locality principle, at the cost of a much higher parameter count. In this paper, we show that the 3D local processing in Pangu-Weather is computationally sub-optimal. We design ArchesWeather, a transformer model that combines 2D attention with a column-wise attention-based feature interaction module, and demonstrate that this design improves forecasting skill. ArchesWeather is trained at 1.5° resolution and 24h lead time, with a training budget of a few GPU-days and a lower inference cost than competing methods. An ensemble of four of our models shows better RMSE scores than the IFS HRES and is competitive with the 1.4° 50-members NeuralGCM ensemble for one to three days ahead forecasting. Our code and models are publicly available at https://github.com/gcouairon/ArchesWeather.

연구 동기 및 목표

• 1.5° 해상도에서 높은 정확도를 유지하면서도 계산 효율성이 뛰어난 AI 기상 예측 모델을 설계한다.
• Pangu-Weather와 같은 모델에서 3D 국소 주의가 최적임을 전제로 하는 가정을 도전하여, 기상 예측 모델링에서 이 주의 방식이 계산적으로 열악함을 입증한다.
• 압력 수준 간 전역적 특징 상호작용을 가능하게 하는 새로운 크로스레벨 주의(Cross-Level Attention, CLA) 메커니즘을 통해 예측 정확도를 향상시킨다.
• 모델 크기나 훈련 비용을 증가시키지 않고도 훈련 비용과 추론 시간을 줄이며, 더 큰 비용이 드는 모델들과 동등하거나 뛰어난 성능을 달성한다.
• ERA5 데이터의 분포 이탈 영향을 조사하고, 최근 데이터로의 피지컬 튜닝을 통해 성능 향상을 도모한다.

제안 방법

• 모델는 Earth 전용 위치 편향을 가진 3D Swin U-Net 트랜스포머 아키텍처를 사용하며, Pangu-Weather를 영감으로 삼았지만 3D 국소 주의를 비국소 크로스레벨 주의(Cross-Level Attention, CLA) 레이어로 대체하여 압력 수준 간 특징 상호작용을 구현한다.
• CLA 레이어는 압력 수준 간 전역 주의를 가능하게 하여, 이웃 수준에만 국한되는 3D 국소 주의의 계산 비효율성을 해결한다.
• 모델은 $X_{t+ au}$가 아니라 $X_{t+ au} - X_t$의 차이를 예측하도록 훈련되어 훈련 안정성과 일반화 능력을 향상시킨다.
• 풍속 성분(U, V, W)은 입력 특징으로 포함되며, 시간 일반화를 향상시키기 위해 일자와 월을 조건으로 삼는다.
• 역사적 데이터의 분포 이탈을 완화하고 최근 대기 조건에서의 성능 향상을 위해 2007–2018년 ERA5 데이터로 피지컬 튜닝 단계를 적용한다.
• 정확도 향상을 위해 여러 개의 독립적으로 훈련된 모델(예: ArchesWeather-M4, ArchesWeather-L2)의 출력을 평균하여 앙상블 예측을 생성한다.

실험 결과

연구 질문

• RQ1 Pangu-Weather와 같은 모델에서 3D 국소 주의가 물리적 타당성은 있지만 AI 기상 예측에 있어 계산적으로 열악한가?
• RQ2 압력 수준 간 비국소 주의 메커니즘이 모델 크기나 훈련 비용을 증가시키지 않고도 예측 정확도를 향상시킬 수 있는가?
• RQ3 최근 ERA5 데이터(2007–2018)로의 피지컬 튜닝이 역사적 데이터의 분포 이탈로 인한 오차를 줄이는가?
• RQ4 1.5° 해상도에서 훈련 비용이 몇 일치 GPU에 불과한 모델이 경쟁력 있는 성능을 낼 수 있는가?
• RQ5 24시간 예측 지연 시간에서 ArchesWeather의 성능이 NeuralGCM, Pangu-Weather, Stormer와 같은 최신 기준 모델과 비교해 어떻게 되는가?

주요 결과

• ArchesWeather-M는 2개의 A100 GPU에서 약 2일간 훈련되어 IFS HRES 대비 평균 상대 RMSE 개선도(Relative RMSE Improvement, RRH)가 -5.0%를 기록하며, 1.4° 50개 구성의 NeuralGCM 앙상블을 능가한다.
• ArchesWeather-L2 앙상블(두 개의 L 모델)은 주요 고공 변수(Z500, T850, Q700, U850, V850)에서 NeuralGCM 앙상블보다 낮은 RMSE를 기록한다.
• CLA 모듈은 CLA 없이 훈련된 기준 모델 대비 Z500 RMSE를 10.5% 감소시켜 IFS HRES 성능에 더 가까워지게 한다.
• 2007–2018년 최근 ERA5 데이터로의 피지컬 튜닝은 전체 1979–2018년 기간 훈련 대비 Z500 RMSE를 1.2% 감소시키고, T2m RMSE는 0.001°C 감소시킨다.
• 단일 A100 GPU에서 ArchesWeather-M 모델의 추론은 1개의 24시간 예측에 약 0.25초로 소요되어 저비용 배포가 가능하다.
• 모델는 약 10 V100일의 훈련 비용으로도 경쟁성 있는 성능을 달성했으며, SphericalCNN 대비 40배, Stormer 대비 256배 적은 비용으로 동등하거나 뛰어난 정확도를 확보한다.