EMFormer: Efficient Multi-Scale Transformer for Accumulative Context Weather Forecasting

Motivation

• Known: 데이터 기반 기상 예측 모델은 NWP 대비 계산 효율과 오차 누적 측면에서 우수하며, multi-step finetuning을 통해 장기 예측 성능을 향상시킬 수 있다는 것이 알려져 있다.
• Gap: 기존 finetuning 기반 장기 예측 방법은 catastrophic forgetting, 오차 누적, 그리고 긴 시퀀스 finetuning에 따른 높은 학습 비용 문제를 동시에 해결하지 못했으며, multi-scale transformer는 표현력은 좋지만 학습 비용이 과도하다는 한계가 있었다.
• Why: 장기 기상 예측은 항공, 해양, 금융 등 사회경제적 의사결정에 직결되는 문제이며, 학습 비용을 낮추면서도 정확도와 시간적 일관성을 유지하는 효율적 프레임워크는 실제 운용 가능한 데이터 기반 예보 시스템 구축에 핵심적이다.
• Approach: 단일 convolution으로 multi-scale 특징을 추출하는 EMFormer 아키텍처, 과거 KV를 재사용하고 메모리를 pruning하는 accumulative context finetuning, 그리고 latitude 및 variable에 적응하는 sinusoidal weighted composite loss를 결합한 pretraining-finetuning-forecasting 파이프라인을 제안한다.

Achievement

Figure 5. Comparison of our method with 4 competitors on denormalized RMSE ↓and ACC ↑in 0.25°ERA5.

1. 연산 효율성 대폭 개선: custom CUDA kernel 기반 multi-convs layer로 순전파·역전파 속도를 기존 multi-scale 구현 대비 5.69배 가속.
2. 단기 정확도 저하 없는 장기 일관성 확보: accumulative context finetuning이 기존 finetuning 대비 6-hour RMSE를 지속적으로 감소시키면서도 장기(5-day) 예측 정확도를 개선.
3. 효율성-정확도 동시 우위: VA-MoE 대비 더 적은 GPU-day(210 vs 430)로 더 낮은 Z500 RMSE(약 280 vs 295 m2/s2) 달성.
4. 범용성 입증: 기상 예측 외 ImageNet-1K 분류, ADE20K 분할 등 비전 벤치마크에서도 경쟁력 있는 성능 확인.

How

Figure 3. Illustration of Multi-Convs Layer within EMFormer.

• Stage 1 (Pretraining): 대기 변수(13개 압력 레벨의 Z, Q, U, V, T + 지표 변수)를 패치화 및 spatial pruning 후, hierarchical pruning-recovering 구조와 EMFormer 블록 스택으로 단일 스텝 예측 학습.
• EMFormer (2.2): multi-convs layer를 도입해 여러 스케일의 convolution을 단일 convolution 연산으로 재구성, 학습·추론 양쪽에서 재파라미터화 없이 multi-scale 특징 추출 가능하도록 설계, Theorem 2.1로 이론적 근거 제시.
• Stage 2 (Accumulative Context Finetuning, 2.3): 이전 스텝의 KV pair를 현재 생성 스텝에 주입하고 memory module로 KV를 선택적으로 pruning하여 메모리 사용량을 제한하면서 장기 의존성을 강화 (Xt+2, KVt+1 = Φ(Xt+1, KVt)).
• Composite Loss (2.4): latitude-adaptive term과 variable-adaptive term을 sinusoidal weighting으로 동적으로 결합하여 학습 전 과정에서 최적화 궤적을 적응적으로 유도, Theorem 2.2로 이론적 분석 제공.
• Stage 3 (Forecasting): 전지구 결정론적 예측(Φg)과 태풍 경로 예측(Φtc)을 자기회귀적으로 수행하며, 각 스텝의 KV 캐시를 다음 스텝에 전달.

Originality

• 기존 re-parameterization 기법이 추론 단계만 가속하는 것과 달리, 학습과 추론 양쪽에서 단일 convolution으로 multi-scale 특징을 얻는 multi-convs layer를 제안한 점.
• KV 캐시 주입과 메모리 pruning을 결합한 accumulative context finetuning을 통해 장기 시계열 일관성과 메모리 사용량 제한을 동시에 달성한 설계.
• 위도에 따른 공간 왜곡과 변수별 이질성을 모두 고려하는 sinusoidal weighted composite loss를 pretraining과 finetuning 전 과정에 일관되게 적용한 최적화 목표 설계.

Limitation & Further Study

• multi-convs layer의 이론적 분석(Theorem 2.1, 2.2)이 제시되었으나 발췌된 내용만으로는 증명의 일반성 및 다양한 기상 변수/해상도에 대한 적용 한계가 명확히 검증되지 않음.
• 태풍 경로 예측 등 극한 사건 예측에 대한 정량적 비교가 abstract 및 발췌 부분에서 충분히 제시되지 않아 극단 사건 성능의 신뢰도 검증이 추가로 필요.
• 후속 연구로는 더 긴 시계열(수주~계절 단위) 예측으로의 확장 가능성과, 다양한 재분석 데이터셋 및 기후 모델 간 일반화 성능에 대한 추가 검증이 필요.

Evaluation

총평: 효율적인 multi-scale 아키텍처와 누적 컨텍스트 finetoning, 적응형 loss 설계를 결합하여 장기 기상 예측의 정확도와 학습 효율성을 모두 개선한 실용적이고 체계적인 연구로, 기상 도메인을 넘어 비전 태스크까지 일반화 가능성을 보여준 점이 인상적이다.