Diffusion model의 score 학습 방법에 대해 알아보자. 디퓨전 모델에 대한 설명은 이 글을 참고하면 되겠다. 잠시 중요한 개념을 복습하고 넘어가자.
Forward SDE
\begin{equation} dx_t = f(x_t,t) dt + g(t) dw_t , 0\leq t \leq T\end{equation}\tag{Forward SDE}(Forward SDE)는 데이터 분포를 따르는 random vector x_0 를 t=0에서부터 t=T까지 옮기며 노이즈를 추가해줘서 노이즈화 시키는 식이다. 여기서 drift term f와 diffusion term g를 설정하는 방식에 따라 노이즈를 추가하는 규칙이 달라진다.
Reverse SDE
\begin{equation} dx_t = \left[f(x_t,t)-g(t)^2 \nabla_{x_t} \log p_t(x_t) \right]dt + g(t) d\tilde{w}_t \end{equation} \tag{Reverse SDE}(Reverse SDE)는 (Forward SDE)의 반대과정이다. (Forward SDE)가 t=0부터 t=T까지 움직이며 데이터를 노이즈화 시키는 것에 반해 (Reverse SDE)는 t=T부터 t=0까지 움직이며 노이즈로부터 데이터 샘플을 생성하는 과정이다. 여기서 f와 g는 설정해두었떤 함수이고 Reverse SDE를 simulation 하기 위해서는 score \nabla_{x_t} \log p_t (x_t) 를 추정해야 한다.
Perturbation Kernel
Score \nabla_{x_t} \log p_t(x_t) 를 추정하기 위해서 학습하는 과정이 필요한데 score에 대한 ground truth를 알 수가 없다. score 대신 사용할 target을 만들어줄 필요가 있고 perturbation kernel p_{0t}(x_t | x_0) 를 사용한다. p_{0t}(x_t|x_0) 는 (Forward SDE)로 나타내는 stochastic process x_t 에 대하여 시작점 x_0 이 주어졌을 때 x_t 의 conditional probability density function이라고 보면 된다.
Denoising Score Matching (DSM) loss
score \nabla_{x_t} \log p_t(x_t) 를 추정하는 neural network를 score model s_\theta(x_t,t) 이라고 하자. 아래의 DSM loss 를 이용하여 score를 추정한다.
\begin{equation} \mathcal{L}_{DSM}(\theta) = E_{t, x_0, p_t(x_t|x_0)} \lambda (t) \lVert s_\theta (x_t, t) - \nabla_{x_t} \log p_{0t} (x_t | x_0) \rVert^2 \end{equation}\tag{DSM loss}여기서 x_0 는 데이터의 분포에서 나오는 샘플이라 보면 되고 t 는 [0,T]사이의 uniform distribution에서 sampling 된다. 그리고 \lambda(t) 는 weight이다. 위와 같이 (DSM loss)를 이용해서 score를 추정한다.
궁금증
- (DSM loss)는 과연 score를 추정할까? perturbation kernel의 log derivative를 추정하는 꼴인데 score추정하는 거 맞어? ( \nabla_{x_t} \log p_t (x_t) 를 추정하고 싶은데 \nabla_{x_t} \log p_{0t} (x_t |x_0) 를 true값으로 사용중) (이것에 대한 대답은 이 글에서 확인)
- perturbation kernel closed form 구할 수 있을까? Diffusion model에서 perturbation kernel이 closed form을 가질 충분조건