Technical NotebookDDangchani's DataLog

Archive

Score-Based Generative Modeling through Stochastic Differential Equations

Stochastic Differential Equation DDPM 등의 Diffusion 모델은 확률미분방정식 Stochastic differential equation, SDE 의 형태로부터 구할 수 있다. 비록 SDE 기반의 diffusion 모델은 DDPM 등의 모델보다 뒤늦게 제안되었지만, SDE의 한 종류로 볼 수 있다. 이번 글에서는 diffusion 모델들의 근간이 되는...

2024. 4. 21.6 min read

Stochastic Differential Equation

DDPM 등의 Diffusion 모델은 확률미분방정식Stochastic differential equation, SDE의 형태로부터 구할 수 있다. 비록 SDE 기반의 diffusion 모델은 DDPM 등의 모델보다 뒤늦게 제안되었지만, SDE의 한 종류로 볼 수 있다. 이번 글에서는 diffusion 모델들의 근간이 되는 확률미분방정식 기반의 생성형 모델을 살펴보도록 하겠다.

ODE and SDE

이전 글에서 살펴본 것과 같이, 시간 tt에서의 확률변수 X(t)X(t)에 대한 상미분방정식Ordinary differential equation, ODE는 다음과 같은 형태로 주어진다.

dXdt(t)=f(X(t),t)(ODE) \frac{dX}{dt}(t) = f(X(t), t) \tag{ODE}

초기값 조건 X(0)=X0X(0)= X_{0} 이 주어지면, 위 ODE의 해는 다음과 같다.

X(t)=X0+∫0tf(X(s),s)ds X(t) = X_{0}+ \int_{0}^{t}f(X(s),s)ds

확률미분방정식(SDE)은, ODE에 다음과 같이 Brownian motion 항이 추가된 것으로 정의된다.

dXt=f(Xt,t)dt+g(t)dWt(SDE) dX_{t}= f(X_{t},t)dt + g(t)dW_{t} \tag{SDE}

여기서 WtW_{t}는 Brownian motion(혹은 Wiener process)이다.

Brownian motion : Brownian motion WtW_t는 다음과 같은 특성을 가지는 확률과정이다.

  1. W0=0W_{0}=0 almost surely
  2. 변화량 Wt+u−WtW_{t+u}-W_{t}는 WtW_{t}와 독립이다. 즉, Wt+u−Wt⊥WtW_{t+u}-W_{t}\perp W_{t}
  3. 변화량 Wt+u−WtW_{t+u}-W_{t}는 N(0,u)N(0,u)를 따른다. 즉, Wt+u−Wt∼N(0,u)W_{t+u}-W_{t}\sim N(0,u)

ODE에서와 마찬가지로 다음과 같이 적분 형태의 해를 생각할 수 있다.

Xt=X0+∫0tf(Xs,s)ds+∫0tg(Xs,s)dWs X_{t}=X_{0}+ \int_{0}^{t}f(X_{s},s)ds + \int_{0}^{t}g(X_{s},s)dW_{s}

이때, SDE의 해를 위와 같은 적분 형태로 생각하기보다는 확률변수열 {Xt}t=0T\lbrace X_{t}\rbrace _{t=0}^{T} 의 결합분포로 생각하는 것이 더 자연스럽다. 또한, 결합분포에서 tt 시점의 주변확률분포를 Xt∼ptX_{t}\sim p_{t} 라고 표현하자.

Fokker-Planck-Kolmogorov Equation

SDE(식 (SDE)(\mathrm{SDE})) 의 주변확률분포 Xt∼ptX_{t}\sim p_{t}에 대해 다음과 같은 편미분방정식partial differential equation이 주어지는데, 이를 Fokker-Plank-Kolmogorov(FPK) equation이라고 한다.

∂tpt(x)=−∂x(f(x,t)pt(x))+g2(t)2∂x2pt(x) \partial_{t}p_{t}(x) = -\partial_{x}(f(x,t)p_{t}(x)) + \frac{g^{2}(t)}{2}\partial_{x}^{2}p_{t}(x)

다변량 확률분포(d>1d>1)의 FPK equation은 다음과 같다.

∂tpt(x)=−∇x⋅(fpt)+12tr(g⊤∇x2ptg) \partial_{t}p_{t}(x) = - \nabla_{x}\cdot(fp_{t}) + \frac{1}{2}{\mathrm{tr}}(g^{\top}\nabla_{x}^{2}p_{t}g)

증명 : 부분적분을 이용한 증명 (참고 : Applied Stochastic Differential Equations Ch 5-2)

Example. Ornstein-Uhlenbeck Process

Ornstein-Uhlenbeck process는 다음과 같은 SDE로 주어진다.

dXt=−βXtdt+σdWtX0∼N(0,σ2β)Xt∣X0∼N(e−βtX0,σ22β(1−e−2βt))\begin{aligned} dX_{t} &= -\beta X_{t}dt + \sigma dW_{t}\\ X_{0} &\sim N\left(0, \frac{\sigma^2}{\beta}\right)\\ X_t \mid X_0 &\sim N\left(e^{-\beta t}X_0, \frac{\sigma^2}{2\beta}(1-e^{-2\beta t})\right) \end{aligned}

이때, 위 SDE의 해는 다음과 같다.

Xt∼N(0,σ22β)pt(Xt)=1πσ2/βexp⁡[−βσ2Xt2]\begin{aligned} X_t &\sim N\left(0, \frac{\sigma^2}{2\beta}\right)\\ p_t(X_t) &= \frac{1}{\sqrt{\pi\sigma^2/\beta}}\exp\left[-\frac{\beta}{\sigma^2}X_t^2\right] \end{aligned}

위의 해가 FPK equation을 만족함을 확인할 수 있다.

0=∂tpt(Xt)=−∂x(f(Xt,t)pt(Xt))+g2(t)2∂x2pt(Xt)=−∂x(−βxpt(x))+σ22∂x2pt(x)=0\begin{aligned} 0 = \partial_t p_t(X_t) &= -\partial_x(f(X_t,t)p_t(X_t)) + \frac{g^2(t)}{2}\partial_x^2p_t(X_t)\\ &= -\partial_x(-\beta xp_t(x)) + \frac{\sigma^2}{2}\partial_x^2p_t(x)\\ &= 0 \end{aligned}

Reverse-time SDE and Data Generation

Diffusion 모델은 주어진 데이터(e.g. image)를 노이즈로 변환하는 과정을 통해 이미지의 특성을 학습한다. 이후 노이즈를 생성하여 원래 이미지를 복원하는 과정을 통해 새로운 이미지를 생성한다 (아래 그림 참고).

Diffusion Source: Song et al. (2021)

즉, X0∼p0X_{0}\sim p_{0}가 주어진 이미지 데이터의 분포에 대응되며 이로부터 노이즈 분포 XT∼pTX_{T}\sim p_T 를 학습한다. 반면, 새로운 데이터를 생성하기 위해서는 노이즈를 랜덤 샘플링하여 pTp_T로 부터 p0p_{0}을 생성해야 한다. 노이즈는 앞선 Ornstein-Uhlenbeck process의 예시처럼 XT∼N(0,σ22βI)X_T \sim N(0, \frac{\sigma^2}{2\beta}I) 를 샘플링하면 된다. 반면, p0p_{0}로부터 pTp_{T}를 생성하기 위해서는 역방향으로 정의된 SDE를 사용해야 하며, 이를 Reverse-time SDE라고 한다.

Theorem (Anderson, 1982)

정방향 확률미분방정식

dXt=f(Xt,t)dt+g(t)dWt,X0∼p0 dX_{t}= f(X_{t},t )dt + g(t)dW_{t},\quad X_{0}\sim p_0

에 대응되는 역방향 확률미분방정식은

dXˉt=(f(Xˉt,t)−g2(t)∇xlog⁡pt(Xˉt))dt+g(t)dWˉt,XˉT∼pT(RSDE) d\bar X_{t}=(f(\bar X_{t},t)-g^{2}(t) \nabla_{x}\log p_{t}(\bar X_{t}))dt + g(t)d\bar W_{t}, \bar X_{T}\sim p_{T}\tag{RSDE}

와 같다. 여기서 Wˉt\bar W_{t}는 역방향 Brownian motion을 의미한다.

증명 : FPK equation을 사용하여 동치를 보일 수 있다.

Reverse-time ODE

정방향 SDE

dXt=fdt+gdWtdX_t = fdt + gdW_t

에 대응하는 역방향 SDE

dXˉt=(f−g2∇xlog⁡pt)dt+gdWˉtd\bar X_t = (f - g^2\nabla_x\log p_t)dt + g d\bar W_t

를 살펴보았는데, 이때 XˉT\bar X_T의 분포 pTp_T는 다음과 같은 역방향 ODE의 해이기도 하다.

dXˉt=(f(Xˉt,t)−g2(t)2∇xlog⁡pt(Xˉt))dt(RODE)d\bar X_t = \left( f(\bar X_t,t)- \frac{g^2(t)}{2}\nabla_x\log p_t(\bar X_t) \right)dt \tag{RODE}

위 ODE는 일종의 flow model로 볼 수 있게 된다. 증명은 마찬가지로 FPK equation을 사용하여 보일 수 있다.

Sample Generation

다음과 같은 정방향 SDE를 고려하자.

dXt=−βXtdt+σdWtd X_t = -\beta X_t dt + \sigma dW_t

이때, TT가 충분히 클 경우 pT≈N(0,σT2I)p_T \approx N(0,\sigma_T^2I) 이다. 또한, 식 (RODE)(\mathrm{RODE})에 따라 역방향 ODE는 다음과 같다.

dXˉt=(σ22∇xlog⁡pt(Xˉt)−βXt)dtd\bar X_t = \left(\frac{\sigma^2}{2}\nabla_x\log p_t(\bar X_t) - \beta X_t \right)dt

실제로 위 RODE로부터 샘플을 얻기 위해서는, discretization(Euler-Maruyama)을 통해 다음과 같이 샘플을 얻을 수 있다.

XˉK∼N(0,σT2I)for k=K,K−1,…,1:Xˉk−1=Xˉk−Δt(−βXˉk−σ22∇xlog⁡pkΔt(Xˉk))end\begin{aligned} &\bar X_K \sim N(0,\sigma_T^2I)\\ &\text{for } k = K, K-1, \ldots, 1:\\ &\quad \bar X_{k-1} = \bar X_k - \Delta t \left(-\beta \bar X_k - \frac{\sigma^2}{2}\nabla_x\log p_{k\Delta t}(\bar X_k)\right) \\ &\text{end} \end{aligned}

Score Matching

(RODE)(\mathrm{RODE})의 샘플링 과정에는 score function ∇xlog⁡pt\nabla_x \log p_{t}에 대한 정보가 필요하다. 그러나 실제로는 분포 ptp_{t}를 직접 구할 수 없기 때문에, 이를 근사하는 신경망 모델(score network)을 생성하게 된다.

∇xlog⁡pt(x)≈sθ(x,t)\nabla_x \log p_{t}(x) \approx s_{\theta}(x,t)

Score network의 구조는 일반적으로 U-Net 구조를 사용한다 (아래 그림 참고).

Score Network Source: paperswithcode.com

이렇게 구한 score network sθs_{\theta}를 이용하여 ptp_{t}의 score function을 근사할 수 있다. 이를 Score matching이라고 하며, 학습을 위해 다음과 같은 손실함수를 이용한다.

L(θ)=∫0Tλ(t)EXt[∥sθ(Xt,t)−∇Xtlog⁡pt(Xt)∥2]dt \mathcal{L}(\theta) = \int_{0}^{T}\lambda(t) \mathrm{E}_{X_{t}}\left[\Vert s_{\theta}(X_{t},t )-\nabla_{X_{t}}\log p_{t}(X_{t})\Vert^{2}\right]dt

이때, 위 손실함수는 아래 두 손실함수와 동치이다.

L(θ)=∫0Tλ(t)EX0[EXt∣X0[∥sθ(Xt,t)−∇Xtlog⁡pt∣0(Xt∣X0)∥2∣X0]]L(θ)=∫0Tλ(t)EXt[∥sθ(Xt,t)∥2+2Eν[ddhν⊤sθ(Xt+hν,t)∣h=0]]dt+C \begin{align} \mathcal{L}(\theta) &= \int_{0}^{T}\lambda(t)\mathrm{E}_{X_{0}}\left[\mathrm{E}_{X_{t}\mid X_{0}}\left[\left\Vert s_\theta(X_{t},t )-\nabla_{X_{t}}\log p_{t\mid 0}(X_{t}\mid X_{0})\right\Vert^{2} \mid X_{0}\right]\right]\\ \mathcal{L}(\theta)&= \int_{0}^{T}\lambda(t)\mathrm{E}_{X_{t}} \left[\Vert s_\theta(X_{t},t)\Vert^{2}+ 2\mathrm{E}_{\nu}\left[\frac{d}{dh}\nu^{\top}s_\theta(X_{t}+h\nu,t)\bigg\vert_{h=0}\right]\right]dt+C \end{align}

여기서 ν∼N(0,I)\nu\sim\mathcal{N}(0,I) 이다. 위 손실함수를 이용하는 방법을 Denoising score matching(DSM) 이라고 하며, 아래의 손실함수를 이용하는 방법을 Sliced score matching(SSM) 이라고 한다. DSM은 조건부 확률 Xt∣X0X_{t}\mid X_{0} 에 대한 분포를 요구하기 때문에, 일반적으로는 SSM이 더 널리 사용된다. 반면 두 방법이 모두 적용가능한 경우에는 DSM이 더 나은 성능을 보인다고 알려져 있다.

Sampling

위의 score matching을 이용하여 sθs_\theta 를 학습한 경우, 역방향 SDE와 ODE로부터 샘플을 얻을 수 있다. 두 경우 모두 대개 Ornstein-Uhlenbeck process를 이용한다. 우선 역방향 SDE의 샘플링은 다음과 같다.

dXˉt=(−βXˉt−σ2sθ(Xˉt,t))dt+σdWˉt=(σ2σtϵθ(Xˉt,t)−βXˉt)dt+σdWˉtXˉT∼N(0,σT2I)\begin{aligned} d\bar X_t &= \left( -\beta \bar X_t - \sigma^2 s_\theta(\bar X_t, t)\right)dt + \sigma d\bar W_t\\ &= \left(\frac{\sigma^2}{\sigma_t}\epsilon_\theta(\bar X_t, t)-\beta \bar X_t\right)dt + \sigma d\bar W_t\\ \bar X_T &\sim N(0,\sigma_T^2I) \end{aligned}

마찬가지로, 역방향 ODE의 샘플링은 다음과 같다.

dXˉt=(−βXˉt−σ22∇xsθ(Xˉt,t))dt=(σ22σtϵθ(Xˉt,t)−βXˉt)dtXˉT∼N(0,σT2I)\begin{aligned} d\bar X_t &= \left( -\beta \bar X_t - \frac{\sigma^2}{2}\nabla_x s_\theta(\bar X_t,t) \right)dt\\ &= \left( \frac{\sigma^2}{2\sigma_t}\epsilon_\theta(\bar X_t,t) - \beta \bar X_t \right)dt\\ \bar X_T &\sim N(0,\sigma_T^2I) \end{aligned}

일반적으로는 SDE 샘플링이 더 고품질의 샘플을 생성하는 것으로 알려져 있다. 반면, ODE 샘플링은 flow model로 볼 수 있기 때문에, 가능도 함수 계산의 측면에서 유용하며 image interpolation 등에도 활용될 수 있다.

Image Interpolation : 두 이미지 X(1),X(2)X^{(1)}, X^{(2)} 사이의 이미지를 생성하는 방법으로, Forward time ODE를 이용해 XT(1)X_T^{(1)}와 XT(2)X_T^{(2)}를 생성한 후, XT(θ)=(1−θ)XT(1)+θXT(2)X_T^{(\theta)} = (1-\theta)X_T^{(1)} + \theta X_T^{(2)}를 생성한다. 이후 역방향 ODE를 이용해 X0(θ)X_0^{(\theta)}를 생성한다. (아래 예시 참고)

Image Interpolation Source: Song et al. (2021)

References

  • 서울대학교 기계지능특강-Generative AI 강의노트
  • Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., & Poole, B. (2021). Score-Based Generative Modeling through Stochastic Differential Equations (arXiv:2011.13456). arXiv. https://doi.org/10.48550/arXiv.2011.13456
  • Murphy, K. P. (2023). Probabilistic machine learning: Advanced topics. The MIT Press.
  • Särkkä, S., & Solin, A. (2019). Applied Stochastic Differential Equations. Cambridge University Press.