Score SDE와 PF-ODE

📄 관련 논문: [ICLR 2021] Score-Based Generative Modeling through Stochastic Differential Equations

초기 디퓨전 모델인 NCSN과 DDPM은 각자의 방식으로 큰 성공을 거두었지만, 두 모델 모두 노이즈를 주입할 때 $t=1,2,\dots, T$ 처럼 이산 시간 (Discrete-time) 기반의 노이즈 스케줄을 사용했다.

반면, Score SDE는 이를 연속 시간 (Continuous-time) 위로 끌어올렸고, 완전히 다른 방식처럼 보였던 DDPM과 NCSN을 하나의 수학적 공식 아래로 통합하였다.

이렇게 연속적인 시간을 도입하면서 데이터의 Noising / Denoising 과정을 미분방정식으로 표현할 수 있게 되었고, 결과적으로 무작위 노이즈에서 이미지를 생성하는 과정이 단순히 시간에 따른 미분방정식을 푸는 문제로 치환되었다.

생성 문제가 미분방정식 풀이로 바뀌면서, 수학자들이 수백 년 동안 만들어둔 수치해석 도구들 (Solver)을 그대로 가져다 쓸 수 있게 되었다.
이 덕분에 디퓨전 모델의 느린 샘플링 속도를 개선할 수 있는 DPM-Solver와 같은 방법들이 생겨날 수 있었다.

Discrete-time Process에서 Continuous-time Process로의 전환

NCSN과 DDPM의 Forward Process를 $\Delta t$ 간격으로 정리하면 아래와 같다.

• NCSN

  \[\vx_{t+\Delta t}=\vx_t+\sqrt{\sigma_{t+\Delta t}^2-\sigma_t^2\vepsilon_t} \approx \vx_t+\sqrt{\frac{d\sigma_t^2}{dt}\Delta t}\vepsilon_t \tag{1}\]

• DDPM

  \[\vx_{t+\Delta t}=\sqrt{1-\beta_t}\vx_t+\sqrt{\beta_t}\vepsilon_t \approx \vx_t-\frac{1}{2}\beta_t\vx_t\Delta t+\sqrt{\beta_t\Delta t}\vepsilon_t \tag{2}\]

이렇게 두 모델의 식을 $\Delta t$ 간격으로 정리해 보면, 아래의 동일한 구조를 공유하고 있다는 것을 알 수 있다.

\[\vx_{t+\Delta t}\approx \vx_t+\vf(\vx_t,t)+g(t)\sqrt{\Delta t\vepsilon_t} \tag{3}\]

Score SDE

식 (3)에 극한을 취하면 $(\Delta t\to0)$, 위의 Discrete-time 과정은 아래와 같은 Continuous-time SDE로 수렴하게 된다.
즉, 시간 간격 $\Delta t$를 무한히 좁혀서 이산적인 과정을 연속적인 시간 흐름으로 확장하면, 이를 SDE로 표현할 수 있다.

\[d\vx(t)=\vf(\vx(t),t)dt+g(t)d\vw(t) \tag{4}\]

여기서 $\vw(t)$는 표준적인 Wiener Process (Brown motion)으로, 무작위성을 의미한다.

Wiener Process란?

항상 0에서 출발하며, 겹치지 않는 시간 구간에서의 변화량은 서로 독립적입니다. 예를 들어 $s$초에서 $t$초 사이의 변화량 $w(t) - w(s)$는 평균이 0이고 분산이 시간 차이($t-s$)와 동일한 정규분포(가우시안 분포)를 따릅

이는 독립적인 가우시안 변동성이 계속 누적되기 때문에 선이 끊어지지는 않지만, 아무리 현미경으로 확대해도 뾰족뾰족하게 지그재그로 움직여 매끄러운 접선을 그릴 수 없다는 뜻

아주 짧은 시간 $dt$ 동안의 변화량 $d\vw(t)$는 분산이 $dt$인 가우시안 노이즈로 모델링된다.

\[d\vw(t)\sim\mathcal{N}(\vzero,dt\vI)\]

이렇게 데이터에 노이즈를 더하는 Forward SDE 식 (4)를 설계하면, 노이즈에서 다시 데이터로 돌아오는 Reverse-time SDE의 형태가 수학적으로 자동 결정된다.

\[d\vx(t)=[\vf(\vx(t),t)-g^{2}(t)\nabla_{\vx}\log p_{t}(\vx(t))]dt+g(t)d\vw(t) \tag{5}\]

식 (5)에서 각 항의 의미는 다음과 같다.

• $-g^{2}(t)\nabla_{\vx}\log p_{t}(\vx(t))$ (Score): 노이즈를 제거하여 데이터를 구조화하는 복원력이다.
• $g(t)d\vw(t)$ (Diffusion): 무작위성을 부여하여 샘플의 다양성을 확보한다.
  
• $\nabla_{x}\log p_{t}(\vx(t))$ (Score Function): 데이터가 많이 모여있는, 즉 밀도가 높은 영역으로 궤적을 강하게 끌어당긴다.

이때 Score 항과 Diffusion 항은 시간에 따라 역할 비중이 달라진다.

• 생성 초기 ($t \approx T$): 이 시점에는 노이즈 스케일 $g(t)$가 매우 크기 때문에, 모델은 강한 노이즈 (무작위성)를 바탕으로 데이터 공간 전체를 넓게 탐색한다.
• 생성 후기 ($t \approx 0$): 목표에 가까워질수록 $g(t)$가 작아지며, Score Function이 깨끗한 데이터 매니폴드에 샘플을 안착시킨다.

이렇게 결정된 식 (5)의 Reverse-time SDE 수식을 보면, 우리가 모르는 미지수는 각 시간대에서의 Score Function $\nabla_{x}\log p_{t}(\vx(t))$ 하나 뿐이다.
하지만 이 값을 직접 계산할 수 없기 때문에, 신경망 $\vs_\phi(\vx(t),t)$를 학습시켜 이 Score Function을 근사하게 만든다.

따라서 모델이 Score Function만 정확히 학습해 내면, 생성 과정은 단순히 Reverse-time SDE를 수치해석적으로 푸는 (적분하는) 과정으로 귀결된다.

출처: The Principles of Diffusion Models

PF-ODE (Probability Flow ODE)

식 (5)의 Score SDE 수식을 보면, Wiener process인 $d\vw$가 포함되어 있다.
이는 데이터에 노이즈를 점진적으로 가하고 다시 복원하는 과정 전반에 무작위성(Stochasticity)이 존재한다는 것을 의미한다.

반면, PF ODE는 Score SDE와 매 step $t$ 마다 정확히 동일한 Marginal distribution $p_t(\vx)$를 공유하면서도 샘플링 경로를 결정론적(Deterministic)으로 설계한 방정식이다.

Reverse-time PF ODE는 다음과 같다.

\[d\vx(t)=[\vf(\vx(t),t)-\frac{1}{2}g^{2}(t)\nabla_{\vx}\log p_{t}(\vx(t))]dt \tag{6}\]

확률적 과정인 식 (5)와 비교하면, 무작위성을 나타내는 $d\vw$ 항이 완전히 사라지고, 대신 Score 항 앞의 계수가 $\frac{1}{2}$로 변형된 것을 확인할 수 있다.

이처럼 방정식 내에 무작위 항이 없기 때문에, Latent space와 Data space 간의 일대일 맵핑이 보장되어 동일한 초기 노이즈에서 출발하면 항상 똑같은 이미지가 확정적으로 생성된다.

결정론적 방식을 채택하면서, 모델은 SDE가 갖지 못했던 두 가지 강력한 이점을 얻게 된다.

• 무작위성이 포함된 SDE는 노이즈와 데이터 간의 일대일 매핑이 보장되지 않아 한 번 지나온 확률적 궤적을 정확하게 되돌릴 수 없다.
  반면, PF-ODE는 확정적 궤적을 가지므로, 공식 하나로 양방향 모두를 자유롭게 풀 수 있어 정확한 역산과 Latent space 조작이 가능해진다.
• ODE는 SDE보다 수학과 물리학 분야에서 훨씬 오랫동안 깊이 연구되어 온 주제이다.
  따라서 이미 최적화가 끝난 기존의 수많은 ODE Solver를 그대로 가져다 사용하여, 디퓨전 모델의 느린 샘플링 속도를 가속할 수 있다.

출처: The Principles of Diffusion Models