기술 요약: 효율적인 기상 예측을 위한 희소 잔차 라우팅 (Sparse Residual Routing)

문제 정의

데이터 기반 기상 예보(Pangu-Weather, ClimaX, FuXi 등)는 전통적인 수치 예보 모델(NWP)의 정확도에 필적하거나 이를 능가하면서도 추론 비용은 훨씬 적게 소모하며 큰 성공을 거두었습니다. 그러나 이러한 모델의 학습은 여전히 주요한 병목 구간입니다. 대부분의 최첨단 모델은 비전 트랜스포머(ViT)를 활용하는데, 이는 모든 공간 토큰에 대해 균일한 셀프 어텐션 연산을 적용합니다. 이러한 방식은 비효율적입니다. 왜냐하면 대기장은 높은 공간적 자기상관성을 보이며, 넓은 영역이 매끄럽고 완만하게 변화하는 조건을 포함하고 있어 중간 단계의 토큰 표현이 중복되는 경우가 많기 때문입니다.

토큰 프루닝(pruning), 병합(merging), 또는 무작위 선택(예: 이미지 생성을 위한 Sprint)과 같은 기존의 토큰 효율화 방법들은 기상 예보에 적합하지 않습니다. 기상 예측은 모든 격자 셀에 대해 물리적으로 의미 있는 출력을 요구하는 시공간적 밀집 예측(dense prediction) 문제이며, 특히 자기회귀적 롤아웃(autoregressive rollout) 과정에서 원본 위도-경도 격자 구조를 유지해야 합니다. 토큰을 제거하거나, 융합하거나, 학습된 마스크 임베딩으로 대체하는 표준 방식들은 기상 예측에 필요한 고정된 격자 구조와 시간적 조건부(temporal conditioning)를 저해합니다.

방법론: Sparse-Reslim

저자들은 기상 예측의 밀집 격자 요구 사항을 보존하면서 희소 토큰 처리를 가능하게 하는 파라미터 프리(parameter-free) 플러그인 라우팅 모듈인 Sparse-Reslim을 제안합니다.

핵심 메커니즘

Sparse-Reslim은 트랜스포머 블록을 세 단계로 분할합니다:

1. 밀집 초기 단계 (Dense Early Stage): 첫 $n_1$개 블록이 $L$개의 모든 공간 토큰을 처리하여 풍부한 초기 표현을 구축합니다.
2. 희소 중간 단계 (Sparse Middle Stage): 중간 $n_2$개 블록은 무작위로 선택된 $K = \lfloor rL \rfloor$개의 토큰 부분 집합만을 처리합니다 (여기서 $r$은 유지 비율로, 통상 0.25).
3. 밀집 후기 단계 (Dense Late Stage): 마지막 $n_3$개 블록은 예측 헤드 전의 공간적 일관성을 정교화하기 위해 전체 시퀀스 $L$을 처리합니다.

잔차 델타 공식화 (Residual Delta Formulation)

핵심 혁신은 희소 단계가 전체 시퀀스와 상호작용하는 방식에 있습니다. 선택되지 않은 토큰을 플레이스홀더나 마스크 임베딩으로 교체하는 대신, Sparse-Reslim은 **잔차 델타(residual delta)**를 계산합니다:

1. 선택된 토큰들이 희소 중간 블록을 통과합니다.
2. 이 블록들의 출력값을 입력값(라우팅 전의 표현)과 비교하여 변화량($\Delta$)을 계산합니다.
3. 이 $\Delta$를 원래의 토큰 위치로 다시 분산(scatter)시킵니다.
4. 선택되지 않은 토큰들은 제로 델타를 받게 되며, 이는 암시적인 항등 스킵 연결(identity skip connection)을 통해 라우팅 전의 표현을 정확히 유지하게 합니다.

이 접근 방식은 어떤 격자 셀도 탈락하거나 합성 토큰으로 대체되지 않도록 보장하여, 네트워크 전체에서 대기 상태 표현의 물리적 무결성을 유지합니다.

아키텍처 통합

• 결정론적 모델 (Res-Slim-ViT): 이 모듈은 표준 ViT 블록 루프에 직접 통합됩니다. 희소 중간 블록은 $O(L^2D)$ 대신 감소된 셀프 어텐션 복잡도 $O(K^2D)$로 동작합니다.
• 생성 모델 (EDM/Diffusion): 저자들은 **비대칭 어텐션 설계(asymmetric attention design)**를 채택했습니다. 타겟 토큰에 대한 셀프 어텐션은 희소화($O(K^2D)$)되지만, 크로스 어텐션은 전체 컨디셔닝 시퀀스($Z_c$)를 키(key)와 값(value)으로 유지합니다($O(KL \cdot D)$). 이를 통해 타겟 시퀀스가 희소화되더라도 디노이징에 필요한 전체 대기 맥락(과거 프레임)을 모델이 유지할 수 있도록 합니다.

토큰 선택 전략

이 방법은 학습된 선택이나 어텐션 점수 기반의 선택 대신 무작위 토큰 선택을 사용합니다. 저자들은 이것이 세 가지 이점을 제공한다고 주장합니다:

1. 점수 산출을 위한 오버헤드가 없습니다.
2. 확률적 정규화제(stochastic regularizer) 역할을 합니다 (토큰 차원에서의 드롭아웃과 유사).
3. 역동적인 영역(예: 폭풍 전선)을 선호하여 정적인 영역(예: 열대 지역)을 과소 표현할 수 있는 체계적인 공간적 편향을 방지합니다.

주요 기여

1. Sparse-Reslim 모듈: 잔차 델타 공식을 사용하여 희소한 중간 블록에서도 밀집 격자 제약을 위반하지 않고 ViT 기반 기상 모델의 희소 토큰 처리를 가능하게 하는 파라미터 프리 라우팅 메커니즘을 제시했습니다.
2. 실증적 검증: 두 가지 아키텍처 패밀리(결정론적 Res-Slim-ViT 및 생성적 EDM)와 세 가지 ERA5 해상도(최대 운영 표준인 0.25°까지)에 걸친 포괄적인 평가를 수행했습니다.
3. 성능 향상: 희소 라우팅이 계산 비용을 줄일 뿐만 아니라 예측 정확도를 향상시킨다는 것을 입증했습니다.
4. 직교성(Orthogonality): 이 방법이 혼합 정밀도(mixed-precision) 훈련 및 해상도 다운샘플링과 호환됨을 보여주었습니다.

실험 결과

저자들은 ERA5 재분석 데이터셋을 사용하여 결정론적 Res-Slim-ViT와 확산 기반 생성 모델에 대해 Sparse-Reslim을 평가했습니다.

효율성

• 학습 속도: Sparse-Reslim은 주요 설정에서 약 2.5배 빠른 학습을 달성했으며, 0.25° 해상도에서는 3.18배의 가속에 도달했습니다.
• 메모리: 테스트된 모든 설정에서 피크 VRAM 사용량이 2.2배 이상 감소했습니다.
• 추론: 추론 시간은 약 1.8배 단축되었습니다.

정확도

일반적인 효율화 방법이 속도를 위해 정확도를 희생하는 것과 달리, Sparse-Reslim은 평가된 모든 변수(Z500, T850, T2m, U10)와 해상도에서 예측 정확도를 향상시켰습니다.

• 1.0° 해상도에서 Z500 RMSE는 868.7(밀집 모델)에서 768.3(Sparse-Reslim)으로 감소했고, ACC는 0.551에서 0.651로 증가했습니다.
• T2m, U10, T850에서도 유사한 개선이 관찰되었습니다.
• 절제 연구(Ablation Studies): 정확도 향상은 중복 계산의 감소와 확률적 정규화의 결합 덕분입니다. 결정론적 희소 라우팅이 이득의 대부분을 차지하며, 무작위 선택은 추가 파라미터 없이 일관된 이득을 제공합니다.
• 안정성: 10일간의 자기회귀적 롤아웃 동안 모델은 안정성을 유지하며, 오차가 누적되기보다 스킬(skill)이 완만하게 감소합니다.

스케일링

Sparse-Reslim의 이점은 모델 크기에 따라 확장됩니다. 더 큰 모델(예: 48.6M 파라미터)은 작은 모델에 비해 더 높은 가속(2.85배)과 더 큰 정확도 개선(RMSE 12.6% 감소)을 달가는데, 이는 더 깊은 네트워크일수록 희소 중간 블록에서 차지하는 계산 비중이 크기 때문입니다.

의의 및 주장

본 논문은 Sparse-Reslim이 결정론적 및 생성적 설정 모두에서 ViT 기반 기상 모델의 효율적인 학습을 위해 희소 토큰 라우팅을 적용한 첫 번째 사례라고 주장합니다. 이 연구의 주요 의의는 토큰 효율성과 기상 예측의 밀집 예측 요구 사항 사이의 충돌을 해결했다는 데 있습니다. 잔차 델타 공식을 사용함으로써, 다른 방법들(예: Sprint)에서 발견되는 합성 플레이스홀더 토큰의 "불일치" 문제를 피하고, 모든 격자 지점에서 물리적으로 의미 있는 표현을 유지할 수 있게 합니다.

저자들은 이 방법이 매우 효과적이지만, 향arian 무작위 선택 대신 기상학적으로 활발한 영역(예: 폭풍 전선)에 계산을 동적으로 할당하기 위해 **적응형 토큰 선택(adaptive token selection)**을 탐구하는 것이 향후 과제라고 겸허히 언급했습니다. 그러나 현재의 연구는 무작위 선택이 잔차 라우팅 아키텍처와 결합되었을 때 상당한 가속과 정확도 향상을 모두 달성하기에 충분하다는 것을 입증했습니다.