AI 기상 예측 모델의 multi-step rollout 학습 전략 — BPTT, Replay Buffer, Online Replay Buffer

Published: June 24, 2026

1. 왜 multi-step 학습이 필요한가

데이터 기반 기상·기후 모델은 대부분 autoregressive(AR) 구조다. 모델 $f_\theta$ 는 한 시점의 상태 $x_t$ 를 받아 다음 시점 $x_{t+1}$ 을 예측하고, 예측을 다시 입력으로 넣어 미래로 굴러간다.

\[\hat{x}*{t+k} = \underbrace{f*\theta \circ f_\theta \circ \cdots \circ f_\theta}_{k\ \text{회}}(x_t)\]

문제는 학습과 추론의 불일치다. 가장 단순한 학습 방식인 single-step teacher forcing은 항상 정답 상태 $x_t^{\text{true}}$ 를 입력으로 받아 one-step loss만 최소화한다.

\[\mathcal{L}*{\text{1-step}} = \big| f*\theta(x_t^{\text{true}}) - x_{t+1}^{\text{true}} \big|\]

하지만 추론 시에는 모델이 자기 자신의 (조금씩 틀린) 출력을 입력으로 받는다. 모델은 한 번도 본 적 없는 분포의 입력을 마주하게 되고, 그 위에서 다시 예측하면서 오차가 누적된다. 결과적으로 스펙트럼이 뭉개지거나(blurring), 며칠 뒤 롤아웃이 발산하는 현상이 나타난다. 이 distribution shift가 AR 모델 불안정성의 핵심이다.

Multi-step rollout 학습의 목표는 한 문장으로 요약된다.

학습 중에 모델이 자기 자신의 출력 분포를 미리 경험하게 만들어, 추론 시 분포와 학습 시 분포를 맞춘다.

이 목표를 어떻게 구현하느냐에 따라 전략이 갈린다. 핵심 축은 세 가지다.

Gradient horizon — gradient가 롤아웃을 얼마나 거슬러 올라가 보는가
On-policy vs off-policy — 모델이 보는 입력 상태가 현재 파라미터로 만든 것인가, 과거 파라미터로 만든 것인가
메모리·연산 비용 — 그래프에 몇 step을 올리고, 별도 저장공간을 얼마나 쓰는가

이제 세 전략을 이 축들 위에서 비교해 보자.

2. BPTT (Backpropagation Through Time)

아이디어

가장 직관적인 방법이다. 모델을 $K$ step 펼쳐서(unroll) 전체 롤아웃을 하나의 계산 그래프로 만들고, 각 step의 loss를 합한 뒤 전체 그래프를 거슬러 backprop한다.

x = x0
loss = 0.0
for k in range(1, K + 1):
    x = f_theta(x)                 # 그래프에 그대로 유지
    loss += criterion(x, target[k])
loss.backward()                    # K step 전체를 거슬러 미분

특징

gradient가 롤아웃 전체를 본다는 게 결정적 장점이다. 모델은 “지금 이 예측이 몇 step 뒤 loss에 어떤 영향을 주는가”까지 고려해서 파라미터를 갱신한다. 즉, 미래까지 안정적으로 굴러가는 표현을 직접 학습한다. multi-step 신호의 품질 면에서는 사실상 gold standard다.

대신 대가가 크다.

메모리가 $K$ 에 선형 비례한다. backprop을 위해 모든 중간 step의 activation을 들고 있어야 한다. 고해상도 전지구 필드에서는 $K$ 를 키우는 순간 GPU 메모리가 터진다. → gradient checkpointing으로 일부 activation을 버리고 backward에서 재계산하는 방식이 거의 필수다(메모리↓, 연산↑ 트레이드오프).
Gradient pathology. 긴 체인을 거치면 vanishing/exploding gradient가 생기고 loss 표면이 거칠어진다. 그래서 보통 작은 $K$(2~12 step 수준)에 머문다.
연산 비용도 update당 $K$ 회 forward + backward로 비싸다.

실제 사례

GraphCast는 처음부터 긴 롤아웃으로 학습하지 않는다. curriculum 방식으로, 1 step부터 시작해 롤아웃 길이를 점진적으로 늘려 가며(최종적으로 12 step 수준까지) 짧은 fine-tuning 단계를 거친다. BPTT의 메모리·안정성 한계를 curriculum과 checkpointing으로 우회한 전형적인 예다.

요약: gradient 신호는 가장 풍부하지만 메모리 $O(K)$ 와 gradient 불안정 때문에 긴 horizon으로 확장하기 어렵다.

3. Replay Buffer

아이디어

BPTT의 메모리 문제는 “전체 롤아웃을 그래프에 올린다”는 데서 온다. Replay buffer는 발상을 뒤집는다.

그래프에는 딱 한 step만 올린다. 대신, 입력 상태 자체를 모델이 이미 여러 번 굴려 만든 예측으로 채워서 distribution shift를 흉내 낸다.

이는 Brandstetter et al.(2022)의 pushforward trick과 강화학습의 replay buffer 개념을 결합한 형태다. 버퍼에는 다양한 롤아웃 깊이(depth)의 예측 상태들이 쌓여 있다. 매 iteration마다 버퍼에서 상태 하나를 뽑아 한 step만 전진시키고, 그 한 step에 대해서만 backprop한 뒤(입력에는 stop-gradient), 새로 만든 예측을 다시 버퍼에 밀어넣는다.

buffer.init_from_data()            # depth 0 상태들로 초기화

for it in range(num_iters):
    x, depth = buffer.sample()     # 과거에 만들어진 예측 상태
    x_next = f_theta(x)            # 단 1 step만 그래프에 올림
    loss = criterion(x_next, target[depth + 1])
    loss.backward()                # 1-step gradient (x는 detach 상태)

    if depth + 1 < max_depth:
        buffer.push(x_next.detach(), depth + 1)   # 한 step 더 깊은 상태로 재투입
    else:
        buffer.push(sample_fresh_data(), 0)       # 깊이 한계 도달 시 리셋

버퍼에 한 번 들어간 상태는 다음에 뽑힐 때 또 한 step 깊어지므로, 시간이 지나면 임의의 깊이까지 도달한다. 매 update의 메모리는 항상 1 step 수준으로 일정하다.

특징

Update당 메모리가 상수다. 그래프에 1 step만 올리니 $K$ 와 무관하게 깊은 롤아웃 경험을 줄 수 있다. BPTT가 못 가는 horizon을 메모리 부담 없이 커버한다.
다양한 깊이의 분포를 자연스럽게 섞어 모델에 노출시킨다.

하지만 두 가지 본질적 약점이 있다.

Off-policy / staleness. 버퍼의 상태는 과거 버전의 파라미터 $\theta_{\text{old}}$ 로 만들어진 것이다. 학습이 진행될수록 현재 모델 $\theta$ 와 버퍼 상태를 만든 모델 사이의 괴리가 커진다. 모델이 실제로 추론 시 만들 분포와, 버퍼가 담고 있는 분포가 어긋나는 것이다. RL의 off-policy 문제와 정확히 같은 구조다.
저장 비용과 부기(bookkeeping). 고해상도 필드 상태를 다수 보관해야 하므로 CPU/GPU 메모리를 따로 먹는다. 또 깊이 분포가 버퍼 동역학에 의해 암묵적으로 결정되어, “롤아웃 깊이 분포를 이렇게 하고 싶다”고 직접 지정하기 어렵다.

요약: 메모리 상수로 깊은 horizon을 커버하지만, 버퍼 상태가 stale(off-policy)하고 저장공간·깊이 제어가 까다롭다.

4. Online Replay Buffer

아이디어

Replay buffer의 가장 큰 흠은 staleness였다. 그렇다면 버퍼를 영구히 들고 가지 말고, 매번 현재 파라미터로 즉석에서 롤아웃을 만들면 되지 않을까? 이것이 online replay buffer의 발상이다.

매 iteration마다 정답 상태에서 출발해 현재 모델로 $k-1$ step을 no_grad로 굴리고(그래프에 안 올림), 마지막 한 step만 그래프에 올려 backprop한다.

for it in range(num_iters):
    x = sample_fresh_data()
    k = sample_rollout_depth()         # 예: 1..K 중 랜덤

    with torch.no_grad():              # 현재 가중치, 그래프 X
        for _ in range(k - 1):
            x = f_theta(x)

    x_next = f_theta(x)                # 마지막 1 step만 그래프에 올림
    loss = criterion(x_next, target[k])
    loss.backward()                    # 1-step gradient

여기서 “버퍼”는 영구 저장소가 아니라 매 step 새로 생성되고 버려지는 임시적 개념이다. 사실상 pushforward trick을 임의 깊이 $k$ 로 일반화한 형태로 볼 수 있다.

특징

On-policy. $k-1$ step의 사전 롤아웃을 현재 $\theta$ 로 만들기 때문에, 모델이 보는 입력 분포가 실제 추론 시 분포와 일치한다. classic replay buffer의 staleness가 사라진다.
메모리는 여전히 1 step. no_grad 구간은 activation을 저장하지 않으니, 그래프 메모리는 BPTT와 달리 $k$ 와 무관하게 작다.
별도 저장공간이 없다. 버퍼를 들고 다니지 않으므로 RB의 CPU/GPU 저장 부담이 없다.
깊이 $k$ 를 명시적으로 샘플링할 수 있어, 롤아웃 깊이 분포를 직접 통제하기 쉽다.

대신 trade-off는 연산이다. 매 iteration마다 $k-1$ 번의 forward(비록 no_grad라 가볍지만)를 새로 돌려야 한다. 깊은 $k$ 를 자주 뽑으면 RB처럼 한 번 만든 상태를 재사용하지 못하고 매번 다시 굴리는 비용이 든다.

요약: 매 step 현재 모델로 롤아웃을 즉석 생성 → on-policy + 저장공간 0 + 깊이 직접 제어. 대가는 매 iteration의 추가 forward 연산.

실제 사례: COAF

이 전략이 단순한 학습 트릭을 넘어 결합(coupled) Earth system 모델의 안정성에 직접 기여한 사례가 COAF(Coupled Ocean-Atmosphere Framework, Chen et al., 2025)다. COAF는 사전학습된 대기·해양 모델을 spectral transformation 기반 coupler로 동적으로 결합하는 구조인데, 두 가지 핵심 기여 중 하나로 바로 online replay buffer 메커니즘을 내세운다.

결합 모델에서 대기와 해양은 시간 스케일이 크게 다르고, 두 컴포넌트가 서로의 (틀린) 예측을 주고받으며 굴러가기 때문에 long-range 롤아웃에서 오차가 더 쉽게 증폭된다. COAF는 online replay buffer로 모델이 자기 자신의 결합 롤아웃 분포를 on-policy로 경험하게 하여 장기 안정성을 끌어올리고, 10일 이상 예측에서 Z500·T2m의 latitude-weighted RMSE를 약 10% 낮췄다고 보고한다. single-step teacher forcing만으로는 결합 시스템의 누적 오차를 잡기 어렵다는 점을 보여주는, 본 글의 §4 전략에 대한 구체적 응용 예다.

참고: COAF는 ICLR 2026에 제출되었다가 철회된(withdrawn) 원고다. 공개 게시 시에는 정식 게재 논문이 아니라 OpenReview preprint(withdrawn submission)로 인용하는 것이 정확하다.

5. 한눈에 비교

기준	BPTT	Replay Buffer	Online Replay Buffer
Gradient horizon	$K$ step 전체	1 step	1 step
On/off-policy	on-policy	off-policy (stale)	on-policy
그래프 메모리	$O(K)$ (높음)	$O(1)$	$O(1)$
별도 저장공간	없음	버퍼 저장 필요	없음
Update당 연산	$K$× fwd+bwd	1× fwd+bwd	$k$× fwd ($k{-}1$은 no_grad) + 1× bwd
깊이 분포 제어	명시적(but 짧음)	암묵적	명시적
핵심 한계	메모리·gradient 불안정	staleness·저장비용	깊은 $k$ 의 재계산 비용

직관적으로 정리하면 이렇다.

BPTT는 “gradient가 미래를 본다”는 품질을 위해 메모리를 지불한다.
Replay Buffer는 깊은 롤아웃을 싸게 재사용하지만 그 대가로 입력이 낡는다(off-policy).
Online Replay Buffer는 낡음을 없애는 대신 매번 다시 굴리는 연산을 지불한다.

세 전략 모두 single-step teacher forcing보다 AR 안정성을 크게 개선하지만, 어느 것도 “공짜 점심”은 아니다. 각자 메모리 ↔ 연산 ↔ on-policy성 ↔ 깊이 제어 사이의 다른 지점에 서 있을 뿐이다.

참고 문헌

Brandstetter, J., Worrall, D., & Welling, M. (2022). Message Passing Neural PDE Solvers. ICLR. — pushforward trick
Chen, K., Ling, F., Fei, B., Han, T., Wang, S., Liu, L., Li, B., & Bai, L. (2025). Bridging Oceans and Atmosphere: A More Comprehensive Weather Model (COAF). OpenReview, ICLR 2026 withdrawn submission. https://openreview.net/forum?id=ybLoYEO9PF — online replay buffer 기반 결합 모델
Lam, R. et al. (2023). Learning skillful medium-range global weather forecasting (GraphCast). Science. — autoregressive fine-tuning curriculum
McCabe, M. et al. (2023). Towards Stability of Autoregressive Neural Operators. — AR 안정성 분석

Jeong-Hwan Kim

1. 왜 multi-step 학습이 필요한가

2. BPTT (Backpropagation Through Time)

아이디어

특징

실제 사례

3. Replay Buffer

아이디어

특징

4. Online Replay Buffer

아이디어

특징

실제 사례: COAF

5. 한눈에 비교

참고 문헌