구모양 퓨리에 신경망 연산자로 지구 대기 모델링하기

Reading Time: 5 minutes

머신 러닝 기반 일기 예보는 기존의 수치 일기 예보(NWP) 모델을 보완할 수 있는 유망한 기술로 부상하고 있습니다. NVIDIA 포캐스트넷(FourCastNet)과 같은 모델은 일기 예보 생성에 걸리는 계산 시간을 몇 시간에서 단 몇 초로 단축할 수 있다는 것을 입증했으며, 이는 현재의 NWP 기반 워크플로우를 크게 개선한 것입니다.

기존 방법은 제1원칙에 따라 공식화되며 일반적으로 기본 수치 방법의 정확성을 보장하기 위해 시간 간격 제한이 필요합니다. ML 기반 접근 방식은 이러한 제한이 없으며, 균일한 메모리 액세스 패턴이 GPU에 이상적으로 적합합니다.

그러나 이러한 방법은 순전히 데이터 기반이므로 다음과 같은 의문이 들 수 있습니다:

이러한 모델을 어떻게 신뢰할 수 있을까요?
얼마나 잘 일반화할 수 있을까요?
첫 번째 원칙에서 공식화되지 않았다면 어떻게 하면 그 기술, 신뢰성 및 설명 가능성을 더욱 높일 수 있을까요?

이 포스팅에서는 구형 퓨리에 신경 연산자(SFNO), 구의 물리적 시스템, 대칭의 중요성, 구형 고조파 변환(SHT)을 사용해 SFNO를 구현하는 방법에 대해 설명합니다. 수학에 대한 자세한 내용은 ICML 논문인 구형 푸리에 신경 연산자를 참조하세요: 구형 퓨리에 신경 연산자: 구형에서 안정된 역학 학습을 참조하세요.

대칭의 중요성

*그림 1. 5개월에 걸친 SFNO의 롤아웃. SFNO와 지상 실측 데이터를 사용한 지표 풍속 예측을 서로 비교하고 있습니다*.

원칙적이고 신뢰할 수 있는 모델을 만들기 위한 잠재적 접근 방식에는 물리 법칙을 공식화하는 것과 유사한 방식으로 모델을 공식화하는 것이 포함됩니다.

물리 법칙은 일반적으로 대칭성을 고려하여 공식화됩니다:

우리는 물리학이 참조 프레임에 의존할 것으로 기대하지 않습니다.
또한 기준 프레임이 변경되더라도 기본 물리 법칙은 변하지 않을 것으로 예상합니다.

구의 물리적 시스템에서 기준 프레임의 변경은 회전을 통해 이루어집니다. 따라서 회전에 따라 등가성을 유지하는 공식을 확립하기 위해 노력합니다.

현재 ML 기반 일기예보 모델은 대기 상태를 시간에 따른 다양한 공간 위치에서 관심 있는 물리적 양을 나타내는 불연속적인 일련의 벡터로 취급합니다. 이러한 벡터는 학습 함수에 의해 업데이트되며, 학습된 함수는 현재 상태를 시퀀스의 다음 상태로 매핑합니다.

쉽게 말해, 신경망에 오늘의 날씨를 보여줄 때 다음 시간 단계의 날씨를 연속적으로 예측하도록 요청하는 것입니다. 이는 전통적인 방법을 사용하여 물리 시스템을 통합하는 것과 비슷하지만, 물리 법칙에서 역학을 도출하는 것이 아니라 순전히 데이터 기반 방식으로 역학을 학습한다는 점에서 주의해야 합니다. 이 접근 방식은 기존 방법과 달리 훨씬 더 큰 시간 단계를 가능하게 합니다.

따라서 당면한 과제는 유한 차원 벡터 공간 간의 이미지 간 매핑을 학습하는 것으로 이해할 수 있습니다.

U-Net과 같은 다양한 신경망 토폴로지가 이 작업에 적용될 수 있지만, 이러한 접근 방식은 문제의 기능적 특성을 무시합니다. 입력과 출력은 모두 함수이며, 그 진화는 편미분 방정식에 의해 제어됩니다.

U-Net과 같은 기존 머신러닝 접근 방식은 고정된 해상도로 지도를 학습하기 때문에 이를 무시합니다. 신경 연산자는 이 문제를 해결하기 위해 신경망을 일반화합니다. 유한 차원 공간 사이의 맵을 학습하는 대신, 한 함수를 다른 함수에 직접 매핑할 수 있는 연산자를 학습합니다.

따라서 퓨리에 신경 연산자(FNO)는 함수 공간 간의 맵을 학습하고 한 상태를 다음 상태로 매핑하는 PDE의 해 연산자를 근사화하기 위한 강력한 프레임워크를 제공합니다.

그러나 고전적인 FNO는 데카르트 공간에서 정의되며, 관련 대칭이 구의 대칭과 다릅니다. 실제로 기하학을 무시하고 지구가 주기적인 직사각형인 것처럼 가정하면 모델의 자동 회귀 특성으로 인해 긴 롤아웃에 누적되는 아티팩트가 발생합니다. 이러한 아티팩트는 일반적으로 극 주변에서 발생하며 모델의 고장으로 이어집니다(그림 2).

이제 구형의 FNO는 어떤 모습일지 궁금하실 것입니다.

그림 2는 적응형 푸리에 신경 연산자(AFNO)를 사용한 온도 예측을 구형 푸리에 신경 연산자(SFNO)와 비교하여 보여줍니다. 구형 기하학 및 관련 대칭을 존중하면 아티팩트를 방지하고 안정적인 롤아웃이 가능합니다.

구형 푸리에 신경 연산자

유니티는 지구의 구형 기하학적 구조를 존중하기 위해 구형 좌표로 직접 공식화된 푸리에 신경 연산자인 구형 푸리에 신경 연산자(SFNO)를 구현했습니다. 이를 위해 구에 대해 공식화된 컨볼루션 정리를 활용했습니다.

글로벌 컨볼루션은 FNO의 핵심 구성 요소입니다. FFT를 통한 계산은 컨볼루션을 퓨리에 변환에 연결하는 강력한 수학적 도구인 컨볼루션 정리를 통해 가능합니다.

마찬가지로 구에 대한 컨볼루션 정리는 구의 컨볼루션을 구에 대한 푸리에 변환의 일반화인 구형 고조파 변환(SHT)에 연결합니다.

미분 구형 고조파 변환 구현하기

SFNO를 구현하기 위해서는 미분 가능한 SHT가 필요했습니다. 이를 위해 미분 SHT를 위한 PyTorch 라이브러리인 토치 하모닉스(torch-harmonics)를 구현했습니다. 이 라이브러리는 확장 가능한 모델 병렬 처리를 위해 CPU뿐만 아니라 단일 및 다중 GPU에서 SHT 계산을 기본적으로 지원하며, 다음 명령어를 실행하여 쉽게 설치할 수 있습니다:

pip install torch-harmonics

토치 하모닉스(torch-harmonics)는 PyTorch와 원활하게 통합됩니다. 차별적인 SHT는 기존 ML 아키텍처에 모듈로 쉽게 통합할 수 있습니다. 임의 함수의 SHT를 계산하려면 다음 코드 예제를 실행합니다:

import torch
import torch_harmonics as th
 
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
# parameters
nlat = 512
nlon = 2*nlat
batch_size = 32
signal = torch.randn(batch_size, nlat, nlon)
 
# create SHT instance
sht = th.RealSHT(nlat, nlon).to(device).float()
 
# execute transform
coeffs = sht(signal)

토치 고조파(torch-harmonics)를 시작하려면 화성 고도 지도의 구형 고조파 계수 계산을 안내하는 시작하기 노트북을 사용하는 것이 좋습니다(그림 3). 이 예제에서는 SHT와 ISHT의 미분성을 모두 사용하여 계수를 계산하는 방법을 보여줍니다.

*그림 3. 토치 고조파로 계산한 화성 고도 지도의 구형 고조파 계수(왼쪽). 역구면 고조파 변환으로 계산된 재구성된 신호(오른쪽).*

머신 러닝 기반 일기 예보에 대한 시사점

우리는 유럽 중기 예보 센터(ECMWF)에서 제공하는 ERA5 데이터 세트로 SFNO를 훈련시켰습니다. 이 데이터 세트는 지난 44년간의 지구 대기 상태를 가장 잘 이해하고 있는 데이터 세트입니다. 그림 2를 보면 SFNO는 극지방에서 인공물 징후를 보이지 않으며, 최대 1년 동안 수천 개의 자동 회귀 단계를 거치면서 롤아웃이 놀라울 정도로 안정적으로 유지된다는 것을 알 수 있습니다(그림 1).

이러한 결과는 머신 러닝 기반 날씨 예측 방법의 배포를 위한 길을 열어줍니다. 이 결과는 계절별 예측이라는 성배에서 ML 기반 방법이 일기 예보와 기후 예측 사이의 간극을 메울 수 있는 열쇠를 쥐고 있음을 엿볼 수 있게 해줍니다.

1460개의 자동 회귀 단계를 포함하는 1년간의 SFNO 단일 롤아웃을 단일 NVIDIA RTX A6000에서 13분 만에 계산할 수 있습니다. 이는 기존의 수치 기상 예측 방법보다 천 배 이상 빠른 속도입니다.

이처럼 훨씬 더 빠른 예측 도구는 기존 NWP를 사용하여 하나의 시나리오를 계산하는 데 걸렸던 시간을 수천 개의 가능한 시나리오를 계산할 수 있는 문을 열어주며, 드물지만 영향력이 큰 기상이변의 위험에 대해 더 높은 신뢰도로 예측할 수 있게 해줍니다.