효율적인 다중스케일 트랜스포머 기반 누적 컨텍스트 기상 예측

초록

본 논문은 기상 예측에서 장기 컨텍스트를 효과적으로 학습하기 위해, 단일 합성곱으로 다중 스케일 특징을 추출하는 EMFormer와 누적 컨텍스트 파인튜닝 기법을 제안한다. 또한, 위도와 변수별 가중치를 동적으로 조절하는 사인파 가중 손실을 도입해 학습 안정성을 높였다. 실험 결과, 기존 멀티스케일 모듈 대비 5.69배 빠른 연산 속도와 장기 예측 정확도 향상을 입증하였다.

상세 분석

EMFormer는 기존 멀티스케일 트랜스포머가 여러 개별 컨볼루션 브랜치를 사용해 다양한 수용 영역을 학습하는 방식을 재구성한다. 논문에서는 1×1, 3×3, 5×5 커널을 하나의 “멀티‑컨브” 레이어에 결합해 단일 합성곱 연산으로 동일한 출력값을 생성하도록 설계했으며, 이를 통해 전·후방 연산을 모두 5.69배 가속화하였다. 핵심은 커널을 제로패딩 후 중앙에 정렬해 합산하고, 역전파 단계에서는 각 스케일별 그래디언트를 독립적으로 계산하도록 커스텀 CUDA 커널을 구현한 점이다. 이 설계는 Theorem 2.1에서 제시된 함수 및 그래디언트 동등성을 수학적으로 증명한다.

누적 컨텍스트 파인튜닝은 자동회귀 예측 과정에서 이전 단계의 Key‑Value(KV) 쌍을 메모리 모듈에 저장하고, 새로운 입력과 결합해 현재 단계의 어텐션에 활용한다. 메모리 용량이 증가하면 캐시를 동적으로 프루닝하는 전략을 적용해 중요한 KV만 유지한다. 이때 스코어는 현재 어텐션 맵과 히스토리 스코어를 λ = 0.9의 가중치로 블렌딩하고, 상위 N‑2개의 토큰을 선택해 다음 단계로 전달한다. 이러한 메커니즘은 장기 예측 시 오류 누적을 억제하고, 단기 정확도 저하 없이 시간적 일관성을 확보한다.

손실 함수는 변수별 손실과 위도‑가중 손실을 사인파 형태로 스케줄링한다. 사인파 가중치는 학습 초기에 전체 손실을 고르게 반영하고, 학습이 진행될수록 특정 변수나 고위도 영역에 더 큰 비중을 부여한다. 이는 대기 변수의 물리적 이질성과 지역별 변동성을 고려한 설계로, Theorem 2.2에서 최적화 경로가 보다 안정적임을 보인다.

실험에서는 Z500(500 hPa 기압면) 예측을 기준으로 6시간 단기와 5일 중기 예측에서 기존 VA‑MoE 대비 RMSE가 5~10% 감소했으며, GPU‑day 기준 210일 대비 430일을 절감했다. 또한, EMFormer는 ImageNet‑1K와 ADE20K에서 기존 멀티스케일 모델과 동등하거나 우수한 성능을 보이며, 비전 분야에서도 일반화 가능성을 입증했다. 한계점으로는 메모리 프루닝 정책의 하이퍼파라미터 선택이 데이터셋에 민감하고, 매우 장기(수주 이상) 예측에서는 여전히 오류 축적이 관찰된다는 점을 언급한다.