[paper, NeurIPS 2023] PreDiff: 시공간 Diffusion Model 기반 기상·기후 예측모델

paper : https://openreview.net/pdf?id=Gh67ZZ6zkS

github : https://github.com/gaozhihan/PreDiff

 

Amazon Web Services (AWS), Boson AI, Hong Kong University of Science and Technology

 

해당 논문에서는 딥러닝 기반 초단기 강수 예측에서 불확실성과 물리 법칙에 부합하는 예측을 동시에 다루려고 시도하였다.

기존 기상 예측은 물리 방정식을 푸는 수치 모델에 의존했지만, 계산 비용이 크고 전문 지식이 필요하다.

또한 최근엔 레이더·위성 기반의 관측 데이터가 폭발적으로 늘어나면서, 딥러닝 기반 예측 모델이 등장해 일부 수치 모델을 능가하기도 하지만,

기존 딥러닝 모델은 불확실성을 잘 표현하지 못하거나 도메인 지식 없이 학습되어, 하나의 평균적인 미래만 예측하거나, 물리적으로 말이 안 되는 결과를 출력하는 경우가 많다.

예를 들어, deterministic한 모델은 여러 가능한 시나리오를 평균내서 예측을 흐리게 만들고, 반대로 generative한 영상 모델은 현실성 없는 강수 패턴을 만들어낼 수 있다.

 

이 한계를 극복하기 위한 새로운 방식이 바로 PreDiff이다.

PreDiff는 Diffusion model을 활용해 불확실성과 울리법칙을 반영해 미래 기후상태를 예측하는 딥러닝모델이다.

 

확산 모델(Diffusion Model)을 기반으로 하여, 높은 품질의 이미지를 생성하면서도 다양한 미래 시나리오를 샘플링할 수 있는 확률적 예측을 가능하게 한다.

 

GAN과는 달리 모드 붕괴(=비슷한 결과만 나오는 문제*) 없이 안정적이다.

 * mode collapse: 자주 나오는 패턴만 생성하고 나머지는 무시
→ diffusion은 다양한 패턴을 천천히 학습하게 되면서, mode collapse 현상이 거의 없음

 

하지만 기존 확산 모델을 그대로 기상 예측에 쓰면 물리적 제약이 부족해 이상한 결과가 나올 수 있다. (확률적 기반이므로)

일부 연구는 이를 막기 위해 별도 구조나 손실 함수를 추가했지만, 새로운 제약마다 모델을 새로 만들어야 하는 비효율이 있었다.

 

PreDiff

PreDiff는 두 단계 프레임워크로 구성된다.

• 1단계는 잠재 확산 모델(Latent Diffusion Model)*을 통해 입력의 시공간 패턴을 잠재 공간에서 학습하고, 다양한 확률적 예측을 생성한다.
• 2단계는 지식 정렬(knowledge alignment)* 메커니즘으로 예측 결과에 물리 법칙이나 도메인 지식을 반영한다.

*Latent Diffusion Model (LDM) : 고해상도 데이터를 직접 생성하는 대신, 압축된 잠재 공간(latent space)에서 샘플링하여 효율적이고 고품질의 생성을 가능하게 하는 확산 기반 모델

고해상도 이미지나 데이터를 먼저 인코더로 압축

→ 잠재 공간(latent space) z로 표현

→ 이 z에 점점 노이즈를 추가하며 학습

→ 훈련 시에는 z에서 노이즈 제거하는 과정을 학습

→ 생성할 때는 순수 노이즈 zₜ에서 시작

→ 노이즈를 점점 제거하며 z₀ 생성

→ 디코더를 통해 원래 데이터 공간으로 복원

→ 최종적으로 고해상도 이미지나 예측 결과 생성

Knowledge Alignment : 생성된 결과가 현실의 지식이나 제약(예: 물리 법칙)과 일치하도록 조정하는 과정으로, 모델 출력의 신뢰성과 타당성을 높이기 위한 기술

생성 모델이 만든 출력 결과 (예: latent z 또는 예측값)

→ 외부 지식(예: 물리 법칙, 시뮬레이션 결과)과 비교

→ 이 차이를 기반으로 정렬(loss 또는 보정) 수행

→ 모델이 현실과 일치하는 방식으로 출력 조정

→ 결과적으로 신뢰도 높은, 물리적으로 타당한 출력 생성

 

해당 논문에서는 PreDiff는 확산 모델로 확률적 예측을 생성하면서, 별도 모듈로 물리 제약까지 반영해 기존 딥러닝 예측의 한계를 넘은 모델이라고 주장한다.

• PreDiff 모델 제안: 조건부 잠재 확산 모델로, 복잡한 강수 데이터를 잠재 공간에서 학습해 다양한 미래 시나리오를 생성
• 두 단계 예측 구조: 예측과 도메인 지식 적용을 분리. 지식 정렬 메커니즘으로 이미 학습된 모델에 제약을 유연하게 추가 가능
• 우수한 성능: 합성(N-body MNIST)과 실제(SEVIR) 데이터에서 기존 최고 성능을 뛰어넘음. 예측 정확도, 영상 품질, 불확실성 표현 모두 향상됨.

PreDiff 모델의 개념 및 구조

PreDiff는 전체 구조에서 잠재 공간 기반의 인코더-디코더 예측 모듈과 지식 정렬 모듈로 구성된다. 이 모델은 LatentDiffusion을 상속하며, 확산 과정에서 노이즈를 제거하고 예측을 수행하는 백본 네트워크(backbone)로 Earthformer-UNet이라는 시공간 U-Net 구조를 사용한다.

입력 데이터는 AutoencoderKL을 통해 잠재 표현 공간으로 인코딩되고, 정렬 모듈로는 SEVIRAvgIntensityAlignment가 사용된다.

여기서 LatentDiffusion은 영상 생성 분야에서 성공적으로 쓰였던 Stable Diffusion의 LDM 개념을 시계열 예측에 적용했으며,

지구과학 데이터의 시공간 구조를 잘 반영하기 위해 Earthformer 기반의 구조가 확산 네트워크로 채택되었다.

Earthformer-UNet은 Amazon에서 제안한 Earthformer 모델을 기반으로, U-Net 스타일로 구성된 시공간 예측 네트워크이며, 지구 시스템 데이터의 복잡한 시공간 패턴을 효과적으로 포착하도록 설계되었다.

해당 논문에서는 동일한 LDM이라도 백본 네트워크를 무엇을 쓰느냐에 따라 성능 차이가 크다는 것을 강조했는데, PreDiff에 Earthformer-UNet을 사용한 것이 기존 단순 U-Net 기반 LDM(Latent Diffusion Model) 대비 성능 향상의 핵심 요인임을 실험적으로 보였다 .

 

*Earthformer-UNet

- Residual Connection 도입 : 정보 손실 방지 (샘플링 과정에서 발생할 수 있는 세부 정보 손실 방지)

- z_{t+1}과 z_cond를 concat하여 z_t를 예측하는 방식으로 반복해 최종 예측 생성

 

잠재 공간 기반 확산 모델 (Latent Diffusion Model)

first_stage_model = AutoencoderKL(...) #  인코더·디코더 정의
latent_model = CuboidTransformerUNet(...) # Earthformer-UNet 정의. 영상 + 시계열 구조를 동시에 이해
super(PreDiffSEVIRPLModule, self).__init__(...) # 전체 latent diffusion 구조를 설정

 

지식 정렬 (Knowledge Alignment)

self.use_alignment = self.alignment_type is not None
self.alignment_obj = SEVIRAvgIntensityAlignment(...) #alignment 모듈 활성화
alignment_kwargs = get_alignment_kwargs_avg_x(context_seq=context_seq, target_seq=target_seq) #정렬 기준 전달
self.sample(...)  # 추론 시 정렬 적용
# 학습 중에는 p_sample() 내부에서 조건적으로 aligned_mean이 호출됨

 

PreDiff 모델 동작 단계

1. 인코딩 단계

# cond_stage_forward = AutoencoderKL.encode()
def encode_conditioning(self, c): # c : 과거 시퀀스
    if isinstance(c, torch.Tensor):
        c = rearrange(c, f"{self.einops_layout} -> {self.einops_spatial_layout}")
    c_encoded = self.cond_stage_forward({"y": c})
    return c_encoded # z_cond = E(y)

프레임 단위 인코더 E를 통해 과거 관측 시퀀스 y (예: 과거 레이더 강수 지도 연속 프레임)를 저차원 잠재 컨텍스트 z_cond로 변환

 

2. 잠재 공간 예측 단계
미래 ground-truth 프레임 x를 인코더로 압축해 latent 표현 z를 얻고, 확산 과정 초기 상태인 z_T를 샘플링

z = self.encode_first_stage(x)  # 미래 정답 x를 잠재공간으로
z_noised = q_sample(z, t, noise=noise) # z: 미래 정답 x를 encode한 후 latent로 만든 것, t: 무작위 timestep, noise: 가우시안 노이즈
# z_noised = z_T 생성 

 

 

3. 역확산(denoising) 과정

확산 모델은 Earthformer-UNet 구조로 구현되며, 초기에는 무작위 가우시안 잡음 z_T에서 시작해 점진적으로 denoise(잡음 제거)하여 미래 잠재 상태 z_{T-1}, z_{T-2}, ..., z_0를 복원

여기서 Earthformer-UNet 구조 내의 StackCuboidSelfAttentionBlock는 입력 텐서를 시공간 큐보이드 단위로 분할하고, 각 블록에 self-attention 수행, global_vector를 함께 사용하면서 로컬 정보 + 전역 정보도 함께 처리 가능

pred = self.torch_nn_module(z_noised, t, context=z_cond) # self.torch_nn_module = CuboidTransformerUNet
def forward(self, x, t, context=None):
    temb = self.time_embed(timestep_embedding(t, self.base_units))
    ...
    x = self.first_proj(x)
    ...
    for i in range(self.num_blocks):
        x = attention_block(x, context, temb) # 각 블록에 StackCuboidSelfAttentionBlock 적용

 

 

4. 조건 기반 예측 반복

각 확산 단계(timestep)에서 이전 단계의 latent z_{t+1}와 조건 z_cond를 입력받아 현재 단계의 z_t를 예측하고, 이를 t=T부터 t=0까지 반복하여 최종적으로 z_0를 획득

# CuboidTransformerUNet이 반복 호출되며 z_T → z_0 생성
for t in reversed(range(self.num_timesteps)):
    z_t = self.p_sample(z_t_plus_1, t, cond=z_cond) 

def p_sample(...):
    model_out = self.apply_model(...)
    mean, log_var = self.p_mean_variance(...)

    if self.use_alignment and self.alignment_fn is not None:
        mean = self.aligned_mean(...)  # 지식 정렬 적용

 

 

5. 디코딩 단계

최종 얻어진 z_0는 디코더 D를 통해 다시 원래 화소 공간(pixel space)의 예측 출력 x̂ (미래 강수 지도 시퀀스)으로 복원

# first_stage_model = AutoencoderKL(...)

def decode_first_stage(self, z):
    z = 1. / self.scale_factor * z
    return self.first_stage_model.decode(z) # decode()는 latent → image 복원 수행

 

데이터 및 실험

1 N-body MNIST (합성 데이터셋)

Moving MNIST에 중력 기반의 천체물리 시뮬레이션을 접목한 복잡한 동적 시계열 데이터로 PreDiff가 복잡한 동역학과 제약까지 학습할 수 있는지 확인하기 위한 데이터셋

10프레임 관측 → 10프레임 미래 예측을 목표로 함

각 프레임 해상도는 64×64, 중력 운동 및 에너지 보존 등 정확한 물리 법칙 존재 → 모델이 물리 제약을 얼마나 잘 지키는지 평가 가능

 

 

PreDiff는 기존 모델(Earthformer)보다 정확도(MSE, MAE)와 영상 품질(SSIM, FVD)에서 모두 뛰어남을 확인

PreDiff는 숫자의 위치, 형태까지 정확하게 예측한 반면, 기존 확률 모델은 영상은 선명하지만 위치나 형태 오류 발생함.

결정론적 모델은 예측 결과가 흐릿해지는 경향 존재.

PreDiff-KA는 지식 정렬(Knowledge Alignment, KA)을 통해 에너지 보존 법칙을 만족하도록 학습.

그 결과 E-MSE는 PreDiff-KA가 0.0039로, PreDiff(0.0226)나 VideoGPT(0.0228)에 비해 현저히 낮게 나타남.

이는 물리적 제약을 잘 반영하고 있음을 의미.

다만, 화질 지표에서는 다소 손해가 있었는데, PreDiff-KA의 MSE(화질 기반)는 21.90으로, PreDiff의 9.49에 비해 크게 높게 나타남.

이는 물리 법칙의 정밀한 준수와 예측 화질 간의 균형을 고려해야 함을 보여줌.

 

*MSE (Mean Squared Error, 평균 제곱 오차): 예측값과 실제값의 차이를 제곱해 평균한 값. 오차가 클수록 크게 패널티를 주며, 낮을수록 좋음

MAE (Mean Absolute Error, 평균 절대 오차): 예측값과 실제값의 차이를 절대값으로 평균한 것. 직관적인 오차의 평균으로, 낮을수록 좋음

SSIM (Structural Similarity Index Measure, 구조적 유사도 지수): 예측 영상과 실제 영상 간의 구조적 유사도를 평가. 1에 가까울수록 시각적으로 유사함

FVD (Fréchet Video Distance, 생성된 영상의 시각적 품질 평가): 생성된 시계열 영상과 실제 영상 간의 전체적인 시각 품질 차이를 측정. 낮을수록 품질이 우수함

 

2. SEVIR (실세계 강수 예측 데이터셋)

미국 내 폭풍 이벤트 관측 시계열 (위성/레이더/VIL/번개 등 포함되어 있으며 실험에선 VIL만 사용).

PreDiff 모델의 실험 결과 정량적(강수 정확도)·정성적(영상 품질) 평가 모두에서 우수하며, 특히 폭우 예측 능력과 불확실성 표현에서 뛰어난 성능을 보임

입력: 10분 간격 7프레임 (70분 관측)

출력: 10분 간격 6프레임 (60분 예측)

해상도: 128×128 (원본은 384km 도메인, 1km 해상도)

 

PreDiff는 FVD에서 가장 낮은 수치를 기록했으며(33.05), 비교 모델인 VideoGPT는 261.6, LDM은 133.0, Earthformer는 690.7로 PreDiff보다 크게 높았음.

확률 기반 예측 정확도를 나타내는 CRPS에서도 PreDiff가 0.0246으로 가장 우수했으며 Earthformer는 0.0304, LDM은 0.0280을 기록.

정답 적중률을 평가하는 CSI에서는 평균 CSI 기준으로 PreDiff가 0.410, Earthformer가 0.442로 약간 뒤처지는 경향을 보임.

그러나 공간 해상도를 줄여 위치 오차를 보정한 Pooling CSI(16×16) 기준에서는 PreDiff가 0.6244로 월등히 높게 나왔으며, 같은 조건에서 VideoGPT는 0.5798, LDM은 0.5522로 나타남.

이를 통해 위치 오차를 무시하면 PreDiff의 강수 영역 예측 정확도가 매우 뛰어남을 확인 가능

 

*CSI(Critical Success Index) : 특정 강수 강도 기준의 예측 적중률

CRPS(Continuous Ranked Probability Score) : 예측 분포의 정확도. 낮을수록 좋음

FVD(Fréchet Video Distance) : 생성된 영상의 시각적 품질 평가. 낮을수록 좋음

 

 

PreDiff-KA는 지식 정렬(Knowledge Alignment, KA)을 통해 에너지 보존 법칙을 학습하며, 평균 강수량을 ±4σ 범위로 조정함으로써 폭우나 가뭄과 같은 다양한 기상 시나리오를 생성할 수 있음.

예를 들어, +4σ는 매우 강한 비, -4σ는 비가 거의 없는 상황을 나타냄.

FVD 평가에서도 PreDiff-KA는 34.18로, PreDiff의 33.05와 큰 차이를 보이지 않음.

이를 통해 PreDiff-KA는 현실적인 예측 품질을 유지하면서도 다양한 극한 상황을 시뮬레이션할 수 있는 능력을 갖췄다는 것을 확인 가능함

 

한계점 및 미래 연구 방향

해당 논문의 결론 부분을 확인하면 PreDiff의 한계와 앞으로의 개선 방향에 대해 다음과 같이 언급하고 있음을 확인할 수 있다.

 

1. 평가 지표 및 벤치마크의 부족 : 강수 Nowcasting 분야에는 아직 표준화된 벤치마크와 평가 지표가 부족함, 지금은 주로 CSI나 CRPS 같은 수치를 사용하지만, 이 지표들이 전문가들이 실제로 느끼는 예측의 품질과 항상 잘 맞는 것은 아님
→ 새로운 평가 체계가 필요

2. 물리 지식 통합의 추가 연구 필요: PreDiff는 지식 정렬(knowledge alignment)을 통해 일부 도메인 지식을 사용했지만, 아직 물리 법칙을 딥러닝 모델 안에 자연스럽게 녹여내는 것은 완전히 해결되지 않은 과제임. 현재는 각 물리 제약마다 별도의 네트워크나 구성 요소가 필요함.
→ 하이브리드 모델로의 발전 필요성 및 AI와 기상/기후 전문가 간 협업 및 물리 기반 학습 구조에 대한 연구가 요구됨

3. 훈련 데이터의 한계와 모델의 외삽(extrapolation)* : PreDiff는 데이터 기반 모델이므로, 학습 데이터의 양과 다양성에 크게 영향을 받음. 데이터 부족 시, 정렬 제약을 맞추려다 비현실적 예측(hallucination)이 나올 수 있음
→ 데이터 확충 또는 샘플 효율이 높은 학습 방식, 일반화 능력 강화 필요

* extrapolation problem : 모델이 훈련 데이터에서 본 적 없는 상황(범위 밖)을 예측할 때 생기는 오류나 불안정성 문제

4. 계산 효율성과 실시간 적용의 과제 : PreDiff는 기존 확산 모델보다 효율적이지만, 실제로 실시간으로 예측을 하려면 속도와 모델 크기 측면에서 더 개선이 필요함.
→ 예측 속도 향상, 모델 경량화 연구 필요

 

PreDiff의 의의 및 적용 가능성

1. 확산 모델 기반의 불확실성 예측 + 확률 분포 샘플링

PreDiff는 Latent Diffusion Model (LDM) 구조를 기반으로 하며, 단일 모델에서 다양한 미래 강수 시나리오를 생성하는 probabilistic prediction이 가능하다. 기존 deterministic 방식과 달리 분산 표현과 다중 샘플링을 통해 기상 예보의 불확실성을 자연스럽게 반영할 수 있다.

2. 지식 정렬(knowledge alignment)을 통한 물리 법칙 통합

PreDiff는 확산 과정에서 직접 물리 기반 제약(예: 에너지 보존, 평균 강수량)을 주입하지 않고, 별도의 정렬 네트워크 U_{phi}를 통해 예측 분포를 간접적으로 보정한다. 이 메커니즘은 에너지 기반 모델링(Energy-based correction)과 유사하며, 확률적 샘플링을 유지하면서 제약 조건을 만족할 수 있도록 만든다.

3. 모듈화된 두 단계 프레임워크로 도메인 확장성 확보

PreDiff는 ① Latent Diffusion + ② 정렬 네트워크로 분리된 구조를 가지며, 후속 단계에 다양한 도메인 지식을 쉽게 추가할 수 있다. 예를 들어 질량 보존(Mass conservation), 온도/염분 범위 제한, 경계 조건 삽입 등 Earth system modeling 분야에서 흔한 제약들을 정렬 네트워크에서 유연하게 통합할 수 있다.

4. 기존 앙상블 예측을 대체 가능한 통합 생성 모델

여러 deterministic 모델을 조합하는 ensemble forecasting 대신, PreDiff는 하나의 generative model로 다양한 시나리오를 생성할 수 있다. 이는 계산 효율성을 높이며, 다양한 “what-if” 상황을 빠르게 테스트할 수 있는 scenario-driven decision support 체계를 구축하는 데 적합하다.