이 세미나를 듣기 전에

아래의 내용들을 정말 잘 알고 있어야 합니다.

• random variable
• random vector
• pdf (probability density function)
• expectation
• mean
• variance
• 가우시안 랜덤벡터 (gaussian random vector)
• independent

자료제공자의 확률미분방정식 관련 경력사항

• 학부생때 문과생들도 듣는 확률과 통계 들음 (고등학생도 푸는 문제 위주이고 엑셀로 정규분포 끄적거리던 수준)
• 석사과정일 때 확률론 들음 (공대에서 듣는 랜덤프로세스보다 범위가 한참 적으나 깊이는 더 깊다. discrete time에서 랜덤프로세스 다룬게 전부)
• 확률론 응용과정으로 보험수학개론 들음
• 확률미분방정식 수업 들은적 없음 애초에 대한민국에서 확률미분방정식만을 위한 수업 열리는 곳도 거의 없음
• 석사과정때 확률미분방정식 존재는 알았고, 거기서 파생된 결과를 이용해 expectation계산하고 학위논문 작성
• 석사과정때 확률과 통계, 실험 통계학, 확률미적분과 금융수학 조교
  • 확률과 통계, 실험 통계학은 1주일에 한번 수업할당 받음
    • 30분 설명, 20분 퀴즈 이것을 위해 공부 하게 됨

Diffusion model through SDE (Stochastic differential equation)

Forward SDE

$$\begin{equation} dx_t = f(x_t,t)dt + g(t) dw_t, 0\leq t \leq T, w_t \text{: Wiener process} \label{eq:forwardsde}\tag{1} \end{equation}$$

위의 수식 ($\ref{eq:forwardsde}$)을 forward SDE 라고 부른다. forward SDE ($\ref{eq:forwardsde}$)를 이용하여 random vector $x_0$를 변형시키고 노이즈를 추가하여 $x_t$를 만든다. $f(x_t, t)$가 변형시키는 역할을 하며 drift term이라 부른다. $dw_t$는 쉽게 말하면 wiener process의 차이인데 이것이 노이즈를 추가하는 역할을 하고 그 정도를 $g(t)$가 조절한다. $g(t)$를 diffusion term이라 부른다.

정리

• $f(x_t,t)$ : drift term이고 $(x_t,t)$에 대한 함수다.
• $g(t)$ : diffusion term이라 부르고 $t$에 대한 함수이다.
• Forward SDE를 이용하여 $x_0$를 변형
• $dx_t$ : $x_t$의 순간적인 변화

질문

• $f(x_t,t)dt$는 $\int f(x_t,t) dt$에서 부터 나온것 같은데 $g(t) dw_t$는 어떻게 정의함? → Ito integration

 

Ito integration 정의

Ito integration정의 하자. 이것을 정의하는 목적은 diffusion model을 구성하는 forward SDE에 나오는 $g(t) dw_t$를 정의하기 위해서이다.

왜 Ito integration를 정의할까?

• Forward SDE ($\ref{eq:forwardsde}$) 에서 등장한  $g(t)dw_t$가 뭔지 알기위해

Stochastic Process의 정의

양수 $T >0$가 있다고 하자. random vector의 모임 $\{x_t\}_{[t\in [0,T]]}$를 stochastic process라 한다.

Wiener process의 정의

Wiener process $\{w_t\}_{[0,T]}$를 정의하자. 위너프로세스 $w_t$는 $\mathbb{R}^d$에서 정의된다. 아래와 같은 성질을 갖는 stochastic process를 의미한다.

1. $t,s$가 있을 때 $w_{t}-w_s \sim N(0,|t-s|I)$
2. $t_1 < t_2< t_3 < t_4 <…<t_N $일 때 $w_{t_i}-w_{t_{i-1}}$들은 independent, time축이 겹치지 않으면 독립이라는 의미이다.
3. 관습적으로 $w_0=0$이라 한다.

 

Wiener process의 중요한 성질

Wienerprocess $w_t$의 중요한 성질에 대해 알아보자. $dt$를 $t$으로 아주 작은 변화라고 하자. 그러면 아래와 같은 식이 성립한다고 보면 된다.

$$\begin{equation} (w_{t+dt} – w_t)(w_{t+dt}-w_t)^T \approx dt I \tag{2}\label{eq:wienerquad} \end{equation}$$

아주 중요한 성질이다. 이 점을 이용하면 아래와 같은 사실을 도출할 수 있다.

$$\begin{equation}(w_{t+dt}-w_t) \approx \sqrt{|dt|} (1,1,…,1)^T \tag{3}\label{eq:wienersqrd} \end{equation}$$

수학적으로 엄밀한 증명은 skip 하도록 하겠다. 위에서 언급한 성질 ($\ref{eq:wienerquad}$)와 ($\ref{eq:wienersqrd}$)로 부터 좋은 성질들이 많이 파생되오니 알아두면 좋다.

대망의 Ito Integration (이토적분) 정의

$w_t$를 wiener process라고 하자. $w_t$의 차원을 $d$차원이라 하자. $A_t$를 stochastic process 인데 $A_t$의 차원이 $m \times d$라고 하자. 실수축에 $S_1 < S_2$가 있다고 하자. 이제 ito integration은 아래와 같이 정의한다.

$\begin{equation} \int_{S_1}^{S_2} A_t dw_t = \lim_{\Delta t \to 0 } \sum_{i=0}^{N-1} A_{t_{i}} (w_{t_{i+1}}-w_{t_i}) \tag{4}\label{eq:itointegration} \end{equation}$

$\ref{eq:itointegration}$에서 $t_0=S_1 < t_1 < t_2 <…<t_N=S_2$, $\Delta t = max\{t_{i+1}-t_i | i=0,1,…,N-1\}$. 주목할점은 $A_{t_{i}}$에서 $t_i$는 $[t_{i}, t_{i+1}]$에서 왼쪽 점이라는 것이다.

Ito integration 정의할 때 partition의 왼쪽점만 사용하는 이유?

왜 $\ref{eq:itointegration}$에서 $[t_{i}, t_{i+1}]$중에 왼쪽점인 $t_{i}$만 사용할까? 그 이유는

• 설정하는 값 $t_{i}^*$에 따라 $\lim_{\Delta t \to 0 } \sum_{i=0}^{N-1} A_{t_{i}^*} (w_{t_{i+1}}-w_{t_i}) $의 값이 달라지기 때문이다.

partition에서 선택하는 점에 따라 적분 값이 달라지는 경우

$w_t$를 1차원에서 wiener process라 하자. $a(b-a) = \frac{1}{2}(b-a)(b+a)- \frac{1}{2}(b-a)^2$라는 식을 사용하면 $[0,T]$사이에서!

$$ \lim_{\Delta t \to 0} \sum_{i=0}^{N-1} w_{t_i} (w_{t_{i+1}}-w_{t_i}) = w_T^2/2 – T/2$$

$$ \lim_{\Delta t \to 0} \sum_{i=0}^{N-1} w_{t_{i+1}} (w_{t_{i+1}}-w_{t_i}) = w_T^2/2 + T/2$$

위와 같이 왼쪽 점을 선택했을때와 오른쪽을 선택했을때 값이 달라진다.

 

Ito integration 다시!

어떤 구간을 partitioning한후 만들어진 subinterval에서 맨 왼쪽 점 값을 이용해 구한다는 것을 기억하자.

$$\begin{equation} \int_{S_1}^{S_2} A_t dw_t = \lim_{\Delta t \to 0 } \sum_{i=0}^{N-1} A_{t_{i}} (w_{t_{i+1}}-w_{t_i}) \end{equation}$$

주목할점: 디퓨전에서 사용하는 이토적분은?

위에서 이토적분을 stochastic process $A_t$에서 정의햇다. 그런데 디퓨전 모델에서 사용하는 이토적분은 $w_t$가 어떠한 dimension을 갖는것에 상관없고 피적분되는 대상은 deterministic한 실수 함수 $g(t) \in \mathbb{R}$를 사용한다.

$$g(t) dw_t \leftrightarrow g(t) I dw_t \leftrightarrow \int g(t) I dw_t$$

주목할점: Ito integration이용해서 SDE 바라보기

아래와 같은 SDE가 $[T_1, T_2]$에서 정의되어있다고 하자.

$$dx_t = f(x_t,t)dt + g(t)dw_t$$

위의 식을 해석하면 stochastic process $x_t$의 $[T_1,T_2]$에서 $t \in [T_1,T_2]$의 순간적이 변화량이라고 볼 수 있다. 이제 Ito integration을 이용하면 위의 식을 적분형태로 표현가능하다.

$$\text{For } S_1 < S_2, x_{S_2}-x_{S_1} = \int_{S_1}^{S_2} f(x_t,t)dt +\int_{S_1}^{S_2} g(t) dw_t $$

위의 식은 $S_1 < S_2$라는 구체적인 숫자가 주어졌을 때 $x_t$가 $t=S_1$에 있다가 $t=S_2$로 가면서 얼마나 변했는지를 보여준다.

 

다시 디퓨전 모델로 돌아와서

디퓨전 모델은 아래의 forward SDE를 이용해서 $x_t$를 변형하고 노이즈를 추가한다고 하였다.

$$dx_t = f(x_t,t)dt+g(t)dw_t, 0\leq t \leq T , t=0 \to t=T\tag{forward SDE}$$

$x_t$는 데이터 분포 $p_{data}(x)$를 따르는 $x_0$부터 시작해서 위의 forward SDE에 의해 변형되서 마지막 점인 $t=T$에서 $x_T$가 노이즈가 낀 분포 (가령 가우시안)를 따르길 원한다. 이것을 좀더 수식을 써서 기술하자.

• $p_t(x)$를 위의 forward SDE에 의해 기술되는 stochastic process $x_t$의 $t$에서 pdf라 하자.
• $x_t$는 $t=0$에서 시작하고 시작점인 $t=0$에서 데이터분포를 따른다. 데이터 분포를 $p_{data}(x)$라 표시하면 $p_0(x) = p_{data}(x)$라는 조건이 붙는다.
• forward SDE를 통해 $x_t$의 $t$에서 분포인 $p_t(x)$가 변화하는데 최종적으로 $p_T(x)$가 가우시안 $N(\mu_\star, \sigma^2_\star I)$이 되길 바란다.
• $f(x,t)$와 $g(t)$를 잘 설정하면 $p_T$는 가우시안분포에 가깝게 된다.
• 그러면!! data sample은 어떻게 생성할까?
• forward SDE는 데이터 → 가우시안 으로 변형하는 것이니 반대 과정을 생각하면 된다.
• 그 반대과정은 reverse SDE를 이용해서 표현된다.

 

Reverse SDE

디퓨전 모델을 이용해서 데이터 샘플을 생성하기 위해 아래의 reverse SDE를 사용한다.

$$dx_t = [f(x_t,t)-g(t)^2 \nabla_{x_t} \log p_t(x_t)] dt + g(t) d\tilde{w}_t, 0\leq t \leq T, t=T \to t=0 \tag{reverse SDE}$$

여기서 $\tilde{w}_t$는 역방향 wiener process를 의미한다. reverse SDE는 $t=T$에서 시작하여 $t=0$위에서 끝난다. forward SDE가 $t=0$에서 시작하고 $t=T$에서 끝나지만 reverse SDE는 방향이 반대라는 얘기이다. 특이점이라면 $g(t)^2 \nabla_{x_t} \log p_t(x_t)$가 생긴다.

왜? Reverse SDE를 이용하면 데이터 샘플을 생성할 수 있을까? 결론만 말하자.

• $p_t$를 forward SDE에 의해 만들어진 stochastic process의 t에서의 pdf라 하자.
• $q_t$를 reverse SDE에 의해 만들어진 stochastic process의 t에서의 pdf라 하자.
• 만약에 $q_T = p_T$라면 즉, reverse SDE의 시작점의 분포가 forward SDE의 끝점의 분포와 같다면 아래와 같이 모든 t에서 pdf가 동일하다.
  $$p_t = q_t \text{ if } p_T = q_T$$
  • 특별히 $p_0 = q_0$이다.
  • $p_0 = p_{data}$였으므로 $q_0 = p_{data}$가 된다.
    • 즉 reverse SDE로 표현되는 $x_t$의 끝점 $t=0$에서의 분포가 data distribution이 된다. 그러므로 Reverse SDE를 이용해서 데이터 샘플을 생성할 수 있다.

그러면 왜? 이와 같은 성질을 만족할까? 아래에 설명할 Fokker Planck Kolmogorov Equation을 사용하면 보일수 있다.

 

Fokker Planck Kolmogorov (FPK) Equation

아래와 같은 SDE가 있다고 하자.

$$dx_t = f(x_t,t)dt + L(x_t,t)dw_t, 0\leq t \leq T \tag{5}\label{eq:generalSDE}$$

여기서 $x_t, f(x_t,t) \in \mathbb{R}^d, L(x_t,t) \in \mathbb{R}^{d \times d}$이다.

여기서 $p(x,t)$를 $\ref{eq:generalSDE}$로 기술되는 $x_t$의 $t$에서의 pdf라고 하자. 그랬을 때 $p(x,t)$는 아래와 같은 미분방정식을 만족하고 이 미분방정식을 Fokker Planck Kolmogorv (FPK) equation이라고 한다.

$$ \frac{ \partial p(x,t)}{\partial t} = – \sum_{i} \frac{\partial}{\partial x_i} [f_i(x,t)p(x,t)] + 1/2 \sum_{i,j} \frac{\partial^2}{\partial x_i \partial x_j} \{ (L(x,t)L(x,t)^T)_{ij} p(x,t) \} \tag{FPK equation}$$

FPK equation을 이용하는 의의?

FPK equation이 등장했는데 왜? 좋을까?

• SDE가 있으면 그에 대응하는 FPK equation이 있고 FPK equation이 있으면 그에 대응하는 SDE가 있다.
• FPK equation을 이용하여 SDE로 표현되는 stochastic process $x_t$의 $t$에서의 분포를 알 수 있다.

 

Reverse SDE와 Forward SDE의 FPK equation은 같다.

믿기지 않겠지만 그렇다. 계산하면 나온다.

추가적으로 끝점 $T$에서 분포가 같으라는 조건이 있다.
→ 이렇게 되면 forward SDE에 의해 만들어진 분포 $p_t$와 reverse SDE에 의해 만들어진 분포 $q_t$는 동일하다!

 

이해를 돕기 위한 간단한 예시

두개의 미분방정식이 있다.

$$dx_t/dt = 3, x_0 = 1 \tag{6} $$

$$dx_t/dt = 3, x_1 = 4 \tag{7} $$

• (6)의 해는 $x_t = 3t +C$ 꼴 인데 $x_0 =1 $이므로 $x_t=3t+1$, $x_1 = 4$
• (7)의 해는 $x_t = 3t +C$ 꼴인데 $x_1 = 4$이므로 $x_t=3t+1$, $x_0 =1$
• 조건 ($t=0$에서 값, $t=1$에서 값)이 달라도 solution은 같을 수 있다.
• 이와 같은 원리로 Reverse SDE와 Forward SDE는 동일한 확률분포 생성

 

디퓨전 모델 학습 방법

다시 반복하자. 디퓨전 모델은 데이터를 sde를 이용해 변형하고 noise를 추정하는 과정 (forward SDE)과 노이즈에서부터 데이터를 생성하는 과정 (reverse SDE)로 구성된다.

dx_t = f(x_t,t)dt+g(t) dw_t, 0\leq t \leq T, t=0 \to t=T (\text{데이터} \to \text{noise}), x_0 \sim ~ p_{data}(x_0) \tag{forward SDE} dx_t = [f(x_t,t)-g(t)^2 \nabla_{x_t} \log p_t(x_t)] dt + g(t) dw_t,0\leq t \leq T, t=T \to t=t (\text{noise} \to \text{데이터})\tag{reverse SDE}

• 여기서 $p_t(x)$는 forward SDE를 따르는 $x_t$의 $t$에서의 pdf이다.
• reverse SDE에서 $\nabla_{x_t} \log p_t(x_t)$가 등장했는데 $\nabla_{x} \log p_t(x)$를 score라 부른다.

Score를 학습해야 한다.

reverse SDE를 사용하기 위해서는 score $\nabla_{x_t} log p_t(x_t)$를 사용해야 한다. 그러면 $\nabla_{x_t} \log p_t(x_t)$를 알던가 추정해야 한다.

Score 학습의 어려움

score를 학습 하기 위해 neural network $s_\theta(x_t,t)$를 학습하려고 할 때 다음과 같은 loss를 생각할 수 있다.

$$L_{SM}(\theta) = E_{t \in U[0,T], x_t \sim p_t(x_t)}\lVert \nabla_{x_t} \log p_t(x_t) -s_\theta(x_t,t) \rVert^2 \tag{score matching loss}$$

그런데 위의 loss에서 사용하고자 하는 참값인 score $\nabla_{x_t} \log p_t(x_t)$와 expectation 에 걸려있는 $p_t(x_t)$를 알 턱이 없다. 분포가 간단하다면 모를까

score matching을 위해 실제로 사용하는 방식

위에서 언급한 score를 알수없는 어려움이 있지만 아래와 같은 식을 이용해서 학습한다.

$$L_{DSM}(\theta)= E_{t \in U[0,T], x_0 \sim p_{data}(x_0), x_t \sim p_t(x_t|x_0)}\lVert \nabla_{x_t} \log p_t(x_t|x_0) -s_\theta(x_t,t) \rVert^2 \tag{denoising score matching loss}$$

• denoising score matching loss $L_{DSM}$에서는 condition을 주고 expectation 하고 condition된 score $\nabla_{x_t} \log p_{t}(x_t|x_0)$ 구하는 방식을 취한다.
• $x_0$의 분포를 모르지만 $x_0$를 고정(conditioning)시켜두고 이것이 어떻게 변하는지만 보자는 얘기이다.

$L_{DSM}$ loss를 최적화해도 되는 이유

• 의문이 들 수 있다. $L_{SM}$은 $\nabla_{x_t} \log p_t(x_t)$를 추정하기 위해 만든 loss이고 $L_{DSM}$은 $\nabla_{x_t} \log p_t(x_t|x_0)$를 추정하기 위해 만든 loss이다.
• 목적이 다른데 왜 이렇게 하느냐?
• 아래와 같은 사실이 있다.

$$\nabla_{\theta} L_{SM}(\theta) =\nabla_{\theta} L_{DSM} $$

• gradient가 같다.
• 딥러닝 학습을 할 때 gradient를 이용해서 업데이트 한다.
• 두 loss의 gradient가 같으므로 $L_{SM}$을 최적화하나 $L_{DSM}$을 최적화 하나 같은 의미이다.
• 그러므로 $L_{DSM}$를 최적화 하는데에 집중하면 된다.

perturbation kernel

$L_{DSM}$을 학습해야 하는것 알겠고 다시 써보자.

$$L_{DSM}(\theta)= E_{t \in U[0,T], x_0 \sim p_{data}(x_0), x_t \sim p_t(x_t|x_0)}\lVert \nabla_{x_t} \log p_t(x_t|x_0) -s_\theta(x_t,t) \rVert^2 \tag{denoising score matching loss}$$

• 결국엔 $p_t(x_t|x_0)$를 알아야한다.
• $p_t(x_t|x_0)$를 perturbation kernel이라 부른다.
• 이거 어떻게 구할까? forward SDE가 linear SDE일 때 쉽게 구할 수 있다.

디퓨전 모델에서는 linear SDE를 사용한다.

$L_{DSM}$에 등장하는 perturbation kernel $p_t(x_t|x_0)$를 구하기 위하여 이것을 쉽게 구할 수 있는 경우인 linear SDE를 사용한다.

dx_t = (F(t)x_t + u(t))dt+g(t) dw_t, 0\leq t \leq T, t=0 \to t=T (\text{데이터} \to \text{noise}), x_0 \sim ~ p_{data}(x_0) \tag{(linear) forward SDE} dx_t = [F(t)x_t+u(t)-g(t)^2 \nabla_{x_t} \log p_t(x_t)] dt + g(t) dw_t,0\leq t \leq T, t=T \to t=t (\text{noise} \to \text{데이터})\tag{reverse SDE}

• $F(t), u(t), g(t)$는 모두 $t$에 관한 deterministic한 함수
• $F(t) \in \mathbb{R}^{d \times d} , u(t) \in \mathbb{R}^d, g(t) \in \mathbb{R}$
• 디퓨전 모델에서는 실제로 linear SDE만 사용한다. perturbation kernel을 구하기 쉽기 때문이다.
• perturbation 과연 어떻게 풀까요?
  • $p_t(x_t|x_0)$구하는 것인데 $x_0$가 주어질때 $x_t$가 어떤 form인지 구하면 되는 것이다.

linear ode 풀이 방법

먼저 linear SDE의 noise없는 버전인 lienar ODE부터 보자.

dx_t =(F(t)x_t +u(t))dt, x_0\text{ given} \tag{linear ODE}

linear ODE 어떻게 풀까? 아래와 같이 어떤 함수를 정의하자.

\psi(t) =\exp( \int_{0}^t -F(\tau)d\tau) 

위의 $\psi(t)$를 linear ODE에 곱해보자. 그리고 정리하자.

$$ \psi(t) dx_t – \psi(t) F(t)x_t dt =\psi(t) u(t) dt$$

위의 식에서 좌변은 chain rule에 의하여 정리가 된다.

$$ d(\psi(t) x_t) =\psi(t) u(t) dt$$

이점을 이용하여 x_t를 구할 수 있다.

$$\int_{0}^t d(\psi(\tau) x_\tau) = \int_{0}^t \psi(\tau) u(\tau) d\tau $$

$$x_t = \psi(t)^{-1} x_0 + \psi(t)^{-1} \int_{0}^t \psi(\tau) u(\tau) d\tau$$

• 위의 linear ODE의 solution을 찾기 위하여 chain rule $ d(\psi(t) x_t) =\psi(t) u(t) dt$을 사용했다.
• calculus에서 chain rule은 아주 핵심적인 역할을 한다.
• 그러면 stochastic calculus에서 chain rule은?

Ito lemma

$\phi(x,t) \in \mathbb{R}$인 함수가 있다고 하자.

x_t 가 deterministic 한 경우에는

$$d\phi(x_t,t) = \frac{\partial \phi}{\partial t} dt +  \sum_{i} \frac{\partial \phi}{\partial x_i} dx_i$$

위와 같이 나온다. chain rule때문이기도 하고 아래의 taylor series에서 2차항 이후는 버리는 것이라 해석해도 된다.

$$\phi(x+\Delta x , t+\Delta t ) -\phi(x,t) = \frac{\partial \phi}{\partial t} \Delta t + \sum_{i} \frac{\partial \phi}{\partial x_i} \Delta x_i + \text{ (2차항 이후는 버림) }$$

그러면 $x_t$가 아래와 같은 sde에 의해 정해진 random process라면 $\phi(x_t,t)$는 어떻게 될까?

$$dx_t = f(x_t,t)dt + L(x_t,t)dw_t $$

아래와 같은 사실을 기억하자.

• dt dt = 0
• $dw_t dw_t^T = dt I$
• $dw_t = \sqrt{dt}$
• $dw_t dt = 0$

여기서 보면 2차항인 $dw_t dw_t= dt$이다. 이 사실 때문에 stochastic calculus에서 chain rule은 달라진다. 이점을 염두해두고 $\phi(x,t)$에 대해 taylor 근사를 하자.

$$\phi(x+\Delta x , t+\Delta t ) -\phi(x,t) = \frac{\partial \phi}{\partial t} \Delta t + \sum_{i} \frac{\partial \phi}{\partial x_i} \Delta x_i + 1/2 \sum_{i,j} \frac{\partial^2 \phi}{\partial x_i \partial x_j} \Delta x_i \Delta x_j + (\text{ 나머지 항들은 무시할만큼 작다})$$

이 사실을 이용하면 아래와 같은 식이 유도되고 이것을 Ito lemma라 부른다. stochastic calculus의 chain rule이라 보면 된다.

$$d\phi(x_t,t)  = \frac{\partial \phi}{\partial t} dt + \sum_{i} \frac{\partial \phi}{\partial x_i} dx_t^i + 1/2\sum_{i,j} \frac{\partial^2 \phi}{\partial x_i \partial x_j} dx_t^i dx_t^j \tag{Ito lemma} $$

혹은

$$ d\phi(x_t,t) = \frac{\partial \phi}{\partial t} dt + dx_t^T \frac{\partial \phi}{\partial x_i}  + 1/2dx_t^T  \frac{\partial^2 \phi}{\partial x_i \partial x_j} dx_t$$

현재 $\phi(x,t)$를 $\mathbb{R}$값은 1차원 함수로 생각했는데 $\phi=(\phi_1,…,\phi_d)^T$가 다차원함수여도 각 coordinate $\phi_i$ 마다 ito lemma 적용하면 된다.

여기서 저엉말 중요한 사실은

• dt dt = 0
• $dw_t dw_t^T = dt I$
• $dw_t = \sqrt{dt}$
• $dw_t dt = 0$

linear sde 풀이 방법

아래와 같은 linear SDE가 있을 때 Ito lemma를 이용하여 $x_t$의 solution을 구해보자.

dx_t = (F(t)x_t + u(t))dt+g(t) dw_t, 0\leq t \leq T, t=0 \to t=T , \tag{linear SDE}

ODE에 대해 하던것처럼 $\psi$를 정의하자.

$$\psi(t) = \exp(\int_{0}^t -F(\tau)d\tau) $$

$\psi(t)$를 (linear SDE)에 곱하고 정렬하자.

$$\psi(t)dx_t – \psi(t) F(t) x_t dt = \psi(t) u(t) dt + \psi(t) g(t) dw_t$$

 

이제 Ito lemma에 의해

$$d(\psi(t) x_t ) = \psi(t) dx_t -\psi(t) F(t)  x_t dt $$ 그러면

$$d(\psi(t) x_t) = \psi(t) u(t) dt + \psi(t) g(t) dw_t$$

$$\psi(t) x_t – x_0 = \int_{0}^t \psi(\tau) u(\tau) d\tau + \int_{0}^t \psi(\tau) g(\tau) dw_\tau$$

$$x_t = \psi(t)^{-1} x_0 +\psi(t)^{-1} \int_{0}^t \psi(\tau) u(\tau) d\tau+ \psi(t)^{-1}\int_{0}^t \psi(\tau) g(\tau) dw_\tau$$

 

perturbation kernel 구하기

perturbation kernel을 구하려다가 linear sde solution까지 왔다. linear sde solution을 다시 써보자.

x_t = \psi(t)^{-1} x_0 +\psi(t)^{-1} \int_{0}^t \psi(\tau) u(\tau) d\tau+ \psi(t)^{-1}\int_{0}^t \psi(\tau) g(\tau) dw_\tau

• linear sde solution으로 부터 perturbation kernel p_{t}(x_t|x_0) 구할 수 있다.
• x_0 가 given 되었다고 하면 \psi(t)^{-1} x_0 +\psi(t)^{-1} \int_{0}^t \psi(\tau) u(\tau) d\tau 상수 취급 가능하다.
• 문제는 마지막 term인 \psi(t)^{-1}\int_{0}^t \psi(\tau) g(\tau) dw_\tau 이다.
  • 결론적으로 이것은 gaussian 형태가 된다.

왜?

• dw_t \sim N(0, dt I)
• \psi(t)g(t) dw_t \sim N(0, \psi(t)\psi(t)^T g(t)^2 dt I)
• \int_{0}^t \psi(\tau) g(\tau) dw \tau  \sim N(0, \int_{0}^t \psi(\tau)\psi(\tau)^T g(\tau)^2  d\tau)

그렇다는 것은

p_t(x_t \mid x_0 ) = N\left(x_t\mid \psi(t)^{-1} x_0 +\psi(t)^{-1} \int_{0}^t \psi(\tau) u(\tau) d\tau, \psi(t)^{-1} \left[ \int_{0}^t \psi(\tau)\psi(\tau)^T g(\tau)^2 d\tau \right] (\psi(t)^{-1})^T \right)

Diffusion model 정리

Diffusion model 구성

• p_{data} : data distribution

Forward SDE

dx_t = f(x_t,t)dt+g(t)dw_t , 0\leq t \leq T, t=0 \to t=T, \text{ 데이터 } \to \text{ noise}

• f(x_t,t)=F(t)x_t+u(t) linear form이고 F(t) 를 scala function으로 잡는다.
• g(t) 는 scala function
• p_t 를 위의 forward SDE에 의해 정의되는 x_t 의 t 에서의 pdf
• p_0 = p_{data}
• p_T \approx N(\tilde{\mu}, \tilde{\sigma}^2 I ) 가 되길 바란다.

Reverse SDE

dx_t = [f(x_t,t)-g(t)^2 \nabla_{x_t} \log p_t(x_t) ]dt+g(t)dw_t , 0\leq t \leq T, t=T \to t=0, \text{ noise} \to \text{ 데이터}

• reverse SDE의 시작점이 forward SDE의 끝점인 $p_T$에서 시작한다면
  • reverse SDE로 표현되는 $x_t$의 분포 = forward SDE로 표현되는 $x_t$의 분포
• 따라서 reverse SDE를 이용하여 $x_T$로부터 $x_0$로 가면 data distribution에서 샘플이 만들어진다.
• \nabla_{x_t} \log p_t(x_t) 를 아는 것이 중요하다.

Diffusion model 학습

DSM loss 최적화

L_{DSM}(\theta) = E_{t \in U[0,T], x_0 \in p_{data}(x_0), x_t \in p_t(x_t \mid x_0)} \lVert \nabla_{x_t} \log p_t(x_t|x_0) - s_\theta(x_t,t) \rVert^2

• p_t(x_t|x_0)를 아는것이 중요하다.

Perturbation kernel

• p_t(x_t|x_0)를 특별히 perturbation kernel이라 부른다.
• 이거 구하는 방법은 잘 나와있다.

Diffusion model을 이용한 generation

reverse SDE에서 score \nabla_{x_t} \log p_t(x_t) 를 학습한 neural network s_\theta(x_t,t) 를 사용한다.

dx_t = [f(x_t,t)-g(t)^2 s_\theta (x_t,t)]dt+g(t)dw_t , 0\leq t \leq T, t=T \to t=0, \text{ noise} \to \text{ 데이터}

• time step 설정: T = t_N < t_{N-1} <t_{N-2} <...< t_0 =0
• 초기화: x_{t_N} \sim N(\tilde{\mu}, \tilde{\sigma}^2 I )
• sampling (Euler Maruyama method)
  • x_{t_{i-1}} = x_{t_{i}} +[f(x_{t_i},t_i)-g(t_i)^2 s_\theta(x_{t_i},t_i)](t_{i-1}-t_i)+g(t_i)\sqrt{|t_{i-1}-t_{i}|} z_{t_i} 여기서 z_{t_i} \sim N(0,I)

Probabillity Flow ODE (PF-ODE), Flow matching의 원조

Diffusion 세미나를 이대로 끝낼려 했었는데 중요한 probability flow ode (PF-ODE)를 놓쳤다. probability flow ODE는 forward SDE가 있으면 이것으로 부터 파생되는 미분방정식이다.

dx_t = [f(x_t,t) - 1/2 g(t)^2 \nabla_{x_t} \log p_t(x_t) ]dt, 0 \leq t \leq T

• 위의 식을 PF-ODE라고 부른다.
• PF-ODE로 표현되는 x_t 의 분포는 t=0 에서 p_{data} 로 시작한다면, PF-ODE, forward SDE, reverse SDE의 t시점에 분포는 동일
• PF-ODE로 표현되는 x_t 의 분포는 t=T 에서 p_T (forward SDE의 마지막점의 분포 혹은, reverse SDE의 시작점의 분포)로 시작한다면, PF-ODE, forward SDE, reverse SDE의 t시점에 분포는 동일
• 이 사실은 Fokker Planck Kolmogorov Equation을 이용해서 증명 된다.

 

PF-ODE를 이용한 데이터 샘플 생성방법

• time step 설정: T = t_N < t_{N-1} <t_{N-2} <...< t_0 =0
• 초기화: x_{t_N} \sim N(\tilde{\mu}, \tilde{\sigma}^2 I )
• sampling (Euler method)
  • x_{t_{i-1}} = x_{t_{i}} +[f(x_{t_i},t_i)-g(t_i)^2/2 s_\theta(x_{t_i},t_i)](t_{i-1}-t_i)

 

PF-ODE와 Flow matching의 관계

Flow matching은 아래의 PF-ODE에서 dt앞의 항 v_t(x_t) := f(x_t,t) - 1/2 g(t)^2 \nabla_{x_t} \log p_t(x_t)  를 근사하는 것을 목표로 한다.

dx_t = \underbrace{[f(x_t,t) - 1/2 g(t)^2 \nabla_{x_t} \log p_t(x_t) ]}_{v_t(x_t)} dt, 0 \leq t \leq T

• 디퓨전은 f,g를 정해두고 하지만
• Flow matching에서는 f,g를 모른다 치고 probabililty path p_t(x_t | x_0) 를 정하고 학습한다.